Scalable Training of Diffusion Samplers via Fixed-point Score Matching

Denis Blessing opens by framing the core problem: sampling from an unnormalized density, where you can evaluate an energy function but have no samples from the target distribution. He stresses that this is a pure sampling problem, which makes it harder than the data-driven generative setting. Target densities are often multimodal, high-dimensional, and expensive to evaluate (a single DFT energy call can take hours), so efficiency matters as much as correctness.

Classical methods and why neural samplers

He reviews MCMC and molecular dynamics. These come with asymptotic guarantees, but multimodal targets trap chains in local modes, producing slow mixing. They also rarely support conditioning: changing temperature or another parameter means rerunning the full simulation. Neural samplers aim to fix both issues by learning a transport map from a simple prior, such as a normalizing flow or diffusion model, trained only from the energy. A well-trained sampler yields independent samples with finite-time guarantees and can be amortized over system parameters like temperature, so one model covers many configurations. When the model is imperfect, importance weights can debias expectations and recover free energies.

A short history of diffusion samplers

He traces diffusion samplers from roughly 2022 to 2023 (PDDS, DDS, and related work). The early approaches minimized the reverse KL divergence on path space. He points out the drawbacks: sample reuse is difficult, the target density must be evaluated many times, gradients must be backpropagated through the SDE solver (a memory problem at scale), and reverse KL is mode-seeking, which invites mode collapse. Log-variance losses followed in 2023 and 2024. They avoid backpropagation through the integrator by detaching parameters and can reuse trajectories from any reference process, but they remain high-variance and still require storing full trajectories, unlike score matching, where random time subsampling is straightforward.

Fixed point score matching

The central contribution is a unifying view of the recent methods (Edge sampling, Adjoint/shooting bridges, flow sampling, and the bridge machine). He presents them under one principle he calls fixed point score matching. He first reviews stochastic interpolants with independent coupling, where the velocity and score are conditional expectations over the interpolant. This yields regression objectives like denoising score matching and flow matching. The obstacle is that these expectations require samples from the target, which are unavailable.

The fix is a target score identity, derived by integration by parts. It expresses the score through an expectation over the target score, which is computable from the energy. The remaining step is to replace the target-distribution expectation with one under the current model. The resulting fixed point iteration is: sample from the current model, store terminal samples and their target scores in a replay buffer, interpolate, regress velocity and score, and repeat. He states that the optimal fields are fixed points of this loop and expresses the hope that it converges to the Boltzmann distribution. The payoff is that training now resembles data-driven optimization, which makes neural samplers efficient to train.

For evaluation he describes two importance-weight routes. The first uses the instantaneous change of variables from the velocity field. These weights are low variance but expensive in high dimensions because they require a divergence per time step, and stochastic estimators tend to be high variance. The second uses forward-backward path weights, which add little overhead but accumulate variance at each diffusion step. He concedes that free energy estimation is still open, particularly at low temperatures.

Materials science applications

The first application is Atlas, which models amorphous materials. These are disordered, with many metastable states, and they are hard to sample because Markov chains get stuck in local minima. No database of structures exists to train a standard flow matching or score matching model, which is what makes the fixed point approach attractive. Atlas starts from a uniform random configuration in a periodic cell and learns to generate metastable structures. The adaptations include a periodic-aware target identity, a periodic-equivariant GNN, and amortization over temperature and composition.

In experiments on the Kob-Andersen glass former, a well-known and difficult benchmark especially at low temperature, the generated structures match the reference in radial distribution functions and energy histograms. For free energies, he uses path weights at high temperature as an anchor and propagates to lower temperatures with classical methods like MBAR, since path weights are unreliable at low temperature. Because the model uses graph neural networks, it transfers from 64-particle training to systems as large as 512 particles. Atlas also needs several orders of magnitude fewer target energy evaluations than parallel tempering (about 10^5 versus 10^8). He briefly mentions further applications: potential of mean force via umbrella sampling, reward-based steering and fine-tuning, and inverse design with an LLM-driven loop.

