One Pass for All: A Discrete Diffusion Model for Knowledge Graph Triple Set Prediction

2026-04-20

AI Summary

This work addresses the challenge of inconsistent predictions in knowledge graph completion by proposing DiffTSP, a novel approach that formulates triple set prediction as a graph generation problem. Unlike conventional methods that generate triples independently and thus fail to capture inter-triple dependencies, DiffTSP introduces a discrete diffusion mechanism for the first time in this task. It perturbs the input graph by masking relational edges and recovers the complete graph structure in a single forward pass during the reverse denoising process. The model employs a structure-aware denoising network that integrates a relational context encoder with a relational graph diffusion Transformer, enabling conditional generation of globally consistent triple sets. Experimental results demonstrate that DiffTSP achieves state-of-the-art performance, significantly outperforming existing methods across three standard benchmarks.

Technology Category

• Data Mining & Knowledge Management: Linked Open Data, Knowledge Graphs & KB Completion
• Machine Learning: Graph-based Machine Learning
• Knowledge Representation and Reasoning: Computational Complexity of Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Knowledge Graphs (KGs) are composed of triples, and the goal of Knowledge Graph Completion (KGC) is to infer the missing factual triples. Traditional KGC tasks predict missing elements in a triple given one or two of its elements. As a more realistic task, the Triple Set Prediction (TSP) task aims to infer the set of missing triples conditioned only on the observed knowledge graph, without assuming any partial information about the missing triples. Existing TSP methods predict the set of missing triples in a triple-by-triple manner, falling short in capturing the dependencies among the predicted triples to ensure consistency. To address this issue, we propose a novel discrete diffusion model termed DiffTSP that treats TSP as a generative task. DiffTSP progressively adds noise to the KG through a discrete diffusion process, achieved by masking relational edges. The reverse process then gradually recovers the complete KG conditioned on the incomplete graph. To this end, we design a structure-aware denoising network that integrates a relational context encoder with a relational graph diffusion transformer for knowledge graph generation. DiffTSP can generate the complete set of triples in a one-pass manner while ensuring the dependencies among the predicted triples. Our approach achieves state-of-the-art performance on three public datasets. Code: https://github.com/ADMIS-TONGJI/DiffTSP.

Problem

• Knowledge Graph Completion
• Triple Set Prediction
• Discrete Diffusion Model
• Generative Task
• Relational Dependencies

Innovation

• discrete diffusion model
• triple set prediction
• knowledge graph completion
• structure-aware denoising
• one-pass generation