💬 NLP

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

This paper provides a mechanistic analysis of masked diffusion language models (DLMs), revealing that they implement a direction-symmetric bidirectional induction circuit for in-context learning and implicitly compute global masking fractions as timesteps, though they only outperform autoregressive models when leveraging full bidirectional context.

Original authors: Andy Catruna, Emilian Radoi

Published 2026-07-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Andy Catruna, Emilian Radoi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to finish a sentence. For years, the best way to do this was to make the robot read a story from left to right, word by word, guessing the next one based only on what it had already seen. This is how most modern AI writers work; they are like a person reading a book who can only see the pages they've already turned. But recently, scientists have started building a different kind of robot. Instead of reading left-to-right, this new robot sees a whole page of text where some words are hidden behind black boxes. It looks at the visible words on both the left and the right, guesses what's inside the box, fills it in, and then repeats the process until the whole page is clear. This is called a "diffusion" model.

The big question for scientists is: How does this new robot learn? When you show a robot a pattern, like "The cat sat on the mat. The cat sat on the [BLANK]," and it figures out the blank should be "mat," we call that "induction." It's the robot's way of saying, "I've seen this before, so I know what comes next." We know exactly how the old, left-to-right robots do this trick, but we have no idea how the new, "look-at-everything" robots pull it off. Do they use the same brain circuits? Do they have a secret superpower because they can see both sides of the missing word? Understanding this is crucial because if we want to build better AI, we need to know how these different machines actually think, not just what they say.

In this paper, the researchers decided to play detective with two very similar robots to see how they solve the same puzzle. They built a standard "left-to-right" robot and a new "look-at-everything" robot, making sure they were twins in every way except for how they read. They gave them a simple game: find a repeated pattern in a long string of random letters and copy the missing letter from the matching spot.

The team discovered that the new robot doesn't just copy the old robot's tricks; it learns a brand new, two-way superpower. While the old robot can only look backward to find its answer, the new robot builds a special "induction circuit" that works in both directions. It's like the robot has two little helpers: one who whispers what was just to the left of the missing word, and another who whispers what was just to the right. These helpers write these clues into the robot's memory. Then, a third "detective" part of the robot uses those clues to scan the entire text—both the past and the future—to find the matching pattern and copy the answer.

Here is the twist: The new robot is only better than the old one when it is allowed to peek at both sides of the missing word. If you block its view of the future and force it to only look at the past (just like the old robot), it performs exactly the same. This proves the new robot isn't just a "smarter" version of the old one; its advantage comes entirely from its ability to see both sides of the gap at the same time.

The researchers also found something surprising about how the robot knows when to guess. Usually, these robots are told exactly what step of the guessing game they are on (like "Step 1 of 100"). But these robots weren't given that information. Instead, the paper shows that the robot secretly counts how many black boxes are still on the page. It uses this count as a hidden timer. If there are many boxes left, it acts one way; if there are few, it acts another. The scientists proved this by poking the robot's brain and showing that changing this "box count" directly changes how confident the robot is in its guesses.

In short, the paper reveals that these new diffusion robots don't just think differently; they think in a symmetrical, two-way street. They gather clues from both sides, match them up, and even keep a secret mental tally of how much work is left to do, all without being explicitly told how to do it. This gives us a clearer map of how these powerful new tools are learning to understand the world.

Technical Summary: Induction in Both Directions

Problem Statement

While the internal mechanisms of autoregressive (AR) transformers have been extensively studied through the lens of mechanistic interpretability, the operational mechanics of Diffusion Language Models (DLMs) remain largely unexplored. DLMs generate text via iterative denoising rather than sequential next-token prediction, utilizing bidirectional attention and different training objectives. A critical gap exists in understanding whether DLMs implement induction—a canonical mechanism for in-context learning where a model identifies a repeated context and copies the subsequent token—in the same manner as AR models. Specifically, it is unclear how DLMs handle induction when the target token is hidden behind a mask token [M] and surrounded by bidirectional context, as opposed to the strictly causal context available to AR models.

Methodology

The authors conducted a controlled mechanistic study comparing matched attention-only AR transformers and absorbing-mask DLMs.

