From Uniform to Learned Graph Priors: Diffusion for Structure Discovery

Abstract: Neural relational inference (NRI) methods discover interaction graphs from trajectories through variational reasoning on discrete potential edges. However, these methods typically rely on oversimplified, factorized graph priors. Such priors, typically nearing uniform distributions, treat edges as independent entities. This systemic misalignment does not match the real-world systems and yields diffuse and indecisive edge posteriors limiting the reliability of structural discovery. To address this, we propose \textit{Diff-prior}, a diffusion-parameterized adaptive prior used to calibrate latent graph distribution rather than generate graphs. Our core insight is to reframe prior integration as a learnable denoising-style calibration that organizes scattered, uncertain edge posteriors into a more reliable overall structure which can be trained by the diffusion model. Diff-prior learns an adaptive structure prior that performs structured calibration on the edge posteriors during inference, guiding it towards a distribution closer to the underlying structure. The diff-prior operates before structural sampling and acts as a denoising calibrator directly on the encoder edge distribution, which provides a generic training paradigm over structured variables. Experiments on standard benchmarks validated our framework, and the results indicate that Diff-prior improves the performance of structure inference and generates more decisive edge posteriors across multiple NRI-family architectures. The code is available on https://github.com/Hardy158118/Diffprior.

• The paper introduces Diff-Prior, a diffusion-based learned structural prior that calibrates encoder edge logits to yield decisively inferred graph structures.
• It employs a denoising diffusion probabilistic model to refine ambiguous edge posteriors, consistently improving AUROC by up to 6.43 points across backbone models.
• The method enhances structural calibration, reduces uncertainty metrics, and maintains compatibility with existing Neural Relational Inference architectures.

Diff-prior: Diffusion-based Learned Structural Priors for Neural Relational Inference

Introduction

The paper "From Uniform to Learned Graph Priors: Diffusion for Structure Discovery" (2606.11831) addresses a critical limitation in Neural Relational Inference (NRI) methods for structure discovery from dynamic trajectory data: the use of oversimplified, factorized (often uniform) graph priors. These independent-edge priors disregard the structured regularities present in real-world systems, resulting in highly uncertain and weakly decisive edge posterior distributions. This paper introduces Diff-prior, a non-factorized, diffusion-parameterized structural prior that adaptively calibrates edge posteriors before discretization, enhancing the decisiveness and reliability of inferred graphs without modifying the encoder or decoder architectures in NRI-family models.

Background and Motivation

Neural Relational Inference (NRI) methods encode observed state trajectories to infer edge-type distributions representing the latent interaction graph, and subsequently decode predicted future dynamics conditioned on graph samples. Traditional NRI and its extensions often deploy a uniform categorical prior on edges, assuming independence between edges due to scalability considerations. However, this abstraction is typically misaligned with realistic structures where phenomena such as sparsity, modularity, and motifs prevail, leading to diffuse and poorly calibrated edge posteriors. Weak priors fail to bias the inference process toward consistent or structured configurations, especially under ambiguous or information-poor regimes.

The key observation motivating this work is that the quality and calibration of posterior edge distributions are substantially influenced by the prior's ability to encode structural correlations. The inadequacy of independent priors becomes exacerbated under complex system dynamics, where trajectory data alone cannot sufficiently disambiguate the underlying interaction topology.

Methodology: Diff-prior as a Latent Diffusion-Based Calibrator

The Diff-prior framework leverages a denoising diffusion probabilistic model (DDPM) to parameterize a global, learnable prior directly on the encoder-produced edge logits ($Z_0$), prior to the final edge discretization step. Rather than acting as a discrete graph generator, the diffusion model here serves as a calibration transform: it takes ambiguous or noisy logit configurations and iteratively refines them towards configurations consistent with learned structural motifs.

The method consists of the following principal components:

1. Encoder: Given observed trajectories $\mathbf{X}$, the encoder infers edge-type logits for all node pairs, forming a latent representation of the candidate edge types for the graph.
2. Latent Diffusion Calibration:
   • The encoder logits $Z_0$ are subjected to a forward Gaussian noising process, producing diffused versions $Z_s$ at various time steps.
   • A learned denoiser $\epsilon_\theta(Z_s, s)$ reconstructs the original clean logits from the noisy versions, optimized with a weighted mean-squared error objective between injected and predicted noise.
   • The denoiser is typically instantiated as a lightweight MLP, but transformer-based denoisers are also considered for explicit modeling of non-factorized dependencies.
   • Critically, the refined logits $Z_0^\mathrm{ref}$—computed via a residual update informed by the denoiser's predictions—are passed forward for edge sampling.

     Figure 2: Diff-prior workflow—diffusion denoising calibrates the encoder logits, enhancing edge posterior sharpness and structural coherence.

