Reinforced Fast Weights with Next-Sequence Prediction

Hee Seung Hwang; Xindi Wu; Sanghyuk Chun; Olga Russakovsky

cs.CLFeb 18, 2026

Reinforced Fast Weights with Next-Sequence Prediction

Authors: Hee Seung Hwang, Xindi Wu, Sanghyuk Chun, Olga Russakovsky

Organizations: Princeton University

Abstract

Fast weight architectures offer a promising alternative to attention-based transformers for long-context modeling by maintaining constant memory overhead regardless of context length. However, their potential is limited by the next-token prediction (NTP) training paradigm. NTP optimizes single-token predictions and ignores semantic coherence across multiple tokens following a prefix. Consequently, fast weight models, which dynamically update their parameters to store contextual information, learn suboptimal representations that fail to capture long-range dependencies. We introduce REFINE (Reinforced Fast weIghts with Next sEquence prediction), a reinforcement learning framework that trains fast weight models under the next-sequence prediction (NSP) objective. REFINE selects informative token positions based on prediction entropy, generates multi-token rollouts, assigns self-supervised sequence-level rewards, and optimizes the model with group relative policy optimization (GRPO). REFINE is applicable throughout the training lifecycle of pre-trained language models: mid-training, post-training, and test-time training. Our experiments on LaCT-760M and DeltaNet-1.3B demonstrate that REFINE consistently outperforms supervised fine-tuning with NTP across needle-in-a-haystack retrieval, long-context question answering, and diverse tasks in LongBench. REFINE provides an effective and versatile framework for improving long-context modeling in fast weight architectures.

Figures & tables

Figure 1: Comparison of standard NTP and ReFINE . Standard NTP ( top ) computes cross-entropy loss at each token position, providing only token-level supervision to fast weight models. ReFINE ( bottom ) provides sequence-level supervision by generating multi-token rollouts at high-entropy positions, assigning sequence-level rewards from hidden states, and optimizing with RL.

Figure 2: Comparison of Standard Transformer and Fast Weight Models , adapted from Zhang et al. (2025) . Fast weight models replace attention with a fixed-size memory implemented as a weight matrix ( $W$ ), and updated according to Eq. 1 .

Training Phase	Source of Training Data
Pre-training	General web-scraped data
Mid-training	Similar data to pre-training
Post-training	Task / instruction / preference data
Test-Time Training	Test data

Table 1: Training Phases for Fast Weight Language Models. ReFINE can be applied across all phases beyond pre-training to improve long-context modeling, using different data sources.

Figure 3: ReFINE . We forward the sequence through the policy model and compute token-level entropy values. Sequences are split into chunks and a target token position is sampled from each chunk based on the entropy ( Entropy-Based Token Selection ). Prefixes are copied from the original sequence up to each target token. The policy model predicts continuations from the prefixes ( Rollout Generation ). Reward is computed based on the generated rollouts and ground truth tokens ( Reward Assignment ). Finally, we update the policy model with GRPO ( Optimization with RL ).

	$c$	$k$	$n$	batch size	$\lambda_{\text{SFT}}$	$\lambda_{\text{RL}}$	reward	lr
MidTr	8	5	1	128	1	0.2	$R^{\varphi}$	1e-6
PostTr	8	5	1	64	1	0.2	$R^{\text{hybrid}}$	1e-6
TTT	8	5	1	8	1	0.4	$R^{\text{binary}}$	1e-6

Table 2: Training configurations. Detailed training configurations for Mid-training ( MidTr ), post-training ( PostTr ) and test-time training ( TTT ) are shown.

	RULER S-NIAH				RULER MK-NIAH				RULER MQ-NIAH				RULER MV-NIAH
	4K	8K	16K	Avg	4K	8K	16K	Avg	4K	8K	16K	Avg	4K	8K	16K	Avg
LaCT-760M	98.8	91.2	95.8	95.3	70.2	44.8	24.2	46.2	39.2	24.9	17.4	27.1	40.6	27.0	17.7	28.4
+ SFT MidTr	98.4	90.8	97.6	95.6	70.6	45.0	24.2	46.6	40.7	25.2	17.4	27.8	41.7	26.3	17.7	28.5
+ ReFINE MidTr	99.0	93.0	96.8	96.3	70.4	46.0	26.6	47.7	40.5	25.8	18.0	28.1	41.2	29.5	18.6	29.8
DeltaNet-1.3B	100.0	100.0	100.0	100.0	23.6	17.8	3.4	14.9	27.7	3.9	1.8	11.1	23.9	5.4	2.0	10.4
+ SFT MidTr	100.0	100.0	100.0	100.0	23.8	19.2	7.8	16.9	33.2	16.8	3.5	17.8	28.3	19.3	3.7	17.1

