DiffTSP: A Conditional Discrete Diffusion Model for Knowledge Graph Triple Set Prediction

The gist One Pass for All: A Discrete Diffusion Model for Knowledge Graph Triple Set Prediction proposes DiffTSP, a novel discrete diffusion model that treats Triple Set Prediction (TSP) as a generative task to capture inter-triple dependencies in knowledge graphs.

How it works

1. DiffTSP operates via a unified ”One-Pass Generation” process where all missing triples are generated holistically rather than as a sequence of independent predictions The core idea of DiffTSP is to capture the intertriple dependencies that are critical for generating a coherent triple set.

2. To learn the dependencies effectively, the model designs a structure-aware denoising network that consists of two synergistic modules: a Relational Context Encoder (RCE) and a Relational Graph Diffusion Transformer (RelDiT).

3. The RCE performs relation-guided message passing to capture local structural dependencies among entities.

4. Built upon the entity representations, the RelDiT further models global structural dependencies via relational attention mechanisms.

5. The forward process systematically corrupts a graph by progressively masking its relational edges.

6. The structure-aware denoising network is trained to reverse this corruption by predicting the original clean graph given a noisy graph and the incomplete graph as a condition.

7. Finally, an iterative sampling process is employed to generate a complete and internally consistent set of triples.

Relation-Balanced Learning Task Generation

The training involves a meta-learning framework to explicitly train DiffTSP to generate query triple sets from support sets For each task, the set of triples T is partitioned into a support set Ts and a query set Tq.

The model uses a relation-balanced split strategy to ensure that the probability distribution of relations in Ts and Tq mirrors that of the original graph T. This ensures that DiffTSP is trained in a variety of combinations of support and query graphs, thereby achieving robust graph generalization.

Structure-Aware Denoising Network

The structure-aware denoising network takes the noisy query graph Gqt, the support graph Gs, and the timestep t as input.

This network consists of a relational context encoder and a relational graph diffusion transformer.

The RCE module maintains a learnable relation embedding matrix R where this matrix stores embeddings for each relation type and its inverse, allowing the model to differentiate edge directions.

The initial feature vector hi(0)i for an entity i is computed by averaging the embeddings of all its adjacent relations.

The RelDiT module models global structural dependencies via relational attention mechanisms, incorporating a relation-aware attention bias Bij that reflects the relational dependency between the two entities.

Reverse Process and Training Objective

The structure-aware denoising network fθ is trained to reverse the forward diffusion process, i.e., predicting the clean graph Gq0 from its noisy version Gqt under the condition Gs.

To learn this, a variational lower bound (VLB) of the log-likelihood is optimized instead of directly optimizing the conditional log-likelihood.

The simplified training objective Lsimple is derived by approximating the KL divergence terms with Binary CrossEntropy (BCEw) loss.

Theoretical Analysis

Theorem 1 states that if the learned denoising network pθ(Gt-1 Gt) perfectly matches the true reverse distribution q(Gt-1 Gt), then the distribution of generated graphs pθ(G0) is identical to the true data distribution, i.e., pθ(G0) = pdata(G0).

Theorem 2 shows that minimizing the VLB provably tightens an upper bound on the probability of generating inconsistent triples.

Experiment Results

DiffTSP outperforms all baselines in terms of FT SP which is highlighted in gray, demonstrating its superiority.

In Wiki79k and Wiki143k, DiffTSP achieves SOTA performance in terms of both RS-POWA and CWA metrics.

In CFamily, our model also achieves SOTA performance on JP recision and FT SP.

Case studies highlight that DiffTSP provides consistent predictions for entity 35, avoiding internal contradictions present in HAKE and AstarNet.

Conclusion

In this work, we treat TSP as a generative task, and propose a novel discrete diffusion model named DiffTSP which could effectively capture the interdependence among the predicted triples. In addition, we design a structure-aware denoising network to combine relation-guided message passing with global structural attention for generating knowledge graphs of high quality. Extensive experiments on multiple datasets showcase the superior performance of DiffTSP compared to strong baselines.

1. No-Pass Knowledge Graph Completion via DiffTSP: The system can infer missing triples in a single pass by treating Triple Set Prediction (TSP) as a generative task, progressively adding noise to the KG through a discrete diffusion process and then reversing it to gradually recover the complete KG conditioned on the incomplete graph.

2. Structure-Aware Denoising Network Integration: The architecture features two synergistic modules, a Relational Context Encoder (RCE) and a Relational Graph Diffusion Transformer (RelDiT), which allows the model to capture both local relational structures via RCE and global structural dependencies via RelDiT.

3. Dependency-Preserving Generation: The framework is designed to ensure consistency by learning the joint distribution of triples, as it aims to capture the intertriple dependencies that are critical for generating a coherent triple set, contrasting with existing methods that score triples in isolation and fail to capture crucial inter-triple dependencies.

4. Robust Training via Relation-Balanced Learning: The model is trained using a meta-learning framework where it generates multiple training tasks, ensuring the probability distribution of relations in T s and T q mirrors that of the original graph T through a relation-balanced split strategy.

5. Contradiction Avoidance During Sampling: During the iterative sampling process, the model ensures logical consistency by constructing the next step as Gt−1 = (m ⊙ G known t−1) + ((1 − m) ⊙ G unknown t−1), and explicitly sampling to ensure that generated edges are not present in the support graph GS.

Triple set prediction (TSP) aims to discover missing facts from an incomplete knowledge graph without predefined queries. Existing approaches typically construct the predicted set through triple-level scoring, sometimes preceded by entity pair filtering that can exclude valid missing facts. We propose DiffTSP, a conditional discrete diffusion model that formulates missing fact discovery as conditional subgraph generation. In each subgraph, DiffTSP directly predicts the missing relational structure without a separate entity pair selection stage. The model learns to generate masked target relation states at different corruption levels, with the observed support subgraph retained as conditioning context. Our structure-aware denoising network integrates a Relational Context Encoder, which captures local relational patterns through relation-guided message passing, and a Relational Graph Diffusion Transformer, which models broader structural interactions through relational attention. During generation, partially generated relations can also serve as provisional context for subsequent predictions. Experiments on three public datasets demonstrate improved TSP performance over strong baselines. Additional analyses show lower conflict rates in the predicted triple sets relative to the evaluated baselines. Code: https://anonymous.4open.science/r/DiffTSP-0F92/.