TNT: Improving Chunkwise Training for Test-Time Memorization
Authors: Zeman Li
Ali Behrouz
Yuan Deng
Peilin Zhong
Praneeth Kacham
Mahdi Karami
Meisam Razaviyayn
Vahab Mirrokni
Venue: ICLR 2026 Poster (2026)
Keywords: Recurrent Neural Networks
Sequence Modeling
Primary area: foundation or frontier models, including LLMs

Abstract: Recurrent neural networks (RNNs) with deep test-time memorization modules, such as Titans and TTT, represent a promising, linearly-scaling paradigm distinct from Transformers. While these expressive models do not yet match the peak performance of state-of-the-art Transformers, their potential has been largely untapped due to prohibitively slow training and low hardware utilization.
Existing parallelization methods force a fundamental conflict governed by the chunksize hyperparameter: large chunks boost speed but degrade performance, necessitating a fixed, suboptimal compromise. To solve this challenge, we introduce TNT, a novel training paradigm that decouples training efficiency from inference performance through a two-stage process. Stage one is an efficiency-focused pre-training phase utilizing a hierarchical memory. A global module processes large, hardware-friendly chunks for long-range context, while multiple parallel local modules handle fine-grained details. Crucially, by periodically resetting local memory states, we break sequential dependencies to enable massive context parallelization. Stage two is a brief fine-tuning phase where only the local memory modules are adapted to a smaller, high-resolution chunksize, maximizing accuracy with minimal overhead. Evaluated on Titans and TTT models, TNT achieves a substantial acceleration in training speed—up to 17$\times$ faster than the most accurate baseline configuration—while simultaneously improving model accuracy. This improvement removes a critical scalability barrier, establishing a practical foundation for developing expressive RNNs and facilitating future work to close the performance gap with Transformers.

