Discrete Feynman-Kac Correctors

2026-01-15

AI Summary

This work addresses the challenge of flexibly controlling the output distribution in discrete diffusion models during generation. The authors propose a training-free, inference-stage control framework based on sequential Monte Carlo (SMC) methods, integrating ideas from Feynman–Kac path integrals to dynamically correct the generative distribution of discrete masked diffusion models. The approach enables temperature annealing, multi-conditional marginal distribution product sampling, and guidance via external rewards. It demonstrates significant effectiveness across diverse tasks, including efficient sampling from Ising model Boltzmann distributions, improved code generation, amortized inference, and the design of high-reward protein sequences.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• discrete diffusion models
• distribution control
• inference-time correction
• reward-guided generation
• masked diffusion

Innovation

• Discrete Diffusion Models
• Feynman-Kac Correctors
• Sequential Monte Carlo
• Distribution Control
• Reward-Guided Generation