Discrete Feynman-Kac Correctors

AI summary

Overview

Research area: Discrete generative modeling — specifically, inference-time control of discrete (masked) diffusion models using Sequential Monte Carlo methods.

Technical level: Advanced. The paper is built around continuous-time Markov chains, the Forward Kolmogorov Equation, rate matrices, and Feynman-Kac formulations, with proofs relegated to appendices.

Scope: The paper derives a general framework — Discrete Feynman-Kac Correctors (DFKC) — that reweights and resamples the trajectories of a trained discrete masked diffusion model so that it samples from annealed, product, geometric-average, or reward-tilted distributions, without any additional training or fine-tuning.

What This Paper Is About

Discrete diffusion models generate discrete sequences (text, code, biological sequences) by gradually denoising or demasking, but the process as trained gives no flexible control over the distribution of the samples it produces. Prior work (Skreta et al., 2025) introduced Feynman-Kac Correctors for this purpose, but derived them only for the Fokker-Planck equation, so they do not apply directly to discrete diffusion models, which are described by a Continuous-Time Markov Chain (CTMC). This paper closes that gap by deriving correctors directly on the Forward Kolmogorov Equation.

Key Contributions

1. A general CTMC framework. The authors state a Feynman-Kac formula (Theorem 2.1) for the Forward Kolmogorov Equation with a rate matrix and a weighting function, and then derive three transformations of an arbitrary CTMC: temperature annealing (Theorem 3.1), product of marginals (Theorem 3.3), and reward-tilting (Theorem 3.5). A geometric-average formulation for any number of marginals is given in Theorem C.1 and Corollary C.2.

2. Specialization to masked diffusion. Corollaries 3.2, 3.4 and 3.6 give the reverse-time rate matrix and weighting function for masked diffusion under each transformation. The new rate matrices and weights are expressed purely in terms of the ratio of marginal probabilities — equivalently the denoising conditional probability — which a trained model already provides, so no extra training or fine-tuning is required.

3. A practical Sequential Monte Carlo algorithm. Algorithm 1 generates a batch of samples with weights, updating each state by sampling from a categorical distribution defined by the corrected rate matrix, updating log-weights by the weighting function, and periodically resampling via Self-Normalized Importance Sampling. This introduces explicit interactions between samples through re-weighting and re-sampling.

4. A multi-domain empirical demonstration. The framework is illustrated on Ising model sampling, code generation, amortized learning, and reward-tilted protein sequence generation.

Main Findings

• Annealing beats direct training at low temperature on the Ising model. Training a discrete diffusion model at β = 0.3 and annealing it at inference time to a target temperature was more efficient than training directly at the target. On a 16×16 lattice with J_ij = 1, h_i = 0 and periodic boundary conditions, for target β = 0.4, DFKC(0.3) achieved Energy-W₂ of 14.24 ± 3.11 versus 69.38 ± 4.25 for the directly trained DDM, Magnetization-W₂ of 0.256 ± 0.052 versus 0.889 ± 0.063, and Correlation-MSE of 0.041 ± 0.013 versus 0.172 ± 0.021. For target β = 0.3, DFKC(0.2) achieved 33.38 ± 0.46, 0.031 ± 0.011 and 0.023 ± 0.007, versus 35.14 ± 0.63, 0.046 ± 0.012 and 0.014 ± 0.009 for the DDM. All metrics are mean ± standard deviation over 3 seeds.

• Annealing extends beyond the critical temperature. The model was trained at β = 0.3, and DFKC controlled temperature over β ∈ [0.3, 0.6] with critical point β_crit ≈ 0.4407. The released LEAPS model (Holderrieth et al., 2025) is trained on the path π_t(σ) ∝ exp(−t β_crit H(σ)) for t ∈ [0,1], so the authors do not expect it to work reliably for β > β_crit. The authors report that their training-free method samples beyond the critical temperature, and that mean energy and magnetization closely follow theoretical values.

• Higher accuracy on code generation. Using pretrained LLaDA-8B-Instruct as the masked diffusion model, annealing with DFKC on HumanEval and MBPP produced higher accuracy than both "Argmax" sampling (taking the most likely token at each step) and "Naive Annealing" (sampling with inverse-temperature β). The specific accuracy values are not reported in the provided content.

• Amortized learning via products of marginals. Because a uniform-prior posterior factors as p(θ|X) ∝ ∏ p(θ|X_k) over dataset subsets, the product formulation applies directly: each factor is conditioned on a different subset C_k = X_k. The task was evaluated on a synthetic dataset generated by a noisy linear predictor f_θ(x) = θ₁x + θ₀ + ε with ε ∼ N(0, 0.1²). Numerical results for this experiment are not reported in the provided content.

• Reward-tilted protein generation. DFKC was used to generate realistic protein sequences while optimizing external reward functions (Wang et al., 2024b). Quantitative results are not reported in the provided content.

• Cost of reward-tilting. Evaluating the reward-tilted rate matrix B_τ^reward(i,j) requires computing the reward function at all states j reachable from the mask state m, and the weighting term requires summing the reward over all such states, which the authors note can be computationally expensive depending on the application.

• Scaling to multiple dimensions. The derivations transfer to any number of dimensions; the reverse-time rate matrix is nonzero only when all coordinates except one match, so one can parameterize the process by predicting (m−1)d values, where d is the number of dimensions (or sequence length) and (m−1) is the vocabulary size for each discrete variable.

Methodology in Plain English

The authors begin from the Forward Kolmogorov Equation, the equation that governs how the probability distribution of a continuous-time Markov chain evolves over time, and which is defined by a rate matrix A_t(i,j) that specifies how likely the state is to jump from i to j.

