Few-Step Diffusion Language Models via Trajectory Self-Distillation

[arXiv]

Overview: We proposed T3D (Trajectory Self-Distillation via Direct Discriminative Optimization), a trajectory self-distillation framework designed to enable high-quality few-step inference for Diffusion Large Language Models (DLLMs).

While DLLMs provide strong potential for parallel token generation, they often suffer from severe quality degradation when the number of denoising steps is aggressively reduced. T3D addresses this challenge by distilling generative trajectories from a full-step teacher model, significantly narrowing the performance gap between few-step and full-step decoding.

Key Ideas

• Label-Free Training: Learns entirely from self-generated trajectories without external labels.

• Trajectory Self-Distillation: Distills on-policy teacher rollouts to reduce train–test mismatch.

• Direct Discriminative Optimization: Uses reverse-KL-style optimization to focus on high-probability teacher modes.

Installation

1. Environment Setup

conda create -n dllm python=3.10
conda activate dllm

2. Install Dependencies

pip install torch==2.6.0
pip install --no-cache-dir \
  https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/\
flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
pip install -r requirements.txt
conda install -c nvidia cuda

3. (Optional) Configure Cache Directory

Useful when using NVMe or shared memory for faster Triton / Torch extension compilation.

export T3D_CACHE_ROOT=/path/to/cache   # default: ~/.cache/t3d

Data Preparation

Place dataset JSON files under:

data/

Example:

data/MATH_train.json

You can download datasets using:

cd data
python download_data.py --dataset MATH_train
python download_data.py --dataset MATH500
cd ..

Workflow

The training pipeline consists of three stages:

1. Rollout trajectory generation
2. Trajectory preprocessing
3. T3D training

Stage 1 — Rollout (Trajectory Generation)

Run from the sample/ directory.

Default behavior:

• Input data: ../data/
• Output trajectories: ../<experiment.project>/temp_data/

cd sample

python sdar_sample.py config=../configs/sdar_sample.yaml \
    evaluation.eval_dataset=MATH_train \
    evaluation.data_type=math \
    rollout.block_size=4 \
    rollout.denoising_steps_per_block=4 \
    rollout.max_token=2000 \
    rollout.temperature=0.1 \
    rollout.top_p=1.0 \
    rollout.top_k=0 \
    rollout.remasking_strategy="low_confidence_static"

cd ..

Stage 2 — Preprocess Trajectories

Move rollout file into data/:

mv path/to/rollouts.json data/SDAR-4B-Chat-MATH_train.json

Then preprocess:

cd data
python rename_key.py --model_name SDAR-4B-Chat --dataset MATH_train
cd ..

Stage 3 — Training (T3D)

(Optional) W&B Login

wandb login

Launch Training

accelerate launch \
  --num_machines 1 \
  --machine_rank 0 \
  --main_process_ip 127.0.0.1 \
  --main_process_port 8889 \
  --config_file accelerate_configs/1_node_8_gpus_deepspeed_zero1.yaml \
  train/self_ddo_sdar_full.py \
  config=configs/ddo_sdar_full_self.yaml

Outputs → experiments/

Training data location:

data/<config.dataset.optimization_data>.json

Checkpoint

We provide a checkpoint: SDAR-4B-Chat trained with T3D on MATH_train:

Evaluation

Modify:

configs/eval.yaml

Then run:

python eval.py config=configs/eval.yaml \
    evaluation.checkpoint_path=/path/to/checkpoint \
    evaluation.eval_dataset=MATH500 \
    evaluation.data_type=math \
    rollout.block_size=4 \
    rollout.denoising_steps_per_block=1 \
    rollout.temperature=0.1

Configuration Guide

Troubleshooting: CUDA / Compilation Issues

If compilation fails, try:

export CUDA_HOME="$CONDA_PREFIX"
export CUDACXX="$CONDA_PREFIX/bin/nvcc"

export PATH="$CONDA_PREFIX/bin:$PATH"
export CPATH="$CONDA_PREFIX/include:$CPATH"

export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$CONDA_PREFIX/lib64:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH

export LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:$LIBRARY_PATH

BibTeX

@misc{zhang2026fewstepdiffusionlanguagemodels,
      title={Few-Step Diffusion Language Models via Trajectory Self-Distillation}, 
      author={Tunyu Zhang and Xinxi Zhang and Ligong Han and Haizhou Shi and Xiaoxiao He and Zhuowei Li and Hao Wang and Kai Xu and Akash Srivastava and Chengzhi Mao and Hao Wang and Vladimir Pavlovic and Dimitris N. Metaxas},
      year={2026},
      eprint={2602.12262},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2602.12262}, 
}

License

The code is released under the MIT License.