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: Welcome to another episode of ResearchPod. Today, we're looking at how AI models generate text — specifically a technique called discrete diffusion, and why one version of it has been underperforming for reasons that turned out to be surprisingly fixable.
Sam: Right. There are two main flavours of these models. One is called a Masked Diffusion Model, and the other is called a Uniform Diffusion Model. Masked models have consistently outperformed uniform ones — and for a long time, researchers assumed the difference came down to the type of noise each model uses. This paper argues that assumption was wrong.
Alex: So the noise wasn't the real problem?
Sam: That's what the research suggests. The actual issue was more subtle — it was about how the model was being asked to make its predictions during training. And to understand why that matters, it helps to think about what these models are actually doing.
Alex: Walk me through it.
Sam: Okay. Imagine you're learning a language by filling in blanks. Someone gives you a sentence with one word scrambled or hidden, and your job is to guess what the original word was. If you're good at this, you've genuinely learned how words relate to each other — grammar, meaning, context. Now imagine a different version of the same exercise: the scrambled word is still visible to you, right there in the blank. You can just copy it. You'd score well on the exercise, but you haven't actually learned anything about the language.
Alex: You're gaming the test rather than learning the material.
Sam: Exactly. That's essentially what was happening with Uniform Diffusion Models. During training, the model was being shown a noisy version of a word — say, a scrambled or randomly replaced token — and then asked to predict what the original clean word was. But the noisy version was still visible at that position. So the model could lean on that local signal rather than learning the deeper structure of the surrounding text.
Alex: And the fix was to take that shortcut away?
Sam: Yes. The researchers argue that the model should be forced to predict each word using only the surrounding context — the words before and after — not the noisy version sitting right there at that spot. They call this the "leave-one-out" approach, because you're leaving out the local observation at the position you're trying to predict. When the model can't cheat, it has to genuinely learn how language fits together.
Alex: That's a meaningful distinction. But if this is the right approach, why wasn't it the standard from the beginning?
Sam: Partly habit, partly simplicity. The default approach — predicting based on the full input including the noisy token — is the most straightforward thing to write down mathematically. The leave-one-out version requires a more careful understanding of what the training objective is actually rewarding. The researchers found that the standard approach was, in a precise mathematical sense, optimising for the wrong goal.
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.