Title: Discrete Feynman-Kac Correctors
Primary area: generative models
Keywords: discrete diffusion, protein generation, language models, sequential monte carlo

Abstract: Discrete diffusion models have recently appeared 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 the product of the marginal with an external reward function producing likely samples from the target distribution that have high reward at the same time. Notably, our framework does not require any training of additional models or finetuning of the original model. We illustrate the utility of our framework on several applications: the efficient sampling from the annealed Boltzmann distribution of the Ising model, extending the context of language models for amortized learning and multi-constraint generation, as well as reward-tilted protein sequence generation.

TL;DR: Sequential Monte Carlo framework for controlling discrete diffusion models at inference time