Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Learn from Your Mistakes".

Jane: Here's the gist of the paper. The gist Masked diffusion models face a fundamental limitation: once tokens are unmasked, they remain fixed, leading to error accumulation and ultimately degrading sample quality The proposed framework addresses this by training a model to perform both unmasking and correction How it works The authors propose a framework called Progressive Self-Correction (ProSeCo) that trains a model to perform both unmasking and correction.

Tom: This is achieved by reusing outputs from the MDM denoising network as inputs for corrector training, enabling the model to recover from potential mistakes During generation, additional corrective refinement steps are applied between unmasking ones in order to change decoded tokens and improve outputs The training involves adding a simple additional cross-entropy loss term to standard MDM objectives.

Jane: Self-Correcting Objective The high level idea of the approach is to train a corrector to minimize error between outputs from the denoiser and clean data The objective for training a corrector alongside the standard MDM denoiser is given by L CMDM, which includes an auxiliary error correction loss L C added to every term in the integrand This auxiliary loss amounts to a simple cross-entropy (CE) term, which encourages the model to identify and correct mistakes from the original denoising network The weighting scheme for this loss is set as λt:= λ(t) < 0 to control the relative weight between the correction and diffusion losses ProSeCo Design Decisions To have a unified model that can both decode and correct, the authors elect to tie the weights ϕ = θ, which eliminates memory overhead for training a separate corrector model The training objective is summarized in L SCMDM, which combines the standard MDM objective with the correction loss L C The weighting scheme λt:= α˙t/(1 − αt) is used because it discounts samples that are more heavily noised, as sequences with heavy masking which are harder to denoise will also be harder to correct and should therefore be down-weighted Sampling with ProSeCo The sampling algorithm interleaves corrective steps in between standard unmasking steps, allowing the model to dynamically refine and “self-correct" outputs In Algorithm 2, the corrector routine is invoked during sampling to potentially update already decoded positions in z 1:L t and provide improved predictions for the unmasking posterior The function corrector samples from x θ(z 1:L t) iteratively for S steps to generate y 1:L t, and then unmasked positions are replaced with corresponding positions in the corrector sequence Experimental Validation Extensive experimental validation across multiple conditional and unconditional tasks demonstrates that ProSeCo yields better quality-efficiency trade-offs (up to ∼4x faster sampling) and enables inference-time compute scaling to further increase sample quality beyond standard MDMs (up to ∼1.2x improvement on benchmarks) For guided generation, recovering from mistaken tokens improves the Pareto frontier of sample quality and property maximization Across all experiments, ProSeCo better trades-off quality and speed, and can also benefit from inference-time scaling to further increase sample quality The method significantly improves standard fine-tuning of large MDMs, enabling up to ∼4x faster generation without quality degradation, and up to ∼1.2x increase in benchmark accuracy In the context of guided generation, ProSeCo can recover from errors when guidance strength is increased, thereby improving the guided generation trade-off of maximizing a property of interest while producing a diverse set of samples The model also shows robustness to the choice of corrector sampling parameters, beating the best performing Baseline SFT results for all choices of ω and S at every level of decoding parallelism Future Directions In follow up work, the authors plan to explore the disentangling of the corrector and unmasking models via weight untying or with completely separate neural network backbones for each model Additionally, while they present a performant sampling algorithm, the ability to correct mistakes opens up the design space to more sophisticated schemes of jointly using corrector and unmasking steps which they leave to future work to explore In this work, the authors presented a framework for jointly training a diffusion model to unmask and selfcorrect They demonstrated that ProSeCo consistently outperforms vanilla MDMs and alternative corrector methods both in terms of speed-quality tradeoffs and in the ability to further scale inference-time compute for improved generation The key finding is that ProSeCo can vary number of corrector steps and attain comparable performance more efficiently (Ours: Balanced; orange star), or it maximize quality by further scaling inference-time compute (Ours: Max; blue star) This work holds the promise for positive impact by accelerating inference and improving quality of language models Data and Datasets The training data for supervised fine-tuning is a modified version of the Llama-Nemotron-Post-Training dataset

5: , comprising approximately 32 million samples across four core domains: mathematical reasoning, coding, science, and instruction following The authors also trained ProSeCo from scratch on the OpenWebText (OWT;

18: ) dataset for 1M steps Software and Resources The work utilizes various software packages including HuggingFace

62: , Hydra

64: , Fast-DLLM

63: , and RDKit

