Information-geometric adaptive sampling for graph diffusion

Abstract: Standard diffusion models for graph generation typically rely on uniform time-stepping, an approach that overlooks the non-homogeneous dynamics of distributional evolution on complex manifolds. In this paper, we present an information-geometric framework that reinterprets the diffusion sampling trajectory as a parametric curve on a Riemannian manifold. Our key observation is that the Fisher-Rao metric provides a principled measure of the intrinsic distance. By analyzing this metric, we derive the Drift Variation Score (DVS), a geometry-aware indicator that quantifies the instantaneous rate of distributional change. Unlike prior heuristic-based adaptive samplers, our DVS solver enforces a constant informational speed on the statistical manifold, automatically maintaining a uniform rate of distributional change along the sampling trajectory. This equal arc-length strategy ensures that each discretization step contributes equally to the information speed. Theoretical analysis verifies that DVS characterizes the local stiffness of the sampling dynamics in the Fisher-Rao sense. Experimental results on molecule and social network generation show that DVS significantly improves structural fidelity and sampling efficiency. Code is at https://github.com/kunzhan/DVS

• The paper introduces the Drift Variation Score (DVS), a training-free metric based on Fisher-Rao geometry that quantifies the rate of distribution evolution.
• It adaptively adjusts the sampling step size by enforcing constant informational progress, dynamically refining steps in high-curvature regions and expanding them in stable areas.
• Empirical results demonstrate that DVS-driven sampling improves structural fidelity and convergence speed, achieving up to 21.3% faster convergence compared to fixed-step techniques.

Information-Geometric Adaptive Sampling for Graph Diffusion

Core Contributions and Theoretical Framework

This paper introduces a principled information-geometric framework for adaptive sampling in graph diffusion models, specifically addressing the inefficiencies and inaccuracies arising from standard uniform and heuristic time-stepping schedules. The central innovation is the Drift Variation Score (DVS), an interpretable, training-free metric derived from the Fisher-Rao geometry of the statistical manifold. DVS quantifies the instantaneous rate of distributional evolution as the reverse-time SDE transitions from a noise prior to the data distribution. The adaptive sampler enforces constant informational progress—formally, uniform arc-length increments under the Fisher-Rao metric—along the sampling trajectory, refining step sizes dynamically in high-curvature regimes and expanding them in stable regions.

This geometric approach contrasts sharply with previous static and error-based schedules, which are either blind to manifold curvature or rely on arbitrary state-space heuristics. Theoretical analysis shows DVS reliably characterizes local stiffness, preventing numerical instability and facilitating efficient allocation of computation. The framework generalizes to multi-component graph states (nodes and edges), employing exponential moving average smoothing and a power-law scaling rule to maintain robust adaptation.

Methodological Innovations

Fisher-Rao Metric and Drift Variation Score

The sampling trajectory is recast as a parametric curve on a statistical manifold, where the Fisher-Rao metric provides an intrinsic, coordinate-invariant measure of statistical distance. The line element:

$$\mathrm{d}s^2 = \frac{\mathrm{d}t}{g_t^2} \|\mathrm{d}\bm{f}_t\|_2^2$$

is dominated by drift variation, as shown by numerical scale analysis in Figure 1.

Figure 1: Numerical scale analysis of information variation components on QM9 using the GruM model. The drift component $V_f$ (red solid line) represents the data-dependent variation of the transition distribution, which dominates the noise component $V_g$ (blue dashed line) by over six orders of magnitude on average.

DVS is defined discretely as:

$$V_k = \frac{\| \bm{f}(\bm{x}_k, t_k) - \bm{f}(\bm{x}_{k-1}, t_{k-1}) \|_2^2}{g_{t_k}^2}$$

The step size is governed by inversely scaling with DVS, guaranteeing that each update covers approximately equal statistical distance—shrinking in stiff regions and expanding in flat regions.

Integration in Graph Diffusion and Coupled Dynamics

For graph-structured data, node and adjacency transitions may exhibit unsynchronized denoising rates. The multi-component DVS monitors both modalities, aggregating variation states via EMA smoothing, applying component-wise adaptation, and aggregating final step sizes by bottlenecking. The resulting algorithm is inherently geometry-aware and robust to high-dimensional structural transitions.

Empirical Results and Numerical Analysis

