UCB-DSP-DiffusionLM

Official implementation of the EACL 2026 paper:

Principled Self-Correction in Discrete Diffusion: A UCB-Guided Framework for Text Generation Masaki Asada and Makoto Miwa EACL 2026, pp. 6678–6692 [Paper]

Overview

This repository provides a framework for discrete diffusion language models that addresses the fundamental training-inference discrepancy: models are trained on corrupted ground-truth tokens, yet at inference time must denoise inputs corrupted from their own predictions.

The framework combines two complementary methods:

• Deeper Self-Prediction (DSP): A multi-step training objective that uses the model's own intermediate outputs as noise sources, teaching robust self-correction behavior.
• UCB-guided Decoding: An inference algorithm that frames token re-masking as a multi-armed bandit problem and applies the Upper Confidence Bound (UCB) criterion to balance exploration and exploitation.

The framework is built on top of LLaMA-3.2 and evaluated on XSum, SQuAD, MSNews, MSQG, and GSM8K, consistently outperforming existing discrete diffusion baselines on both automatic metrics and LLM-as-a-Judge evaluations.

Requirements

pip install torch transformers datasets evaluate accelerate deepspeed wandb

Flash Attention 2 is required:

pip install flash-attn --no-build-isolation

Tested on CUDA 12.6.

Training

torchrun --nproc_per_node=<NUM_GPUS> finetune_text_gen.py \
    --run_name <RUN_NAME> \
    --output_dir OUTPUT_DIR/<RUN_NAME> \
    --dataset_dir DATASET_DIR/<TASK_NAME> \
    --model_name meta-llama/Llama-3.2-1B \
    --num_train_epochs 3 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --deepspeed ds_zero2_bf16.json \
    --lr_scheduler_type cosine \
    --learning_rate 1e-4 \
    --num_diffusion_steps 32 \
    --max_new_tokens 64 \
    --task_name XSum \
    --num_losses 2

Key hyperparameters:

Inference

torchrun --nproc_per_node=<NUM_GPUS> infer_text_gen.py \
    --run_name <RUN_NAME> \
    --output_dir OUTPUT_DIR/<RUN_NAME> \
    --finetuned_model_path <FINETUNED_MODEL_PATH> \
    --per_device_eval_batch_size 16 \
    --num_inference_steps 16 \
    --remask_value random \
    --ucb False

To enable UCB-guided Decoding, pass --ucb True. Results are saved to ./results/<RUN_NAME>.json.

Evaluation

Automatic metrics (ROUGE, BERTScore, METEOR)

Computed automatically during inference by infer_text_gen.py.

LLM-as-a-Judge

cd evaluation/llm
python prepare_openai_jsonl.py   # Prepare prompts
python submit_openai_jsonl.py    # Submit via OpenAI Batch API
python calc_score_for_sum.py     # Aggregate scores for summarization
python calc_score_for_qg.py      # Aggregate scores for question generation

Classic metrics (BLEU, DISTINCT, etc.)

cd evaluation/classic
bash run_calc.sh

Repository Structure

.
├── finetune_text_gen.py      # Training script
├── infer_text_gen.py         # Inference and evaluation script
├── diff_args.py              # Diffusion-specific hyperparameter definitions
├── ds_zero2_bf16.json        # DeepSpeed ZeRO-2 config (bfloat16)
├── job_ft_text_gen.sh        # Training job script (PBS)
├── job_infer_text_gen.sh     # Inference job script (PBS)
├── our_transformers/         # Modified Transformers library with diffusion-aware LLaMA
└── evaluation/
    ├── classic/              # BLEU, ROUGE, METEOR, Distinct evaluation
    └── llm/                  # LLM-as-a-Judge evaluation via OpenAI Batch API

Citation

@inproceedings{asada-miwa-2026-principled,
    title     = "Principled Self-Correction in Discrete Diffusion: A {UCB}-Guided Framework for Text Generation",
    author    = "Asada, Masaki and Miwa, Makoto",
    booktitle = "Proceedings of the 19th Conference of the European Chapter of the Association for Computational Linguistics (Volume 1: Long Papers)",
    month     = mar,
    year      = "2026",
    address   = "Rabat, Morocco",
    publisher = "Association for Computational Linguistics",
    url       = "https://aclanthology.org/2026.eacl-long.314/",
    doi       = "10.18653/v1/2026.eacl-long.314",
    pages     = "6678--6692",
}