PaperlayerView exact revision on arXiv

Source-linked AI summary

Test-Time Training Done Right

Tianyuan Zhang, Sai Bi, Yicong Hong, Kai Zhang, Fujun Luan, Songlin Yang, Kalyan Sunkavalli, William T. Freeman, Hao Tan

arXiv:2505.23884v1cs.LGcs.CLcs.CV

TL;DR

Existing test-time training methods face inefficient GPU use and limited suitability for structured, long-context data. LaCT addresses these issues with large-chunk updates and window attention, achieving superior or competitive performance across novel view synthesis, language modeling, and autoregressive video diffusion, including sequences up to 56,000 visual tokens and 1M-token novel-view inputs.

  • Problem

    Existing test-time training methods are inefficient on modern GPUs, while small online chunks are poorly suited to data modalities with intra-chunk structure and locality.

  • Method

    LaCT uses large chunks to update fast weights and integrates local window attention to model structure within each chunk.

  • Results

    LaCT achieves superior or competitive performance across novel view synthesis, language modeling, and autoregressive video diffusion, processing up to 1M novel-view tokens and 56,000 visual tokens.

  • Takeaways & Limitations

    LaCT provides an efficient, scalable framework for long-sequence modeling across diverse modalities using native PyTorch implementations and supports larger fast-weight states and flexible test-time optimizers.

  • Takeaways & Limitations

    Autoregressive video diffusion lacks a reliable, distinguishable metric for measuring scalability, so the paper reports validation loss as a common evaluation choice.

Abstract

from arXiv · show

Test-Time Training (TTT) models context dependencies by adapting part of the model's weights (referred to as fast weights) during inference. This fast weight, akin to recurrent states in RNNs, stores temporary memories of past tokens in the current sequence. Existing TTT methods struggled to show effectiveness in handling long-context data, due to their inefficiency on modern GPUs. The TTT layers in many of these approaches operate with extremely low FLOPs utilization (often <5%) because they deliberately apply small online minibatch sizes (e.g., updating fast weights every 16 or 64 tokens). Moreover, a small minibatch implies fine-grained block-wise causal dependencies in the data, 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 using 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). It improves hardware utilization by orders of magnitude, and more importantly, facilitates scaling of nonlinear state size (up to 40% of model parameters), hence substantially improving state capacity, all without requiring cumbersome and error-prone kernel implementations. It also allows easy integration of sophisticated optimizers, e.g. Muon for online updates. We validate our approach across diverse modalities and tasks, including novel view synthesis with image set, language models, and auto-regressive video diffusion. Our approach can scale up to 14B-parameter AR video diffusion model on sequences up to 56K tokens. In our longest sequence experiment, we perform novel view synthesis with 1 million context length. We hope this work will inspire and accelerate new research in the field of long-context modeling and test-time training. Website: https://tianyuanzhang.com/projects/ttt-done-right

1 Introduction

LaCT addresses the poor GPU efficiency and limited data-structure flexibility of existing test-time training by updating fast weights over very large chunks and combining them with window attention. The approach is validated across image-set view synthesis, language modeling, and autoregressive video diffusion, including contexts up to 1M tokens and 56K visual tokens.

  • Motivation: Existing TTT methods often use small minibatches that achieve below 5% peak FLOPs utilization, limiting efficient long-context scaling.Frequent updates provide poor parallelism and low compute intensity, especially for large nonlinear fast weights.
  • LaCT: LaCT updates fast weights over chunks ranging from 2048 to 1M tokens, reaching up to 70% GPU utilization on NVIDIA A100s.Its pure PyTorch implementation also supports integrating online optimizers such as Muon.
  • LaCT: Window attention complements large-chunk updates by modeling local dependencies while LaCT captures longer-range context.This hybrid design accommodates data whose chunks contain meaningful internal structure.
  • Experiments: LaCT is evaluated on novel view synthesis, language modeling, and autoregressive video diffusion across different modalities and data structures.The evaluations include image sets, text sequences, and video.
  • Experiments: The models process up to 1M tokens for novel view synthesis, achieve competitive language-modeling performance, and generate consistent videos up to 56,000 visual tokens.The video result uses a 14-billion-parameter autoregressive diffusion model.
  • Conclusion: The paper presents LaCT as an efficient and scalable framework for long-sequence modeling across diverse modalities.The authors emphasize that its implementation avoids dependence on low-level hardware-specific kernels.

