Fixed-Point Score Matching: A New Paradigm for Training Diffusion Samplers

The Sampling Problem: Generating Samples from an Energy Function

In scientific computing, there is a class of problems that is both common and notoriously difficult: we only have an unnormalized density function that we can evaluate and differentiate, but we have no samples. This is the core setup of the sampling problem. Target densities often have multimodal structure and are expensive to evaluate—for instance, a single energy calculation using density functional theory may take hours—and the problem dimensionality is typically extremely high.

Such problems are widespread in molecular simulation, materials science, Bayesian inference, and reinforcement learning. Classical methods like Markov Chain Monte Carlo (MCMC), despite their asymptotic guarantees, easily get trapped in local modes when facing multimodal distributions, with prohibitively long mixing times. More critically, MCMC lacks flexibility: changing the temperature or system parameters requires re-running the entire simulation.

The idea behind neural samplers is to use machine learning to learn a transport map from a simple prior to the target distribution. If trained well, it can generate independent samples, amortize costs, and debias via importance weights. Diffusion samplers are among the most expressive methods in this class.

The relationship between unnormalized density and energy function: In statistical physics and Bayesian inference, the target distribution is typically given in Boltzmann form $p(x) \propto e^{-E(x)}$, where $E(x)$ is the energy function or negative log-density. The normalization constant (partition function) $Z=\int e^{-E(x)}dx$ is often intractable, so we can only evaluate the unnormalized density or its gradient with respect to $x$ (i.e., the score $\nabla_x \log p(x) = -\nabla_x E(x)$). The so-called "sampling problem" is to generate samples that follow $p(x)$ when only $E(x)$ is known and neither $Z$ nor samples are available. This setup runs through molecular simulation, materials science, and Bayesian inference, and is the fundamental difficulty that diffusion samplers aim to solve.

Early Struggles of Diffusion Sampling

Diffusion models perform impressively in data-driven settings, but score matching requires samples from the target distribution—precisely what is missing in the sampling problem. Starting in late 2022, researchers attempted to train diffusion samplers using variational inference by minimizing the reverse KL divergence over path space.

These methods have several fatal problems. Reverse KL is zero-forcing and tends toward mode-seeking, leading to severe mode collapse in practice. More practically, the loss function is defined over the entire trajectory, requiring backpropagation through the SDE solver with enormous memory overhead; samples cannot be reused, meaning repeated evaluation of the costly target density.

The log-variance loss proposed in 2023–2024 alleviated some issues—it replaces the expectation under the current model with an arbitrary reference process, allowing sample reuse and gradient truncation. However, high-variance gradients make optimization highly unstable, and storing the full trajectory is still required. Subsequent benchmarks even concluded directly that these methods do not work well.

Reverse KL divergence and zero-forcing: In variational inference, the reverse KL divergence $D_{KL}(q|p)$ is commonly used to measure the difference between an approximate distribution $q$ and the target $p$. A key property is "zero-forcing": when $p(x)=0$ but $q(x)>0$, the KL divergence imposes a huge penalty, forcing $q$ to vanish in regions where $p$ is zero. This causes $q$ to cover only a few high-probability modes of $p$, i.e., mode collapse, and to fail to correctly assign relative weights among modes. In contrast, forward KL $D_{KL}(p|q)$ is "mass-covering", encouraging $q$ to cover all regions where $p$ is nonzero, but it requires samples from the target distribution. Early diffusion samplers used reverse KL precisely to avoid the need for target samples, but inherited the mode-seeking deficiency.

Fixed-Point Score Matching: A Unified Framework

In his talk, Dennis Blessing proposed that the many diffusion sampling algorithms that have emerged in recent years—from marginal shooting to Schrödinger bridges, flow matching, and bridge matching sampling—appear disparate but actually share the same underlying principle. He unified them under the name fixed-point score matching.

The core insight is: although we do not have samples from the target distribution, we do have samples generated by the current model. By replacing the denoising score identity in stochastic interpolation with the target score identity, the expectation can be computed under the current model distribution instead of the target distribution. The resulting iterative loop is: generate terminal samples from the current model into a replay buffer, compute the energy gradient (i.e., score) at these samples, then supervise with standard score matching and flow matching losses, update the model, and repeat.

Theoretically, the optimal velocity field and score field are precisely the fixed point of this iteration. The elegance of this framework is that the training procedure is almost identical to data-driven diffusion model training, allowing efficient subsampling of time and states for optimization, free from the constraint of the entire trajectory.

Stochastic interpolation and the denoising score identity: Stochastic interpolation constructs a probability path from a prior distribution $p_0$ to a target distribution $p_1$, typically defined as $x_t = \alpha_t x_0 + \sigma_t x_1$, where $x_0 \sim p_0$ and $x_1 \sim p_1$. The denoising score identity states that, given $x_t$, there is an analytical relationship between the conditional expectation $\mathbb{E}[x_1|x_t]$ and the marginal score $\nabla_{x_t} \log p_t(x_t)$. In standard diffusion models, $p_1$ is the data distribution, so the identity can be estimated from data samples. In the sampling problem, however, $p_1$ is unknown and no samples are available. The key innovation of fixed-point score matching is to replace the denoising score identity with the target score identity: using samples $x_1 \sim q_\theta$ generated by the current model in place of true target samples, thereby transferring the expectation from the target distribution to the model distribution and forming an iteratively optimizable fixed-point equation.

