T5: Bridging the Gap Between Softmax and Linear Attention via Test-Time Training
Linearizing Vision Transformer with Test-Time Training
The paper introduces T5 (Transformer To Test-Time Training), a novel approach to linearizing pretrained Vision Transformers by replacing Softmax attention with Test-Time Training (TTT) layers. By aligning architectures and representational properties, the authors successfully convert large-scale models like Stable Diffusion 3.5 into linear-complexity variants with minimal fine-tuning.
TL;DR
Researchers have unlocked a "shortcut" to linearize heavy-duty Transformers like Stable Diffusion 3.5 and DiT. By using Test-Time Training (TTT) as a structural proxy for Softmax attention, the proposed T5 (Transformer To Test-Time Training) framework allows for full weight inheritance. The result? Linear complexity, a 1.47x speedup at 2K resolution, and conversion that takes only one hour of fine-tuning.
Problem: The "Representational Gap" in Linearization
The machine learning world is currently obsessed with sub-quadratic architectures (SSMs, Linear Attention, etc.), but there is a catch: training them from scratch is incredibly expensive.
Naturally, we want to "linearize" existing giants like Llama or Stable Diffusion. However, standard linear attention (using kernels) usually fails to inherit weights effectively. The authors argue this is due to a fundamental mismatch:
- Softmax Attention is essentially a two-layer dynamic MLP with nonlinearities.
- Standard Linear Attention is a single-layer dynamic linear transform.
This structural gap makes it impossible for linear models to "understand" the weights learned by a Softmax-based Transformer, leading to poor performance or divergence during fine-tuning.
Methodology: The TTT Breakthrough
The core insight of this paper is that Test-Time Training (TTT) is the "missing link." TTT treats sequence modeling as an online learning problem where internal weights are updated for every input.
1. Structural Alignment
Unlike kernel-based linear attention, TTT allows for a multi-layer internal model. A two-layer TTT-MLP matches the mathematical structure of Softmax attention (shown in the figure below). This allows the TTT layer to "swallow" the pretrained Q, K, and V weights directly.
Figure 1: Structural similarity between Softmax Attention and TTT enables direct weight inheritance.
2. Representational Alignment: The "Secret Sauce"
Even with the right structure, two hurdles remain:
- Shift-Invariance: Softmax is naturally invariant to constant shifts in keys (it subtracts the max). TTT is not. The authors found that pretrained keys have a massive "bias" (ratio ≈ 0.5) that causes TTT gradients to explode. They solve this using Instance Normalization on the keys.
- Locality: Softmax naturally focuses on local neighbors. To give TTT this same "vision," the authors added Depthwise Convolutions (DWC) on queries and keys.
Experiments: Rapid Conversion, Linear Scaling
The authors put T5 to the test on ImageNet classification and Diffusion Transformers (DiT/SD3.5).
- Stable Diffusion 3.5-Medium: Conventionally, linearizing such a model would take weeks. T5 achieved it in 1 hour on 4x H20 GPUs. At 2048 resolution, it is 1.47x faster than the original Softmax model.
- DiT-XL/2: The model recovered nearly full FID performance with only 8 epochs of fine-tuning (0.57% of the original training budget).
Figure 2: The efficiency gap between Softmax (DeiT/DiT) and T5 widens significantly as image resolution grows.
Detailed Results on SD 3.5
In text-to-image benchmarks (DPG-Bench and GenEval), SD3.5-T5 actually improved over the base model in some metrics (e.g., moving GenEval from 0.66 to 0.69) while drastically reducing latency.
Table 1: Competitive quality vs. SD3.5-Medium with significantly lower Latency.
Critical Analysis & Conclusion
The value of this work lies in its minimalism. It proves that we don't need complex distillation or multi-stage "activation alignment" to escape the quadratic bottleneck.
Takeaways:
- TTT is a powerful proxy: Its internal model capacity allows it to mimic Softmax better than any fixed kernel.
- Key Normalization is critical: Without centering the keys, weight inheritance from pretrained Transformers will likely crash.
Limitations: While T5 is faster at high resolutions, at low resolutions, the overhead of the TTT internal gradient updates can slightly mitigate the speed gains. Future work may focus on optimizing the TTT "inner loop" for even faster inference at all scales.
Future Outlook: T5 provides a blueprint for migrating the entire industry's library of pretrained Transformers to linear-time architectures without the "compute tax" of retraining.