Table 3: Performance on Long-Context Retrieval Tasks. We evaluate mid-trained ( MidTr ) models on the NIAH tasks in RULER at 4K, 8K, and 16K context lengths (standard SFT vs. ReFINE ). Highest scores in each category are highlighted in bold .

				RULER SQuADQA				RULER HotpotQA
	MidTr	PostTr	TTT	4K	8K	16K	Avg	4K	8K	16K	Avg
LaCT 760M	-	-	-	17.5	14.0	6.5	12.7	20.5	14.5	12.0	15.7
	SFT	-	-	19.5	14.0	6.0	13.2	21.5	15.0	12.0	16.2
	ReFINE	-	-	20.5	15.5	8.0	14.7	23.0	16.0	12.5	17.2
	ReFINE	SFT	-	30.5	20.0	7.5	19.3	25.0	16.5	11.5	17.7
	ReFINE	Nested SFT	-	38.0	20.0	7.5	21.8	24.0	16.0	12.5	17.5

Table 4: Performance on Multi-Doc QA Tasks. We evaluate various training strategies during mid-training ( MidTr ), post-training ( PostTr ), and test-time training ( TTT ) for multi-document question and answering tasks. For Nested SFT and Nested ReFINE , we train the model with the training method on the prompt portion of the sample as described in § 4.3 . Higher is better. First and second highest scores on each task are highlighted in bold and underline , respectively.

			Single-doc QA			Multi-doc QA		Summarization		Few-shot QA			Coding
	MidTr	TTT	NQ	QR	MF	HP	2W	QM	MN	SS	TC	TQ	LC	RP	Avg
LaCT 760M	-	-	6.5	10.5	7.2	11.7	9.8	13.6	9.2	14.2	10.5	8.0	26.7	29.7	13.1
	SFT	-	5.8	10.1	7.4	12.6	9.2	13.1	10.5	13.8	7.5	12.2	29.8	30.1	13.5
	ReFINE	-	6.5	11.1	12.6	19.6	18.0	14.3	15.9	17.0	11.0	12.9	32.9	31.1	16.9
	ReFINE	SFT	5.1	12.6	13.5	15.9	18.4	13.2	16.2	17.4	12.5	16.2	32.0	31.4	17.0
	ReFINE	ReFINE	6.7	14.5	14.1	18.4	22.8	13.9	15.9	17.4	15.5	11.8	32.2	32.3	18.0

Table 5: Performance on Long-Context Tasks in LongBench. We study the impact of the learning algorithm during mid-training ( MidTr ) and test-time training ( TTT ) on tasks with long-context. SFT denotes the supervised fine-tuning with next-token prediction. We evaluate on 12 tasks in LongBench, filtered for samples with at most 16K tokens. Details are similar to Table 4 .

Panel of Figure 4 — Figure 4: NTP Accuracy on Booksum. ReFINE mid-training on DeltaNet-1.3B (a) and LaCT-760M (b) leads to a consistent increase in NTP accuracy on the validation dataset while that of SFT mid-training is stagnant. The error bars show the minimum and maximum values from three independent trials.

Figure 5: Ablation on

k

and

c

. We mid-train models with different numbers of tokens per rollout

k

( left ) and numbers of chunks per sequence

c

( right ). We evaluate on 16K-context samples from 12 tasks in LongBench ( Bai et al., 2024 ) . With cosine similarity reward, there is an optimal

k

. Higher

c

leads to more NSP training per sequence, which leads to better overall performance.

	Reward	Avg. Score
LaCT-760M + ReFINE MidTr	$\mathcal{R}^{\text{binary}}$	16.6
LaCT-760M + ReFINE MidTr	$\mathcal{R}^{\varphi}$ (ours)	16.9
DeltaNet-1.3B + ReFINE MidTr	$\mathcal{R}^{\text{binary}}$	16.5
DeltaNet-1.3B + ReFINE MidTr	$\mathcal{R}^{\varphi}$ (ours)	17.0

Table 6: Impact of reward function on mid-training. We compare

\mathcal{R}^{\text{binary}}

(binary exact match) and

\mathcal{R}^{\varphi}

(ours) reward strategies on 12 tasks in LongBench.

	Sampling	Avg. Score
LaCT-760M + ReFINE MidTr	Uniform	16.2
	$\arg\max H$	16.4
	$\arg\min H$	16.6
	Ours	16.9
DeltaNet-1.3B + ReFINE MidTr	Uniform	15.9
	$\arg\max H$	16.7

