Test-Time Training Done Right

## Abstract

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 https://tianyuanzhang.com/projects/ttt-done-right/ .

## 1 Introduction

The demand for handling long contexts is rapidly growing. While softmax attention vaswani2017attention has become the de facto solution for modeling various types of data, its computational cost grows quadratically with sequence length, motivating extensive research into more
efficient long-context modeling.

Recently, Test-Time Training (TTT) sun2024learning has emerged as a promising approach for efficient sub-quadratic sequence modeling. TTT extends the concept of recurrent states in RNNs to a small, online-adapted sub-network.
The parameters of this sub-network also referred to as fast weight schlag2021linear, as they are rapidly adapted online via self-supervised objectives to memorize in-context information.
Numerous recent studies wang2025testtimeregressionunifyingframework; behrouz2024titans; behrouz2025itsconnectedjourneytesttime; karami2025lattice 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 to extremely low hardware utilization in their TTT layers (often below 5% peak FLOPS on modern GPUs).
This inefficiency is because of the usage of
small mini-batch 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-efficient
implementation, especially when using large, nonlinear fast weights, making it difficult to achieve non-trivial (above 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 to 1M tokens) as the basic unit to 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 A100s) with just a few dozen lines of pure PyTorch code (see Appendix [A.1](https://arxiv.org/html/2505.23884v1#A1.SS1)).
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 jordan2024muon.

Furthermore, LaCT’s large-chunk design is also natural to model diverse N-dimensional data as 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
at a resolution of $960\!\times\!536$ leading to a maximum of $1$M tokens, and outperforms 3D Gaussian Splatting kerbl20233d in terms of rendering quality under such input scale.
- •
Language Modeling.
Our model achieves competitive performance compared to SoTA methods such as DeltaNet yang2024parallelizing, even though a chunk structure is not explicitly present in language data.
- •
Autoregressive Video Diffusion.
We adapt a 14-billion-parameter bidirectional video diffusion transformer into an autoregressive model by incorporating LaCT with sliding window attention. This adapted model generates consistent videos up to 56,000 visual tokens.

To summarize, our approach establishes an efficient, scalable, and highly performant framework for long
sequence modeling across diverse modalities.
By removing the dependency on low-level, hardware-specific 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.

## 2 Preliminary

### 2.1 Test-Time Training

Consider a one-dimensional sequence of $N$ tokens
$\mathbf{x}=[x_{1},x_{2},\dots,x_{N}]$, where each token $x_{i}\in\mathbb{R}^{d}$.
Following attention formulation,
each input tokens $x_{i}$ is projected into query ($q_{i}$), key ($k_{i}$), and value ($v_{i}$) vectors.
For clarity, we assume all these vectors $q_{i},k_{i},v_{i}\in\mathbb{R}^{d}$.

Test-Time Training (TTT) sun2024learning introduces a neural network with rapidly adaptable weights—called fast weights schlag2021linear—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 a neural network:
$f_{W}(\cdot):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ parameterized by the fast weights $W$, and it involves two primary operations:

$$ $\textbf{Update operation:}\quad W\leftarrow W-\eta\nabla_{W}\mathcal{L}\big{(} f_{W}(k),v\big{)}$ (1) $$

where $\mathcal{L}(\cdot,\cdot)$ is a loss function between the transformed key $f_{W}(k)$ and the value $v$, commonly Mean Squared Error, designed to encourage the network to associate keys with corresponding values.
$\eta$ is the learning rate.
Intuitively, this learning objective is to encode the KV cache into a neural memory with fixed state size as *accurate* as possible wang2025testtimeregressionunifyingframework.

$$ $\textbf{Apply operation:}\quad\quad\quad o=f_{W}(q),\quad$ (2) $$

where the updated fast weights $W$ are used to compute the output vector $o$ given the query $q$.
The per-token TTT layer iteratively perform the update and apply operations on each token $x_{i}$ in sequence.

### 2.2 Challenges in Efficient Implementation

Frequent online update of fast weights is inefficient due to memory bandwidth limitations. Consequently, previous works sun2023retentive; gu2023mamba; yang2024gated; qinvarious; yangparallelizing often employ customized kernels that keep fast weights in SRAM across updates to reduce memory load.
However, this strategy typically requires fast weights to evolve mostly independently within SMs to reduce communications, which is not valid for large nonlinear states (e.g., the nonlinear SwiGLU fast weight in Sect. [3.1](https://arxiv.org/html/2505.23884v1#S3.SS1) and the Muon update in Sec. [3.2](https://arxiv.org/html/2505.23884v1#S3.SS2)). Moreover, developing such kernel code is cumbersome, with far longer development cycles than native PyTorch code, hindering rapid research exploration.

Figure: Figure 1: Using larger chunk sizes significantly improves GPU utilization compared to 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 BF16 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 Sec. [C.4](https://arxiv.org/html/2505.23884v1#A3.SS4).
Refer to caption: x1.png

On the other hand, a PyTorch-based implementation, while simpler, is typically bounded by memory speed. As an illustration, consider a PyTorch implementation of simple MLP fast weight, the core of which is a matrix multiplication between fast weight (e.g., $h\times h$ matrix) and the mini-batch input ($b\times h$ where b is the chunk size).
The ideal compute-to-memory ratio is:

$$ $\displaystyle r$ $\displaystyle=\frac{2h^{2}b}{2h^{2}+4hb}=\frac{h/2}{1+\frac{h}{2b}}=\frac{b}{1 +\frac{2b}{h}}\leq\min(h/2,b).$ (3) $$

Here, $2h^{2}b$ is the FLOPs to for matrix multiplication, the denominator $2h^{2}+4hb$ is the memory workload for two input matrices and the output in BF16 (2 bytes).
Small fast weight size (e.g., $h=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 $1$M).
This allows us to achieve higher throughput (Fig. [1](https://arxiv.org/html/2505.23884v1#S2.F1)a) leading to better performance in less training wall-clock time(Fig. [1](https://arxiv.org/html/2505.23884v1#S2.F1)d).
Our design also allows the state size to be
scaled up efficiently(Fig. [1](https://arxiv.org/html/2505.23884v1#S2.F1)b), leading to significant results improvement with such scaling
(Fig [1](https://arxiv.org/html/2505.23884v1#S2.F1)c, Fig. [7](https://arxiv.org/html/2505.23884v1#S5.F7)a).
Our architecture achieves a state-to-parameter size ratio $\geq\!40\%$, which is an order of magnitude
larger than previous methods’ ratio of $0.1\%$ to $5\%$. Detailed pseudocode is provided in Appendix [1](https://arxiv.org/html/2505.23884v1#algorithm1).

Parallelism over the sequence length dimension, in addition to the batch and head dimensions, is crucial to achieve high occupancy when handling long sequences (where the batch size is often small). Linear Attention variants like Mamba gu2023mamba, Gated Linear Attention yang2024gated and DeltaNet yangparallelizing enable such parallelism by utilizing the associative property of linear recurrence. Attention vaswani2017attention; dao2023flashattention can be parallelized along the sequence length dimension using online softmax milakov2018online, a key improvement in FlashAttention-2 dao2023flashattention over FlashAttention-1 dao2022flashattention. 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 a 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 [3.4](https://arxiv.org/html/2505.23884v1#S3.SS4), with pseudocode in Appendix [3](https://arxiv.org/html/2505.23884v1#algorithm3).

## 3 LaCT Model Architecture

Figure: Figure 2: The basic diagram for a LaCT block. The large-chunk TTT layer updates the fast weight $W$ to 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. [4](https://arxiv.org/html/2505.23884v1#S4) use different chunk sizes and window attention types according to the specific data structure. Additionally, window attention and large-chunk TTT layers can be combined within the same layer by sharing the QKV and summing their outputs; this in-layer mixing is used in our language modeling and video generation experiments (see Appendix [2](https://arxiv.org/html/2505.23884v1#algorithm2) for such pseudocode).
Refer to caption: x2.png

As shown in Fig. [2](https://arxiv.org/html/2505.23884v1#S3.F2),
LaCT block consists of three types of layers: a window attention layer, a large-chunk TTT layer, and a feed-forward layer.
Each layer is equipped with residual connections he2015deep following the practice in Transformer vaswani2017attention.
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 key vectors $K$ and value $V$), and latest weight is ‘applied’ to the current query vector (Q) for computing its corresponding output.
The feed-forward layer performs channel mixing as in Transformer.
We omit several linear and normalization layers in Fig. [2](https://arxiv.org/html/2505.23884v1#S3.F2) for clarity and details are in Appendix [A.1](https://arxiv.org/html/2505.23884v1#A1.SS1).
Our framework offer great flexibility in handling diverse data types.
In this section, we
present the general designs in our approach and later describe data-specific variations
in Sec. [4](https://arxiv.org/html/2505.23884v1#S4).

### 3.1 Large-Chunk TTT Layer

Different from the per-token update in Eqn. [1](https://arxiv.org/html/2505.23884v1#S2.E1), the chunk-wise update computes the gradient of the summed loss over all keys $\{k_{i}\}$ and values $\{v_{i}\}$ 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. [3.2](https://arxiv.org/html/2505.23884v1#S3.SS2)) and amortizes the update cost.
The ‘update’ operation for the fast weight is:

$$ $\displaystyle g$ $\displaystyle=\nabla_{W}\sum_{i=1}^{b}\eta_{i}\mathcal{L}\big{(}f_{W}(k_{i}),v _{i}\big{)}$ (4) $\displaystyle W$ $\displaystyle\leftarrow\mathrm{weight\mbox{-}update}(W,g),$ (5) $$

where $b$ is the chunk size, $g$ is the gradient of the fast-weight loss function,
and $\eta_{i}$ is the learning rate of each token (usually predicted from input tokens).
The ‘apply’ operation $o_{i}=f_{W}(q_{i})$ is the same as Eqn. [2](https://arxiv.org/html/2505.23884v1#S2.E2) and
all query vectors $\{q_{i}\}$ in the chunk share the same updated fast weight $W$.

Motivated by recent
LLMs touvron2023llama, we adopt SwiGLU-MLP shazeer2020glu without bias terms as the fast-weight
network.
Our fast weights consists of three weight matrix $W=\{W_{1},W_{2},W_{3}\}$, and the network
is:

$$ $f_{W}(x)=W_{2}\left[\mathrm{SiLU}(W_{1}x)\circ(W_{3}x)\right]$ (6) $$

where $\circ$ is an elementwise multiplication.
We apply a simple dot product loss as our loss function:

$$ $\displaystyle\mathcal{L}\big{(}f_{W}(k_{i}),v_{i}\big{)}=-f_{W}(k_{i})^{\top}v _{i}$ (7) $$

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 in different
orders; this allows us to model diverse kinds of data dependencies, similar to different
attention masks in self-attention.
Figure [3](https://arxiv.org/html/2505.23884v1#S3.F3) illustrates this concept.
In Figure [3](https://arxiv.org/html/2505.23884v1#S3.F3)a, 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
a block-wise causal mask (Fig. [3](https://arxiv.org/html/2505.23884v1#S3.F3)b), where the block size corresponds to the chunk size.
Switching the order between the two operations results in the a shift in the mask (Fig. [3](https://arxiv.org/html/2505.23884v1#S3.F3)c).
This shifted mask
does not leak future information within the chunk and is important when building the full causal mask
in Language Modeling (Sec. [4.2](https://arxiv.org/html/2505.23884v1#S4.SS2)).
Moreover, only updating on a subset of chunks
and applying to all
(Figure [3](https://arxiv.org/html/2505.23884v1#S3.F3)d) is analogous to strided
block-wise causal mask.

Figure: 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 output depends on the j-th token.
Refer to caption: x3.png

### 3.2 Non-Linear Update of Fast-Weight

Fast-weight updates in TTT repeatedly accumulate gradients, and thus suffer from magnitude explosion or decayed memory.
Large-chunk TTT allows non-linear updates to improve stability and effectiveness while preserving efficiency.
For the ‘weight-update’ operation in Eqn. [5](https://arxiv.org/html/2505.23884v1#S3.E5), our vanilla implementation involves gradient descent followed by weight normalization:

$$ $\displaystyle\mathrm{weight\mbox{-}update}(W,g)$ $\displaystyle=\mathrm{L2\mbox{-}Normalize}(W-g).$ (8) $$

We have also explored a more robust nonlinear Muon jordan2024muon update rule (^1^11Muon requires weights in matrix form, and our current fast-weight function SwiGLU-MLP has three matrices as the weights (i.e., $W_{1}\mbox{,}W_{2}\mbox{,}W_{3}$ in Eqn. [6](https://arxiv.org/html/2505.23884v1#S3.E6)).) with weight normalization:

$$ $\displaystyle\mathrm{weight\mbox{-}update}(W,g)$ $\displaystyle=\mathrm{L2\mbox{-}Normalize}(W-\mathrm{Muon}(g))$ (9) $$

Fast-weight normalization.
We apply L2 weight normalization salimans2016weight to the updated fast weights along the input dimension. We do not use explicit weight-decay term as in previous methods behrouz2024titans; dao2024transformers; yang2024gated; sun2023retentive.
When the network is 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 he2015deep. In 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.
In short, let $g=USV^{T}$ be the Singular Value Decomposition(SVD) of the gradient $g$, then Muon operator approximately converts the gradient as:

$$ $\displaystyle\mathrm{Muon}(g)$ $\displaystyle\simeq UV^{T}$ (10) $$

Muon also improves the numerical stability in our setup.
For example,
the learning rate ($\eta_{i}$ in Eqn. [4](https://arxiv.org/html/2505.23884v1#S3.E4)) now only reflects the relative importance of tokens within a chunk as Muon normalizes the absolute scale. See jordan2024muon and Appendix [A](https://arxiv.org/html/2505.23884v1#A1) for analysis of its computational cost.

### 3.3 Window Attention

The large-chunk TTT layer treats data as sequences of sets because its fast weight updates inherently disregard token order and spatial locality within each chunk.
However, many data modalities—such as videos (sequences of grids), image collections (sets of grids), or text (1D sequences)—do not fully align with this set-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 to focus its fixed-size fast weight capacity on modeling non-local dependencies. This hybrid strategy is also employed in other notable works like BASED arora2024simple, GAU hua2022transformer and InifinitAttention munkhdalai2024leavecontextbehindefficient.
In summary, LaCT is a hybrid architecture with the quadratic-compute attention for local structure and linear-compute TTT for non-local context.

### 3.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 feed-forward 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.
Suppose each shard contains $s$ tokens,
the fast weight gradient of the chunk is the sum over all shard’s gradients given the linearity of the gradients:

$$ $\displaystyle g$ $\displaystyle=\nabla_{W}\sum_{j=1}^{\text{shards}}\sum_{i=1}^{\text{s}}\eta_{i }\mathcal{L}_{i}=\sum_{j=1}^{\text{shards}}\nabla_{W}\sum_{i=1}^{\text{s}}\eta _{i}\mathcal{L}_{i}$ (11) $$

This can be implemented through distributed all-reduce-sum 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 in training the novel view synthesis task (see Sec. [4.1](https://arxiv.org/html/2505.23884v1#S4.SS1))
and observe minimal throughput overheads (1% to 3%).
LaCT architecture is compatible with other parallelism strategies (e.g., data parallelism, pipeline parallelism, and tensor parallelism). See Appendix for pseudocode on implementing context parallelism(Alg. [3](https://arxiv.org/html/2505.23884v1#algorithm3)) and tensor parallelism(Alg. [4](https://arxiv.org/html/2505.23884v1#algorithm4)) for LaCT.

## 4 LaCT for N-Dimensional Data

In this section, we introduce the three tasks we address using LaCT—novel 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 [B](https://arxiv.org/html/2505.23884v1#A2).

### 4.1 Novel View Synthesis - Image Set

Novel view synthesis (NVS)mcmillan2023plenoptic; levoy2023light aims to render images of a static scene from previously unseen viewpoints. Formally,
given a set of $N$ input posed images $\{(I_{i},P_{i})\}_{i=1}^{N}$ of a static scene,
where $I_{i}\in\mathbb{R}^{H\times W\times 3}$ is an RGB image and $P_{i}$ 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 as it requires spatial compression, dense retrieval, and basic physical reasoning.
Secondly, NVS can be formulated as a non-generative 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, we 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. [3](https://arxiv.org/html/2505.23884v1#S3).
Both the posed input images and poses of the target novel views are tokenized by patchify and linear layers, following LVSM jin2024lvsm.
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. [3](https://arxiv.org/html/2505.23884v1#S3.F3)d), which updates the fast weight using all input image tokens, and applies to both the input and target tokens. The update step resembles a prefill stage, while the apply operation resembles parallel decoding. During rendering of novel views, each test-time training layer functions as a static weight layer, making the entire model a static vision transformer dosovitskiy2020image. We illustrate this design in Figure [9](https://arxiv.org/html/2505.23884v1#A2.F9).

### 4.2 Language Modeling - Text Sequence

Autoregressive language models predict the probability distribution of the next token given preceding tokens, $p_{\theta}(x_{n}|x_{1},\dots,x_{n-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 as in Fig. [3](https://arxiv.org/html/2505.23884v1#S3.F3)(c) for the TTT apply-update sequence to avoid seeing future tokens in a chunk.
Since LaCT lacks per-token causality within each chunk, we employ sliding window attention—with window size equal to the chunk size—to efficiently model per-token causal dependencies. The sliding window is integrated into the same TTT layer with shared QKV similar to GAU hua2022transformer. We illustrate the detailed architecture in Fig. [10](https://arxiv.org/html/2505.23884v1#A2.F10) and pseudocode [2](https://arxiv.org/html/2505.23884v1#algorithm2).

### 4.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 teacher-forcing training by interleaving noisy and clean frame chunks.
Specifically, a video of N frame chunks is structured as:

$$ $S=[X_{1}^{\text{noise}},X_{1},X_{2}^{\text{noise}},X_{2},\ldots,X_{N}^{\text{ noise}}]$ (12) $$

where each noisy chunk $X_{i}^{\text{noise}}$ is produced by adding unit Gaussian noise $\epsilon$ to the $i$-th clean video chunk as $X_{i}^{\text{noise}}=X_{i}(1-t_{i})+\epsilon t_{i}$ and $t_{i}\in[0,1]$ denotes the strength of chunk-independent noise.

To handle such a data structure, we employ the strided block-wise causal mask in Fig. [3](https://arxiv.org/html/2505.23884v1#S3.F3)d 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 a non-overlapping window with 2 consecutive chunks (i.e., $[X_{i},X_{i+1}^{\text{noise}}]$) to build temporal and spatial locality.
Within each window, the attention from $X_{i}$ to $X_{i+1}^{\text{noise}}$ is excluded.
We incorporate the first noisy chunk by shifting all attention and TTT masking patterns similar to Fig. [3](https://arxiv.org/html/2505.23884v1#S3.F3)c.
The details of this hybrid architecture and more efficient trainings are in the Appendix [B.3](https://arxiv.org/html/2505.23884v1#A2.SS3).

## 5 Experiments

In this section, we present our experiment results on novel view synthesis (Sec. [5.1](https://arxiv.org/html/2505.23884v1#S5.SS1)),
language modeling (Sec.[5.2](https://arxiv.org/html/2505.23884v1#S5.SS2)), and autoregressive video generationo (Sec. [5.3](https://arxiv.org/html/2505.23884v1#S5.SS3)), and an in-depth analysis (Sec. [5.4](https://arxiv.org/html/2505.23884v1#S5.SS4)) of different design choices.
Tab. [1](https://arxiv.org/html/2505.23884v1#S5.T1) 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 in Appendix [C](https://arxiv.org/html/2505.23884v1#A3).

**Table 1: Summary of our experiments on three different data structures. ‘d’ denotes model dimension. The state size denotes the size of the fast weight per model block.**
| Task name | Data Structure | Chunk Size | State Size | Model Size | Max Length | Context Parallelism |
| --- | --- | --- | --- | --- | --- | --- |
| Novel View Synthesis | Image set | Full sequence | $6d^{2}$ | 0.3B | 1M | Within-chunk parallel |
| AR Video Diffusion | Image sequence | Three frames | $3d^{2}$, $0.75d^{2}$ | 1.3B, 14B | 56160 | Head-dim parallel |
| Language Models | 1D Sequence | 2K, 4K tokens | $0.75d^{2}$ | 0.7B, 3B | 32768 | N/A |

### 5.1 Novel View Synthesis

Datasets & metric. We evaluate our approach on
both object-level and scene-level datasets.
We use Objaverse dataset deitke2023objaverse for object-level training,
following the setup from LVSM jin2024lvsm and GS-LRM zhang2024gs.
After training, we perform evaluations on the Google Scanned Objects (GSO) dataset downs2022google, at resolutions of $256\times 256$ and $512\times 512$. Each evaluation involves 4–48 input views and 8 novel views per object.
For scene-level evaluations, we adopt the challenging DL3DV scene dataset ling2024dl3dv, with over 11K training scenes and 140 testing scenes, each with approximately 300 views.
Evaluations are at a resolution of $960\times 536$.
Performance is measured by Peak Signal-to-Noise Ratio (PSNR) at novel views, with additional metrics provided in the Appendix [C.1](https://arxiv.org/html/2505.23884v1#A3.SS1).

Model details.
Each block of model
has a per-image window attention layer, a SwiGLU-MLP large-chunk TTT layer,
and a feed-forward layer.
The default model totals 312M parameters, including 84M fast weights ($6d^{2}$ per block).

Baselines.
For object-level evaluation, we use two baselines: a full-attention model and a Perceiver-style register-attention model jaegle2021perceiver. The full-attention baseline replaces TTT layers with block-wise causal attention layers, enabling bidirectional interaction among input tokens and cross-attention from novel views. The Perceiver-style baseline compresses input tokens into 4096 registers, decoding novel views via cross-attention to these registers.
For scene-level evaluation, we compare with LongLRM ziwen2024long, a state-of-the-art model combining Mamba gu2023mamba and full attention for 3D Gaussian splat predictions, as well as pure optimization-based 3D Gaussian splatting methods. Table [2](https://arxiv.org/html/2505.23884v1#S5.T2) summarizes the computational complexities of all models.

Training details.
For object dataset, we train all models with $1.25$ trillion tokens with progressive resolutions. For scene dataset, we train our model with $1.8$ trillion tokens with progressively higher resolutions and more views, at a maximal sequence length of 1 million tokens. High-resolution models are trained with inner-chunk context parallelism (Sec. [3.4](https://arxiv.org/html/2505.23884v1#S3.SS4)). See more details in Sec. [C.1](https://arxiv.org/html/2505.23884v1#A3.SS1).

Results.
Experimental results and analysis are presented in Figure [4](https://arxiv.org/html/2505.23884v1#S5.F4).

**Table 2: Complexities of methods on novel view synthesis w/ $n$ input. Prefill and rendering speed are measured on A100 with 48 512$\times$512 input images (196K input tokens, 4K decoding tokens).**
|  | State Size | Prefill Compute | Decoding Compute | # Params | Prefill speed | Rendering FPS |
| --- | --- | --- | --- | --- | --- | --- |
| Full attention | $O(n)$ | $O(n^{2})$ | $O(n)$ | 284M | 16.1 s | 2.3 FPS |
| Perceiver Attention | $O(1)$ | $O(n^{2})$ | $O(1)$ | 287M | 16.8 s | 34.4 FPS |
| Ours | $O(1)$ | $O(n)$ | $O(1)$ | 312M | 1.4 s | 38.7 FPS |

Figure: 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 LongLRM, limited to 32 views, and outperforms 3D Gaussian Splatting with sparse views, remaining competitive up to 128 input views (1M total tokens).
Refer to caption: x4.png

### 5.2 Language Modeling

Datasets & Metrics.
We train our models on the Long-Data-Collections dataset long_data_collection, using approximately 60B tokens from its total 68.8B tokens. For evaluation, we employ the per-token loss metric from linforgetting, assessing models’ ability to effectively use the full context. A monotonically decreasing loss indicates successful context utilization, whereas plateauing suggests limited context usage. Additionally, we report retrieval accuracy hsiehruler at various sequence lengths.

Model details. We remove the window-attention layer from the original the LaCT block, integrating a sliding window-attention(SWA) layer directly into the Large-Chunk TTT layer. Following GAU hua2022transformer, SWA shares Q, K, and V vectors with the fast-weight network, with additional per-channel scaling and shifting on Q and K.
The pseudocode for this design is in Algorithm [2](https://arxiv.org/html/2505.23884v1#algorithm2).

Baselines.
We compare against full attention, Gated Linear Attention (GLA) yang2024gated, DeltaNet schlag2021linear; yangparallelizing. To ensure fairness, we enhance both GLA and DeltaNet with the same sliding window attention. Based on prior work linforgetting; xiong2023effective; men2024base highlighting the importance of a large RoPE su2023roformer base for long-context transformer training, we adopt a RoPE base of 1 million for training with 32K token contexts.
Tab. [3](https://arxiv.org/html/2505.23884v1#S5.T3) summarize the mechanism and training throughput of all methods.

**Table 3: 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-sequence length on A100-40GB GPUs.**
|  | State size | Train TPS | Update Rule | Memory read-out |
| --- | --- | --- | --- | --- |
| Transformer | – | 4.1K | – | – |
| Transformer SWA | – | 6.4K | – | – |
| Per-token recurrence |  |  |  |  |
| GLA SWA | $384d$ | 5.0K | $\displaystyle\mathbf{S}_{t}\leftarrow\mathbf{S}_{t-1}\mathrm{Diag}(\boldsymbol<br>{\alpha}_{t})+\mathbf{v}_{t}\mathbf{k}_{t}^{\top}$ | $\displaystyle\mathbf{o}_{t}=\mathbf{S}_{t}\mathbf{q}_{t}$ |
| DeltaNet SWA | $128d$ | 5.1K | $\displaystyle\mathbf{S}_{t}\leftarrow\mathbf{S}_{t-1}(\mathbf{I}-\beta_{t}<br>\mathbf{k}_{t}\mathbf{k}_{t}^{\top})+\beta_{t}\mathbf{v}_{t}\mathbf{k}_{t}^{\top}$ | $\displaystyle\mathbf{o}_{t}=\mathbf{S}_{t}\mathbf{q}_{t}$ |
| Large-chunk recurrence |  |  |  |  |
| Ours GD | $2304d$ | 5.0K | $W\leftarrow\mathrm{L2norm}(W-\sum_{i}^{b}\eta_{i}\nabla_{W}\mathcal{L}_{i})$ | $\displaystyle\mathbf{o}_{t}=f_{W}(\mathbf{q}_{t})$ |
| Ours Momentum | $2304d$ | 4.9K | $M\leftarrow\beta M+\sum_{i}^{b}\eta_{i}\nabla_{W}\mathcal{L}_{i};\;W\leftarrow<br>\mathrm{L2norm}(W-M)$ | $\displaystyle\mathbf{o}_{t}=f_{W}(\mathbf{q}_{t})$ |
| Ours Muon | $2304d$ | 4.3K | $M\leftarrow\beta M+\sum_{i}^{b}\eta_{i}\nabla_{W}\mathcal{L}_{i};\;W\leftarrow<br>\mathrm{L2norm}(W-\mathrm{Muon}(M))$ | $\displaystyle\mathbf{o}_{t}=f_{W}(\mathbf{q}_{t})$ |

Training details. We trained models at two scales using a 32768-token sequence length: a 760M-parameter model trained for 40B tokens with a 2048-token sliding window, and a 3B-parameter model trained for 60B tokens with a 4096-token sliding window. See more details in Sec. [C.2](https://arxiv.org/html/2505.23884v1#A3.SS2).

Results. Please refer to Fig. [5](https://arxiv.org/html/2505.23884v1#S5.F5) for experiment results and analysis. See more results in Sec. [C.2](https://arxiv.org/html/2505.23884v1#A3.SS2).

Figure: Figure 5: Language Model results. (a, c) Our model achieves lower 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.
Refer to caption: x5.png

### 5.3 Autoregressive Video Diffusion

We fine-tune the pretrained Wan 2.1 wang2025wan text-to-video diffusion model into an autoregressive video diffusion model. Specifically, we replace all bidirectional attention layers with our LaCT layers combined with sliding window attention. The sliding window attention uses a window size spanning two autoregressive chunks.

Datasets. We fine-tune the model using an internal, filtered proprietary collection of videos, each accompanied by a short text prompt generated by a visual language model.

Training details.
Following esser2403scaling; wang2025wan, we employ time-step shifting and denoising loss weighting using a logit-normal distribution. we train on 5-second videos at 16 FPS and 480$\times$832 resolution, autoregressively denoising in 3 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-second clip contains 56,160 visual tokens, resulting in interleaved noisy-clean chunks totaling 107K tokens under teacher-forcing training. We use sequence parallelism for MLP layers and tensor parallelism (sharding heads across devices) for TTT and window attention layers. Full details are listed in Sec. [C.3](https://arxiv.org/html/2505.23884v1#A3.SS3).

Baselines.
We compare our method against three baselines: sliding window attention (SWA) alone, Mamba2 dao2024transformers combined with SWA (using a similar parallel combination strategy as our 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](https://arxiv.org/html/2505.23884v1#S5.F6) plots the chunk-wise denoising loss across evaluated video frames. We only measure validation loss up to training sequence length. See project website for our autoregressively generated videos (^2^22See video results in project website: [https://tianyuanzhang.com/projects/ttt-done-right/](https://tianyuanzhang.com/projects/ttt-done-right/)).

Figure: 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 on longer videos (c).
Refer to caption: x6.png

### 5.4 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. [7](https://arxiv.org/html/2505.23884v1#S5.F7)a),
test-time optimizers (Fig. [7](https://arxiv.org/html/2505.23884v1#S5.F7)b),
linear versus nonlinear fast weights (Fig. [8](https://arxiv.org/html/2505.23884v1#S5.F8)a), and
per-token recurrence versus chunk-wise recurrence (Fig. [8](https://arxiv.org/html/2505.23884v1#S5.F8)b).
Overall, we find that a large state size, advanced optimizers such as Muon, and nonlinear fast weights significantly improve our model’s performance. For comparing chunk recurrence with per-token recurrence, in a controlled NVS experiment, our linear large-chunk 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 DeltaNet—surpasses them when combined with a large nonlinear state and the Muon optimizer.
We refer the readers to each figure and its caption for more detailed analysis.

Figure: 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 $12d^{2}$ per block, totaling 40% of model weights as fast weights. (b) Comparison of test-time optimizers demonstrates Muon’s surprising effectiveness over Vanilla Gradient Descent and Momentum.
Refer to caption: x7.png

Figure: 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.
Refer to caption: x8.png

#### State size scaling.

These controlled experiments utilize a 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$. To change the state size, we keep the head dimension fixed as model dimension. i.e. single head, and vary the intermediate dimension of SwiGLU MLP, such that the intermediate dimension ranges from $192$ to $3072$. The largest configuration results in a state size per model block as $12d^{2}$, 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](https://arxiv.org/html/2505.23884v1#S5.F7)(a) 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 on momentum implementation are in Appendix [A](https://arxiv.org/html/2505.23884v1#A1.SS0.SSS0.Px6). 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](https://arxiv.org/html/2505.23884v1#S5.F7)(b) 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, $f_{W}(x)=Wx$. Both are updated using the same online dot product loss for key-value association. Figure [8](https://arxiv.org/html/2505.23884v1#S5.F8) (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 671B tokens with 24 blocks and $d=768$. Language modeling used the 760M parameter setup.

#### Large-chunk v.s. Per-token recurrence.

Figure [8](https://arxiv.org/html/2505.23884v1#S5.F8)(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: $f_{W}(x)=Wx$ and is benchmarked against a Mamba-2 baseline (a linear 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 $d^{2}$ per block. Both of these two approaches employs a per-image window attention within each model block. Under this fair comparison, our linear large-chunk recurrence achieves significantly better view synthesis performance.

For the language modeling experiments also shown in Figure [8](https://arxiv.org/html/2505.23884v1#S5.F8)(b), the blue line “Our Linear" variant uses the same state size ($0.25d^{2}$) as the GLA SWA baseline. It 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 is equipped with a larger non-linear state ($1.5d^{2}$) and Muon updates, we significantly outperform these per-token recurrence baselines.

## 6 Related Work

#### Test-time training.

Test-Time Training (TTT) sun2024learning 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 to learn in-context. Existing methods typically employ a 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 sun2024learning; wang2025testtimeregressionunifyingframework has opened a vast design space for new recurrent model architectures. For instance, many recent works have developed novel test-time optimizers behrouz2024titans; karami2025lattice and online training objectives behrouz2025s. However, current 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-magnitude higher hardware utilization without relying on error-prone custom kernel implementations. Furthermore, it 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 attention with linear recurrence, such as Gated Attention Unit (GAU) hua2022transformerqualitylineartime, MEGA ma2023mega, MEGALODON ma2024megalodon, and InfiniAttention munkhdalai2024leavecontextbehindefficient. Among these, InfiniAttention is conceptually closest to our work, as it incorporates recurrence at the chunk level using the delta rule—interpreted 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 hutchins2022blockrecurrent 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. [5.1](https://arxiv.org/html/2505.23884v1#S5.SS1), Table [2](https://arxiv.org/html/2505.23884v1#S5.T2)) is conceptually similar to the Block-Recurrent Transformer in its use of register tokens for context compression. As shown in Figure [4](https://arxiv.org/html/2505.23884v1#S5.F4), our method significantly outperforms this approach in both speed and quality, with a comparable state size.

#### 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-based approaches, such as NeRF mildenhall2021nerf and 3D Gaussian Splatting kerbl20233d, have achieved significant breakthroughs.
These methods optimize a 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 can render novel views photorealistically, and the optimized parameters form a 3D representation of the input scene.

Recently, data-driven approaches zhang2024gs; jin2024lvsm; ziwen2024long; liu2023zero have also shown promising results. These methods can either directly render novel views or predict 3D representations given input images. Although successful on simpler object datasets, these methods often struggle with densely sampled scenes (e.g., scenes with over 100 input images). Our experiments demonstrate that our large-chunk test-time training approach outperforms or achieves comparable performance to 3D Gaussian Splatting on challenging scene datasets with up to 128 input images with $960\times 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 bidirectional diffusion transformers operating in latent space videoworldsimulators2024; yang2024cogvideox; polyak2410movie; wang2025wan. These methods factorize the video distribution into a sequence of conditional distributions based on noise levels, following diffusion processes sohl2015deep; song2020score or flow matching lipman2022flow, then use a diffusion transformer to jointly learn all the conditional distribution. Autoregressive video diffusion alonso2024diffusion; jin2024pyramidal; valevski2024diffusion; ruhe2024rolling; yin2024slow; song2025history 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 as condition alonso2024diffusion; jin2024pyramidal; valevski2024diffusion, 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 as progressive noise injection ruhe2024rolling or the use of frame-independent noises (sometimes in a diffusion-forcing style) yin2024slow; chen2024diffusion; magi1 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 [12](https://arxiv.org/html/2505.23884v1#S4.E12)). This strategy achieves over 50% token utilization and integrates effectively with our large-chunk TTT implementation, by only changing a few lines to constrain fast-weights are only updated on clean frame chunks.

## 7 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 as RoPE su2023roformer), our SwiGLU and Linear Fast Weight components do not exhibit this property. The practical implications of this absence remain underexplored.

We conduct our experiments on three tasks.
Although the tasks are diverse and cover different modalities, the effectiveness of our method would request of more tasks.
For example, the novel-view synthesis task is essentially a 3D reconstruction with input pose information.
The task of unposed reconstruction is more challenging and is not explored in this paper.

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 the parameter size.
Previous papers showed that a 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 diffusion, 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 our paper and it is a common choice in evaluating the scalability of video generation.
This is a general problem for the video generation evaluation and is not specific to our paper.

## 8 Conclusion

We presented LaCT, a 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—and demonstrate its effectiveness
by achieving superior or competitive performance when compared to state-of-the-art baselines.
LaCT achieves high GPU efficiency even with native PyTorch implementation with dozens of lines of code and supports efficient scaling
up of 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.

## 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 on 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 on Muon.

## References

- [1]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.
Attention is all you need.
Advances in neural information processing systems, 30, 2017.
- [2]
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al.
Learning to (learn at test time): Rnns with expressive hidden states.
arXiv preprint arXiv:2407.04620, 2024.
- [3]
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber.
Linear transformers are secretly fast weight programmers.
In International Conference on Machine Learning, pages 9355–9366. PMLR, 2021.
- [4]
Ke Alexander Wang, Jiaxin Shi, and Emily B. Fox.
Test-time regression: a unifying framework for designing sequence models with associative memory, 2025.
- [5]
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni.
Titans: Learning to memorize at test time.
arXiv preprint arXiv:2501.00663, 2024.
- [6]
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, 2025.
- [7]
Mahdi Karami and Vahab Mirrokni.
Lattice: Learning to efficiently compress the memory.
arXiv preprint arXiv:2504.05646, 2025.
- [8]
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.
- [9]
Bernhard Kerbl, Georgios Kopanas, Thomas Leimkühler, and George Drettakis.
3d gaussian splatting for real-time radiance field rendering.
ACM Trans. Graph., 42(4):139–1, 2023.
- [10]
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, 2024.
- [11]
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 arXiv:2307.08621, 2023.
- [12]
Albert Gu and Tri Dao.
Mamba: Linear-time sequence modeling with selective state spaces.
arXiv preprint arXiv:2312.00752, 2023.
- [13]
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, pages 56501–56523. PMLR, 2024.
- [14]
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong.
Various lengths, constant speed: Efficient language modeling with lightning attention.
In Forty-first International Conference on Machine Learning, 2024.
- [15]
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim.
Parallelizing linear transformers with the delta rule over sequence length.
In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
- [16]
Tri Dao.
Flashattention-2: Faster attention with better parallelism and work partitioning.
arXiv preprint arXiv:2307.08691, 2023.
- [17]
Maxim Milakov and Natalia Gimelshein.
Online normalizer calculation for softmax.
arXiv preprint arXiv:1805.02867, 2018.
- [18]
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.
Flashattention: Fast and memory-efficient exact attention with io-awareness.
Advances in neural information processing systems, 35:16344–16359, 2022.
- [19]
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.
Deep residual learning for image recognition. arxiv e-prints.
arXiv preprint arXiv:1512.03385, 10:9, 2015.
- [20]
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian 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, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas
Scialom.
Llama 2: Open foundation and fine-tuned chat models, 2023.
- [21]
Noam Shazeer.
Glu variants improve transformer, 2020.
- [22]
Tim Salimans and Durk P Kingma.
Weight normalization: A simple reparameterization to accelerate training of deep neural networks.
Advances in neural information processing systems, 29, 2016.
- [23]
Tri Dao and Albert Gu.
Transformers are ssms: Generalized models and efficient algorithms through structured state space duality.
arXiv preprint arXiv:2405.21060, 2024.
- [24]
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 arXiv:2402.18668, 2024.
- [25]
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le.
Transformer quality in linear time.
In International conference on machine learning, pages 9099–9117. PMLR, 2022.
- [26]
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal.
Leave no context behind: Efficient infinite context transformers with infini-attention, 2024.
- [27]
Leonard McMillan and Gary Bishop.
Plenoptic modeling: An image-based rendering system.
In Seminal Graphics Papers: Pushing the Boundaries, Volume 2, pages 433–440. 2023.
- [28]
Marc Levoy and Pat Hanrahan.
Light field rendering.
In Seminal Graphics Papers: Pushing the Boundaries, Volume 2, pages 441–452. 2023.
- [29]
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 arXiv:2410.17242, 2024.
- [30]
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 arXiv:2010.11929, 2020.
- [31]
Matt Deitke, Dustin Schwenk, Jordi Salvador, Luca Weihs, Oscar 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, pages 13142–13153, 2023.
- [32]
Kai Zhang, Sai Bi, Hao Tan, Yuanbo Xiangli, Nanxuan Zhao, Kalyan Sunkavalli, and Zexiang Xu.
Gs-lrm: Large reconstruction model for 3d gaussian splatting.
In European Conference on Computer Vision, pages 1–19. Springer, 2024.
- [33]
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 household items.
In 2022 International Conference on Robotics and Automation (ICRA), pages 2553–2560. IEEE, 2022.
- [34]
Lu Ling, Yichen Sheng, Zhi Tu, Wentian Zhao, Cheng Xin, Kun Wan, Lantao Yu, Qianyu Guo, Zixun Yu, Yawen Lu, et al.
Dl3dv-10k: A large-scale scene dataset for deep learning-based 3d vision.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22160–22169, 2024.
- [35]
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira.
Perceiver: General perception with iterative attention.
In International conference on machine learning, pages 4651–4664. PMLR, 2021.
- [36]
Chen Ziwen, Hao Tan, Kai Zhang, Sai Bi, Fujun Luan, Yicong Hong, Li Fuxin, and Zexiang Xu.
Long-lrm: Long-sequence large reconstruction model for wide-coverage gaussian splats.
arXiv preprint arXiv:2410.12781, 2024.
- [37]
Together AI.
Long data collections database, 2024.
- [38]
Zhixuan Lin, Evgenii Nikishin, Xu He, and Aaron Courville.
Forgetting transformer: Softmax attention with a forget gate.
In The Thirteenth International Conference on Learning Representations, 2025.
- [39]
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg.
Ruler: What’s the real context size of your long-context language models?
In First Conference on Language Modeling, 2024.
- [40]
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.
- [41]
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.
- [42]
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu.
Roformer: Enhanced transformer with rotary position embedding, 2023.
- [43]
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.20314, 2025.
- [44]
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.03206, 2, 2024.
- [45]
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.
arXiv preprint arXiv:2504.13173, 2025.
- [46]
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc V. Le.
Transformer quality in linear time, 2022.
- [47]
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.
- [48]
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 Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
- [49]
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, editors, Advances in Neural Information Processing Systems, 2022.
- [50]
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):99–106, 2021.
- [51]
Ruoshi Liu, Rundi Wu, Basile Van Hoorick, Pavel Tokmakov, Sergey Zakharov, and Carl Vondrick.
Zero-1-to-3: Zero-shot one image to 3d object.
In Proceedings of the IEEE/CVF international conference on computer vision, pages 9298–9309, 2023.
- [52]
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 as world simulators.
2024.
- [53]
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 an expert transformer.
arXiv preprint arXiv:2408.06072, 2024.
- [54]
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://arxiv. org/abs/2410.13720, page 51, 2024.
- [55]
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
Deep unsupervised learning using nonequilibrium thermodynamics.
In International conference on machine learning, pages 2256–2265. pmlr, 2015.
- [56]
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.
Score-based generative modeling through stochastic differential equations.
arXiv preprint arXiv:2011.13456, 2020.
- [57]
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le.
Flow matching for generative modeling.
arXiv preprint arXiv:2210.02747, 2022.
- [58]
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos J Storkey, Tim Pearce, and François Fleuret.
Diffusion for world modeling: Visual details matter in atari.
Advances in Neural Information Processing Systems, 37:58757–58791, 2024.
- [59]
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.05954, 2024.
- [60]
Dani Valevski, Yaniv Leviathan, Moab Arar, and Shlomi Fruchter.
Diffusion models are real-time game engines.
arXiv preprint arXiv:2408.14837, 2024.
- [61]
David Ruhe, Jonathan Heek, Tim Salimans, and Emiel Hoogeboom.
Rolling diffusion models.
In International Conference on Machine Learning, pages 42818–42835. PMLR, 2024.
- [62]
Tianwei Yin, Qiang Zhang, Richard Zhang, William T Freeman, Fredo Durand, Eli Shechtman, and Xun Huang.
From slow bidirectional to fast causal video generators.
arXiv preprint arXiv:2412.07772, 2024.
- [63]
Kiwhan Song, Boyuan Chen, Max Simchowitz, Yilun Du, Russ Tedrake, and Vincent Sitzmann.
History-guided video diffusion.
arXiv preprint arXiv:2502.06764, 2025.
- [64]
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, 37:24081–24125, 2024.
- [65]
Sand-AI.
Magi-1: Autoregressive video generation at scale, 2025.
- [66]
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 pattern recognition, pages 586–595, 2018.
- [67]
Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen.
Query-key normalization for transformers.
In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 4246–4253, 2020.
- [68]
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.
- [69]
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, pages 24185–24198, 2024.
- [70]
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, 37:62432–62457, 2024.
- [71]
Horace He, Driss Guessous, Yanbo Liang, and Joy Dong.
Flexattention: The flexibility of pytorch with the performance of flashattention.
PyTorch Blog, 8, 2024.
- [72]
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu.
Unipc: A unified predictor-corrector framework for fast sampling of diffusion models.
Advances in Neural Information Processing Systems, 36:49842–49869, 2023.
- [73]
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 arXiv:2309.14509, 2023.
- [74]
Jiarui Fang and Shangchun Zhao.
Usp: A unified sequence parallelism approach for long context generative ai.
arXiv preprint arXiv:2405.07719, 2024.

## Appendix A LaCT Model Implementation Details

Figure: Algorithm 1 Large Chunk Test-Time Training Layer Pseudocode

Figure: Algorithm 2 LaCT Layer with In-Layer Hybrid Window Attention Pseudocode

#### State Size calculation.

Motivated by recent progress in LLM, we adopt SwiGLU-MLP shazeer2020glu without bias terms as the fast-weight
network.
Our fast weights consists of three weight matrix $W=\{W_{1},W_{2},W_{3}\}$, and the forward pass of the fast weight model is:

$$ $f_{W}(x)=W_{2}\left[\mathrm{SiLU}(W_{1}x)\circ(W_{3}x)\right]$ (13) $$

where $\circ$ is an elementwise multiplication. We define $\mathit{hd}$ as the head dimension, $\mathit{nh}$ as the number of heads, and the intermediate dimension of the SwiGLU-MLP as $\mathit{hd}\times r$, where $r$ is a scaling multiplier.When $r>1$, the intermediate dimension surpasses the input head dimension, which is the current common practice in LLMs. TThus, matrices $W_{1},W_{2}\in\mathbb{R}^{\mathit{hd}\times\mathit{hd}}$ and $W_{3}\in\mathbb{R}^{\mathit{hd}\times\mathit{hd}\times r}$. Consequently, the total state size becomes $\mathit{nh}\times\mathit{hd}\times\mathit{hd}*r$. Given that typically the total head dimension across all heads equals the model dimension $d$ (i.e., $\mathit{nh}\times\mathit{hd}=d$), the total state size simplifies to:

$$ $\text{State Size}=\frac{d^{2}}{\mathit{nh}}*r.$ (14) $$

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 to compute the final results of $f_{W}(v)$. We only need to compute $W_{1}v,W_{3}v$ 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:

$$ $\text{FLOPs}=4n\frac{d^{2}}{\mathit{nh}}r+8n\frac{d^{2}}{\mathit{nh}}r+6n\frac {d^{2}}{\mathit{nh}}r=18n\frac{d^{2}}{\mathit{nh}}r=6*\text{State Size}$ (15) $$

#### 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/\sqrt{\text{fan-in}}$. Specifically, in the SwiGLU FFN fast weights, the matrices $w_{1}$ and $w_{3}$ have their fan-in set as the head dimension, while the fan-in of $w_{2}$ is the intermediate dimension of the SwiGLU FFN fast weights. Additionally, when local window attention is incorporated within the LaCT layer, we 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 jordan2024muon is a 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 $G_{0}=G/|G|F$, then iteratively applies:

$$ $\mathbf{G}_{k}=a\mathbf{G}_{k-1}+b(\mathbf{G}_{k-1}\mathbf{G}_{k-1}^{\mathrm{T }})\mathbf{G}_{k-1}+c(\mathbf{G}_{k-1}\mathbf{G}_{k-1}^{\mathrm{T}})^{2} \mathbf{G}_{k-1},$ (16) $$

where the constants $a,b,c$ are carefully chosen for optimal convergence. Following the original implementation, we set $a=3.4445$, $b=-4.7750$, $c=2.0315$, and perform five iterations.

Each Muon iteration requires three matrix multiplications, resulting in a computational cost per fast weight head of $2\mathit{hd}^{3}r+2\mathit{hd}^{3}+2\mathit{hd}^{3}r=\mathit{hd}^{3}(4r+2)$ FLOPS. Hence, the total computation for five iterations across all fast weights is:

$$ $5\times\mathit{nh}\times\mathit{hd}^{3}\times(4r+2).$ (17) $$

For the case where $r=1$ (head and intermediate dimensions are equal), the total computational cost simplifies to:

$$ $30\times\mathit{nh}\times\mathit{hd}^{3}=30\times\mathit{hd}\times\text{State Size}.$ (18) $$

This indicates that the computational overhead of Muon becomes less significant than computing token outputs only if the online chunk size exceeds $\frac{5}{3}\mathit{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 su2023roformer. In contrast, our SwiGLU and Linear Fast Weight components do not possess this property.

#### Implementing momentum for test-time optimizers.

Muon uses momentum by default. Following Titans behrouz2024titans, we implement momentum in the test-time optimizer by predicting a scalar momentum $\beta_{i}$ from each token:

$$ $\beta_{i}=\sigma(\text{Linear}(\boldsymbol{x_{i}})),$ (19) $$

where $\sigma$ is the sigmoid function. This $\beta_{i}$ is then averaged over all tokens in the chunk, and the average momentum is applied as follows:

$$ $\displaystyle g\leftarrow$ $\displaystyle\sum_{i}^{b}\eta_{i}\nabla_{W}\mathcal{L}(f_{W}(k_{i}),v_{i}),$ (20) $\displaystyle M\leftarrow$ $\displaystyle M(\sum_{i}^{b}\beta_{i}/b)+g,$ $\displaystyle W\leftarrow$ $\displaystyle\text{weight-update}(W,M),$ $$

where the weight-update can be simple subtraction followed by L2 normalization normalization (as in Equation [8](https://arxiv.org/html/2505.23884v1#S3.E8)). or Muon update before subtraction (as in Equation [9](https://arxiv.org/html/2505.23884v1#S3.E9)).

### A.1 Pseudocode

See Algorithm [1](https://arxiv.org/html/2505.23884v1#algorithm1) for pseudocode of a full LaCT layer.
For details on how to mix local window attention inside each layer with shared query, value, embedding, see Algorithm [2](https://arxiv.org/html/2505.23884v1#algorithm2) for pseudocode.

## Appendix B LaCT Architecture Details for N-Dimensional Data

### B.1 LaCT Architecture for Novel View Synthesis

Novel view synthesis (NVS) renders images of a static scene from novel viewpoints.
Formally, given a set of $N$ input posed images $\{(I_{i},P_{i})\}_{i=1}^{N}$ of a static scene,
where $I_{i}\in\mathbb{R}^{H\times W\times 3}$ is an RGB image and $P_{i}$ is its corresponding camera pose,
the model needs to synthesize new images from novel camera poses $P_{\text{novel}}$ 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 [3.1](https://arxiv.org/html/2505.23884v1#S3.SS1)) 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=(\text{rays}_{o},\text{rays}_{d})$.
$\text{rays}_{o}\in\mathbb{R}^{H\times W\times 3}$ is the 3D coordinate for the origins of the ray, and $\text{rays}_{d}\in\mathbb{R}^{H\times W\times 3}$ is the direction of the ray.
We follow GS-LRM zhang2024gs to use 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 a concatenation of the ray’s origin, the ray’s direction, and the cross product of the above two:
$[\text{rays}_{o},\text{rays}_{d},\text{rays}_{o}\times\text{rays}_{d}]$.
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](https://arxiv.org/html/2505.23884v1#A2.F9).
We first apply the attention for each image (either input or the novel target).
The attention is bidirectional for the tokens belonging to the same image, and is 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. [9](https://arxiv.org/html/2505.23884v1#A2.F9) 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 a single data point for better training efficiency.
Given the design, the novel images are independent to each other, which is illustrated in the ‘Overall LaCT Mask’ in the right of Fig. [9](https://arxiv.org/html/2505.23884v1#A2.F9).
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 of the encoder-decoder model in LVSM jin2024lvsm, except we use TTT in replace of transformer for long-context modeling.

Figure: 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 a single image, either the input image or the novel target image. TTT updates over all input tokens and apply to all tokens.
Refer to caption: x9.png

For actually using this model for NVS task during inference, we first get the updated fast weight with all input images.
Then, we 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 a 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](https://arxiv.org/html/2505.23884v1#S3.F2)) stored the world knowledge like physical rendering rules.

Figure: Figure 10: Detailed LaCT model for language models. Dashed line indicates flow of fast weight. Solid line indicates flow of tokens. We illustrate with TTT chunk size 4 or 2, and the actual chunk size is 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.
Refer to caption: x10.png

### B.2 LaCT Architecture for Language Models

Autoregressive Language Models (LM) predicts the distribution of the next tokens $p_{\theta}(x_{n}|x_{1},\dots,x_{n-1})$ from its history context.
It is a factorization of the full sequence distribution $p_{\theta}(x_{1},\dots,x_{n})$ through chain rule $p_{\theta}(x_{1},\dots,x_{n})=p_{\theta}(x_{1})p_{\theta}(x_{2}|x_{1})\ldots p
_{\theta}(x_{n}|x_{1},\dots,x_{n-1})$.
Thus it requires a token-level causal mask (demonstrated in the topright of Fig. [10](https://arxiv.org/html/2505.23884v1#A2.F10))
and this is the main difficulty for the large-chunk design in LaCT.
We use a combination of TTT layer with ‘Shifted Causal Block Mask’ (introduced before in Fig. [3](https://arxiv.org/html/2505.23884v1#S3.F3)c) and a sliding window attention to facilitate it.
By shifting the mask of TTT, it excludes the information leakage from future tokens.
As shown in the right part of Fig. [10](https://arxiv.org/html/2505.23884v1#A2.F10), 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 in Sec. [2.2](https://arxiv.org/html/2505.23884v1#S2.SS2)).

As illustrated in the left most of Fig. [10](https://arxiv.org/html/2505.23884v1#A2.F10), 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 is an attention with constant window size over the past history, starting from the target tokens.
For the TTT layer, we start with an 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 see 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. [10](https://arxiv.org/html/2505.23884v1#A2.F10).
For details of the parallel design, please refer to the Pseudocode Algorithm [2](https://arxiv.org/html/2505.23884v1#algorithm2).

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) to enable larger state size, as state size is propotional to the head dimension in our design (Sec. [A](https://arxiv.org/html/2505.23884v1#A1), and Equation [14](https://arxiv.org/html/2505.23884v1#A1.E14)). By default, we use four heads in langugae model experiments. For positional encoding, we use the same RoPE as the window attention branch.

Figure: Figure 11: 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 a 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.
Refer to caption: x11.png

### B.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 noisy and clean frame chunks. Specifically, a video of N frame chunks is structured as:

$$ $S=[X_{1}^{\text{noise}},X_{1},X_{2}^{\text{noise}},X_{2},\ldots,X_{N}^{\text{ noise}}]$ (21) $$

where each noisy chunk $X_{i}^{\text{noise}}$ is produced by adding unit Gaussian noise $\epsilon$ to the $i$-th clean video chunk as $X_{i}^{\text{noise}}=X_{i}(1-t_{i})+\epsilon t_{i}$ and $t_{i}\in[0,1]$ denotes the strength of chunk-independent noise. However, compared to previous methods that employs progressive ruhe2024rolling or frame-independent noise strategies yin2024slow like diffusion forcing chen2024diffusion,
our teacher-forcing formulation in Equation [21](https://arxiv.org/html/2505.23884v1#A2.E21) 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_{1}^{\text{noise}},X_{1}^{\text{noise}_{*}},X_{1},X_{2}^{\text{noise}},X_ {2}^{\text{noise}_{*}},X_{2},\ldots,X_{N}^{\text{noise}},X_{N}^{\text{noise}_{ *}}],$ (22) $$

where $X_{i}^{\text{noise}}$ and $X_{i}^{\text{noise}_{*}}$ represent two different noise levels applied to each clean video chunk $X_{i}$. 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 a window size of two frame chunks and employs a 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 [11](https://arxiv.org/html/2505.23884v1#A2.F11) illustrates this design for autoregressive video generation. The input sequence depicted in the figure follows Equation [22](https://arxiv.org/html/2505.23884v1#A2.E22) 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.

## Appendix C Experimental Details

### C.1 Novel View Synthesis

Datasets & metric. We evaluate our approach on
both object-level and scene-level datasets.
We use Objaverse dataset deitke2023objaverse for object-level training,
and render 32 random views per object,
following the setup from LVSM jin2024lvsm and GS-LRM zhang2024gs.
After training, we perform evaluations on the Google Scanned Objects (GSO) dataset downs2022google, at resolutions of $256\times 256$ and $512\times 512$.
To ensure stablized evaluation results, we render at fixed view points instead of random view points as 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 ling2024dl3dv, which contains over 11K training scenes and 140 testing scenes, each with approximately 300 views.
Evaluations are performed at a resolution of $960\times 536$ (^3^33The original DL3DV 960p frames released in resolution of $960\times 536$. To accommodate the patch-size $8$ in our modeling, we crop it to $960\times 536$ and the camera parameters are changed accordingly.).
We report Peak Signal-to-Noise Ratio (PSNR) in the paper’s main figures, and other metrics Structural Similarity Index Measure (SSIM) and LPIPS zhang2018unreasonable can be found in Tables below.
For DL3DV evaluation, we follow the original paper ling2024dl3dv and LongLRM ziwen2024long 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 ziwen2024long.

Model details.
Our models consist of 24 stacked LaCT blocks, each with a model dimension of 768.
The detail of such block is illustrated in Sec. [C.1](https://arxiv.org/html/2505.23884v1#A3.SS1):
Unless otherwise specified, we use a single-headed fast-weight SwiGLU-MLP with a hidden dimension of 1536.
The window attention has 12 heads with head dimension 64, and is equipped with QK-normalization henry2020query.
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., $6d^{2}$ per block).
We use an fast-weight lr initialization of $0.01$ by setting ‘const_lr_bias’ in Algorithm [1](https://arxiv.org/html/2505.23884v1#algorithm1) to $\mathrm{softplus}(\text{const_lr_bias})=0.01$.
As we used Muon in fast weight update for NVS, LaCT is not sensitive to lr scale as discussed in Sec. [3.2](https://arxiv.org/html/2505.23884v1#S3.SS2).

Baselines.
For object-level evaluations, we compare against two baselines, including a full-attention model,
and a register-attention model in a Perceiver style jaegle2021perceiver,
In the full-attention baseline, we
replace the TTT layer in our model with a block-wise causal attention layer, where
the input tokens interact bidirectionally and the novel view tokens cross-attend
to the input tokens. Such a design resembles our method’s prefill and parallel
decoding strategy described in Section [4.1](https://arxiv.org/html/2505.23884v1#S4.SS1), 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
input-to-register full-attention layers and the remaining half with register-to-novel-view
cross-attention layers. Such a model first compresses the input tokens into a constant
set of register tokens and then decoding the novel view tokens by attending to the registers.
For scene-level evaluations, we compare against a state-of-the-art long-sequence 3D
reconstruction work LongLRM ziwen2024long that applies Mamba gu2023mamba hybrid with full attention to predict
3D Gaussian splats kerbl20233d.
We also include comparisons with pure optimization-based 3D Gaussian splatting methods.
Tab. [2](https://arxiv.org/html/2505.23884v1#S5.T2) compares the computational complexity of the
baseline models and our models.

Training details.
For object-level experiments, we first train all the model with $671$B tokens at 8 input view and 8 novel view setting at a resolution of $256\times 256$.
We then finetune them with $512\times 512$ resolution for an additional $587$B tokens.
For scene dataset, we first pre-train our model first with 32 input views and 32 novel views at $128\times 128$ resolution for $1.5$T 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 to 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\times 72$, $256\times 144$, $512\times 288$, and $960\times 536$.
The training tokens for each stage is around $100$B.
High-resolution models (starting from $512\times 288$) are trained with inner-chunk context parallelism (Sec. [3.4](https://arxiv.org/html/2505.23884v1#S3.SS4)).

At each training stage, we always use AdamW with linear learning rate warmup and weight decay of $0.05$.
The peak learning rate of the pre-training is $4e-4$.
During fine-tuning, we use smaller learning rate (usually $1e-5$ to $5e-5$).

The training is completed with 64 A100 GPUs.
The pre-training takes $8$ days, and each fine-tuning stage is about $12$hours (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\times 256$ in Table [4](https://arxiv.org/html/2505.23884v1#A3.T4), $512\times 512$ in Table [5](https://arxiv.org/html/2505.23884v1#A3.T5)) and DL3DV evaluations (at resolution $960\times 536$ in Table [6](https://arxiv.org/html/2505.23884v1#A3.T6)).

**Table 4: 256-Res object-level novel view synthesis results on GSO. Both the input and output are with resolution $256\times 256$. $\uparrow$: higher is better, $\downarrow$: lower is better.**
| Input | # Input Tokens | LaCT | Full Attention | Perceiver Attention |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Views |  | PSNR ($\uparrow$) | LPIPS ($\downarrow$) | SSIM ($\uparrow$) | PSNR ($\uparrow$) | LPIPS ($\downarrow$) | SSIM ($\uparrow$) | PSNR ($\uparrow$) | LPIPS ($\downarrow$) | SSIM ($\uparrow$) |
| 4 | 4,096 | 32.4 | 0.030 | 0.962 | 32.6 | 0.029 | 0.964 | 30.3 | 0.039 | 0.950 |
| 8 | 8,192 | 35.3 | 0.019 | 0.976 | 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.014 | 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 5: 512-Res object-level novel view synthesis results on GSO. Both the input and output are with resolution $512\times 512$ comparison across methods. $\uparrow$: higher is better, $\downarrow$: lower is better.**
| Input Views | # Input Tokens | LaCT | Full Attention |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| PSNR ($\uparrow$) | LPIPS ($\downarrow$) | SSIM ($\uparrow$) | PSNR ($\uparrow$) | LPIPS ($\downarrow$) | SSIM ($\uparrow$) |  |  |
| 4 | 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 | 38.9 | 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 |

**Table 6: 960P scene-level novel view synthesis results for LaCT on DL3DV. Both the input and output are with resolution $960\times 536$ (width x height). $\uparrow$: higher is better, $\downarrow$: lower is better.**
| Input Views | # Input Tokens | PSNR ($\uparrow$) | LPIPS ($\downarrow$) | SSIM ($\uparrow$) |
| --- | --- | --- | --- | --- |
| 16 | 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 |

### C.2 Language Modeling

Datasets & Metrics.
We train our models on the Long-Data-Collections dataset long_data_collection, containing approximately 68.8B tokens tokenized using Mixstra tokenizer (32,000 codebook size).
The dataset is a mix of 41.4% General Data (e.g., RedPajama-Book, RedPajama-ArXiv, 1B tokens from RedPajama, and a Pile subsample) and 58.6% Instruction Data (e.g., UL2 Oscar, NI, and P3).
To evaluate long-context capabilities, we utilized the per-token loss metric from linforgetting. A consistently decreasing per-token loss across the input sequence indicates effective use of the entire context, while a plateau suggests an inability to leverage information beyond that point. Specifically, we evaluated next-token prediction loss on 2.5B tokens from the Book-3 dataset gao2020pile for our 760M-parameter model, and on 5B tokens for our 3B-parameter model. Additionally, we measured retrieval accuracy using the RULER benchmark hsiehruler across various sequence lengths to 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 a causal sliding window-attention(SWA) layer directly into the Large-Chunk 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 hua2022transformer). 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-weight lr initialization of $0.001$ by setting ‘const_lr_bias’ in Algorithm [1](https://arxiv.org/html/2505.23884v1#algorithm1) to $\mathrm{softplus}(\text{const_lr_bias})=0.001$.
We illustrate this archtecture in Figure [10](https://arxiv.org/html/2505.23884v1#A2.F10).

To ensure a fair comparison with baselines in terms of trainable parameters.we adjusted the LaCT block’s extra learnable initial fast weights $W=[W_{1},W_{2},W_{3}]$. o reduce parameters, we employed a low-rank version for $W_{1},W3$ with a rank of 32. For instance, if $W_{1}\in\mathcal{R}^{d\times d}$, its low-rank initial fast weight is $W_{1}=L\cdot R+0.5*I_{d}$, where $L\in\mathcal{R}^{d\times 32},R\in\mathcal{R}^{32\times d}$, and $I_{d}$ is identity matrix. This reduces the extra trainable parameters for the fast weights in each block to $128*\text{model-dim}+\frac{1}{\text{num-heads}}(\text{model-dim}^{2})$. Additional minor parameters for learning rate projection, per-head scalers, an extra RMSNorm, and the SWA’s learnable scale and shift are of order $O(\text{model-dim})$. Standard blocks typically have $12\cdot\text{model-dim}^{2}$ parameters. Our approach adds approximately $\frac{1}{\text{num-heads}}(\text{model-dim}^{2})$ 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 $\frac{3}{4}(\text{model-dim}^{2})$.

Baselines.
We compare our approach with full attention, Gated Linear Attention (GLA) yang2024gated, DeltaNet schlag2021linear; yangparallelizing. To ensure fairness, we enhance both GLA and DeltaNet with the same sliding window attention. As pointed out in previous work linforgetting; xiong2023effective; men2024base, a large RoPE su2023roformer base is critical for transformers in long-context training, thus we adopt a large RoPE base of 1 million for training with 32K token contexts whenever softmax attention is used.
Tab. [7](https://arxiv.org/html/2505.23884v1#A3.T7) 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-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 a 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$.

**Table 7: 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-sequence length on A100-40GB GPUs.**
|  | State size | Train TPS | Update Rule | Memory read-out |
| --- | --- | --- | --- | --- |
| Transformer | – | 4.1K | – | – |
| Transformer SWA | – | 6.4K | – | – |
| Per-token recurrence |  |  |  |  |
| GLA SWA | $384d$ | 5.0K | $\displaystyle\mathbf{S}_{t}\leftarrow\mathbf{S}_{t-1}\mathrm{Diag}(\boldsymbol<br>{\alpha}_{t})+\mathbf{v}_{t}\mathbf{k}_{t}^{\top}$ | $\displaystyle\mathbf{o}_{t}=\mathbf{S}_{t}\mathbf{q}_{t}$ |
| DeltaNet SWA | $128d$ | 5.1K | $\displaystyle\mathbf{S}_{t}\leftarrow\mathbf{S}_{t-1}(\mathbf{I}-\beta_{t}<br>\mathbf{k}_{t}\mathbf{k}_{t}^{\top})+\beta_{t}\mathbf{v}_{t}\mathbf{k}_{t}^{\top}$ | $\displaystyle\mathbf{o}_{t}=\mathbf{S}_{t}\mathbf{q}_{t}$ |
| Large-chunk recurrence |  |  |  |  |
| Ours GD | $2304d$ | 5.0K | $W\leftarrow\mathrm{L2norm}(W-\sum_{i}^{b}\eta_{i}\nabla_{W}\mathcal{L}_{i})$ | $\displaystyle\mathbf{o}_{t}=f_{W}(\mathbf{q}_{t})$ |
| Ours Momentum | $2304d$ | 4.9K | $M\leftarrow\beta M+\sum_{i}^{b}\eta_{i}\nabla_{W}\mathcal{L}_{i};\;W\leftarrow<br>\mathrm{L2norm}(W-M)$ | $\displaystyle\mathbf{o}_{t}=f_{W}(\mathbf{q}_{t})$ |
| Ours Muon | $2304d$ | 4.3K | $M\leftarrow\beta M+\sum_{i}^{b}\eta_{i}\nabla_{W}\mathcal{L}_{i};\;W\leftarrow<br>\mathrm{L2norm}(W-\mathrm{Muon}(M))$ | $\displaystyle\mathbf{o}_{t}=f_{W}(\mathbf{q}_{t})$ |

Training Details.
We trained models at two scales using a sequence length of 32,768 tokens:

- •
760M parameters: We use 24 stack blocks, with model dimension as $1536$. All models are trained for 40B tokens (40,960 steps) with a sliding window of 2048 tokens and a batch size of 1 million tokens. Each experiment ran on 32 A100-40GB SXM GPUs for approximately 20 hours.
- •
3B parameters: We use 25 stack blocks, with model dimension as $3072$. All models are 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 on 64 A100-40GB SXM GPUs for approximately 50-60 hours.

For both scales, we used a base learning rate of $1\times 10^{-3}$
with a cosine decay scheduler and 1024 warmup steps. All models were randomly initialized with a standard deviation of 0.02.

Results. Detailed results on the RULER benchmark hsiehruler are presented in Tables [8](https://arxiv.org/html/2505.23884v1#A3.T8) and  [9](https://arxiv.org/html/2505.23884v1#A3.T9). We evaluated models on S-NIAH-1, S-NIAH-2, and S-NIAH-3 tasks, which represent varying difficulties of the single "needle in a haystack" retrieval. We also report performance on NIAH-MultiKey-1, NIAH-MultiQuery, and NIAH-MultiValue. Other RULER tasks are not reported as the full attention baseline also achieved trivial results beyond a 16K sequence length.

**Table 8: RULER benchmark results for Single Needle in a Haystack (S-NIAH) tasks. * Our method with two heads (default is four).**
|  | S-NIAH-1 | S-NIAH-2 | S-NIAH-3 | Average |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | 4K | 8K | 16K | 32K | 4K | 8K | 16K | 32K | 4K | 8K | 16K | 32K | 4K | 8K | 16K | 32K |
| 760M parameters |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| Transformer | 99.2 | 96.6 | 85.2 | 68.0 | 100 | 100 | 85.8 | 82.2 | 81.0 | 73.8 | 74.8 | 36.8 | 93.4 | 90.1 | 81.9 | 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 |
| GLA + SWA | 51.8 | 26.2 | 14.4 | 8.6 | 55.8 | 26.4 | 15.8 | 7.8 | 58.0 | 23.8 | 16.2 | 5.0 | 55.2 | 25.5 | 15.5 | 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 |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| Transformer | 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 | 100 | 52.8 | 26.0 | 13.2 | 100 | 51.8 | 29.6 | 14.4 | 98.0 | 54.4 | 27.6 | 12.4 | 99.3 | 53.0 | 27.7 | 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 Momentum | 99.4 | 97.0 | 98.6 | 93.4 | 100 | 75.6 | 39.6 | 15.0 | 91.8 | 63.0 | 27.8 | 13.4 | 97.1 | 78.5 | 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 |

**Table 9: Performance on Multi-Key (MK-NIAH), Multi-Query (MQ-NIAH), and Multi-Value (MV-NIAH) Needle in a Haystack tasks from the RULER benchmark. * Our method with two heads (default is four).**
|  | MK-NIAH | MQ-NIAH | MV-NIAH | Average |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | 4K | 8K | 16K | 32K | 4K | 8K | 16K | 32K | 4K | 8K | 16K | 32K | 4K | 8K | 16K | 32K |
| 760M parameters |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| Transformer | 63.8 | 72 | 71.4 | 54 | 33.4 | 28.9 | 24 | 23.1 | 27.95 | 24 | 20.5 | 27.35 | 41.7 | 41.6 | 38.6 | 34.8 |
| DeltaNet+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 | 35 | 23.65 | 14.1 | 7.45 | 20.7 | 22.05 | 12.7 | 6.85 | 38.8 | 26.8 | 14.5 | 7.1 |
| Ours Momentum | 62 | 41 | 21.2 | 10.4 | 35.3 | 24.95 | 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 |
| Ours Muon | 62.8 | 46.6 | 22 | 8.6 | 37.7 | 26.55 | 15.7 | 7.1 | 28.35 | 23.15 | 13.6 | 6.85 | 42.9 | 32.1 | 17.1 | 7.5 |
| 3B parameters |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| Transformer | 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 |
| GLA 3B | 78 | 45.8 | 28.6 | 14.4 | 50.05 | 28.05 | 19 | 10.7 | 29.4 | 21.4 | 16.75 | 9.9 | 52.5 | 31.8 | 21.4 | 11.7 |
| DeltaNet SWA | 75.8 | 57.4 | 34.2 | 17.8 | 66.25 | 33.05 | 21.45 | 13.45 | 43.7 | 23.2 | 18.85 | 13.2 | 61.9 | 37.9 | 24.8 | 14.8 |
| Ours Momentum | 96.2 | 59.6 | 35 | 17.2 | 87.05 | 40.25 | 25.6 | 13.2 | 88.08 | 30.65 | 21.9 | 12.3 | 90.4 | 43.5 | 27.5 | 14.2 |
| Ours Muon | 75.2 | 69.2 | 46.2 | 25.2 | 44.75 | 39.1 | 24.9 | 19 | 26.55 | 29.1 | 25.05 | 19.3 | 48.8 | 45.8 | 32.0 | 21.2 |

### C.3 Autoregressive Video Diffusion

We fine-tune the pretrained Wan 2.1 wang2025wan text-to-video diffusion model into an autoregressive video diffusion model, that generates videos by iteratively denoising successive chunks of video frames.

Model details. The original Wan 2.1 is a 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 a $2\times 2\times 1$ patchification layer to 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 a 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 a 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 are combined within 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-initialized 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-weight lr initialization of $0.001$ by setting ‘const_lr_bias’ in Algorithm [1](https://arxiv.org/html/2505.23884v1#algorithm1) to $\mathrm{softplus}(\text{const_lr_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.

Datasets. We fine-tune the model using an internal, filtered proprietary collection of videos, each accompanied by a short text prompt generated by a visual language modelchen2024internvl.

Training details.
Following esser2403scaling; wang2025wan, we use time-step shifting (scale factor 3.0) and logit-normal denoising loss weighting (mean=0.5, std=1.0). We also apply an exponential moving average with a decay rate of 0.995 to the model weights. Each 5-second video (16 FPS, 480×832 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-forcing with an interleaved noisy-clean chunk sequence (see Section [4.3](https://arxiv.org/html/2505.23884v1#S4.SS3)).

- •
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\times 10^{-5}$ with a 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 an interleaved sequence of 81 latent frames (approximately 126K tokens including noisy chunks). Training for 5-second videos takes $\sim$20 seconds per iteration on 64 A100 80GB SXM GPUs (or $\sim$10 seconds on 64 H100 80GB SXM GPUs).
- •
14B Parameter Model:. To manage GPU memory usage, noisy chunks are not repeated in this setting. We train the model on five-second videos with a batch size of 64 for 5000 iterations with a base learning rate of $5\times 10^{-6}$, and use a sequence parallel size of 2 GPUs. This phase takes $\sim$80 seconds per iteration on 64 A100 GPUs. The model is then fine-tuned on 8.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 $\sim$80 seconds per iteration on 64 H100 GPUs.

Baselines.
We compare our method with three baselines: sliding window attention, Mamba2 dao2024transformers with sliding window attention, and full block-wise causal attention, where the window attention in the baselines is implemented the same as in our model.
For the Mamba2 layer, we follow wang2024mambainllama to apply the original projected $k$, $q$, and $v$ as $B$, $C$, and $x$, respectively. The Mamba2’s state is updated token-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 FlexAttention he2024flexattention.

Evaluation. We validate 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](https://arxiv.org/html/2505.23884v1#S5.F6). Figure [6](https://arxiv.org/html/2505.23884v1#S5.F6)(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 [6](https://arxiv.org/html/2505.23884v1#S5.F6)(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 [6](https://arxiv.org/html/2505.23884v1#S5.F6)(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 zhao2023unipc sampler with 50 steps, classifier-free guidance of 5.0, and a timestep shift of 3.0.

### C.4 Experiment details in Figure 1

Fig. [1](https://arxiv.org/html/2505.23884v1#S2.F1)(c) shows results for training a 760M-parameter LaCT language model. We employ a SwiGLU MLP fast weight with the Muon test-time optimizer. To scale the fast weight size, we fix the intermediate dimension of the fast-weight 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 32,768-token sequence, averaged over 76K sequences from the Book3 dataset.

Fig. [1](https://arxiv.org/html/2505.23884v1#S2.F1)(d) uses the object-level novel view synthesis experiment. All models consist of 14 stacked blocks with a fixed model dimension of 768 and were trained for 167 billion tokens. Training time (wall-clock) is measured on an A100-40GB SXM GPU.

## Appendix D Details for Mamba Baselines

Mamba is an efficient model architecture, it is logically similar to a linear TTT taking per-token linear update rule of 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](https://arxiv.org/html/2505.23884v1#S5.F8) and Fig. [6](https://arxiv.org/html/2505.23884v1#S5.F6).
In this section, we detailed the experimental setup.

We take the official Mamba-2 implementation(^4^44https://github.com/state-spaces/mamba) 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:

$$ $\displaystyle X,B,C,\delta$ $\displaystyle=\mathrm{Linear}(u)$ $\displaystyle\delta$ $\displaystyle=\mathrm{softplus}(\delta+\delta_{\mathit{init}})$ $\displaystyle H_{t}$ $\displaystyle=\exp(-\delta_{t})H_{t-1}+\delta_{t}B_{t}^{T}X_{t}$ $\displaystyle y_{t}$ $\displaystyle=C_{t}H_{t}$ (23) $$

where $X$, $B$, $C$ is of shape $(L,d)$, and $\delta$ is of shape $(L,1)$. $H_{t}$ is a matrix state of shape $(d,d)$.

Transferring the above formula to a standard linear-attention / TTT / DeltaNet notations, it is equivalent to:

$$ $\displaystyle V,K,Q,\mathit{lr}$ $\displaystyle=\mathrm{Linear}(\mathit{input})$ $\displaystyle\mathit{lr}$ $\displaystyle=\mathrm{softplus}(\mathit{lr}+\mathit{lr}_{\mathit{init}})$ $\displaystyle W_{t}$ $\displaystyle=\exp(-\mathit{lr}_{t})W_{t-1}+\mathit{lr}_{t}K_{t}^{T}V_{t}$ $\displaystyle O_{t}$ $\displaystyle=Q_{t}W_{t}$ (24) $$

We will denote the above equations as $O=\mathrm{Mamba}(\mathit{input})$.

We use the multi-head design as in Transformer’s multi-head attention.
Multiple independent Mamba-2 layer are run in parallel and their outputs are concatenated.
Suppose the number of heads is $\mathit{nh}$, the formula is:

$$ $\displaystyle O^{k}$ $\displaystyle=\mathrm{Mamba^{k}}(\mathit{input})$ $\displaystyle O$ $\displaystyle=[O^{1},\ldots,O^{\mathit{nh}}]$ (25) $$

where each $\mathrm{Mamba^{k}}$ is a Mamba with its own parameters.
In Mamba-2’s terminology, this design is equivalent to setting the number of ‘groups’ to be the same as 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-head 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 use a similar ‘apply’ operation as in LaCT that only updates 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~{}\text{(num heads)}\times 192^{2}~{}\text{(head dim)}\times 2~{}\text{(
bidir)}$, matches LaCT with a standard large-chunk large-weight linear attention of dimension $768$ ($768~{}\text{input dim}\times 768~{}\text{intermediate dim}$ in Fig. [8](https://arxiv.org/html/2505.23884v1#S5.F8).
We take $\mathit{lr}_{\mathit{init}}{=}-4.6$, which corresponds to a $0.01$ initialized learning rate (i.e., $\mathrm{softplus}(-4.6)=0.01$).

For the autoregressive video diffusion task, we apply a unidirectional Mamba over the flattened video tokens. As mentioned in Sec [5.3](https://arxiv.org/html/2505.23884v1#S5.SS3), we follow wang2024mambainllama to inherit the Wan’s self-attention projected $k$, $q$, 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~{}\text{(num heads)}\times 128^{2}~{}\text{(head dim)}$.8*19
We take $\mathit{lr}_{\mathit{init}}{=}-4.6$, which corresponds to a $0.01$ initialized learning rate.

## Appendix E 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 a 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 on each device, logically the same as 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 [3](https://arxiv.org/html/2505.23884v1#algorithm3) 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.

Figure: Algorithm 3 Large Chunk Test-Time Training Layer with Context Parallel Sharded inside chunk Pseudocode

## Appendix F Details of LaCT Tensor Parallelism Implementation

Beyond Context Parallelism, our large-chunk Test-Time 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 in methods like DeepSpeed Ulysses jacobs2023deepspeed.

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 a gather-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 of TTT operation. Algorithm [4](https://arxiv.org/html/2505.23884v1#algorithm4) provides pseudocode detailing this tensor parallelism implementation, omitting minor details like padding. While this gather-then-scatter method effectively enables head-sharded tensor parallelism, more sophisticated communication strategies fang2024usp; jacobs2023deepspeed could potentially be employed to further optimize communication overhead.

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.

Figure: Algorithm 4 Large Chunk Test-Time Training Layer with Tensor Parallelism by sharding heads Pseudocode

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 4
presentation: 2
questions: Most of the major weaknesses presented are questions. If the authors can clarify these questions, then I will gladly increase my rating.

Why the related work and many other ablations presented in the appendix?
soundness: 3
strengths: 1. This particular method enables large chunk updates enabling higher FLOPs utilization. 
2. The authors claims to implement the LaCT layer simply using PyTorch without custom CUDA kernels.
3. The authors present metrics highlighting the throughput.
4. The empirical results shows the validity of the method on various tasks requiring long context updates.
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.
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.
2.  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.
3. The positional embedding interaction with the fast weights seems ambiguous, authors need to clarify this for the fast weight layer.
4. 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.

Minor Concerns:
1. The paper needs improvement in the writing, it is quite difficult to follow the flow of the paper. 
2. The paper should have compared against other methods such as the current state  RNNs etc. 
3. Few of the abbreviations are assumed to be known, the authors should explicitly state them eg: SMs in Section 2.2.

[A]. https://arxiv.org/pdf/2407.04620

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
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.

2. 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.
soundness: 4
strengths: 1. The proposed approach is general and can be readily adapted to a variety of domains and applications.

2. The evaluation across three domains validates its scalability, with the view synthesis task demonstrating near-transformer performance at much higher efficiency.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
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?

2. 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)?
soundness: 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.

2. 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.

3. 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.
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.
The 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.
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. 
2. 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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. How does LaCT compare against previous TTT methods, given its much higher propensity for scalability?
2. Why was LaCT trained on proprietary data when any video diffusion dataset could have sufficed? 

Please 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.
soundness: 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.
- **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).
- **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.
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.  

LaCT 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.
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.  
- **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.  
- **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.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 3
questions: 1. Is it possible to ablate the effect of the chunk size on performance?
2. Can you clarify the end-to-end complexity where the chunk size equals the full sequence length?
soundness: 3
strengths: 1. The suggested solution of combining large chunks with SWA is effective, simple and easy to implement.
2. Evaluation is very robust spanning multiple modalities and tasks, comparing to relevant baselines.
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.
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.
2. 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.

[1] An Empirical Study of Mamba-based Language Models


Decision: Accept (Poster)