Mechanism Shift During Post-training from Autoregressive to Masked Diffusion Language Models

AI summary

Overview

Research area: Mechanistic interpretability of language models, specifically the internal computation of masked diffusion models (MDMs) created by post-training pretrained autoregressive models (ARMs).

Technical level: Advanced. The paper assumes familiarity with circuit analysis, edge attribution patching, logit lens probing, and diffusion-based text generation.

Scope: A circuit-level, component-level, and neuron-level comparison of two paired 7B ARM–MDM model families across four controlled diagnostic tasks, reporting where and how inherited autoregressive computation is preserved or reorganized.

What This Paper Is About

Post-training a pretrained autoregressive model into a masked diffusion model is a cheap alternative to training a diffusion language model from scratch, and it works well empirically. What is not known is whether the resulting model simply keeps using the autoregressive computation it inherited, or whether it rewires that computation for non-autoregressive, bidirectional generation. The authors compare ARMs and the MDMs post-trained from the same backbones at the level of task-specific circuits, influential components, and early-layer neurons, and they report that the answer depends on the structure of the task.

Key Contributions

1. A paired-model comparison framework. The authors compare each ARM with the MDM post-trained from the same autoregressive backbone (Qwen2.5-7B with Dream-Base-7B, and LLaMA-2-7B with DiffuLLaMA-7B), reducing confounds from model scale, initialization, and base architecture. They analyze matched top-attribution subgraphs at the edge, component, and layer levels using metrics including Jaccard overlap, a component score s(v), and a layer-wise Center of Gravity (CoG).

2. An entropy-based operationalization of task regime. They distinguish "causal" from "global" tasks using an MDM conditional-entropy gap, defined as H_MDM(x_t | x_<t) − H_MDM(x_t | x_∖t), where x_∖t is the sequence with x_t masked. Positive values indicate a global regime; non-positive values indicate a causal or prefix-dominant regime. The grouping is computed with Dream and its sign-based grouping is verified with LLaDA.

3. Evidence for a task-dependent mechanism shift. On prefix-dominant tasks, MDMs largely preserve inherited high-attribution pathways or change them only modestly. On globally constrained tasks, reorganization is substantially stronger and task-relevant computation shifts toward earlier layers. The authors frame this as a working "task-compatibility hypothesis."

4. Component-level and neuron-level characterization of the shift. Diagnostic probes indicate that ARMs rely more strongly on sharply specialized components, whereas MDMs exhibit weaker single-component specialization and more diffuse output-space alignment, with more even coverage across semantic categories in early-layer neuron explanations.

Main Findings

• Circuit reuse depends on task structure. Edge overlap between ARM and MDM top-1,000 attribution subgraphs was 0.193 [0.174, 0.214] for Qwen/Dream on IOI and 0.088 [0.074, 0.104] for LLaMA/DiffuLLaMA on IOI. On Countdown it dropped to 0.008 [0.005, 0.012] and 0.032 [0.025, 0.041] respectively. Greater-Than fell in between (0.290 [0.268, 0.313] and 0.054 [0.044, 0.066]), and Semantic Infilling was 0.225 [0.191, 0.263] and 0.018 [0.010, 0.031].

• Depth-wise localization shifts strongly on global tasks. Center of Gravity moved from 16.5 [16.1, 16.9] in the ARM to 4.8 [4.5, 5.1] in the MDM for Qwen/Dream on Countdown, and from 17.1 [16.6, 17.5] to 5.3 [4.9, 5.7] for LLaMA/DiffuLLaMA. Semantic Infilling showed the same direction (13.7 to 3.8, and 13.1 to 3.9). On IOI the CoG was comparatively stable (17.5 to 20.4; 18.2 to 19.8). Greater-Than showed a smaller earlier shift (18.7 to 14.4; 17.3 to 11.7).

• Shared edges on IOI are functionally important under ablation. Ablating the intersection of the ARM and MDM subgraphs reduced IOI accuracy to 28.1%, whereas ablating a random subset of non-shared edges of equal size retained 92.4 ± 1.2% accuracy across five random seeds.

• Depth patterns survive resampling and budget changes. Paired bootstrap ΔCoG intervals remained strongly negative for Countdown (−11.7 [−12.2, −11.1] for Q/D; −11.8 [−12.4, −11.2] for L/DL) and SI (−9.9 [−10.8, −9.0]; −9.2 [−10.1, −8.3]), while IOI showed no early-layer relocation (+2.9 [+2.4, +3.5]; +1.6 [+1.0, +2.1]). Recomputing with equal-size top-n subgraphs for n ∈ {50, 100, …, 1000} preserved the qualitative depth ordering, though the paper states it makes no robustness claim above 1,000 or below 50 edges.

• ARMs retain sharper single-component specialization. For IOI, median ΔLME values were 0.90 for Qwen and 0.89 for LLaMA, versus 0.31 for Dream and 0.05 for DiffuLLaMA. Although Dream placed person-name tokens among its top aligned candidates more often than Qwen (NameFrac@10 of 13.2% versus 7.9%, with DiffuLLaMA at 8.6% and LLaMA at 3.4%), the authors note that raw name frequency does not imply specialization.

• Reduced concentration generalizes to GT and SI. Qwen showed LogitGap values of 3.107 on GT and 3.050 on SI, whereas Dream's gaps were 0.085 and 0.042. DiffuLLaMA remained at the low-gap, high-entropy end, especially on SI (LogitGap 0.076; entropy 2.263).

