Linearizing Vision Transformer with Test-Time Training

Large AI models like Stable Diffusion are slow because they use a complex "attention" mechanism. This mechanism becomes significantly more computationally expensive as images get larger. This quadratic scaling creates a massive bottleneck for high-resolution generation. While researchers have attempted to swap this for faster, linear-complexity alternatives, doing so usually requires expensive retraining. This often leads to a loss of the original model's intelligence.

A new paper proposes a way to bypass this. They treat the conversion as a structural alignment problem. By using Test-Time Training (TTT)—a method where the model updates its internal weights dynamically during inference—the authors found they could "inherit" weights from a pre-trained Softmax Transformer. This allows for a rapid transition to a linear-complexity architecture. It maintains high quality while slashing inference latency.

The Problem

The status quo in Vision Transformers (ViTs) relies on Softmax attention. This computes pairwise interactions between every token in a sequence. This is highly expressive but hits a wall as sequence length $N$ grows. The complexity scales at $O(N^2)$, meaning doubling the resolution quadruples the cost. Current efforts to move toward linear-complexity ($O(N)$) models typically fall into two traps.

First, there is a fundamental representational gap. Standard linear attention uses a single-layer dynamic transformation. In contrast, Softmax attention behaves like a two-layer dynamic MLP (Multi-Layer Perceptron) because of its non-linear activation. Second, transferring knowledge from a pre-trained Softmax model to a linear one is notoriously difficult. Existing methods often only inherit MLP weights. This discards the learned intelligence in the attention layers themselves. Practitioners must choose between a slow, high-quality Softmax model or a fast, low-quality linear model.

How It Works

The authors argue that we should stop approximating Softmax with simple kernels. Instead, they use Test-Time Training (TTT) as the replacement architecture. They approach the conversion through two distinct alignment phases:

1. Structural Alignment: The authors identify that TTT can be formulated as a two-layer dynamic MLP. As shown in, this structure is mathematically similar to Softmax attention.

This symmetry allows the TTT layer to inherit the actual weights from the pre-trained Softmax attention layers. This avoids the need for complex distillation (the process of transferring knowledge from a large teacher model to a smaller student model).

1. Representational Alignment: Even with matching structures, the math does not perfectly align. The authors identify two specific discrepancies:
2. Key Shift-Invariance: Softmax is invariant to constant shifts in keys. It effectively "absorbs" biases. TTT is not. The authors find that pre-trained keys carry a significant bias (a ratio of $\approx 0.5$, as seen in [, bottom]).

This causes TTT to diverge or hit NaNs (Not-a-Number errors) during fine-tuning. They solve this by applying Instance Normalization (a technique to normalize features across the sequence dimension) to center the keys. * Locality: Softmax attention has a strong inherent bias toward local pixel relationships. TTT tends to be too global. To fix this, they introduce DWCQK. This module applies depthwise convolutions (a type of efficient convolution that processes channels separately) to the queries and keys. This injects a local inductive bias back into the model. This is verified by looking at implicit attention scores, where TTT originally showed much weaker local patterns than Softmax .

Numbers

The most impressive result is the efficiency gain in text-to-image generation. The authors applied their method to Stable Diffusion 3.5, creating "SD3.5-T5." They report that with only 1 hour of fine-tuning on 4$\times$ H20 GPUs, the model achieves comparable quality to the original. Most importantly, they report a 1.32$\times$ speedup at 1K resolution and a 1.47$\times$ speedup at 2K resolution [Table 7]. This means high-resolution generation becomes significantly faster on existing hardware.

On the DiT-XL/2 model, the efficiency of the transfer is equally stark. The T5 architecture can recover performance using only 0.57% of the original training steps [Table 1]. For image classification on DeiT-Tiny, the T5 architecture reaches 71.19% Top-1 accuracy. This outperforms other linearization attempts like LiT, CLEAR, and Hedgehog under the same 30-epoch fine-tuning budget [Table 5]. As resolution increases, the FLOPs (Floating Point Operations, a measure of computational cost) advantage of the T5 architecture becomes increasingly dominant .

What's Missing

While the results are strong, some areas remain unaddressed. First, the TTT inner model introduces a small number of new parameters compared to vanilla linear attention. While the authors suggest this is negligible, the cumulative overhead might affect the absolute theoretical limit of compression in massive production models.

Second, the dependency on a pre-trained Softmax model is a prerequisite. This is a "conversion" paper, not a "pre-training" paper. If you are building a new architecture from the ground up, this method does not help you avoid the $O(N^2)$ cost of initial training. It only helps you migrate once you have a working Softmax baseline.

Finally, the paper focuses heavily on vision and diffusion tasks. While TTT has been explored in language modeling, the "structural alignment" argument is tailored to spatial inductive biases in vision. It remains to be seen if this recipe translates to the complex, long-context dependencies found in Large Language Models.

Should You Prototype This

Yes, specifically if you are managing high-resolution generative workloads. If your team struggles with the inference latency of Stable Diffusion or large-scale Vision Transformers, this is a highly actionable path. Unlike previous linearization papers that require heavy distillation, this method allows you to take an existing checkpoint and "convert" it with minimal compute.

Code is reportedly available; see the paper for the canonical link. Given that the authors achieved the SD3.5-T5 results with just one hour of fine-tuning on four H20s, the barrier to entry is low. Try it on your smallest high-res model first. Be sure to implement the Instance Normalization step to ensure training stability.

Figures from the paper