PaperlayerView exact revision on arXiv

Source-linked AI summary

CORE: Context-Robust Remasking for Diffusion Language Models

Kevin Zhai, Sabbir Mollah, Zhenyi Wang, Mubarak Shah

arXiv:2602.04096v3cs.LG

TL;DR

Masked diffusion decoding can preserve early, context-dependent errors because static confidence signals overlook tokens that become brittle as context changes. CORE reframes remasking as context-robust revision through targeted masked-context perturbations, and it consistently improves benchmark performance at matched compute, including up to +9.2% on MBPP.

  • Problem

    Static confidence-based revision can miss context-brittle tokens because predictions that appear reliable under early partial context may become incompatible as surrounding context evolves.

  • Method

    CORE is a training-free remasking framework that stress-tests token predictions under masked-context perturbations and prioritizes unstable tokens using an efficient approximation to robust optimization.

  • Results

    +9.2% on MBPP is the largest reported gain, with CORE consistently improving generation across coding, reasoning, and other benchmarks at equivalent inference cost.

  • Takeaways & Limitations

    Context-aware instability assessment can allocate limited revision compute more effectively than standard confidence, random, or margin-based remasking.

  • Takeaways & Limitations

    CORE improves internal structural consistency, while external factual correctness remains an open direction.

Abstract

from arXiv · show

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 can appear confident to the model itself. We propose Context-Robust Remasking (CORE), a training-free framework for inference-time revision. Rather than trusting static token probabilities, CORE identifies context-brittle tokens by probing their sensitivity to targeted masked-context perturbations. We formalize revision as a robust optimization objective over context shifts and efficiently approximate this objective to prioritize unstable tokens for revision. On LLaDA-8B-Base, CORE delivers consistent improvements across reasoning and code benchmarks, outperforming compute-matched baselines and improving MBPP by up to 9.2 percentage points.

1 Introduction

Masked diffusion decoding can lock in early errors because tokens selected under incomplete context are rarely revisited. CORE replaces static confidence-based revision with context-robust stress tests, improving quality at matched inference cost, especially for code.

  • Background: MDMs iteratively predict masked positions from partially unmasked context and select positions to unmask.Early selections therefore condition later generation on incomplete context.
  • Motivation: Static confidence and top-k-margin heuristics can miss tokens that later become incompatible as surrounding context evolves.Such tokens may appear reliable under an ambiguous early context despite being context-brittle.
  • CORE: CORE stress-tests token predictions under targeted masked-context perturbations and prioritizes tokens with the largest support drop for revision.The method treats robustness to context shifts, rather than static uncertainty, as the revision criterion.
  • CORE: CORE provides an efficient training-free remasking framework using a lightweight stress test and a single forward pass over a candidate subset.This approximates a worst-case instability objective with minimal overhead.
  • Results: +9.2% on MBPP is the largest reported gain, while CORE also delivers consistent performance on reasoning benchmarks at equivalent inference cost.The strongest improvements occur on structure-sensitive code generation.

2 Related Work

Related work has advanced masked diffusion language modeling and generation, but revision strategies still rely on unreliable context-local signals or require additional training. CORE addresses this gap with context-aware inference-time revision.

  • Diffusion language modeling: Discrete diffusion language models evolved from early state-space formulations to scalable masked language modeling systems such as MDLM and LLaDA.Subsequent work also supports variable-length generation and latent dependencies.
  • Revision strategies: Existing revision methods use stale historical confidence, potentially misleading current probabilities, or additional training requirements.These approaches do not reliably assess whether tokens remain compatible with updated context.

3 Problem Formulation

Masked diffusion decoding is context-rigid: early tokens become anchors even when selected from sparse context, and standard uncertainty proxies become stale. The formulation therefore distinguishes brittleness from uncertainty by testing sensitivity to context changes.

  • MDM decoding: The diffusion state represents each response position as either a vocabulary token or [MASK], with masked positions progressively replaced over steps.At each step, the model computes token distributions for masked positions before the sampler chooses positions to unmask.
  • Context rigidity: Standard sampling treats unmasked tokens as immutable constraints, allowing early predictions from sparse context to anchor later generation.Suboptimal early tokens can force subsequent tokens into self-reinforcing inconsistencies.
  • Revision challenge: Inference-time revision can reset selected unmasked positions to [MASK], but identifying which tokens to revisit remains the central challenge.Existing low-confidence and small-margin signals reflect transient ambiguity at selection time.
  • Brittleness versus uncertainty: A token is stable when it remains strongly predicted after parts of its surrounding context are masked, and brittle when its probability collapses.This changes the criterion from past uncertainty to persistence under context perturbation.

