TriDLM-124M-causal

TriDLM-124M-causal

A 124M-parameter masked-diffusion language model (MDLM-style absorbing-state denoiser) with all heads causal (lower-bound control), 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)

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

import torch
from transformers import AutoModel, AutoTokenizer
tok = AutoTokenizer.from_pretrained("EER6/TriDLM-124M-causal")                      # gpt2 BPE + [MASK] = 50257
model = AutoModel.from_pretrained("EER6/TriDLM-124M-causal", 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.