Dynamic Context Selection for Noise-Robust Generation in Masked Diffusion Models
Takato Kinoshita ⋅ Masaki Asada ⋅ Makoto Miwa
Abstract
Masked diffusion language models (MDMs) condition on external context through static concatenation, exposing the model equally to relevant evidence and irrelevant distractor tokens throughout the entire reverse diffusion process. We identify this as a source of robustness degradation and propose **Dynamic Context Selection**, which *dynamically* selects context visibility at each denoising step. A lightweight Context Selector reads the current noisy state $x_t$ and the previous-step hidden states $\mathbf{h}_{t+1}$, and emits a per-token gate score that softly interpolates each context embedding toward the mask embedding, suppressing distractors while preserving gold evidence. The gate is trained using oracle gold/distractor annotations available in multi-hop QA through two complementary training-time components---an *oracle curriculum* that starts from oracle gating and gradually transfers control to the learned gate, and direct *oracle supervision* via binary cross-entropy---both of which are absent at inference. On four benchmarks spanning multi-hop question answering (HotpotQA, MuSiQue, 2WikiMultihopQA) and distractor-injected mathematical reasoning (GSM-IC), DCS consistently outperforms the static-context baseline on two MDM backbones (LLaDA-MoE and Dream-7B), with an especially pronounced gain on unanswerable questions where the static baseline struggles to abstain.
Chat is not available.
Successful Page Load