29: for chemical informatics The models were fine-tuned using 256 NVIDIA B200 GPUs, and all evaluation runs were conducted on 8 NVIDIA H100 GPUs Impact Statement Our work holds the promise for positive impact by accelerating inference and improving quality of language models While our method is also subject to the dangers of misuse present in language modeling, especially in the application to biological sequences, given the limited scale of models that we explore, the experiments described in this work do not pose any significant risks LLM Usage Disclosure Large Language Models (LLMs) were used to format and edit this manuscript and for formatting figures Additionally, integrated LLMs in the code editing software used while developing this work were used for auto-complete suggestions and in debugging LLMs were not used for the core development of our methodology, derivation of key results, nor the programming of our main code and experiments References

1: Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, and Volodymyr Kuleshov.

Tom: Block diffusion: Interpolating between autoregressive and diffusion language models.

Title and authors: Tom: So we're looking at this paper now called "Learn from Your Mistakes: Self-Correcting Masked Diffusion Models". It tackles a problem where masked diffusion models get stuck once they unmask a token, and that error just keeps adding up, which ruins the final output quality.

Jane: Right. So the main idea here is proposing a way to train an AI model to do both the unmasking and the correction simultaneously. It seems like they're trying to build in a way for the model to recover from its own mistakes during generation, instead of just letting those mistakes pile up as we go.

Lu: Exactly. They call this framework Progressive Self-Correction or ProSeCo, and it works by training a corrector model using the outputs from the main diffusion network as inputs for that corrector training. This allows the model to learn how to fix potential errors when it sees them during sampling.

Meng: But what does that mean for deployment? If we have to train a whole separate corrector model, isn't that just adding complexity and memory overhead? What’s the trade-off there?

Tom: Well, the authors address that by tying the weights together so they use the same parameters for both tasks, which eliminates any extra memory cost for having two separate models. They summarize their training objective as L SCMDM, which combines standard masked diffusion objectives with this new correction loss L C.

Jane: That correction loss is key here. They introduce an auxiliary error correction loss, L C, which encourages the model to spot and fix the mistakes made by the original denoiser network during training. This loss is essentially a simple cross-entropy term added to every part of their standard masked diffusion objective.

Lu: The authors use a specific weighting scheme for this loss called lambda t, which they set as lambda t less than zero, to control how much emphasis the model puts on the correction versus the original diffusion losses. They also use lambda t equal alpha dot t over one minus alpha t to discount samples that are more heavily noised during training.

Title and authors: Meng: Okay, so during generation, they aren't just doing unmasking once and you're done; they interleave corrective refinement steps between those unmasking steps to change the decoded tokens and improve the final output. That sounds like it adds complexity to the inference pipeline.

Tom: It does, but they claim this allows for a significant improvement in quality-efficiency trade-offs, mentioning up to four times faster sampling compared to standard masked diffusion models. They also suggest that this approach can benefit from inference-time scaling to further increase sample quality beyond what vanilla masked diffusion models can do.

Jane: That scaling aspect is interesting because it means you could potentially improve the generation quality by increasing the compute time at inference, which is a different way to handle performance than just making the model bigger or faster.

Lu: And for guided generation, they found that recovering from mistaken tokens actually helps improve the Pareto frontier of sample quality and property maximization. This means if you are trying to get a specific property, like a certain ring count in chemistry, this method helps you find better samples.

Meng: I'm wondering about the practical application on real traffic. If we have to run these extra corrective steps during sampling, how much of that speedup is actually realized when we’re dealing with production-scale inputs?

Tom: The experiments show that ProSeCo can significantly improve standard fine-tuning of large masked diffusion models, enabling up to four times faster generation without any quality degradation. And they also saw up to a one point two times improvement on benchmarks.

Jane: So, when we look at the results from this paper, it seems they've shown that you can balance quality and speed in a way that vanilla methods couldn't easily achieve with the same compute budget. They are also robust to the choice of corrector sampling parameters.

Lu: Their future work looks like they plan to explore disentangling the corrector and unmasking models, maybe using separate neural network backbones or weight untying, which would be a way to make it even more flexible in architecture.

Title and authors: Meng: I'm skeptical about that separation idea for now. From what I see here, we need something that actually runs well on a standard GPU setup without needing massive architectural changes just to get the benefit.