Model Architecture: The study utilized small-scale models (depths $L \in \{1, 2, 3\}$ , hidden dimension $d_{model}=512$ , 16 attention heads) trained on FineWeb-Edu sequences. Both model families shared identical architectures and training recipes, differing only in their attention masks (causal for AR, bidirectional for DLM) and training objectives (next-token prediction vs. absorbing-mask diffusion).
Evaluation Task: Induction capabilities were measured using synthetic sequences of length 512, where the first half was repeated as the second half. The task involved predicting masked tokens in the second half based on the unmasked source in the first half (forward induction) and vice versa (reverse induction).
Analysis Techniques:
- Causal Ablations: Mean ablations were performed on individual attention heads to determine their causal necessity for induction.
- Weight Decomposition: The study employed QK (Query-Key) and OV (Output-Value) decompositions to analyze how attention logits are formed and how information flows from source to target.
- Linear Probing: A linear probe was trained on the residual stream to detect the presence of global mask fraction information.
- Patch Experiments: "Patching" techniques were used to causally intervene on the representation of the global mask rate to observe changes in prediction entropy.
Model Variants: For mechanistic analysis, the authors utilized "folded no-LayerNorm" DLMs, converting trained LayerNorm models to linear maps to facilitate additive decomposition of the residual stream, while verifying that performance remained comparable to the original LayerNorm models.

Key Contributions and Results

1. Bidirectional Induction Circuit

The study identifies that DLMs learn a bidirectional induction circuit that is direction-symmetric. Unlike AR models, which rely on a single causal pathway, DLMs utilize two distinct pathways:

Previous-Token and Next-Token Heads: In the first layer (Layer 0), specific heads attend to the immediate left and right neighbors of the masked token, writing these local context features into the residual stream.
Induction Heads: In subsequent layers (Layer 1+), induction heads utilize QK weights to match the features written by the Layer 0 heads at the masked position with the corresponding features at the source answer position.
Mechanism: The QK decomposition reveals that induction heads use Q-composition (query built from the mask's local context) and K-composition (key built from the source's local context) to locate the matching source. The OV path then copies the source token to the masked position.

2. Context Access vs. Mechanism Strength

The authors demonstrate that the superior induction performance of DLMs compared to AR models is not due to a fundamentally stronger induction mechanism, but rather to access to bidirectional context.

When restricted to left-context only (mimicking the AR setting), DLMs perform at or below the level of AR models.
DLMs outperform AR models only when both the left and right neighbors of the masked token are visible.
The induction score in the full bidirectional setting approximates the sum of the scores from the left-only and right-only settings, confirming that the advantage stems from utilizing information from both directions.

3. Implicit Timestep via Global Mask Rate

The paper provides causal evidence that DLMs compute and utilize the global fraction of masked tokens as an implicit timestep, despite lacking explicit timestep embeddings.

Representation: A linear probe can recover the global mask fraction from the residual stream after the first attention layer with high accuracy ( $R^2 \approx 0.82-0.91$ ).
Causal Utility: Patching the "mask-rate direction" from a high-corruption state into a low-corruption state causally increases the prediction entropy at the masked position. This confirms the model uses the global corruption level to modulate its uncertainty, effectively treating the mask rate as a timestep.

4. Circuit Robustness and Emergence

Emergence: Induction emerges abruptly during training (around 8–24k steps) and requires at least two layers to function, consistent with the need for one layer to write features and a subsequent layer to retrieve them.
Local Window: The induction mechanism relies heavily on a local window of approximately five tokens (the masked token plus two neighbors on each side). Performance drops significantly if these immediate neighbors are masked, even if the global corruption rate is low.

Significance and Claims

This work represents a foundational step in understanding the internal mechanisms of DLMs, moving beyond the assumption that they simply replicate AR behaviors. The authors claim that:

Mechanistic Divergence: DLMs implement induction through a symmetric, bidirectional circuit that differs structurally from the causal induction circuits in AR models, specifically in how query and key features are constructed from masked positions.
Contextual Advantage: The performance gap between DLMs and AR models in in-context learning is attributable to the availability of bidirectional context rather than an inherent superiority in the induction algorithm itself.
Implicit State Management: DLMs internally derive temporal information (timestep) from the global state of the sequence (mask fraction) without explicit embeddings, a mechanism that is causally linked to prediction entropy.

The paper concludes that these findings provide a basis for understanding in-context learning in non-autoregressive models, though it acknowledges limitations regarding the use of small, attention-only models and synthetic tasks, noting that further research is needed to determine if these mechanisms scale to larger, MLP-containing DLMs and natural language tasks.