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

Pretraining a diffusion language model (DLM) from scratch is expensive, and existing open DLMs still trail autoregressive language models (ARLMs) of comparable scale on standard benchmarks. Rather than start over, we ask whether the capabilities already learned by a pretrained ARLM can be transferred to a DLM. A natural candidate is on-policy distillation (OPD), which supervises a student on its own rollouts and has proven effective in post-training. This motivates the central research question of our work:

Can we convert a pretrained ARLM into a DLM while preserving its prior? And can OPD do it?

Applying OPD here runs into a chicken-and-egg problem, i.e, the teacher needs to be a capable DLM in order to score the masked, partially masked states the student visits, but a capable DLM is exactly what we are trying to build. OPDLM bypasses this by querying the ARLM directly as the teacher, reading out its prior through causal prefixes of the student's rollouts. This gives a self-distillation setup:

• Teacher: the frozen pretrained ARLM, queried only for token-level distributions over causal clean blocks.
• Student: a block-diffusion LM initialized from the same ARLM weights, trained to predict masked tokens under blockwise bidirectional attention.

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 <opdlm-repo-url>
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