Research map / Large language models / Reasoning and chain of thought

CORE: Context-Robust Remasking for Diffusion Language Models

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Test-time compute scaling · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes CoRe, a training-free inference-time remasking method for masked diffusion LMs that revises tokens whose likelihood collapses when nearby context is masked, improving reasoning and code benchmarks on LLaDA-8B-Base.

Problem

Masked diffusion language models commit to early unmasked tokens under sparse context ("context rigidity"), and existing remasking heuristics rely on static confidence or margin scores that miss context-brittle tokens, hurting structure-sensitive generation such as code.

Approach

Select revision targets by stress-testing unmasked tokens: mask a small margin-selected candidate set in one extra forward pass, score each token's negative log-likelihood under the perturbed context, and remask the most unstable tokens, framed as a robust (worst-case context shift) optimization.

Key results

• On LLaDA-8B-Base with Low-Confidence unmasking, CoRe raises MBPP from 15.60 to 24.80 (+9.2 points) and HumanEval from 12.20 to 17.07.
• ReMDM-conf drops MBPP accuracy by 6.4 points under Top-k Margin unmasking (21.20 to 14.80), while CoRe improves it to 29.60.
• Compute-matched Random and Margin Remask controls yield negligible gains (e.g., MBPP 16.60 and 17.40 vs. CoRe 24.80).

Surprising findings

• Standard confidence-based remasking (ReMDM-conf) degrades code generation performance, including MBPP under Top-k Margin unmasking.
• Margin-based targeting alone gives negligible gains, suggesting low uncertainty is not a proxy for context brittleness.

Limitations

• Gains on math reasoning (GSM8K) are modest or absent compared with code.
• Evaluated only on LLaDA-8B-Base with a fixed set of hyperparameters (N=128, m=32, E=8, k_rm=1).
• Improves internal structural consistency but does not guarantee external factual correctness.

Experimental setup

Benchmarks: GSM8K MATH BBH HumanEval MBPP

Metrics: accuracy pass@1 exact match strict-match accuracy

Baselines: Low-Confidence unmasking Top-k Margin unmasking ReMDM-conf Random Remask Margin Remask

Builds on: ReMDM LLaDA Top-k margin unmasking

Closest papers by meaning

• Re-evaluating Confidence Remasking in Masked Diffusion Language Models (NeurIPS 2026 · similarity 0.94)
• Remasking Discrete Diffusion Models with Inference-Time Scaling (NeurIPS 2025 · similarity 0.91)
• Learn from Your Mistakes: Self-Correcting Masked Diffusion Models (NeurIPS 2026 · similarity 0.91)
• Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models (NeurIPS 2026 · similarity 0.90)
• A Theoretical Analysis of Why Masked Diffusion Models Mitigate the Reversal Curse (NeurIPS 2026 · similarity 0.90)
• Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions (ICML 2025 · similarity 0.90)