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.
翻译:离散扩散模型最近作为一种有前景的替代方案出现,用于生成离散序列,以取代自回归方法。通过逐步去噪或去掩码过程的样本生成,使它们能够捕捉数据中分层的非序列相互依赖关系。然而,这些定制过程并未对生成样本的分布提供灵活的控制。我们提出了离散费曼-卡克校正器,这是一个允许在推理时控制离散掩码扩散模型生成分布的框架。我们推导了序贯蒙特卡洛算法,该算法在给定一个已训练的离散扩散模型的情况下,能够控制采样分布的温度(即执行退火),从多个扩散过程(例如不同条件的过程)的边缘分布乘积中采样,以及从边缘分布与外部奖励函数的乘积中采样,从而从目标分布中生成既具有高可能性又具有高奖励的样本。值得注意的是,我们的框架不需要训练任何额外模型或对原始模型进行微调。我们在多个应用中展示了我们框架的实用性,包括:从伊辛模型的退火玻尔兹曼分布中进行高效采样,提高代码生成和摊销学习的语言模型性能,以及奖励倾斜的蛋白质序列生成。