Published as a conference paper at ICLR 2026
TNT: IMPROVING CHUNKWISE TRAINING FOR TEST- TIME MEMORIZATION
Zeman Li Ali Behrouz² Yuan Deng" Meisam Razaviyayn Vahab Mirrokni
Peilin Zhong?
Praneeth Kacham²
Mahdi Karami²
University Southern California
"Googl Research
{zemanli i, al ibehrouu, dengyuan, pkacham, nahdika, razaviyayn, mi mirrokni ]0googl com
peilin. thongeco comeogealumoi..dd
ABSTRACT
Recurrent neural networks (RNNs) with deep test-time memorization modules, such as Titans and TTT. represent a promising, linearly-scaling paradigm distinct from Transformers. While these expressive models do not yet match the peak performance of state-of-the-ard Transformers, their potential has been largely un- tapped due to prohibitively slow training and low hardware utilization. Existing parallelization methods force a fundamental conflict governed by the chunksize hyperparameter: large chunks boost speed but degrade performance, necessitat- ing a fixed, suboptimal compromise. To solve this challenge, we introduce TNT. novel training paradigm that decouples training efficiency from inference per- formance through a two-stage process. Stage one is an efficiency- training phase utilizing hierarchical memory. A global module processes large, hardware-friendl chunks for long-range context, while multiple parallel local modules handle fine-grained details. Crucially, by periodically resetting local memory states, we break sequential dependencies to enable massive context par- allelization. Stage two brief fine-tuning phase where only the local memory modules are adapted to smaller, high-resolution chunksize, maximizing accu- racy with minimal overhead. Evaluated on Titans and TTT models, TNT achieves a substantial acceleration in training speed to 17x faster than the most ac- curate baseline configuration-wwe -hile simultaneously improving model accuracy. This improvement removes critical scalability barrier, establishing practical foundation for developing expressive RNNs and facilitating futuree work to close the performance gap with Transformers.
pre-
INTRODUCTION
The demand for modeling long sequences highlights fundamental limitation of standard softmax attention (Vaswani al. 2017): its quadratic complexity bottlenecks scaling. This has spurred extensive research into more efficient architectures.
Among these emerging paradigms, a particularly powerful approach is rooted in test-time memo- rization (Sun et al. 2024 Architectures leveraging this principle, which we refer to as deep mem- ory modules, utilize deep, online-adapted sub-network whose parameters encode in-context information. Prominent examples include Titans (Behrouz et al. (2025d) and At- las (Behrouz al. 2025a).. This method stands in sharp contrast linear memory modules (Yang et 2024a b. Dao & Gu 2024) Sun et al. 2023). which, despite their efficiency, are constrained by matrix-valuce hidden states and linear state transitions. By leveraging expressive non-linear objec- tives and update rules, deep memory modules can theoretically overcome these limitations. While these methods generally not yet achieve the oatt--of-the-at performance of Transformers, they represent potentially promising paradigm for efficient sequence modeling, provided their training bottlenecks can be resolved.
rapidly updated to
Despite their expressive advantages, deep memory modules lack the efficient training algorithms of their linear counterparts, leading to low hardware utilization. Unlike linear memory modules,
Work done while interning Google Research.

---

Published as a conference paper at ICLR 2026
Sequence 1x D
Global
Global Menory Chink Six
Memory
Local Mensrr Window Size
Local
Memory
Local
Memory
1-0
1-1
1-2
1-2
Local
Memory
Figure I:
The basic diagram for illustraing TNT memory hierarchy. In each row, the updates at the same value of t ran at the same time (run in parrallel). = 0 is the initialization of the memory.
which utilize hardware-efficient parallelization, deep memory modules face challenges stemming from non-linear recurrences (e.g.. LayerNorm between chunks) and the complexity of their deep structures. In practice, these challenges constrain their training to more frequent online updates on small data segments, resulting in poor computational throughput in training. This creates an in- herent tension, as these models typically rely on a fixed, small chunk size (e.g., 16 to 64 tokens) to balance memory layer expressiveness against training efficiency. Consequently, this trade-off between in-context learning capability and computational performance has become a critical bottle- neck preventing the application of these models to truly long sequences practice. Resolving this fundamental tension is the primary goal of this work.
Recent work attempts to mitigate this issue. Zhang et al. (2025) combines large chunks with local attention to enhance parallelism. However, this circumvents the inefficiency rather than solving it, complicates the analysis by mixing memory and attention, and neglects the need for small chunks (ideally 1) during inference. Concurrently, Guo al. (2025) proposed hierarchical memory sys- tem, but it is limited to linear memory modules and does not support short-term memories.
To resolve this tension, introduce TNT novel training paradigm for deep memory modules. Our core insight is that different components of the model should process information at different granularities during distinct training stages. TNT begins with an efficiency-focused pre-training stage designed to maximize throughput. This is achieved via a hierarchical memory system: global memory module operates on large, hardware-friendly chunks capture long-range context, while multiple local memory modules handle fine-grained details in parallel. Crucially, we introduce a periodic reset mechanism for the local memory states. This breaks the sequential dependencies herent even in non-linear RNNs (e.g., those with normalization between steps), enabling massive context parallelization. This is a key innovation, as efficiently parallelizing non-linear recurrences across the sequence length is a long-standing challenge, largely unsolved outside of Transformers and specialized linear RNNs (where parallel scans apply). Subsequently, a performance-fouurr fine-tuning stage adapts the model for optimal inference. During this stage, only the local memory modules are adjusted to use smaller chunk sizes, achieving high-resolution accuracy with minimal additional computational cost. This two-stage approach effectively decouples training efficiency from inference performance, significantly improving training scalability while addressing a key lim- itation of prior architectures. Furthermore, the local memory system itself be hierarchical, em- ploying multiple modules operating at different resolutions. This multi-resolution approach allows the model to capture complex, multi-scale temporal dynamics more effectively than single fixed chunk size.
TNT is general training paradigm applicable any deep memory module rather than a specific architecture. By decoupling training throughput from inference accuracy, we resolve fundamen- tal tension constraining prior work. This removes dependency on hardware-specine optimizations for small chunks and enables flexible exploration of the architectural design space. We believe this paradigm will open new research avenues towards replacing softmax attention. Our main contribu- tions are summarized as follows:
'TNT can be viewed as an abbreviation of Two-stage Non-linear Training TTT iNside 7TT. It also hints its "explosive" impact on training efficiency.
2

---

Published as a conference paper at ICLR 2026
We identify three fundamental challenges limiting the scalability and performance of deep mem- ory modules: 1) domain mismatch between memory compression and retrieval; 2) tradeoff be- tween memory performance and computational efficiency; 3) chunksize mismatch between train- ing and inference (Section
We introduce Q-K Projection, efficient mechanism to resolve the domain mismatch between memory compression and retrieval (Section 4.1.2).
We introduce novel hierarchical memory architecture with periodic state resets, enabling context parallelism for non-linear deep memory modules (Section 4.1).
We introduce efficient fine-tuning mechanism address chunksize mismatch between training and inference in deep memory modules (Section 4.2
Putting all above together, introduce TNT, a general two-stage training paradigm that de- couples training efficiency from inference performance by combining efficient pre-training with high-resolution fine-tuning (Figure Figure Section
We validate TNT on the Titans architecture, achieving up to a 17.37x training speedup while improving accuracy, significantly advancing the practicality of expressive RNNs (Section 5p.
Problem Definition and Notations We aim to train a neural network with parameters 0 E Rd to perform next-token prediction. Given a sequence to predict each token using the context of its preceding tokens (f,....,TT_l... attention formulation, each token I; is represented by d-dimensional vector. Each input token X is projected into query, key, and value vectors: qt, kt, 2 E Rdd For ease of notation in subsequent chunkwise operations, we define a function {(i,j) == i (i mod j). which finds the beginning of the chunk containing index i for chunk size
.11....,2L).. the model's objective is Following the
2 PRELIMINARY
This section reviews preliminaries. Expanded related work is in Appendix A
2.1 DEEP MEMORY MODULES VIA TEST-TINE TIME MEMORIZATION
A powerful paradigm for sequence modeling is Test-Time Memorization (Sun et al. 2024). which enhances models by incorporating a secondary, rapidly adaptable neural network. Unlike the pri- mary model parameters, or "slow weights" (0) updated only during training, this approach introduces "fast weights" (Schlag et al. 2021). These fast weights, denoted by W, parameterize a sub-network, f(W-)) Rdd + R. that is updated online-during both training and inference-based on incom- ing tokens to dynamically store contextual information. While these modules do not yet achieve SOTA results compared to Transformers (Arora et al. 2024 Behrouz et al.) 2025a). improving their training efficiency is crucial for enabling the wider experimentation needed to close this gap.
In this work, focus on similar/relevant principle: deep memory modules (Irie et al. 2021 Sun et al. 2024 Behrouz et al. 2025d a c). In contrast to linear memory modules (Sun et al, 2023 Yang al. 2024b Dao & Gu 2024 Karami & Mirrokni 2025 Hu al. 2025). which
characterized by linear state transitions, deep memory modules employ non-linear recurrence rules
and complex memory structures.
The core mechanism of a deep memory module can be distilled into two sequential operations for each input token: 1. Memory Compression and 2. Memory Retrieval. These are formally defined
Memory Compression:
W, < W 7ı V (( kz),
(1)
Memory Retrieval: O2 = (WW.g))
(2)
In Memory Compression, the fast weights W are updated via gradient descent, guided by self- supervised loss L(..-) (e.g., MSE) and a learned learning rate 7/. The objective associates a trans- formed key, fWW-1 1-1. kt). with its value, Ut, compressing information into the fixed-size neural memory (Wang et al. 2025; Behrouz et al.; (2025)). In Memory Retrieval, the updated W, processes query qe produce o,. These two operations performed iteratively for each token.
3

---

Published as a conference paper at ICLR 2026
2.2 CHUNKWISE PARALLEL TRAINING
The sequential dependency (W; depends on W in Eqs. 12 prevents parallelization across the sequence length. To address this, deep memory modules adopt chunkwise parallel training (Hua) et al. 2022 Sun et al. 2023) to enable hardware-efficient training.
The core principle to divide the input sequence into non-overlapping chunks of size C. Within each chunk, approximation of the gradient of the loss for every token is computed with respect to the fast-weight state from the beginning of that chunk. This formulation breaks the sequential token-to-token dependency for gradient calculation, which allows the updates for all tokens within the chunk to be computed in parallel. The formal operations for a token at time step are as follows:
Chunkwise Memory Compression: We -- WE(L.C)
Σ T=E(CCC
1 Vw ((W,
(3)
Chunkwise Memory Retrieval:
-((W.
(4)
Here, Wg(t.c) denotes the state of the fast weights at the start the chunk containing token (See definition of -(:-) at the end Section 1). Although the update to obtain W, still depends on prior tokens within its chunk, the summation of gradients can be implemented efficiently using par- allel operations (c.g.. cumulative summation), significantly improving hardware utilization during training. However, a sequential dependency remains: the final state of the fast weights from the k-th chunk, WKC. is used as the initial state for the (k + 1)-th chunk.
CHALLENGES IN DEEP MEMORY MODULES
While chunkwise parallelization enables deep memory modules train on long sequences, this paradigm introduces significant challenges that limit their practical performance and scalability. In this section, we outline three fundamental challenges with deep memory modules.
Challenge I: Lack of Efficient Training Implementations. primary challenge for deep mem- ory modules is the inefficiency of their training process, which leads poor hardware utilization. While chunkwise parallelization theoretically enables sub-quadratic scaling, in practice, the training throughput lags significantly behind that of linear memory modules. This discrepancy arises from fundamental tension between model expressiveness and computational efficiency.
To maintain a fine-grained learning signal, deep memory modules require small chunk sizes (e.g., 16-64 tokens) (Sun al.) 2024), which fail to saturate accelerators, making training memory-bound (rather than compute-bound), While linear memory modules use customized kernels (e.g., leveraging SRAM) (Sun et al.
40
Validation Perplexity
36.45
Train Chunksize
35
34.15
30
25
24.23
22.4
2023
Gu & Dao 2023 Qin et al. 2024| Yang
20
17.88
et al.
2024a c). this relies on linear state tran-
15.5
sitions and is incompatible with the large, non- linear states of deep memory modules.
15
13.78
10
The consequence is that deep memory mod- ules suffer from extremely low FLOPs uti- lization, often falling below 5-10% of peak hardware performance (Zhang et al.) 2025).. This severe inefficiency makes pre-train pro- hibitively slow and costly, creating major bot- tleneck that undermines the practical applica- tion of these expressive models to truly long sequences.
64
128
256
512
Inference Chunksize (C)
Figure 2: Sensitivity of inference chunk size on a 550M Titans model pre-trained with 64. Performance is optimal when the inference chunk size matches the training one.
Challenge 2: Inconsistency Between Memory Compression and Retrieval. A fundamental in- consistency exists between how the memory sub-network trained and how it is utilized. During Memory Compression (Eq. [], the sub-network f(W,,) is optimized to learn mapping from the
4

---

Published as a conference paper at ICLR 2026
key space the value space by associating keys (k²) with values (v)). However, during Memory Retrieval (Eq. 2) the network queried using a query vector (qt) instead of a key. This substitu- tion violates the intended input domain of the learned function, creating discrepancy between the training objective and the retrieval task. This domain shift can degrade the integrity of the learned mapping and limit the model's retrieval performance. Our empirical validation be found in Section 5.4
Challenge Performance Sensitivity to a Fixed Pre-training Chunksize. The chunk size hy- perparameter, C, governs the trade-off between training throughput and model expressiveness. Cur- rent practice for deep memory modules is to use the same fixed chunk size for both pre-training and inference. However, observe that inference-time performance is highly sensitive to this pre- training choice. For example, as shown in Figure 2 a model pre-trained with a chunk size of 64 achieves optimal perplexity only when evaluated with that same chunk size.
This result reveals a critical train-test mismatch and contradicts the intuition that smaller chunks at inference should yield superior performance by capturing fine-grad tieeggaieed dependencies with "fresher" learning signal. Instead, the model becomes over-specialized to the specific chunk resolution seen during training. This inflexibility a significant limitation; ideally, model pre-trained with large, hardwarefriendly chunk size should be adaptable enough to perform even better with smaller, more precise chunk sizes at inference. Current deep memory modules fail to achieve this adaptability.
TNT: AN IMPROVED TRAINING FRAMEWORK FOR DEEP MEMORY
To address the challenges outlined in Section g we introduce TNT, an im- proved training paradigm for deep mem- ory modules. Our framework is structured around two-stage process designed to re- solve the inherent tension between train- ing efficiency and inference performance: an Efficiency-focused Pre-training Stage and Performance-focuser Fine-tuning Stage.
Global
Compresion
Retrieval
Feod-
Nt: Local Memoey
The first stage maximizes training throughput by introducing novel hierar- chical memory architecture that enables unprecedented parallelism, directly ad- dressing the challenges of low hardware utilization and inconsistent memory objectives (Challenges 1-2). The second stage employs an efficient fine-tuning strategy that adapts the model to high- resolution, small-chunk inference, resolving the sensitivity the pre-training chunk size (Challenge 3). This two-stage approach effectively decouples training efficiency from inference performance, overcoming key limitation of prior deep memory architectures.
Compression
Retricual
Figure 3: Architectural overview of TNT Stage 1.
4.1 TNT STAGE 1: EFFICIENT-FOCUSED PRE-TRAINING
4.1.1 TNT MEMORY COMPRESSION: HIERARCHY MEMORY
Sequential state dependency prevents context parallelism (processing sequence shards in parallel across devices). To enable this, we propose that all parallel shards initialize their local memory with the same learned state, Wini.. This breaks inter-chunk dependency, allowing massive parallelization. However, this causes local memory modules to lose the global context. To solve this, we introduce global memory module, parameterized by V, that operates in parallel with the sharded local memories. The global memory processes the sequence with relatively large chunk size (c.g., 2048 or greater), allowing it to efficiently capture long-range dependencies while maintaining high
5

---

Published as a conference paper at ICLR 2026
hardware utilization. This creates hierarchical system where local memories handle fine- grained information within parallel shards, while the global memory provides the overarching context.
This hierarchical structure is flexible; model be designed with global and N local memory modules, each operating at a different resolution. For clarity of illustration, we will assume the simplest case where N 1. We defer the generalized formulation of TNT to Appendix E We formally define our memory compression mechanism.
TNT Memory Compression Rule. The hierarchical memory is updated as follows:
Global Memory Update. with large chunk size Cg.
The global memory state V evolves sequentially across the input
(A+1)CCa
ViK+ICC & ViCa
Σ t=C
k = (0......//CC. (5)
Local Memory Update. The local memory W operates in parallel sequence shards of length St. Within each shard, updates use smaller chunk size Ct.
W. - {
Wini
if = (mod St)
W E0-,0, Tr Nw (+(W
Otherwise
(6)
The global memory update (Eq. 5) follows standard chunkwise formulation where the state is carried over sequentially between large chunks. To maximize training throughput, the gradient for all tokens within global chunk is computed with respect to the initial state of that chunk, allowing for a highly parallelized update.
In contrast, the local memory update (Eq. () introduces key innovation: periodic state reset. This rule enforces that the local memory state, We, is reset to a shared, Icarnable initial state Winit at the beginning of each segment of length Sr, This periodic reset is the critical mechanism that breaks the long-range sequential dependency across the input, thereby enabling true context parallelism for the fine-graine -reigra local memory modules.
The hicrarchical design of deep memory modules boosts training efficiency through two-pronged approach. Global modules create hardware-saturatingg compute-intensin operations by processing large chunks. Concurrently, the local memory's reset mechanism enables context parallelism, where the sequence is processed as independent chunks that can be distributed across devices or stacked on single accelerator to substantially increase training throughput. 4.1.2 TNT MEMORY RETRIEVAL: Q-K PROJECTION
As identified in Challenge 2, the memory compression step (Eq. 6 optimizes f(W.-) to map the key space to the value space. However, retrieval, the network is queried using a query vector, qt. which may lie outside the learned key domain, degrading performance. To resolve this, we propose Q-K Projection: projecting the query q. onto the subspace spanned by previously observed keys. This ensures the input to the memory function is in the space memory was trained on. The final output combines retrieval from the global memory (raw query) and the local memory (projected query). We apply projection only locally its fine-grained nature makes it more sensitive to the mismatch
TNT Memory Retrieval Rule. The hierarchical memory is retrieved as follows:
(( (We. (w.
k kT
o, f(Veu.c))) +
|| 1 qt
r=E(t,C2)
Crucially, this Q-K Projection does not require storing all past keys, which would be computationally and memory prohibitive. Instead, the projection matrix, Σi=1 ||= E [Raxd can be maintained as running sum. This results in a constant-size state that is updated efficiently in chunkwise parallel manner. Since many modern deep memory modules normalize the query (q:) and key (k:) vectors by their L2 norm, the denominator in the Q-K projection can simplify to Σt= k, We provide further details this efficient implementation in Appendix C
k-

---

Published as a conference paper at ICLR 2026
Table 1: TNT reaches the target training loss up 17x faster than the baseline Titans.
The table
compares the time required for different 150M models to reach the same target loss 3.20.
Models
Implementation
C or C₁
Training Time (hrs)
Speedup
Titans
JAX
19.48
1.00x
Titans
JAX
16
10.79
181x
Titans
JAX
32
6.45
3.02
Titans
JAX
64
4.18
4.67x
Titans
JAX
3.25x
Transformer (wlo guting)
JAX
1.74
11.18x
Transformer (w guting)
JAX
1.38
14.10x
Transformer (wlo gating)
Fltentttentionn (Pallas)
1.23
1
Transformer (w gating)
FlashAttention (Pallas)
0.96
20.22
TNT TNT TNT
JAX
(8)
2.54
7.68x
JAX
(16)
1.65
11.78x
JAX
(32)
1.22
x
TNT
JAX
(64)
1.12
17.37x
TNT
JAX
(128)
1.16
16.75x
4.2 TNT STAGE 2: PERFORMANCE-ROCCUEED FINE-TUNING AT FINER RESOLUTION
Having addressed training efficiency in Stage 1, we now turn to optimizing for inference perfor- mance. An intuitive approach to enhance model resolution would be to evaluate the pre-trained model using a smaller chunk size. However, as established in Challenge 3, a direct mismatch be- tween the pre-training and evaluation chunk sizes leads to significant performance degradation.
Our key insight is that this train-test discrepancy can be rectified with minimal computational over- head. We empirically observe that brief fine-tuning phase, where the pre-trained model is updated for a small number of steps using a smaller local memory chunk size, not only recovers but often surpasses the original performance.
Based on this finding, we introduce Stage 2 of our TNT framework: a Performance-focusee Fine- tuning Stage. In this stage, continue training the efficiently pre-trained model with smaller local chunk size (C% < Cz.). This process adapts the model to the higher resolution required for optimal inference Challenge 3, bridging the gap between the large chunk sizes required for efficient training and the small chunk sizes that yield the best performance at inference.
fraction of the cost of pre-training. By doing so, Stage 2 directly resolves
This two-stage process decouples pre-training efficiency from inference requirements. The bulk of training uses maximum throughput (large chunks), while the final model is produced with minimal overhead. Furthermore, fine-tuning specializes the model for the ideal inference scenario: local paradigm of autore-
chunk size of one (C), = 1). This aligns with the standard prefill-and-decon gressive generation. The global memory handles the context prefill, and the optimized local memory handles iterative decoding.
EXPERIMENTS
We empirically evaluate our two-stage training framework, TNT. While TNT is model-agnostic, we instantiate it with a strong deep memory model, Titans (Behrouz et al. 2025d), to demonstrate its effectiveness. We validate claims about training time and model accuracy in our experiments.
5.1 EXPERIMENTAL SETUP
Baselines. We compare against several strong long- long-context architectures. Our primary comparison is Titans (Behrouz et al., 2025d), our base model. We also benchmark against vanilla Transformer (Vaswani et al. 2017). Gated Transformer (Qiu et al. 2025), and TTT (Sun et al. 2024).
Training and TNT Configuration. We train 150M parameter models following (Behrouz et al. 2025d), using a T5 tokenizer (32k vocab). We use the Adam optimizer (Loshchilov & Hutter 2017) with 0.1 weight decay and cosine schedule (peak LR 1 x 10 10-3), Experiments on a TPUv4 pod (2x2x2 topology, model parallelism 2). For TNT, the local modules configuration denoted by their chunksizes, Ct {CI,1, .... CL.N }. For instance, Cz (8, 16) indicates two local modules with chunksizes and 16. The global memory uses CG = 2048.
conducted
Experimental Configurations. For efficiency benchmarks (Sec. 5.2). vary context length (2k-32k) with a 0.5M token batch size and local window S, = 2048. For performance evaluation (Sec 5.3). use 16k context length, 1M token batch size, and St. 4096.

---

Published as a conference paper at ICLR 2026
5.2 FASTER MEMORY TRAINING WITH TNT
Linear Runtime Scaling with Sequence Length. We first analyze single-step runtime perfor- mance by varying the sequence length while keeping the number of tokens per batch fixed. As shown in inFiguree 4 TNT's runtime grows linearly with sequence length, in contrast to the quadratic growth of Titans and standard attention. This scaling advantage is significant at long contexts. At 32K quence length, TNT is 5.1x x faster than a comparable Titans model with the same memory chunksize (C, = 16). We also observe that larger local chunk sizes consistently improve TNT's speed; with CL = {128), TNT x.3x faster than the highly optimized PlashAttention kernel (Dao 2024
TNT's highly parallelizable architecture achieves runtime that scales linearly with sequence length, a key advantage over the quadratic complexity of standard attention. Although models like Titans are also the- oretically linear, their inherent sequential dependencies impede effective paralleliza- tion, resulting in poor hardware utilization and slower wall-clock times on long sequences. As sequence length increases, TNT's superior scalability creates a crossover point where it becomes significantly faster. This efficiency most pronounced very long contexts; for instance, at sequence length of 32K, native JAX implementation of TNT (Cz 128) outperforms even the highly optimized FlashAttention kernel, confirming its suitability for demanding long-context training scenarios.
Actemiitoiioon
Attention
Thans (C-26]
Runtime (ms)
(C2-16)
TNT (C2m128)
103
2048
4096
8192
16384
32768
Sequence Length
Figure 4: Runtime comparison of different mod- els and implementations across varying sequence lengths, with the number tokens per batch fixed at 0.5M. Additional results are presented in FFigure
Timalto-Quality Comparison. We next translate these single-step runtime gains into practical time-to-quality setting. As shown in Table I, our TNT framework significantly accelerates the to- tal training time required reach target model quality. Our best config configuration achieves this up to 17.4x faster than the original Titans baseline. This efficiency gain is fundamental to the archi- tecture; for instance, using an identical local memory chunksize 8, TNT already 7.7x x faster than its Titans counterpart. While competitive with standard vanilla Transformers in JAX, our im- plementation does not yet outperform highly optimized baselines like the Gated Transformer with FlashAttention (Dao al. 2022). This is expected result, as TNT currently lacks custom kernel, which we leave for future work. Nonetheless, these results establish TNT as an efficient foundation for research on recurrent models, with a clear path toward matching the speed of state-of-the-art Transformers.
5.3 TNT IMPROVES MODEL QUALITY
Our TNT framework significantly enhances model quality, outperforming strong RNN-based base- lines and standard Transformer implementations. As detailed in Table 2, the initial Stage 1 pre- training highly effective its own. Our best Stage I model achieves an average perplexity of 23.13, a marked improvement over the best best-performing Titans model (25.07) and the vanilla Trans- former (23.58). While TNT does not fully match the perplexity of the state-of-the-un ohr Gated Trans- former (22.39), it achieves a higher average accuracy on common-sense reasoning tasks (41.0% % vs. 39.7%). At this scale, we consider perplexity more stable metric for language modeling capability, downstream task accuracy can be subject to higher variance.
Furthermore, the Stage 2 fine-tuning process offers an efficient method to further boost perfor- mance. This stage is computationally inexpensive, requiring only additional 5% of the original pre- pre-training compute (see Table 4), yet consistently lowers the average perplexity to a final value of 23.09. These results validate TNT as effective framework for producing high-quality models that surpass the limitations of prior RNN-based architectures and stand as a strong alternative to standard Transformers.
8

---

Published as a conference paper at ICLR 2026
Performance of TNT (150M parameters) and baselines language modeling and common-sense reasoning tasks, trained on 10B tokens, For TNT models, we only use global memory and use Cg to denote the global chunksize and the N local modules configuration is de- noted by their chunksizes, Cz [CI,1,.... CL.N). For instance, Cz = {8, 16) indicates two local
Table 2:
modules with chunksizes 8 and 16. The best results within block are highlighted The detailed
training time is reported Table 4
Model
Co
Cor
C4 ppl 1
FineWeb ppl
PG19 ppl
Avg. ppl
PIQA acc 7
Hella. 800
ARC-e ace
CSQA acc
acc 1
150M params/ 10B tokens
Transformer (wlo gating)
20.98
20.59
29.18
23.58
62.0
30.9
34.8
25.5
38.3
Transformer (w gating)
19.82
19.61
27.75
32.2
36.8
26.7
DeltaNet 2024c
22.49
22.36
31.84
25.56
62.6
32.2
35.9
27.0
39.4
Gated DelraNet 2025
21.25
21.37
31.1
35.2
27.8
TTT 2024
256
24.18
24.31
34.36
27.62
60.6
30.8
34.1
26.9
38.1
Titans 2025d
256
23,53
24.13
33.73
27.13
61.3
30.8
35.1
27.8
38.8
22.25
32.0
27.8
TNT Stage I: Etticiency-Foeue Focuuse Pre-training
TNT Siage
2048
[8]
21.04
21.01
30.24
24.10
61.8
32.8
37.4
30.3
40.6
2048
{8.16)
20.74
20.73
29.94
23.80
63.5
32.4
37.4
30.6
41.0
2048
(4.8.16)
20.47
20.43
29.43
23.44
32.4
36.4
28.9
40.2
(4,8,
20.15
32.0
36.7
30.3
INT Stage 2: Performance- Focusse Fine-tuning Stage
models
TNT Stage 2
2048
[1]
20.86
20.91
30.21
23.99
63.2
32.8
37.4
30.1
40.9
2048
(2,4)
20.65
20.70
29.97
23.77
63.4
32.5
37.3
30.2
40.9
2048
{2.4.8
20.32
20.35
29.42
23.36
64.0
32.0
36.9
28.1
40.3
2048
(2,4,8.16)
20.10
20.13
29.05
23.09
32.3
37.4
30.2
40.9
5.4
ABLATION STUDY
We conducted ablation study to validate TNT's key design choices, with results summa- rized in Table
Table Ablation study on TNT. The results show the contribution of each proposed change to the deep memory modules. TNT
Hierarchical Memory. The effectiveness of hierarchical design is evident. Incremen- tally adding local memory modules consis- tently improves performance over the Titans baseline. reducing perplexity from 23.53 to
N
Language Modeling
C.S. Reasoning
ppl
ace 1
Base Model (Titans)
23.53
38.8
TNT Stage (1 Global Momory)
Local Kemozy
+2 Local Memory
20.74
41.0
+3
Local Kemory
20.47
40.2
20.15 with four local modules. Conversely, re-
Kemozy
moving the global memory is detrimental (PPL
Stage
40.6
increases 25. 25.60), confirming its critical role in capturing long-range dependencies that are otherwise lost due to the local memories' reset mechanism.
wlo global memory
25.60
35.5
w/o Q-K projection
22.01
36.4
Stage
Q-K Projection. The query-key projection proves essential for performance. Its removal incurs substantial penalty (PPL increases from 21.04 to 22.01), validating our hypothesis that it is necessary to mitigate the compression-retrieval mismatch (Challenge 2).
Stage 2 Fine-tuning. Applying Stage 2 fine-tuning further enhances model capabilities, improving both language modeling (20.86 PPL. and common-sene reasoning (40.9%). This demonstrates its effectiveness in adapting the pre- trained models for high-resolution inference.
CONCLUSION
We introduce TNT, a two-stage training framework that resolves the fundamental conflict between training efficiency and inference performance in deep memory modules. By leveraging a hierar- chical memory architecture with periodic state resets, TNT enables massive context parallelism during pre-training, followed by efficient fine-tuning for high-resolution inference. Our experiments demonstrate up to 17x speedup compared to the most accurate RNN baselines while simultane- ously improving performance. TNT removes critical scalability bottleneck, significantly improv- ing the practicality of deep memory modules and facilitating future research close the performance gap with Transformers.

---

Published as a conference paper at ICLR 2026
REFERENCES
Simran Arora, Sabri Byuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Re. Simple linear attention language models balance the recall-throughput tradeoff. In Forty-first International Conference Machine Learning, 2024. URL https:
//openzeview. wttp://ww.eett/oommmdeddeeee
Ali Behrouz, Zeman Li, Pranceth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Raza- viyayn, and Vahab Mirrokni. Atlas: Learning to optimally memorize the context at test time. arXiv preprint arXiv:2505.23733 2025a.
Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, and Vahab Mirrokni. It's all connected: A journey through test-time memorization, attentional bias, retention, and online optimization, 2025b. URL
fatt:s://axxxv.orr//ab//////
13173
Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, and Vahab Mirrokni. It's all connected: A jour- ney through test-time memorization, attentional bias, retention, and online optimization. arXiv preprint arXiv:2504.13173, 2025c.
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time. Ad- vances Neural Information Processing Systems, 38, 2025d.
Róbert Csordás, Christopher Potts, Christopher D Manning, and Atticus Geiger. Recurrent neural networks learn to store and generate sequences using non-linear representations. In Proceedings of the 7th BlackboxNI.P Workshop: Analyzing and Interpreting Neural Networks for NLP, pp- 248 262, 2024.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openreview net. wwt. fort/ m2n22y9
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:245..20660, 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances neural information processing systems, 2022.
Xavier Gonzalez, Andrew Warrington, Jimmy Smith, and Scott Linderman. Towards scalable and stable parallelization of nonlinear rnns. Advances in Neural Information Processing Systems, 37: 5817-5849, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023.
Han Guo, Songlin Yang. Tarushii Goel, Eric P Xing, Tri Dao, and Yoon Kim. Log-linear attention. arXiv preprint arXiv:2506. 040600,0446..... 2025.
Ramin Hasani, Mathias Lechner,
Wang. Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. In The Eleventh International Confer- ence Learning Representations, 2023. URL https: /ttp://////peeeiee... net forum?id
Donald Olding Hebb. The organization of behavior: A neuropsychological theory. Psychology press, 2005.
John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8):2544225 2558, 1982.
Jerry Yao-Chieh Hu, Dennis Wu, and Han Liu. Provably optimal memory capacity for modern hopfield models: cansformer-compattiie dense associative memories as spherical codes. arXiv preprint arXiv:2410.23126 2024.
10

---

Published as a conference paper at ICLR 2026
Jiaxi Hu, Yongqi Pan, Jusen Du, Disen Lan, Xiaqiang Tang, Qingsong Wen, Yuxuan Liang, and Weigao Sun. Comba: Improving nonlinear mns with closed-loop control. arXiv preprint
S0MV::2506002475. 2025.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoe V. Le. Transformer quality in linear time, 2022. URL httptttts//:aiiiii org/ abs/2202. 10:02.104400 10447
Kazuki Irie, Imanol Schlag, Robert Csordas, and Jurgen Schmidhuber. Going beyond linear trans- formers with recurrent fast weight programmers. Advances in neural information processing systems, 34:7703-7717, 2021.
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bemsteen. Muon: An optimizer for hidden layers in neural networks, 2024. URL https:
//kellerrjordee www.erroordnnggithhtta wttpbi.oooosts/mmmo...
M. Karami and V. Mirronni. Lattice: Learning efficiently compress the memory, 2025.
Mahdi Karami, Ali Behrouz, Praneeth Kacham, and Vahab Mirrokni. TRELLIS: Learning to com-
press key- memory in attention models. In Second Conference Language Modeling, 2025.
URL http::t://pppeeeeiiiiiiieae/ wwt.iiewww..eee.....ee..ee.eee.e wtte://ooomiidddrllllle
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980. 2014.
Dmitry Krotov. Hierarchical associative memory. arXiv preprint arXiv:2107.06446 2021.
Dmitry Krotov and John J Hopfield. Dense associative memory for pattern recognition. Advances in neural information processing systems, 29, 2016.
Xiaoyu Li, Yuanpeng Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. On the expressive power of modern hopfield networks. arXiv preprint arXiv:2412.05562, 2024.
Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non-linear sequential models over the sequence length. In The Twelfth International Confer- ence Learning Representations, 2024. URL htllh:tts///ecc.. httsss//oppeeeieww net wwt.:/rrretrrrre//// forum?id E34AIVLNOV
Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, and Qiang Liu. Longhorn: State space models are amortized online learners. arXiv preprint arXiv:2407. 14207, 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17Xiv:1711.05101. 2017.
Carlo Lucibello and Mare Mézard. Exponential capacity of dense associative memories. Physical Review Letters, 132(7):077301, 2024.
William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space mod- cls. In Forty-first International Conference Machine Learning, 2024. URL https:
7 //oppnneview.net////.
Tsendsuren Munkhdalai and Hong Yu. Neural semantic encoders. In Proceedings the conference. Association for Computational Linguistics. Meeting, volume 1, pp. 397. NIH Public Access, 2017.
Tsendsuren Munkhdalai, Alessandro Sordoni, Tong Wang, and Adam Trischler. Metalearned neural memory, Advances Neural Information Processing Systems, 32, 2019.
Bo Peng. Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, al. Rwkv: Reinventing for the transformer era. arXiv preprint 13448, 2023.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemystaw Kazienko, et al. Eagle and finch: Rwky with matrix-valued states and dynamic recurrence. arXiv preprint arXiv:2404 05892, 2024.
11

---

Published as a conference paper at ICLR 2026
Bo Peng, Ruichong Zhang, Daniel Goldstein, Eric Alcaide, Haowen Hou, Janna Lu, William Merrill, Guangyu Song, Kaifeng Tan, Saiteja Utpala, et al. Rwkv-7" goose" with expressive dynamic state evolution. arXiv preprint vxXv:2553. 2025.
Hao Peng, Jungo Kasai, Nikolaos Pappas, Dani Yogatama, Zhaofeng Wu, Lingpeng Kong, Roy Schwartz, and Noah A. Smith. ABC: Attention with bounded memory control. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Anmual Meet- ing of the Association for Computational Linguisties (Volume 1: Long Papers), pp. 7469-7483, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 10.18653/112022. acl-long.515. URL htttt:aa//aaaaaattllllo///////// wwtp:///oooooooo222222222 aww-llanl...... 515/
wrg/
DL Prados and SC Kak. Neural network capacity using delta rule. Electronics Letters, 25(3): 197-199, 1989.
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong, Various lengths, constant speed: Efficient language modeling with lightning attention. In Forty-first International Conference on Machine Learning, 2024.
Zihan Qiu, Zekun Wang, Bo Zheng, Zeyu Huang, Kaiyue Wen, Songlin Yang. Rui Men, Le Yu, Fei Huang. Suozhi Huang, et al. Gated attention for large language models: Non-linearity, sparsity, and attention-sink-fiee arXiv preprint arXiv:2505.06708 2025.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gru- ber, Markus Holzleitner, Thomas Adler, David Kreil, Michael K Kopp, Günter Klambauer, Jo- hannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In International Conference on Learning Representations, 2021. URL/ttps: http:////pppeeei.../mmomm.mmm... net forum?
1d=t LBGtLL89R2IICC.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Jürgen Schmidhuber. Learning to control fast- memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
Jürgen Schmidhuber. Reducing the ratio between learning complexity and number of time varying variables in fully recurrent nets. In ICANN'93: Proceedings of the International Conference on Artificial Neural Networks Amsterdam, The Netherlands 13-16 September 1993 3, pp. Springer, 1993.
Mark Schöne, Babak Rahmani, Heiner Kremer, Fabian Falck, Hitesh Ballani, and Jannes Gladrow. Implicit language models are mns:: Balancing parallelization and expressivity. arXiv preprint arXiv:2502.07887, 2025.
Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Increasing the expressivity of deltanet through products of householders. arXiv preprint axXiv:2502.10297, 2025.
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for se- quence modeling. In The Eleventh International Conference on Learning Representations, 2023.
URL http:::///peeeevieeee net wwwpeeewwwwweeeeeeoneeeeteeeeeee Bww.:eiiAAAAiWwAAAAe.. 3AXqks
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): Rons with expressive hidden states. arXiv preprint arXiv:2407.04620 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuging Xia, Jilong Xue, Jianyong Wang. and Furu Wei. Retentive network: A successor transformer for large language models. arXiv preprint arXiv:2307.08622) 2023.
Matteo Tiezzi, Michele Casoni, Alessandro Betti, Tommaso Guidi, Marco Gori, and Stefano Melacci. On the resurgence recurrent models for long sequences: Survey and research op- portunities in the transformer era. arXiv preprint arXiv:2402.088 033,, 2024.
12

---

Published as a conference paper at ICLR 2026
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Ke Alexander Wang, Jiaxin Shi, and Emily B. Fox. Test-time regression: unifying framework for designing sequence models with associative memory, 2025. URL httth:t//axxi>>>>> org/
abs/2501.1239 12352
Songlin Yang. Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. In International Conference Machine Learning, pp. 56501-56523, PMLR, 2024a.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- with the delta rule over sequence length. Advances in neural information processing systems, 37:115491-11522 2024b.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024c.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule, arXiv preprint arXiv:2412.06464 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.
Yu Zhang, Songlin Yang, Rui-Jie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, al. Gated slot attention for efficient linear-time sequence modeling. Advances in Neural Information Processing Systems, 37:1 16870-1 16898, 2024.
13

---

Published as a conference paper at ICLR 2026
A RELATED WORK
Modern Linear Recurrent Neural Networks Due to the quadratic complexity of transformers, recently developing alternative architectures have gained attention, which led to the development of efficient recurrent alternatives (Tiezzi al. 2024). Initial advancements in this domain, starts with models such as RetNet (Sun et al. 2023), RWKV (Peng et al. 2023). and S5 (Smith et al.) 2023), which employed data-independent transition matrices coupled with Hebbian-like update mechanisms. Subsequently, second generation of models emerged, incorporating input-dependent parameters within these linear architectures (e.g., linear RNNs (Hasani et al. 2023 Smith et al.) 2023). RWKV6 (Peng al.) 2024)). These models also explored more expressive memory updat- ing rules, notably those based on the delta rule (Peng et al., 2025: Schlag et al. 2021 Yang et al.) 2024b Liu al. 2024). Further evolution in this line research has extended these memory architectures to deeper models, while concurrently utilizing delta-rule-like update mechanisms (Sun al. 2024) or data-dependent momentum- based update rules with forget gating (Behrouz et al., 2025d). More recently, to augment the performance of delta-rulc-base sequential models, (Siems al. (2025) have proposed the application of multiple gradient descent updates per token, thereby yielding more expressive sequence models, particularly in state tracking tasks. In addition to the above fast linear recurrent sequence models, several studies have focused on RNNs with non-linear recurrence (Behrouz et al. 2025d b a Csordás et al. 2024 Merrill et al., 2024 Lim et al. [2024]
Schöne et 2005 Karami& & Mirrokni | Gonzalez et al. 2024), and how their training can be
faster (Gonzalez al. 2024 Lim et al. 2024 Schöne et al. 2025).
Fast Weight Programs The conceptualization of linear layers as key-value associative memory systems can be traced back to Hopfield networks (Hopfield 1982). This concept was subsequently developed in the context of fast weight programmers, wherein dynamic fast programs are integrated into recurrent neural networks to serve writable memory stores (Schlag et al. [222] (chhiddhuber 1992 1993). Among the learning paradigms for such systems, Hebbian learning (Hebb) (2005) and the delta rule (Prados & Kak 1989) have emerged as the most prominent. Both learning rules have been the subject of extensive investigation within the existing literature (Munkhdalai & Yu
2017 [Schmidhuber 1992 Munkhdalai et al. 2019 Schlag et
2021|
Irie et al.) 2021] Yang al.,
2024b a).
Hopfield Networks Our formulation is architecturally founded upon the broad concept of associa- tive memory. wherein the primary objective is to learn an underlying mapping between keys and values. Seminal work by HHopiied (1982) on Hopfield Networks introduced one of the earliest neu- ral architectures explicitly based on associative memory, defining it through the minimization of an energy function for storing key-value pairs. Although traditional Hopfield networks have seen di- minished applicability in recent years, primarily due to constraints in vector-valued memory capac- ity and the nature of their energy function, several contemporary studies have focused on enhancing their capacity through various methodologies. These include efforts by Krotov (2021), Li et al. (202),, and Krotov& & Hopfield (2016). Notably, extensions to the energy function of these models, often incorporating exponential kernels, have been explored (Krotov & Hopfield 2016 Lucibello & Mézard 2024). Furthermore, the relationship between these modernized Hopfield networks and Transformer architectures has been subject of recent investigation (Ramsauer al. 2021 Hu al.
2024).
B CONNECTION OF QK-PROJECTION WITH MEMORY BOUNDED TRANSFORMERS
Revisiting our QK-Projection retrieval process, we first project the query vector into the space of stored keys and then retrieve its corresponding value by forward pass on the memory module. In particular, given query ge and stored keys of { the output corresponds to query ge can be calcualted as:
k-k |
0: =
WT,
96
(8)
14

---

Published as a conference paper at ICLR 2026
Given normalization of keys, i.e. | 1, this formulation, can be re-writtnn
=
Wv.
k k
in which the second element, Σx= Ki Ki Gt, equivalent a simple forward pass for query qi over a linear memory module of M, Ei=1 ki with recurrence of:
M; = M + kick
(10)
Such formulation of OK-Projection can remind the two-pass process of memory bounded Trans- formers (Peng et al., 2022 Zhang et al.] 2024; Karami et al. (025). where in the simple linear attention form (Peng et 2022). the retrieval process can be written as:
W. W spe
(11)
(E kippT
0: W, softmax
(12)
Comparing with above two-pass process of ABC (Peng et al. 2022). our QK -Kopojecciion method is applicable both deep and linear memory. Furthermore, parameters of Se as well h tied and so the learning process is considerably easier, making the model more adaptable to new data/tasks. Moreover, when employ this projection the local memory, we only do the summations starting from the "reser" state of the memory rather than starting from 1.
C EFFICIENT IMPLEMENTATION OF QK OK-PROJECTION
This section details the efficient, parallelizable implementation of the QK-Projection mechanism. We demonstrate that this projection can be integrated into the chunkwise training paradigm without introducing sequential bottlenecks, thereby preserving the training efficiency of the TNT architec- ture.
The QK Projection relies on projection matrix, Mr, which accumulates the outer products of keys within the current local memory shard (length St). Assuming normalized keys K| 1), this matrix is defined by the following recurrence:
Kirk Me MM- + K.AT
ift= (mod SL)
otherwise
This rule ensures that the projection state M, is reset at the beginning of each shard, mirroring the reset of the local memory state W, The local memory retrieval then computed as Miq:).
Chunkwise Parallel Computation. To maintain training efficiency, M₁ is computed in a chunk- parallel manner. For any time step & within chunk of size CL, the projection matrix can be decom- posed into two components:
M, M
Σ k
((t.C.)
Carry-wver State
Intru- chunkchuukkcco Sum
The first term is the final state from the previous chunk, which carried over. The second term, the intra-chunk sum, is computed efficiently for all steps in the chunk simultancously using a parallel prefix sum (scan) operation over the sequence of outer products (k, }.
This implementation preserves end-to-end parallelism. The state passed between chunks is sin- gle, constant-size matrix (d x d). incurring minimal overhead. The periodic reset handled by re-initializing this carry-over state shard boundaries. Thus, QK-Projection enhances the model's retrieval mechanism without compromising the training efficiency fundamental to the TNT archi- tecture.
15

---

Published as a conference paper at ICLR 2026
D TNT APPLICABILITY
In this paper, we have focused on showing the effectiveness and efficiency of TNT and so for the sake of clarity, we use simple memory module that optimizes its inner-loop with gradient descent. However, TNT recipes are applicable to different deep memory and non-linear architectures. For example, one can adapt the gating formulation Titans (Behrouz et al.) 2025d) or Mamba2 (Dao & Gu 2024) for each of the local memories as well as the global memory. Another potential ex- ploration is to incorporate closed feedback loop in the objective of the inner-loop employ more expressive optimizers as the inner-loop optimizers such as gradient descent with momentum, Adam W (Kingma & Ba 2014). or muon (Jordan et al,, 2024 as it has been done by Behrruzz et al. (2025a); Zhang et al. (2025). While exploring all these combinations with TNT is a promising direction, for the sake of clarity and space constraint, we leave them for future studies.
has done in
Hu et al. (2025). Similarly,
TNT GENERALIZATION FORMULATION
The TNT architecture can be generalized hierarchical system comprising one global memory and N parallel local memories. This allows the model to capture information at multiple timescales and resolutions simultancously. Each local memory, denoted by w/() for E {I,...,N] N). operates with its distinet chunk size Cr. shard length St.y, and learnable initial state W((
mni
GENERALIZED MEMORY UPDATE
The update rules extended follows: the single global memory evolves sequentially, while the N local memories are updated in parallel, each with its independent schedule.
Global Memory Update. The global memory state V is updated sequentially with large chunk size Cc, identical to the base case.
(k+1)CC
V(k+1JCa + Vkoa
Σ t=kCo
([V
(13)
N-Local Memories Update. Each of the N local memories W( operates in parallel. The state of each memory reset periodically according to its specific shard length St., enabling multi- resolution context parallelism.
W
10= (mod St,)
init
W
(
-E d7 Wooh (S(W't) E(I.C2,)) kt),
Otherwise
(14)
where 1, N.
E.2
GENERALIZED MEMORY RETRIEVAL
During retrieval, the final output is a composition the outputs from the global memory and all N local memories. The global memory uses the raw query qr, while each local memory uses Q-K projection tailored to its specific context window, determined by its chunk size Cz,-
TNT Generalized Retrieval Rule. butions from each memory module.
The hierarchical memory is retrieved by summing the contri-
T
0, = S(V q:)
W
4t |
(15)
This formulation allows the network to integrate long-range dependencies from the global mem- ory with fine-graine ane-grained, prraesel-processedd information from diverse set of local memories, each specializing in different temporal patterns.
16

---

Published as a conference paper at ICLR 2026
10
Runtime (ms)
10a
2048
4096
8192
16384
3276B
Sequence Length
FlashArtention
Titans (C=8)
Titans (C=32)
TNT (Gz=16)
TNT (Ok=64)
JAX Attention
Titans (C-16)
TNT (Cz-0) TNT (C2-32)
TNT (CL-128)
Figure 5: Runtime comparison of different models and implementations across varying sequence lengths, with the number of tokens per batch fixed at 0.5M.
Table 4: Training time for 150M parameter models trained on 10B tokens. For TNT models, the global chunksize is fixed at CG = 2048, and C₁ denotes the set chunksizes for the local memory modules.
Model
Co Cz
Training Time (hrs)
150M params/ / 10B tokens
Transformer (wfo gating)
0.80
Trans former (w gating)
0.82
TTT 2024
256
1.69
Titans
256
1.99
Titans
8.44
TNT Stuge
TNT Stage
(8)
3.06
(8.16)
4.24
[4,8.16)
5.00
14.8.16.32)
5.55
TNT Stage
TNT Stage
(1)
0.15
[2.4)
0.23
224,8)
0.26
(24488166)
0.46
Table 5: The scaling behavior of local memory depth (N) model performance. The depth is varied from = 1 to
9. A clear improvement in performance is observed as the depth
increases.
Model
Co
C. Cz.
C4
FineWeb ppl
PG19 ppl.
PIQA acc
Hella. acc
ARC- Bec
CSQA
Avg.
ppl
ppl
acc
acc
150M purams/ 10B
TNT Stage
2048
[8]
21.04
21.01
30.24
24.10
61.8
32.8
37.4
30.3
40.6
[8,16]
37..4
2048
(48816)
20.47
20.43
29.43
23.44
62.9
32.4
36.4
28.9
40.2
1048
(4.3.66332
20.15
20.17
29.08
23.13
63.2
32.0
36.7
30.3
40.6
14..66.32.64)
36.8
148.16.32,44128
20:08
29.66
64.0
37.3
2048
(48.16332.111222266
19.96
20.10
29.56
23.21
63.2
32.6
37.5
32.6
41.5
2048
(48.16..2.64122.2666373
19.84
19.98
29.27
23.03
62.7
32.5
36.5
32.5
41.0
2048
(48.16.222611128.66552..00
19.74
19.92
29.24
22.97
63.7
32.9
37.0
32.9
41.6
17

---

Published as a conference paper at ICLR 2026
Table 6: Impact of global memory chunk size the TNT model's performance. This is evaluated at four fixed local memory sizes (CL = {4, 8, 16, 32}). No clear correlation was observed between
global chunk size and performance.
Model
C
C4
FineWeb
PG19
Avg. ppl 4
PIQA acc
Hella. aCc 1
ARC-e acc
CSQA
Avg.
ppl
ppl
p
ace
acc f
150M params
10B tokens
TNT Stage
32
(4,8,16,32)
20.28
20.31
29.26
23.28
63.1
32.8
37.6
29.6
40.8
(4.8,6.33)
20.27
20.29
29.16
23.24
63.4
33.2
30.3
41.1
256
(6,8,16,32)
20.22
20.25
29.17
23.21
62.9
32.9
38.4
30.0
41.0
512
(4.8,16.32)
20.26
20.28
29.26
23.27
64.1
32.8
36.8
29.6
40.8
10124
[4.8.16,32)
20.25
20.31
29.34
23.30
64.0
32.4
36.8
29.6
40.7
2048
[4,8,16,32
20.20
20.23
29.26
23.23
63.4
32.7
37.0
29.6
40.7
4096
(4.8.16.32
20.26
20.30
29.25
23.27
63.2
32.4
37.9
29.9
40.8
16,8,
20.29
20.30
29.24
23.28
32.1
37.3
40.3
Table 7:
Impact of local memory chunk size the TNT model's performance. We use fixed
global memory (Cg = 2048) and vary the chunksize size CL single local memory. Smaller
local memory sizes are shown to yield improved performance.
Model
Ca
C' or Cr.
FineWeb
PG19
Avg.
PIQA ace
Hella. acc 1
ARC-e acc f
CSQA acc
Avg. acc f
ppl
ppl J
ppl
ppl
150M params/ 10B tokens
TNT Stage
2048
(2)
20.86
20.90
30.16
23.97
63.3
32.4
36,2
28.3
40.1
2048
(4)
21.02
21.04
24.15
32.5
37.5
29.9
41.0
(8)
21.04
21.01
30.24
24.10
61.8
32.8
37.4
40.6
2048
(16)
21.08
21.03
30.28
24.13
63.3
31.5
36.4
30.3
40.4
2048
(32)
21.06
21.04
30.24
24.11
62.8
32.0
36.4
27.6
39.7
2048
(64)
21.17
21.15
30.48
24.26
62.8
32.8
36.5
29.3
40.4
2048
(128)
21.39
21.39
30.80
24.53
63.5
31.8
35.6
28.7
39.9
2048
(256)
21.53
21.57
31.12
24.74
62.6
31.9
38.0
30.1
40.6
2048
[512)
21.74
21.81
31.45
25.00
62.3
31.4
36.1
28.0
39.4
2048
(1024)
22.18
22.24
32.21
25.54
61.9
30.9
36.6
28.9
39.6
2048
(2048)
22.79
22.83
33.06
26.23
62.2
31.5
35.8
26.7
39.1
Table 8: This analysis investigates different local memory configurations. We observe that the best performance is achieved with an optimal configuration of heterogeneous local memories (i.e.. having different resolutions), which aligns with the core hypothesis of TNT.
Model
Co
C or Cr,
C4
FineWeb
PG19
Avg. acc
PIQA
Hella. acc
ARC-e acc
CSQA acc
Avg.
ppl
ppl
ppl
ppl
acc
150M
10B tokens
TNT Stage
2048
14,8,16)
20.47
20.43
29.43
23.44
62.9
32.4
36.4
28.9
40.2
2048
(4,8.32)
20.36
20.33
29.37
23.35
62.6
32.6
36.9
28.2
40.1
2048
(4,8664)
20.38
20.39
29.45
23.40
63.1
32.2
36.8
29.4
40.4
2048
(4,8,128)
20.46
23.49
39.3
2048
(4,16.64)
20.37
20.36
29.35
23.36
62.7
32.9
37.3
29.5
40.6
2048
[4,16,128)
20.43
20.44
29.44
23.44
63.5
32,5
38.1
29.7
41.0
2048
(4,16,256)
20.54
20.56
29.67
23.59
63.2
32.5
36.2
28.1
40.0
2048
(4,16,512)
20.60
20.62
29.79
23.67
32.0
39.5
F ADDITIONAL. EXPERIMENT RESULTS
In this section, we provide extended ablation studies on the memory hyperparameters of the TNT framework, including local memory depth and chunk sizes, as well as comparisons against additional modern RNN baselines.
Scaling of Local Memory Depth (N). To evaluate the benefits of deeper local modules, perimented with increasing the local memory depth from N = 1 to N 9. As shown in Table 55 performance steadily improves with depth: average perplexity decreases from 24.10 22.97, and
18

---

Published as a conference paper at ICLR 2026
downstream accuracy increases from 40.6 41.6. This demonstrates that deeper local memory continuously improves performance, highlighting the scalability of the TNT framework.
Impact of Global Memory Chunk Size (Ca). We evaluated the effect of varying the global chunk size Ca from 32 to 8192 (Table 6). We observe that while a smaller Ca 128 yields slightly better average perplexity (23.21) compared to our default setting of Cg 2048 (23.23) the improvement is marginal. This minor gain does not justify the additional computational cost, as smaller chunks reduce opportunities for parallclization. Therefore, our original configuration provides a strong balance between performance and computational efficiency.
Impact of Local Memory Chunk Size (CL). Fixing the global memory (CG = 2048), we also varied the local memory chunk size Cz. As shown in Table our results indicate that smaller local memory chunk sizes generally yield improved performance.
Optimal Local Memory Configurations. To determine the optimal chunk sizes for multi-local setups, we experimented with different memory hierarchies (Table We find that the optimal con- figuration utilizes heterogeneous local memories with varying resolutions. As a practical heuristic, we suggest starting with an exponential progression (e.g.. 2 2 .... 2N 4 4 4 for N local memories). This rule of thumb generally yields strong and stable performance across different setups.
G LLM USAGE
We acknowledge the use of large language model (LL.M) solely for improving the linguistic quality and clarity of this manuscript. The model was not used for ideation, research methodology, or generating the scientific content presented in this work.
19


Reviews:
[{"review_id": "fJ7iHcaBYU", "reviewer": "Reviewer_g52o", "summary": "This paper, proposes TNT, a two-stage training paradigm for deep memory modules in RNNs to address the conflict between training efficiency and inference performance. Stage 1 uses a hierarchical memory (global for long-range context, local with periodic resets for parallelism) and Q-K Projection to boost efficiency. Stage 2 fine-tunes local modules for small chunks. Experiments show up to 17.37× faster training and better accuracy on Titans. Its contributions include identifying three challenges, introducing Q-K Projection, hierarchical memory with resets, efficient fine-tuning, and the TNT paradigm.", "questions": "1. Could you clarify if adding more than 4 local modules brings further performance gains or overhead, and is there a method to select optimal chunk sizes for multi-local setups?\n2. Have you done preliminary tests on custom kernels for TNT, and would TNT’s speed advantage be more pronounced if Transformer baselines use only JAX (no custom kernels)? \n3. Have you validated TNT on sequences longer than 32K (e.g., 64K, 128K) to check linear scaling, and why not test it on complex long-sequence tasks like document summarization?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper shows strong originality by proposing the two-stage TNT paradigm, decoupling deep memory modules’ training efficiency and inference performance to break the balancing bottleneck in existing work. Its hierarchical memory (with locally periodic resets) and Q-K Projection solve non-linear module parallelization and memory compression-retrieval mismatch, as targeted innovations.\n2. It features rigorous technicality and thorough experiments, defining core mechanisms via clear formulas for reproducibility. Covering multiple baselines and scenarios, with ablation studies verifying key components, its 17.37× training speedup and performance gains are supported by detailed tables/figures, ensuring high reliability. \n3. It has research and application value: breaking deep memory modules’ scalability bottleneck to enable RNNs as Transformer alternatives, identifying three core challenges. Its linear runtime and low fine-tuning overhead fit long-sequence tasks, and it applies to various deep memory modules", "weaknesses": "1. The paper only tests up to 4 local memory modules and does not analyze performance saturation points or optimal chunk size selection for multi-local configurations, leaving gaps in guiding practical hierarchical memory setup.\n2. TNT lacks custom kernel optimization, and the speed comparison with optimized Transformer baselines (e.g., Gated Transformer with FlashAttention) is unfair due to hardware optimization mismatch, failing to highlight inherent efficiency advantages.\n3. The paper does not validate TNT on sequences longer than 32K or complex long-sequence tasks (e.g., document summarization), limiting demonstration of its practical utility.\n4. It fixes global memory chunk size ($C_G=2048$) without testing variations or analyzing the impact of ($C_G$)-($C_L$) ratio, leaving gaps in understanding optimal hierarchical memory coordination.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZQszUhksyS", "reviewer": "Reviewer_NKDU", "summary": "This paper introduces a two-stage training method for test-time memorization modules, which typically suffer from low hardware efficiency during training. In the first stage, a global memory module and several local memory modules are employed to model long-sequence dependencies while maintaining efficiency. The second stage further fine-tunes the memory modules to achieve finer context lengths for the local memories, thereby enhancing inference performance. Experimental results demonstrate the effectiveness and efficiency of the proposed TNT method, showing its potential to achieve performance comparable to standard Transformer models.", "questions": "See weaknesses", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well structured and easy to follow.\n- The proposed two-stage training method effectively balances performance and training efficiency.\n- Extensive experiments across different model architectures demonstrate the method’s effectiveness and robustness.", "weaknesses": "- There is no hyperparameter study on $C_G$. How does the global chunk size influence performance?\n- TNT only outperforms the FlashAttention method at a 32K sequence length. Is this due to an under-optimized kernel or the limitation of maintaining additional memory modules?\n- What are the sizes of the global and local memory modules? Since TNT introduces additional parameters for the global memory module, comparable parameter sizes should be used for Titans when comparing performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OJ8JO5giqd", "reviewer": "Reviewer_2mkd", "summary": "This paper presents TNT, a novel training paradigm designed to address the efficiency–accuracy trade-off in Recurrent Neural Networks (RNNs) with deep test-time memorization modules such as Titans and TTT. Existing parallelization approaches rely on a fixed chunk size, where larger chunks improve training speed but degrade performance, creating a fundamental bottleneck. TNT overcomes this limitation through a two-stage training process. In the first, efficiency-oriented pretraining stage, a hierarchical memory structure is employed: a global module processes large chunks to capture long-range context efficiently, while multiple parallel local modules handle fine-grained details. By periodically resetting local memory states, TNT removes sequential dependencies and enables large-scale context parallelization. In the second, fine-tuning stage, only the local memory modules are adapted to smaller chunk sizes, restoring high-resolution accuracy with minimal computational overhead. Experiments on Titans and TTT models demonstrate that TNT achieves up to 17× faster training while simultaneously improving accuracy, effectively removing a key scalability barrier and establishing RNNs as a promising alternative to Transformers.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The work could bring some advantages: 1) stage 1 in the method can increase training throughput by introducing a novel hierarchical memory architecture that enables unprecedented parallelism. By the clear framework overview figure and TNT Memory Compression Rule, authors provide clear explanation for their method. 2) stage 2 can bridge the gap between the large chunk sizes required for efficient training and the small chunk sizes that yield the best performance at inference. 3) By the experiments given by authors, the results are impressive that can support the main points in the paper.", "weaknesses": "1) The presentation of the paper still needs more improvement. For example, in the main experimental results, I can not understand the meaning of column in table 2.\n2) The datasets used by the paper are not clear. I am confused on this.\n3) The four baselines seems to be not enough to better support the efficacy of method.\n4) I think the table1 can show the effectiveness of method. But other results in paper can not show obviously better performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)