Extensive benchmarks on molecular (QM9, ZINC250k) and synthetic/real-world graphs (Planar, SBM, Ego-small) consistently demonstrate that DVS-driven sampling improves structural fidelity, distributional alignment, and sampling efficiency compared to fixed-step and quadratic schedules.

Key quantitative results:

• On QM9 (GruM model), DVS-Euler achieves a Validity of 99.53%, FCD of 0.095, and NSPDK of 0.0002, outperforming both baselines with a 21.3% reduction in required steps. This improvement is statistically significant across multiple independent runs.
• DVS-driven methods often surpass heuristic schedules and even higher-order integrators (e.g., Fixed-Step Heun), indicating that geometric allocation of step-size is more impactful than local error order for graph-structured data.
• On Planar and SBM datasets, DVS reduces maximum mean discrepancy (MMD) on spectral and orbit metrics, underscoring superior preservation of both global topology and local motifs.

The convergence pattern of FCD shown in Figure 2 reveals faster and smoother trajectory for DVS, validating constant informational speed.

Figure 2: Convergence behavior of FCD during sampling on QM9 using GRUM. By dynamically sensing the drift variation, our method converges 21.3\% faster than the fixed-step baseline while attaining a higher distributional fidelity (lower FCD). The inset confirms that the DVS-driven trajectory maintains smoother convergence during the final phase of molecule generation.

The informational increment ($\Delta s^2$) progression, visualized in Figure 3, demonstrates that fixed-step schedules accumulate negligible progress in early phases but are forced into unstable jumps during late-stage crystallization, whereas DVS maintains near-uniform arc-length progression.

Figure 3: Evolution of the information-geometric increment ($\Delta s^2$) on QM9 using GRUM. We compare the information progress of the fixed-step Euler-Maruyama sampler and the DVS-driven Euler-Maruyama sampler. For the fixed-step sampler, the information increment remains small at early stages and increases sharply near the end, whereas the adaptive sampler maintains an approximately equal arc-length progression throughout sampling.

Sample visualizations provide further qualitative confirmation: in molecular graphs (Figure 4), DVS resolves complex chemical motifs and multi-ring systems, while in structural evolution snapshots (Figures 6 and 7), DVS more accurately crystallizes sparse planar skeletons and community clusters.

Figure 4: Visualizations of generated molecules on QM9 (top row) and ZINC250k (bottom row). The samples are generated using the GruM model with the DVS-Euler--Maruyama sampler. The top row shows organic molecules from QM9 with realistic geometric configurations. The bottom row displays complex drug-like molecules from ZINC250k, successfully capturing multi-ring systems and diverse heteroatoms (e.g., O, N, S, Cl).

Figure 5: Snapshots of structural evolution on the Planar dataset. The top row represents the standard fixed-step Euler sampler, while the bottom row represents our DVS-Euler--Maruyama sampler. As $t$ increases, the DVS-driven sampler resolves the sparse planar skeleton more cleanly, especially in the late stages ($t \ge 0.6$).

Figure 6: Snapshots of structural evolution on the SBM dataset. The top row (fixed-step) and bottom row (DVS-Euler--Maruyama) both start from a dense noise state. However, the DVS-driven sampler achieves clearer community separation and fewer spurious inter-cluster edges by the final step ($t=1.0$).

Ablation studies on the aggregation scaling factor $\gamma$ confirm that DVS remains efficient and robust, achieving optimal FCD with fewer steps compared to baselines.

Practical and Theoretical Implications

DVS-driven adaptive sampling is robust, training-free, and easily integrable with existing graph diffusion pipelines. It yields superior structural fidelity and convergence without additional neural network evaluations. As the step count is governed by the intrinsic geometry, the method remains agnostic to dataset-specific dynamics and model architecture.

The ability to achieve constant informational speed on statistical manifolds opens avenues for generalization to domains involving disparate dynamical rates, such as 3D molecular modeling and multimodal graph generative processes. The geometric criterion is theoretically grounded, invariant under sufficient statistics, and conforms to rigorous information-theoretic principles.

Conclusion

The paper formally establishes an information-geometric paradigm for adaptive sampling in graph diffusion, delivering an efficient and stable mechanism via Drift Variation Score. DVS realizes adaptive control of step size based on manifold curvature, yielding superior performance across molecular and general graph benchmarks. The approach is theoretically motivated, numerically validated, and offers practical flexibility for plug-and-play integration. Future directions include extension to complex domains with asynchronous evolution and further exploration of information-geometric control in generative modeling.