softmax-209m-1k-rope-s2357
A 209M-parameter research language model from the DilutionAttention release: softmax attention in every layer with RoPE (base 5e5); the control. Trained at a 1,024-token context on 4.19B tokens (1x Chinchilla), seed 2357. 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.1259 nats/token |
| NIAH retrieval at 1,024 tokens (in context; chance 1/7) | 1.00 |
| NIAH retrieval at 1,920 tokens (1.9x the trained context, NTK scaling) | 0.29 |
| NIAH retrieval at 3,072 tokens (3.0x the trained context, NTK scaling) | 0.33 |
| NIAH retrieval at 3,840 tokens (3.8x the trained context, NTK scaling) | 0.30 |
| NIAH retrieval at 6,144 tokens (6.0x the trained context, NTK scaling) | 0.16 |
| NIAH retrieval at 7,936 tokens (7.8x the trained context, NTK scaling) | 0.14 |
| NIAH retrieval at 12,288 tokens (12.0x the trained context, NTK scaling) | 0.16 |
| NIAH retrieval at 16,000 tokens (15.6x the trained context, NTK scaling) | 0.14 |
| NIAH retrieval at 24,000 tokens (23.4x the trained context, NTK scaling) | 0.16 |
| NIAH retrieval at 32,000 tokens (31.2x the trained context, NTK scaling) | 0.14 |
| PG-19 loss at positions 512-1,024 (in context) | 3.158 nats/token |
| PG-19 loss at positions 1,024-2,048 (2x, NTK scaling) | 4.413 nats/token |
| PG-19 loss at positions 2,048-4,096 (4x, NTK scaling) | 6.845 nats/token |
| PG-19 loss at positions 4,096-8,192 (8x, NTK scaling) | 7.382 nats/token |
| PG-19 loss at positions 16,384-32,768 (32x, NTK scaling) | 7.525 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-1k-rope-s2357", 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 1,024 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_1024/all_softmax_rope_1k_seed2357/in the release repository (metrics, resolved config, data provenance). Source checkpoint SHA-256:3d6851a15783d500631a9e0b60f15509cdb296ee7463f9b9d7e8c62f5112c8ed, 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.