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
λSFT
λRL
reward
lr
MidTr
8
5
1
128
1
0.2
Rφ
1e-6
PostTr
8
5
1
64
1
0.2
Rhybrid
1e-6
TTT
8
5
1
8
1
0.4
Rbinary
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 .
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
Rbinary
16.6
Rφ (ours)
16.9
DeltaNet-1.3B + ReFINE MidTr
Rbinary
16.5
Rφ (ours)
17.0
Table 6: Impact of reward function on mid-training. We compare Rbinary (binary exact match) and Rφ (ours) reward strategies on 12 tasks in LongBench.
Sampling
Avg. Score
LaCT-760M + ReFINE MidTr
Uniform
16.2
argmaxH
16.4
argminH
16.6
Ours
16.9
DeltaNet-1.3B + ReFINE MidTr
Uniform
15.9
argmaxH
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 ( argmaxH) or minimum entropy ( argminH) , and our entropy-weighted sampling.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Symbol
Description
Sequences and Tokens
S=(x1,…,xT)
Input sequence with T tokens
xt
Token at position t
x<t , x≤t
Token sequence up to (and including) position t
xt+1:t+k
k -token ground-truth continuation starting from position t+1
x^t+1:t+k
Generated rollout token sequence from position t+1
Appendix
Table A.1 : Glossary and notation.
Symbol
Description
Rewards
R(t)
Self-supervised reward at position t (shorthand for R(x^t+1:t+k,xt+1:t+k) )
Rkφ(t)
Cosine similarity reward over k -token rollout
Rkbinary(t)
Binary exact match reward over k -token rollout
Rkhybrid(t)
Hybrid reward ( Rkφ+Rkbinary )
RT∗
Set of rewards for all sampled positions in T∗
Appendix
Table A.2 : Glossary and notation (continued).
Phase
Dataset
Metric
Context
Size
Mid-training
Long-Data-Collections
-
16K
∼ 200M tokens
RULER NIAH
recall
4K/8K/16K
500 per context
Booksum
NTP Accuracy, CE loss
≤ 16K
9600
Post-training
RULER SQuADQA
recall
4K/8K/16K
1600 train / 200 test
RULER HotpotQA
recall
4K/8K/16K
1600 train / 200 test
Test-time
NarrativeQA (NQ)
F1
≤ 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 (β1,β2)
(0.9,0.999)
Post-Train batch size
64
Weight decay
0.01
Test-Time-Train batch size
8
Sampling temperature τ
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.
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.
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.
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 ↓
ppl ↓
recall ↑
recall ↑
acc ↑
acc ↑
acc_n ↑
acc ↑
acc ↑
acc_n ↑
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
-
Rφ
-
16.9
ReFINE
SFT
Rφ
-
17.0
ReFINE
ReFINE
Rφ
Rφ
17.5
ReFINE
ReFINE
Rφ
Rbinary
18.0
DeltaNet-1.3B
ReFINE
-
Rφ
-
17.0
ReFINE
SFT
Rφ
-
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 ( Rbinary ), and cosine similarity of the hidden states of predicted and ground truth completions ( Rφ ).
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.
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.
Contemporary language models are dominated by the transformer architecture, which leverages self-attention mechanisms to enable more efficient, parallelized training across a wide set of documents and corpora. This has allowed transformers to effectively model data across a wide range of modalities and contexts. However, transformers, along with their conventional counterparts such as recurrent neural networks (RNNs) and convolutional neural networks (CNNs), often struggle to maintain efficiency when processing long contexts. We introduce ResonatorLM, a new mechanism that replaces attention with a physics-derived alternative. ResonatorLM treats token sequences as a single, driven one-dimensional latent field and replaces attention dot products with causal functions of damped resonators. We implement ResonatorLM on a traditional network architecture and test it on standard long-context modeling tasks. We find that in a small, 6M matched setting, training and prefill speedups increase with sequence length, decode speed reaches 6.47x compared to that of a standard, optimized transformer at 32K tokens, and accuracy reaches 61.31 percent (compared to 55.32 percent) on WikiText.
Large language models can now process increasingly long inputs, yet their ability to effectively use information spread across long contexts remains limited. We trace this gap to how attention budget is spent during supervised fine-tuning (SFT) on long sequences: positional biases and attention sinks cause the model to allocate most of its attention to positionally privileged tokens rather than semantically relevant content. This training-time attention dilution (the starvation of content tokens in the attention distribution) weakens the gradient signal, limiting the model's ability to learn robust long-context capabilities. We introduce FocuSFT, a bilevel optimization framework that addresses this problem at training time. An inner loop adapts lightweight fast-weight parameters on the training context to form a parametric memory that concentrates attention on relevant content, and the outer loop performs SFT conditioned on this sharpened representation. Both loops apply bidirectional attention over context tokens while preserving causal masking for responses, reducing the causal asymmetry that gives rise to attention sinks and aligning inner-outer behavior. On BABILong, FocuSFT improves accuracy by up to +14pp across 4K--32K context lengths; on RULER, it raises CWE aggregation from 72.9% to 81.1% at 16K; and on GPQA with agentic tool use, it yields a 24% relative gain in pass@1. Attention analysis shows that FocuSFT reduces attention sink mass by 529× and triples context engagement during training. Code: https://github.com/JarvisPei/FocuSFT
Zehua Pei, Hui-Ling Zhen, Xianzhi Yu +3
The Chinese University of Hong Kong · Huawei Technologies Co., Ltd