Llama-3.1-8B-NHS
Llama-3.1-8B continued-pretrained with TTT-NTP (test-time training with next-token prediction): chunk-parallel causal rank-one fast-weight writes into the MLP down-projections, supervised by the model's own next-position hidden state.
Code: https://github.com/yancyou/TTT-NTP
Recipe
| base | meta-llama/Llama-3.1-8B |
| TTT layers | 0, 6, 12, 18, 24, 30 |
| inner optimizer | specnorm (per-chunk write normalized by its spectral norm) |
| inner lr (eta) | 0.3 |
ttt_proj init | small — N(0, initializer_range / sqrt(d)) |
ttt_norm_preserve | true (row-wise renorm back to ||W_down||_row) |
| conv / predict mode | no conv, predict-next, target = hidden states |
| chunk | 1024 |
| corpus | Long-Data-Collections, 32768 seq len |
| steps | 200, global batch 64, outer lr 5e-6, warmup 0.05 |
Two details differ from a naive reading of the paper and both matter:
ttt_proj is small-init, not identity-init (this matches the released
reference checkpoints), and ttt_norm_preserve is on.
RULER (Full-13, 100 samples/task, our DP eval pipeline)
| 4k | 8k | 16k | 32k | avg | |
|---|---|---|---|---|---|
| Llama-3.1-8B base | 65.13 | 57.16 | 47.01 | 54.19 | 55.87 |
| CPT-only control (no TTT, same tokens) | 60.38 | 54.85 | 48.73 | 51.74 | 53.92 |
paper recipe, sgd inner opt | 62.82 | 62.07 | 57.60 | 55.69 | 59.55 |
this checkpoint (specnorm) | 70.26 | 70.64 | 65.54 | 64.89 | 67.83 |
+11.96 over base. The CPT-only control (-1.95) rules out extra training
tokens as the cause; the gain is attributable to TTT. specnorm beats the
paper's sgd inner optimizer by 8.28.
Gains concentrate in multi-key / multi-value associative retrieval
(niah_multikey_2 18 -> 93 at 8k, niah_multikey_3 3 -> 53 at 8k), with
regressions on niah_multikey_1 and vt.
Caveats, stated plainly: single seed; niah_single_1 scores near zero for both
base and this checkpoint in our RULER build, so all averages are depressed by
roughly the same amount.
Usage
Requires trust_remote_code (the TTT MLP lives in the modeling file shipped
with the repo above).
from transformers import AutoModelForCausalLM, AutoTokenizer
m = AutoModelForCausalLM.from_pretrained("YoungXuan/Llama-3.1-8B-NHS",
trust_remote_code=True,
dtype="bfloat16", device_map="auto")
Setting ttt_mode: false in config.json disables the fast-weight path and
recovers the plain backbone.