Composing diffusion priors with explicit physical context via generative Gibbs sampling

The paper introduces GG-PA, a training-free framework that combines pretrained diffusion priors with explicit physical context via a generative Gibbs sampler to achieve asymptotically exact sampling and recover context-induced distribution shifts in scientific systems without retraining.

The Gist

Imagine you are trying to bake the perfect cake, but you have two different tools: a magic recipe book and a real kitchen.

• The Magic Recipe Book (The Diffusion Prior): This is a pre-trained AI model. It has "read" millions of photos of isolated cake layers. It knows exactly what a perfect, standalone cake layer looks like. However, it has never seen a cake with frosting, or a cake sitting next to a bowl of fruit, or a cake in a humid kitchen. It only knows the "pure" cake layer.
• The Real Kitchen (The Physical Context): This is the actual environment where you are baking. It includes the humidity, the weight of the frosting, the heat of the oven, and how the cake interacts with the fruit.

The Problem:
If you just use the Magic Recipe Book, you get a perfect cake layer, but it won't fit in your real kitchen. If you try to force the kitchen rules onto the book, you might break the book's understanding of what a cake is. Scientists often face this: they have great AI models for specific parts of a system (like a protein backbone), but they need to simulate the whole system (protein + water + ions), and the AI doesn't "know" about the water.

The Solution: GG-PA (Generative Gibbs for Physics-Aware Sampling)
The authors created a new method called GG-PA. Think of it as a smart dance between the Magic Recipe Book and the Real Kitchen.

Instead of trying to rewrite the recipe book or ignore the kitchen, GG-PA makes them work together in a loop:

1. The "Denoising" Step (Consulting the Book): The system looks at the current state of the cake in the kitchen. It asks the Magic Recipe Book: "Given this messy kitchen situation, what does a perfect cake layer look like?" The book gives a suggestion based on its training.
2. The "Aggregation" Step (Listening to the Kitchen): The system then takes that suggestion and asks the Real Kitchen: "Okay, but does this suggestion actually fit with the frosting and the humidity? Let's adjust the cake to make sure it obeys the laws of physics in this specific room."

They repeat this dance over and over. The book keeps the cake looking like a cake, and the kitchen keeps the cake fitting the environment.

The Secret Sauce: The "Noise" Dial
The paper introduces a clever trick involving a "Noise Dial" (called Diffusion Time).

• Low Noise (Strict Mode): The Magic Recipe Book is very strict. It demands the cake look exactly like its training data. This is accurate, but the dance becomes stiff and slow. The cake gets stuck in one spot and can't explore new shapes.
• High Noise (Relaxed Mode): The Magic Recipe Book is more relaxed. It says, "Okay, the cake can look a bit messy." This makes the dance fast and energetic, allowing the system to explore many different cake shapes quickly.

The "Replica Exchange" Trick
To get the best of both worlds, GG-PA runs multiple copies (replicas) of the dance at the same time.

• Some copies dance with the Strict Book (Low Noise) to ensure accuracy.
• Some copies dance with the Relaxed Book (High Noise) to explore quickly.
• Every now and then, they swap places. The strict copy gets a turn to be relaxed and explore, and the relaxed copy gets a turn to be strict and refine the shape.

This is like having a team of bakers: some are perfectionists who double-check every detail, and others are fast explorers who try wild new ideas. They swap roles so the team gets both speed and accuracy.

What They Proved
The authors tested this on three things:

1. A Simple Math Puzzle: A system with two valleys (like a ball rolling between two hills). They showed that when the math is simple (quadratic), their method is perfectly exact, even with the noise dial turned up.
2. A Grid of Interacting Particles: They showed that even if the AI only learned about single particles, this method could combine many of them to create complex, collective behaviors (like a crowd moving together) that the AI never saw during training.
3. Real Molecules (Peptides): They used the method to simulate a small protein (Alanine Dipeptide) interacting with a sodium ion and another protein. The AI knew the protein shape, but not the ion. GG-PA successfully combined them, showing the protein changing shape to fit the ion, something the AI couldn't do on its own.

In Summary
GG-PA is a way to use a specialized AI (which knows a lot about one part of a system) and combine it with real-world physics rules (which know about the rest of the system) without having to retrain the AI. It uses a "dance" of alternating updates and a "team swapping" strategy to ensure the result is both scientifically accurate and computationally efficient.

Technical Summary

Technical Summary: Composing Diffusion Priors with Explicit Physical Context via Generative Gibbs Sampling

Problem Statement
Pretrained diffusion models offer powerful learned priors for scientific sampling, yet they often describe only a selected subset of a system's degrees of freedom (e.g., a protein backbone or a molecular fragment) rather than the full system state. In scientific applications, the target distribution frequently depends on physical context—such as solvents, ions, external fields, or interactions with other subsystems—that is not fully represented by a single generative model. Standard inference-time approaches like guidance or posterior sampling typically require all context to be expressed in terms of the generative model's variables. This necessitates marginalizing unrepresented degrees of freedom into an effective free energy term, which is often intractable for high-dimensional environments or redundant when other subsystems are already well-modeled by separate priors or force fields. The core challenge addressed is the inference-time composition of multiple partial learned priors with explicit system-level physical context without retraining the models.

