Test-Time Training Done Right
Authors: Tianyuan Zhang
Sai Bi
Yicong Hong
Kai Zhang
Fujun Luan
Songlin Yang
Kalyan Sunkavalli
William T. Freeman
Hao Tan
Venue: ICLR 2026 Poster (2026)
Keywords: Test-Time Training
Sequence Model
Long Context Model
Primary area: foundation or frontier models, including LLMs
TL;DR: Large chunk test-time training enables efficient training of large nonlinear RNN.

Abstract: Test-Time Training (TTT) models context dependencies by adapting part of the model's weights (often referred to as fast weights) at inference time. This adapted fast weight, similar to recurrent states in RNNs, stores temporary memories of past tokens in the current sequence. Existing TTT methods have struggled to demonstrate effectiveness in handling long-sequence data, due to their computational inefficiency on modern GPUs. The TTT layers in many of these approaches operate with extremely low FLOPs utilization (often below 5%) because they deliberately apply small online mini-batch sizes (e.g., updating fast weights every 16 or 64 tokens). Moreover, a small mini-batch implies fine-grained block-wise causal dependencies in the data, making them unsuitable for data beyond 1D ordered sequences, like sets or N-dimensional grids such as images or videos. In contrast, we pursue the opposite direction by proposing an extremely large chunk update, ranging from 2K to 1M tokens across tasks of varying modalities, which we refer to as Large Chunk Test-Time Training (LaCT). This approach improves hardware utilization by orders of magnitude, and more importantly, facilitates scaling of nonlinear state size (up to 40% of model parameter size), hence substantially improving state capacity, all without requiring cumbersome and error-prone custom kernel implementations. It also allows easy integration of sophisticated optimizers like Muon for online memory updates. We validate our approach across diverse data modalities and tasks, including novel view synthesis from image sets, language models, and auto-regressive video diffusion models. Our approach can scale up to 14-billion-parameter auto-regressive video diffusion models handling sequences of up to 56K tokens. In our longest sequence experiment, we perform novel view synthesis with more than one million context length. Our results highlight the computational and performance benefits of large-chunk test-time training, paving the way for more efficient and scalable long-context sequence modeling. We hope that this work will inspire and accelerate new research in the field of long-context modeling and test-time training. See visual results on project website \url{https://tianyuanzhang.com/projects/ttt-done-right/}.

