모든 것을 한 번에: 지식 그래프 삼중항 집합 예측을 위한 이산 확산 모델

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

지식 그래프(KG)는 삼중항으로 구성되며, 지식 그래프 완성(KGC)의 목표는 누락된 사실 기반의 삼중항을 추론하는 것입니다. 기존의 KGC 작업은 삼중항의 일부 요소가 주어졌을 때 누락된 요소를 예측합니다. 보다 현실적인 작업인 삼중항 집합 예측(TSP) 작업은 관찰된 지식 그래프에만 의존하여 누락된 삼중항 집합을 추론하며, 누락된 삼중항에 대한 어떠한 부분 정보도 가정하지 않습니다. 기존의 TSP 방법은 삼중항 단위로 누락된 삼중항 집합을 예측하며, 예측된 삼중항 간의 의존성을 파악하여 일관성을 보장하는 데 한계가 있습니다. 이러한 문제를 해결하기 위해, 우리는 TSP를 생성적 작업으로 취급하는 새로운 이산 확산 모델인 DiffTSP를 제안합니다. DiffTSP는 관계 기반 연결을 마스킹하여 이산 확산 과정을 통해 KG에 점진적으로 노이즈를 추가합니다. 역 과정에서는 불완전한 그래프를 조건으로 하여 완전한 KG를 점진적으로 복원합니다. 이를 위해, 우리는 지식 그래프 생성을 위해 관계 컨텍스트 인코더와 관계 그래프 확산 트랜스포머를 통합한 구조 인식 디노이징 네트워크를 설계했습니다. DiffTSP는 예측된 삼중항 간의 의존성을 보장하면서 한 번의 과정으로 완전한 삼중항 집합을 생성할 수 있습니다. 우리의 접근 방식은 세 가지 공개 데이터 세트에서 최첨단 성능을 달성했습니다. 코드: https://github.com/ADMIS-TONGJI/DiffTSP.

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.