Lalam: Based on my analysis of the literature and this paper, ProSeCo consistently outperforms vanilla masked diffusion models and alternative corrector methods both in terms of speed-quality trade-offs and in the ability to further scale inference-time compute for improved generation. I'm very sure about those claims because they are backed by the results they show here.

Tom: Exactly. The key finding is that ProSeCo can vary the number of corrector steps and still get comparable performance, or it can maximize quality by further scaling inference-time compute as shown in their blue star result. This work holds the promise for accelerating inference and improving quality of language models.

Jane: So, to wrap up this discussion on "Learn from Your Mistakes: Self-Correcting Masked Diffusion Models", it’s clear they've given us a principled way to make diffusion models not just generate but also actively fix their own errors. It seems like the real power here is in how you can dial the number of corrector steps or scale up compute for a better result.

Lu: I think what they are showing is that treating model-generated outputs as corrupted versions of the true data where errors are noise that can be undone through a nested correction process fundamentally changes how we train these models to be more robust.

Meng: For us, the main thing is that it shows a way to get better results without needing a complete overhaul of our existing production infrastructure just yet. The ability to tune the compute scaling seems like something we can actually start thinking about for efficiency gains.

Lalam: I'm confident that ProSeCo offers a practical path forward by showing how to jointly train a diffusion model to unmask and self-correct, which is genuinely useful for making these models more reliable in practice.

Tom: It's interesting stuff. We can’t wait to see what the next set of papers brings.

The paper's summary: Tom: So, to recap what we just heard, this paper introduces ProSeCo which is basically training an AI model to fix its own mistakes while it’s unmasking tokens in a masked diffusion process.

Jane: Right. So instead of just letting the error pile up as the model goes from noisy pixels back to a clean image or text, they're teaching it how to use the results it already generated as an input for a correction step.

Lu: The core idea is this joint training objective, L SCMDM, which mixes the standard diffusion loss with this auxiliary error correction loss, L C. This forces the model to learn from its own bad outputs by treating those outputs as corrupted data that needs fixing.

Meng: But what does it mean for deployment? When we run this at scale on production traffic, does adding these extra correction steps actually give us a real speed boost, or are we just making things slower?

Tom: The authors claim that ProSeCo can be up to four times faster sampling compared to vanilla masked diffusion models. It’s not just theoretical; they show concrete numbers.

Jane: That speedup comes from interleaving corrective steps right between the standard unmasking steps during generation, which helps the model dynamically refine those tokens as it goes.

Lu: And they also suggest that we can actually scale inference-time compute to further increase sample quality beyond what standard diffusion models can do, up to a one point two times improvement on benchmarks.

Meng: Scaling compute at inference time is something I can get behind if it means better output without needing a whole new, massive training pipeline just to get the base model right. But how much extra compute are we actually talking about?

Tom: The key finding they highlight is that ProSeCo lets you control the quality-efficiency trade-off by changing how many corrector steps you run, or by scaling up compute for a higher quality result. You can even choose between a "Balanced" setting for a good middle ground, or a "Max" setting if you want the highest accuracy possible.

Jane: And it’s not just about speed and quality anymore; they also showed this works well for guided generation tasks too, where you're trying to make the model adhere to specific properties.

Lu: Yes, in guided generation, recovering from those mistaken tokens actually helps improve the Pareto frontier of what you can achieve—that is, maximizing a property of interest while still getting diverse samples.

Meng: So if we’re talking about real-world use cases where accuracy matters and inference time is tight, this looks like a way to get better results without having to completely redesign our entire infrastructure from scratch.

Tom: Exactly. This work suggests that we can build models that are inherently more robust because they learn how to fix their own errors during the sampling process, which is a big step forward for reliable AI deployment.

Jane: So, the main takeaway here is that ProSeCo provides a principled way to make diffusion models not just generate outputs but also actively fix those outputs as they happen.

Lu: It seems like treating model-generated sequences as corrupted versions of the true data where errors are noise that can be undone through this nested correction process fundamentally changes how we train these models to be more robust.

Meng: For us, the main thing is that it shows a way to get better results without needing a complete overhaul of our existing production infrastructure just yet. The ability to tune the compute scaling seems like something we can actually start thinking about for efficiency gains.

Tom: It's interesting stuff. We can’t wait to see what the next set of papers brings on this topic.

The paper's improvements: Tom: So, we're looking at how ProSeCo actually improves things beyond just having a correction mechanism in the training loop.

Jane: Right, so it’s not just about adding a loss term; it’s about how they structure the whole sampling process to make those corrections happen dynamically during generation.