2 Preliminary

TTT stores context in fast weights updated online, but frequent small updates are memory- and compute-inefficient on modern GPUs. LaCT instead uses large chunks, combining fast-weight updates with window attention to improve parallelism, state scaling, and modeling of structured data.

  • Test-Time Training: TTT uses rapidly adaptable fast weights to store in-context information while slow model weights remain frozen during inference.Fast weights are updated through self-supervised objectives that associate keys with values.
  • Test-Time Training: The TTT update minimizes a loss between transformed keys and values, then applies the updated fast weights to queries to produce outputs.The standard per-token procedure iteratively performs update and apply operations across the sequence.
  • Efficient Implementation: Customized kernels can reduce memory traffic for frequent updates but are difficult to use with large nonlinear states and require lengthy development cycles.The cited constraints include communication requirements and incompatibility with nonlinear SwiGLU fast weights or Muon updates.
  • Efficient Implementation: Small fast-weight matrices or chunks make PyTorch TTT memory-bound because their compute-to-memory ratio remains far below GPU theoretical peaks.For example, h = 64 or b = 16 can limit compute usage relative to a 290 FLOPs-per-byte H100 peak.
  • Large-Chunk Design: Large chunks from 2048 to 1M tokens improve throughput, support state-to-parameter ratios of at least 40%, and can improve results as state size grows.The reported ratio is an order of magnitude larger than previous methods’ 0.1% to 5%.
  • LaCT Block: LaCT blocks combine window attention for local dependencies with large-chunk TTT for compressing history into fast weights.Sequence-length parallelism is important for high occupancy when long-context workloads have small batch sizes.

3 LaCT Model Architecture

LaCT separates large-chunk fast-weight updates from applications, combining TTT with window attention to model non-local context while preserving local structure. Its infrequent, normalized updates improve parallelism, support nonlinear optimizers, and enable context parallelism.

  • Large-Chunk TTT Layer: LaCT splits sequences into large chunks, computes one summed-loss gradient per chunk, and applies the resulting fast weight to all queries in that chunk.All query vectors share the same updated fast weight, while updates occur infrequently and amortize their cost.
  • Large-Chunk TTT Layer: LaCT uses a SwiGLU-MLP with three matrices as its fast-weight network and a dot-product loss for online adaptation.The fast weights are W = {W1, W2, W3}.
  • Large-Chunk TTT Layer: LaCT decouples fast-weight update and apply operations, allowing adaptive chunk sizes and different orders that encode distinct attention-mask dependencies.Alternating operations produce block-wise causal masks, while switching their order shifts the mask without leaking future information.
  • Non-Linear Update of Fast-Weight: L2 normalization and nonlinear Muon updates improve fast-weight stability while preserving the efficiency of large-chunk TTT.Muon normalizes gradient spectral scale and makes token learning rates reflect relative importance within a chunk.
  • Window Attention: Window attention handles intra-chunk locality, allowing TTT to focus its fixed-size fast-weight capacity on non-local dependencies across diverse data structures.The hybrid architecture combines quadratic-compute local attention with linear-compute TTT for non-local context.
  • Context Parallelism: Large-chunk TTT supports context parallelism by sharding tokens within a chunk and aggregating shard gradients with all-reduce-sum.Novel view synthesis training observed 1% to 3% throughput overheads under this parallelism.

4 LaCT for N-Dimensional Data