Materials Science Applications: From Amorphous to Alloys

The research team applied this framework to two materials science problems. The ATLAS project models amorphous materials, starting from a periodic uniform distribution and learning to generate glassy structures with numerous metastable states. In the classic Kob-Andersen glass-forming system experiments, the model not only accurately reproduced the energy histogram but also matched reference values in free energy estimation.

Thanks to the graph neural network architecture, the model exhibits excellent size transferability: trained on 64 particles, it directly infers systems of 512 particles while maintaining good accuracy. Even more impressive is the computational cost—ATLAS requires orders of magnitude fewer target energy evaluations than parallel tempering methods.

The MANNERS project further handles the more complex semi-grand canonical ensemble, simultaneously modeling continuous channels (atomic positions and volume) and discrete channels (atomic species). The discrete part employs a masked diffusion model, and the team derived a label cross-entropy loss whose conditional expectation exactly corresponds to the true denoising posterior, providing theoretical guarantees for fixed-point iteration. On the Ising model and copper-nickel and copper-silver alloy systems, the model amortizes temperature and chemical potential, covering the entire phase diagram parameter space with a single model.

Semi-grand canonical ensemble and masked diffusion: In open systems such as alloys, the species and number of atoms can vary, requiring description by the semi-grand canonical ensemble, where continuous variables (atomic positions, volume) and discrete variables (atomic species) coexist. Masked diffusion models handle discrete variables as absorbing states: during the diffusion process, atomic species are replaced with a mask token with some probability, and the model learns to recover the original species from the masked state. This design allows defining a diffusion process in discrete space and deriving the label cross-entropy loss. Its conditional expectation exactly corresponds to the true denoising posterior, seamlessly integrating with the fixed-point score matching framework, enabling a single amortized model to cover different temperatures, chemical potentials, and compositions, substantially reducing repeated simulation costs.

Reinforcement Learning: New Possibilities for Deterministic ODE Policies

The same framework also has unexpected applications in reinforcement learning. The optimal policy in maximum entropy reinforcement learning is proportional to the exponential of the Q-function, which is itself a sampling problem. The research team derived the velocity score identity, combined it with the target score identity, and obtained the target flow matching loss—its minimum is the same as the standard flow matching loss, but it depends only on energy gradients rather than target samples.

This means deterministic ODE policies can be trained directly, without SDE conversion or adjoint ODE backpropagation. In behavior-constrained reinforcement learning, KL regularization can be achieved by accessing only the gradient of the reference policy, and the reference policy's gradient can also be obtained through the velocity score transformation.

The algorithm also introduces two practical extensions. Trust region updates ensure training stability through L2 constraints between old and new policies and primal-dual descent; cross-state correlated exploration exploits the deterministic mapping property of flow policies to impose correlations on initial Gaussian noise across different environment states, significantly improving exploration efficiency without changing marginal distributions.

On DeepMind Control and Humanoid tasks, the method with correlated exploration achieved the best performance among all compared methods. In offline-to-online reinforcement learning settings, the flow policy's final online performance is also excellent.

Open Questions

This framework also has clear limitations. Scores are inherently local quantities, and when target distribution modes are widely separated, the method struggles to correctly weigh the relative importance of each mode. Although free energy estimation has two approaches—path weights and ODE weights—the former's variance grows cumulatively with the number of steps, and the latter's Hessian computation cost is too high for large systems. In practice, the team combined diffusion weights at high temperature with classical methods for downward propagation, which remains an open problem.

The use of independent coupling discards relevant information in the pairing of x0 and x1, and better couplings such as shooting coupling, while theoretically appealing, perform poorly in practice due to instability from alternating optimization. Even though optimization itself requires no simulation, the replay buffer still needs to continuously generate samples, and simulation costs may dominate training for large models. Theoretically, the uniqueness and convergence of fixed-point iteration have not been proven—these methods perform well in practice, but the mathematical foundations remain to be refined.

One-line Summary

Fixed-point score matching unifies diffusion sampling algorithms, trains solely from energy functions, and applies efficiently in materials science and reinforcement learning.

Paragraph Summary

This article introduces the core setup of the sampling problem: only an unnormalized density (energy function) can be evaluated, with no samples available. Traditional MCMC and early diffusion samplers suffer from slow mixing, mode collapse, and large memory requirements. Dennis Blessing's fixed-point score matching framework combines the target score identity with stochastic interpolation, computes energy gradients on samples generated by the model, and iteratively trains diffusion samplers without target samples. The framework unifies algorithms such as marginal shooting, Schrödinger bridges, and flow matching. In materials science, the ATLAS and MANNERS projects handle amorphous and alloy systems, reducing computational costs by orders of magnitude and supporting size transferability and parameter amortization. In reinforcement learning, the target flow matching loss directly trains deterministic ODE policies, and combined with correlated exploration achieves optimal performance. Limitations include difficulty in mode weight estimation, the trade-off between variance and cost in free energy estimation, and the lack of convergence theory.