Test-Time Training Done Right (LaCT)

TL;DR

LaCT (Large-Chunk Test-Time Training) inverts the prevailing TTT design choice — instead of updating fast weights every 16–64 tokens with custom GPU kernels, it updates once per 2K–1M-token chunk in plain PyTorch. The result is a 10–20× jump in FLOPs utilization (sub-5% → up to 70% on A100), enabling a SwiGLU-MLP fast-weight network whose state can reach 40% of model parameters (vs. 1–10% in prior TTT). The recipe ships as a hybrid block (window attention + large-chunk TTT + FFN with optional Muon update rule) and is validated on novel view synthesis (up to 1M-token context), 760M/3B language models, and a 14B autoregressive video diffusion model adapted from Wan 2.1 generating 56K-token (8.8s) clips.

Key claims

Existing TTT methods get <5% FLOPs utilization on modern GPUs because small (16–64 token) mini-batches make the fast-weight update memory-bound; the matmul of fast weight (d × d) by mini-batch input (b × d) needs b ≫ 1 to escape the bandwidth bound [§2.2, Eq. 3].
Increasing the chunk size to 2K–1M tokens raises utilization to ~70% on A100 in pure PyTorch — no custom kernels needed [§1, Fig. 1a].
Improved utilization unlocks scaling the non-linear fast-weight state size to 40% of total model parameters, vs. 1–10% in prior methods like DeltaNet / GLA / Titans [§3, Table 3].
The update/apply ordering of the chunk-wise TTT layer is a knob analogous to attention masks: full-sequence apply-then-update ≈ full attention; alternating update/apply ≈ block-wise causal; update-then-apply gives the shifted mask needed for strict causal LM training; update-on-subset, apply-on-all gives a strided block-wise causal mask used for AR video diffusion [§3.1, Fig. 3].
The fast-weight network is a 3-matrix SwiGLU-MLP with L2 weight normalization (post-layer-norm analog along the “time-as-depth” axis); explicit weight decay is omitted [§3.2, Eq. 6, 8].
A Muon-style spectral-norm-normalized update rule (Newton-Schulz on the fast-weight gradient) plugs into LaCT and consistently outperforms momentum-based updates on long-context LM retrieval [§3.2, Eq. 9, 10; §5.2 Fig. 5].
A window-attention layer is concatenated with the large-chunk TTT layer (sharing QKV in the LM variant, GAU-style) so the TTT layer can specialize on non-local context while local intra-chunk structure is handled exactly [§3.3].
Large chunks support context parallelism by sharding tokens within a chunk and all-reduce-summing the fast-weight gradient (linearity of the gradient operator); 1–3% overhead vs. single-device, with the same logical semantics as DDP applied to fast weights [§3.4, Eq. 11].
On novel view synthesis (object-level Objaverse → GSO; scene-level DL3DV), LaCT-312M outperforms a Perceiver-register baseline and a full-attention baseline at matched parameters; 1.4 s prefill and 38.7 FPS rendering vs. 16.1 s / 2.3 FPS for full attention at the same compute scale [§5.1, Table 2].
On scene-level DL3DV with up to 1M-token context, LaCT exceeds 3D Gaussian Splatting rendering quality at the same input-view scale [§5.1].
On 760M / 3B language models trained on 60B tokens at 32K context, LaCT-Muon achieves lower per-position loss at long positions than GLA+SWA and DeltaNet+SWA, with better RULER retrieval accuracy at matched throughput (4.3K vs 5.0K TPS for the Muon variant) [§5.2, Table 3].
For autoregressive video diffusion, the bidirectional Wan 2.1 (14B) is converted by replacing every bidirectional attention layer with a LaCT layer + sliding-window attention; teacher-forced training on 8.8 s / 480×832 clips at 16 FPS (56,160 visual tokens) yields lower chunk-wise denoising loss than SWA-only and Mamba2+SWA baselines [§4.3, §5.3, Fig. 6].
The same architecture works across novel-view-synthesis (image set), language modeling (1D sequence), and video diffusion (image sequence) by choosing chunk size to match the inherent data grouping (full sequence / 2K–4K tokens / three video latent frames) [Table 1].