4 Method

CORE identifies context-brittle tokens by masking selected surrounding positions, measuring instability under the perturbed context, and revising the most sensitive tokens. Its efficient approximation uses token contention to choose probe positions and performs selection and revision with limited overhead.

  • Context-Robust Remasking: CORE frames remasking as robust optimization over context perturbations, identifying tokens whose predictions are unstable when surrounding information is masked.The framework constructs worst-case perturbations and revises tokens exhibiting instability under those perturbations.
  • Revision Operation: CORE masks selected unmasked positions to create a perturbed context, then uses the resulting predictions to update the highest-instability tokens.The illustrated example changes “a” to “an” under the perturbed context, while base unmasking predicts the next token “icy.”
  • Instability Scoring: For each candidate token, the instability score measures the negative log-likelihood of its current token under the perturbed context.Higher scores indicate greater sensitivity to masked-context changes and higher priority for revision.
  • Efficient Approximation: The worst-case perturbation is approximated by a deterministic candidate subset because exhaustive search over all mask configurations is combinatorially intractable.The selected subset provides a certified lower bound on the worst-case instability objective.
  • Efficient Approximation: Candidate positions are chosen using the smallest top-2 probability margins, while revision itself is determined solely by instability under the perturbed context.This separates where CORE probes from which tokens it ultimately revises.
  • Efficient Approximation: A temperature-scaled Softmax allocates the expected masking budget across currently unmasked positions, and the final probabilities are clipped at 1.In the implementation, τ →0 deterministically selects the m smallest-margin positions.

5 Experiments

Experiments show that CORE improves benchmark performance by selecting context-brittle tokens for revision, with especially strong gains on structurally constrained code tasks. Analyses indicate that instability-based selection, compute-matched targeted revision, and balanced revision settings drive these improvements.

  • Benchmark results: CORE consistently improves performance across BBH and MATH, with its largest gains on code-generation benchmarks such as MBPP.The evaluation compares Low-Confidence, Top-k Margin, and ReMDM-conf baselines under compute-constrained greedy decoding.
  • Benchmark results: +9.2% on MBPP, while ReMDM-conf drops MBPP accuracy by 6.4% under Top-k Margin.These results contrast CORE’s gains with the degradation observed for a confidence-based revision baseline.
  • Selection signal: Instability-based selection, unlike random or margin-based remasking, produces consistent gains under the same m=32 and E =8 settings.Compute-matched controls show negligible gains for random and margin-based target selection, while CORE improves notably on MBPP.
  • Compute matching: Under a matched budget of 136 forward passes, targeted revision outperforms simply adding decoding steps, while ReMDM-conf degrades on MBPP.CORE uses 128 decoding steps plus 8 auxiliary passes; baselines use all 136 passes for decoding or confidence-based revision.
  • Efficiency and sensitivity: m=32 and E =8 provide the best accuracy–compute trade-off, whereas m=64 degrades performance by increasing false-positive remasking.More frequent revision doubles revision overhead for only ≈0.6% MBPP gain, while less frequent revision reduces revision opportunities.
  • Instability analysis: Instability scores separate stable tokens near zero from brittle revised tokens in a heavy tail, targeting a specific outlier population rather than general uncertainty.The proxy is computed after simultaneously masking candidate positions and approximates the distributionally worst-case objective.
  • Qualitative analysis: CORE resolves structural inconsistencies by revising context-brittle operators, brackets, or other tokens that standard decoding leaves locked into broken contexts.The qualitative examples describe recovery from syntax errors and selective correction of tokens whose incompatibility emerges as context evolves.

6 Conclusion

CORE mitigates context rigidity in masked diffusion language models by framing revision as robustness optimization. It selectively revises brittle tokens, improving generation accuracy with minimal computational overhead while leaving external factual correctness as future work.

  • 6 Conclusion: CORE targets brittle tokens that disrupt structural and contextual consistency through robustness-based revision.The framework combines efficient margin-based screening with instability assessment.
  • 6 Conclusion: CORE improves generation accuracy with approximately 6% more forward passes.The reported overhead is characterized as minimal computational overhead.
  • 6 Conclusion: Guaranteeing external factual correctness remains an open direction beyond CORE’s demonstrated internal structural consistency.

