Discrete Feynman-Kac Correctors

TL;DR

We address the lack of controllability over samples from discrete diffusion models by introducing Discrete Feynman-Kac Correctors (DFKC), a principled inference-time framework based on CTMCs and the Feynman-Kac formalism. DFKC transforms the forward marginals via temperature annealing, product/geometric averaging, or reward tilting and implements these via redesigned rate matrices, enabling unbiased sampling with Sequential Monte Carlo and resampling. The approach is training-free and applicable to masked discrete diffusion, with demonstrations on Ising model sampling, language-model code generation and amortized learning, and reward-guided protein sequence design. The results show efficient temperature control, improved task performance, and enhanced generation quality, pointing to broad practical impact for constrained, reward-aware discrete generation. Future work includes extending to joint continuous-discrete models and combining inference-time corrections with reward fine-tuning.

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.

Figures (8)

• Figure 1: Discrete Feynman-Kac Correctors allow sampling from annealed distributions, product (or geometric average), and reward-tilted distributions. Panel (a) depicts the schematic of DFKC compared to the standard inference of masked discrete diffusion. Panel (b) demonstrates how DFKC, given trained discrete diffusion models and the reward function, samples from modified distributions at inference time.
• Figure 2: Results for annealing on the Ising model.
• Figure 3: Amortized learning task: Mean squared error (MSE) between predicted and true parameters reported for DFKC (1 and 5 samples), and joint prompting, across different dataset sizes. ** indicates $p \leq 0.02$, * indicates $p \leq 0.05$ (one-sided Student's t-test).
• Figure 4: Accuracy on coding tasks, with standard error reported over 5 seeds.
• Figure A1: Increasing the number of SMC samples for DFKC improves over no SMC resampling; gain is largest with 4 or 8 samples. Taking the product has a lower (better) mean squared error (MSE) than joint prompting, and resampling with DFKC significantly improves this further.

Theorems & Definitions (25)

• Theorem 2.1
• Theorem 3.1
• Corollary 3.1
• Theorem 3.2
• Corollary 3.2
• Theorem 3.3
• Corollary 3.3
• Theorem B.1
• proof
• Theorem C.1