softmax-209m-16k-rope-1x-s1337
A 209M-parameter research language model from the DilutionAttention release: softmax attention in every layer with RoPE (base 5e5); the control. Trained at a 16,384-token context on 4.19B tokens (1x Chinchilla), seed 1337. It is a base model for studying the attention operator, not an assistant: it is not instruction-tuned or safety-tuned.
Results
| measure | value |
|---|---|
| validation loss (common 1M-token FineWeb-Edu slice) | 3.0883 nats/token |
| NIAH retrieval at 16,000 tokens (in context; chance 1/7) | 0.79 |
| NIAH retrieval at 24,000 tokens (1.5x the trained context, NTK scaling) | 0.17 |
| NIAH retrieval at 32,000 tokens (2.0x the trained context, NTK scaling) | 0.14 |
| NIAH retrieval at 48,000 tokens (2.9x the trained context, NTK scaling) | 0.17 |
| NIAH retrieval at 64,000 tokens (3.9x the trained context, NTK scaling) | 0.14 |
| PG-19 loss at positions 8,192-16,384 (in context) | 3.537 nats/token |
| PG-19 loss at positions 16,384-32,768 (2x, NTK scaling) | 3.710 nats/token |
| PG-19 loss at positions 32,768-65,536 (4x, NTK scaling) | 4.517 nats/token |
| PG-19 loss at positions 65,536-131,072 (8x, NTK scaling) | 5.338 nats/token |
NIAH is likelihood-scored keyword retrieval from repetitive filler (see the release's reproducing.md), not long-context reasoning. Every number here, and the matched controls, are in the release's evidence ledger.
Use
pip install "dilution-attention[hub] @ git+https://github.com/Royer-Research-Labs/Dilution-Attention" tiktoken
import tiktoken, torch
from dilution.hub import from_pretrained
model = from_pretrained("Royer-Research-Labs/softmax-209m-16k-rope-1x-s1337", device="cuda") # fused Triton kernels on CUDA
enc = tiktoken.get_encoding("gpt2")
ids = torch.tensor([enc.encode("The history of the printing press")], device="cuda")
with torch.autocast("cuda", dtype=torch.bfloat16):
out = model.generate(ids, 64, temperature=1.0, top_p=0.9, vocab_limit=50257)
print(enc.decode(out[0].tolist()))
Inputs longer than 16,384 tokens need dilution.extend.ExtendedContext(model, length, ntk=True) around the call.
Training
- Architecture: 24 layers, width 768, 12 heads, SwiGLU, no biases, tied embeddings, GPT-2 vocabulary; position encoding: rope.
- Data: FineWeb-Edu
sample-100BT(revision fc9850d), a 40B-token GPT-2-tokenized corpus rebuilt withdilution-prepare-data(reproducing.md). - Schedule: 65,536 tokens per step, LR 6e-4 cosine over a 2x-Chinchilla (8.4B-token) schedule; this is the step-64,000 (1x) checkpoint, where the learning rate is 55.7% of peak.
- Run record:
runs/scale_16384/all_softmax_rope_seed1337/in the release repository (metrics, resolved config, data provenance). Source checkpoint SHA-256:b2132bf6fa7fe940e98034cac7304fee9046ef68688683f914789f0b8275a747, step 64,000.
License and data
Weights: Apache-2.0. Training data: FineWeb-Edu (ODC-By 1.0), used with attribution. The model can produce incorrect, biased or offensive text; it is a research artifact.
Citation
Cite the software release: Royer, N. (2026). DilutionAttention. Zenodo. https://doi.org/10.5281/zenodo.23073037 (BibTeX in the README and CITATION.cff). Technical report: https://doi.org/10.5281/zenodo.23073151.