Composing diffusion priors with explicit physical context via generative Gibbs sampling

Abstract: Pretrained diffusion models provide powerful learned priors, but in scientific sampling the target distribution often depends on physical context that is not fully represented by one generative model. We introduce Generative Gibbs for Physics-Aware Sampling (GG-PA), a training-free framework that formulates the composition of learned partial priors and explicit physical context as inference over a joint target distribution in an augmented state space. We derive a Gibbs sampler for this joint target, show that it is asymptotically exact as the diffusion time approaches zero, and prove that in settings with quadratic interactions it remains exact at finite diffusion times. We further introduce replica exchange over diffusion time to accelerate mixing. Experiments on a double-well system, a $φ^4$ lattice model, and atomistic peptide systems show that GG-PA recovers context-induced distribution shifts and emergent collective behavior in interacting systems using partial priors without retraining. These results demonstrate GG-PA as a practical approach for combining pretrained generative priors with explicit physical context.

• The paper introduces GG-PA, which composes pretrained diffusion priors with explicit physical context to enable accurate equilibrium sampling in scientific systems.
• GG-PA employs generative Gibbs sampling with replica exchange to facilitate modular, inference-time composition and efficient phase space exploration.
• Empirical evaluations on double-well, 2D lattice, and atomistic models demonstrate GG-PA's ability to recover system-specific distributions and mitigate sampling bias.

Composing Diffusion Priors with Explicit Physical Context via Generative Gibbs Sampling

Overview and Motivation

The paper "Composing Diffusion Priors with Explicit Physical Context via Generative Gibbs Sampling" (2605.10642) introduces Generative Gibbs for Physics-Aware Sampling (GG-PA), addressing the critical challenge of deploying pretrained diffusion generative models in scientific systems where the equilibrium target distribution depends on explicit physical context, often not fully represented by a single generative model. This context is typically provided by environmental, field, or interaction variables. The target distribution is generally not matched by the prior training distribution, resulting in difficulties when performing posterior sampling using standard approaches. GG-PA facilitates modular, inference-time composition of one or more partial priors trained on subsystems with explicit system-level context, enabling accurate and efficient equilibrium sampling.

Framework Formulation

GG-PA operates by defining an augmented state space that comprises the explicit system configuration and the variables associated with each pretrained prior. The joint target density is indexed by diffusion time $t$, with couplings enforced by forward diffusion kernels. As $t\to 0$, strict consistency is imposed between system states and prior variables via $\delta$-functions; at finite $t$, the coupling relaxes. The mathematical formulation guarantees asymptotic correctness as $t\to 0$ for decomposable systems, with exactness at finite $t$ in the case of quadratic/linear-Gaussian interactions. The framework generalizes plug-and-play (PnP) split samplers, handling cases where the priors and context factorize nontrivially.

The paper further introduces replica exchange over diffusion time, allowing high-noise replicas to explore the phase space efficiently, while low-noise replicas maintain fidelity to the priors. The swap criterion cancels intractable learned prior densities and depends only on explicit context and known forward kernels, providing computational simplicity.

Theoretical Results

The theoretical analysis establishes GG-PA's asymptotic exactness for decomposable systems. For systems with quadratic interactions, an explicit matching condition ensures that the marginal distribution remains uncorrupted at finite diffusion time, provided the generative noise is below a critical threshold determined by intrinsic thermal fluctuations. This construction reveals a covariance correction for split Gibbs samplers in linear inverse problems, mitigating a bias present in existing methods and allowing exact posterior recovery without requiring the auxiliary noise to vanish.

Empirical Evaluation

GG-PA is empirically evaluated on three increasingly complex systems to demonstrate its modularity, correctness, and practical advantages:

1. Coupled Double-Well System

A benchmark where the prior is trained on the symmetric double-well, and at inference, environmental context induces an asymmetry. GG-PA recovers the asymmetric density accurately and demonstrates, through a tight Jensen-Shannon error analysis, that replica exchange accelerates mixing relative to fixed-$t$ and standard molecular dynamics.

Figure 1: Coupled double-well system. (a) The environment induces an asymmetric potential. (b) GG-PA captures the context-induced asymmetry in the marginal density. (c) JS divergence versus $t$ highlights the critical diffusion time bound for exactness.

2. 2D Ginzburg-Landau $\phi^4$ Lattice Model

GG-PA composes many partial priors (one per lattice site, trained on the local double-well) to generate the full many-body distribution. The approach produces correct collective behavior, including symmetry breaking and susceptibility peak at the critical coupling, verified via scaling collapse and order parameter measurement. Replica exchange enhances mixing near criticality, confirmed via dramatic reductions in autocorrelation time.

Figure 2: Representative field configurations, thermodynamic observables, and autocorrelation time across the phase transition. GG-PA tracks MC benchmarks and exhibits universal scaling.

3. Atomistic Alanine Dipeptide Systems

For non-quadratic and explicit molecular systems, GG-PA composes monomer priors with MD-based context. In the AD--Na$^+$ case, GG-PA captures ion-coordination-induced distribution shifts. For AD dimers, zero-shot composition from isolated priors recovers topology-dependent occupancy and emergent symmetry breaking. GG-PA-RE samples rare organized branches and offers state-assignment diagnostics with structural interpretation.

Figure 3: AD–Na$t\to 0$0: GG-PA recovers the ion-induced distribution shift. AD dimer: GG-PA-RE captures symmetry-broken organization and generates rare states.

Physical Interpretation of Replica Exchange

Analysis along the diffusion-time axis reveals a noise-induced phase transition among replicas. High-noise replicas exhibit softened interactions and local potentials, facilitating rapid exploration. This mechanism parallels Hamiltonian parallel tempering and is evidenced in the $t\to 0$1 lattice case, where the replica coordinate traverses from ordered to disordered regimes as generative noise increases.

Figure 4: Order parameter and susceptibility across the diffusion-time replica ladder, revealing a noise-induced transition and accelerated mixing.

Limitations, Practical Implications, and Future Directions

GG-PA provides modular inference-time composition that overcomes intractable marginalization required by standard posterior sampling, avoiding retraining monolithic models for environmental changes. Exact finite-$t\to 0$2 recovery is currently restricted to quadratic settings; in non-quadratic cases (e.g., atomistic systems), GG-PA remains asymptotically correct but incurs finite-$t\to 0$3 approximation error. Mixing relies on MCMC dynamics, with replica exchange providing practical mitigation. Residual double counting and projection design matter for practical accuracy.

Future work should extend context schedules to more general settings, develop efficient aggregation and tempering schemes, integrate multiscale priors, and refine interactions to avoid double-counting in force-field contexts. The modular paradigm enables reusable subsystem priors and explicit environmental composition, with direct applications to biomolecular simulation, protein–environment coupling, and complex molecular assembly.

Conclusion

GG-PA establishes a formal and computationally efficient method for combining pretrained diffusion priors with explicit physical context. The framework exhibits asymptotic and, in certain regimes, exact finite-time correctness, robust modularity, and strong empirical performance across scientific domains. Diffusion-time replica exchange substantially improves sampling efficiency in stiff regimes. The modular, inference-time composition paradigm outlined by GG-PA promises scalable generative modeling in scientific systems where environmental context and subsystem interactions are essential and difficult to anticipate during training.