Diffusion Step Efficiency

Minimal-scale research project comparing discrete diffusion language models against autoregressive (AR) transformers on IMDB. The core question:

How many denoising steps does a PLM-initialized diffusion LM actually need to match an equally sized AR baseline?

One-sentence summary: We implement matched ~55 M parameter AR, diffusion, and LLaDA-style hybrids on IMDB to map quality, diversity, and latency as we reduce diffusion/refinement steps down to the smallest count that still matches the AR baseline.

Artifacts at a glance

• 📄 Full LaTeX paper with AI-generated disclaimer: docs/paper.tex
• 📊 Executive summary + tables: docs/results_summary.md / .pdf
• 📁 Raw metrics + samples: results/ablation_results.{json,csv,md}
• 🖼 Plots ready for publication: results/plots/*.pdf

---

1. Setup

1.1 Environment

python -m venv venv
.\venv\Scripts\activate       # PowerShell on Windows
pip install -r requirements.txt
pip install -e .

1.2 Dataset cache

python -c "from data.prepare_dataset import prepare_dataset; prepare_dataset()"

First run tokenizes IMDB and saves to data/cache/. Reuses cached tensors afterward.

---

2. Training

2.1 Autoregressive LM

python train.py --model ar

Saves checkpoints/ar_model.pt after 10 epochs (config: configs/ar_lm.yaml). Current settings use CPU-friendly batch 8 + gradient accumulation 4.

2.2 Diffusion LM

python train.py --model diffusion

Saves checkpoints/diffusion_model.pt (config: configs/diffusion_lm.yaml). Trains once with 1000-step cosine schedule; inference later sweeps smaller step counts.

2.3 Both sequentially

python train.py --model both

---

3. Testing (TDD)

python -m pytest tests/ -v

Specific suites:

python -m pytest tests/test_ar_model.py -v
python -m pytest tests/test_diffusion_model.py -v
python -m pytest tests/test_metrics.py -v

---

4. Evaluation & Ablations

4.1 Metrics & generation

python evaluate.py --num-samples 10

Outputs results/evaluation.json with samples plus BLEU/ROUGE/self-BLEU/distinct-$n$ and throughput statistics. Mirrors the metrics used in the paper.

4.2 Step-count sweep

python experiments/run_ablations.py

Will produce JSON/CSV tables covering diffusion steps {1,5,10,50,100,500,1000} plus sanity-check logging (flag if trends violate PRD expectations).

4.3 Publication plots

evaluation/plots.py renders:

• Quality vs. diffusion steps
• Diversity vs. diffusion steps
• Inference time vs. diffusion steps
• Diffusion vs. AR trade-off curve

All saved as .pdf in results/plots/ for paper inclusion.

---

5. Project Structure

diffusion-step-efficiency/
├── README.md
├── configs/
├── data/
├── docs/
│   └── research_paper.md
├── evaluation/
├── experiments/
├── inference/
├── models/
├── tests/
├── training/
├── train.py
├── evaluate.py
└── checkpoints/ (runtime)

---

5. Results Snapshot

Model / Setting	BLEU	ROUGE-L	Distinct-2	Tokens/sec
AR baseline	0.021	0.062	0.293	13.5
Diffusion (100 steps)	0.030	0.082	0.983	14.6
Diffusion (500 steps)	0.058	0.149	0.726	4.0
LLaDA (8 iterations)	0.042	0.121	0.260	334.4
LLaDA (12 iterations)	0.043	0.121	0.249	240.1

See docs/results_summary.md or the LaTeX paper for complete tables (all step counts/iterations, latency, self-BLEU).

---

6. Current Status

• ✅ AR model trained on CPU (checkpoints/ar_model.pt)
• ✅ Diffusion model trained + step-count ablations logged (results/ablation_results.*)
• ✅ Evaluation scripts report BLEU/ROUGE/distinct/self-BLEU/tokens-sec with sanity checks
• ✅ Plotting module + CSV/Markdown export finished (results/plots/, results/ablation_results.md)
• 📝 Research paper + summary docs at docs/paper.tex, docs/results_summary.md

---

7. Hardware Notes

• Works on CPU (16 GB RAM) albeit slowly (~12–16 h per model)
• GPU configs available (set hardware.device: cuda, adjust batch size)
• Cached dataset occupies ~3 GB under data/cache/

---

8. References

• Ho et al., Denoising Diffusion Probabilistic Models
• Li et al., Diffusion-LM
• LLaDA: Latent Language Diffusion for Autoregressive Decoding
• GPT-style AR transformers for baseline matching