3. Decoder: The calibrated logits are discretized via the Gumbel-Softmax estimator, yielding a sampled interaction graph $\mathbf{G}$. This graph is used by a message-passing GNN to decode the predicted dynamics, allowing full differentiability and joint end-to-end training.
4. Objective and Optimization: The total training objective combines trajectory reconstruction loss with a diffusion denoising regularizer, upper-bounding the KL divergence between the approximate posterior and the learned diffusion prior via a tractable joint-augmentation approach. A single-timestep Monte Carlo strategy yields efficient, unbiased optimization.

Experimental Results

Benchmarks and Settings: The methodology is evaluated on the StructInfer suite, using standard binary link-inference tasks on synthetic and real datasets spanning various network families and dynamics (e.g., Spring, Netsims). Diff-prior is implemented as a plug-in module across multiple NRI-style backbones (NRI, ACD, MPM), with no modifications to the base architectures.

Structure Inference Performance: In both Spring (SP) and Netsims (NS) dynamics, Diff-prior consistently improves AUROC for structure recovery across all tested backbone models and network families. The performance gain is most prominent under complex dynamics (e.g., NS), where average AUROC increases by up to 6.43 points (MPM backbone).

Posterior Calibration: Posterior confidence diagnostics demonstrate that Diff-prior produces sharply decisive edge posteriors, as visualized by a significant rightward shift in the confidence complementary CDFs. This effect is consistently observed across all network-dynamics pairs.

Figure 1: Confidence (complementary CDF) distributions of edge posteriors under different priors, showing Diff-prior sharply concentrates posterior mass near confident (0 or 1) predictions.

Uncertainty Metrics: Diff-prior reduces both average edge entropy and expected calibration error (ECE) in nearly all tested settings, indicating not only greater confidence but also improved calibration. Importantly, Diff-prior maintains superior calibration even for the most confident/low-entropy predictions.

Structural Motif Recovery: Beyond edge-level accuracy, Diff-prior leads to consistently improved recovery of higher-order graph statistics (density, clustering, modularity, average degree) compared to independent priors. This indicates successful internalization of global topological patterns through the learned prior.

Robustness: Systematic robustness experiments, involving trajectory noise, truncated time sequences, and random missingness, demonstrate that Diff-prior maintains positive gains in entropy, ECE, and AUROC across all perturbations, reflecting enhanced resilience to signal degradation.

Ablation Studies: The calibration intensity parameter ($\gamma$) is essential—removing or diminishing this residual step sharply worsens calibration. Control experiments with parameter-matched MLP priors (removing diffusion) fail to reproduce the gains, isolating the improvement to the diffusion-based calibration mechanism.

Theoretical and Practical Implications

This work fundamentally advances the regularization of structural uncertainty in deep generative models of interacting systems by moving away from independent, edge-wise priors toward fully non-factorized, learnable graph priors. The proposed denoising-diffusion calibration paradigm preserves plug-and-play modularity: it can be incorporated into any NRI-style inference pipeline with negligible additional training overhead and does not increase model complexity substantively.

Practical implications include:

• Improved decisiveness and reliability in inferred graph structures, benefitting applications where scientific interpretability and posterior confidence are critical (e.g., system identification, multi-agent analysis, causality discovery).
• Robustness to ambiguous or limited observations, reducing dependence on carefully calibrated regularizers or hand-crafted priors.
• Backward-compatibility with existing NRI-family encoders and decoders, facilitating adoption in downstream structure discovery pipelines.

Theoretical implications extend to the broader understanding of posterior calibration in models with combinatorial latent states. By demonstrating that global adaptive priors can be efficiently learned via diffusion in high-dimensional latent spaces, the methods suggest future directions for multi-scale and hierarchical priors over more abstract graph characteristics.

Future Directions

Potential avenues for further research include:

• Scaling to larger, sparser, or more complex graph structures and real-world data modalities.
• Extending adaptive diffusion priors to multi-relational or continuous-valued edge inference.
• Incorporating domain-specific physical, causal, or hierarchical constraints into the diffusion prior.
• Developing more sophisticated diagnostics for posterior uncertainty and reliability in combinatorial latent spaces.

Conclusion

Diff-prior introduces a principled, diffusion-based approach to learning non-factorized, adaptive structural priors for graph structure discovery via NRI-style models. This framework systematically improves inference performance, decisiveness, and reliability of latent graph estimates, especially under ambiguous or information-constrained regimes. The methodology is fully modular and scales to diverse backbone architectures, marking a substantive step toward robust and interpretable relational modeling in complex dynamical systems.