LaCT adapts its chunking, masking, and local-attention choices to image sets, text sequences, and video chunks. These designs align the model with each modality’s data structure while retaining large-chunk fast-weight updates.

  • Task Overview: LaCT is evaluated on novel view synthesis, language modeling, and autoregressive video generation across different data structures.The tasks span image sets, 1D text sequences, and video-frame chunks.
  • Novel View Synthesis: Novel view synthesis tests online memory and compression through spatial compression, dense retrieval, physical reasoning, and redundant multi-view inputs.Its non-generative formulation reduces training computation and the need for extensive model parameters to store world knowledge.
  • Novel View Synthesis: The NVS model tokenizes posed images and target-view poses, applies per-image window attention, and updates fast weights using all input-image tokens.A strided block-wise causal mask applies the updated fast weight to both input and target tokens.
  • Language Modeling: For language modeling, LaCT uses chunk sizes such as 2048 or 4096 tokens, shifted block-wise causality, and sliding-window attention to preserve per-token causal dependencies.Text lacks an inherent chunk structure, so chunk size is treated as a hyperparameter.
  • Autoregressive Video Diffusion: Video training interleaves noisy and clean frame chunks, with each noisy chunk formed by adding unit Gaussian noise at a chunk-specific strength.The noise strength is represented by t_i ∈ [0, 1].
  • Autoregressive Video Diffusion: For video diffusion, LaCT updates fast weights only on clean chunks and applies them sequentially so denoising accesses previously cleaned frames.Non-overlapping windows contain two consecutive chunks and exclude attention from a clean chunk to the following noisy chunk.

5 Experiments

The experiments evaluate LaCT across novel view synthesis and language modeling, comparing efficiency, quality, long-context use, and design choices. Results show strong performance with large chunks, scalable fast-weight states, Muon updates, and nonlinear fast weights.

  • Experimental setup: Experiments cover novel view synthesis and language modeling, with comparisons against full attention, Perceiver-style attention, LongLRM, 3D Gaussian Splatting, GLA, and DeltaNet.The evaluations also report state size, throughput, update rules, and memory read-out mechanisms.
  • Novel View Synthesis: Up to 1M tokens, the novel-view-synthesis model processes 128 input images at 960×536 resolution and outperforms 3D Gaussian Splatting in rendering quality.On the high-resolution scene dataset, it also surpasses LongLRM, which is limited to 32 views.
  • Novel View Synthesis: LaCT achieves quality comparable to full-attention models with significantly lower prefill latency and clearly outperforms Perceiver-attention baselines.The comparison is reported in Figure 4 for novel view synthesis.
  • Language Modeling: At both 760M and 3B scale, LaCT has lower per-position loss at larger token indices and higher retrieval accuracy than GLA and DeltaNet.The Muon variant consistently outperforms the Momentum variant.
  • Analysis on Design Choices: Larger state sizes consistently improve performance, with the gap between small and large states widening as sequence length increases.The language-model experiment allocates up to 40% of model weights to fast weights.
  • Analysis on Design Choices: Muon consistently outperforms vanilla gradient descent and gradient descent with momentum in the reported test-time optimizer comparisons.The comparison includes both novel view synthesis and language-model settings.
  • Analysis on Design Choices: Nonlinear fast weights outperform linear fast weights despite smaller state sizes, while large-chunk recurrence outperforms matched-state-size per-token recurrence for view synthesis.For language modeling, large-chunk recurrence surpasses per-token baselines when combined with a larger nonlinear state and Muon updates.

6 Related Work

