Induction in Both Directions: A Mechanistic Analysis of In-Context Learning in Masked Diffusion Language Models

Andy Catruna ,

Emilian Radoi

LLM Other LLM GenAI Diffusion ML DM LM AI Safety / AI Ethics XAI

尽管自回归（AR）Transformer的内部机制已被广泛研究，但关于扩散语言模型（DLM）——一种通过迭代去噪生成文本的新兴替代范式——的内在工作机制，目前所知甚少。本文聚焦于DLM如何实现“归纳”（induction）这一机制；该机制是上下文内学习（in-context learning）的核心环节，即模型识别出重复出现的上下文，并复制紧随其后的词元（token）。我们的分析对比了仅依赖注意力机制的自回归模型与采用吸收掩码（absorbing-mask）策略的DLM，二者在模型架构上严格匹配。我们发现，DLM习得了一种双向归纳回路（bidirectional induction circuit）：其中“前序词元头”（previous-token heads）与“后续词元头”（next-token heads）将局部上下文写入残差流（residual stream），随后的归纳头则利用这些信息，在匹配的源位置定位并复制答案。该回路具有方向对称性——无论源上下文出现在待预测词元之前还是之后，该机制均能有效运作。当仅向模型提供左侧上下文（即与AR模型所见一致）时，DLM在归纳能力上并不优于其AR对应模型；然而，当被掩码词元的左右两侧上下文均可见时，DLM展现出更强的归纳能力，这表明其优势源于对双向上下文的访问能力，而非单侧机制本身的增强。除归纳机制外，我们还提供了因果性证据，表明DLM能够计算当前输入中被掩码词元所占的全局比例，并将其作为隐式的“时间步”（timestep）加以利用——尽管模型并未接收任何显式的时间步嵌入（timestep embedding）。

While the internal mechanisms of autoregressive (AR) transformers have been studied extensively, much less is known about diffusion language models (DLMs), an emerging alternative that generates text by iterative denoising. In this work, we study how DLMs implement induction, a mechanism behind in-context learning in which the model finds a repeated context and copies the token that followed it. Our analysis compares attention-only AR models and absorbing-mask DLMs with matched architectures. We find that DLMs learn a bidirectional induction circuit, where previous-token and next-token heads write local context into the residual stream and later induction heads use it to find and copy the answer from the matching source position. The circuit is direction-symmetric, working whether the source appears in the past or in the future. When only left context is visible, matching what an AR model sees, the DLM does not outperform its AR counterpart in induction capabilities. However, we observe it has stronger induction when both sides of the masked token are visible, pointing to bidirectional context access rather than a stronger one-sided mechanism. Beyond induction, we provide causal evidence that DLMs compute the global fraction of masked tokens and use it as an implicit timestep, even though they are given no explicit timestep embedding.

许愿