Methodology: Generative Gibbs for Physics-Aware Sampling (GG-PA)
The authors propose 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.

1. Augmented State Space: The method maintains an explicit representation of the full system state $s$ (e.g., all-atom coordinates including solvent) and couples it to $K$ pretrained diffusion priors via projection operators $\Phi_i: S \to X_i$. The augmented state is $Z = S \times \prod X_i$.
2. Joint Target Distribution: A family of joint target densities is defined indexed by diffusion time $t$:
   $$\pi_t(s, \{x_i\}) \propto q_{\text{ctx}}(s, t) \prod_{i=1}^K \left[ p_i(x_i) \cdot q^{(i)}_t(\Phi_i(s) | x_i) \right]$$
   Here, $p_i$ are the pretrained priors, $q^{(i)}_t$ are the forward diffusion kernels acting as couplings, and $q_{\text{ctx}}$ is the explicit physical context factor (e.g., a Boltzmann factor). As $t \to 0$, the coupling kernels enforce strict consistency ($\Phi_i(s) = x_i$), recovering the composed distribution where priors govern specific subsets and the context governs the rest.
3. Generative Gibbs Sampler: Sampling alternates between two steps:
   • Parallel Denoising: Each prior variable $x_i$ is updated by sampling from the posterior induced by the prior $p_i$ and the current projected state $\Phi_i(s)$ treated as a noisy observation. This is performed by running the pretrained reverse-time sampler.
   • Context-Aware Aggregation: The full system state $s$ is updated conditioned on the current $x_i$ values and the explicit context. This step minimizes an effective potential $U_{\text{eff}}$ derived from the context and the log-likelihoods of the forward kernels.
4. Replica Exchange: To address the trade-off between fidelity (small $t$) and mixing (large $t$), the authors introduce replica exchange over diffusion time. Multiple replicas run at different $t$ values, with swap moves proposed based on a tractable acceptance ratio where the intractable prior densities cancel out.

Theoretical Properties

• Asymptotic Exactness: For decomposable systems, the marginal target recovers the true physical distribution as $t \to 0$.
• Finite-Time Exactness: In settings where interactions are quadratic (linear-Gaussian), the method remains exact at finite $t$ provided the context schedule is parameterized to satisfy specific moment-matching conditions (Gaussian deconvolution). This yields a critical bound on the maximum admissible diffusion time $t_{\text{max}}$.
• Connection to Split Gibbs: The framework generalizes split Gibbs samplers for linear inverse problems, offering a covariance correction that avoids the bias present in standard implementations.

Experimental Results
The authors evaluate GG-PA on three systems of increasing complexity:

1. Coupled Double-Well System: A 2D quadratic system used to verify finite-time exactness and the efficacy of replica exchange. GG-PA successfully recovered the environment-induced asymmetry. Replica exchange significantly accelerated mixing in the stiff, low-$t$ regime compared to fixed-$t$ sampling and molecular dynamics (MD).
2. $\phi^4$ Lattice Model: A 2D Ginzburg-Landau model testing the composition of many-body collective behavior absent from the training distribution. The model was trained only on local on-site double-well factors. GG-PA successfully reproduced the equilibrium phase transition, spontaneous symmetry breaking, and critical exponents. Replica exchange provided orders-of-magnitude speedups near the critical point.
3. Alanine Dipeptide Systems: Atomistic models involving non-quadratic interactions.
   • AD–Na+: GG-PA captured the distribution shift in carbonyl oxygen distances induced by ion coordination, outperforming a vacuum-trained prior used directly.
   • AD Dimer: Two copies of a monomer prior were composed to model hydrogen-bonded dimers. GG-PA-RE recovered the qualitative symmetry-broken organization (anti-parallel vs. parallel topologies) and conditional torsional state preferences, despite the non-quadratic nature of the interactions and the lack of exact finite-$t$ guarantees.

Key Contributions

1. Formulation: A novel formulation of partial diffusion prior composition as inference over an explicit full-system state, bypassing intractable marginalization.
2. Algorithm & Theory: Derivation of the GG-PA sampler with proofs of asymptotic exactness, finite-time exactness for quadratic interactions, and a covariance correction for split Gibbs samplers.
3. Practical Demonstration: Numerical demonstration of modular multi-prior composition in systems with and without quadratic interactions, showing the ability to recover context-induced shifts and emergent collective behavior without retraining.

Significance and Claims
The paper positions GG-PA as a practical approach for combining pretrained generative priors with explicit physical contexts. The authors claim that this modular paradigm allows learned priors and explicit physics to be applied where most appropriate, avoiding the need to retrain monolithic models when system environments change. The method is particularly valuable for scientific systems with high-dimensional environmental degrees of freedom that are readily handled by force fields or separate priors. The authors acknowledge limitations, including the reliance on quadratic structures for finite-$t$ exactness and the computational cost of maintaining multiple replicas, but emphasize the method's ability to handle complex, context-dependent sampling tasks that are difficult for standard posterior sampling or guidance techniques.