The related work situates LaCT among test-time-training, chunk-attention, novel-view-synthesis, and autoregressive-video-diffusion approaches. It emphasizes LaCT’s focus on large-chunk updates and long-context modeling across modalities.

  • Test-Time Training: Test-Time Training adapts fast weights online to memorize in-context information, but existing approaches often have low hardware utilization and limited state sizes.LaCT addresses these challenges with extremely large online minibatches for fast-weight updates.
  • Test-Time Training: Earlier TTT methods commonly rely on customized kernels to keep fast weights in SRAM, but these kernels are cumbersome and do not support large nonlinear states or Muon updates well.The limitation arises from communication and independence constraints during fast-weight evolution.
  • Chunk Attention and Recurrence: Chunk-attention and recurrence models such as GAU, MEGA, MEGALODON, and InfiniAttention are related, but InfiniAttention’s delta-rule update is described as limited in expressivity.Block-Recurrent Transformer likewise uses chunk-level register-token memory updates, making it conceptually similar to the Perceiver-style baseline.
  • Novel View Synthesis: Data-driven novel-view-synthesis methods often struggle with densely sampled scenes containing over 100 input images, whereas the paper reports LaCT outperforming or matching 3D Gaussian Splatting on challenging scenes.This positions LaCT for high-view-count scene reconstruction.
  • Autoregressive Video Diffusion: Autoregressive video diffusion methods address video generation with sequential conditioning, including teacher forcing, progressive noise injection, and frame-independent noise strategies.The paper places LaCT within this broader effort to improve autoregressive video modeling.

7 Limitation

The paper identifies scope and evaluation limitations across rotation invariance, task diversity, language-model reasoning and scaling, and video-diffusion scalability measurement.

  • LaCT’s SwiGLU and Linear Fast Weight components lack the rotation invariance exhibited by softmax and linear attention.The practical implications of this absence remain underexplored.
  • The evaluation covers three tasks, while unposed reconstruction—a more challenging setting—is not explored.The novel-view synthesis task uses input pose information.
  • Language-model experiments do not explore reasoning capacity or parameter-size scalability because of computation limitations.The paper notes that reasoning ability requires a certain amount of training compute beyond its budget.
  • Video-diffusion scalability lacks a reliable and distinguishable evaluation metric.The paper reports validation loss, a common choice for evaluating video-generation scalability, but contrasts this with language-model perplexity and NVS PSNR.

8 Conclusion

The paper presents LaCT, combining large-chunk test-time training with window attention for long-context modeling. Across three modalities, it reports superior or competitive performance, high GPU efficiency, scalable state size, and flexible design using native PyTorch.

  • LaCT combines large-chunk test-time training for long context with window attention for local structure.
  • LaCT achieves superior or competitive performance across novel view synthesis, language modeling, and autoregressive video diffusion.
  • Native PyTorch implementation provides high GPU efficiency while supporting larger state sizes and more flexible test-time-training designs.

A LaCT Model Implementation Details

The implementation uses SwiGLU fast weights, chunkwise updates, optional Muon optimization, and hybrid local window attention. These choices define the state size, computation, initialization, and sequence-processing behavior of LaCT.

  • Fast-weight network: The fast-weight network is a bias-free SwiGLU-MLP with three matrices W1, W2, and W3.Its forward pass applies SiLU to one projection, multiplies elementwise with another projection, and maps through W2.
  • State size: The total state size is nh × hd^2 × r, and it can increase by reducing the number of heads or increasing the intermediate-dimension multiplier.
  • Chunkwise updates: Large-chunk updates process key, value, query, and learning-rate tensors iteratively, applying fast weights to queries and updating them with keys and values.
  • Muon optimizer: Muon orthogonalizes matrix gradients through Newton-Schulz iterations, while the implementation uses five iterations with constants a = 3.4445, b = −4.7750, and c = 2.0315.
  • Hybrid attention: The layer combines local softmax window attention with LaCT outputs, then applies RMSNorm, per-head scaling, head merging, and output projection.
  • Momentum: Per-token momentum coefficients βi are predicted through a sigmoid-linear function, averaged across each chunk, and used in the optimizer update.

B.1 LaCT Architecture for Novel View Synthesis

LaCT adapts fast weights from posed input images, then applies the compressed scene representation to render novel views. Its image-wise bidirectional attention and all-input-token updates align the architecture with set-structured visual data.

  • LaCT first compresses posed input images into fast weights during an Update Operation, then renders novel views through an Applying Operation.
  • Novel-view synthesis converts posed input images and target views into tokens, with camera poses represented through dense ray information and Plücker ray embeddings.
  • Window attention is bidirectional within each image and independent across images, while TTT updates use all input-image tokens and apply the updated weight to input and target tokens.
  • During rendering, the fast weight remains fixed, so its feed-forward network stores scene information while the slow-weight network stores world knowledge such as physical rendering rules.

