Linearizing Vision Transformer with Test-Time Training

Hugging Face Daily Papers Papers
vision-transformer test-time-training linear-attention weight-transfer text-to-image efficient-inference model-compression

Summary

This paper proposes a method to convert pretrained Softmax attention models into linear-complexity Test-Time Training (TTT) architectures, achieving comparable text-to-image quality to fine-tuned Softmax models while significantly accelerating inference. The approach is validated by linearizing Stable Diffusion 3.5, resulting in SD3.5-T^5 with 1.32x speedup at 1K resolution.

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 4timesH20 GPUs, SD3.5-T^5 achieves comparable text-to-image quality to the fine-tuned Softmax model, while accelerating inference by 1.32times and 1.47times at 1K and 2K resolutions. Code is available at https://github.com/LeapLabTHU/Transformer-to-TTT.
Original Article
View Cached Full Text

Cached at: 06/02/26, 03:24 AM

Paper page - Linearizing Vision Transformer with Test-Time Training

Source: https://huggingface.co/papers/2605.02772

Abstract

Researchers develop a method to convert pretrained Softmax attention models to linear-complexity Test-Time Training architectures through architectural and representational alignment, achieving fast inference with minimal fine-tuning.

Whilelinear-complexity attentionmechanisms offer a promising alternative toSoftmax attentionfor overcoming the quadratic bottleneck, training such models from scratch remains prohibitively expensive. Inheriting weights from pretrained Transformers provides an appealing shortcut, yet the fundamentalrepresentational gapbetween Softmax and linear attention prevents effectiveweight transfer. In this work, we address this conversion challenge from two perspectives: architectural alignment and representational alignment. We identifyTest-Time Training(TTT) as a linear-complexity architecture whosetwo-layer dynamic formulationis structurally aligned withSoftmax attention, enabling direct inheritance of pretrained attention weights. To further align representational properties, including key shift-invariance and locality, we introducekey instance normalizationand a lightweightlocality enhancement module. We validate our approach by linearizingStable Diffusion 3.5and introduce SD3.5-T^5 (Transformer To Test Time Training). With only 1 hour of fine-tuning on 4timesH20 GPUs, SD3.5-T^5 achieves comparable text-to-image quality to the fine-tuned Softmax model, while accelerating inference by 1.32times and 1.47times at 1K and 2K resolutions. Code is available at https://github.com/LeapLabTHU/Transformer-to-TTT.

View arXiv pageView PDFGitHub7Add to collection

Get this paper in your agent:

hf papers read 2605\.02772

Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash

Models citing this paper0

No model linking this paper

Cite arxiv.org/abs/2605.02772 in a model README.md to link it from this page.

Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2605.02772 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2605.02772 in a Space README.md to link it from this page.

Collections including this paper0

No Collection including this paper

Add this paper to acollectionto link it from this page.

Similar Articles

Modular TTT: Rethinking Test-Time Training as Composable Modules

Hugging Face Daily Papers

This paper introduces Modular TTT, a framework that represents test-time training inner learners as directed acyclic graphs, enabling systematic ablation and composition of components. The authors train 410M and 1.45B parameter models on 100B tokens, achieving performance comparable to GatedDeltaNet.

LLT: Local Linear Transformer for PDE Operator Learning

arXiv cs.LG

Introduces LLT, a transformer-based neural operator that combines linear global attention with local spatial mixing for PDE learning. It achieves competitive accuracy and faster training compared to baselines on multiple PDE problems.

Retrofitting Linear Attention into Diffusion Language Models

arXiv cs.LG

This paper introduces block-hybrid attention, which combines exact softmax attention within active denoising blocks and linear attention over previous blocks, to accelerate inference in pretrained diffusion language models. The authors retrofit this hybrid attention into LLaDA 2.1, achieving up to 1.7x higher decoding throughput with minimal post-training.