Class posteriors enable control and interpretation in discrete diffusion models
Online Class Posteriors for Interpretable and Controllable Discrete Diffusion
Machine Learning
Summary
Discrete diffusion models produce outputs one step at a time but are less understood than their continuous versions. The authors create a way to track how likely different classes are during generation, offering clear insights into which parts of the output support or contradict a category. They also show how this tracking can help combine or exclude certain categories logically, giving precise control over what the model produces. Their approach was tested on data types like images, text, and antibody sequences.
What this means in practice
- •For natural language engineers: Control text generation by selectively encouraging or excluding specific token classes during masked language modeling tasks.
- •For bioinformatics developers: Guide discrete diffusion models in antibody sequence generation to include or exclude biological classes for targeted protein design.
Authors
Felix Koulischer, Gianni Van De Velde, François Remy, Chris Develder, Thomas Demeester
Abstract
Discrete diffusion models remain far less understood than their continuous counterparts, despite recent breakthroughs across modalities, most notably in text. Leveraging tools from continuous diffusion, we formulate a theoretical framework to track the class-conditional posterior in an online manner by combining conditional and unconditional model predictions. In masked language modeling, this posterior offers highly interpretable dynamics: tokens aligned with a class boost the posterior, while contradictory ones suppress it. We also show that the posterior can be used for logical guidance operations (AND, OR, NOT) in discrete diffusion. These operations enable precise control over generation, such as avoiding specific classes. In addition, we discuss a mismatch in the formulation of classifier-free guidance between discrete and continuous settings: discrete models mix the data distributions, while continuous models mix the marginals. We validate our methods and discuss their potential applications on synthetic data, with all quantities analytically tractable, as well as on MNIST, text, and antibody sequences.