B.3 LaCT Architecture for Autoregressive Video Diffusion

For autoregressive video diffusion, LaCT processes frame chunks with interleaved noisy and clean tokens, updating fast weights from clean chunks and combining them with local causal attention. The design supports chunk-level denoising while preserving within-chunk bidirectional dependencies.

  • Autoregressive video diffusion iteratively denoises sequential frame chunks conditioned on previously generated clean frames, with each chunk spanning thousands of visual tokens.
  • Each noisy video chunk is formed by adding unit Gaussian noise with a chunk-specific strength to the corresponding clean chunk.
  • Repeating each clean video chunk at two noise levels increases training token utilization from 50% to about 67%, while repetition beyond twice would reduce sample diversity.
  • LaCT fast weights update exclusively from clean chunks and are applied to the current clean chunk and subsequent noisy chunks, while local attention uses a two-chunk block-causal window.
  • The NVS experiments use 24 LaCT blocks with model dimension 768 and a 312M-parameter model, trained with specified multi-stage datasets and optimization settings.

C.2 Language Modeling

LaCT language models combine large-chunk TTT with sliding-window attention while preserving causal next-token prediction through shifted block-wise masking. The experiments evaluate long-context loss and retrieval across 32K-token training sequences.

  • The language-modeling evaluation measures per-token loss over sequences and RULER retrieval accuracy across lengths up to the trained context.
  • RULER results cover single-needle, multi-key, multi-query, and multi-value retrieval tasks, excluding tasks where full attention already achieved trivial results beyond 16K tokens.

E Details of LaCT Context Parallelism Implementation

LaCT supports context parallelism by sharding tokens within each large chunk across devices. Local fast-weight gradients are aggregated with differentiable all-reduce before the global update, enabling million-token NVS training.

  • Context parallelism partitions the sequence-length dimension across devices, while local window attention supports this partitioning directly.
  • The implementation assumes k, v, and learning-rate tensors are sequence-sharded, making the effective chunk size equal to the local length multiplied by the context-parallel group size.
  • After local fast-weight gradients are computed, differentiable all-reduce-sum aggregates them before Muon or other weight updates.
  • LaCT shards tokens within each large chunk, allowing the update to span multiple devices while per-token apply operations remain independently parallelizable.
  • The authors use intra-chunk context parallelism for view synthesis with chunks exceeding half a million tokens and training sequences longer than one million tokens.

F Details of LaCT Tensor Parallelism Implementation

LaCT supports tensor parallelism by sharding TTT heads across devices while gathering sequence-sharded inputs to reconstruct full context for each head subset. The approach was used in autoregressive video generation to train 14-billion-parameter diffusion models on sequences exceeding 100K tokens.

  • LaCT supports Tensor Parallelism by sharding TTT heads across multiple devices, similarly to DeepSpeed Ulysses.
  • TTT inputs initially sharded by sequence length are gathered into full-sequence context and scattered across local heads within the tensor-parallel group.The transformed tensors include Q, K, V, and TTT learning rates.
  • Four TTT heads were sharded across four local GPUs, enabling training of 14-billion-parameter diffusion models with sequence lengths exceeding 100K tokens.
  • The gather_scatter operation transforms [B, N_full, L_local, D] into [B, N_local, L_full, D] by all-gathering one dimension and slicing another.Slicing indices are determined from the device rank, group size, and scatter dimension.
  • The implementation represents fast weights as sharded per-head tensors, with w1 and w3 shaped [nh, d, dh] and w2 shaped [nh, dh, d].Here, nh denotes the number of local heads.
  • After TTT computation on sharded heads, the reverse transformation gathers the head dimension and scatters the sequence dimension to restore sequence sharding.This reconstructs the full head dimension while returning each device to its local sequence shard.
Loading 2505.23884v1…