Table 7: Impact of token selection. We compare various token selection strategies on 12 tasks in LongBench: uniform random, selecting the token with maximum entropy (

\arg\max H)

or minimum entropy (

\arg\min H)

, and our entropy-weighted sampling.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Sequences and Tokens
Symbol	Description
$S=(x_{1},\ldots,x_{T})$	Input sequence with $T$ tokens
$x_{t}$	Token at position $t$
$x_{<t}$ , $x_{\leq t}$	Token sequence up to (and including) position $t$
$x_{t+1:t+k}$	$k$ -token ground-truth continuation starting from position $t+1$
$\hat{x}_{t+1:t+k}$	Generated rollout token sequence from position $t+1$

Appendix

Table A.1 : Glossary and notation.

Rewards
Symbol	Description
$R(t)$	Self-supervised reward at position $t$ (shorthand for $R(\hat{x}_{t+1:t+k},x_{t+1:t+k})$ )
$R^{\varphi}_{k}(t)$	Cosine similarity reward over $k$ -token rollout
$R^{\text{binary}}_{k}(t)$	Binary exact match reward over $k$ -token rollout
$R^{\text{hybrid}}_{k}(t)$	Hybrid reward ( $R^{\varphi}_{k}+R^{\text{binary}}_{k}$ )
$\mathcal{R}_{\mathcal{T}^{*}}$	Set of rewards for all sampled positions in $\mathcal{T}^{*}$

Appendix

Table A.2 : Glossary and notation (continued).

Phase	Dataset	Metric	Context	Size
Mid-training	Long-Data-Collections	-	16K	$\sim$ 200M tokens
	RULER NIAH	recall	4K/8K/16K	500 per context
	Booksum	NTP Accuracy, CE loss	$\leq$ 16K	9600
Post-training	RULER SQuADQA	recall	4K/8K/16K	1600 train / 200 test
Post-training	RULER HotpotQA	recall	4K/8K/16K	1600 train / 200 test
Test-time	NarrativeQA (NQ)	F1	$\leq$ 16K	56

Appendix

Table C.1 : Summary of datasets and benchmarks used across training phases.

Parameter	Value	Parameter	Value
Actor gradient clip	0.2	Learning rate	$10^{-6}$
Mid-Train batch size	128	Adam $(\beta_{1},\beta_{2})$	$(0.9,\,0.999)$
Post-Train batch size	64	Weight decay	0.01
Test-Time-Train batch size	8	Sampling temperature $\tau$	1.0
Mid-Train PPO mini-batch size	32	Max prompt length	16,384
Post-Train PPO mini-batch size	16	Entropy loss coefficient	0

Appendix

Table D.1 : Training hyperparameters.

Panel of Figure 17 — Figure E.1 : Validation Loss on Booksum Dataset. We track the NTP loss on the Booksum validation dataset in Long-Data-Collections ( TogetherAI, 2024 ) throughout mid-training on DeltaNet-1.3B (a) and LaCT-760M (b) . The validation loss for SFT on LaCT-760M does not decrease, as the model has already been pre-trained on the mid-training dataset. The error bars show the minimum and maximum values from three independent trials.

Panel of Figure 18 — Figure E.2 : NTP Entropy Distribution. We compute the token-level NTP entropy of a randomly selected sample using LaCT-760M (a) and DeltaNet-1.3B (b) . We use the same sample to extract the entropy distribution from both models.

Panel of Figure 19 — Figure E.3 : Reward Distribution. We report the mean and standard deviation of the cosine similarity reward during mid-training for different values of $k$ . As the rollout length increases, the mean reward ( a , b ) decreases and the standard deviation ( c , d ) also decreases.

Models	Wiki	LMB.	FDA	SWDE	LMB.	PIQA	Hella.	Wino.	ARC-e	ARC-c	Avg
	ppl $\downarrow$	ppl $\downarrow$	recall $\uparrow$	recall $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc_n $\uparrow$
LaCT-760M	20.8	29.8	36.7	66.0	32.4	67.5	41.5	53.4	48.4	27.6	45.1
+SFT MidTr	20.7	28.9	37.2	66.2	32.9	67.4	41.4	52.5	48.4	27.1	45.0
+ ReFINE MidTr	20.8	30.2	36.1	66.5	32.3	67.3	41.4	52.8	48.8	27.5	45.0
DeltaNet-1.3B	16.7	14.7	18.0	54.4	43.0	70.8	50.5	53.7	58.3	25.9	50.4
+SFT MidTr	16.8	15.3	17.7	54.5	42.7	71.1	50.1	53.8	58.4	26.4	50.4

Appendix

