Technical Summary: Don't Retrain—Align: Adapting Autoregressive LMs to Diffusion LMs via Representation Alignment

Problem Statement

Diffusion Language Models (DLMs) offer distinct advantages over standard autoregressive (AR) models, particularly in non-sequential generation, bidirectional editing, and iterative refinement. However, scaling DLMs is computationally expensive. Theoretically, DLMs must learn $L!$ generation paths compared to the single left-to-right path of AR models, often requiring significantly more compute.

Recent approaches attempt to mitigate this cost by converting pretrained AR checkpoints into DLMs. While these methods successfully initialize DLMs from AR weights, they primarily rely on continued denoising training with modifications to objectives, attention patterns, or sampling conventions. Crucially, existing recipes do not explicitly constrain the converted DLM to preserve the internal representation geometry learned during the original AR pretraining. This raises a central question: Is it necessary to treat AR and DLM as disjoint paradigms requiring the relearning of language representations, or can the semantic structure learned by AR pretraining be explicitly preserved during conversion?

Methodology: REPR-ALIGN

The authors propose REPR-ALIGN, a minimalist adaptation method that treats DLM training as an alignment problem rather than a representation learning problem. The core hypothesis is that the "hard part" of language generation is learning semantic and syntactic structures (representations), which AR pretraining has already mastered. The remaining task is merely adapting these features to an iterative, any-order decoding mechanism.

The method involves two models with identical architecture:

1. Teacher: A frozen, pretrained autoregressive (AR) transformer with causal attention.
2. Student: A masked diffusion language model (DLM) with bidirectional attention, initialized from the same AR weights.

Training Procedure

During training, the student model consumes a corrupted sequence (with tokens masked) and attempts to predict the masked tokens using a standard masked denoising objective. Simultaneously, the frozen teacher model processes the clean sequence under causal attention.

To preserve the AR model's internal geometry, the authors introduce a layer-wise representation alignment loss. For every layer $\ell$ in the transformer stack, the hidden states of the student ($h^{(\ell)}_D$) are aligned with the hidden states of the teacher ($h^{(\ell)}_{AR}$) using a cosine similarity loss.

The total training objective is:
$$\mathcal{L}(\theta) = \mathcal{L}_{diff}(\theta) + \lambda \mathcal{L}_{align}(\theta)$$
Where:

• $\mathcal{L}_{diff}$ is the standard masked denoising cross-entropy loss.
• $\mathcal{L}_{align}$ is the mean cosine distance between student and teacher hidden states across all layers and masked positions.
• $\lambda$ is a weighting hyperparameter (default 10).

Key design choices include:

• No Adapters: The method requires no additional parameters or architectural changes beyond switching the attention mask from causal to bidirectional.
• Frozen Teacher: The AR teacher remains frozen, providing a stable representational coordinate system.
• Shift Convention: The method adheres to the standard shift convention used in AR-to-DLM adaptation to reconcile next-token indexing with position-wise denoising.

Key Contributions

1. Identification of Representation Preservation: The paper identifies that explicit preservation of the AR model's internal representation geometry is a missing ingredient in current AR→DLM conversion recipes.
2. A Minimalist Conversion Recipe: REPR-ALIGN introduces a simple, adapter-free method that anchors the DLM student to the frozen AR teacher via layer-wise cosine alignment. This approach improves conversion quality and sample efficiency across model scales.
3. Efficiency in Data and Compute: The authors demonstrate that AR→DLM conversion is not inherently data-hungry. With representation alignment, training on a tiny subset of data (0.8B tokens) can outperform training on the full dataset (50B tokens) under the same step budget. Furthermore, freezing large blocks (embeddings and MLPs) improves training throughput by up to 2$\times$ without degrading quality.
4. Scale-Up Validation: The authors trained a 4B parameter "oDLM" using this recipe. Compared to Dream-7B (a strong public DLM also built on AR initialization), oDLM achieved better HumanEval and HumanEval+ pass@10 scores despite using fewer parameters and a significantly lighter data-and-compute budget.

Experimental Results

The method was evaluated on code generation benchmarks (HumanEval, HumanEval+, MBPP, MBPP+) using Qwen3 checkpoints at 0.6B, 1.7B, and 4B scales.

• Adaptation Speed: REPR-ALIGN consistently accelerated adaptation. On HumanEval pass@10, the aligned model outperformed both AR fine-tuning and scratch training throughout early conversion. At 200k steps, the 0.6B model saw a 6.1 point gain, and the 1.7B model saw a 9.4 point gain over the baseline.
• Scaling: The benefits of alignment grew with model size. The 1.7B model showed larger absolute improvements than the 0.6B model, suggesting that larger AR pretraining geometries are more valuable to preserve.
• Data Efficiency: Training on a 0.8B-token random subsample with alignment yielded higher pass@10 scores than training on the full 50B-token stream at the same step count.
• Freezing: Freezing token embeddings and MLP blocks increased throughput by $\sim$2$\times$ while maintaining or slightly improving performance, supporting the view that the primary learning signal is in adapting the denoising mechanism, not relearning representations.
• Ablations:
  • Metric: Cosine alignment outperformed L2 distance, suggesting the useful signal is geometric (direction) rather than metric (magnitude).
  • Weight: $\lambda=10$ provided the optimal trade-off between anchoring and adaptation.
  • Layers: Aligning all layers provided the best pass@1, while middle and upper layers were particularly important for pass@10.

Significance and Claims

The paper claims that linguistic representations are universal regardless of generation order. By reframing DLM training as an alignment problem, the authors argue that the high cost of diffusion pretraining is largely due to the unnecessary relearning of representations that AR pretraining has already provided.

The significance of this work lies in offering a "go-to technique" for training diffusion language models that:

• Reduces Compute: Achieves competitive performance with a fraction of the standard pretraining cost.
• Simplifies Training: Removes the need for complex adapters or heavy fine-tuning, relying instead on a simple alignment objective.
• Democratizes Access: Makes high-quality DLM training feasible in low-data regimes and with constrained compute resources.

The authors position their work as a shift from "retraining" to "aligning," suggesting that future advances in discrete generation may come from better mechanism alignment rather than training paradigms from scratch.