Lu: The authors present Progressive Self-Correction sampling, which interleaves those corrective refinement steps between the standard unmasking steps in Algorithm two. This means you’re not just correcting at the end; you're constantly refining as you go.

Meng: How does that dynamic refinement actually translate to real-world performance improvements compared to just running a static correction step after the whole process is done?

Tom: They show that this approach can lead to up to one point two times improvement on benchmarks, which is significant because it means better quality without needing a massive increase in generation time.

Jane: That scaling capability at inference time is what I find really interesting for people who use these models in production pipelines; it lets you bump up the quality when you need it most.

Lu: And they also discuss how this method benefits guided generation tasks, where you’re trying to get a model to follow specific rules or properties.

Meng: So if I’m trying to get a chemical structure with a specific ring count, and I use this ProSeCo approach, the model is better at recovering from errors it makes during the guidance process.

Tom: It means that when you increase the guidance strength, you're less likely to get stuck on those wrong tokens because there’s a built-in mechanism to fix them mid-way through.

Jane: The authors are also very flexible here; they show that this method can be tuned by varying the number of corrector steps or by scaling up your inference compute for a better result.

Lu: And they found that you can choose between different configurations, like a "Balanced" setting for a good middle ground, or a "Max" setting if you want the highest accuracy possible.

Meng: So I can choose my setup based on whether I need to be fast and okay, or if I'm willing to spend more compute for better quality. That level of control over the trade-off is valuable for deployment.

Tom: The authors also demonstrated that this method is quite robust; they tested it against different corrector sampling parameters and found it beats baseline results across all those choices at every level of decoding parallelism.

Jane: It sounds like the main implication here is that we can now design diffusion models not just as a generation engine but as a system capable of self-repair during inference.

Lu: And their future work suggests they want to explore ways to separate the corrector and unmasking models, maybe with completely different neural network backbones or by untying the weights.

Meng: That separation idea sounds like it would make training much more flexible for me in the long run, if that actually works out cleanly.

Tom: The paper’s limitations are that they don't detail exactly how those corrector steps interact with other complex reasoning models, but for what they show, the results are solid.

Jane: So this work really pushes the idea that diffusion models can be made more reliable by giving them a built-in feedback loop to fix their own mistakes.

Conclusion: Tom: So to wrap up, ProSeCo gives us a way to train diffusion models to unmask and self-correct their own errors by using an auxiliary cross-entropy loss during training and interleaving correction steps during sampling.

Jane: It really boils down to giving the AI a mechanism where it learns from its own mistakes and fixes them as it generates outputs, which improves quality without needing massive retraining for every new task.

Lu: The key finding is that ProSeCo consistently outperforms vanilla masked diffusion models in terms of speed-quality trade-offs and its ability to scale inference-time compute to increase sample quality beyond what standard diffusion models can do.

Meng: I see how it helps the efficiency side, especially with that scaling potential, but I still wonder about the computational overhead of running those iterative corrector loops during actual live inference.

Tom: The paper shows it can be tuned by adjusting the number of corrector steps or scaling compute for a better result, so you get a quality-efficiency frontier you can navigate yourself.

Jane: That flexibility is what makes it practical for different use cases, from fast prototyping to high-fidelity generation where accuracy is paramount.

Lu: The authors are also looking ahead by suggesting they might separate the corrector and unmasking models, maybe with weight untying or different neural network backbones in future work.

Meng: If they can figure out a clean way to separate those two parts, that would be huge for deployment because it means I wouldn't have to manage two coupled training processes.

Lalam: I’m confident that ProSeCo offers a practical path forward by showing how to jointly train a diffusion model to unmask and self-correct, which is genuinely useful for making these models more reliable in practice.

Tom: It’s interesting stuff. The key finding is that ProSeCo can vary the number of corrector steps and attain comparable performance or maximize quality by scaling inference-time compute.

Jane: So this work really pushes the idea that diffusion models can be made more reliable by giving them a built-in feedback loop to fix their own mistakes during generation.

Lu: It seems like treating model-generated sequences as corrupted versions of the true data where errors are noise that can be undone through this nested correction process fundamentally changes how we train these models to be more robust.

Meng: For me, it means I have a better tool for balancing my budget and my quality needs right now, rather than just betting on the biggest model available.

Tom: It's interesting stuff. We can’t wait to see what the next set of papers brings on this topic.