The second project is Yanos, which studies alloys and crystals in the semi-grand canonical ensemble. Here the volume is free and the pressure is fixed. The state space is hybrid: a continuous part for atomic positions and volume, and a discrete part for species. Continuous variables use the same velocity and score machinery with physics-informed features. For the discrete species, he uses a masked diffusion model. Because the naive substitution of samples from the current model yields a self-consistent but uninformative objective, he derives single-site conditionals from substitution energies and trains with a label cross-entropy loss. These conditionals are shown to equal the true denoising posterior, so the loss is consistent with the target. Importance weights combine the continuous path weight with the discrete autoregressive weight.

Validation includes the Ising model and binary alloys (copper-nickel and copper-silver). The phase-diagram-like outputs match the reference well. Once amortized over temperature and chemical potential, one model replaces the hundreds of simulations that phase diagrams would otherwise require. The cost reduction is roughly 10^9 energy evaluations for parallel tempering versus about 10^6 for the amortized model.

Reinforcement learning

The second application is a paper on reinforcement learning via fixed point score matching. He casts two settings as sampling problems: maximum-entropy RL, where the optimal policy is proportional to the exponential of the Q-function, and behavior-constrained RL, where the optimal policy additionally tilts a reference policy. He uses expressive flow policies for this. Because flow matching needs expectations under the target, he derives a target flow matching loss by combining the velocity-score relation with the target score identity. It shares the minimizer of standard flow matching but depends on the gradient of the target instead of samples. A notable advantage is that the behavior-constrained variant needs only the gradient of the reference policy’s log density, not its likelihood, so no ODE-based likelihood or adjoint backpropagation is required.

He adds two extensions. The first is a trust-region update, enforced by a KL constraint solved with dual gradient descent. The second is correlated exploration: noise draws are correlated across states with an Ornstein-Uhlenbeck process while each marginal stays Gaussian, which produces more sustained, directed exploration.

On 20 DeepMind Control tasks and humanoid tasks, the method performs competitively, and with correlated exploration it achieves the best results among the compared methods (including PPO and diffusion RL baselines). In offline-to-online RL, a pretrained flow policy is fine-tuned offline and then online. Offline performance lags the baseline on some tasks, but online performance is strong. He notes that baseline variance is omitted from the plots for readability. In a cherry-picked example, correlated exploration lifts returns from near zero to much higher and stabilizes training.

Limitations

He closes with several limitations. Because the score is local, these methods are insensitive to the relative weights of well-separated modes. Path-based free energy estimates are cheap but high variance. The independent coupling used in the target score identity discards information from simulated (x0, x1) pairs. Simulation cost still dominates for large models, even though the optimization itself is simulation-free. Uniqueness and convergence of the fixed point iteration remain open theoretical questions.

Discussion highlights

In the Q&A, a participant asked how the method compares to adjoint matching. Blessing agreed the objectives coincide in simple linear cases. For nonlinear reference policies, his target score identity can be more computationally efficient. He noted the reference score can be much larger than the reward gradient, which may increase variance. He mentioned that Efficient Action Matching uses the terminal score at arbitrary times, which is less principled but works well in practice. Another participant raised the importance of sampling cost and whether the variance problems of path weights are overstated, since molecular applications usually require sampling anyway. Blessing agreed that no universal solution exists yet. On non-stationary environments he said he had not studied them. Asked about couplings beyond independent pairs, he mentioned that Edge/shooting-based matching learns a shooting coupling, but the alternating optimization proved unstable in his experiments, so independent coupling remains his practical default.

• Tags:
• Reinforcement Learning
• Feature
• Diffusion Models
• Flow Matching
• Materials Science
• Neural Samplers
• Fixed Point Score Matching
• Boltzmann Distribution
• Amortized Inference
• Importance Sampling
• Free Energy Estimation