Diffusion Model with Selective Attention for Temporal Knowledge Graph Reasoning

Abstract

Temporal knowledge graph reasoning aims to predict missing entities at future time steps, and as a critical task, it has attracted widespread attention in recent years due to its impressive ability to capture historical correlations and forecast future events. Although existing approaches, such as graph learning and logic rules, have partially addressed this problem, they still face limitations in modeling the uncertainty of future events—especially when predicting rare or unseen facts. To address these challenges, we propose a diffusion model based on a selective attention mechanism (DMSA) for temporal knowledge graph reasoning. In our method, the encoder incorporates selective attention to emphasize key information, while the diffusion module introduces noise to enhance the model’s capability to predict unseen events. By integrating selective attention with the diffusion module, our model improves both its memory and its ability to predict future, unseen events. Experimental results on five public datasets demonstrate that our proposed model achieves state-of-the-art performance across multiple evaluation metrics.