Table E.1: Performance on Short-Context Tasks. We evaluate mid-trained models on short-context benchmarks to verify that ReFINE does not cause catastrophic forgetting.

	MidTr	TTT	MidTr Reward	TTT Reward	Avg. Score
LaCT-760M	ReFINE	-	$\mathcal{R}^{\varphi}$	-	16.9
	ReFINE	SFT	$\mathcal{R}^{\varphi}$	-	17.0
	ReFINE	ReFINE	$\mathcal{R}^{\varphi}$	$\mathcal{R}^{\varphi}$	17.5
	ReFINE	ReFINE	$\mathcal{R}^{\varphi}$	$\mathcal{R}^{\text{binary}}$	18.0
DeltaNet-1.3B	ReFINE	-	$\mathcal{R}^{\varphi}$	-	17.0
	ReFINE	SFT	$\mathcal{R}^{\varphi}$	-	16.5

Appendix

Table E.2: Impact of reward function on TTT. We compare different TTT reward strategies on 12 tasks in LongBench: binary exact match between predicted and ground truth completion (

\mathcal{R}^{\text{binary}}

), and cosine similarity of the hidden states of predicted and ground truth completions (

\mathcal{R}^{\varphi}

Example	Type	Text	Cosine Reward
Ex. 1	GT	very laid back lodge	–
	P1	perfect hill	0.508
	P2	unique hotel to spend a	0.498
	P3	great place to stay	0.535
	P4	top-notch l	0.418
Ex. 2	GT	loved every minute of it	–

Appendix

Table F.1: Qualitative examples of cosine similarity reward assignment during mid-training. GT denotes the ground truth continuation, and P1–P4 denote the four predicted continuations generated by the model.

Fast Weights

See all 8 papers

Test-Time Training with Next-Token Prediction
Jun 19, 2026Xuan Ouyang, Zefan Cai, Junjie HuLong-Context Language Modeling Next-Token Prediction
University of Wisconsin–Madison
ResonatorLM: Causal Resonant Field Mixing for Efficient Long-Context Language Modeling
Jul 6, 2026Archie ChaudhuryLong-Context Language Model Inference Long-Context Language Modeling
Axionic Labs
FocuSFT: Bilevel Optimization for Dilution-Aware Long-Context Fine-Tuning
May 11, 2026Zehua Pei, Hui-Ling Zhen, Xianzhi Yu +3Supervised Fine-Tuning Self-Attention
The Chinese University of Hong Kong · Huawei Technologies Co., Ltd

Jun 19, 2026J/K move · Enter open · S save

Test-Time Training with Next-Token Prediction

Xuan Ouyang, Zefan Cai, Junjie Hu

University of Wisconsin–Madison

Long-Context Language Modeling Next-Token Prediction Test-Time Training Sequence Modeling Inference-Time Adaptation of Language Models

Open paper AskPDF arXiv

Next-token prediction is the self-supervised signal that trains language models, and every observed prompt token provides the same signal at test time. We study whether this signal can define the inner-loop objective for test-time training (TTT) in pretrained long-context language models. Many TTT architectures require models to be trained with test-time adaptation in mind, limiting their direct applicability to released LLM checkpoints. While recent in-place TTT methods make fast-weight adaptation possible for pretrained LLMs without redesigning the backbone, they leave a central question unresolved: what should each test-time write store? Existing recipes train the fast weight to match a learned local value proxy but they are not directly tied to the self-supervised next-token prediction signal. We introduce Test-Time Training with Next-Token Prediction (TTT-NTP), a drop-in fast-weight adaptation method for pretrained LLMs that instead supervises updates using the model's own next contextual hidden state. This makes each local write follow the same causal computation that supports next-token prediction: the value target is a pointwise linear projection of a single next-position contextual state. On RULER Full-13, averaged over 4k to 32k contexts, TTT-NTP is the only method that consistently improves the released backbone across four models spanning three families and a 0.6-8B size range, by 3.9 points on Llama-3.1-8B, 3.0 on Mistral-7B-v0.3, 4.1 on Qwen3-4B, and 2.9 on Qwen3-0.6B. On the real-world LongBench-v2 long-document QA benchmark, TTT-NTP improves over the base model by 5.6 points on Llama-3.1-8B and 3.7 on Mistral-7B-v0.3, while preserving commonsense and knowledge performance. Our code is publicly available at https://github.com/yancyou/TTT-NTP.

Abstract

Fast Weights

Explore similar work

Test-Time Training with Next-Token Prediction

ResonatorLM: Causal Resonant Field Mixing for Efficient Long-Context Language Modeling

FocuSFT: Bilevel Optimization for Dilution-Aware Long-Context Fine-Tuning