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

TL;DR

Post-training autoregressive language models into masked diffusion models fundamentally reshapes internal computation rather than simply re-tuning parameters. Using a circuit-centric pipeline with automated discovery, EAP-IG edge analysis, and logit-lens/neuron explanations, the study contrasts ARM and MDM circuits across tasks demanding local causal vs global planning reasoning. The key finding is a mechanism shift: MDMS preserve autoregressive circuitry for locally dependent tasks but rewire early layers and distribute semantic roles for global planning, enabling non-sequential refinement. This mechanistic insight clarifies how diffusion objectives unlock bidirectional reasoning and informs design choices for diffusion-based language models in tasks requiring global coordination.

Abstract

Post-training pretrained Autoregressive models (ARMs) into Masked Diffusion models (MDMs) has emerged as a cost-effective strategy to overcome the limitations of sequential generation. However, the internal algorithmic transformations induced by this paradigm shift remain unexplored, leaving it unclear whether post-trained MDMs acquire genuine bidirectional reasoning capabilities or merely repackage autoregressive heuristics. In this work, we address this question by conducting a comparative circuit analysis of ARMs and their MDM counterparts. Our analysis reveals a systematic "mechanism shift" dependent on the structural nature of the task. Structurally, we observe a distinct divergence: while MDMs largely retain autoregressive circuitry for tasks dominated by local causal dependencies, they abandon initialized pathways for global planning tasks, exhibiting distinct rewiring characterized by increased early-layer processing. Semantically, we identify a transition from sharp, localized specialization in ARMs to distributed integration in MDMs. Through these findings, we conclude that diffusion post-training does not merely adapt model parameters but fundamentally reorganizes internal computation to support non-sequential global planning.

Figures (10)

• Figure 1: Overview of the mechanism shift analysis pipeline. We extract task-specific circuits for both the Autoregressive Model (ARM) baseline and the post-trained Masked Diffusion Model (MDM). We then identify the Top-$K$ components that exhibit the highest topological divergence between the two architectures. Finally, we interpret the algorithmic nature of these shifts using Logit Lens and Neuron Explanation.
• Figure 2: Circuit comparison across tasks and architectures. From left to right: IOI (Qwen2.5-7B), IOI (Dream-Base-7B), Countdown (Qwen2.5-7B), and Countdown (Dream-Base-7B).
• Figure 3: Layer-wise difference in unique attention component usage (MDM minus ARM). Rows:IOI (top) vs. Countdown (bottom). Columns: DiffuLLaMA vs. LLaMA-2 (left) and Dream vs. Qwen (right). Green regions ($>0$) indicate that the diffusion model utilizes more attention heads, while red regions ($<0$) indicate greater usage by the autoregressive model.
• Figure 5: step-wise circuit visualization of Dream on the Countdown task. Steps 1--12 are shown from left to right and top to bottom.
• Figure 6: step-wise circuit visualization of DiffuLLaMA on the Countdown task. Steps 1--12 are shown from left to right and top to bottom.