Method

A LaCT block is a 3-layer hybrid: (1) a window-attention layer covering local structure, (2) a large-chunk TTT layer with a SwiGLU-MLP fast-weight network of size up to 40% of model parameters, and (3) a feed-forward layer, all with residual connections. The TTT layer splits the sequence into chunks (full-sequence for image-set NVS, 2K–4K tokens for LM, three latent frames for video) and performs two operations per chunk: an update step that runs gradient descent on a dot-product key/value loss summed over all tokens in the chunk (Eq. 4–5), followed (or preceded — see Fig. 3) by an apply step that maps queries through the updated fast weights. The update rule can be plain GD + L2 weight-norm or a Muon-style Newton-Schulz spectral-norm normalization that improves numerical stability and lets the learning rate represent only relative token importance. The order of update vs. apply, and whether updates run on every chunk or only a subset, recovers full / block-causal / shifted-causal / strided masks, which is what lets the same block specialize to image-set, autoregressive-sequence, and noisy-clean-interleaved AR-video data structures. Context parallelism is exposed by sharding tokens within a chunk and all-reduce-summing the linear-in-shards fast-weight gradient.

Results

Novel view synthesis (Objaverse → GSO and DL3DV scenes): at 312M params (84M fast weights / block), LaCT beats full-attention and Perceiver baselines on PSNR while running 11× faster prefill (1.4 s vs 16.1 s) and 17× faster rendering (38.7 vs 2.3 FPS) than full attention. Scene-level training reaches 1M-token context and surpasses optimization-based 3D Gaussian Splatting at the same input scale [§5.1, Table 2].
Language modeling (Long-Data-Collections, 32K context, RoPE base 1M): 760M and 3B LaCT models trained for 40B / 60B tokens get monotonically decreasing per-position loss out to 32K tokens, beating GLA+SWA and DeltaNet+SWA at matched throughput; Muon update beats Momentum and GD ablations [§5.2, Table 3].
Autoregressive video diffusion (Wan 2.1 14B → LaCT-AR): on 8.8 s / 16 FPS / 480×832 clips (56,160 visual tokens; 107K tokens including interleaved noisy/clean chunks under teacher-forcing), lower chunk-wise denoising loss than SWA-only, Mamba2+SWA, and full block-wise causal attention across all five evaluation timesteps [§5.3, Fig. 6].
GPU utilization: ~70% peak on NVIDIA A100s in pure PyTorch at large chunk sizes, vs. the <5% reported for kernel-customized small-chunk TTT baselines [§1, Fig. 1a].

Why it’s interesting

The headline observation — “small chunk + custom kernels” was a local optimum and “large chunk + plain PyTorch” wins by an order of magnitude — flips the implicit dataflow assumption underlying every recent TTT / linear-attention variant (Mamba, GLA, DeltaNet, Titans). For Luma’s interests this is most actionable on the AR-video side: LaCT is the only filed approach in Autoregressive Video Generation where the core layer architecture itself (not just the cache / curriculum / distillation) is rebuilt to support long context, and it has been demonstrated on a 14B Wan 2.1 conversion at 56K visual tokens — directly competitive with Context Forcing: Consistent Autoregressive Video Generation with Long Context‘s slow/fast cache partitioning and MALT Diffusion: Memory-Augmented Latent Transformers for Any-Length Video Generation‘s recurrent-memory-vector approach. The hardware-utilization argument also reframes the IO-Aware Kernel Design story: the same “characterize the bottleneck, then redesign the algorithm” template that FA4 / Flash-KMeans / FlashSampling apply at the kernel level, LaCT applies one level up at the layer level — and arrives at a design that avoids custom kernels entirely.