CORE: Context-Robust Remasking for Diffusion Language Models

Standard decoding in Masked Diffusion Models (MDMs) is hindered by context rigidity: tokens are retained based on transient high confidence, often ignoring that early predictions lack full context. This creates cascade effects where initial inconsistencies misguide the remaining generation. Existing revision strategies attempt to mitigate this by relying on static confidence scores, but these signals are inherently myopic; inconsistent tokens frequently appear confident to the model itself. To address this, we propose Context-Robust Remasking (CORE), a training-free framework for inference-time revision. We introduce a new selection paradigm: rather than trusting static token probabilities, we identify *context-brittle* tokens by probing their sensitivity to adversarial perturbations. We formalize revision as a robust optimization problem targeting worst-case context shifts. CORE efficiently approximates this objective using theoretically bounded probability margins to expose and revise unstable tokens. On LLaDA-8B-Base, CORE delivers consistent improvements across reasoning and code benchmarks, outperforming compute-matched baselines and boosting performance on code generation (MBPP) by up to $+9.2\\%$, with comparable gains generalizing to the Dream-7B architecture.

Problem: Traditional AI writes text left-to-right, but newer diffusion models write by guessing words across an entire document simultaneously. While fast, this parallel approach has a major flaw: the AI commits to an incorrect early guess before the surrounding sentence is finished. Because the AI operates with high confidence in this mistake, it builds the rest of the computer code or text around that flawed foundation, causing a cascade of errors. Solution: We developed Context-Robust Remasking (CORE) to fix these errors during generation. CORE acts as an active stress test. Instead of trusting the AI's initial confidence, CORE temporarily hides the surrounding text to check the stability of a specific word. If the AI changes its prediction when the surrounding context is hidden, CORE flags the original word as unstable and forces the model to revise it using the most recent information. Impact: By targeting words that are vulnerable to changing context, CORE stops structural errors before they propagate. This self-correction boosted the accuracy of AI-generated computer code by over 9%, ensuring these parallel models produce reliable outputs.