Their starting point is that when you add an extra re-weighting term g_t(i) to this equation, the process can be simulated by carrying weights alongside the states — this is the Feynman-Kac formula (Theorem 2.1). The weights are accumulated multiplicatively as the trajectory evolves, and the weighted samples can be turned into an estimate of the target distribution using Self-Normalized Importance Sampling (SNIS).

They then take each desired target distribution — the annealed marginal p_t(i)^β, the product p_t¹(i)p_t²(i), the geometric average, and the reward-tilted p_t(i)exp(β_t r(i)) — and work out algebraically what rate matrix and what weighting function would reproduce it. The key practical point is that the resulting formulas depend only on ratios of marginal probabilities, p_t(j)/p_t(m), which are exactly what a trained masked diffusion model already computes. For masked diffusion, the ratio relates to the conditional de-masking probability p(x₀ = j | x_t = m), the standard output of the network. Each corrector therefore requires only a single forward pass through each model at each step.

Simulation follows Algorithm 1: initialize samples from the noise distribution p_{t=1}(i) with uniform weights 1/K, and for each step in the discretized reverse process, sample a new state from the corrected rate matrix, update the log-weights by the weighting function, and — when resampling is triggered — renormalize the weights, draw new sample indices from the categorical distribution over weights, and reset all weights to 1/K.

Experiments follow the same recipe across domains. For the Ising model, training data was produced with the Swendsen-Wang algorithm, the model was a UNet on a 16×16 lattice, and sampling used 10,000 samples following a long burn-in period of Glauber dynamics. For language modeling, the pretrained LLaDA-8B-Instruct model was used as-is.

Why This Matters

Impact on research. The paper transports inference-time distribution control from continuous diffusion (where it was formulated for the Fokker-Planck equation and SDEs) into the discrete CTMC setting that actually describes masked diffusion models of text, code, and biological sequences. Since the correctors need no retraining and require only the quantities a trained model already produces, they turn any existing discrete diffusion model into a controllable sampler, and they connect discrete diffusion to the sequential Monte Carlo machinery used in continuous generative modeling.

Real-world applications:

• Sampling from physical distributions. Efficiently sampling configurations of the Ising model at various temperatures, including beyond the critical temperature, is directly relevant to statistical physics and to using generative models as samplers for Boltzmann distributions.
• Code generation. Annealing a pretrained language model at inference time improved accuracy on HumanEval and MBPP, suggesting a drop-in way to sharpen a code model's outputs without any fine-tuning.
• Amortized Bayesian inference through text. Sampling from a posterior over parameters by prompting a language model with dataset subsets, using the product-of-marginals formulation, points toward doing statistical inference entirely through a model's text interface.
• Protein design. Generating realistic protein sequences while optimizing an external reward function is the standard setting for property-conditioned biological sequence design.

Industry relevance. Because DFKC is training-free and inference-time only, it can be layered onto already-deployed discrete diffusion models (including large pretrained ones such as LLaDA-8B-Instruct) as a sampling-time wrapper. The cost is a single forward pass per model per step plus resampling overhead, roughly doubling the per-step network evaluations in the two-model product case (the rate matrix is scaled by 2), which matters for deployment budgets.

Future Directions

• Cheaper reward evaluation. Computing the weighting term for reward-tilting requires summing the reward over all states reachable from the mask. The authors explicitly leave as future work the possibility of using alternative functions that evaluate the reward difference on transitions from m to j, i.e. r(j) − r(m), instead.
• Completing and quantifying the remaining experiments. The provided content reports qualitative outcomes for amortized learning and reward-tilted protein generation but no numerical results; those results, and the effect of varying the number of data subsets K in the amortized learning setup, are natural open questions.
• Extending the geometric-average formulation. The general geometric average over any number of marginals is stated in Theorem C.1 and Corollary C.2, but the main text presents the product of two marginals in detail; scaling to many simultaneously combined models is a natural extension.
• Practical guidance on resampling and temperature ranges. The method introduces interactions between samples through re-weighting and resampling; how resampling schedules, number of samples K, and annealed temperature ranges interact with sample quality across domains remains to be characterized.

Target Audience

Researchers and practitioners in generative modeling who work with discrete diffusion or masked diffusion models; machine learning scientists interested in inference-time control, Sequential Monte Carlo, and Feynman-Kac methods; computational physicists using generative models as Boltzmann samplers; and applied researchers in code generation, amortized Bayesian inference, and protein sequence design who want to steer a pretrained discrete model toward a target distribution without retraining. Readers need comfort with continuous-time Markov chains, rate matrices, and importance sampling to follow the derivations.

Authors’ abstract

Discrete diffusion models have recently emerged as a promising alternative to the autoregressive approach for generating discrete sequences. Sample generation via gradual denoising or demasking processes allows them to capture hierarchical non-sequential interdependencies in the data. These custom processes, however, do not assume a flexible control over the distribution of generated samples. We propose Discrete Feynman-Kac Correctors, a framework that allows for controlling the generated distribution of discrete masked diffusion models at inference time. We derive Sequential Monte Carlo (SMC) algorithms that, given a trained discrete diffusion model, control the temperature of the sampled distribution (i.e. perform annealing), sample from the product of marginals of several diffusion processes (e.g. differently conditioned processes), and sample from the product of the marginal with an external reward function, producing likely samples from the target distribution that also have high reward. Notably, our framework does not require any training of additional models or fine-tuning of the original model. We illustrate the utility of our framework in several applications including: efficient sampling from the annealed Boltzmann distribution of the Ising model, improving the performance of language models for code generation and amortized learning, as well as reward-tilted protein sequence generation.