Samson Gourevitch, Yazid Janati, Dario Shariatian, Umut Simsekli, Eric Moulines, Eric P. Xing, Alain Durmus
5 min
Abstract
Discrete diffusion models are often trained through clean-data prediction, but the prediction can be used in different ways to define the reverse dynamics. In Masked Diffusion Models (MDM) these choices largely coincide, whereas in Uniform Diffusion Models (UDM) they do not. We show that the standard plug-in bridge parameterization for UDM is not optimized by the denoising posterior, but by a leave-one-out posterior that predicts each clean token without using its own noisy observation. This identifies a mismatch between the plug-in ELBO and the usual cross-entropy denoising objective. We characterize the leave-one-out target and derive exact conversions between the denoiser, the leave-one-out posterior, and the score. These conversions allow us to disentangle parameterization and training objective. Our results also lead to inference improvements without any additional training through an informed predictor-corrector sampler and improved temperature sampling based on the leave-one-out predictor. We further introduce an absorbing-state reformulation of uniform diffusion that preserves the UDM joint law while decomposing it into masked-diffusion-like sampling operations, with simpler denoising posteriors, carry-over unmasking, and a natural remasking mechanism. On language modeling, leave-one-out parameterizations consistently improve UDM generation, while the absorbing construction matches or surpasses masked diffusion. These results suggest that the empirical gap between masked and uniform diffusion is driven less by the choice of marginals themselves than by parameterization and sampling design. The code and models can be found at https://github.com/samsongourevitch/rev_udm.
Alex: So the model was being trained to do one thing, but the math was quietly pushing it toward something slightly different.
Sam: That's a good way to put it. The stated goal and the actual optimisation target were misaligned. And once you see that gap, the path forward becomes clearer.
Alex: What does the fix actually look like in practice? Do you have to retrain everything from scratch?
Sam: That's the notable part — you don't. The authors provide conversion formulas that let you reinterpret what an existing model's predictions mean, without retraining it. You're essentially changing the lens through which you read the model's output, rather than rebuilding the model itself. They also introduced what they call an "absorbing-state reformulation," which is a way of restructuring the math so that uniform diffusion models behave more like the masked models that have been performing well.
Alex: So it's less like fixing the engine and more like realising you had the instrument panel wired up incorrectly.
Sam: That's a fair analogy. The underlying machinery was capable of more — it just wasn't being read correctly.
Alex: And the results held up?
Sam: Consistently, yes. By aligning the training with the leave-one-out target, generation quality improved across the board. The paper's central claim is that the performance gap between uniform and masked diffusion models was not fundamentally about the noise process — it was about how the prediction task was structured. Change that structure, and the gap narrows significantly.
Alex: Though I imagine it's not a perfect solution. There must be trade-offs.
Sam: There are. Asking the model to predict without seeing the local token is a harder task. In very large models, or in settings with many dimensions to track, the model doesn't always fully learn to ignore that local signal — even when the training objective tells it to. The researchers have a diagnostic for this, essentially a test for whether the model has genuinely learned the leave-one-out invariance or is still partly relying on the shortcut.
Alex: So the theory and the practice don't always line up cleanly.
Sam: Not always. There are architectural approaches that try to enforce this by design — structures that physically prevent the model from seeing its own input at a given position. But those designs tend to restrict how information flows through the model, which can create other problems. So for now, the authors suggest the most practical path is using these conversion formulas to align the objective, rather than imposing rigid architectural constraints.
Alex: It's a tension between mathematical purity and what actually works at scale.
Sam: Exactly. And what the paper ultimately demonstrates is that the most efficient gains don't always come from building bigger models or collecting more data. Sometimes the clearest improvement comes from asking more carefully what the model is actually being trained to do — and whether that matches what we want it to learn.
Alex: That's a useful reminder. The logic underneath the training loop matters just as much as the scale on top of it. Thanks for listening to ResearchPod.