• Early-layer neuron patterns differ in evenness. ARMs concentrated explained units in a smaller set (266 units) but distributed them unevenly across semantic categories (σ = 369.78). MDMs spread interpretable labels across more units (426) with lower variance (σ = 202.52). MDM early layers typically had roughly 50–200 active neurons per top category, corresponding to broad, genre-level cues. In IOI, many ARM early-layer neurons were strongly activated by descriptive adjectives and modifier-related tokens (e.g., 2,821 neurons in LLaMA); in Countdown the pattern shifted toward numerical or technical content (over 1,000 neurons).

• Cross-family differences are reported descriptively. Dream tends to preserve more of its Qwen backbone on causal tasks, while DiffuLLaMA shows greater structural decoupling from LLaMA across regimes. The authors state that the two pairs differ in backbone capability, pretraining history, post-training data, and diffusion-training recipe, so magnitude differences cannot be attributed to a single factor.

Methodology in Plain English

The authors take two pairs of models where the second member of each pair was created by further training the first with a diffusion objective on the same backbone. They then run both members of a pair on the same four tasks and ask whether they solve those tasks using the same internal wiring.

To do this, they use Edge Attribution Patching with Integrated Gradients to score every connection in the model by how much it contributes to the task output. They keep the 1,000 connections with the highest scores in each model, giving a matched-size subgraph for a fair comparison. They then compare three things: which specific connections survive in both models (edge overlap), which components receive the most attribution mass (a component score aggregated over incoming and outgoing high-attribution edges, with Top-K set to 100), and how the attribution mass is distributed across network depth (Center of Gravity, where a lower value means computation sits in earlier layers).

All of this is computed on 500 examples drawn from the pool each model pair both answered correctly, so the analysis describes successful computation on that shared-success distribution rather than failure cases or the full task distribution. Full evaluation sets contained 700 examples for IOI and Greater-Than, 2,500 for Countdown, and 1,000 for Semantic Infilling.

To understand what changed at a finer grain, they project the activations of the Top-100 components into vocabulary space using each model's own unembedding matrix (a component-wise logit lens), and they examine individual neuron activations in layers 0–4 by collecting maximally activating input contexts and generating automated explanations. They also run targeted ablations to check whether the structures they identified actually matter for task performance under the tested intervention protocols.

Why This Matters

Impact on research. The paper argues that diffusion post-training does not uniformly overwrite autoregressive mechanisms; it selectively preserves or reorganizes them depending on task structure. This suggests that mechanistic analyses built around a few localized, sharply specialized components may be insufficient for diffusion language models, and that ensemble-, interaction-, and routing-level views may be needed. The paired design also gives a template for studying other adaptation processes where the base model is known.

Real-world applications (framed as implications of the findings, not as claims tested in the paper):

• Choosing when to convert an existing autoregressive model into a diffusion model rather than training from scratch, given that post-training is described as cheaper than from-scratch diffusion training.
• Diagnosing where diffusion post-training has changed a model's behavior, for example when a converted model regresses on tasks that depend on left-to-right prefix processing.
• Designing post-training curricula and data mixes, since the authors report that reorganization strength tracks whether a task is prefix-resolvable or globally constrained.
• Auditing converted models for safety-relevant behavior, since a converted model may compute a behavior using different internal pathways than the original.

Industry relevance. Practitioners who convert pretrained autoregressive checkpoints into masked diffusion models can use the depth-wise and component-level diagnostics described here as a monitoring signal for how much internal computation has moved, rather than relying on benchmark accuracy alone. The paper notes that diffusion post-training retains parallel decoding, iterative refinement, and bidirectional attention, which are the properties that motivate the conversion in the first place.

Future Directions

• Disentangle the causes. The paired-model design captures holistic ARM-to-MDM adaptation, but the diffusion objective, denoising interface, post-training data, and optimization trajectory all change together. Separating these factors is left to future work.
• Test generalization beyond the four controlled tasks. The authors state that whether the patterns hold for open-ended linguistic tasks and naturalistic generation remains to be tested.
• Test settings requiring stronger non-local coordination. The paper names verb-final languages such as Japanese, Korean, Turkish, and Basque as candidate settings.
• Check whether the same patterns arise in diffusion language models trained from scratch, rather than obtained by post-training an autoregressive checkpoint.

Target Audience

Mechanistic interpretability researchers studying circuits, attribution methods, or diffusion language models; machine learning engineers and researchers considering post-training an autoregressive checkpoint into a diffusion model; and model-evaluation practitioners interested in diagnostics that go beyond benchmark accuracy. The paper is written for readers already comfortable with transformer internals, attribution-based circuit discovery, and diffusion objectives, and its conclusions are explicitly scoped to the four tasks and two 7B ARM–MDM families studied rather than to diffusion language models generally.

Authors’ abstract

Post-training pretrained autoregressive models (ARMs) into masked diffusion models (MDMs) provides an efficient route to diffusion language modeling, but it remains unclear whether the resulting models reuse inherited autoregressive computation or reorganize it for non-autoregressive generation. We compare two 7B ARM-MDM families across four controlled diagnostic tasks and find a task-dependent mechanism shift. On prefix-dominant tasks, MDMs largely preserve inherited high-attribution pathways or exhibit only modest changes in where computation occurs. On globally constrained tasks, the reorganization is substantially stronger, with task-relevant computation shifting toward earlier layers. This depth-wise pattern persists across prompt resampling, circuit budgets, and tested inference budgets, while targeted ablations support the functional importance of the identified structures under the tested intervention protocols. At the component level, diagnostic probes suggest that ARMs rely more strongly on sharply specialized components, whereas MDMs exhibit weaker single-component specialization and more diffuse output-space alignment. Together, these results suggest that diffusion post-training selectively preserves or reorganizes inherited computation according to task structure, rather than uniformly replacing autoregressive mechanisms.