Modelpublic

YoungXuan/Llama-3.1-8B-NHS

sourceHugging Facellama3.1updated 1mo agoView on Hugging Face

0likes56downloads

Model Card

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

python

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.