Data-Efficient Autoregressive-to-Diffusion Language Models via On-Policy Distillation

Pretraining a diffusion language model (DLM) from scratch is expensive. Recent AR-to-DLM conversion methods reduce this cost by starting from pretrained autoregressive language models (ARLMs), but they still face two key mismatches: changing the training objective can weaken knowledge inherited from the ARLM, and standard DLM training uses randomly masked states that differ from the partially denoised states encountered during inference with confidence-based samplers.

OPDLM asks whether we can convert a pretrained ARLM into a DLM as a lightweight post-training procedure while preserving the capabilities learned during autoregressive pretraining and while addressing the divide between train and inference state distributions.

A natural tool for this is on-policy distillation (OPD), which trains a student on states sampled from its own generation process while using a teacher to provide token-level supervision. But applying OPD to DLM conversion creates a chicken-and-egg problem: a DLM student visits partially masked diffusion states, so a direct OPD setup would require a capable DLM teacher to score those states.

OPDLM bypasses this requirement by using the original frozen ARLM as the teacher. The student is initialized from the same ARLM weights, converted into a block-diffusion model, and trained on its own reverse diffusion trajectories. For each partially denoised student state, OPDLM uses the terminal denoised sequence to query the ARLM teacher on causal prefixes, producing token-level target distributions for the masked positions.

This creates a self-distillation setup:

• Teacher: the original frozen ARLM, queried for token-level distributions on causal prefixes of the student's generated sequence.
• Student: a block-diffusion LM initialized from the same ARLM, trained to predict masked tokens under blockwise bidirectional attention.

By training directly on the student's inference-time states while distilling from the original ARLM, OPDLM reduces the train-inference gap in DLM training and improves knowledge retention during conversion.

1. Roll out the student: sample a reverse unmasking trajectory from the current DLM and fixed sampler.
2. Select an on-policy state: choose a non-terminal partially denoised state from the realized trajectory.
3. Query the ARLM teacher: build causal prefixes from the terminal sequence and retrieve token-level teacher distributions.
4. Optimize KL: align the DLM prediction at masked positions with the frozen ARLM distribution.

Rollout-Length Curriculum

Early in training, the terminal sequences the student generates by running its reverse diffusion process are low-quality, since the converted student is being queried with masked-token inputs and blockwise bidirectional attention for the first time. To address this, OPDLM begins by generating shorter sequences and gradually increases their length, helping training stability and convergence.

General-Purpose DLM Results

OPDLM converts Qwen3 into a diffusion language model for general-purpose reasoning across knowledge, mathematics, science, and code. OPDLM-4B and OPDLM-8B reach performance competitive with existing DLMs while training on only 0.076B and 0.066B tokens, orders of magnitude fewer than the baselines, and at substantially lower FLOPs.

Zero-Shot Results

Zero-Shot Extended Thinking

Modern ARLMs can reason through a problem inside a <think>...</think> trace before committing to an answer. We never train OPDLM to do this, yet the converted DLM does it zero-shot: when prompted to think, OPDLM-8B improves on the hardest reasoning benchmarks, raising AIME-24 from 14.7 to 18.6 and AIME-25 from 12.4 to 19.4. The base ARLM's reasoning ability survives on-policy conversion intact, emerging as a capability we never explicitly trained for.

Multilingual Results

OPDLM keeps the multilingual ability of the base ARLM after conversion. without any multilingual-specific training. It holds performance across MMMLU-lite, INCLUDE-lite, and MLogiQA, and even improves on multilingual Math (MT-AIME 2024).

Specialized DLM Results

Since OPDLM is a form of post-training applied to ARLMs, we can also build specialized DLMs. Below, we train OPDLM specifically for math to obtain OPDLM-MATH, using the same on-policy distillation setup. Additionally, we train OPDLM-MATH-Thinking for extended reasoning.

We show two controls on inference throughput: lowering the decoding threshold lets OPDLM produce more tokens per denoising step, while the training block size sets the upper bound on the parallelism it can expose at inference time.

OPDLM converts a pretrained Qwen3 autoregressive model into a BD3LM student with on-policy distillation. Data and model artifacts are hosted in the divelab/opdlm Hugging Face collection.

Environment

git clone https://github.com/divelab/OPDLM.git
cd OPDLM

conda create -n opdlm python=3.10.19 -y
conda activate opdlm

# Install torch first.
pip install torch==2.6.0+cu124 --index-url https://download.pytorch.org/whl/cu124

# Install project dependencies.
pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu124

# Install flash-attn last so it builds against the active torch install.
pip install flash-attn==2.7.4.post1 --no-build-isolation

If DeepSpeed rejects a CUDA 12.x minor-version mismatch while compiling CPU Adam, set DS_SKIP_CUDA_CHECK=1.

Data

# Evaluation data: 19 of the 20 paper benchmarks.
huggingface-cli download divelab/opdlm_eval_data --local-dir data/ --repo-type dataset

# Training data: opdlm_train.json, 61,816 rows.
huggingface-cli download divelab/opdlm_train_data --local-dir data/ --repo-type dataset

# Paper eval and DAPO math data that live outside the OPDLM collection.
python data/prepare_codeforces.py
huggingface-cli download BytedTsinghua-SIA/DAPO-Math-17k --local-dir data/ --repo-type dataset

Models

# Teacher ARLMs.
huggingface-cli download Qwen/Qwen3-4B --local-dir $HF_HOME/Qwen3-4B
huggingface-cli download Qwen/Qwen3-8B --local-dir $HF_HOME/Qwen3-8B

# Student initializations with bidirectional attention.
huggingface-cli download divelab/Qwen3-4B-a2d-init --local-dir $HF_HOME/Qwen3-4B-a2d-init
huggingface-cli download divelab/Qwen3-8B-a2d-init --local-dir $HF_HOME/Qwen3-8B-a2d-init

Smaller Qwen3-0.6B and Qwen3-1.7B init models can be regenerated with convert_qwen_to_bd3lm.py.

Train

python rl.py config=configs/rl_bd3lm.yaml \
    model.pretrained_model=$HF_HOME/Qwen3-4B-a2d-init \
    model.teacher_model=$HF_HOME/Qwen3-4B \
    dataset.train_dataset=opdlm_train

Reference launchers with the paper hyperparameters live in scripts/general_pre_train/ and scripts/post_train_dapo/. Edit DATA_PATH, STUDENT, TEACHER, and the SBATCH header for your cluster.

Evaluate

python pure_inference/eval.py \
    --models <path-to-your-trained-opdlm-ckpt> \
    --model_bases bd3lm \
    --datasets HumanEval MBPP MATH500 GSM8K AIME2024 \
    --max_token 2048 \
    --remasking_strategy low_confidence_static \
    --dynamic_threshold 0.9 \
    --temperature 0.0 \
    --block_size 4 --denoising_steps_per_block 4 \
    --out_dir pure_inference/results