Paper deep dive
Learning What to Remember: Test-Time Training via Context Distillation
Zixuan Wang, Xingyu Dang, Rui-Jie Zhu, Zixin Wen, Hengyu Fu, Wenhao Chai, Jason D. Lee
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Effective long-context modeling is not merely about retaining more of the past, but about preserving the information that may prove relevant later. Test-time training (TTT) is an appealing approach that performs online parameter updates for long-context modeling, yet existing TTT methods only optimize either reconstruction or online adaptation objectives without considering the future utility of retained information. In this work, we propose \textbf{T}est-\textbf{T}ime \textbf{C}ontext \textbf{D}istillation (TTCD), a TTT framework that introduces a self-supervised objective for allocating limited memory capacity for future use. Specifically, TTCD uses a long-window teacher to supervise the fast weights of a short-window student, where the hidden-state discrepancy between them offers a dense, self-supervised signal guiding the model to memorize the contextual information crucial for future token predictions. We focus on an in-place variant: In-Place TTCD (IP-TTCD), which uses the existing MLP parameters as the fast weights. Experiments on long-context language modeling tasks show IP-TTCD consistently outperforms DeltaNet, Gated DeltaNet, sliding-window attention, and TTT when pre-trained from scratch. Furthermore, IP-TTCD allows pre-trained transformer models to adapt their parameters during inference through continual pre-training, gaining long-context capabilities with only a lightweight architectural augmentation. Our results position TTCD as a step toward architectural continual learning.
Tags
Links
- Source: https://arxiv.org/abs/2608.01672v1
- Canonical: https://arxiv.org/abs/2608.01672v1
Trouble viewing inline? Open PDF directly →
Full Text
75,447 characters extracted from source content.
Expand or collapse full text
@key newfloatplacement @keynewfloatname @keynewfloatfileext @keynewfloatwithin Learning What to Remember: Test-Time Training via Context Distillation Zixuan Wang1,2 Xingyu Dang111footnotemark: 1 Rui-Jie Zhu3 Zixin Wen4 Hengyu Fu2 Wenhao Chai1 Jason D. Lee2 1Princeton University 2UC Berkeley 3UC Santa Cruz 4Carnegie Mellon University Equal Contribution. Code available at https://github.com/dangxingyu/ttcd. Abstract Effective long-context modeling is not merely about retaining more of the past, but about preserving the information that may prove relevant later. Test-time training (T) is an appealing approach that performs online parameter updates for long-context modeling, yet existing T methods only optimize either reconstruction or online adaptation objectives without considering the future utility of retained information. In this work, we propose Test-Time Context Distillation (TTCD), a T framework that introduces a self-supervised objective for allocating limited memory capacity for future use. Specifically, TTCD uses a long-window teacher to supervise the fast weights of a short-window student, where the hidden-state discrepancy between them offers a dense, self-supervised signal guiding the model to memorize the contextual information crucial for future token predictions. We focus on an in-place variant: In-Place TTCD (IP-TTCD), which uses the existing MLP parameters as the fast weights. Experiments on long-context language modeling tasks show IP-TTCD consistently outperforms DeltaNet, Gated DeltaNet, sliding-window attention, and T when pre-trained from scratch. Furthermore, IP-TTCD allows pre-trained transformer models to adapt their parameters during inference through continual pre-training, gaining long-context capabilities with only a lightweight architectural augmentation. Our results position TTCD as a step toward architectural continual learning. 1 Introduction As language models are increasingly used for long-chain reasoning and agentic trajectories, long-context language modeling [41, 13, 18] has become a central direction for scaling models beyond short, isolated inputs. Full attention [47, 37] preserves all previous context, enabling exact retrieval and utilization of the past information but incurring computation and memory costs that grow quadratically with context length. For humans, in contrast, effective ‘long-context’ reasoning rarely requires the entire history: when processing a codebase or a book, people decide which information is likely to matter for future predictions and compress the rest away. Therefore, efficient alternatives face a harder problem than lossless retention: they must compress the past into a bounded state. Modern recurrent architectures approach this challenge by replacing the growing key–value cache with a compact recurrent state [38, 21, 20, 50, 22], achieving constant cost per token. Recent variants, including DeltaNet [38], Gated DeltaNet [50] and KDA [45], have narrowed the gap to softmax attention while showing favorable scaling. Test-time training (T) [42, 43] offers a complementary perspective by making a subset of the model’s parameters trainable at inference time, so that the weights themselves act as contextual memory. Several recent methods [55, 15, 44] demonstrate that this paradigm is beneficial and practical for long-context language modeling. However, under a fixed-capacity bottleneck, these alternative models must effectively decide at every step what to write, preserve, or forget to achieve better performance in a longer context. Existing T objectives provide only partial solutions to this issue. Reconstruction-based objectives [43, 4, 55] ask the state to store all history uniformly, but do not prioritize the storage of predictive information compared to irrelevant detail. The next-token objectives [15, 44] identify what helps the current prediction, but do not explicitly attribute that predictive computation to remote history rather than to the local context. Instead, we desire an objective that directly teaches which past information should be retained for potential future use. Therefore, we ask the following research question: Can we design a T objective that compresses past history according to its utility for future predictions? We propose TTCD (Test-Time Context Distillation), a test-time training method based on context distillation [40, 10, 27, 14, 6]. The context distillation objective captures the counterfactual discrepancy between two views of the same model with different context lengths. Specifically, we calculate the difference between a long-context self-teacher with more remote history tokens and a short-window student without the context when predicting the next token. The hidden-state difference indicate what the remote information contributes beyond the local context. The self-supervised objective therefore enables TTCD to update the fast weight, encouraging the limited parametric state to store distant information that the student would miss without the context, providing a dense representation-level signal of the essential past for potential future prediction. @picture Local context (recent tokens) Context of past histories (longer context) Student Model (short context window) Teacher Model (long context window) shared and frozen MLP Down-Projection fast weight down W_down ℒtCD=12‖(yt(T))−fdown(ht(S))‖22L^CD_t= 12 \| stopgrad(y_t^(T))-f_ W_down\! (h_t^(S) ) \|_2^2ht(S)h_t^(S)(yt(T)) stopgrad(y_t^(T))T update @picture Figure 1: Overview of TTCD. A shared and frozen backbone is run with two context lengths: a short-window student that observes only recent tokens and a long-window teacher that observes a longer history. TTCD treats the discrepancy between the teacher representation and the student’s projected hidden state as a context-distillation loss, and uses this loss to update the down-projection fast weight at test time. The predictive contribution of the self-teacher relative to a short-context student is distilled into fast weights, which persists after the evidence leaves the local window. The learned information can be retrieved by later related queries, benefiting future predictions. We validate the effectiveness of TTCD with extensive experiments on synthetic and long-context language modeling tasks across different model scales. When trained from scratch, TTCD consistently outperforms sub-quadratic and T baselines including sliding-window attention, Gated DeltaNet, DeltaNet, and IP-T, achieving the best sliding-window perplexity over context lengths up to 32K tokens, and its gains increase with context length. At 760M parameters, TTCD also more than doubles IP-T on the long context RULER benchmark with the gap widening as context length grows. Beyond from-scratch training, applying TTCD as continual pre-training to SmolLM2-360M/1.7B and LLaMA-3.1-8B improves RULER performance at 32K and 64K, demonstrating drop-in extensibility for existing pre-trained transformer-based models. Our results show a further step forward to the continual learning architecture in language models. 2 Preliminary and Related Work In this section, we introduce the notation, mechanisms and background used throughout the paper. We recall the Test-Time Training (T) framework as fast-weight memory, and describe context distillation as a supervision principle, which transfers information from a stronger teacher with more contextual information to a weaker student with a smaller window size. (a) Sliding Window PPL of 340M Model (b) Sliding Window PPL of 760M Model Figure 2: Sliding window perplexity on the Books [17] at increasing context lengths. We show that IP-TTCD consistently achieves lower perplexity than all the T-related baselines. 2.1 Test-Time Training Recurrent neural networks have a long history in AI [24, 8, 3]. Recent recurrent and test-time training architectures use fixed-size states or fast-weight updates as sub-quadratic alternatives to full attention [21, 16, 35, 33, 20, 12, 52, 50]. DeltaNet and Gated DeltaNet use delta-rule key-value state updates [52, 50]; Mamba uses selective state-space recurrences [20]; and Log-Linear Attention and RATTENTION interpolate between linear and softmax attention to improve in-context recall [22, 48]. While these models differ in their parameterization of the recurrent state, they share the view that long-context computation can be compressed into a compact dynamical memory rather than explicitly stored in the attention cache. Test-time training (T) treats a set of sequence-specific fast weights as the hidden state of a sequence model [2, 38, 43, 44, 15]. The fast weights are not frozen after training, but serve as dynamical states that encode the contextual information by consistently evolving under certain update rules. Specifically, let θ denote the slow weights shared across sequences, and let t W_t denote fast weights initialized at the beginning of each sequence and ℒ(,x)L( W,x) be the online loss function for input x. Given an input sequence representation (x1,x2,…,xt)(x_1,x_2,...,x_t) with current token xtx_t, the canonical T step consists of an update operation and a apply operation, where the update operation is computed by t=t−1−η∇ℒ(t−1,xt), W_t= W_t-1-η _ WL ( W_t-1,x_t ), (1) thus t W_t becomes a recurrent state over the time axis. The apply operation ot=ft(xt)o_t=f_ W_t(x_t) is done through a feature map ft(⋅)f_ W_t(·) associated with the fast weight t W_t, which produces the final output. Under the T framework, many design choices are investigated, including different loss functions, optimizers and memory parameterizations [50, 48, 4, 43]. Recently, there are two major competitive and efficient variants of T architectures111For computation efficiency, we only discuss the T variants without multi-layer backward propagation, e.g. T-E2E [44], in this section.: Large Chunk T (LaCT) [55] and In-Place T (IP-T) [15]. Large Chunk T. For computation efficiency, practical T methods often replace per-token recurrence with mini-batch [43] or chunk-wise updates [55]. By partitioning the sequence into chunks of size C, and letting ℓ(,⋅) ( W,·) denote the T loss at position t, the chunk update is i=t−1−ηC∑t∈[(i−1)C,iC]∇ℒ(i−1,xt), W_i= W_t-1- ηC _t∈[(i-1)C,iC] _ WL( W_i-1,x_t), (2) for i=1,…,T/Ci=1,…,T/C when sequence length T is divisible by the chunk size C. In this form, the choice of ℓ(⋅,⋅) (·,·) determines what information each chunk writes into the fast weights for later use. In-Place T. In-Place T adapts this framework to standard transformer language models by choosing existing model parameters as the fast weights [15], shedding light on “drop-in" designs for T without a fundamental architecture modification. In particular, In-Place T provided a bridge between fast and slow weights by reusing the native parameters (the down projection weights in the MLP layers), which is a fundamental step towards connecting temporary memory and long term memory. Our method follows this in-place formulation. 2.2 Context Distillation and KV compression Context distillation [40] is a long-held method in language model post-training, aiming to internalize in-context knowledge into the model’s parameters. Specifically, context distillation tries to match the logit outputs or hidden states of a teacher model to those of a student model, where the teacher perceives more context than the student is allowed to see. Formally, let ht(T)h_t^(T) denote the teacher’s latent representations and ht(S)h_t^(S) the student’s. A standard hidden-state distillation objective is ℒt=‖ht(S)−(ht(T))‖22L_t= \|h_t^(S)- stopgrad (h_t^(T) ) \|_2^2 (3) where stopgrad(⋅) stopgrad(·) detaches back-propagation of the teacher. Existing context-compression methods often use distillation to produce an offline prompt, memory, or compressed KV representation before generation [7, 28, 53, 32, 19, 34, 56, 30, 31]. Earlier work optimizes soft prompts for efficient adaptation [30, 31] or drops KV-cache entries to trade accuracy for efficiency [32, 19, 34, 56]. Newer long-context methods target million-token settings with more memory-efficient compression: Cartridges trains small document weights from synthetic QA pairs [14] and performs offline, corpus-specific optimization to compress contexts; KV-Distill learns token dropping via full-cache distillation [6]; Finch prunes caches using attention scores and calibration prompts [10]; and prior analyses find extractive compression competitive across tasks [27]. These methods typically follow an offline two-stage pipeline, producing a fixed surrogate before generation and decoupling compression from the model’s current computation. In contrast, TTCD is pretrained end-to-end to learn an online test-time update. These fast weights function as an evolving recurrent state, rather than as a fixed representation prepared before querying. 3 Test-Time Context Distillation In this section, we introduce Test-Time Context Distillation (TTCD), which updates the model’s fast weights using the context distillation objective. We first present our methodology on the context distillation training objective using shared weights teacher and student attentions (Section˜3.1), followed by a theoretical analysis on the benefit of our objective (Section˜3.2). Finally, we introduce our practical implementation details for architectural design. 3.1 Methodology Prior T objectives often adopt a reconstruction objective to compress key-value pairs in the T fast weight [43, 4, 55] or an equivalent state space [20, 50, 51, 50]. Recent works [15, 44] noticed the suboptimality of the reconstruction loss and proposed to use next-token prediction (NTP) loss. However, the two objectives face different drawbacks: • Reconstruction-type objectives often compress contexts uniformly, rather than allocating its capacity for more informative tokens. • NTP-type objectives often exhibit a recency bias, overemphasizing local signals near the prediction target. Our idea is to combine context distillation with a forward-looking objective, encouraging the fast weights to retain the information most useful for future predictions. Following the prior works [55, 15] which place parallel efficiency as top priority, we consider a chunk-wise formulation of T as well. Let X=(x1,…,xT)X=(x_1,…,x_T) be processed in chunks iC_i of size C. For each chunk, we run two causal sliding-window views of the same model. The teacher uses a long window of size wTw_T, while the student uses a shorter window of size wS<wTw_S<w_T. Thus, at the same token position, the teacher can condition on earlier prefix tokens that are hidden from the student, and therefore contains predictive information that depends on remote history. Matching the student to the teacher turns this difference in the hidden states into a self-supervised memory target. Hidden-state context distillation. For a given layer and position t, let ht(T)h_t^(T) and ht(S)h_t^(S) denote the pre-MLP hidden states produced by the teacher with window size wTw_T, and student with window size wSw_S, and let yt(T)=MLP(ht(T))y_t^(T)=MLP(h_t^(T)) and yt(S)=MLP(ht(S))y_t^(S)=MLP(h_t^(S)) be the corresponding MLP outputs. We use the stop-gradient teacher output as the target for the student to match: ℒtCD=12‖yt(S)−(yt(T))‖22,ℒiCD=1|i|∑t∈iℒtCD.L_t^CD= 12\|y_t^(S)- stopgrad (y_t^(T) )\|_2^2, _i^CD= 1|C_i| _t _iL_t^CD. (4) where LiCDL_i^CD is the chunk-wise aggregated loss. Note that TTCD performs online distillation that updates the model during test time, which is different from the traditional context distillation methods. IP-TTCD: Unifying slow and fast weights. Throughout this paper, we focus on an instance of TTCD: In-Place Test-Time Context Distillation (IP-TTCD). It follows In-Place T [15] that uses the MLP down-projection matrix as the fast weight for online adaptation. In a typical SwiGLU MLP layer, suppose the input hidden states are denoted by ht(r)h_t^(r) at step t. The output yt(r)y_t^(r) is computed by yt(r) y_t^(r) ≜MLP(ht(r)) (h_t^(r)) ≡down(upht(r)⊙σ(gateht(r))),r∈T,S. ≡ W_down ( W_uph_t^(r) σ( W_gateh_t^(r)) ), r∈\T,S\. (5) By considering the down projection weight down W_down as the online updated fast weight in chunk-wise T eq.˜2, combined with hidden state context-distillation objective eq.˜4, we can obtain the update rule for IP-TTCD. To be more specific, let’s express the intermediate hidden states of MLP by zt(r)=upht(r)⊙σ(gateht(r)),yt(r)=downzt(r),t∈i z_t^(r)= W_uph_t^(r) σ( W_gateh_t^(r)), y_t^(r)= W_downz_t^(r), t _i Let’s write i(r)=(zt(r))t∈i Z_i^(r)=(z_t^(r))_t _i and i(r)=(yt(r))t∈i Y_i^(r)=(y_t^(r))_t _i to represent the chunks of yty_t and ztz_t. At the i-th chunk iC_i, the current fast weight down(i−1) W_down^(i-1) is used to compute both teacher and student MLP outputs i(r),r∈T,S Y_i^(r),r∈\T,S\. The negative gradient of the chunk loss with respect to the student fast weight is −∇downℒiCD - _ W_downL_i^CD =−1|i|∑t∈i∇down12‖yt(S)−(yt(T))‖22 =- 1|C_i| _t _i _ W_down 12 \|y_t^(S)- stopgrad(y_t^(T)) \|_2^2 =−1|i|[i(S)−(i(T))][i(S)]⊤ =- 1|C_i| [ Y_i^(S)- stopgrad( Y_i^(T)) ][ Z_i^(S)] The IP-TTCD update for chunk iC_i is therefore down(i)=down(i−1)−ηi|i|[i(S)−(i(T))][i(S)]⊤. W_down^(i)= W_down^(i-1)- _i|C_i| [ Y_i^(S)- stopgrad( Y_i^(T)) ][ Z_i^(S)] . (6) The update direction is determined by the representation gap between long-window and short-window computation, while the write key is the student’s own MLP activation. The teacher–student difference determines what information is written, while the student activation determines when similar information can be retrieved later. The updated down-projection matrix therefore acts as a sequence-specific memory rather than only as an adapted model parameter. 3.2 Theoretical Analysis After introducing the future utility intuition behind our method, the following analysis shows how in-place context distillation writes a long-context signal from an earlier query into the fast weights, which could be reusable at later related queries. For simplicity, we analyze a single final-layer MLP block. Following the chunked update in Equation (6), if position n lies in chunk j, then all preceding chunks induce the fast-weight update Δ<j=η∑t∈ℐ<j(yt(T)−yt(S))zt(S)⊤, W_<j=η _t _<j (y_t^(T)-y_t^(S) )z_t^(S) , (7) where ℐ<jI_<j denotes the token positions in chunks strictly before chunk j. Constant factors from the squared distillation loss are absorbed into the learning rate η. Therefore, the student hidden state at position n is shifted by yn(S)←yn(S)+Δ<jzn(S)y_n^(S)← y_n^(S)+ W_<jz^(S)_n (8) Now consider a sequence containing an in-context learning episode, =[⋯,xc1,…,xc2⏟context,⋯,xq1,⋯,xq2], X=[·s, x_c_1,…,x_c_2_context,·s,x_q_1,·s,x_q_2], where (xc1,…,xc2)(x_c_1,…,x_c_2) is the supporting context span, q1q_1 is a representative answer-bearing position for the first query q1q_1, and q2q_2 is the analogous position for a later related query q2q_2. The long-window teacher can use the context span to answer q1q_1, while the short-window student cannot. Assumption 1. We make the following two assumptions. • Non-trivial in-context signal: q1−wT<c1≤c2≤q1−wSq_1-w_T<c_1≤ c_2≤ q_1-w_S and yq1(T)⊤(yq1(T)−yq1(S))≥cgap>0y_q_1^(T) (y^(T)_q_1-y^(S)_q_1 )≥ c_ gap>0. • Chunk causality and related-query alignment: if chunk j contains q2q_2, then q1∈ℐ<jq_1 _<j, zq1(S)⊤zq2(S)≥calign>0z^(S) _q_1z^(S)_q_2≥ c_ align>0, and zt(S)⊤zq2(S)=0z_t^(S) z^(S)_q_2=0 for all t∈ℐ<j∖q1t _<j \q_1\. The first assumption states that the teacher has a useful long-context signal at q1q_1 that the student lacks. The second isolates a later query whose student activation reads the earlier signal from the fast weights. The following theorem guarantees the effective memorization of the distilled hidden state within the one-step update, making it possible for future utilization. Detailed proof see Proposition˜2. Proposition 1 (Context distillation transfers an in-context solution). Under the setup and ˜1, for any learning rate η>0η>0, the update induced by preceding chunks satisfies yq1(T)⊤Δyq2(S)≥ηcgapcalign. y_q_1^(T) y_q_2^(S)≥η\,c_ gap\,c_ align. (9) Equivalently, the chunk update moves the later student hidden state at q2q_2 toward the teacher hidden state that already encodes the in-context solution for q1q_1. Remark 1. Proposition˜1 shows that the distillation update stores the teacher-student residual from q1q_1 and replays it at an aligned later query q2q_2. Thus, fast weights can preserve a useful long-context computation after the supporting context leaves the student’s attention window. 3.3 Practical Implementation Based on the theoretical guarantees, we verify the feasibility of our method conceptually. To further improve the expressivity and efficiency of our method, we further consider certain practical implementations and add some lightweight components. Off-policy update for context parallelism Equation (6) describes the ideal recurrent update, where each chunk is evaluated with the fast weight produced by all previous chunks. Directly materializing this recurrence limits parallelism, since the MLP outputs in chunk i depend on the fast weight produced by chunk i−1i-1. To improve parallelism over the context, our implementation uses an off-policy approximation: all teacher and student activations used to construct the update are computed with the base down-projection down(0) W_down^(0), while the resulting updates are accumulated causally. Practical feature transformation To further improve the performance of our model, we add some lightweight augmentation to the current TTCD update. We first adopt lightweight short causal convolutions convTconv_T and convSconv_S applied to teacher and student MLP activations, enhancing the model capability by introducing local mixing. For chunk iC_i, let us recall that i(r),r∈T,S Z_i^(r),r∈\T,S\ are the intermediate activations of the SwiGLU MLP layers, we transform them by ¯i(T)=convT(i(T)),¯i(S)=convS(i(S)). Z_i^(T)=conv_T( Z_i^(T)), Z_i^(S)=conv_S( Z_i^(S)). (10) We then form a value residual and normalized features, i=down(0)proj(¯i(T)−¯i(S)),i=ℓ2-(¯i(S)), R_i= W_down^(0) W_proj( Z_i^(T)- Z_i^(S)), K_i= _2- norm( Z_i^(S)), (11) where proj W_proj is a learned projection and ℓ2-(⋅) _2- norm(·) normalizes the student hidden states. The practical fast-weight correction is Δdown(i)=βiii⊤, W_down^(i)= _i R_i K_i , (12) followed by the causal parallel scan down(i)=down(0)+∑1≤j≤iΔdown(j). W_down^(i)= W_down^(0)+ _1≤ j≤ i W_down^(j). (13) This preserves the update structure of Equation˜6: the residual i R_i specifies the discrepancy between the teacher and the student, while i K_i specifies where it should be written in the student activation. NIAH (Avg) Common Sense Reasoning BABILong Model 4K 8K 16K 32K HellaSwag ARC-E ARC-C MMLU Average 4K 8K 16K 32K DeltaNet 41.67 38.25 24.83 19.33 45.81 48.61 28.67 26.44 37.38 16.99 8.85 3.90 1.94 Gated DeltaNet 32.21 29.88 22.29 18.62 48.15 54.76 29.44 25.25 39.40 19.24 11.93 6.08 2.84 SWA 44.88 22.58 10.79 4.21 47.32 52.53 30.72 25.74 39.08 25.42 15.45 7.91 4.40 IP-T 59.83 33.88 23.92 9.29 47.26 53.45 29.35 26.89 39.24 21.26 14.87 7.59 4.46 IP-TTCD 65.79 46.92 41.75 21.96 46.83 55.77 30.20 26.78 39.90 23.32 16.69 10.37 6.73 Table 1: Common-sense and long-context evaluation for 760M models. NIAH reports average RULER NIAH accuracy; Common Sense Avg is computed over HellaSwag, ARC-Easy, ARC-Challenge, and MMLU; BABILong reports average accuracy over five QA tasks. (a) MAD Benchmark Results. In-context Recall Noisy Recall Fuzzy Recall Compress Memorize Selective Copy Overall DeltaNet 99.9 86.7 25.1 42.3 55.5 100.0 67.8 GDN 99.9 99.7 18.2 47.3 86.9 99.9 74.4 LaCT 100.0 100.0 18.6 48.3 80.7 99.8 73.8 IP-T 98.1 97.1 52.4 43.4 82.9 96.6 78.4 IP-TTCD 99.3 98.8 56.2 50.0 88.9 99.2 82.1 (b) 760M S-NIAH-2 Figure 3: Table (a) Results on the MAD benchmark [36]. IP-TTCD attains the best overall score (82.1) and leads on the most memory-sensitive tasks (Fuzzy Recall, Compress, Memorize) that have not been saturated by prior baselines. Figure (b) Single-Needle-in-a-Haystack (S-NIAH-2) accuracy of 760M models from 4K to 32K. IP-TTCD consistently outperforms all baselines and degrades far more gracefully as the context grows. 4 Experiments In this section, we conduct extensive experiments to compare our method with various competitive T-related architectures and show the advantage of IP-TTCD. Specifically, we aim to answer: • Can IP-TTCD selectively preserve future-predictive past information under limited capacity? • Can IP-TTCD outperform prior T approaches when trained from scratch? • Can continual pre-training install the adaptability of TTCD and enhance pre-trained models? We first train IP-TTCD on synthetic tasks (Section˜4.1) requiring selective compression as a proof of concept. We then compare our method with the baselines on long-context language modeling pre-training tasks (Section˜4.2). We also explore the possibility of enabling pre-trained transformer models to acquire the test-time adaptability capability through continual pre-training (Section˜4.3). We finally examine the effects of the architecture components. 4.1 Synthetic Tasks We first evaluate some synthetic tasks as a proof of concept to show whether our model can selectively preserve useful information under limited capacity. Specifically, we compare IP-TTCD with other architectures on the MAD benchmark, a suite of synthetic sequence modeling tasks designed to probe whether an architecture can recall, denoise, compress, memorize, and selectively copy information under limited-context or recurrent-state constraints. The results are shown in Table (3). The MAD benchmark evaluates synthetic long-context abilities including recall, compression, memorization, and selective copying. On MAD, IP-TTCD achieves the best overall score and outperforms all baselines, showing clear gains on the most memory-sensitive or fuzzy recall tasks. 4.2 Pre-training From Scratch We further evaluate the approach’s long context performance when training a model from scratch. We analyze the long-context language modeling capabilities at two scales: 340M and 760M. Setup We use Long-Data-Collections [46] for pre-training with 32K context length. We compare IP-TTCD against various baselines including (1) In-Place T with the LM-aligned Conv1D target constructed from input embeddings[15].222This setting uses IP-T’s original recommended implementation and the LM-aligned formulation. We observe instability in IP-T’s training when the target is set as the hidden state of the current layer. Using input embedding as a target achieves clearly better performance. (2) standard Transformer with sliding window attention (SWA) (3) DeltaNet and Gated DeltaNet [38, 50]. All models are trained on sequences with a 32k sequence length. Please refer to Section˜A.1 for implementation details. Evaluation We first evaluate the long-context language modeling performance using sliding window perplexity [15] on the validation set, which measures perplexity on a fixed final block of tokens when extending the preceding context. For the 760M model, we further evaluate general downstream capabilities on four commonsense and knowledge benchmarks, including HellaSwag [54], ARC-Easy and ARC-Challenge [9], and MMLU [23]. We also evaluate long-context performance on RULER [25] and BABILong [29]. For RULER, we evaluate on six retrieval tasks across different sequence lengths, including Single and Multi-Key/Query/Value Needle in a Haystack tasks. Results Figure˜2 shows the sliding-window perplexity of models at different context lengths. Across both model scales, IP-TTCD consistently achieves the lowest perplexity among all compared architectures, outperforming sliding-window attention, DeltaNet, Gated DeltaNet, and In-Place T across nearly the entire context range. The advantage becomes especially clear as the context length increases: while several baselines saturate or degrade at longer contexts, IP-TTCD continues to benefit from additional context and reaches its best performance at the full 32K length. Beyond perplexity, IP-TTCD also demonstrates stronger long-context retrieval while maintaining competitive general capability. As shown in Table˜1, IP-TTCD achieves the best average commonsense score among all compared architectures, suggesting that the proposed adaptation does not compromise short-context reasoning. The advantage is especially clear on RULER NIAH, where IP-TTCD obtains the highest average accuracy at every context length and more than doubles IPTTT at 32K, improving from 9.29 to 21.96. It also consistently outperforms the baselines on BABILong across all evaluated lengths beyond 4k4k, confirming that the context-distillation objective better supports retention and reuse of distant information. SmolLM2-360M SmolLM2-1.7B Model 4K 8K 16K 32K 64K Avg 4K 8K 16K 32K 64K Avg Base 58.13 39.75 0.00 0.00 0.00 19.58 68.73 50.64 0.03 0.00 0.00 23.88 Base (Yarn-8x) 35.59 7.21 6.28 6.18 4.23 11.90 43.54 14.74 9.69 8.61 7.07 16.73 CPT 50.53 42.74 39.92 27.24 10.09 34.10 63.74 57.35 54.74 44.18 13.38 46.68 IP-T 50.20 42.43 40.47 29.32 10.35 34.55 67.41 60.13 54.89 41.53 14.44 47.68 IP-TTCD 49.79 42.13 40.59 33.55 9.79 35.17 67.61 59.58 55.50 42.96 19.17 48.96 Base Model Method 4K 8K 16K 32K 64K Avg Llama-3.1-8B Base 94.26 90.59 81.71 19.74 0.03 57.26 CPT 93.87 92.58 87.32 65.19 25.30 72.85 IP-T 93.95 92.31 88.04 66.17 26.50 73.39 IP-TTCD 93.72 92.32 88.59 66.48 27.65 73.75 Table 2: Evaluation results on RULER benchmark [25] based on continual pre-trained SmolLM2-360M/1.7B (Full attention), and LLaMA-3.1-8B. The scores are the average accuracy of all the 13 RULER tasks. 4.3 Continual Pre-training Since training LLMs from scratch is highly resource-intensive, we further show that TTCD can be applied as a lightweight augmentation to pre-trained transformer-based LLMs. We start with open-sourced base models without intensive context extension training to verify the performance of our method, showing improved performance gain by continual training on long-context datasets. Setup. We consider SmolLM2-360M/1.7B [1] as the main base models, and additionally evaluate on Qwen3-0.6B/1.7B in Appendix A.2.2. For data, we use the 64K-context book subset from ProLong [18] for long-context extension training. For more practical verification, we try our method on a larger scale model LLaMA-3.1-8B following the same continual learning protocol. We evaluate the long-context performance of the models on the RULER benchmark [25] with context length ranging from 4k to 64k. For baselines, we compare the IP-TTCD enhanced model against two baselines (1) full attention transformers (with YaRN length extrapolation) and (2) IP-T enhanced full attention transformers. More details of the settings are included in Appendix A.1. Results. The results, summarized in Table˜2 , show that the proposed adaptation methods consistently improve the long-context performance of SmolLM models over the attention-only tuning baseline. IP-TTCD achieves the best overall average long-context performance and substantially improves performance at 32K in our 360M model and 64K in our 1.7B model. These results indicate that IP-TTCD is particularly effective in extending the usable context range of the model. Additional experimental results and details can be found in Appendix A.1 and Appendix. A.2.2. 4.4 Ablation Studies We finally conduct the ablation studies on the design choices or hyperparameters of our method. Unless otherwise noted, we conduct the experiments on our 340M pre-training setting on Long Data Collections [46]. Figure 4: Ablation studies on IP-TTCD design choices. We evaluate three key factors using 340M pre-training experiments on Long Data Collections. Left: ablating the causal convolutional feature maps and learned projection in the fast-weight update shows that combining both components yields the best long-context perplexity. Middle: increasing the number of IP-TTCD layers consistently improves sliding-window perplexity, indicating that larger fast-weight capacity better supports long-context modeling. Right: increasing the chunk size from 512 to 2048 does not hurt performance and even slightly improves it, suggesting that IP-TTCD remains effective under large-chunk updates. Figure 5: Prefill throughput with our fused inference kernels. FlashAttention-3 stack (IP-TTCD uses the LSE-combination kernel); mean of four repeats on H100. Across all three settings, IP-TTCD’s overhead relative to IP-T is 1.3%1.3\%–6.6%6.6\% (Appendix˜B). Convolution and projection We first investigate the effectiveness of the additional short convolution on the sequence ¯i(T)=convT(i(T)) Z_i^(T)=conv_T( Z_i^(T)), and the learnable linear projection proj W_proj on the test-time gradient update. The ablation experiments in Figure˜4 (Left) show that the performance is the best when both parameters are used. The model with only convolution added in the update has a similar downward tendency in long-context perplexity as the model with both the convolution and the projection, but has much higher perplexity. Linear projection alone does not improve performance much, but it offers better training gain for convolution update. We conjecture that the phenomenon is related to the training dynamics but did not investigate further. Number of IP-TTCD layers We then consider how the performance scales with the number of IP-TTCD layers, which accounts for the state size of the implicit fast weight. Figure˜4 (Mid) shows that the sliding window perplexity decreases clearly when the number of layers increases, and the steeper reduction also indicates a better utilization of models with more IP-TTCD layers. Chunk size We finally ablate on the chunk size of the update, which controls the efficiency and parallelism level of our method. Surprisingly, we found that within 512,1024,2048\512,1024,2048\, larger chunk size does not hurt the performance. Instead, the model with chunk size 2048 even has the best performance. This further confirms the scalability of large chunk update mentioned in [55, 15]. Inference efficiency At inference, IP-TTCD computes two sliding-window attentions per IP-TTCD layer (teacher and student), which we fuse into dedicated dual-window kernels; fusion speeds up end-to-end prefill by 10%10\%–14%14\% over the unfused implementation. As Figure˜5 shows, the resulting prefill-time gap to IP-T is only +1.3%+1.3\%–+6.6%+6.6\% on the FlashAttention-3 stack [39] (−1.3%-1.3\%–+3.2%+3.2\% with FlashAttention-2 [11]); the remaining cost stems from the student MLP path rather than attention. Kernel designs, full measurements, and numerical-equivalence validation are given in Appendix˜B. 5 Conclusion We introduced TTCD, a test-time training framework that unifies past-context compression and online adaptation through a single context-distillation objective. Rather than asking the fast weights to store all the past information, TTCD uses the hidden-state discrepancy between a long-window teacher and a short-window student as a dense self-supervised signal, encouraging the limited parametric state to store distant information that is most useful for future predictions. Our in-place instantiation, IP-TTCD, realizes this idea as a low-rank correction to the existing MLP down-projection. We validate our method in both pre-training from scratch and continual pre-training, and find that IP-TTCD consistently improves long-context language modeling, retrieval, and reasoning over strong baselines. By repeatedly distilling what the short-window model misses into fast weights, IP-TTCD allows a model to acquire and reuse information as a sequence unfolds. We therefore view TTCD as a step toward architectural continual learning: the model not only adapts at test time, but also learns which parts of its experience should remain available for future predictions. Limitations. While IP-TTCD improves long-context performance, several limitations remain. A naive two-kernel implementation of the teacher–student attention slows down prefill by up to 1.3×1.3× at small teacher windows. With the fused inference kernels of Appendix˜B, the 64K prefill overhead relative to IP-T drops to −1.3%-1.3\%–+3.2%+3.2\% on the FlashAttention-2 stack and +1.3%+1.3\%–+6.6%+6.6\% on FlashAttention-3; training remains 1.1×1.1×–1.2×1.2× slower than standard Transformer baselines. Our from-scratch pre-training experiments are limited to 760M parameters, although our continual-pretraining experiments extend to 8B models. Scaling IP-TTCD from scratch to substantially larger models remains open. References Allal et al. [2025] Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blázquez, Guilherme Penedo, Lewis Tunstall, Andrés Marafioti, Hynek Kydlíček, Agustín Piqueres Lajarín, Vaibhav Srivastav, et al. Smollm2: When smol goes big–data-centric training of a small language model. arXiv preprint arXiv:2502.02737, 2025. Ba et al. [2016] Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin Ionescu. Using fast weights to attend to the recent past. Advances in neural information processing systems, 29, 2016. Bai et al. [2018] Shaojie Bai, J Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint arXiv:1803.01271, 2018. Behrouz et al. [2024] Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time. arXiv preprint arXiv:2501.00663, 2024. Biderman et al. [2024] Stella Biderman, Hailey Schoelkopf, Lintang Sutawika, Leo Gao, Jonathan Tow, Baber Abbasi, Alham Fikri Aji, Pawan Sasanka Ammanamanchi, Sidney Black, Jordan Clive, et al. Lessons from the trenches on reproducible evaluation of language models. arXiv preprint arXiv:2405.14782, 2024. Chari et al. [2025] Vivek Chari, Guanghui Qin, and Benjamin Van Durme. Kv-distill: Nearly lossless learnable context compression for llms. arXiv preprint arXiv:2503.10337, 2025. Chuang et al. [2024] Yu-Neng Chuang, Tianwei Xing, Chia-Yuan Chang, Zirui Liu, Xun Chen, and Xia Hu. Learning to compress prompt in natural language formats. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 7756–7767, 2024. Chung et al. [2014] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555, 2014. Clark et al. [2018] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018. Corallo and Papotti [2024] Giulio Corallo and Paolo Papotti. Finch: Prompt-guided key-value cache compression for large language models. Transactions of the Association for Computational Linguistics, 12:1517–1532, 2024. Dao [2024] Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations, 2024. Dao and Gu [2024] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060, 2024. Dao et al. [2022] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344–16359, 2022. Eyuboglu et al. [2025] Sabri Eyuboglu, Ryan Ehrlich, Simran Arora, Neel Guha, Dylan Zinsley, Emily Liu, Will Tennien, Atri Rudra, James Zou, Azalia Mirhoseini, and Christopher Re. Cartridges: Lightweight and general-purpose long context representations via self-study. arXiv, 2025. URL https://arxiv.org/abs/2506.06266. Feng et al. [2026] Guhao Feng, Shengjie Luo, Kai Hua, Ge Zhang, Wenhao Huang, Di He, and Tianle Cai. In-place test-time training. In The Fourteenth International Conference on Learning Representations, 2026. URL https://openreview.net/forum?id=dTWfCLSoyl. Fu et al. [2022] Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212.14052, 2022. Gao et al. [2020] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. Gao et al. [2025] Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. How to train long-context language models (effectively). In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7376–7399, 2025. Ge et al. [2023] Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310.01801, 2023. Gu and Dao [2023] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. Gu et al. [2021] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021. Guo et al. [2025] Han Guo, Songlin Yang, Tarushii Goel, Eric P Xing, Tri Dao, and Yoon Kim. Log-linear attention. arXiv preprint arXiv:2506.04761, 2025. Hendrycks et al. [2020] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300, 2020. Hochreiter and Schmidhuber [1997] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735–1780, 1997. Hsieh et al. [2024] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What’s the real context size of your long-context language models? arXiv preprint arXiv:2404.06654, 2024. Hu et al. [2024] Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06395, 2024. Jha et al. [2024] Siddharth Jha, Lutfi Eren Erdogan, Sehoon Kim, Kurt Keutzer, and Amir Gholami. Characterizing prompt compression methods for long context inference. arXiv preprint arXiv:2407.08892, 2024. Jiang et al. [2023] Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Llmlingua: Compressing prompts for accelerated inference of large language models. In Proceedings of the 2023 conference on empirical methods in natural language processing, pages 13358–13376, 2023. Kuratov et al. [2024] Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of llms with long context reasoning-in-a-haystack. Advances in Neural Information Processing Systems, 37:106519–106554, 2024. Lester et al. [2021] Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 conference on empirical methods in natural language processing, pages 3045–3059, 2021. Li and Liang [2021] Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 4582–4597, 2021. Li et al. [2024] Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. Advances in Neural Information Processing Systems, 37:22947–22970, 2024. Munkhdalai et al. [2024] Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention. arXiv preprint arXiv:2404.07143, 101:15, 2024. Oren et al. [2024] Matanel Oren, Michael Hassid, Nir Yarden, Yossi Adi, and Roy Schwartz. Transformers are multi-state rnns. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 18724–18741, 2024. Peng et al. [2023] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derczynski, et al. Rwkv: Reinventing rnns for the transformer era. In Findings of the association for computational linguistics: EMNLP 2023, pages 14048–14077, 2023. Poli et al. [2024] Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, et al. Mechanistic design and scaling of hybrid architectures. arXiv preprint arXiv:2403.17844, 2024. Radford et al. [2018] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018. Schlag et al. [2021] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9355–9366. PMLR, 2021. URL https://proceedings.mlr.press/v139/schlag21a/schlag21a.pdf. Shah et al. [2024] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. In Advances in Neural Information Processing Systems, 2024. Snell et al. [2022] Charlie Snell, Dan Klein, and Ruiqi Zhong. Learning by distilling context. arXiv preprint arXiv:2209.15189, 2022. Su et al. [2024] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024. Sun et al. [2020] Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International conference on machine learning, pages 9229–9248. PMLR, 2020. Sun et al. [2024] Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (learn at test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.04620, 2024. Tandon et al. [2025] Arnuv Tandon, Karan Dalal, Xinhao Li, Daniel Koceja, Marcel Rød, Sam Buchanan, Xiaolong Wang, Jure Leskovec, Sanmi Koyejo, Tatsunori Hashimoto, et al. End-to-end test-time training for long context. arXiv preprint arXiv:2512.23675, 2025. Team et al. [2025] Kimi Team, Yu Zhang, Zongyu Lin, Xingcheng Yao, Jiaxi Hu, Fanqing Meng, Chengyin Liu, Xin Men, Songlin Yang, Zhiyuan Li, et al. Kimi linear: An expressive, efficient attention architecture. arXiv preprint arXiv:2510.26692, 2025. TogetherAI [2024] TogetherAI. Long data collections database, 2024. Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. Wang et al. [2025] Bailin Wang, Chang Lan, Chong Wang, and Ruoming Pang. Rattention: Towards the minimal sliding window size in local-global attention models. arXiv preprint arXiv:2506.15545, 2025. Wen et al. [2025] Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them. arXiv preprint arXiv:2509.02046, 2025. Yang et al. [2023] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023. Yang et al. [2024a] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464, 2024a. Yang et al. [2024b] Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. Advances in neural information processing systems, 37:115491–115522, 2024b. Yuan et al. [2025] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Yuxing Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 23078–23097, 2025. Zellers et al. [2019] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th annual meeting of the association for computational linguistics, pages 4791–4800, 2019. Zhang et al. [2025] Tianyuan Zhang, Sai Bi, Yicong Hong, Kai Zhang, Fujun Luan, Songlin Yang, Kalyan Sunkavalli, William T Freeman, and Hao Tan. Test-time training done right. arXiv preprint arXiv:2505.23884, 2025. Zhang et al. [2023] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661–34710, 2023. Appendix A Experiment Details A.1 Details of Training And Evaluation Datasets For the pre-training experiments, we use the open-sourced Long Data Collection [46]. Specifically, since the original dataset on huggingface is down, we used a public open-sourced copy 333https://huggingface.co/datasets/emozilla/Long-Data-Collections-Pretrain-Without-Books. To make up the missing part and keep the data mixing ratio similar, we find the open-source Books444https://huggingface.co/datasets/Geralt-Targaryen/books3/tree/main dataset and mix it with the long-data-collection dataset with ratio 2:8. For evaluation, we also used another leave-out subset of the open-sourced Books dataset. For continual pre-training experiments, we used the tokenized 64K book subset from ProLong [18]. Training Details All models are trained on NVIDIA H100 GPUs, detailed hyperparameters are listed below. For pre-training we apply warmup-stable-decay (WSD) learning rate schedule [26] which couples better with Muon [49]. For pre-training, the teacher convolution is zero initialized, and student convolution is random initialized. For continual pre-training, we use a cosine learning rate decay to 0.1 times the peak learning rate after warm-up steps. Hyperparameter 340M Model 760M Model Continual Training Optimizer Muon Muon AdamW Learning Rate 3×10−33× 10^-3 3×10−33× 10^-3 1×10−51× 10^-5 Batch Size 500K tokens 1M tokens 1M tokens Weight Decay 0.1 0.1 0.1 Gradient Clipping 1.0 1.0 1.0 Warmup Steps 500 500 200 Sequence Length 32768 32768 65536 Tokens Trained 20B 40B 10B Sliding Window Size 2048 2048 Full (8192 for SWA) Student Window Size 1024 1024 2048 Hidden State Size 1024 1536 See base model Number of T layers 4 4 4 T learning rate 0.3 0.3 0.3 Table 3: Pre-training hyperparameters for 340M and 760M models, and continual pre-training. *For continual pre-training, we remove the normalization on i(S) Z_i^(S) for stability. Evaluation Details We use the official lm-evaluation-harness [5] for all commonsense reasoning, BABILong, and RULER benchmarks. For commonsense reasoning, we evaluate on HellaSwag, ARC-Easy, ARC-Challenge, and MMLU, and report accuracy following the default evaluation protocol of lm-evaluation-harness. For RULER, we evaluate long-context retrieval performance at different context lengths and report average accuracy over the corresponding Needle-in-a-Haystack tasks. For BABILong, we report the average accuracy over the evaluated QA tasks at each sequence length. For sliding-window perplexity, we follow the evaluation protocol of Feng et al. [15]: given a fixed final block of tokens, we progressively increase the amount of preceding context and compute the perplexity only on the final block. For the MAD benchmark, we follow the standard MAD evaluation protocol and report the best test accuracy after sweeping learning rates and weight decay for each task configuration. A.1.1 Synthetic MAD Benchmark Setup Model architecture. All models follow the same two-layer architecture to ensure a fair comparison. Each layer consists of a sequence mixing sublayer followed by a channel mixing sublayer (SwiGLU MLP), wrapped in a Pre-LN residual block with RMSNorm. The hidden dimension is d=128d=128 throughout. The SwiGLU MLP uses an intermediate dimension of 8d3 8d3 rounded to the nearest multiple of 256, resulting in dff=256d_f=256. All models use 4 attention heads with head dimension 32. Sliding window configuration. Since several MAD tasks use short sequences (e.g., Compression and Memorization use L≤256L≤ 256 and L=32L=32, respectively), a single window size does not yield a meaningful comparison across all tasks. For IP-TTCD, IP-T and LaCT, we use task-dependent window sizes: • ICR, Noisy-ICR, Fuzzy-ICR, Selective-Copying: Teacher window size W=32W=32, student/T chunk size C=16C=16, student visible window =1×C=16=1× C=16. • Compression, Memorization: Teacher window size W=16W=16, student/T chunk size C=8C=8, student visible window =1×C=8=1× C=8. GDN and DeltaNet are recurrent models and do not use a sliding window. Training. All models are trained with AdamW and a cosine learning rate schedule (minimum LR 10−610^-6) for 200 epochs with batch size 128 in BF16 mixed precision. We sweep over learning rates 10−4,5×10−4,10−3\10^-4,5× 10^-4,10^-3\ and weight decays 0,0.1\0,0.1\, yielding 6 hyperparameter configurations per task setting. We report the best test accuracy across the hyperparameter sweep for each (task, vocab_size, seq_len) configuration, following the standard MAD protocol. The total benchmark comprises 396 configurations across 6 tasks. All experiments use seed 12345. A.2 Additional Results A.2.1 Pre-training Per-task RULER NIAH breakdown. Table˜4 reports per-subtask accuracy for the 760M models from Section˜4.2. IP-TTCD attains the highest average at every context length, and the gap to the baselines widens with length: at 32K, IP-TTCD reaches 21.96, while IP-T-Embedding drops to 9.29 and all recurrent baselines stay below 20. The advantage is most pronounced on the harder retrieval tasks (S2, S3, MK1, MV, MQ), which require disambiguating distractors or composing multiple keys or values. SWA degrades sharply beyond its 2K training window, as expected. Model Context S1 S2 S3 MK1 MV MQ Avg IPTTCD 4K 48.00 100.00 91.00 71.00 43.50 41.25 65.79 8K 23.00 90.00 75.00 41.00 27.00 25.50 46.92 16K 14.00 70.00 74.00 42.00 25.00 25.50 41.75 32K 1.00 24.00 29.00 33.00 21.50 23.25 21.96 In-Place T 4K 96.00 68.00 69.00 70.00 30.25 25.75 59.83 8K 51.00 32.00 27.00 52.00 25.25 16.00 33.88 16K 23.00 23.00 10.00 45.00 24.25 18.25 23.92 32K 0.00 6.00 7.00 21.00 15.00 6.75 9.29 DeltaNet 4K 100.00 72.00 6.00 40.00 14.75 17.25 41.67 8K 100.00 47.00 15.00 36.00 15.00 16.50 38.25 16K 100.00 10.00 3.00 21.00 6.75 8.25 24.83 32K 100.00 2.00 3.00 6.00 1.25 3.75 19.33 Gated DeltaNet 4K 100.00 60.00 3.00 17.00 6.00 7.25 32.21 8K 100.00 51.00 2.00 12.00 5.00 9.25 29.88 16K 100.00 9.00 1.00 12.00 8.50 3.25 22.29 32K 98.00 2.00 2.00 4.00 2.75 3.00 18.62 SWA 4K 48.00 51.00 49.00 57.00 32.00 32.25 44.88 8K 23.00 24.00 24.00 27.00 19.00 18.50 22.58 16K 11.00 17.00 6.00 10.00 10.50 10.25 10.79 32K 1.00 6.00 3.00 5.00 5.25 5.00 4.21 Table 4: Detailed RULER synthetic NIAH results for 760M models. S1, S2, and S3 denote the three single-needle retrieval tasks; MK1 denotes multikey retrieval; MV denotes multivalue retrieval; MQ denotes multiquery retrieval. Scores are averaged over 100 examples per task and reported as percentages. Avg is the mean over the six NIAH subtasks at each context length. Higher is better. A.2.2 Continual Pre-training Base models. We use SmolLM2-360M and SmolLM2-1.7B [1] as base Transformer models. Both are originally trained with an 8K context window. We compare the original base model, YaRN-based length extrapolation, attention-only continual pre-training (CPT), IP-T [15], and our IP-TTCD augmentation. We additionally evaluate continual pre-training on Qwen3-0.6B and Qwen3-1.7B, whose base models already support long context natively, as a complementary stress test of our method. Training setup. For continual pre-training, we use the tokenized 64K book subset of ProLong [18]. All models are trained with sequence length 64K for 10B tokens. We use AdamW with learning rate 1×10−51× 10^-5, batch size 1M tokens, weight decay 0.1, gradient clipping 1.0, and 200 warmup steps, followed by a cosine decay schedule to 0.1 times the peak learning rate. For IP-T and IP-TTCD, we insert 4 T layers and use a T learning rate of 0.3. The teacher shortconv is initialized as identity, while the student convolution is initialized to zero in order to keep the model behavior the same before training. Evaluation setup. We evaluate long-context performance on RULER [25] with context lengths 4K, 8K, 16K, 32K, and 64K. For Qwen3, whose native context length is 32K, we report results up to 32K. We report the average accuracy over all 13 RULER tasks. All evaluations are conducted with the official lm-evaluation-harness implementation [5]. Additional results. In addition to the averaged RULER results reported in the main text, we provide the full breakdown for SmolLM2 in Table˜6, covering both full attention and sliding window attention (SWA) inference settings. We further report results for Qwen3 base models in Table˜7. Together, these results show that IP-TTCD extends the usable context range of base Transformer models, especially at longer context lengths where the original models and simple length-extrapolation baselines degrade sharply, and that the gains transfer across base models with very different long-context starting points. IP-TTCD recovers long-context performance under full attention. Under full-attention inference, IP-TTCD achieves the highest RULER average on both SmolLM2 scales, with the most visible gains at the longest contexts (e.g. 33.55 vs. 27.24 / 29.32 at 32K on 360M, and 19.17 vs. 13.38 / 14.44 at 64K on 1.7B), suggesting that the long-window distillation signal is most useful exactly where short-window supervision becomes least informative. Under SWA inference, where the long-range attention path is removed at test time, the three continual pre-training variants (CPT, IP-T, IP-TTCD) perform within roughly one RULER point of each other, which is consistent with the design of IP-TTCD: the distillation signal is informative when the long-window path is available, and benign when it is not. Robustness across different base models. We also evaluate IP-TTCD on Qwen3-0.6B and Qwen3-1.7B, which already have strong long-context capability out of the box. The pre-training data of Qwen3 is known to be of high quality, and continual pre-training on a new long-context corpus slightly degrades Qwen3-1.7B’s performance. Within this regime, IP-TTCD retains the highest average among continual pre-training variants on both Qwen3-0.6B and Qwen3-1.7B, indicating that it behaves consistently across base models with different long-context starting points. A.2.3 Additional Experiments: Scaling up to 8B Models To extend our results to a larger scale, we start from LLaMA-3.1-8B and conduct the additional continual pre-training experiments. Hyperparameter Value Base model LLaMA-3.1-8B CPT architecture IP-T and IP-TTCD Transformer layers 32 Hidden size 4096 MLP intermediate size 14336 Context length 65,536 RoPE base θ 500,000 T layers 0,8,16,24\0,8,16,24\ T chunk size 2048 T learning rate 0.3 T target Hidden state Tokens per optimizer step 2,097,152 Learning-rate schedule Cosine decay Peak / final learning rate 5×10−65× 10^-6 / 0 Warmup 200 optimizer steps Target training budget 10,000,000,000 tokens Table 5: Architecture and IP-T configuration of the 64K continual-pre-training run. Training setup. We use 10B tokens in the tokenized ProLong [18] dataset with 64K context length, but use all subsets as a more practical settings. The training tokens are mixed based on the original datasets’ (token-count) ratio. All models are trained with sequence length 64K for 10B tokens. We use AdamW with learning rate 5×10−65× 10^-6, batch size 1M tokens, weight decay 0.1, gradient clipping 1.0, and 200 warmup steps, followed by a cosine decay schedule to 0.1 times the peak learning rate. For IP-T and IP-TTCD, we insert 4 T layers. The training details are included in the Table˜5 below. Model Attn. 4K 8K 16K 32K 64K Avg4-32K_4-32K Avgall_all SmolLM2-360M Base — 58.13 39.75 0.00 0.00 0.00 24.47 19.58 Base + YaRN-8x — 35.59 7.21 6.28 6.18 4.23 13.82 11.90 CPT Full 50.53 42.74 39.92 27.24 10.09 40.11 34.10 IP-T Full 50.20 42.43 40.47 29.32 10.35 40.61 34.55 IP-TTCD Full 49.79 42.13 40.59 33.55 9.79 41.52 35.17 CPT SWA 56.33 38.36 18.23 12.42 8.16 31.34 26.70 IP-T SWA 58.40 38.83 18.26 12.90 8.10 32.10 27.30 IP-TTCD SWA 57.41 38.35 19.55 13.00 7.21 32.08 27.10 SmolLM2-1.7B Base — 68.73 50.64 0.03 0.00 0.00 29.85 23.88 Base + YaRN-8x — 43.54 14.74 9.69 8.61 7.07 19.14 16.73 CPT Full 63.74 57.35 54.74 44.18 13.38 55.00 46.68 IP-T Full 67.41 60.13 54.89 41.53 14.44 55.99 47.68 IP-TTCD Full 67.61 59.58 55.50 42.96 19.17 56.41 48.96 CPT SWA 68.58 47.85 25.22 13.44 8.66 38.77 32.75 IP-T SWA 69.84 48.39 25.00 13.18 7.36 39.10 32.75 IP-TTCD SWA 69.29 46.53 24.67 13.32 8.16 38.45 32.39 Table 6: Complete RULER evaluation of continual pre-training on SmolLM2-360M and SmolLM2-1.7B, under both full attention and sliding window attention (SWA) inference settings. Base and Base + YaRN-8x are evaluated without continual pre-training and serve as references. Avg4-32K_4-32K averages over 4K–32K, and Avgall_all averages over all five lengths. Bold marks the best result within each model size, attention setting block. Model Method 4K 8K 16K 32K Avg Qwen3-0.6B Raw 80.98 72.23 65.91 59.69 69.70 CPT 81.25 71.92 66.98 60.76 70.23 IP-T 80.44 72.69 68.03 59.60 70.19 IP-TTCD 82.27 72.52 66.97 60.89 70.66 Qwen3-1.7B Raw 89.17 83.63 78.27 70.32 80.35 CPT 87.21 81.47 74.87 67.69 77.81 IP-T 87.35 81.51 75.60 66.84 77.83 IP-TTCD 87.20 82.65 74.56 67.94 78.09 Table 7: RULER evaluation of continual pre-training on Qwen3-0.6B and Qwen3-1.7B. Raw denotes the original base model without continual pre-training, which already supports long context natively. The remaining rows are continual pre-trained on the same long-context corpus as the SmolLM2 experiments. Bold marks the best result among continual pre-trained variants within each model size. Appendix B Inference Efficiency: Fused Dual-Window Attention This appendix details the inference kernels summarized in Section˜4.4, together with complete measurements and the numerical-equivalence validation. Why prefill costs extra. During prefill, every IP-TTCD layer runs two sliding-window attentions over shared Q,K,VQ,K,V—the wTw_T-window teacher and the wSw_S-window student (wS=wT/2w_S=w_T/2 in all our recipes)—followed by the chunk-wise fast-weight scan in the MLP. A naive implementation issues two FlashAttention calls per layer and recomputes the student’s attention from scratch, even though the student’s key range is a strict subset of the teacher’s. The extra work (student attention, student MLP activations, short convolutions, Δ W accumulation) is linear in sequence length and independent of wTw_T, while the baseline attention cost scales with wTw_T, so the relative overhead grows as the teacher window shrinks: with off-the-shelf kernels we measure +18%+18\% prefill time versus IP-T at wT=2048w_T=2048 and +13%+13\% at wT=8192w_T=8192 under the wS=wT/2w_S=w_T/2 recipe (64K context, H100), and up to +28%+28\% with the larger student windows of earlier configurations—the origin of the “up to 1.3×1.3×” slowdown quoted in the limitations. Fused dual-window attention. The student’s window is contained in the teacher’s, so every score the student needs is already computed during the teacher’s pass. Our kernel, implemented in Triton, maintains two sets of online-softmax accumulators (output, running maximum, normalizer) per query block and sweeps the teacher’s KV range once, backward from the causal diagonal. Within the shared range the student’s PVPV product is recovered from the teacher’s by the row-wise rescaling identity PS=PTemT−mSP_S=P_T\,e^m_T-m_S, so keys and values are loaded once, QK⊤QK is computed once, and the second output is produced with a handful of extra vector operations. The loop is split into five segments (causal diagonal, shared mask-free interior, student boundary, teacher-only interior, teacher boundary) so the bulk of the sweep runs without any masking; numerics follow the CUTLASS FlashAttention-2 implementation step by step (base-2 exponentials, the same guard for fully-masked rows, normalization by reciprocal multiplication), yielding outputs that are bit-identical to FlashAttention-2 on 99.799.7–99.9%99.9\% of elements and within 1–2 units in the last place (ulps) in bf16 elsewhere. We additionally replace the depthwise short convolutions of the T layers with a dedicated kernel that operates directly on the chunked layout; it is bitwise-equal to the cuDNN path and 11.6×11.6× faster once the layout copies are included. Interestingly, this kernel is faster than FlashAttention-2 even while doing strictly more work: FlashAttention-2’s local-attention path applies the sliding-window mask on every KV block of the inner loop, whereas ours masks only diagonal and boundary blocks. As Table˜8 shows, the fused kernel produces both windowed outputs in less time than a single FlashAttention-2 SWA call, at every batch size. LSE-combination variant. When maintaining a custom kernel is undesirable, the same saving is available from stock kernels: attention over disjoint key sets can be merged exactly through the log-sum-exp that FlashAttention already tracks, o=eL1−Lo1+eL2−Lo2o=e^L_1-Lo_1+e^L_2-Lo_2 with L=log(eL1+eL2)L= (e^L_1+e^L_2). We obtain the student output from a wSw_S-window call and the teacher output by merging it with a shifted “far-band” call over the teacher-only keys beyond the student window, Attn(qwS:,k:−wS,v:−wS)Attn(q_\,w_S:,\,k_:-w_S,\,v_:-w_S) with window wT−wSw_T-w_S, using a one-pass merge kernel. The student output is bitwise identical to the stock kernel by construction. On FlashAttention-3, whose SWA path does not suffer from the per-block masking issue, this variant computes both outputs in roughly 30%30\% less time than two FA3 calls (Table˜8). Batch FA2 1×1× FA2 2×2× FA3 1×1× FA3 2×2× Fused (ours) LSE (ours, FA3) 1 13.2 20.3 5.7 9.9 11.3 6.6 4 56.1 84.4 24.8 40.5 45.0 30.7 Table 8: Attention-only prefill time (ms) at 64K context with wT=8192w_T=8192, wS=4096w_S=4096 (Qwen3-0.6B shape: 16 query / 8 KV heads, head dimension 128, bf16, H100). “1×1×” is a single teacher-window call; “2×2×” is the naive teacher-plus-student pair. Our fused kernel produces both outputs faster than a single FlashAttention-2 call; the LSE variant is the fastest option on the FlashAttention-3 stack. Median over 10 timed iterations after warmup. End-to-end prefill throughput. Figure˜5 in the main text and Figure˜6 report end-to-end prefill throughput of the continually pre-trained Qwen3-0.6B configuration (28 layers, 4 T layers) against the SWA Transformer and IP-T baselines, at 64K/batch 1, 64K/batch 4, and 256K/batch 1, sweeping the teacher window with wS=wT/2w_S=w_T/2. Every cell is the mean of four independent repeats with all three models measured back-to-back on the same GPU. On the FlashAttention-2 stack (all models use FA2 for their standard attention; IP-TTCD uses the fused dual-window and convolution kernels on its T layers), IP-TTCD is within −1.3%-1.3\% to +3.2%+3.2\% of IP-T—at the largest window (16K) it is slightly faster, because the fused kernel more than repays the dual-window cost. Relative to the naive two-kernel implementation, the fused kernels shorten end-to-end prefill by 10%10\%–13%13\% across windows (11%11\%–14%14\% higher throughput). On the FlashAttention-3 stack (all models on FA3; IP-TTCD uses the LSE variant plus the convolution kernel), throughput of every model rises by 25%25\%–90%90\% (more at larger windows) and the residual IP-TTCD overhead is +1.3%+1.3\% to +6.6%+6.6\%; at this point attention is no longer the bottleneck and the remaining cost is the window-independent student MLP path (student activations, convolutions, and the Δ W scan). A CUTLASS-level port of the dual-window design to FlashAttention-3 would recover most of this remainder; we leave it to future work. Figure 6: End-to-end prefill throughput on the FlashAttention-2 stack. Mean of four repeats on H100; IP-TTCD uses the fused dual-window and convolution kernels on its T layers. Across all settings the gap to IP-T is −1.3%-1.3\% to +3.2%+3.2\%. Numerical equivalence. Because IP-TTCD’s fast-weight recursion can amplify kernel-level rounding differences, we validate the kernels at three levels on the continually pre-trained checkpoint. (i) Kernel level: as reported above, the fused kernel matches FlashAttention-2 bit-for-bit on 99.799.7–99.9%99.9\% of output elements; the convolution kernel and the LSE student path are exactly bitwise. (i) Distribution level: over 10M tokens of held-out book text at 64K, the per-token KL divergence between the two implementations’ next-token distributions has median 7.8×10−47.8× 10^-4 nats (p99 =4.0×10−3=4.0× 10^-3), total-variation distance has median 1.25%1.25\% (p99 =3.8%=3.8\%), and greedy top-1 decisions flip on 1.76%1.76\% of positions. (i) Task level: RULER NIAH accuracy is identical in all 8 evaluated (task, context) cells (four NIAH tasks at 32K and 64K). We caution that for such recursive architectures, kernel changes should be signed off with task metrics rather than logit comparisons. Appendix C Theoretical Analysis C.1 Proof of Proposition˜1 We first restate the proposition. Proposition 2 (Context distillation transfers an in-context solution). Under the setup and assumptions above, for any learning rate η>0η>0, the update induced by preceding chunks satisfies yq1(T)⊤Δyq2(S)≥ηcgapcalign. y_q_1^(T) y_q_2^(S)≥η\,c_ gap\,c_ align. (14) Equivalently, the chunk update moves the later student hidden state at q2q_2 toward the teacher hidden state that already encodes the in-context solution for q1q_1. Interpretation. The vector q1(T) Y_q_1^(T) is the representation that solves q1q_1 using the long context. Equation (9) shows that the distillation update stores the teacher-student residual from q1q_1 in the fast weights and replays it at a later aligned query q2q_2. Thus, the fast weights can carry over an in-context learning result even after the original supporting context is no longer available to the short-window student. Proof. By Equations (7) and (8), yq1(T)⊤Δyq2(S) y_q_1^(T) y_q_2^(S) =yq1(T)⊤Δ<jzq2(S) =y_q_1^(T) W_<jz^(S)_q_2 =ηyq1(T)⊤∑t∈ℐ<j(yt(T)−yt(S))zt(S)⊤zq2(S) =η\,y_q_1^(T) _t _<j (y_t^(T)-y_t^(S) )z_t^(S) z^(S)_q_2 =ηyq1(T)⊤(yq1(T)−yq1(S))zq1(S)⊤zq2(S), =η\,y_q_1^(T) (y_q_1^(T)-y_q_1^(S) )z_q_1^(S) z^(S)_q_2, where the last equality follows from the assumptions. Applying the positive teacher-student gap and the related-query alignment gives yq1(T)⊤Δyq2(S)≥ηcgapcalign,y_q_1^(T) y_q_2^(S)≥η\,c_ gap\,c_ align, which proves Equation (9). ∎