A Notation

The notation used throughout the paper is summarized in Table 5. The table is identified as a summary of notation.

  • A Notation: Table 5 summarizes the notation used throughout the paper.
  • A Notation: The notation section serves as a reference for symbols used across the paper.
  • A Notation: The notation summary is presented as a dedicated table.

B Additional Implementation Details

The implementation details specify hardware, numerical precision, remasking limits, and benchmark evaluation protocols. These settings are intended to support numerical stability, efficiency, and rigorous evaluation.

  • Hyperparameters and Settings: Experiments use single NVIDIA H100 GPUs with bfloat16 precision.The stated purpose is numerical stability and efficiency.
  • Hyperparameters and Settings: The remasking limit is fixed to krm = 1 across all experiments.This keeps interventions minimally invasive and isolates the selection mechanism’s impact.
  • Evaluation Protocols: GSM8K uses strict-match accuracy, while MATH uses Minerva-style evaluation with Math-Verify.The protocols are described as rigorous benchmark evaluation procedures.

B.1 ReMDM Evaluation

The ReMDM-conf comparison uses a revision-based decoding setup aligned with CORE, holding timing, candidate selection, and compute allocation comparable. The main difference is stale-confidence selection versus context-incompatibility selection.

  • B.1 ReMDM Evaluation: ReMDM-conf and CORE share revision timing, candidate sets, and compute allocation.Both operate within the same mid-trajectory window and use matched decoding structures.
  • B.1 ReMDM Evaluation: ReMDM-conf samples lower-confidence tokens using cached probabilities from when tokens were last unmasked.These cached values are called stale confidence because they do not reflect the evolved context.
  • B.1 ReMDM Evaluation: CORE selects revisions based on incompatibility with evolved context rather than stale confidence.

C Theoretical Consistency: Computed Instability Lower-Bounds Worst-Case Risk

CORE’s computed instability score provides a certified lower bound on a token’s worst-case instability, even when the selected perturbation set is not optimal.

  • The framework defines perturbations by masking subsets S of eligible unmasked non-prompt indices C_t and measuring token instability under the resulting context.The perturbation context is denoted ỹ^(t)(S), and Y_j represents the token at position j.
  • CORE evaluates instability on a feasible candidate set S_t with |S_t| ≤ m using one auxiliary pass.The selected set need not maximize the robust objective.
  • A computed score ℓ_t(j; S_t) lower-bounds index j’s worst-case instability under size-m perturbations.Because S_t is feasible in the maximization, the computed instability cannot exceed the corresponding worst-case value.

D Sensitivity to Stochastic Decoding

CORE improves decoding consistency under stochastic sampling, with gains across symbolic reasoning and coding benchmarks beyond the baseline’s error margins.

  • 25.24% mean MBPP performance exceeds the baseline’s 18.40% by nearly 4 standard deviations under stochastic decoding.The evaluation uses temperature 1.0 and five random seeds.
  • CORE consistently improves logic and code benchmarks BBH, HumanEval, and MBPP beyond the baseline’s error margins under temperature 1.0 across 5 seeds.Table 6 reports mean ± standard deviation for the stochastic-decoding comparison.

E Qualitative Examples

Qualitative examples show CORE revising context-brittle tokens across mathematical reasoning, output formatting, and programming, while ReMDM-conf often remasks irrelevant tokens.

  • Mathematical reasoning and formatting: CORE changes the erroneous GSM8K answer 151 to 51 by replacing the extra leading 1 with a space.ReMDM-conf instead focuses on the unrelated <|endoftext|> token and fails to resolve the formatting error.
  • Mathematical reasoning and formatting: CORE progressively corrects a malformed mathematical expression and then revises the mistaken answer -6 to -2.The staged revisions repair both an intermediate syntax error and the final answer, whereas ReMDM-conf leaves the incorrect trajectory unresolved.
  • Programming examples: CORE sequentially removes an invalid newline and an unwanted equal sign when they become identifiable in a programming solution.ReMDM-conf instead remasks a stable or nearby token without fixing the structural flaw.
  • Programming examples: CORE repairs a corrupted function signature by revising the context-brittle :) token to num, producing a coherent program.ReMDM-conf repeatedly masks non-erroneous tokens, leaving the code syntactically incorrect.
Loading 2602.04096v3…