From Uniform to Learned Graph Priors: Diffusion for Structure Discovery

2026-06-10

AI Summary

Existing neural relational inference methods rely on simplified independent-edge priors, leading to ambiguous and unreliable structure discovery. This work proposes Diff-prior, which, for the first time, leverages a diffusion mechanism not for generating graph structures but for calibrating them. By applying a learnable denoising refinement to the edge posteriors produced by the encoder, Diff-prior constructs a more structured and adaptive prior that better aligns with real-world systems. Embedded within a variational inference framework, this approach significantly enhances the determinacy and accuracy of structural inference. Empirical results demonstrate consistent improvements across multiple standard benchmarks, where Diff-prior markedly sharpens edge posteriors and boosts their reliability for diverse neural relational inference architectures.

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.

Problem

• graph priors
• neural relational inference
• structure discovery
• edge posteriors
• uniform prior

Innovation

• diffusion prior
• neural relational inference
• structured calibration
• adaptive graph prior
• edge posterior refinement