vedkdev/text-diffusion-en-de

1---2tags:3- text-diffusion4- machine-translation5- en-de6- masked-diffusion7- from-scratch8language:9- en10- de11datasets:12- wmt/wmt1413license: apache-2.014---15 16# Text Diffusion Model for EN→DE Translation17 18A **masked discrete diffusion** model for English-to-German machine translation, trained from scratch on WMT14 EN-DE.19 20## Architecture21 22| Component | Detail |23|---|---|24| **Type** | Masked Discrete Diffusion |25| **Backbone** | DiT (Diffusion Transformer) with adaLN |26| **Parameters** | ~72M |27| **Blocks** | 12 DiT blocks |28| **Hidden dim** | 512, 8 attention heads |29| **Attention** | Bidirectional (no causal mask) with RoPE |30| **Conditioning** | Timestep via sinusoidal embeddings + adaLN; Segment embeddings for src/tgt |31| **Weight tying** | Input embeddings tied to output projection |32| **Tokenizer** | [Helsinki-NLP/opus-mt-en-de](https://huggingface.co/Helsinki-NLP/opus-mt-en-de) (~58K vocab) |33| **Max sequence** | 128 src + 128 tgt tokens |34 35### Inspired by36- **[MDLM](https://arxiv.org/abs/2406.07524)** — DiT backbone architecture, masked diffusion objective37- **[LLaDA](https://arxiv.org/abs/2502.09992)** — Conditional generation via SFT (keep prompt unmasked, mask only target), 1/t ELBO weighting38- **[DiNoiSer](https://arxiv.org/abs/2302.10025)** — Noise manipulation for conditional seq2seq diffusion39 40## How It Works41 42### Training (Forward Diffusion)431. Source (EN) and target (DE) tokens are concatenated: `[source | target]`442. A random masking rate `t ~ Uniform(0, 1)` is sampled per example453. Each target token is independently masked with probability `t`464. The bidirectional DiT predicts all masked tokens simultaneously475. Loss = cross-entropy on masked positions only, weighted by `1/t` (continuous-time ELBO)48 49### Inference (Reverse Diffusion)501. Start with source tokens + fully masked target: `[source | MASK MASK ... MASK]`512. Over 50 denoising steps, iteratively predict and unmask tokens523. At each step `t → s`: predict all masked tokens, randomly re-mask a fraction `s/t`534. Final step: all remaining masks are filled with predictions54 55## Training Details56 57| Setting | Value |58|---|---|59| **Dataset** | WMT14 EN-DE (~4.5M parallel sentence pairs) |60| **Optimizer** | AdamW (lr=3e-4, β₁=0.9, β₂=0.98, wd=0.01) |61| **Schedule** | Cosine with 4K linear warmup |62| **Effective batch size** | 256 (64 × 4 gradient accumulation) |63| **Max steps** | 200,000 |64| **Mixed precision** | FP16 |65| **Gradient clipping** | max_norm=1.0 |66| **Evaluation** | SacreBLEU on WMT14 test set every 20K steps |67 68## Quick Start69 70### Install dependencies71 72```bash73pip install torch transformers datasets trackio sacrebleu sacremoses sentencepiece protobuf74```75 76### Train77 78```bash79git clone https://huggingface.co/vedkdev/text-diffusion-en-de80cd text-diffusion-en-de81python train.py82```83 84The script will:85- Download WMT14 EN-DE automatically86- Train for 200K steps with logging via [Trackio](https://huggingface.co/docs/trackio)87- Evaluate SacreBLEU periodically88- Push checkpoints to this repo89 90### Adjusting for your hardware91 92Edit the `TRAIN_CONFIG` dict in `train.py`:93 94| GPU VRAM | Recommended `batch_size` | `gradient_accumulation_steps` |95|---|---|---|96| 24GB (A10G/3090/4090) | 64 | 4 |97| 16GB (T4/V100) | 32 | 8 |98| 12GB (3060) | 16 | 16 |99| 8GB (3070) | 8 | 32 |100 101### Inference (after training)102 103```python104import torch, json105from train import DiffusionTranslator, DiffusionTranslatorConfig, generate106from transformers import AutoTokenizer107 108# Load checkpoint109config = DiffusionTranslatorConfig(**json.load(open("checkpoints/best/config.json")))110model = DiffusionTranslator(config)111model.load_state_dict(torch.load("checkpoints/best/model.pt", map_location="cpu"))112model.eval()113 114tokenizer = AutoTokenizer.from_pretrained("checkpoints/best/")115 116# Translate117text = "The weather is nice today."118src = tokenizer(f"translate English to German: {text}",119                max_length=128, truncation=True, padding="max_length",120                return_tensors="pt")121 122gen_ids = generate(model, src["input_ids"], torch.zeros_like(src["input_ids"]),123                   config, num_steps=50, device="cpu")124print(tokenizer.decode(gen_ids[0], skip_special_tokens=True))125```126 127## Expected Results128 129Based on published literature for similar architectures on WMT14 EN→DE:130 131| Model | BLEU | Reference |132|---|---|---|133| Autoregressive Transformer | ~27 | Vaswani et al. |134| DiNoiSer (continuous diffusion) | 24.6 | Ye et al. 2023 |135| SeqDiffuSeq | 19.8 | Yuan et al. 2022 |136| E2D2 (discrete diffusion) | 24.8 | Kuleshov et al. 2024 |137| **This model (target)** | **15-20** | ~72M params, no KD |138 139> Note: Text diffusion models typically score 2-5 BLEU below autoregressive transformers of similar size. Knowledge distillation (KD) from an AR teacher can close the gap by ~1-2 BLEU.140 141## Citation142 143If you use this model, please cite the foundational papers:144 145```bibtex146@article{sahoo2024mdlm,147  title={Simple and Effective Masked Diffusion Language Models},148  author={Sahoo, Subham Sekhar and Arriola, Marianne and Schiff, Yair and Gokaslan, Aaron and Marroquin, Edgar and Kuleshov, Volodymyr},149  journal={NeurIPS},150  year={2024}151}152 153@article{nie2025llada,154  title={Large Language Diffusion Models},155  author={Nie, Shen and Zhu, Fengqi and You, Chao and Zhang, Xiaojun and Ou, Zhenguo and Zhu, Jun},156  journal={arXiv preprint arXiv:2502.09992},157  year={2025}158}159 160@article{ye2023dinoiser,161  title={DiNoiSer: Diffused Conditional Sequence Learning by Manipulating Noises},162  author={Ye, Jiasheng and Zheng, Zaixiang and Bao, Yu and Qian, Lihua and Gu, Quanquan},163  journal={ACL},164  year={2023}165}166```167