Linearizing Vision Transformer with Test-Time Training

While linear-complexity attention mechanisms offer a promising alternative to Softmax attention for overcoming the quadratic bottleneck, training such models from scratch remains prohibitively expensive. Inheriting weights from pretrained Transformers provides an appealing shortcut, yet the fundamental representational gap between Softmax and linear attention prevents effective weight transfer. In this work, we address this conversion challenge from two perspectives: architectural alignment and representational alignment. We identify Test-Time Training (TTT) as a linear-complexity architecture whose two-layer dynamic formulation is structurally aligned with Softmax attention, enabling direct inheritance of pretrained attention weights. To further align representational properties, including key shift-invariance and locality, we introduce key instance normalization and a lightweight locality enhancement module. We validate our approach by linearizing Stable Diffusion 3.5 and introduce SD3.5-T$^5$ (Transformer To Test Time Training). With only 1 hour of fine-tuning on 4×H20 GPUs, SD3.5-T$^5$ achieves comparable text-to-image quality to the fine-tuned Softmax model, while accelerating inference by 1.32× and 1.47× at 1K and 2K resolutions. Code is available at this URL: https://github.com/LeapLabTHU/Transformer-to-TTT.

Modern vision Transformers and text-to-image generation models achieve strong performance, but their attention mechanism can be slow and expensive when processing high-resolution images. Replacing it with a more efficient mechanism is promising, yet training such models from scratch is costly and direct conversion from pretrained models often causes performance loss. In this paper, we study how to convert pretrained vision Transformers into efficient models from two perspectives: architectural alignment and representational alignment. We identify Test-Time Training (TTT) as an efficient architecture whose two-layer dynamic formulation is naturally similar to the original attention mechanism, making it possible to directly inherit pretrained attention weights. To further make the converted model behave like the original pretrained model, we introduce key normalization for better stability and a lightweight locality module. We validate our method on image classification and image generation. On Stable Diffusion 3.5, our converted model, SD3.5-T$^5$, achieves comparable image quality after only about one hour of additional training on 4 H20 GPUs, while running 1.32× faster at 1K resolution and 1.47× faster at 2K resolution.