This software project accompanies the research paper, Learning Unmasking Policies for Diffusion Language Models.

To summarize the work very briefly: Diffusion LLMs generate text by iteratively unmasking tokens. Most prior work use heuristics to decide which tokens to unmask at each step. This work instead learns a lightweight policy (via GRPO) that makes these decisions autonomously, formalizing the unmasking problem as Markov Decision Process where the frozen dLLM serves as the environment. The policy, which we implement as a DiT-style, single-block transformer, observes token confidences and outputs per-position unmasking probabilities. If you are interested in learning more, please follow the link to the paper above.

If you find this work useful, please cite the paper:

@misc{jazbec2025learningunmaskingpoliciesdiffusion,
      title={Learning Unmasking Policies for Diffusion Language Models},
      author={Metod Jazbec and Theo X. Olausson and Louis Béthune and Pierre Ablin and Michael Kirchhof and Jo\~ao Monteiro and Victor Turrisi and Jason Ramapuram and Marco Cuturi},
      year={2025},
      eprint={2512.09106},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2512.09106},
}

The code requires Python 3.12. To install:

uv pip install -e .   # or just `pip install -e.` if uv is not available

You will also need to set your Hugging Face token for model access:

export HF_TOKEN=<your-token>

To train a policy on LLaDA-8B-Instruct with the GSM8k+MATH mixture:

python -m train.train --config configs/experiment_configs/llada_8b_instruct_dit_confidence_BL32_mixture.yaml

To evaluate a trained policy:

python -m eval.pipeline ./outputs/my_experiment \
    configs/experiment_configs/llada_8b_instruct_dit_confidence_BL32_mixture.yaml \
    --checkpoints last \
    --datasets gsm8k \
    --temperatures 1.0 \
    --sampling_mode bernoulli-argmax

Here, ./outputs/my_experiment should be a directory containing checkpoint-* subdirectories (as created by training).

See more details below.

Models:

• LLaDA-8B-Instruct
• Dream-Instruct-7B

Training datasets: GSM8k, MATH, KodCode (or mixtures thereof).

Evaluation datasets: GSM8k, MATH-500, HumanEval, MBPP.

As stated in the Quick Start section, the main training interface is train.train, which takes all of its parameters in the input yaml config file.

Key config parameters are:

See configs/experiment_configs/ for more full-fledged examples.

For multi-GPU training, you can control the sharding using accelerate. Since the policies are small in size, we use a simple DDP setup: accelerate launch --config_file configs/accelerate_configs/8gpu_ddp.yaml -m train.train --config ....

Checkpoints are saved to output_dir specified in the config. This also supports automatically pushing the checkpoints to an s3 bucket when the output_dir starts with s3://. To use this functionality, you must implement the function common/s3.py:configure_s3(path), which we have left as a stub for your convenience.

The recommended way to evaluate is using eval.pipeline, which handles checkpoint resolution, multi-seed evaluation, and result aggregation:

python -m eval.pipeline ./outputs/my_experiment \
    configs/experiment_configs/llada_8b_instruct_dit_confidence_BL32_mixture.yaml \
    --checkpoints last \
    --datasets gsm8k,math \
    --temperatures 1.0 \
    --seeds 42,43,44 \
    --sampling_mode bernoulli-argmax \
    --save_path ./eval_results

Key arguments:

• First positional arg: path to directory containing checkpoint-* subdirectories
• Second positional arg: path to experiment config
• --checkpoints: comma-separated checkpoint numbers, or first/last for automatic resolution
• --datasets: comma-separated list, or all for gsm8k,math,humaneval,mbpp
• --seeds: comma-separated random seeds for multiple evaluation runs

Results are saved to --save_path as JSON files containing generations, with aggregated metrics in CSV format.

For more control, you can use eval.eval directly:

python -m eval.eval \
    --policy_path ./outputs/my_experiment/checkpoint-1000/model.safetensors \
    --config configs/experiment_configs/llada_8b_instruct_dit_confidence_BL32_mixture.yaml \
    --dataset gsm8k \
    --seed 42 \
    --temperature_policy 1.0 \
    --sampling_mode bernoulli-argmax \
    --output_dir ./eval_results