Set Diffusion: Interpolating Token Orderings between Autoregression and Diffusion for Fast and Flexible Decoding

August 9, 2026 · View on GitHub

deploy deploy deploy

0. Setup

Setup environment

Install conda:

# For conda: https://docs.conda.io/projects/conda/en/stable/user-guide/install/linux.html
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh && \
bash miniconda.sh -b -p /opt/conda

Create the locked conda environment:

conda env create -f requirements.yaml
conda activate dllm-dev

The conda file installs Python, pip, the pinned dependencies in requirements-lock.txt, and this package in editable mode. To validate a fresh environment, run:

python -m pip check
python - <<'PY'
import torch
import transformers
import src

print(torch.__version__)
print(transformers.__version__)
PY

For a pip-only install, use the pinned direct requirements:

python -m pip install -r requirements.txt -e .

Use requirements-lock.txt when byte-for-byte dependency reproduction is required. Regenerate the lockfile only as an intentional release step.

Activate an existing environment with:

conda activate dllm-dev

We also include a setup_env.sh script for runtime shell variables on compute nodes. Run it from the repository root after the environment has been created:

source setup_env.sh

Credentials are intentionally opt-in. If you want the setup script to source private W&B or Hugging Face settings, point DLLM_PRIVATE_ENV at a local, untracked shell file:

export WANDB__SERVICE_WAIT=600
export WANDB_ENTITY="<WANDB_ENTITY>"
export WANDB_API_KEY="<WANDB_API_KEY>"
export HUGGINGFACE_TOKEN="<HF_TOKEN>"

Then run DLLM_PRIVATE_ENV=/path/to/private_env.sh source setup_env.sh.

  • WandB token can be found here.
  • Hugging Face token can be created here.

1. Code Organization

  1. bash_scripts: These shell scripts can be used to reproduce the experiments from our work.
  2. configs: We use Hydra config files to organize experiments.
    1. config.yaml: Entry point for launching training experiments.
    2. eval_config.yaml: Entry point for evaluations.
  3. scripts: The main training and evaluation scripts
    1. scripts/composer_scripts/train_discrete_denoiser.py: This script is the main training entry point.
    2. scripts/eval: These scripts run evaluation for the translation, summarization, and math reasoning datasets, as well as any likelihood evaluation.
  4. src:
    1. src/denoiser: During training, denoisers take in "noisy" inputs and predict clean signals. At inference, starting from a purely noisy signal, through iterative denoising, these classes produce samples that resemble data.
      1. AR: We can view autoregressive models within this paradigm. Noise is applied by masking tokens one at a time from right-to-left. Denoising is done one token at a time, left-to-right.
      2. Diffusion: We implement masked diffusion models:
        • MDLM: Standard masked diffusion.
        • BD3LM: Block diffusion models.
        • SetDLM: Set diffusion models.
    2. src/backbone: These are the underlying neural networks that take in noisy inputs and produce logits. Each denoiser is parameterized by a backbone. The denoiser can optionally post-process the logit outputs of the backbone to produce log-probs over the clean sequence.

2. Reproducing Experiments

The shell scripts provided in bash_scripts can be used to reproduce the training and evaluations from our work.

  • For training, the files follow a convention where the dataset and denoiser class are specified. For example, to train SetDLM on the GSM8K dataset, use run_train_setdlm_gsm8k.sh.
  • Once models have been trained or downloaded, the provided evaluation scripts can be used to reproduce the reported metrics from our work. For example, to evaluate GSM8K models, use run_lm_eval_harness.sh. The task-specific wrappers below document the supported evaluation entry points. Plotting utilities are kept out of this release repo; generate plots from exported metrics/TSV artifacts downstream.

Evaluation scripts resolve checkpoint keys through bash_scripts/eval_model_paths.sh. Select a known checkpoint with EVAL_MODEL_KEY, or pass an explicit HF id or local checkpoint path in MODEL_PATH:

# Resolve to kuleshov-group/setdlm-gsm8k-smax32.
EVAL_MODEL_KEY=gsm8k:setdlm-smax32 bash bash_scripts/run_lm_eval_harness.sh

# Explicit HF ids and local paths are also accepted.
MODEL_PATH=kuleshov-group/cnndm-setdlm-smax16 bash bash_scripts/run_seq2seq_eval_cnndm.sh
MODEL_PATH=/share/kuleshov/ma2238/runs/dllm-dev/<run-dir> bash bash_scripts/run_lm_eval_harness.sh

Dataset configs read from DLLM_DATA_DIR and default to data/. Set DLLM_DATA_DIR=/path/to/datasets when cached datasets live elsewhere. Evaluation scripts write outputs under outputs/ by default and accept checkpoint-related overrides such as MODEL_PATH, EVAL_MODEL_KEY, CKPT_FILE or CKPT, and USE_EMA. LM1B_MODEL_KEY and LM1B_MODEL_PATH are accepted by the LM1B likelihood wrapper.

