Looped Uniform Diffusion Language Models
Final OpenWebText checkpoints for Looped Uniform Diffusion Language Models
(looped-UDLM). Code: looped_udlm.
Training logs: looped_udlm.
Each model is a small DUO++ uniform diffusion transformer. It has 12 blocks, width 768, and a sequence length of 1024, and uses the GPT-2 tokenizer. The looped models share the weights of blocks 1β2 and replay them $S$ times per denoiser call, so the executed depth is $12 + 2(S-1)$ blocks with no extra parameters. During training, $S \in {1, \dots, 12}$ is drawn at random each optimizer step. Every looped run stops at the step whose training FLOPs match 1M steps of the non-looped baseline.
Looping is training-compute efficient. The looped models reach the non-looped baseline's best validation NLL with 2.79Γ (uniform) and 2.70Γ (flop-proportional) fewer training FLOPs on OpenWebText.
Checkpoints
| File | Loop-count schedule | Mean $S$ | Steps |
|---|---|---|---|
owt/baseline_step1000000.ckpt |
non-looped ($S = 1$) | 1 | 1,000,000 |
owt/uniform_step655420.ckpt |
uniform: $S \sim \mathcal{U}{1,\dots,12}$ | 6.5 | 655,420 |
owt/flop_proportional_step688915.ckpt |
flop-proportional: $p(S) \propto 1/F(S)$ | 5.72 | 688,915 |
$F(S)$ is the forward FLOPs of one sequence at depth $S$. Under the flop-proportional schedule, every depth receives the same share of training compute.
Each file is a full PyTorch Lightning checkpoint. It contains the model weights, the EMA weights (decay 0.9999), the optimizer state, and the training config. The evaluation scripts use the EMA weights.
All runs used the same recipe on one 4ΓH100 node: a global batch of 512 sequences, the DUO++ poly9 curriculum over the first 500k steps, and seed 1.
Usage
git clone https://github.com/SaitejaUtpala/looped_udlm && cd looped_udlm
pip install -r requirements.txt
bash scripts/download_checkpoints.sh # -> checkpoints/owt/*.ckpt
Or download a single file:
from huggingface_hub import hf_hub_download
ckpt = hf_hub_download("saitejautpala/looped_udlm", "owt/uniform_step655420.ckpt")
Every evaluation script takes a checkpoint and a depth policy. Use baseline
for a non-looped checkpoint, an integer $S$ for a fixed depth, or adaptive
for per-sample adaptive looping:
# validation likelihood at a fixed depth, or with adaptive depth
bash scripts/eval_ppl.sh checkpoints/owt/uniform_step655420.ckpt 12
bash scripts/eval_ppl.sh checkpoints/owt/uniform_step655420.ckpt adaptive
bash scripts/eval_ppl.sh checkpoints/owt/baseline_step1000000.ckpt baseline
# generation: Gen. PPL under GPT-2 Large, plus entropy
STEPS=256 bash scripts/gen_ppl.sh checkpoints/owt/flop_proportional_step688915.ckpt adaptive
# zero-shot likelihood and zero-shot multiple choice
bash scripts/zero_shot.sh checkpoints/owt/uniform_step655420.ckpt 8
bash scripts/mcq.sh checkpoints/owt/uniform_step655420.ckpt adaptive
A looped checkpoint runs at any test-time depth $S$, including depths beyond the training range such as $S = 24$. Under adaptive looping, sequence $i$ stops at the first traversal $s$ where
with $\epsilon = 0.1$ and a cap of $S_{\max} = 12$. Converged sequences leave the recurrent batch, so each sequence runs only the traversals it needs.
Training data
The corpus is pre-tokenized with GPT-2 and packed into 1024-token blocks,
ready for the dataloader to read directly:
saitejautpala/openwebtext-gpt2-packed-1024.
Acknowledgements
The code and training recipe build on DUO.