Technical Summary: Generating Graph-like Rules for Knowledge Graph Reasoning via Diffusion Models

1. Problem Statement

Knowledge Graph (KG) reasoning relies heavily on logical rules to infer missing facts with interpretability. However, existing rule mining methods face two primary limitations:

1. Structural Restriction: Current approaches predominantly focus on chain-like rules (linear sequences of relations). While efficient, these structures fail to capture richer relational patterns such as cycles, branches, and overlapping constraints found in complex KG schemas. This leads to ambiguous predictions in scenarios requiring joint multi-relational constraints.
2. Computational Bottleneck: Extending rule discovery to general graph-like rules (where the rule body forms a directed graph) results in a combinatorial explosion of the search space ($O((L|R|)^L)$ compared to $O(|R|^L)$ for chains). Traditional search-based paradigms struggle with this exponential complexity.
3. Generative Model Misalignment: While diffusion models have succeeded in other domains, they cannot be directly applied to rule mining. Standard diffusion training optimizes differentiable reconstruction losses, whereas rule quality is evaluated using non-differentiable discrete metrics (e.g., confidence, coverage, PCA confidence). This creates a mismatch between the training objective and the semantic goal of generating high-quality rules.

2. Methodology: The GRiD Framework

The authors propose GRiD (Generating Graph-like Rules with Diffusion), a framework that reformulates rule discovery as a conditional discrete generative process. GRiD employs a two-phase training strategy to overcome the challenges of discrete search and non-differentiable metrics.

2.1 Discrete Diffusion Formulation

GRiD models the rule body as a directed adjacency matrix $A \in \{0, 1, \dots, |R|, \emptyset\}^{S \times S}$, where $S$ is the maximum number of nodes and $R$ is the relation set.

• Conditioning: The generation is conditioned on the target head relation $r_h$ and the diffusion timestep $t$.
• Forward Process: A clean rule matrix $A_0$ is progressively corrupted into noise $A_T$ using a categorical Markov transition process (D3PM framework), where edges are either preserved or resampled uniformly from all relation types.
• Reverse Process: A Graph Transformer denoising network $f_\theta$ learns to reconstruct the clean matrix from noise, conditioned on the target relation and timestep via Feature-wise Linear Modulation (FiLM).

2.2 Two-Stage Training Strategy

To address the non-differentiability of rule quality metrics, GRiD integrates Supervised Learning (SL) and Reinforcement Learning (RL):

• Phase A: Supervised Pre-training (SL):

  • Goal: To learn structural priors and provide a warm start.
  • Method: The model is trained to reconstruct subgraphs sampled from the KG meta-graph. The objective combines a weighted cross-entropy loss for structural reconstruction with an auxiliary loss to predict a scalar quality score (based on coverage and confidence). This aligns the learned representations with rule-quality signals without requiring direct optimization of discrete metrics.

• Phase B: Reinforcement Learning (RL) Fine-tuning:

  • Goal: To directly optimize for non-differentiable rule quality metrics.
  • Method: The denoising process is framed as a Markov Decision Process (MDP) where the denoising network acts as a stochastic policy.
  • Reward Signal: Upon generating a complete rule, a reward is computed as the geometric mean of Coverage, Confidence, and PCA Confidence.
  • Optimization: The model is fine-tuned using the REINFORCE algorithm with a baseline to minimize variance, maximizing the expected episode-level reward. Entropy regularization is used to encourage exploration.

2.3 Inference Pipeline

During inference, the model performs iterative conditional denoising starting from a noisy adjacency matrix to generate a discrete graph-like rule body. The resulting graph is parsed into a logical rule (identifying head and tail variables) and filtered against KG structural constraints.

3. Key Contributions

1. Framework Innovation: The proposal of GRiD, which shifts rule discovery from iterative combinatorial search to a conditional discrete diffusion process, enabling the generation of complex graph-like structures.
2. Training Strategy: The introduction of a two-stage training pipeline that combines SL pre-training for structural pattern recognition with RL fine-tuning to align the generative process with non-differentiable rule-quality metrics.
3. Empirical Validation: Extensive experiments on six benchmark datasets (including FB15K-237, YAGO3-10, and domain-specific datasets) demonstrating that GRiD achieves competitive performance in Knowledge Graph Completion (KGC) tasks.
4. Insight on Rule Types: Ablation studies confirming that graph-like rules complement chain-like rules. While chain-like rules capture dominant logical regularities, graph-like rules provide necessary disambiguation in complex schemas, particularly when fusion weights are low (approx. 0.1–0.2).

4. Experimental Results

• Performance: GRiD outperforms most existing rule-based baselines (e.g., AnyBURL, AMIE, Neural-LP, RNNLogic) and remains competitive with embedding-based and GNN-based methods across six datasets. Notably, it achieves the best or second-best results on large-scale KGs where traditional rule-based methods often face scalability issues.
• Ablation Findings:
  • Graph vs. Chain: Combining graph-like and chain-like rules consistently outperforms using either type alone, validating their complementary nature.
  • RL Necessity: RL fine-tuning yields approximately a 10% improvement in metrics over SL-only pre-training, proving the necessity of direct optimization for rule quality.
  • Graph Size: A maximum graph size of $S=6$ nodes is sufficient for complex KGs, balancing expressiveness with computational efficiency.
  • Fusion Weight: Optimal performance is achieved when graph-like rules contribute a small, controlled proportion (10-20%) to the final scoring, acting as a complementary signal rather than the dominant reasoning paradigm.

5. Significance and Claims

The paper claims that GRiD addresses the fundamental bottleneck of searching for graph-like rules by leveraging the generative capabilities of diffusion models. By successfully bridging the gap between discrete rule quality metrics and continuous generative modeling through RL fine-tuning, the framework enables the efficient discovery of graph-like rules that capture complex relational patterns (cycles, branches) previously inaccessible to standard mining methods.

The authors emphasize that while graph-like rules are not a replacement for chain-like rules, they serve as a crucial complementary mechanism to resolve ambiguity in KGs with rich, complex schemas. The work suggests a shift from exhaustive search to learnable optimization procedures for graph-structured combinatorial problems. Future work is noted to explore generating rules directly from underlying KG distributions to mitigate sampling bias and developing more precise structural evaluation metrics.