GSM8K Distillation Data

The GSM8K training wrappers for AR, MDLM, BD3LM, and SetDLM use distilled teacher answers saved as preprocessed Hugging Face datasets. Build them before launching GSM8K distillation training:

bash bash_scripts/run_build_gsm8k_distill.sh

By default this writes:

outputs/distillation/Qwen3-32B-AWQ/gsm8k_train
outputs/distillation/Qwen3-32B-AWQ/gsm8k_eval

Set DLLM_DISTILL_DATA_DIR if you want the datasets somewhere else, and use the same value when launching training:

DLLM_DISTILL_DATA_DIR=/path/to/Qwen3-32B-AWQ bash bash_scripts/run_build_gsm8k_distill.sh
DLLM_DISTILL_DATA_DIR=/path/to/Qwen3-32B-AWQ NUM_VISIBLE_DEVICES=8 bash bash_scripts/run_train_setdlm_gsm8k.sh

The builder defaults to Qwen/Qwen3-32B-AWQ as the teacher and Qwen/Qwen3-1.7B-Base as the student tokenizer. It saves the same artifact shape used by the local L=1024 GSM8K distillation datasets: input_ids plus index, final sequences capped at 1024 tokens, and chat rows terminated with <|im_end|>. Override the teacher and student tokenizer with DISTILL_TEACHER_MODEL and DISTILL_STUDENT_TOKENIZER. For a fast smoke test that validates the preprocessing path without loading the teacher model, set DISTILL_USE_GOLD_ANSWERS=true; this does not reproduce the paper distillation data.

Evaluation scripts are provided for the following tasks:

For full experiment sweeps, scripts/eval/repro_suite_runner.py builds the evaluation matrix and uses the same HF-first checkpoint resolver, so generated commands and expected output paths agree with the shell wrappers.

Paper H100 Throughput Reproduction

For the H100 throughput numbers reported in the paper, use the paper-specific wrappers:

All three wrappers use 50 warmup examples. CNN/DailyMail and ROCStories use 1000 measured examples; the GSM8K helper defaults to 200 measured examples and accepts EVAL_MODEL_KEY, MODEL_PATH, CONFIDENCE_THRESHOLD, BLOCK_SIZE, and MAX_WINDOW_SIZE overrides for the other Pareto points.

The MCQA evaluation flow uses Hydra overrides in the same style as the other eval frameworks, for example +eval/mcqa@task=all or +eval/mcqa@task=hellaswag. It evaluates the validation splits of HellaSwag, PIQA, and Social IQa with continuation scoring rather than greedy generation, writes per-example predictions and option scores to predictions.json, and saves aggregate accuracies plus a summary table to metrics.json / metrics.txt. By default, answer options are ranked by average log-probability per answer token to reduce length bias.

3. Hugging Face Checkpoints

SetDLM Hugging Face repo ids use smax{2 * desired_block_size} naming. For example, SetDLM-smax32 is published as *-setdlm-smax32. The GSM8K SetDLM release ids are:

  • kuleshov-group/setdlm-gsm8k-smax8
  • kuleshov-group/setdlm-gsm8k-smax16
  • kuleshov-group/setdlm-gsm8k-smax32

The corresponding CNN/DM, OWT, and LM1B ids are:

  • kuleshov-group/cnndm-setdlm-smax8, kuleshov-group/cnndm-setdlm-smax16, kuleshov-group/cnndm-setdlm-smax32
  • kuleshov-group/owt-setdlm-smax8, kuleshov-group/owt-setdlm-smax16, kuleshov-group/owt-setdlm-smax32
  • kuleshov-group/lm1b-setdlm-smax8, kuleshov-group/lm1b-setdlm-smax16, kuleshov-group/lm1b-setdlm-smax32

Other evaluated checkpoints use the kuleshov-group/<dataset>-<model> naming scheme, for example kuleshov-group/owt-bd3lm-s16. The resolver accepts either these HF ids or compact keys such as cnndm:setdlm-smax16, owt:bd3lm-s16, and lm1b:ar. For exact GSM8K SetDLM Pareto reproduction, use the repo evaluation loader/scripts rather than plain AutoModel.from_pretrained, because the loader normalizes legacy checkpoint config and eval-time SetDLM noise/cache-order settings.

Citation

@article{arriola2026setdiffusion,
  title={Set Diffusion: Interpolating Token Orderings Between Autoregression and Diffusion for Fast and Flexible Decoding},
  author={Arriola, Marianne and Kuleshov, Volodymyr},
  booktitle={The Forty-Third International Conference on Machine Learning},
  year={2026},
  url={https://arxiv.org/abs/2607.01775}
}