README.md

---
license: apache-2.0
language: en
library_name: transformers
tags: [masked-diffusion, diffusion-language-model, gpt2, openwebtext, triangular-attention]
datasets: [Skylion007/openwebtext]
---

# TriDLM-124M-split

A 124M-parameter **masked-diffusion language model** (MDLM-style absorbing-state denoiser) with **every layer: even heads causal, odd heads anti-causal**, one of four arms of the
TriDLM experiment: can a diffusion LM keep its quality when every attention head is *triangular* (causal or anti-causal), so that
attention costs ~L²/2 instead of L²? Code, training logs and the full experiment record: https://github.com/AntonXue/TriDLM (`_claude/`).

## The four arms (identical model, data order, masks and recipe; only the per-head attention mask differs)

| repo | attention | val NELBO (nats/token) | ppl | train tok/s (8×GH200) |
|---|---|---|---|---|
| `EER6/TriDLM-124M-bidir` | all heads bidirectional (the stock MDLM control) | 3.374 | 29.2 | 2.60M |
| `EER6/TriDLM-124M-split` | every layer: even heads causal, odd heads anti-causal | 3.395 | 29.8 | 2.73M |
| `EER6/TriDLM-124M-alt` | even layers all-causal, odd layers all-anti-causal | 3.406 | 30.2 | 2.55M |
| `EER6/TriDLM-124M-causal` | all heads causal (lower-bound control) | 4.940 | 139.7 | 2.73M |

NELBO = the exact-count masked-diffusion ELBO (equal in expectation to the MDLM 1/t objective) on all 110,451 OWT-validation blocks.
`split` trails `bidir` by 0.020 nats (~2 % ppl) and `alt` by 0.032; `causal` (no right context) is the expected lower bound.
Generation metrics (gen-ppl under gpt2-large, MAUVE-256) per sampler and prompting condition are in the GitHub record.

## Model

GPT-2-small denoiser ported from nanoGPT (12 layers, 768 d, 12 heads, learned absolute positions, LayerNorm without bias, GELU,
tied embeddings), vocab 50304 = gpt2's 50257 + `[MASK]` (id 50257) + pads, **no next-token shift: position i predicts token i**, and
**per-head QK-norm** (LayerNorm on q and k over the head dim). QK-norm is essential: without it every triangular arm diverged from
attention-logit growth (max |q·k/√d| in layers 1–3 grew from ~40 to 10⁶ by 25–30k steps at both lr 3e-4 and 1e-4, while the
bidirectional arm stayed at ~42); with it all four arms train cleanly.

## Training

OpenWebText (gpt2 BPE, 1024-token packed blocks with EOS separators), **100k steps × 512 × 1024 = 52.4B tokens** (5.9 epochs),
AdamW lr 1e-4 (β 0.9/0.95, wd 0.1, clip 1.0), warmup 1000 → 89,000 stable → 10,000 cosine decay to 0, bf16, seed 1. Loss = mean over
rows of the mean cross-entropy over the masked positions of each row, with the mask count n ~ U{1..L} stratified across the batch
and the mask set a uniform n-subset; `[MASK]`/pad logits are −∞ (SUBS). Trained on 8 GH200 nodes (1 GPU each) in ~5.5 h.

## Usage

```python
import torch
from transformers import AutoModel, AutoTokenizer
tok = AutoTokenizer.from_pretrained("EER6/TriDLM-124M-split")                      # gpt2 BPE + [MASK] = 50257
model = AutoModel.from_pretrained("EER6/TriDLM-124M-split", trust_remote_code=True, attn_impl="sdpa").cuda().eval()
# attn_impl="sdpa": dense masks, any torch >= 2.1; "flex" (block-sparse, needs torch >= 2.5 + Triton) is numerically identical.

x0 = tok("The quick brown fox", return_tensors="pt")["input_ids"].cuda()          # (B, L <= 1024)
mask = torch.rand(x0.shape, device=x0.device) < 0.5
xt = torch.where(mask, torch.full_like(x0, tok.mask_token_id), x0)
logits = model(xt).logits                                                         # (B, L, 50304); position i predicts token i
logits[..., 50257:] = float("-inf")                                               # SUBS: never predict [MASK] or pads
```
Sampling: use the random-order ancestral sampler (`eval.py: ancestral_sample`) — gen-ppl 56 ± 3 under gpt2-large at NFE 1024 for
`bidir`. Confidence-ordered decoding (argmax *or* sampled tokens, threshold or one-at-a-time) collapses into repetition on these
base models; MAUVE catches it. The `causal`/`split`/`alt` checkpoints must be evaluated through their own `attn_mode` (saved in
the config); never run them with a full mask.

## Notes

- The training set (mdlm-vibes' OWT prep) contains one block with a stray `[MASK]` id; it makes one logged loss value `inf` per epoch
  and has no effect on the gradients.
- No EMA; the final checkpoint is the post-decay endpoint. The end-of-stable-phase full checkpoint (step 90k, optimizer included) is
  kept offline for WSD extensions.
- Ported model code (MIT, nanoGPT) is included as `modeling_tridlm.py`; this repo is apache-2.0.