Published as a conference paper at ICLR 2026
TEST-TIME TRAINING DONE RIGHT
Tianyuan Zhang'
Sai Bi
Yicong Hong²
Kai Zhang"
Fujun Luan"
Songlin Yang'
Kalyan Sunkavalli
William T. Freeman'
Hao Tan?
Massachusetts Institute of Technology
Adobe Research
ABSTRACT
Test-Time Training (TTT) models context dependencies by adapting part of the model's weights (often referred to fast weights) at inference time. This adapted fast weight, similar to recurrent states in RNNs, stores temporary memories of past tokens in the current sequence. Existing TTT methods have struggled to demon- strate effectiveness in handling oong-eeuuence data, due to their computational inefficiency on modern GPUs. The TTT layers many of these approaches operate with extremely low FLOPs utilization on modern GPUs(often below 5%) because they deliberately apply small online mini-batch sizes (e.g., updating fast weights every 16 or 64 tokens). Moreover, a small mini-batch implies fine- aine-grained block- wise causal dependencies in the data, making them unsuitable for data beyond 1D ordered sequences, like sets or N-dimensional grids such as images or videos. In contrast, we pursue the opposite direction by proposing an extremely large chunk update, ranging from 2K to IM tokens across tasks of varying modalities, which we refer to as Large Chunk Test-Time Training (LaCT). This approach improves hardware utilization by orders of magnitude, and more importantly, facilitates scaling of nonlinear state size (up to 40% of model parameter size), hence substan- tially improving state capacity, all without requiring cumbersome and error- prror-prone custom kernel implementations. It also allows easy integration of sophisticated optimizers like Muon for online memory updates. We validate our approach across diverse data modalities and tasks, including novel view synthesis from image sets, language models, and auto-regressive video diffusion models. Our approach can scale up to 14- 44-billion-p bbiliion-parameter parameter regrrsegressive video diffusion models handling sequences up to 56K tokens. In our longest sequence experiment, we perform novel view synthesis with more than one million context length. Our results high- light the computational and performance benefits of large-chunk test-time training, paving the way for more efficient and scalable long-context sequence modeling. We hope that this work will inspire and accelerate new research in the field long-context modeling and test-time training. See visual results on project website htwpaaaywaaaaaaannccooooooooeeee cow http:p/aggcccccpppeeeeeetttttttt dowett---one-righhh/
INTRODUCTION
The demand for handling long contexts is rapidly growing. While softmax attention (Vaswani al. 2017) has become the de facto solution for modeling various types of data, computational cost grows quadratically with sequence length, motivating extensive research into more efficient long-context modeling.
Recently, Test-Time Training (TTT) (Sun et al. 2024 has emerged promising approach for efficient sub-quadratic sequence modeling. TTT extends the concept of recurrent states in RNNs to small, anline-adapted sub-network. The parameters of this sub-network also referred to fast weight (Schlag et al. (021). as they are rapidly adapted online via self-supervised objectives to memorize in-context information. Numerous recent studies (Wang et al. 2025b Behrouz al. 2024; 2025b] Karami & Mirrokni] 2025) have explored various online objectives, optimizers, and architectures for fast weight networks.
Despite these efforts, existing TTT methods struggle to scale effectively to long contexts, primarily due extremely low hardware utilization in their TTT layers (often below 5% peak FLOPS

---

Published as a conference paper at ICLR 2026
modern GPUs). This inefficiency is because of the usage of small mini- sizes, i.e. updating fast weights every token or every 16 to 64 tokens, which is conventionally assumed to be more effective for in-context learning. Such small mini-batch results in poor parallelism and low compute intensity, and presents significant challenges for hardware-efficiees implementation, especially when using large, nonlinear fast weights, making it difficult achieve non-trivial (ahove 10%) FLOPS utilization.
In this paper, we adopt the opposite strategy and introduce Large Chunk Test-Time Training (LaCT). LaCT leverages extremely large chunk (from 2048 1M tokens) the basic unit update the fast weight. Since the tokens within each large chunk are treated as an unordered set, we further integrate window attention into LaCT to capture local dependencies within the chunk. LaCT significantly enhances parallelism, leading to substantially improved GPU utilization (up to 70% on NVIDIA A A100s) with just a few dozen lines of pure Py code (see Appendix C.T). This efficiency enables the scaling of non-linear fast weights to enhance the memory capacity. And simple implementation allows easy integration of more effective test-time optimizers, such as Muon (Jordan et al.) 2024).
Furthermore, LaCT's large-chunk design is also natural to model diverse N-dimensional data we can align chunk-size with the internal structure of the data (e.g., grouping tokens within an image or consecutive video frames as a chunk).
We extensively validate LaCT on three tasks spanning different modalities and data structures:
Novel View Synthesis. Our model is capable of processing up to 128 input images a resolution of 960 x 536 leading to maximum of 1M tokens, and outperforms 3D Gaussian Splatting (Kerbl et al., 2023) in terms of rendering quality under such input scale.
Language Modeling. Our model achieves competitive performance compared SoTA methods such as DeltaNet (Yang et al. 2024b), even though a chunk structure is not explicitly present in language data.
Autoregressive Video Diffusion. We adapt 14-billion- parameter bidirectional video diffu- sion transformer into an autoregressive model by incorporating LaCT with sliding window attention. This adapted model generates consistent videos up 56,000 visual tokens.
4-billlon-pprammeter
To summarize, approach establishes an efficient, scalable, and highly performant framework for long sequence modeling across diverse modalities. By removing the dependency on low-level, hardware-specifie implementations, LaCT enables broader exploration of the architectural design space. We believe this can democratize research in efficient long-context modeling and inspire the development of more novel and effective designs.
RELATED WORK
Test-time training. Test-Time Training (TTT) (Sun al. 2024) is an emerging concept in sequence modeling that extends the concept of recurrent states in RNNs to online-adapted neural network components. In TTT models, a subset of weights, termed "fast weights," " are updated learn in- context. Existing methods typically employ self-supervised loss that encourages these fast weights to memorize key-value associations from in-context tokens, using variants of gradient descent for online adaptation. TTT (Sun et al. 2024 Wang et al. 2025b) has opened vast design space for new recurrent model architectures. For instance, many recent works have developed novel test-time optimizers (Bchrouz et al., 2024 Karami& & Mirrokni 2025 Behrouz et al. 2025a) and online training objectives (Behrouz et al. 2025b). However, current nonlinear TTT approaches often suffer from low hardware utilization and limited state sizes, and consequently have not yet demonstrated their full potential. Our work primarily addresses these challenges by advocating for a new paradigm of using extremely large online minibatch (chunk) sizes for updating the fast weights. This paradigm can achieve orders-of-man -faagniud- higher hardware utilization without relying on error-prone custom kernel implementations. Furthermore. enables efficient scaling of nonlinear state sizes and offers the flexibility to use diverse fast weight neural networks and optimizers, thereby accelerating research progress in this area.
Combining chunk attention with recurrence. Several recent models combine local chunk at- tention with linear recurrence, such as Gated Attention Unit (GAU) (Hua et al., 2022b), MEGA (Ma et al. 2023), MEGALODON (Ma et al. 2024), and InfiniAttention (Munkhdalai et al. 2024.
2

---

Published as a conference paper at ICLR 2026
Among these, InfiniAttentionn is conceptually closest to work, as it incorporates recurrence the chunk level using the delta interpreet as an online linear regression objective from the perspective of Test-Time Training (TTT). However, this update rule is limited in expressivity. In contrast, we employ a significantly more expressive update mechanism derived from a more general TTT framework, and demonstrate the substantial gains this brings.
Block-Recurrent Transformer (Hutchins et al. 2022) also explores large chunk memory updates, where memory tokens act as recurrent states that can self-attend and cross-attend with input tokens during each chunk update via attention mechanisms. The Perceiver-style register-token attention baseline used in our novel view synthesis experiments (Sec.66.1 Table 2 is conceptually similar to the
Transformer in its use of register tokens for context compression. As shown in Figure 4 our method significantly outperforms this approach in both speed and quality, with a comparable state size.
Further related work on novel view synthesis and video generation is discussed in Appendix
PRELIMINARY
3.1 TEST- TIME TRAINING
Consider one-dimensional sequence of N tokens [x], 22, Tx., where each token xy Rd. Following attention formulation, each input tokens x; is projected into query (q:), key (k:), and value (04) vectors. For clarity, we assume all these vectors q:: k2 U/ E Rd.
Test-Time Training (TTT) (Sun et al. (024] introduces a neural network with rapidly adaptable caliss-clllldeed fast weights (Hinton & Plaut 1987 Schmidhuber 1992, [chlag et al. 202T] that are updated during both training and inference to dynamically store context information. This contrasts with the slow weights (i.e., model parameters) that are frozen during inference. Formally, TTT defines fast weights in the form of neural network: fw(-): Rd + Rd parameterized by the fast weights W, and it involves two primary operations:
Update operation: W + W -V-vVV VV--VWCC(Jv (k).
where -)) is a loss function between the transformed key fw (k) and the value v, commonly Mean Squared Error, designed encourage the network to associate keys with corresponding values. 1) is the learning rate. Intuitively, this learning objective is to encode the KV cache into neural memory with fixed state size as accurate as possible (Wang al. (025b).
Apply operation:
(W)
where the updated fast weights W are used to compute the output vector given the query 9. The per-token TTT layer iteratively perform the update and apply operations on each token m; in sequence.
3.2 CHALLENGES IN EFFICIENT IMPLEMENTATION
Frequent online update of fast weights inefficient due to memory bandwidth limitations. Conse- quently, previous works (Sun al. 2023 Gu & Dao 2023 Yang et al., 2024at Qin et al. 2024 Yang et al. 2024c) often employ customized kernels that keep fast weights in SRAM across updates reduce memory load. However, this strategy typically requires fast weights
evolve mostly independently within SMs to reduce communications. which is not valid for large nonlinear states (e.g., the nonlinear SwiGLU fast weight in Sec. 4.1 and the Muon update in Sec. 4.2). Moreover, developing such kernel code is cumbersome, with far longer development cycles than native Py PyTorch code, hindering rapid research exploration.
On the other hand, a PyTorch-basee implementation, while simpler, is typically bounded by memory speed, As an illustration, consider PyTorch implementation of simple MLP fast weight, the core of which is a matrix multiplication between fast weight (e.g., h x h matrix) and the mini-batch input (b x where b the chunk size). The ideal compute-to-memory ratio is:
2h2b
h/2
b
T
24² + 4hb
+
h
+
25
< min (h/2, b).
(3)
25
3

---

Published as a conference paper at ICLR 2026
1000
H108
107
2:39
505
515
Thkees secord
223
Thioage tan (Tr.g/oo.TT...... (P.CPP
V9 ton ronden .0%
TTTP
Mamba
0U15
220
203
2.33
state Size (MS, log sCa
Maxtaac lime 1A300 tog S0d4:
Onglasl
Mambe
Burs
Orighnat MLP-MLP
0.234
1he
Mambez
Oure hontnnar
(a) GPU Throughput
Efficient state scaling
(c) Effectivee state scaling
(d) Training Efficiency
Figure 1: Using larger chunk sizes significantly improves GPU utilization compared the original test-time training (TTT) method that even uses customized kernels (a). This enhanced utilization enables efficient and effective scaling to larger state sizes (b). (c), leading to better overall performance in less wall-clock time (d). The dotted line in (a) is the theoretical peak BFI6 throughput of the GPU. Panel (c) measure average validation loss of the last 2K tokens in sequences processed by a LaCT language model across varying state sizes, demonstrating benefits of larger state size. Panel (d) compares performance versus training time across different baselines on the novel view synthesis benchmark. Further experimental details can be found in Appendix E.4
Here, 2 is the FLOPs to for matrix multiplication, the denominator 2h2 + 4hb the memory workload for two input matrices and the output in BF16 (2 bytes). Small fast weight size (e.g.. = 64) or small chunk size (e.g.. b 16) will bound the ratio r far below the theoretical peak (e.g.. 290 FLOPs per byte on H100), making the operation memory-bound and limiting compute usage.
In light of this, we advocate for using large chunk sizes (from 2048 to 1M). This allows us to achieve higher throughput (Fig. [a) leading better performance in less training wall-clock time(Fig. Dt). Our design also allows the state size to be scaled up efficiently(Fig.. [p). leading to significant results improvement with such scaling (Fig. [.. Fig. (7)). Our architecture achieves a topaatto-arr size ratio 2 40%, which is an order of magnitude larger than previous methods' ratio of 0.1% to 5%. Detailed pseudocode is provided Alg.
Parallelism over the sequence length dimension, in addition to the batch and head dimensions, crucial achieve high occupancy when handling long sequences (where the batch size is often small). Linear Attention variants like Mamba (Gu & Dao 2023). Gated Linear Attention (Yang et al. 2024a) and DeltaNet (Yang et al. 2024c) enable such parallelism by utilizing the associative property of linear recurrence. Attention (Vaswani et al. (017)) can be parallelized along the sequence length dimension due to independent computation between different queries. Supporting such parallelism is key improvement of FlashAttention-2 (Dao 2023 over FlashAttention- 1 (Dao et al. (022). For test-time training with non-linear updates, sequence dimension parallelism can only be implemented within online chunks, further motivating the use of extremely large chunk sizes. When implementing large-chunk TTT with PyTorch, this sequence dimension parallelism within device across multiple thread blocks is automatically handled by PyTorch and low-level compilers. An example of such sequence parallelism across multiple devices is provided in Section with pseudocode in Alg.
LACT MODEL ARCHITECTURE
As in Fig. LaCT block consists of three types of layers: window attention layer, large-chunk TTT layer. and a feed-forward layer. Each layer is equipped with residual connections (He et al. 2015) following the practice Transformer (Vaswani af.) 2017). The window attention layer performs local self-attention to capture the local dependency. In the TTT layer, we split the sequence into large chunks. The history context is gradually compressed into the fast weights through an "update" operation (regarding vector key K and value V), and latest weight is 'applied' to the current query vector (Q) to compute the output. The feed-forward layer performs channel mixing as in Transformer. We omit several linear and normalization layers in Fig. for clarity and details are in Appendix Our framework offer great flexibility in handling diverse data types. In this section, present general designs in approach and later describe data-specific variations in Sec.
4

---

Published as a conference paper at ICLR 2026
V
Update
Window Attention
Split into large" chunks
Feed-
Forward
Apply
Causal direction (e.g.. time)
Large-Chunk TTT Layer
Figure 2: The basic diagram for a LaCT block. The large-chunk TTT layer updates the fast weight W store historical context information, while the window attention handles the locality and internal structures within the chunk. The solid line denotes the information flow over model depth and the dashed line denotes the information flow over time (i.e., the fast weight W passing through chunks). Various instantiations in Sec. Sec.5]un use different chunk sizes and window attention types according the specific data structure. Additionally, window attention and large-chunk TTT layers can combined within the same layer by sharing the QKV and summing their outputs; this in in-layer mixing is used in our language modeling and video generation experiments (see Alg. for such pseudocode).
4.1 LARGE-CHUNK TTT LAYER
Different from the per-token update in Eqn. the chunk-wiss update computes the gradient of the summed loss over all keys {k:} and values {u, within the chunk. As the chunk size is large, weight updates are performed infrequently. This enables more sophisticated weight-update rule designs (discussed in Sec. 4.2) and amortizes the update cost. The 'update' operation for the fast weight is:
9 - Vw Σ (fw())
(4)
W < weight-update(W..gg
(5)
where is the chunk size, g is the gradient of the fast-weight loss function, and TJ: is the learning rate of each token (usually predicted from input tokens). The 'apply' operation o; fw (q:) is the same as Eqn. 2 and all query vectors {q:} the chunk share the same updated fast weight W.
Motivated by recent LLMs (Touvron al. 2023), we adopt SwiGLU-MLP (Shazeer 2020 without bias terms as the weith network. Our fast weights consists of three weight matrix (W1, W2, Wa). }. and the network is:
(w) Wa (SiL.U(W1a) o (W3=)
where o is an elementwise multiplication. We apply simple dot product loss as our loss function:
((Jw(ks), v) = -fw(kk)
Execution orders for 'apply' and 'update'. Note that the 'update" operation and 'apply' operation of TTT are decoupled, and we can set the chunk size adaptively and apply these operation different orders: this allows us to model diverse kinds of data dependencies, similar to different attention masks in self-attention. Figure illlsttaats this concept. Figure when the chunk size equals the full sequence length, performing the apply followed by the update operation is conceptually similar to full attention. Using update and apply alternately leads to block-wise causal mask (Fig. (p), where the block size corresponds to the chunk size. Switching the order between the two operations results in the a shift the mask (Fig. Bb. This shifted mask does not leak future information within the chunk and is important when building the full causal mask in Language Modeling (Sec. (5.2). Moreover, only updating on a subset of chunks and applying to (Figure 3 is analogous to strided block-wise causal mask.
5

---

with weight Published as a conference paper at ICLR 2026
TTT Update- Apply Order
Apply
Apply
Apply
Apply
Apply
Apply
Applyy
Apply
Apply
W
Upores
W --o
Updale
Update
W
Update
Updsee
W--r
Updaate
Update
Equivalent
Attention
Mask
(a) Full Mask
(b) Block-Wise Causal
(c) Shifted Black-Wise Causal
(d) Strided Block-Wise Causal
Figure 3: Different "Update' and 'Apply' orders and their equivalent attention mask. A blue mask in i-th row and j-th column means the i-th token's 's output depends on the j-th token.
4.2 NON-LINEAR UPDATE OF FAST- WEIGHT
Fast-weight updates in TTT repeatedly accumulate gradients, and thus suffer from magnitude ex- plosion decayed memory. Large-chunk TTT allows non-linear updates improve stability and effectiveness while preserving efficiency. For the 'weight-update" operation in Eqn. our vanilla implementation involves gradient descent followed by weight normalization:
weight-update W. = T22-NormalizeeW - g)-
(8)
We have also explored a more robust nonlinear Muon (Jordan et al. 2024 update rule normalization:
weight-update((M (W, g) = 1.2-Normalize(WW (W- MMun(gg)
Fast-weight normalization. We apply L2 weight normalization the updated fast weights along the input dimension. We do not use explicit weight-decat term as in previous methods (Behrouz al. 2024 Dao & Gul [202)) Yang et al., [2244a] Sun et al. 2023). When the network conceptually rotated 90 degrees, treating the sequence dimension as the depth of a virtual model, the test-time training updates act as residuals over time (He et al. 2015). this view, our fast-weight normalization is analogous to the post-layer norm in Transformer architectures, which constrains activation scales within the residual path.
Muon-update rule. Essentially, Muon normalizes the spectral norm of matrix gradient using Newton- Schulz iterations. short, let y = USV be the Singular Value Decomposition(SVD) the gradient g. then Muon operator approximately converts the gradient as:
Muon(g) = UVT
(10)
Muon also improves the numerical stability in our setup. For example, the learning rate (7: in Eqn. now only reflects the relative importance of tokens within a chunk as Muon normalizes the absolute scale. See (Jordan et al. 2024) and Appendix for analysis of its computational cost.
4.3 WINDOW ATTENTION
The large-chunk TTT layer treats data sequences of sets because its fast weight updates inherently disregard token order and spatial locality within each chunk. However, many data modalities such videos (sequences of grids), image collections (sets of grids), text (1D sequences)-da not fully align with this set- bet-based perspective. For these modalities, intra-chunk structure and locality are vital for capturing the overall data structure. We therefore integrate local window attention (either causal or bidirectional) alongside TTT layers to handle data structure within a chunk. Moreover, window attention efficiently handles localities in the data, enabling the TTT layer focus its fixed-size fast weight capacity on modeling non-local dependencies. This hybrid strategy is also
'Muon requires weights in matrix form, and our current fast-weight function SwiGLU-MI.P has three matrices as the weights (i.c., Ws in Eqn.

---

Published as a conference paper at ICLR 2026
employed in other notable works like BASED (Arora al. 2024). GAU (Hua al. 2022a) and InifinitAttention (Munkhdalai et al. 2024). In summary, LaCT is a hybrid architecture with the quadratic-compp attention for local structure and linear-compute TTT for non-local context.
4.4 CONTEXT PARALLELISM
Context Parallelism (CP) partitions the sequence along the context length dimension and distributes the shards across multiple devices for parallel computing. The ferd-fowward layer and window attention are local operators thus natively support CP. For TTT layer, small chunks hardly support CP thus tensor parallelism (i.e., parallel over the heads) is preferred. Our large-chunk TTT layer allows CP by sharding the tokens within a chunk. This can be implemented through distributed all-reduce-see and is logically the same as Distributed Data Parallelism (DDP), except that the parameters are the fast weights and input data are the tokens in the chunk. We adopt such parallelism training the novel view synthesis task (Sec. (. and observe minimal throughput overheads (1% to 3%). See Appendix for pseudocode (Alg- and other parallelism(AIg
LACT FOR N-DIMENSIONAL DATA
In this section, we introduce the three tasks we address using LaCT-noves view synthesis, language modeling, and autoregressive video generation. These tasks have different inherent data structures and we address them with corresponding design choices. The full model architecture details for these data types are provided in Appendix D
5.1 NOVEL VIEW SYNTHESIS IMAGE SET
Novel view synthesis (MMillan & Bishop 1995 Levo & Hanrahan 1996) aims to render images of static scene from previously unseen viewpoints. Formally, given a set of N input posed images {T,, P.) of static scene, where I, E R "x x3 is an RGB image and P. is its corresponding camera pose, the model needs to synthesize new images from novel camera poses that typically do not overlap with the input views.
We find that NVS is an effective test bench for evaluating a model's online memory and compression capabilities. Firstly, NVS is challenging it requires spatial compression, dense retrieval, and basic physical reasoning. Secondly, NVS can formulated as non-generattiv task, significantly reducing training computation and the need for extensive model parameters to store world knowledge, thereby enabling rapid experimentation. Thirdly, the substantial redundant information in dense input views incentivizes the model to learn effective compressions. Given these observations, use NVS for our initial research iterations. We find that some of the insights gained are transferrable to other tasks.
Our NVS model follows the basic LaCT diagram in Sec. Both the posed input images and poses of the target novel views are tokenized by patchify and linear layers, following LVSM (Jin et al.) 2024a). The window attention exactly covers the tokens from a single image. The LaCT layer adapts a single-round of strided block-wise causal mask (Fig. 3d), which updates the fast weight using all input view tokens, and applies to both the input and novel view tokens. The update step resembles prefill stage, while the apply operation resembles parallel decoding. During rendering of novel views, each test-tima training layer functions static weight layer, making the entire model static vision transformer (Dosovitskiy et al. 2020). We illustrate this design in Figure
5.2 LANGUAGE MODELING TEXT SEQUENCE
Autoregressive language models predict the probability distribution of the next token given preceding tokens, Po (I., . -1) Text sequences lack inherent chunk structures, so for LaCT, we define chunk size as a hyperparameter (e.g.. 2048 or 4096 tokens). We utilize the shifted block-wise causal mask in Fig. (() for the TTT apply-update sequence to avoid seeing future tokens in chunk. Since LaCT lacks per-token causality within each chunk, we employ sliding window attention with window size equal to the chunk size-to to efficiently model per-token causal dependencies. The sliding window is integrated into the same TTT layer with shared QKV similar to GAU (Hua et al.) 2022a). We illustrate the detailed architecture in Fig. and pseudocode
7

---

Published as a conference paper at ICLR 2026
Table I: Summary of our experiments three different data structures, 'd' denotes model dimension.
The state size denotes the size of the fast weight per model block.
Task
Data Siructure
Chunk Size
State Size
Model Size
Max Length
Context Parailelism
Novel View Symthesis
Image set
Full sequence
6et
0.3B
IM
Within Within-chnk purallel
AR Video Diffusion
Image sequence
Three fromes
3d 0.75d
1,3B,
56160
Head-dim parallel
Language Models
ID Sequence
2K. 4K tokens
0.75d
0.7B, 3B
32768
N/A
5.3 AUTOREGRESSIVE VIDEO DIFFUSION IMAGE SEQUENCES
Chunkwise autoregressive video diffusion iteratively denoises a number of subsequent video frames, conditioned on the previously generated clean frames, where each chunk can contain thousands of visual tokens. We use feacher-forcing training by interleaving noisy and clean frame chunks. Specifically, a video of N frame chunks is structured as:
S
[ X1, noise X2,
moüse N
(11)
where each noisy chunk X. is produced by adding unit Gaussian noise to the i-th clean video chunk X X.(1 (4) + el, and 1, E [o, 1] denotes the strength of chunk-independent noise.
To handle such data structure, we employ the strided block-wise causal mask in Fig. 3d for LaCT. Specifically, it applies fast weights to each chunk sequentially while only updating fast weights on clean chunks. This simple strategy ensures that each denoising operation only accesses previously cleaned frames. The windowed attention uses non-overlapping window with 2 consecutive chunks (i.e., [X,, xxfft)) to build temporal and spatial locality. Within each window, the attention from X; to X X is excluded. We adopt all attention and TTT masking patterns similar to Fig. 3 The details of this hybrid architecture and efficient trainings are in the Appendix D.3 D.3
6 EXPERIMENTS
In this section. we present our experiment results on novel view synthesis (Sec. 6.1 language modeling 6.2), and autoregressive video generation (Sec. 6.3). and an in-depth analysis (Sec. A of different design choices. Tab. [ summarizes key factors in each experiment. When comparing with linear-cost baselines, we augmented them with the same window attention for fair comparisons. The full experimental details for all tasks are provided Appendix
6.1 NOVEL VIEW SYNTHESIS
Datasets & metric. We evaluate method on both object-level (Deitke al. 2023) and scene-level (Ling et al. 2024) datasets. For object-level experiments, models are trained on Objaverse and tested on the Google Scanned Objects (GSO) dataset (Downs et al.) 2022), 256 256 and 512 512 resolutions. Each evaluation uses 4-48 input views and 8 novel views per object. For scene-level experiments, adopt 140 test scenes from DL3DV, evaluated 960 536 resolution. Performance is reported in Peak Signal-to-Noise Ratio (PSNR), with additional metrics provided in Appendix E.]
Model & Training details. Our default model has 312M parameters, including 84M fast weights (6d2 per block). We train on 1.25T tokens for object-level datasets and 1.ST tokens for scene-level datasets, using a maximum sequence length of one million tokens. High-resolution models employ inner-chunk context parallelism (Sec. More details are provided in Appendix
Baselines. For object-level evaluation, we compare against two representative attention designs: one without in-context compression and one with compression: (I) the non non-compressed bascline is a full- attention model, where TTT layers are replaced by block-wise causal attention enabling bidirectional input interactions and cross-attention from novel views, and (2) the attention- ditention-with-come with-compresssion baseline is a perceiver-style register-attention model (Jaegle et al. 2021). which compress inputs to 4096 registers and decodes novel views via cross-altention. For scene-level evaluation, we benchmark against LongLRM (Ziwen et al. 2024), which integrates Mamba (Gu & Dao 2023) with full attention for 3D Gaussian splat prediction, well pure aptimization-basse 3D Gaussian splatting methods.
Table 2 summarizes computational complexity across all models. Figure 4 present experimental
results and analysis. See more results and details in Appendix. F.T
8

---

Published as a conference paper at ICLR 2026
Table 2: Complexities methods on novel view synthesis w/ n input. Prefill and rendering speed are measured on A A100 with 48 512x512 imput images (196K input tokens, 4K decoding tokens).
State Size
Prefill Compute
Decoding Compute
# Params
Prefill speed
Rendering FPS
Full attention
O
O(n")
O(n)
284M
16.1
2.3 FPS
Perceiver Attention
0(1)
O(n²)
0(1)
287M
16.8
34,4 FPS
Ours
0(1)
O(n)
0(1)
312M
1.4
38.7 FPS
GSO 256
GSO 512x512
DL3DV X
30
vieete
vieax
36
3 Viless
visss
Ours
Coussian
Attention
Res512
Attention
FullAtta Res512
30 444
2/X
20100k
40/40x
14/20K
12/250K Number of Views "otua oaans
G4/51KK
128/M
Number views / This Votens
Profil Timats) log scale
(a) Object dataset
(b) Performance Prefill speed
(c) Scene dataset with 1M tokens
Figure 4: (a, b) our method achieves quality comparable to full-attention models with significantly
lower prefill latency, and it clearly outperforms perceiver-attention baselines. (c) On the high resolution scene dataset, our approach surpasses LongL.M, limited to 32 views, and outperforms 3D
Gaussian Splatting with sparse views, remaining competitive up 128 input views (IM total tokens).
(a) 760M Validation loss
(b) 760M S-NIAH-1
(c) 3B Validation loss
(d) 3B S-NTAH-2
2.60
Srgeniode (4)
50
Strgle Moldl 283
800
Loss 1o0
E7.35
2.58
40
60
2.56
20
2:35
2.54
0x
16t
2k
4k
Token index)
Sequence Length
Token inee.
Sequente Lengr>
Transfomeer
Transformer SWA
GLA SWA
DettaNor SWA
Ours Momantum
Oura Muon
Figure 5: Language Model results. (a, e) Our model achieves lower per- per-position loss at larger token indices compared to GLA and DeltaNet at both 760M and 3B scale, indicating stronger long-context modeling capability. (b, d) Our model consistently outperforms GLA and DeltaNet in retrieval accuracy. Furthermore, our Muon variant consistently outperforms our Momentum variant.
6.2 LANGUAGE MODELING
Datasets Metrics. We train our models on the Long-Data-Collectin dataset (Al 2024). using approximately 60B tokens. For evaluation, we employ the per-token loss metric from (Xiong et al. 2023 Lin et al. 2025). assessing models' ability effectively use the full context. A monotonically decreasing loss indicates successful context utilization, whereas plateauing suggests limited context usage. Additionally, we report retrieval accuracy (Hsieh et al. 2024) at various sequence lengths.
layer directly into the Large-Chunk TTT layer by sharing Q, K, and V vectors with the weitht network, following GAU Hua et al. (2022a). The pseudocode for this design is in Algorithm 2 We trained models at two scales using a 32768-token sequence length: a 7600-pprammn window and 3B- 1--parameter model with 4096-token sliding window. See more details in App. E.2
Model & Training details. We integrate the sliding mondow-attentionte
model with a 2048-token sliding
We compare against transformer, Gated Linear Attention (GLA) Yangg et al. 2024a). DeltaNet (Schlag et al.) [202]] Yang al. 2024c). To ensure faimesss we enhance both GLA and DeltaNet with the same SWA. Tab. 6 summarize the mechanism and training throughput of all methods. Fig. 5 presents experimental results and analysis. See results and details in App. E.2
Baselines.

---

Published as a conference paper at ICLR 2026
(a) Comparison with baselines
(b) Ablate window size
(c) Eval longer videos
0.30
framess per window
10
Idatton oss
4 frames per window
Va idation loss
Va idation
05
0.00
0,05
10
Video chunk index
Video chunk Index
Video chunk index
Transformer
Transformer SWA
Mamba SWA
Ours
Figure 6: (a) We achieve comparable validation loss to the full-attention baseline and outperform both Mamba with sliding window and sliding window attention baselines. This improvement over SWA is consistent across different window sizes (b) and when evaluating longer videos (c).
6.3 AUTOREGRESSIVE VIDEO DIFFUSION
We fine-tune the pretrained Wan 2.1 (Wagg et al. 2025a) text-to-video diffusion model into an autoregressive video diffusion model. Specifically, replace all bidirectional attention layers with our LaCT layers combined with sliding window attention, then fine-tune the model using an internal proprictary collection of videos, each accompanied by a short text prompt.
Training details. We train on 5-second videos at 16 FPS and 430 resolution, autoregressively denoising latent-frame chunks. Later we fine-tune the 1.3 billion parameter model with 10 second videos and 14 billion parameter model with 8.8 second videos. Each 8.8 8.8-second clip contains 56, 160 visual tokens, resulting interleaved noisy-clean chunks totaling 107K tokens under teacher-forcing training. Full details are listed in App. E.3
Baselines. We compare our method against three baselines: sliding window attention (SWA) alone, Mamba2 (Dao & Gu 2024) combined with SWA (using similar in-layer hybrid strategy method), and full block-wise causal attention.
Evaluation. We evaluate all models on a collection of 2,000 videos after 5,000 training iterations by computing the denoising loss at five timesteps (550, 650, 750, 850, 950). Figure 6 plots the chunk- wise denoising loss across evaluated video frames. See Appendix EE.3 for results on VBench (Huang
ct al. 2024). See the anonymous website in abstract for our autoregressively generated videos.
CONCLUSION
We presented LaCT, novel model architecture that integrates large-chunk test-time training for capturing long context with window attention for modeling local structure. We validated LaCT across three diverse tasks spanning different modalities novel view synthesis, language modeling. and autoregressive video diffusion-ann demonstrate its effectiveness by achieving superior or competitive performance when compared state-of-the-arr baselines. LaCT achieves high GPU efficiency even with native PyTorch implementation with dozens of lines of code and supports efficient scaling up the state size and more flexible designs in test-time training models and optimizers. By open-sourcing the code and weights, we hope that LaCT can advocate future research explorations into more performant architectures for long-context modeling.
Reproducibility statement. We provide pseudocode for large-chunk test-time training (Algo- rithm(]), the in-layer hybrid strategy with sliding-window attention (Algorithm 2 and context/tensor- parallel implementations (Algorithm 3 Algorithm 4). We provided the implementation details and hyperparameters in Appendix Eoro our main results. The source code, included in the supplementary materials, contains the model and full training details. also provides a benchmarking code to evaluate the running speed of our method against two baselines.
Acknowledgment. We thank Ziwen Chen for processing the DL3DV dataset and providing the K-means clustering. We thank Nathan Carr, Feng Liu, Jianming Zhang, and Hailin Jin for their generous support this project. We thank Haian Jin and Zexiang Xu for leading the LVSM project. We thank Baqiao Liu and Haoran Cai for the video data loader. We thank Guo Han for details on the language dataset. We thank Jeremy Bernstein for discussions Muon.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Together AI. Long data collections database, 2024. URL https://huggingfaet co/
datasets ortraa/s//oocccarrmmmmee//////go
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in atari. Advances Neural Information Processing Systems, 37:58757-58891, 2024.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint 2tfi:2440.116664 2024.
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time, arXiv preprint rrXiv:2501.00663 2024.
Ali Behrouz, Zeman Li, Praneeth 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 arXivv2505.233735 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
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models world simulators. 2024. URL httph:ttt////m https:://ppeees com/ research/
vi sideo-generat on monemode s-as----- wwr-wor---- d-s mu lators
Boyuan Chen, Diego Martí Monsó, Yilun Du, Max Simchowitz, Russ Tedrake, and Vincent Sitzmann. Diffusion forcing: Next-token prediction meets full-sequence diffusion. Advances in Neural Information Processing Systems, 17:24081-241255 2024a.
Zhe Chen, Jiannan Wu, Wenhai Wang, Weijie Su, Guo Chen, Sen Xing, Muyan Zhong, Qinglong Zhang, Xizhou Zhu, Lewei Lu, et al. Internvl: Scaling up vision foundation models and aligning for generic visual-linguistic tasks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 24185-24198, 2024b.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.00 0080000..6611..1 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, arXiv preprint arXiv:2405.2 1060, 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, 35: (6344-16359) 2022.
Matt Deitke, Dustin Schwenk, Jordi Salvador, Luca Weihs, Osear Michel, Eli VanderBilt, Ludwig Schmidt, Kiana Ehsani, Aniruddha Kembhavi, and Ali Farhadi. Objaverse: A universe of annotated 3d objects. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 13142-13153, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint 10Xi::0000.11929 2020.
Laura Downs, Anthony Francis, Nate Koenig, Brandon Kinman, Ryan Hickman, Krista Reymann, Thomas B McHugh, and Vincent Vanhoucke. Google scanned objects: A high-quality dataset of 3d scanned houschold items. 2022 International Conference on Robotics and Automation (ICRA), pp. 2553-2560 IEFE, 2022.
11

---

Published as a conference paper at ICLR 2026
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis, 2024. URL https://arxiv. org/abs/2403.cmm 032006 2, 2024.
Jiarui Fang and Shangchun Zhao. Usp: A unified sequence parallelism approach for long context generative ai. arXiv preprint arXiv:2405.07719 2024.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. URL
https::/arxii: http:///wwaaa.ooaa////2.0......0
Leo Gao, Jonathan Tow, Stella Biderman, Shawn Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jasmine Hsu, Kyle McDonell, Niklas Muennighoff, et al. A framework for few-shot language model evaluation. Version vO. 0.1. Sept, 10:8-9, 9, 2021.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arXiv:23/2.00752 2023.
Junlin Han, Filippos Kokkinos, and Philip Torr. Vfusion3d: Learning scalable 3d generative models from video diffusion models. In European Conference on Computer Vision, pp. 333-350. Springer, 2024.
Horace He, Driss Guessous, Yanbo Liang, and Joy Dong. Flexattention: The flexibility of pytorch with the performance of flashattention. PyTorch Blog, 8, 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. arxiv e-prints. arXiv preprint arXiv:1512.0338 10:9, 2015.
Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen. Query- normalization for transformers. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4246 4253, 2020.
Geoffrey E. Hinton and David Plaut. Using fast weights deblur old memories. In Proceedings of the ninth annual conference of the Cognitive Science Society, pp. 177-186, 1987.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's 's the real context size of your long-context language models? In First Conference Language Modeling, 2024.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality linear time. In International conference machine learning, pp. 9099-9117. PMLR, 2022a.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc V. Le. Transformer quality in linear time, 2022b.
URL htttss://axxii org abs 2202 10447
Ziqi Huang, Yinan He, Jiashuo Yu, Fan Zhang, Chenyang Si, Yuming Jiang, Yuanhan Zhang, Tianxing Wu, Qingyang Jin, Nattapol Chanpaisit, et al. Vbench: Comprehensive benchmark suite for video generative models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 21 21807-21818 2024.
DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur, Block-recurrent transformers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: /tttps://operrviii
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint 2rXiv:2330.14500 2023.
Andrew Jaegle, Felix Gimeno, Andy Brock. Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. International conference on machine learning, pp. 4651 16644. PMLR, 2021.
12

---

Published as a conference paper at ICLR 2026
Haian Jin, Hanwen Jiang, Hao Tan, Kai Zhang, Sai Bi, Tianyuan Zhang. Fujun Luan, Noah Snavely, and Zexiang Xu. Lvsm: A large view synthesis model with minimal 3d inductive bias. arXiv preprint ar arXiv:2410.17244 2024a.
Yang Jin, Zhicheng Sun, Ningyuan Li, Kun Xu, Hao Jiang, Nan Zhuang, Quzhe Huang, Yang Song, Yadong Mu, and Zhouchen Lin. Pyramidal flow matching for efficient video generative modeling. arXiv preprint arXiv:2410.0599. 2024b.
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https
7/kel Terellerorddn titoo.. wttppstts//mnn/
Mahdi Karami and Vahab Mirrokni. Lattice: Learning to efficiently compress the memory. arXiv preprint arXiv:2504.05647 2025.
Bernhard Kerbl, Georgios Kopanas, Thomas Leimküüler, and George Drettakis. 3d gaussian splatting for real-time radiance field rendering. ACM Trans. Graph., 42(4):139-1, 2023.
Mare Levo and Pat Hanrahan. Light field rendering. 1996.
Zhixuan Lin, Evgenii Nikishin, Xu He, and Aaron Courville. Forgetting transformer: Softmax atten- tion with a forget gate. In The Thirteenth International Conference on Learning Representations, 2025.
Lu Ling, Yichen Sheng, Zhi Tu, Wentian Zhao, Cheng Xin, Kun Wan, Lantao Yu, Qianyu Guo, Zixun Yu, Yawen Lu, et al. D13dv-10k: A large-scale scene dataset for deep learning- basegg--bbedddee 3d vision. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22160-22169, 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Ruoshi Liu, Rundi Wu, Basile Van Hoorick, Pavel Tokmakov, Sergey Zakharov, and Carl Vondrick. Zero-1-to-3: Zero-shot one image 3d object. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 9298-9309, 2023.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In The Eleventh International Conference on Learning Representations, 2023. URL URL httt:ttt/ppppeeeeeeee net nett/Eorumiid-qNeee
Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, LILI YU, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient LLM pretraining and inference with unlimited context length. In The Thirth-eighth Annual Conference on Neural Information Processing Systems, 2024. URL
wwt.eviwwwwweeeeeeeeeeeeeeeeeeee Eorum wtwp:ummm-id-AAAAuuuuaaa
Leonard McMillan and Gary Bishop. Plenoptic modeling: An image-based rendering system. 1995.
Xin Men, Mingyu Xu, Bingning Wang, Qingyu Zhang, Hongyu Lin, Xianpei Han, and Weipeng Chen. Base of rope bounds context length. arXiv preprint arXiv:2405, 14591, 2024.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1)999-106, 2021.
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention, 2024. URL http::t://aiii..oo//////
24C4.07144
A Polyak, A Zohar, A Brown, A Tjandra, A Sinha, A Lee, A Vyas, B Shi, CY Ma, CY Chuang, et al. Movie gen: A cast of media foundation models, 2025. URL https:/aarxii. org/abs/2410.13720, pp. 51, 2024.
13

---

Published as a conference paper at ICLR 2026
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 Machine Learning, 2024.
David Ruhe, Jonathan Heek, Tim Salimans, and Emiel Hoogeboom. Rolling diffusion models. In International Conference on Machine Learning, pp. 42818 428885.. PMLR, 2024.
Sand-AL. Magi-1: Autoregressive video generation at scale, 2025, URL
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355 9366. PMLR, 2021.
Jürgen Schmidhuber. Learning to control fast weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1): 1992.
Noam Shazeer. Glu variants improve transformer, 2020.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp. 2265.. pmlr, 2015.
Kiwhan Song, Boyuan Chen, Max Simchowitz, Yilun Du Russ Tedrake, and Vincent Sitzmann. History-guided video diffusion. arXiv preprint arXiv:2502.067644 2025.
Yang Song. Jascha Sohl-Dickstein, Diederik Kingma, Abhishck Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 13456, 2020.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.
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): Runs with expressive hidden states. arXiv preprint arXiv:2407.04620 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint 20Xiv:2307.08621, 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, lliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunoy, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Dani Valevski, Yaniv Leviathan, Moab Arar, and Shlomi Fruchter. Diffusion models are real-time game engines. arXiv preprint arXiv:2408. /4837, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Ang Wang. Baole Ai, Bin Wen, Chaojie Mao, Chen Wei Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang, Jianyuan Zeng, et al. Wan: Open and advanced large-scale video generative models. arXiv preprint arXiv:2503.2033 2025a.
14

---

Published as a conference paper at ICLR 2026
Junxiong Wang, Daniele Paliotta, Avner May, Alexander Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models. Advances in Neural Information Processing Systems, (1162432-624557 2024.
Ke Alexander Wang, Jiaxin Shi, and Emily Fox. Test-time regression: unifying framework for designing sequence models with associative memory, 2025b. URL
abs/ abb/2501.. 12352
Wenhan Xiong, Jingyu Liu, Igor Molybog. Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. Effective long-context scaling of foundation models, 2023.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. In International Conference on Machine Learning, pp. 56501-56523, PMLR, 2024a.
Songlin Yang. Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. arXiv preprint arXiv:2406.06484 2024b.
Songlin Yang. Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024c.
Zhuoyi Yang, Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang, Jiazheng Xu, Yuanming Yang, Wenyi Hong, Xiaohan Zhang, Guanyu Feng, et al. Cogvideox: Text-to-video diffusion models with expert transformer. arXiv preprint arXiv:2408.0600 06072, 2024d.
Tianwei Yin, Qiang Zhang, Richard Zhang, William Freeman, Fredo Durand, Eli Shechtman, and Xun Huang. From slow bidirectional to fast causal video generators. arXiv preprint 2rX1v:2412.07772. 2024.
Kai Zhang, Sai Bi, Hao Tan, Yuanbo Xiangli, Nanxuan Zhao, Kalyan Sunkavalli, and Zexiang Xu. Gs-Irm: Large reconstruction model for 3d gaussian splatting. In European Conference on Computer Vision, pp. 1-19. Springer, 2024.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and paitern recognition, pp. 586-595, 2018.
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. Unipe: A unified predictor- corrector framework for fast sampling of diffusion models. Advances in Neural Information Processing Systems, 36:49842-498 49869, 2023.
Chen Ziwen, Hao Tan, Kai Zhang, Sai Bi, Fujun Luan, Yicong Hong, Li Fuxin, and Zexiang Xu. Long-Irm: Long-sequence large reconstruction model for wide-coverage gaussian splats. arXiv preprint arXiv:2410.122 12781, 2024.
15

---

Published as a conference paper at ICLR 2026
Below is the Appendix of the submission.
A ANALYSIS ON DESIGN CHOICES
In this section, we analyze several key design choices in our model, focusing on both the novel view synthesis and language modeling tasks, where good metrics exist. Specifically, we evaluate the impact of state size (Fig- ( test-time optimizers (Fig. ( 7p), linear versus nonlinear fast weights Fig..B., and per-token recurrence versus chunk-wise recurrence (Fig. 8p). Overall, we find that a large state size, advanced optimizers such Muon, and nonlinear fast weights significantly improve our model's performance. For comparing chunk recurrence with per-token recurrence, in a controlled linear large-chunt recurrence strategy outperforms linear per-token recurrence with the same state size. For language modeling, where chunk structures are not inherent, our linear large-chunk recurrence variant-while initially underperforming per-token methods like GLA and Surpasess them when combined with large nonlinear state and the Muon optimizer. We refer the readers each figure and its caption for more detailed analysis.
NVS experiment,
novel view synthesis
language model
novel view synthesis
language model
36
20
2,60
2555
34
2.50
28
&ivens 22 clews
4l views
2V
4
16k
32
siewss
12 vieers
vieos
2
4.
2
16
32k
Tecam indac.
Toleen indee
0. 0.376
0.76.
1.6
3,
Edt
12tt'
Vanilla Gradant Desicunt
Momentum
Muon
(a) State Size Scaling
(b) Different Test-Time optimizer
Figure 7: (a) Scaling up the state size consistently improves performance in both novel view synthesis and language modeling tasks. Note, the largest version has state size of 12dl per block, totaling 40% of model weights as fast weights. (b) Comparison of test-time optimizers demonstrates Muon's s surprising effectiveness over Vanilla Gradient Descent and Momentum.
novel view synthesis
language model
novel view synthesis
language model
30
8.62
15ML
Lo05100
PSWR
9.60
34
2.58
2.58
vilwss
48 views
2k 4k 8k 16 32 Tokan indox
4 4vicees
32 viows
2.54
2kk 4.K 8K 15 32 Token indox
Number of Inout Views
Namber Irput Views
Lineer
Mamite
SWA
Oura
SwILLL Mu?
SwiGLL MLP 0,575
DullaNei SWA
Oura SwiGLU a Lange State Muon
(a) Linear v.s. NonLinear Fast weight
(b) Large- Chinee v.s. Per-token Recurrence
Figure 8: (a) Nonlinear fast weights consistently outperform linear fast weights despite using smaller state sizes. (b) Our linear large-chunk recurrence approach significantly outperforms linear per-token recurrence (bidirectional Mamba2) for view synthesis tasks at the same state sizes. In language tasks, linear large-chunk recurrence of the same state size underperforms the GLA baseline, but when combined with larger nonlinear states and Muon test-time optimizer, it surpasses all per-token recurrence methods.
State size scaling, These controlled experiments utilize SwiGLU MLP for fast weights and the Muon as the test-time optimizer. For NVS, experiments were conducted on the object dataset. All models were trained for 167B tokens, using 14 stacked blocks and a model dimension d 768.
16

---

Published as a conference paper at ICLR 2026
To change the state size, keep the head dimension fixed model dimension. i.e. single head, and vary the intermediate dimension of SwiGLU MLP, such that the intermediate dimension ranges from 192 3072. The largest configuration results in state size per model block 12df, totaling 40% of model weights as fast weights. For the language model experiments, we use the 760 milion parameter setup, where the chunk size and sliding window attention (SWA) window size were set to 2048 tokens. We keep the intermediate dimension of the fast weight SwiGLU MLP the same as the head dimension. We increase the state size while proportionally decreasing the number of heads to maintain a fixed model dimension. Figure 7 () demonstrates that larger state sizes consistently improve performance. Notably, the performance gap between small and large state sizes widens with increasing sequence length.
Test-Time optimizer comparison. We compare Muon with vanilla Gradient Descent (GD) and GD with momentum. Details momentum implementation are in Appendix([] For NVS, we train all compared approaches for 671 tokens using model specs of 24 stacked blocks with model dimension of 768. Language modeling experiments used the 760M parameter setup. Figure 7 shows Muon consistently outperforming other optimizers.
Linear v.s. NonLinear fast weight. Our default fast weight function is a SwiGLU MLP without bias terms (nonlinear). We compare this against a simple linear fast weight, (w (x) = Wx. Both are updated using the same online dot product loss for key-value association. Figure 8 (a) presents this comparison for NVS and language modeling. Although the linear fast weights were configured with a larger state size than the nonlinear SwiGLU, they achieved lower performance. NVS models were trained for 67 IB tokens with 24 blocks and d = 768. Language modeling used the 760M parameter setup.
Large-chunk v.s. Per-token recurrence. Figure 8(b presents controlled experiments comparing our large-chunk recurrence with per-token recurrence. In the novel view synthesis (NVS) task, "Our Linear" variant employs a linear fast weight: fw (z) Wa x and is benchmarked against Mamba-2 baseline (a linear per-tobe per-token recurrence model) with an identical state size. To accommodate the bidirectional context required by NVS over input image tokens, the Mamba-2 baseline uses two Mamba-2 layers applied in opposite directions within each model block. Both our linear variant and this bidirectional Mamba-2 have state size of d2 per block. Both of these a per-image window attention within each model block. Under this fair comparison, our linear large-chunk recurrence achieves significantly better view synthesis performance.
approaches employs
For the language modeling experiments also shown in Figure Bb), the blue line "Our Linear" variant uses the same state size (0.25d?) as the GLA SWA baseline. initially underperforms GLA SWA (blue line underperforms yellow line), likely because language data lacks the inherent chunk structures that benefit our basic linear chunk recurrence. However, when LaCT equipped with larger non-linear state (1,5d² and Muon updates, we significantly outperform these per-token recurrence baselines.
B ADDITIONAL RELATED WORK
Novel view synthesis. Novel view synthesis (NVS) is a long-standing task at the intersection of computer vision, graphics, and computational photography, requiring algorithms to render images of a static scene from previously unobserved viewpoints. Optimization-baser approaches, such as NeRF (Mildenhall et al. 202T) and 3D Gaussian Splatting (Kerbl et al., 2023), have achieved significant breakthroughs. These methods optimize set of parameterized graphics primitives (i.e.. explicit or implicit representations of radiance fields) through differentiable volumetric rendering to minimize reconstruction loss on input images. After an optimization process typically lasting tens of minutes, these approaches render novel views photorcalistically, and the optimized parameters form a 3D representation of the input scene.
Recently, data-driven approaches (Zhang al. 2024 Jin et al. 2024a Ziwen et al. 2024 Han et al.
2024 Liu et al., 2023) have also shown promising results. These methods can either directly render novel views predict 3D representations given input images. Although successful on simpler object datasets, these methods often struggle with densely sampled scenes (c.g., scenes with over 100 input images). Our experiments demonstrate that our large-chunk test-time training approach outperforms
17

---

Published as a conference paper at ICLR 2026
or achieves comparable performance to 3D Gaussian Splatting on challenging scene datasets with up to 128 imput images with 960 x 536 resolution at challenging scene datasets. We hope our method will inspire further research into effectively scaling data-driven NVS methods to longer and more complex input sequences.
Autoregressive video diffusion. Current state-of-the-art video generation is dominated by bidirec- tional diffusion transformers operating in latent space (Brooks et al. 2024] Yang et al.) 2024d] Polyak
al. 2024; Wang et al. 2025a). These methods factorize the video distribution into sequence of conditional distributions based on noise levels, following diffusion processes (Sohl-Dickster
et al.
2015 Song et al. 2020) or flow matching (Lipman et al. 2022), then use a diffusion transformer to
jointly learn all the conditional distribution. Autoregressive video diffusion (Alonso et al. 2024 Jin
ct al. 2024b Valevski et al. 2024; Ruhe et al. 2024 Yin et al. 2024 Song et al. 2025) introduces
an additional temporal dimension to this factorization, where the neural networks learns to model the
conditional probability of the next chunks of videos at different noise levels, conditional on previous
videos and noisier version of current video frames.
During training, some autoregressive methods employ teacher forcing, supervising the model on noisy video frames given previous clean context frames condition (Alonso et al., 2024 Jin et al.. 2024b) Valevski et al., 2024), though this can lead to low token utilization, i.e. only a small portion of tokens get supervision. To improve token efficiency, other techniques such progressive noise injection (Ruhe et al. 2024) or the use of frame-independent noises (sometimes in diffusion-forcii style) (Yin et al.] 2024 Chen al. 2024a; Sand-Al 2025) have been proposed. When applying our large-chunk design to autoregressive video generation, we format the input sequence with interleaved clean and noisy chunks (see Equation IT]. This strategy achieves over 50% token utilization and integrates effectively with our large-chunk TTT implementation, by only changing few lines to constrain feithweights are only updated on clean frame chunks.
C LACT MODEL IMPLEMENTATION DETAILS
State Size calculation. Motivated by recent progress in LLM, we adopt SwiGLU-MLP (Shazeer 2020) without bias terms as the fast-weight network. Our fast weights consists of three weight matrix [W1, W2, Wa}, and the forward pass of the fast weight model is:
fw) W2 [W1) o (Wgz)]
(12)
where o is an elementwise multiplication. We define hd as the head dimension, nb as the number of heads, and the intermediate dimension of the SwiGLU-MLP as hd x r, where is scaling multiplier. When > 1, the intermediate dimension surpasses the input head dimension, which is the current common practice in LLMs. Thus, matrices Wi, W2 R and W3 Rddxddxr Consequently, the total state size becomes nA x hd x hd = r. Given that typically the total head dimension across all heads equals the model dimension (i.e., nh x hd = d), the total state size simplifies to:
de
State Size
nh * r
(13)
Therefore, we can increase the state size either by reducing the number of heads or by increasing the intermediate dimension multiplier.
FLOPs calculation. When using then negative dot product loss as the online test-time training objectives, we don't need compute the final results of Sw (w) We only need to compute Wi v, Wzy when running forward pass with keys k; thus there are two matmuls in the forward pass with keys. When computing the gradients, there are four matmuls. And in the final forward pass the queries. there would be three matmuls. So the total FLOPs with n tokens would be:
d2 r' Sn nh ruh
d2
d2
d² nh
FLOPs = 4n nh
6n
18n
6 State Size
(14)
Model initializations. We randomly initialize the linear layers using a standard deviation of 0.02.
For the learnable initial fast weights, we initialize them with a standard deviation of 1.0/ / //aann in. Specifically, in the SwiGLU FFN fast weights, the matrices wi and ua have their fan-in set as the head
18

---

Published as a conference paper at ICLR 2026
Algorithm 1 Large Chunk Test-Time Training Layer Pseudocode
def apply_ weight,, g):
w1, w2, w3 faat_
hidden = silu (matmul (9, wll) matmul (q, *3) $ [b, 1, dh] = [b, 1, d] [b, d, dh]
return
(hiddon, w2)
dof update (fast_weight, V, Ir, :
wast-weight update for SWIGLU MLP using chunk of tenaors. Args:
rast weat_weight tuple (w1, w2, wa) with shapes: wl, w3: (b, d, dhly w2: [b, da, d] v Xey value tenaor shape [b, 1,
Irt : per-token learaning rates of shape [b, 1, 3] (1rl, 1r2, 1r3)
use_muon weather to apply Muon to orthogonalize the update
Note:
The head dimension for input tersers k, Ir are assumed to be merged into the batch dimension. This simplifies shape annotation this pasudocade.
# Porward with k;
tete__eefoee_eee act matmul (2, wil [b, i, dhl [b, 1, d) 1b, d, dh] hi bddee__efore_eeeeee natmul ik, w3) [b, ch] [b, d] (b, d, dh]
hidden bilutgattebbeooeeeeeee
hidden before_gearr gate
Backward:
dhidden matmul (v, w2. (2.transpoan (-1, -2)1 t [b, 1, dh] [b, 1, d] x [b, 6, dh]
dhidden silu (gase_before_act) dgate dhidden hidden baforre_aateeee
ggate_before__ee
silu backprep (dgate, yate_basfore_ss))
# Compute gradients:
grad -matmul (bidden.
-2), 1r2) bb, dh, d]
[8, d, da) - [b, d, 1] x (b, 1, dh)
wl wl.grad -matmul (k 1rll transpose (-1, -2), caate_before_act
w3.grad -motmul
1r3] ttrnsppose (-1,
Weight update
if use_mues mue_m
for w in [ast_weight:
k.grad zoropowas_viae zergpowes_jii.rootooosshhhle (w.grad)
for w in fast_weight
[W N.grad)
.ngr)d)nndddddee w.norm(dim-1)
reture fast_weight
def bal__backpropddy x): sigma si sigmoid
reture dy sigma
sigma):
**АНАААИНИААА x: input sequence (b, 1, d), is the batch dim, is length, is model dimension faat_weight: tuple of initial fast weighta- twl, w2, m3); wl, w3 of shape [nh, d, dh [nh,
MultiHead Lace layer #**.**11111*************
1,
tt: tt=_config: 1ist of (operation, begin, end) tuples [b, 1, d
<kv (iin 4, k, v rearrance (qku, "b (nh hd) -> (b nh) hd", (h=num_haads) sppitt
9, X q. q. softplus (Id near k.norm (-1) const_lr_bias,
f [b, 3, num_heads]
ir rearrange (Ir, "b (nh 3) -> (a nh) 3", ah-num heads)
fast_weight repeat flast_weight dim-o, zepeat-b) [nh, ...] [b nh, ...l
zeros_like(v) tb nh, 1, hd] for mode, begin, end in ttt cot_config: qi, ki, vi, lri qlt, beginsend), klt, beginsend), v(t, beginsend), Ir[t, begin: endl
1- made == supdate_then. thtp_apppe___pppppppppppppppttrt fiqure-3 (a, bidi rectional attantion fast wright update weight, ki, vi, iri, uso_muon]
olt, beqin: end] amply_Ewifast_wweinh qi)
elif mode - then
Eiguze-3(b) shifted block-wise causal ol:, begin: end] aggpy_ qi)
fast_weight update update (fast_ ki, vi, Iri, use_muon]
elif mode -- prodate_oniy'; fast_weight = update (fast_e(fast ki, vi, Iri, use_muon)
slif mode == hmply_only';
ol:, begin: end] fpplyyfwlfastt_eeigg
RX5Ner= 4 per-head norm
LinearCutput trearsange (b nh) hd hd)", clamum_hesdaa
return
19

---

Published as a conference paper at ICLR 2026
Algorithm 2 LaCT Layer with In-Layer Hybrid Window Attention Pseudocode
Input::
x: input sequence [b, 2, d], is batch dim, is length, is model dimension Iaat_Weight weight: tuple of initial fast weighta- tw1, w2, wa) wh, w3 of shape [d, dal, w2 of shape [dh, d]
ttt_config: 1ist of (operasion, begin, end) tuples
9, k, v LinearcKV (3) (ppll
#### Local quadratio-cost window attention
attn_a q = 1earnabl e_s_aal learrable_g_offatt f per-chernal reacal la and shift learnable_ learrable_k k_fffast f per-chennel rescals and shift attn_o losal Looootaay_mltthhheeeattttt atttaaa attn_k, 9, attn_mask)
atta_k
test-time training &, x - rearrange (q, k, "b 1 (nh hdi -> (b nh) 1 hd', ah-nun_heads)
+4# large
long memory
5 silu(q), sklu(k)
9. X - q q.norm(-1), & k.norm(-1) [b, 1, num_ heada]
softplus (1))
num
rearrange "b
3) -> mh) 3', nhensm_heads)
Perform update and apply_fe operation iterativaly lact_ a last (fast weiht, 9, k, v, lr, tti_conffg)
churks of tokens.
lact_o /lacoo
scale_per__iee yer yer_heed - rearsange (silu (x)), "b 1 nh -> (b nh) 1 17, mheaun lact_o lact_o acale_per_her
scale
laot_o rearrange (lact_o, (b nh) 1 hd -> (nh hd)", nh-num_heada)
#4# Merge atcention results (shape: [b, d])
attn_c lact_o
LinearCutput (o)
return
dimension, while the fan-in of w2 is the intermediate dimension of the SwiGLU FFN fast weights. Additionally, when local window attention is incorporated within the LaCT layer, introduce four extra learnable embeddings: two scales and two reshifts for queries and values. We initialize the scale embeddings as ones and the reshift embeddings as zeros.
Details of Muon. Muon (Jodann al. 2024 is recently proposed optimizer that orthogonalizes the matrix gradients during updates of matrix weights. It utilizes Newton-Schulz iterations to achieve orthogonalization. Given a matrix gradient G, Muon first normalizes it as Go G/ then iteratively applies:
= =
G
(15)
where the constants a, b, are carefully chosen for optimal convergence. Following the original implementation, we set a = 3.4445, -4. -47750, e = 2.0315, and perform five iterations.
Each Muon iteration requires three matrix multiplications, resulting in computational cost per fast weight head of 2hd + 2hd + 2hd hd (4r FLOPS. Hence, the total computation for five iterations across all fast weights is:
x nh hd (4r + 2).
(16)
For the case where r 1 (head and intermediate dimensions equal), the total computational cost simplifies to:
30 x wh x hd? = 30 hd x State Size.
(17)
This indicates that the computational overhead of Muon becomes less significant than computing token outputs only if the online chunk size exceeds 5 3 hd.
Rotation invariance. Softmax attention and linear attention exhibit rotation invariance: rotating the queries and keys by the same rotation matrix does not alter the output. This property is also used in developing relative positional encodings, like RoPE (Su et al. 2023). In contrast, our SwiGLU and Linear Fast Weight components do not possess this property.
20

---

Published as a conference paper at ICLR 2026
Implementing momentum for test-time optimizers. Muon uses momentum by default. Following Titans (Behrouz et al. 2024), we implement momentum in the test-tima optimizer by predicting a scalar momentum B; from each token:
B, = a Linear(::)):
(18)
where the sigmoid function. This Bg is then averaged over all tokens in the chunk, and the average momentum is applied as follows:
V w (()),
(19)
M M Br/b) + 9.
W + weight--ppatteeWWWW M),
where the weight-update can be simple subtraction followed by L2 normalization normalization (as in Equation 8). or Muon update before subtraction (as in Equation 9
C.1 PSEUDOCODE
See Algorithm ! for pseudocode of full LaCT layer. For details on how to mix local window attention inside each layer with shared query, value, embedding. see Algorithm for pseudocode.
D LACT ARCHITECTURE DETAILS FOR N-DIMENSIONAL DATA
D.1 LACT ARCHITECTURE FOR NOVEL VIEW SYNTHESIS
Novel view synthesis (NVS) renders images of static scene from novel viewpoints. Formally, given a set of N input posed images {(T, P,)) of a static scene, where J, E H x3 is an RGB image and P. is its corresponding camera pose, the model needs to synthesize new images from novel camera poses Pnovel that typically do not overlap with the input views.
Traditional methods in 3D vision usually solved the NVS task by reconstruction and rendering, The reconstruction compresses the posed input into a compact representation and then the render renders the novel view from it. Our method mimics such pipeline, where we first compress the posed input images into fast weights by the 'Update Operation' (Sec in LaCT. Then we render the novel view images from the information in the compressed weights with the 'Applying Operation'.
In details, we first convert the input and output into tokens. The camera pose P for each view is represented in dense ray information for each pixel (usually from the camera intrinsics and extrinsic), i.e., P = (rays,. rays,). rays, E H x3 is the 3D coordinate for the origins of the ray, and rays, E R xW x3 is the direction of the ray. We follow GS-LRM ZZhang et al. 2024) to the Plücker ray embedding for the rays. Plücker ray embedding computes the cross product between the ray origin and ray direction for a normalization. The final positional embedding is concatenation of the ray's origin, the ray's direction, and the cross product of the above two: [rays,, rays rays, X rayszl. We add ray's origin into the embedding since different origins in the same ray can results in different colors due to the occlusions. We then use patchifying and two different Linear layers to convert the RGB map and ray map (i.e., the positional embedding) into tokens. For the posed RGB input images, we simply the sum the RGB embedding and pose embedding as model input. For the novel view cameras, we only use the pose embedding.
We illustrate the design of NVS's LaCT block in Fig. 9 We first apply the attention for each image (either input the novel target). The attention is bidirectional for the tokens belonging to the same image, and independent among different images. Then, the TTT update operation is applied to all input tokens, i.e., all tokens that belonging to all input images. The updated weight then is applied to all tokens. The two updates blocks in Fig. take the same updated fast weight thus can be combined, and we left two 'Apply' block for clearance. Note that the original NVS task definition renders novel views independently. We here supervise multiple novel image poses and their corresponding images in single data point for better training efficiency. Given the design, the novel images
21

---

Published as a conference paper at ICLR 2026
Novel Image
Novel
Image
Leyers
Apply
Apply
Mask
W-
Update
Overall LaCT Mask
Win Aftn
Win Atli
Win Alln
Win Altn
Window-Attention Mask
Novel
Novel
Image
Image
Figure 9: Detailed LaCT model for our Novel-View Synthesis. Dashed line indicates flow of fast weight. Solid line indicates flow of tokens. Window attention is bidirectional within single image, either the input image or the novel target image. TTT updates over all input tokens and apply to all tokens.
Marge
Concater ono Large Chunks
Oveell LaCT:
Overal LaCT:
Full Causal
Full Causel
Sadkig Window
Updite
Tx.
Span into Lange Choniics
Rescala Shitt
Attentions Bliding Window Misi
Aftention: Skddng Window Mask
Causell Mask
Mussration of "Window Size" "TTT Chunk Size"
Illustratiion of "Window Size" "TIT Chunk Size" 2
Figure 10: Detailed LaCT model for language models. Dashed line indicates flow fast weight. Solid line indicates flow of tokens. We illustrate with TTT chunk size 4 or 2, and the actual chunk size over 2048 in LaCT. We take the parallel design for the window attention and TTT block with shared QKV. The overall mask is the causal mask.
independent to each other, which is illustrated in the 'Overall LaCT Mask' in the right of Fig. The layer normalization layer, the residual connections, and the feed-forward network is omitted for clarity. The block is repeated by number of layers times to formulate the full model. The general model largely follow the design the encoder-decoder model in LVSM (Jin et al. 2024a). except we TTT in replace of transformer for long-context modeling.
For actually using this model for NVS task during inference, we first get the updated fast weight with all input images. Then, would not change the fast weight during the rendering process (i.e., the process to convert novel camera poses to the novel images). The LaCT during rendering would be similar to ViT (Vision transformer) architecture despite having two feed-forward networks: the feed-forward network from the fast weight stores the scene information, and the feed-forward network from the slow weight (i.e., the FFN in Fig. 2 stored the world knowledge like physical rendering rules.
D.2 LACT ARCHITECTURE FOR LANGUAGE MODELS
Autoregressive Language Models (LM) predicts the distribution the next tokens
(r [11, " from its history context. It is factorization of the full sequence distribu-
22

---

Published as a conference paper at ICLR 2026
Overal LaCT Mask
Donobse Donokse's
Dunoise, Denose's
Dennin,, Denatie's
Apply
Windew Alarann
Window Ataer
Ataenion
Upsse
Aputsss
Notsoy
Nolse's
Cleans
Noloe
Noloo's
Cleaty
Nolso,
Notoe's
Shihed
Spitimno orga Chuntor
wolt: groups
Detailed LaCT model for autoregressive video generation by diffusion model. The purple tokens are noisy frame chunk for training the diffusion model. The green tokens are clean frame chunk. Each token is large chunk in TTT, e.g., 3 consecutive frames with 4680 tokens in total. Noise and Noise" are two noisy frames with independent Gaussian noise and time stamps over the same clean frame. We denoise them simultaneously to improve the utilization. Dashed line indicates flow of fast weight. Solid line indicates flow of tokens. We take the parallel design for the window attention and TTT block with shared QKV. The TTT mask is shifted (i.e., started with apply) strided causal. The window attention excludes the attention from clean frame to future noisy frame, and also excludes the attention between the independent noisy frames.
Figure 11:
1). Thus requires token-level causal mask (demonstrated in the topright of Fig. T0) and this the main difficulty for the large-chunk design in LaCT. We use combination of TTT layer with 'Shifted Causal Block Mask' (introduced before in Fig. 33) and a sliding window attention to facilitate By shifting the mask of TTT, it excludes the information leakage from future tokens. As shown in the right part of Fig. [0] the overall dependency mask is the union of the TTT mask and the sliding window attention mask. To achieve a token-level causal mask without bubbles, the only requirement is that the window size of the sliding window attention is greater or equal to the chunk size from the TTT. We illustrate two example of such mask with 'Window Size' 'TTT Chunk Size' 2 or 4. The actual chunk size and attention window size is above 2048 in our implementation for better utilization and state size scaling (discussed Sec. 3.2).
tion pa (= Fnn) through chain rule po (o In) (12222z ... ... (2 [1
As illustrated in the left most of Fig. T0 we employ a parallel design of the TTT layer and sliding window attention to save the number of model parameters and computation FLOPs. In details, the query (Q), key (K) and value (V) are shared between the TTT layer and window attention. Sliding window attention an attention with constant window size over the past history, starting from the target tokens. For the TTT layer, start with apply operation over the first chunk using the initialized fast weight (i.e., unupdated yet). The 'apply' operation is followed by the 'update' operation over the first chunk. In this way, the 'apply" operation would not information inside the current chunk to avoid leaking future token information inside the chunk. Alternatively using 'apply' followed by 'update" over subsequent blocks completes the desired 'shifted block-wise causal mask' illustrated in Fig. TO For details of the parallel design, please refer to the Pseudocode Algorithm
We use the multi-head design for both TTT layer and window attention, although their number of heads are different. We empirically take less number of heads for TTT layer (i.e., larger head dimension) enable larger state size, (Sec. C and Equation 13). By default, we use four heads in langugae model experiments. For positional encoding, we use the same RoPE as the window attention branch.
size is propotional the head dimension in our design
D.3 LACT ARCHITECTURE FOR AUTOREGRESSIVE VIDEO DIFFUSION
Chunkwise autoregressive video diffusion generates videos by iteratively denoising sequential chunks of video frames, conditioned on previously generated clean frames. Each chunk can contain several video frames and span thousands of visual tokens. We use teacher-forcing training by interleaving
23

---

Published as a conference paper at ICLR 2026
noisy and clean frame chunks. Specifically, video of N frame chunks is structured as:
S [x X1, X2 X .2,,..., X noise]
(20)
where each noisy chunk X Xpois is produced by adding unit Gaussian noise to the i-th clean video
chunk as Xprais= (1 f4) et, and t; E [0. 1] denotes the strength of
noise. However, compared to previous methods that employs progressive (Ruhe et al. (024) or Trame-independern noise strategies (Yin et al.) 2024) like diffusion forcing (Chen et al. [0224))
teacher-foreing formulation in Equation only uses around 50% of the tokens of the entire sequence
to compute the denoising loss. To improve token utilization, we consider an alternate approach by repeating each video chunk with two noise levels in the training sequence as:
S [x X X1, X. A2 X .2......
X X
(21)
where X.roie and X represent two different noise levels applied to each clean video chunk X;- This increases token utilization from 50% to about 67%. While more repetition could further increase token utilization, it would also reduce training sample diversity; thus, we limit the repetition to twice. We use such repeating strategy when training the 1.3 billion parameter video diffusion model on five seconds videos.
To process these interleaved noisy and clean chunks, LaCT fast weights are updated exclusively using the clean video chunks. These updated weights are then applied to the current clean chunk and subsequent noisy chunks. The integrated local window attention uses window size of two frame chunks and employs block-wise causal mask. This mask allows noisy chunks to attend only to themselves and the immediately preceding clean chunk. By doing this, we main bidirectional dependies within each chunk and causal dependency across chunks.
Similarly to our language model experiments, we integrate the TTT and local window attention into the same layer. Figure M iilustrates this design for autoregressive video generation. The input sequence depicted in the figure follows Equation 21 with each video chunk repeated twice with different noise levels. The pink color in the figure indicates noisy chunks and the green color indicates clean chunks.
E EXPERIMENTAL DETAILS
E. NOVEL VIEW SYNTHESIS
Datasets metric. We evaluate our approach on both object-level and scene-level datasets. We use Objaverse dataset (Deitke et al. 2023) for object-level training, and render 32 random views per object, following the setup from LV SM (Jin et 2024a) and GS-LRM (Zhang et al., 2024). After training, we perform evaluations on the Google Scanned Objects (GSO) dataset (Downs et al. 2022) at resolutions of 256 x 256 and 512 x 512. To ensure stablized evaluation results, we render at fixed view points instead of random view points in training. Each evaluation involves 4-48 input views and 8 novel views per object. For scene-level evaluations, we adopt the challenging DL3DV scene dataset (Ling et al. 2024), which contains over IK training scenes and 140 testing scenes, each with approximately 300 views. Evaluations are performed at a resolution of 960 x 5 We report Peak Signal-to-Noise Ratio (PSNR) the paper's main figures, and other metrics Structural Similarity Index Measure (SSIM) and LPIPS (Zhang et al. 2018) can be found in Tables below. For DL3DV evaluation, we follow the original paper (Ling et al., 2024) and LongLRM (Ziwen et al. 2024) to use one frame from every 8 frame in the full video sequence as the target frames. The input frames are from the K-means clustering of all frames as in (Ziwen et al., 2024).
Model details. Our models consist of 24 stacked LaCT blocks, each with model dimension of 768. The detail of such block is illustrated in Sec. E. Unless otherwise specified, we use a single-headed fast-weight SwiGLU-Me ULP with a hidden dimension of 1536. The window attention has 12 heads with head dimension 64, and is equipped with noomommaainntion (Henry et al.) 2020). The Feed-forward Network has 3072 as its intermediate hidden dimension. The model has a total of 312M parameters, of which 84M are fast weights (i.e., 6d2 per block). We use an fest-weight Ir initialization of 0.01 by
"The original DL33D 960p frames released in resolution of 960 536. To accommoaate the patch-size 8 in modeling, we crop to 960 x 536 and the camera parameters changed accordingly.
24

---

Published as a conference paper at ICLR 2026
setting 'const Ir bias' in Algorithm to softplus(const Ir bias) 0.01. As used Muon in fast weight update for NVS, LaCT is not sensitive to scale as discussed in Sec. (.2
Baselines. For object-level evaluations, compare against two basclines, including full-attention model, and a register-attention model in a Perceiver style (Jagle et al., 202T). In the full-attention baseline, replace the TTT layer in model with block-wise causal attention layer, where the input tokens interact bidirectionally and the novel view tokens cross-attend to the input tokens. Such design resembles our method's prefill and parallel decoding strategy described in Section 5.1 and the key-value caches of the input tokens server as scene representations for novel view renderings. In the perceiver-style model. we replace half of the TTT layers with imput-to-registe full-attention layers and the remaining half with register- register-to-novelee toveovll-view cross-attention layers. Such a model first compresses the input tokens into constant set of register tokens and then decoding the novel view tokens by attending to the registers. For scene-level evaluations, we compare against state-of-the-art long-sequence 3D reconstruction work LongLRM (Ziwen et al. 2024) that applies Mamba (Gu & Dao 2023) hybrid with full attention to predict 3D Gaussian splats (Kerbl et al. 2023). We also include comparisons with pure optimization-base 3D Gaussian splatting methods. Tab. 2 compares the computational complexity of the baseline models and our models.
Training details. For object-level experiments, we first train all the model with 671B tokens at input view and 8 novel view setting at resolution of 256 256. We then finetune them with 512 x 512 resolution for an additional 587B tokens. For scene dataset, we first pre-train our model first with 32 input views and 32 novel views 128 128 resolution for 1.5T tokens, then progressively finetune at larger resolutions, larger field-of-views, and more input views. The finetuning is always go with a non-squared FoV match the raw data. Non-squared FoV has larger view range than the squared FoV, thus is harder. The input and novel views in fine-tuning are both 64 to support better view coverage. The curriculum of the fine-tuning resolution is set as 128 72, 256 144. 512 288. and 960 x 536. The training tokens for each stage is around 100B, High-resolution models (starting from 512 x 288) are trained with inner-chunk context parallelism (See. 4.4).
At each training stage, we always use A with linear learning rate warmup and weight decay of 0.05. The peak learning rate of the pre-training 4e During fine-tuning, we use smaller learning rate (usually le 5 to 5e 5).
The training is completed with 64 A100 GPUs. The pre-training takes days, and each fine-tuning stage is about 12hours (thus 2 days in total).
Detailed Result Numbers We here provided the detailed number for object-level results on the GSO dataset (at resolution 256 x 256 in Table 512 x 512 in Table 4 and DL3DV evaluations (at resolution 960 x 536 in Table
Table 3: 256 Res object-level novel view synthesis results on GSO. Both the input and output are with resolution 256 256. t: higher better, J: lower is better.
Input Views
Input Tokens
LaCT
Full Attention
Perceiver Attention
PSNR (t)
LPIPS (.)
SSIM (f)
PSNR (f)
LPIPS (L)
SSIM (*)
PSNR (f)
LPIPS (.)
SSIM (f)
4,096
32.4
0,030
0.962
32.6
0.029
0.964
30.3
0.039
0.950
8,192
35.3
0.76
35.6
0.018
0.978
32.8
0.026
0.967
12
12.288
36.3
0.017
0.980
36.6
0.015
0.982
33.6
0.023
0.971
20
20.480
37.2
0.015
0.982
37.5
0.984
34,3
0.021
0.974
32
32,768
37.5
0.014
0.982
37.9
0.013
0.985
34.2
0.021
0.974
48
49,152
37.6
0,014
0.983
37.9
0.013
0.985
33.7
0.022
0.972
Table 4: 512-Res object-level novel view synthesis results on GSO. Both the input and output are with resolution 512 x 512 comparison across methods. higher is better, J: lower is better.
LaCT
Full Attention LPIPS (L)
Input Views
Input Tokens
PSNR (*)
LPIPS (L)
SSIM (*)
PSNR (1)
SSIM (f)
16.384
33.4
0.029
0.969
33.6
0.027
0.971
8
32,768
36.6
0.020
0.979
36.7
0.017
0.982
12
49.152
37.7
0.017
0.983
37.9
0.015
0.985
20
81,920
38.6
0.016
0.984
0.013
0.987
32
131.072
39.0
0.015
0.985
39.3
0.013
0.988
48
196,608
39.1
0.015
0.985
39.3
0.012
0.988
25

---

Published as a conference paper at ICLR 2026
Table 5: 960P scene-level novel view synthesis results for LaCT on DL3DV. Both the input and
output are with resolution 960 x 536 (width x height) t: higher is better, U: lower is better.
Input Views # Input Tokens
PSNR (1)
LPIPS ())
SSIM (1)
128,640
24.7
0.224
0.793
32
257.280
26.9
0,185
0.837
64
515,520
28.3
0.169
0.857
128
1,031,520
28.9
0.166
0.861
E.2 LANGUAGE MODELING
Datasets & Metrics. We train our models on the Long-Date Collections dataset AI 2024). ). con- Long-Data-Collections taining approximately 68.8B tokens tokenized using Mixtral tokenizer (32,000 codebook size). The dataset is a mix of 41 4%% General Data (e.g., RedPajama-Bont BBok, RedPaajma-ArXiv 1B tokens from RedPajama, and a Pile subsample) and 58.6% Instruction Data (e.g., UL2 Osear, NI, and P3), To evaluate long-context capabilities, we utilized the per-token loss metric from (Lin et al. 2025). A consistently decreasing per-token loss across the input sequence indicates effective use of the entire context, while a plateau suggests inability to leverage information beyond that point. Specifically, we evaluated next-token prediction loss on 2.5B tokens from the Book-3 dataset (Gao et al., 2020) for our 760M - model, and 5B tokens for our 3B-parameter model. Additionally, we measured retrieval accuracy using the RULER benchmark (Hsieh et al.] (024) across various se- quence lengths assess context memorization and information retrieval, evaluating up to the trained sequence length.
Model details. We modified the original LaCT block by removing its window-attention layer. Instead, we incorporate causal sliding
(SWA) layer directly into the Large-Chunl TTT layer. The SWA layer shares the same Q. K, and V vectors as the fast-weight network, except that a per-channel leranable scale and shift is applied to Q and K before they are fed to the SWA layer (as done in GAU (Hua et al. 2022al). We sum up the output of the SWA layer and that of the TTT layer, where the output of the TTT layer is scaled by another per-head learnable scalar. We use an fast- Ir initialization of 0.001 by setting 'const Ir bias' in Algorithm to softplus(const bias) = 0.001. We illustrate this archtecture in Figure 10
To ensure a fair comparison with baselines in terms of trainable parameters, we adjusted the LaCT block's i extra learnable initial fast weights W = [W, W2, Ws). o reduce parameters, we employed a low-rank version for W1, W3 with a rank of 32. For instance, if W1 E Rddd its low-rank initial fast weight is Wi = L R + 0.5 Id. where L. Rdx32 R E R32xd and Ta is identity matrix. This reduces the extra trainable parameters for the fast weights in each block to 128 * model-dim + nun-heads (model-dim²). Additional minor parameters for learning rate projection, per-head scalers, an extra RMSNorm, and the SWA's learnable scale and shift are of order (modelldimm) Standard blocks typically have 12- model-dim² parameters. Our approach adds approximately num beads (model-dim' extra parameters, which is less than 3% of total trainable weights with four heads, and below 4.5% with two heads. By default, LaCT use four heads in the experiments, unless noted otherwise, which means that the default state size per block is 3 (model-dim").
4
Baselines. We compare our approach with full attention, Gated Linear Attention (GLA) Yang et al. 2024a), DeltaNet (Schlag et al. [2021: Yang et al.) 2024c). To ensure fairnsss, we enhance both GLA and DeltaNet with the same sliding window attention. As pointed out in previous work (Lin ct al. 2025; Xiong et al. 2023 Men et al. 2024). a large RoPE (Su et al., 2023) base is critical for transformers in long-contexx training, thus we adopt a large RoPE base of I million for training with 32K token contexts whenever softmax attention is used. Tab, compares the mechanism and computing complexity of the baseline methods and our method. Training throughput (tokens per second per GPU, TPS) was using a 3B- 1B-parameter model on eight A100-40GB SXM4 GPUs with activation checkpointing and FSDP. At the 3 billion parameter scale, all models use 24 softmax attention heads. The GLA baseline has eight linear attention heads with heads dimension as 384, resulting in a total state size of 384d, with d = 3072 representing the model dimension. DeltaNet employs 24 linear attention heads, each with dimension of 128, leading to a total state size of 128d. Our approach uses four TTT heads with head dimension as 768, and since each block has three fast weights, the total state size is 2304d.
26

---

Published as a conference paper at ICLR 2026
Table 6: Comparison of baseline methods in terms of state size, training throughput (measured in tokens per second. TPS), update rules, and memory read-out mechanisms. Training throughput is evaluated using a 3B-parameter model with 32K- -equueneece length on A100-40GB GPUs.
State
Train TPS
Update Rule
Memory read-out
Transformer Transformer SWA
4.1K
6.4K
Per-token recurrence
GLA SWA
384d
8t Ding((a:) tvrk
o, S.qp
DeliaNet SWA
128d
5.1K
S.-1II A.kik.) Avik?"
O. S.q.
Large chunk
Ours GD
2304d
5.0K
W 1.2n0m((W Σ. CW 3M + VWw W L2uorm(W M)
fw (4.)
Ours Momentum
2304d
4.9K
fw (qc)
2304d
+ Σ A ViyC L2norm[W Muon (mon(M) (M)
o, fw (q.)
Training Details. We trained models at two scales using a sequence length of 32,768 tokens:
760M parameters: We 24 stack blocks, with model dimension 1536. All models trained for 40B tokens (40,960 steps) with a sliding window of 2048 tokens and a batch size of million tokens. Each experiment ran 32 A100-40GB SXM GPUs for approximately 20 hours.
3B parameters: We use 25 stack blocks, with model dimension as 3072. All models trained for 60B tokens (30,000 steps) with a sliding window of 4096 tokens and a batch size of 2 million tokens. Each experiment ran 64 A100-40GB SXM GPUs for approximately 50-60 hours.
For both scales, used base learning rate of 1 10-3 with cosine decay scheduler and 1024 warmup steps. All models were randomly initialized with a standard deviation of 0.02.
Results. Detailed results the RULER benchmark (Hsich al. 2024) are presented in Tables 7
and g We evaluated models on S-NIAH-1, S-NIAH-2, and S-NIAH-3 tasks, which represent varying difficulties of the single "needle in haystack" retrieval. We also report performance NIAH-MultKey-LL. MIAH-Multiuuryy, and NIAH-MultiValue Other RULER tasks are not reported the full attention baseline also achieved trivial results beyond a 16K sequence length. In addition to these long context evaluations, we also report some major language model benchmark in Table 9
Table 7: RULER benchmark results for Single Needle in Haystack (S-NIAH) tasks. Our method with two heads (default is four).
S-NIAH-1
S-NIAH-2
S-NIAH-3
Average
Midel
4K
8K
16K
32K
8K
16K
32K
4K
8K
16K
32K
4K
8K
16K
32K
760M parameters
Transformr
96.6
85.2
85.8
36.8
93.4
62.3
DeltaNet SWA
84.0
85.2
87.8
86.8
62.8
29.4
14.2
7.8
53.8
21.8
11.2
5.8
66.9
45.5
37.7
33.5
26.2
14.4
26.4
7.8
55.2
7.1
Ours
94.8
53.2
26.0
14.8
74.0
28.0
14.2
7.8
42.8
26.6
14.4
6.8
70.5
35.9
18.2
9.8
Ours Momentum
95.6
84.8
83.4
84.8
91.4
73.4
22.8
7.8
82.6
34.8
16.6
6.6
89.9
64.3
40.9
33.1
Ours Momentum *
59.0
30.0
12.4
8.4
93.4
50.0
18.2
7.8
60.2
25.6
14.2
6.8
70.9
35.2
14.9
7.7
Ours Muon
98.0
95.0
92.2
92.4
86.6
60.2
17.0
7.8
49.2
26.2
10.9
5.2
77.9
60.5
40.0
35.1
3B parameters
Transformrr
100
100
100
100
100
99.8
100
98.6
98.6
95.8
90.8
75.0
99.5
98.5
96.9
91.2
GLA SWA
52.8
26.0
13.2
14.4
27.6
12.4
13.3
DeltaNet SWA
100
89.6
76.2
54.8
100
76.4
42.2
17.0
90.6
57.6
27.4
13.4
96.9
74.5
48.6
28.4
Ours Monentum
98.6
39.6
15.0
13.4
55.3
40.6
Ours Muon
98.8
99.2
98.6
93.4
100
99.0
83.2
30.8
95.4
90.8
55.6
19.8
98.1
96.3
79.1
48.0
E.3 AUTOREGRESSIVE VIDEO DIFFUSION
We fine-tune the pretrained Wan 2.1 (Wang et al. 2025a) text-to-video diffusion model into autoregressive video diffusion model, that generates videos by iteratively denoising successive chunks of video frames.
27

---

Published as a conference paper at ICLR 2026
Table 8: Performance Multi-Key (MK-NIAH), Multi-Query (MQ-NIAH), and Multi-Value (MV- NIAH) Needle in Haystack tasks from the RULER benchmark. Our method with two heads (default is four).
MK-NIAH
MQ-NTAH
MV-NIAH
Average
32K
8K
32K
16K
16K
32K
760M parameters
Trunsformer
63.8
72
71.4
54
33.4
28.9
23.1
27.95
24
20.5
27.35
41.7
41.6
38.6
34.8
DellaNet-SWA
41.2
30
14.6
8.2
33
22.45
7.5
4.3
32.4
22.8
9.15
6.6
35.5
25.1
10.4
6.4
GLA + SWA
45.4
28.4
15.8
6.6
26.1
17.75
10.2
5.85
25.4
16.85
10.1
6.6
32.3
21.0
12.0
6.3
Ours
60.8
34.6
16.8
7
23.65
7.45
20.7
22.05
12..7
6.85
38.8
14.5
7.1
Ours Moncatum
62
41
21.2
10.4
35.3
24.99
17.7
8.6
27.9
23.15
16.65
8.2
41.7
29.7
18.5
9.1
Ours Momentum*
59.8
37.8
19.2
8.8
36.65
20.45
12.5
7.4
24.45
16.95
11.6
6.8
40.3
25.1
14.4
7.7
Muea
62.8
26.55
23.15
13.6
17.1
3B parameters
Trunsformer
95
90.4
81.6
65.2
86.45
81.55
71,70
40.85
61.8
42.8
30.75
22.9
81.1
71.6
61.4
43.0
GL.A 3B
78
14.4
28.00
10.7
29.4
21.4
16.25
52.5
31.8
21.4
11.7
DeltaNet SWA
75.8
57.4
34.2
17.8
66.25
35.05
21.45
13:45
43.7
23.2
18.
13.2
61.9
37.9
24.8
14.8
Ours Momentum
96.2
39.6
17.2
87.05
40.25
25.6
13.2
88.
30.65
21,9
12.3
90.4
43.5
27.5
14.2
Ours Muon
75.2
69.2
25.2
44.75
39.1
19
29.
25.05
48.8
32.0
Table 9: Language model results across common reasoning//nowlldgee benchmarks (higher is better unless noted). All results are obtained through 1m-eval wwt.-rhaaeliiiiinneeeeeeeee (Go et al. 202T)
ARC-c (acc)
ARC-
Hella. (acc_norm)
PIQA
BoolQ (acc)
OpenlBook (acc)
LAMBADA
SeiQ (acc)
LAMBADA
Wiki ppl..
Avg.
Model
(ace)
(acc_norm)
(acc)
(ace)
ppl.
(acc)
760M models
Trunsformer
0.247
0,465
0:405
0.664
0.588
0.526
0.190
0.375
0.864
22.574
20.750
0..9%
Trunsformer SWA
0.234
0.479
0.418
0.663
0.456
0,502
0.172
0.407
0.879
19,537
21.215
0.496
Ours Momentum
0.247
0.474
0.413
0.645
0.611
0.523
0.188
0.393
0.861
20.519
21,300
0.500
Muon
0.477
0,420
0.653
0.392
20.584
GLA SWA
0.238
0.475
0,420
0.649
0,501
0.528
0.184
0.395
0.897
20.621
20.695
0.502
DeltnNer SWA
0.253
0.466
0.417
0.650
0.608
0.515
0.184
0.386
0.888
20.415
23.285
0.502
3B+ models
Transformer
0.258
0,488
0.515
0.655
0.204
0.479
12.092
Transformer SWA
0.243
0,462
0.522
0.629
0.557
0,554
0.214
0,490
0.905
11,751
15.737
0.513
Ours Momentum
0.273
0.496
0.529
0.628
0.556
0.533
0.218
0.480
0.905
11.419
15.439
0.528
Muon
0.263
0.473
0.523
10.988
15.313
GLA SWA
0.259
0,466
0.517
0.628
0.508
0.535
0.190
0,465
0.897
12,547
17.284
0.506
DeltaNet SWA
0.255
0.475
0.504
0.648
0.491
0.547
0.190
0.453
0.891
12.682
17.950
0.498
Model details. The original Wan 2.1 is bidirectional diffusion transformer operating on the latent space of a causal video VAE, which performs 8x spatial and 4x temporal downsampling. The diffusion transformer uses x 2 x patchification layer convert VAE video latents to tokens. Each block of the diffusion transformer comprises an MLP layer, a bidirectional self-attention layer for visual tokens, and cross-attention layer for visual and text tokens.
Our primary modification is to the bidirectional self-attention. We first replace it with block-causal sliding window attention (SWA), using window size of two chunks of video frames. We then integrate our LaCT into the same layer. We initialize learnable fast weights for LaCT. Consistent with our language modeling experiments, SWA and our test-time training mechanism each layer: Q and K vectors are rescaled and shifted before input to the test-time training operation. The outputs of SWA and the test-time training layer are summed, with a per-head learnable scalar (from a zero-ine irroiinitallzeed linear projection) applied to the latter. We do not use Muon in the fast-weight update, as it showed no significant difference in validation loss empirically. We use an fast waith initialization of 0.001 by setting "const_ br_b in Algorithm to softplus(const. Ir bias) = 0.001. This allows small update to the fast weight in the beginning of the fine-tuning. To maintain minimal changes to the original Wan architecture, LaCT layers utilize the original RoPE from the Wan model, and we remove the SiLU activation function previously applied to queries and values.
combined within
Datasets. We fine-tune the model using internal, filtered proprietary collection of videos, each accompanied by a short text prompt generated by visual language model(Chen et al. 2024b).
Training details. Following (Esser al. 2024 Wang et al.) 2025a). use time-step shifting (scale factor 3.0) and logit-norman denoising loss weighting (m5 std=1.0). We also apply an exponential moving average with decay rate of 0.995 the model weights. Each 5-second video
28

---

Published as a conference paper at ICLR 2026
(16 FPS, 480x832 resolution) is encoded by the Wan VAE into a [21,60,104] latent representation. Denoising is performed autoregressively in chunks of three latent frames (4680 visual tokens each). We employ teacher-foreing with an interleaved noisy-clean chunk sequence (see Section (5.3.
1.3B Parameter Model:
For initial training on 5-second videos, noisy chunks are repeated twice. This results in sequences of 60 latent frames (14 noisy. 6 clean chunks), totaling 93,600 tokens. We finetune the model with a batch size of 64 for 5000 iterations. The base learning rate is set to 2 x 10 with linear warm-up of 1000 iterations and linear decay. Subsequently, the model is fine-tuned on 10-second video clips for 1,000 iterations. These clips correspond to 42 latent frames for the clean video portion, forming interleaved sequence of 81 latent frames (approximately 126K tokens including noisy chunks). Training for 5-second videos takes 20 seconds per iteration on 64 A100 80GB SXM GPUs (or ~0 seconds on 64 H100 80GB SXM GPUs).
To manage GPU memory usage, noisy chunks are not repeated in this setting. We train the model on five-second videos with batch size of 64 for 5000 iterations with a base learning rate of 5 x 10-6 and use a sequence parallel size of 2 GPUs. This phase takes ~80 seconds per iteration on 64 A100 GPUs. The model then fine-tuned on 8. second video clips (36 latent frames for the clean portion) for an additional 600 iterations, using sequence parallelism (4 GPUs). This fine-tuning takes ~80 seconds per iteration on 64 H100 GPUs.
14B Parameter Model:
Baselines. We compare our method with three baselines: sliding window attention, Mamba2 (Dao &
Gu 2024) with sliding window attention, and full block-wise causal attention, where the window
attention the baselines implemented the same as in our model. For the Mamba2 layer. we follow (Wang et al., 2024) to apply the original projected k, q. and v as B, C, and x, respectively. The Mamba2's state is updated boken-by-token, we revert the state after processing a noise chunk of frames to ensure only clean chunk state updates propagate. The full block-wise causal attention baseline is implemented with Flex Attention (He et al. 2024).
Evaluation. We compute validation loss for all models on a collection of 2,000 videos after 5,000 training iterations by computing the denoising loss at five timesteps (550, 650, 750, 850, 950). The denoising losses are measured with respect to each video frame chunk and plotted in Figure 6 Figure 6 a) compares validation loss (up to 5s videos) of LaCT against SWA, Mamba2 with SWA, and full block-wise causal attention. Our LaCT is comparable to full attention and outperforms other baselines. Figure(]]b) shows comparisons with the SWA baseline using different window sizes for both our method and the baseline (up to 5s videos). The default window covers six latent frames (two chunks). An additional experiment used a four-frame window. Results indicate that increasing window size from four to six frames improves validation loss, but this improvement is smaller than that achieved by incorporating LaCT. Figure ((c) presents validation loss (up to 10s videos) after fine-tuning LaCT and the SWA baseline on 10-second videos for 1,000 iterations.
Generated video samples from our model are provided in an appended folder. Each video chunk is sampled following the original Wan method, using a UniPC (Zhao et al., 2023) sampler with 50 steps. classifier- free guidance of 5.0, and timestep shift of 3.0.
Results on VBench. VBench (Huang et al. 2024) offers a comprehensive suite of metries for video generation. Using 942 prompts from the standard set, we generate two videos per prompt for our 1.3B model and one video per prompt for our 14B model. Summary scores are reported in Table I0 with detailed results across 16 dimensions in Table [2]
Because the standard prompts are short, the Wan team also provides augmented prompts with richer descriptions. We evaluate on this augmented set well, reporting summary results in Table I and detailed scores in [3]
Overall, we observe two key trends: (1) autoregressive models do not yet match their full-sequence diffusion counterparts, and (2) among autoregressive approaches, our test-time training method consistently outperforms sliding-window
window and full-attention variants in temporal quality, semantic score, and total score, while maintaining parity on other metrics. The performance gap is even larger under the standard prompt set.
29

---

Published as a conference paper at ICLR 2026
Table 10: Summarized VBench scores using the standard prompt set. (higher is better).
Temporal Quality
Frame Quality
Text Alignment
Quality Score
Semantic Score
Type
Method
Total Score
Full Seq
Original Wan (1.3B)
91.94%
64.27%
25.06%
83.59%
67.13%
80.30%
AR
Ours (1.3B)
92.50%
63.60%
24.75%
82.51%
62.17%
78.44%
AR
Transformer SWA (1,3B)
91.14%
63.03%
24.77%
81.62%
60.09%
77.31%
AR
Transformer (1.3B)
92.32%
62.19%
24.78%
82.28%
60.40%
77.90%
Full Soq
Original Wan (14B)
93.08%
64.52%
25.79%
84.31 %
69.53%
81.35%
Ours (14B)
92.79%
63.30%
25.67%
82.89%
65.86%
79.49%
Table Summarized VBench scores using the WAN-augmented prompt set. (higher is better).
Temporal Quality
Frame Quality
Text Alignment
Quality Score
Semantic Score
Type
Method
Total Score
Seq
Original Wan (1.3B)
92.20%
66.52%
26.16%
84.48%
75.98%
82.78%
AR
Ours (1.3B)
93.97%
66.58%
26.04%
84.41%
74.25%
82.38%
AR
Transformer SWA (1.3B)
93.20%
66.53%
26.09%
84.48%
73.00%
82.18%
AR
Transformer (1.3B)
93.62%
66.37%
26.00%
84.57%
72.95%
82.25%
Full Seq
Original Wan (14B)
93.18%
66.85%
26.27%
85.13%
77.07%
83.52%
AR
Ours (14B)
94.60%
66.17% 17%
26.16%
84.99%
72.85%
82.57%
E.4
EXPERIMENT DETAILS FIGURE 1
Fig. Mc) shows results for training a 760M 760M-parameter LaCT language model. We employ SwiGLU MLP fast weight with the Muon test-time optimizer. To scale the fast weight size, we fix the intermediate dimension of the fast- MLP to match the head dimension, then increase the head dimension from 128 to 1536 while proportionally decreasing the number of heads to maintain a constant total model dimension. Validation loss is computed on the last 2,048 tokens of each 22.768- token sequence, averaged over 76K sequences from the Book3 dataset.
Fig. d) uses the object-level novel view synthesis experiment. All models consist of 14 stacked blocks with fixed model dimension of 768 and were trained for 167 billion tokens. Training time (wall-clock) is measured on an A100-40GB SXM GPU.
F DETAILS FOR MAMBA BASELINES
Mamba is an efficient model architecture, it is logically similar to linear TTT taking per-token linear update rule the fast weight (i.e., state in Mamba's context). Thus it serves as a baseline to understand the gap between the chunk-wise update and per-wise update in Fig. 8 and Fig. 6 In this section, we detailed the experimental setup.
We take the official Mamba-2 implementatior" in all our experiment. The original Mamba-2 has multiple components, and we largely simplify its implementation to keep a measurable architecture while still maintaining the performance. In detail, our Mamba-2's formulation in experiment is:
X, 5 Linear(u)
5 = softplus(8 + 8ynit)
H, = (op(-6;)H,-1 + 8, X,
y/ C,H,
(22)
where X, B, C is of shape (L, d), and is of shape (L, 1). H, matrix state of shape (d. d).
30

---

Published as a conference paper at ICLR 2026
Table 12: VBench subscores using standard text prompt set. (higher is better).
Suhject
Background
Temporal Flickering
Motion Smoothness
Dynamic Degree
Aesthetic Quality
Imaging Quality
Object
Models
Type
Comsistecyy
Consistency
Class
Full
Original Wan
95.36%
96.6%%
99. 42%
98. 24%
63.19%
61.10%
67.44%
76.9%%
AR
Ours (1.3B)
93.94%
94.47%
96. 98%
98. 20%
74.31%
89.81%
67. 47%
69.66%
AR
Transformer SWA (1.3B)
92.23%
91 49%
97. 62%
98. 11%
74.31%
58.59 58.59%
67.46%
66. 10%
AR
Transformer (1.3B)
92.90%
93.68%
98. 08.05%
98.00%
77.08%
58.36%
66.02%
67.76%
Seq
Original Wan (14B)
95.04%
96.89
99. 2.%
98.44%
70.83%
62.27%
66.78%
81.96%
AR
Ours (14B)
93.57%
94.69%
97.65%
98.31%
76.39%
59.83%
66 68%
73.26%
Multiple Objects
Human Action
Spatial Relationship
Appearance Style
Temporal Style
Overall
Type
Models
Color
Scene
Consistency
Full
Original Wan (1.38)
60.86 00.8
77.50%
91.91%
72.69%
19.44%
20 24%2
23.62%
23.64%
AR
Ouns (1.3B)
44.74%
73.50%
81.39%
66.01%
21.69%
20.41%
23.12%
22.92%
AR
Transformer SWA (I.3B)
35.75%
75.50%
81.49%
58.91%
21.18%
20.26%
23.14%
22.86%
Transformer (1.3B)
70.00%
85.17%
61.99%
23.06%
Seq
Original Wan (14B)
77.00%
90.24%
72.96%
27.33%
21.36%
23.38%
AR
Ours (14B)
49.85%
85.00%
86.18%
63.62%
22.24%
21.86%
23.28%
24.51%
Table 13: VBench subscores using the WAN-augmented prompt set. (higher is better).
Subject Consistency
Background
Temporal Flickering
Motion Smoothness
Dynamie Degree
Aesthetic Quality
Imaging Quality
Object
Type
Method
Consisteney
Class
Full Seq
Original Wan (1.3B)
94.87%
96.64%
99.31%
98.55%
65.28%
65.71%
67.34%
87.22%
AR
(1.3B)
92.87%
94.77%
97.49%
98. 199
86.11%
65.36%
67.81%
AR
Transformer SWA (1.3B)
91.69%
94.63%
98.75%
98.22%
83.33%
64.81%
68.26%
88.37%
AR
Transformer (1.3B)
92.26%
94.68%
98.39%
98.35%
84.72%
64.72%
68.02%
82.04%
Full Seq
Original Wan (14B)
95.05%
97.13%
99.27%
98.55%
70.83%
66.18%
67.52%
91.38%
AR
Ours (14B)
92.39%
95.35%
98.33%
98.21%
90.28%
64.18%
68. 68.17%
85.05%
Multiple Objects
Human Action
Spatial Relationship
Appearance Style
Temporal Style
Overall
Type
Methud
Color
Scene
Consistency
Full soq
Original Wan (1.3B)
73.21%
95.00%
90.59%
74.53%
45.09%
21.58%
23.02%
25.36%
AR
Ours (1.38)
63.03%
92.00%
89. 8.%3%
70.95%
47.09%
21 35%
22.70%
25.45%
AR
Transformer SWA (1.3B)
62.65 65%
95 95.00%
86, 24%
66, 65%75%
41.79%
21.31%
22.72%
25.56%
AR
Transformer (1.3R)
66.23%
96.00%
82.58%
72.45%
40.70%
21.07%
22.93%
25.87%:
Full Seq
Original Wan (14B)
72.87%
94.00%
90.45%
75.29%
49.85%
22.66%
21.83%
25.37%
Ours (14B)
93.00%
83.47%
68.039
21.94%
Transferring the above formula to standard linear-attention TTT/ / DeltaNet notations, it is equivalent to:
V, K, Q, Ir = Lincar(input)
Ir softplus(Ir + brinst)
W. = exp(-Ir.)W.. br br.KVV
O = Q.We
(23)
We will denote the above equations as O = Mamba (input)..
We use the multi- head design in Transformer's 's multi- attention. Multiple independent Mamba- 2 layer are run in parallel and their outputs are concatenated. Suppose the number of heads is nh, the formula is:
0 Mamba (impat(input) (input)
0 10 ...., Onnh
(24)
where each M a Mamba with its own parameters. Mamba-2's terminology, this design is equivalent to setting the number of "groups' to be the same the number of heads.
For the novel view synthesis task, we take a bidirectional Mamba over the input image tokens. In detail, we take two independent multi-heed Mamba with one reading from left to right and the other reading from right to left. The bidirectional model builds a better connection among input tokens and also doubles the state size. We similar 'apply" operation as in LaCT that only updates
31

---

Published as a conference paper at ICLR 2026
the state for input tokens, and the state is static for the target tokens. We also tested with 'update' for the target image tokens, but it empirically leads to worse results. We use a head dimension of 192 and 8 heads. The overall state size, 8 (num heads) x 192 (head dim) x 2 (bidir), matches LaCT with a standard large-chunk large-weight linear attention of dimension 768 (768 input dim x 768 intermediate dim in Fig. 8 We take Inint= 4.6, which corresponds
0.01 initialized
learning rate (i.c., softplus( 100--440 4.6) = 0.01).
For the autoregressive video diffusion task, apply unidirectional Mamba over the flattened video tokens. As mentioned in Sec Sec we follow (Wang el al.) 2024) to inherit the Wan's self-attention projected 9, and v as B. C, and x in the Mamba layer, respectively. Unlike in the NVS task, each token will 'update" the state, which will be 'applied' to the current output and future tokens. Our Mamba uses 12 heads, each of dimension 128, matching the original multi-head self-attention in Wan. The overall state size is 12 (num heads) x 1282 (head dim). We take Iime== 4.6, which corresponds to a 0.01 initialized learning rate.
G DETAILS OF LACT CONTEXT PARALLELISM IMPLEMENTATION
Context Parallelism(CP) partitions the input sequence along its sequence length dimension and distributed the shards across multiple devices for parallel computing. The feed-forward layer and window attentions are local operations thus support CP naively. Our large-chunk Test-Time Training (TTT) approach facilitates CP by sharding tokens within each large chunk.
Within our large-chunk TTT mechanism, the per-token apply operation naively supports CP due to its independent nature. The update allows CP by shading tokens within chunk over multiple devices. This CP can be easily implemented by adding a few lines of distributed all-reduce-sum after computing the local fast weight gradients each device, logically the same the Distributed Data Parallellism. Note that the distributed all-reduce-sum is a differentiable operator and its backward is all-reduce-sum over the gradient, thus the network can be trained end-to-end. Algorithm presents the pseudocode detailing this intra-chunk context parallelism specifically for the large-chunk TTT update operation. We employed this parallelism in our view synthesis experiments, handling maximum chunk sizes exceeding half a million tokens and maximum sequence lengths over one million tokens during training.
H DETAILS OF LACT TENSOR PARALLELISM IMPLEMENTATION
Beyond Context Parallelism, our large-chunk Teme Training (TTT) mechanism also supports Tensor Parallelism (TP). This is primarily achieved by sharding the TTT heads across multiple devices, a strategy similar to that employed methods like DeepSpeed Ulysses (Jacobs et al. 2023).
Specifically, while static feed-forward layers in the model might process inputs sharded along the sequence dimension (Context Parallelism), for the TTT operations within our LaCT layer, the data undergoes gather-then then-scatter transformation. Input tensors (Q, K. V. and learning rates for TTT) that are initially sharded by sequence length are first gathered along the sequence dimension to reconstruct the full sequence context on each device within the tensor-parallel group. Then, these full-sequence tensors are scattered along the head dimension. As a result, each device processes the complete sequence but operates on only its assigned subset of TTT heads during the TTT update and apply iterations. The reverse transformation (gather heads, scatter sequence) is applied to the output provides pseudocode detailing this tensor parallelism implementation, omitting minor details like padding. While this gather- gather-then- then-scatter method effectively enables head- sharded tensor parallelism. more sophisticated communication strategies (Fang & Zhao 2024 Jacobs et al. 2023) could potentially be employed to further optimize communication overhead.
of TTT operation. Algorithm 4
We utilized this tensor parallelism strategy in our autoregressive video generation experiments, sharding, for example, four TTT heads across four local GPUs. This enabled us to train 14-billion- parameter diffusion models with sequence lengths exceeding 100K tokens.
32

---

Published as a conference paper at ICLR 2026
Algorithm 3 Large Chunk Test-Time Training Layer with Context Parallel Sharded inside chunk Pseudocode
def update weight k, Vy Le, cp_group, =ue_m))
Past-weight update for a SwiGLU MLP usiny a context-paral shunk.
Args:
fast_weight tuple (wl, w², w3) with shapea: wl, w3: (b, d, dh1y w2: [b, dh, d]
V xey value tenaor shape [b, 1,
Irt : per-token learaning rates of shape [b, 1, 3] => (1rl, 1r2, 1r3)
epagroup process group metadata for cortext parsllalism use_mwon weather to apply Muon to orthogenalize the update
The Ingut tenaora k, v, Ir are assumed to be already partitioned (sharded) along the dimension over maltiple devices. 1 represents the local sharded sequence length on each device.
total effective chunk size precessed 1 op_aroup.size.
# Forward with k;
gate bete_beffore_act metmul (s, wl)
dh] [b, 1, d) [b, dh] hidden before gate natmul ih, w3) [br chl [b, 1, d] x [b, d, dh]
hidden - beffrgaatebbfore_eeee
sidden_beforue
# Backward:
chi dden matmul w2.transposa -211 [b, 1, dh] [b, d] [b, 6, dh]
thiidee__bffor__eee dqate dhidden silu silu_backprop Idgate, gate_bafore_act) agate_before_act
dhidden siluigate_before_aac
# Compute gradients: w2.grad -matmul (hidden, craiddennrennssseeee- -2), v 1r2) [b, dh, d] [b, de, 11 x
1,
[b, d, dh) [b, d, 1 [b, 1, dh)
w1 w1.grad -matmul (1X 1ell Lransposee (-1, -2),
w3.grad matmuu
1r31 .transpeepp (-1, -2), dhidden
[Standard forward pass and local backward gradient computations are performed above,
# resulting in local u.grad for each device.]
BEGIN CONTEXT PARALLELISM SPECIFIC MODIFICATION: Global Gradient Aggrogation
The following AllReduce operation is the key step insroduced for context
parallelism. operations before this point compute local gradients; operations affeer this point use the globally aggregated gradienss.
fast_weight:
w.grad distri iustrribute_dri all (etrredd_l_rrrdeeeeeeeee.. reduce cp co_aroup, op="SUM*)
END CONTEXT PARALLELISM SPECIFIC MCDIFICATION.
Subsequent operations (Muan, weight updates) now use the globally bly surmed w.grad. same
formulas for non-parallel version, but they act upon these aggregated gradients. www.t1ttttiiiit Eerfetiiii
aubsequent operations
#
# Weight update
uss_muon:
for m in Eest Fest_weight:
w.grad sarrpove__via_nestoraahh ww.d)dd)
for w in fast weightt
(W w.grad)
w.grad) norm (dim=l) w.norm (dim=1)
return west_weight
LIMITATION
One limitation of our method is the absence of rotation invariance. Unlike softmax attention and linear attention, which remain invariant under uniform rotations of queries and keys (a property leveraged by relative positional encodings such RoPESSu et al. (2023)), our SwiGLU and Linear Fast Weight components do not exhibit this property. The practical implications of this absence remain underexplored.
On the language modeling task, some key aspects are not explored due to computation limitation. These aspects include the reasoning capacity of our LaCT model and also the scalability regarding
33

---

Published as a conference paper at ICLR 2026
Algorithm 4 Large Chunk Test-Time Training Layer with Tensor Parallelism by sharding heads Pseudocode
def gather_ (x, gather_dim, ceatter_ding
Gathers tensur x along gasher_ dim across process process then scatters the zeault along scatter_dis Example: Transform [B, MLfull, I_local, D. with gather_dim-2, scatter_dim-1 to (R, N_local, 1_full, D] on each device.
to each device lccally.
***
x all_gather (x, gather_dim, (rocusa_group)
# Calculate slicing indices for the scatter operation
local__ane group_size = process_grouy grrpiggrr.p..a... scatter_stride (acaatiesceerr start_idk local_rank * acatter_stride
process_grrup..tt
// group_
end_idx (local_rank
scatter_stride
# slice the tensor to get the local shard for the current device x slice_tensor (x, scatter_dim, start_idx, end_idx) reture x
1*A9ttA1 KultiRead LeCT Layer with Tensor Parallulism (sharding TTT heads) E*tAAE***
Iaput:
x: input sequence sharded by sequeunce length (CP), Shape [b, 1, d], b is the batch model dimensicn.
dim, local sequence length, faat_weight tuple of aharded initial faat weights, sharede among heada. (wl, w2, ; wl, of shape [nh, d, dh], w2 shape [nh, dh, d]. nh: number of local heads.
gkv (LinearrKK (x)) t [b, 1, + 3]
gkv zearrange (gkv, 'b1 (ah hd) -> b nh 1 hd', ahaaua ssplit (3, dimm-1)
9, q n-rmm(--1,, k.norm (-1)
iI softplus (LineariR(x) const_ir bias) F [b, 1, 3 aum (uadhheads)
ir rearrange "b (nh -> an 3', heads!
BEGIN
PARALLELISM SPECIFIC TRANSFORVATION
Gather along Sequence Lengch (dim 21, then Scatter along Head Dimension Idim 1), Transforms [b, nh fulf,l, lose,, X] [b, nt_lasal, 1ffll, X] Each device new has the fall sequence for subset of heads.
k, 1r (1ambda x:
gather_dim=2, scatter_dim=l), (a, k,
END TENSOR PARALLESISM SPECIFIC TRANSFORMACION
[b, nh local, Lfull, X]
o local heads ... Placeholder for actual TTT computation on shardad
0 o_local_heee heads (MsNo-mic_oocaa_heedd)s + per-heer head norm
BEGIN TENSOR PARALCELISM SPECIFIC REVERSE CRANGPORMATION
Gathez along Baad Dimension Idim 1), then Scatter along Sequence Dimension (dim 2)
Transforms [b, localc full, X] -> [b, ff_f,,l, 1_local, X] reconstructa the (o_ END TENSOR PARallelis ism SPSCIFIC REVERSE TRANSFORMATIGN 44444444*444 (44.444444444444444444444 444444444444448444444444444 rearrange (o, "b nh hd -> b inh hd)", nh-num_headas lanearcut to)
head dimension
shards sequence back.
fither_dim-l acatter_dim-2)
return
the parameter size. Previous papers showed that main weakness of the state-based model (where LaCT belongs to) is its reasoning ability. However, the reasoning ability is only gained with certain amount of training compute thus it is beyond our budget.
Lastly, for the autoregressive video diffuuioon, it is hard to find a reliable and distinguishable metric to measure the model's scalability. It is in contrast to the language modeling with perplexity (i.e., log likelihood loss) and the novel-view synthesis with PSNR. We show the validation loss in paper and it a common choice evaluating the scalability of video generation. This is a general problem for the video generation evaluation and is not specific
paper.
34

---

Published as a conference paper at ICLR 2026
J LLM USAGE
We employ large language models (LLMs) in three main scenarios. During the experimental stage, LLMs to generate portions of tedious code based on clear instructions and guidelines to acclerate experimental iteration. During the writing stage, we leverage LLMs to polish the text and identify grammatical issues. Finally, we also use LLMs generate Python code for plotting figures with matplotlib.
35


Reviews:
[{"review_id": "uNfIbOfutx", "reviewer": "Reviewer_gUWh", "summary": "In this work the author presents LaCT, Large Chunk Test Time Training. This work tries to solve the problem by TTT frameworks where they operate under extreme low FLOPSs (<5%) utilization on the modern GPUs due to low mini-batches (16 or 64 tokens). This work produces large chunk updates consisting 2k-1M tokens across different tasks. The authors validates their claim on three different tasks a) novel view synthesis (NVS), b) language modeling c) auto-regressive video diffusion.", "questions": "Most of the major weaknesses presented are questions. If the authors can clarify these questions, then I will gladly increase my rating.\n\nWhy the related work and many other ablations presented in the appendix?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. This particular method enables large chunk updates enabling higher FLOPs utilization. \n2. The authors claims to implement the LaCT layer simply using PyTorch without custom CUDA kernels.\n3. The authors present metrics highlighting the throughput.\n4. The empirical results shows the validity of the method on various tasks requiring long context updates.", "weaknesses": "1. One of the major concerns for me is how is the work different from the Fast Weight Programmer and Expressive RNNs. In both the cases the objective is to store the short term memory as the fast weights, next the hidden state in Expressive RNNs [A] basically are used as the fast weights similar to this work. Next the synaptic weight modifications in Expressive RNNs uses a network that uses gradient descent while the update in this particular work is by calculating a gradient $ g $  from a self-supervised loss $ \\mathcal{L} $ over the entire chunk. Thus both this work uses update by a learning algorithm Although this particular work uses MUON as an optimizer.\n2.  Can the authors clarify that this method is not compared against the Mamba based LM models, as well as the use of the other architectures that handles long context. The authors also might try to use memory-augmented attentions.\n3. The positional embedding interaction with the fast weights seems ambiguous, authors need to clarify this for the fast weight layer.\n4. It is not very clear how does the authors address the alignment problem, of TTT. The self-supervised loss is not perfectly aligned with the main task, the gradient can be harmful.\n\nMinor Concerns:\n1. The paper needs improvement in the writing, it is quite difficult to follow the flow of the paper. \n2. The paper should have compared against other methods such as the current state  RNNs etc. \n3. Few of the abbreviations are assumed to be known, the authors should explicitly state them eg: SMs in Section 2.2.\n\n[A]. https://arxiv.org/pdf/2407.04620", "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": "qK4A22quHF", "reviewer": "Reviewer_NZ1F", "summary": "This paper revisits Test-Time Training (TTT) for long-context modeling and introduces Large Chunk Test-Time Training (LaCT), which updates fast weights over large token chunks instead of small mini-batches. The method aims to improve GPU utilization, support scaling of non-linear fast-weight networks, and integrate window attention for local dependency modeling together with the Muon optimizer for stable fast-weight updates. The authors evaluate LaCT on three modalities: novel view synthesis, language modeling, and autoregressive video diffusion.", "questions": "1. The authors report 40B-token training for the 760M model and 60B tokens for the 3B model. Have the authors considered scaling experiments with additional training data (~500B-1T) to better evaluate the method’s true language modeling capability? Given the limited token counts, the model’s performance on ARC-c is nearly at random-guess level, which raises questions about the significance of these results.\n\n2. Additionally, given the training setup and results in Table 9, would the 760M model outperform the 3B model if both were trained on the same number of tokens? The current configuration, where the larger model is exposed to only slightly more data, raises the question of whether the observed results reflect true scaling behavior or simply data inefficiency.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "1. The proposed approach is general and can be readily adapted to a variety of domains and applications.\n\n2. The evaluation across three domains validates its scalability, with the view synthesis task demonstrating near-transformer performance at much higher efficiency.", "weaknesses": "There are some minor weaknesses that do not affect my rating but do affect readability. Some parts of the paper are hard to follow, especially for readers unfamiliar with TTT, as several key details and related work are placed in the appendix. A few figures are incorrectly rendered with shading artifacts, which should be fixed.", "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": "BrkqK3VnlM", "reviewer": "Reviewer_vmAu", "summary": "This paper addresses a critical inefficiency in existing Test-Time Training (TTT) methods when applied to long-sequence data. The authors identify that conventional TTT approaches, which utilize very small online mini-batch sizes for frequent weight updates, suffer from extremely low FLOPs utilization on modern GPUs. To overcome this, the paper proposes a novel paradigm called Large Chunk Test-Time Training (LaCT). The core idea is to move in the opposite direction of prior work by using extremely large chunks of tokens (from 2K to 1M) as the basic unit for updating the model's \"fast weights\". The authors claim that this approach: (1) Improves hardware utilization by orders of magnitude, leading to significant computational efficiency gains. (2) Enables the effective scaling of large, non-linear fast weights, thereby substantially increasing the model's state capacity. (3) Can be implemented with just a few dozen lines of native PyTorch code, obviating the need for cumbersome and error-prone custom kernel implementations and thus democratizing research in this area. (4) Facilitates the easy integration of more sophisticated test-time optimizers, such as Muon.\nThe authors embed LaCT within a hybrid architecture that also uses windowed attention to capture local dependencies. The effectiveness of this framework is demonstrated on three diverse and challenging tasks: novel view synthesis, language modeling, and autoregressive video diffusion.", "questions": "1. Did the authors observe any specific failure cases in the experiments that could be attributed to the lack of strict causality within a chunk?\n\n2. How sensitive is LaCT's performance to the chunk size? In practice, how would you advise balancing an increase in chunk size (for more compute intensity) against an increase in state size (for more model capacity and memory)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper correctly identifies a critical bottleneck in existing TTT methods when applied to long-sequence data, and proposes the LaCT using massive chunks for updates instead of tiny ones.\n\n2. This paper shows thorough and rigorous Experiments. (1) Validation on novel view synthesis (image sets), language modeling (1D sequences), and video generation (high-dimensional sequences) powerfully demonstrates the method's generality. (2) The scale of the experiments, such as handling context lengths of over 1M tokens and scaling to a 14B parameter video model, is remarkable and convincingly showcases the method's scalability. (3) The appendix contains ablation studies that systematically analyze the impact of key design choices, including state size, optimizer choice, linear vs. non-linear fast weights, and chunk-wise vs. per-token recurrence. This analysis greatly strengthens the credibility of the paper's conclusions.\n\n3. The paper is well-organized and clearly written. The figures are highly effective to provide the summary of LaCT's advantages, the basic framework and the design choices.", "weaknesses": "1. LaCT fundamentally treats tokens within a large chunk as an unordered set. While the use of windowed attention mitigates this by handling local structure, the inherent limitations of this design choice are not fully explored. For tasks that are extremely sensitive to fine-grained causal ordering, LaCT might struggle. \n2. According to the experimental results in the appendix, LaCT's performance on the VBench benchmark for autoregressive video generation does not show a significant advantage over other AR methods.", "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": "mINAcCWUXU", "reviewer": "Reviewer_YxH9", "summary": "This paper proposes Large-Chunk Test-Time Training (LaCT), a novel test-time training (TTT) method that updates “fast weights” in large chunks rather than small batches. TTT dramatically improves GPU utilization and scalability of nonlinear state size without requiring custom CUDA kernels.  \n\nLaCT integrates windowed attention for local structure, TTT for long-range context, and is flexibly compatible with N-dimensional data. The method is validated on novel view synthesis, language modeling, and autoregressive video diffusion.", "questions": "1. How does LaCT compare against previous TTT methods, given its much higher propensity for scalability?\n2. Why was LaCT trained on proprietary data when any video diffusion dataset could have sufficed? \n\nPlease address question in the Weaknesses section as well. Overall while the paper is quite dense, the method is interesting, applicable, and thoroughly ablated and tested.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- **Extremely thorough experimentation**: LaCT is tested on three different downstream tasks across two different modalities (text, image/video).  Throughput, scaling, and efficiency experiments are also provided to exhibit the feasibility of LaCT.\n- **Clear writing and strong figures**: Lots of preliminaries and related works are provided, as well as a clear linear flow of the paper introducing the research problem and solution (poor TTT GPU utilization caused by small-chunk updated).\n- **Scalability and practicality**: Runs efficiently in PyTorch with strong FLOP utilization, making it generally adaptable - especially given the lack of  reliance on custom CUDA kernels.", "weaknesses": "- **Core results deferred to appendix**: Most quantitative comparisons and ablations are in the appendix, leaving the main paper heavy on discussion. For example, the experiments section provides very surface-level details regarding datasets, training details, etc. which are pushed down to the appendix. This makes the actual quantitative impact of LaCT on downstream datasets less apparent. For example, Tables 3 and 4 in the appendix highlight the near identical performance between full attention and LaCT, however without any discussion highlighting that LaCT is more efficient, the takeaway is unclear.  \n- **Sparse dataset/training details for autoregressive video diffusion**: The authors very briefly mention using a an internal, proprietary video dataset with minimal transparency, limiting reproducibility.  \n- **Limited comparison to alternative TTT methods**: Comparisons against other TTT methods (e.g., Titans, Atlas, which are cited in the main paper under Behrouz) are not provided. While they may be much too inefficient to compare to the scale and downstream tasks of LaCT, explicit quantitative results highlighting this point would be useful.", "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": "hsgyq8g6F5", "reviewer": "Reviewer_21cR", "summary": "The paper addresses the problem of low FLOP utilization with TTT layers. They proposes Large-Chunk Test-Time Training (LaCT): updating fast weights over very large chunks compared to prior work to reach high FLOP utilization. As updates of fast weights treats token in the current chunk as sets, TTT layers are paired with sliding window attention to recover locality. The method is validated on novel view synthesis (NVS, up to ~1M tokens), language modeling (per-position loss & RULER/S-NIAH retrieval), and autoregressive video diffusion showing competitive performance with linear time models and Transformers.", "questions": "1. Is it possible to ablate the effect of the chunk size on performance?\n2. Can you clarify the end-to-end complexity where the chunk size equals the full sequence length?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The suggested solution of combining large chunks with SWA is effective, simple and easy to implement.\n2. Evaluation is very robust spanning multiple modalities and tasks, comparing to relevant baselines.", "weaknesses": "1. Combining linear time sequence models with attention is not new and has been shown to improve the performance, see e.g., [1], although not motivated from low throughput.\n2. The method strongly relies on large chunk sizes yet the effect of chunk sizes on performance (while maintaining high throughput) is not clear from the results.\n\n[1] An Empirical Study of Mamba-based Language Models", "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)