Llama-3.1-8B-NHS

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

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)

+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.