Induction in Both Directions: A Mechanistic Analysis of In-Context Learning in Masked Diffusion Language Models
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
The authors studied how diffusion language models (DLMs), which generate text by gradually fixing noisy input, learn to copy repeated patterns—a process called induction—compared to traditional autoregressive (AR) transformers that generate text one token at a time. They found that DLMs create a two-way (bidirectional) system that can use context from both before and after a word to find and copy matching tokens, unlike AR models limited to past context. When DLMs only see past context, they perform similarly to AR models, but they do better when they can see both sides. The authors also discovered that DLMs keep track of how many tokens are masked globally, using this as a kind of internal timer without being explicitly told to do so.
Authors
Andy Catruna, Emilian Radoi
Abstract
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.