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

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.

Training, Reading, and Editing Legible Transformers

arXiv cs.LG

This paper introduces methods to train transformers with legible-by-construction operators using a crispness penalty that avoids collapse via a per-channel variance floor, achieving high legibility and quality parity with conventional baselines.

U-TTT: Towards Generalizable PET Image Denoising via Test-Time Training

Hugging Face Daily Papers

This paper introduces U-TTT, a U-shaped deep learning model with test-time training layers and dual-domain adaptation for robust PET image denoising under distribution shifts, achieving state-of-the-art performance across different dose levels and scanner types.

@VukRosic99: Long-context Transformers hit two walls: quadratic attention compute and a KV cache that reaches hundreds of GB at 1M t…

X AI KOLs Timeline

MiniCPM-SALA is a 9B-parameter hybrid attention model that interleaves sparse and linear attention to overcome the quadratic compute and large KV cache bottlenecks of long-context Transformers. It achieves 3.5x faster inference than Qwen3-8B at 256K tokens and supports up to 1M tokens on consumer GPUs, with a cost-effective continual training approach that reduces training costs by ~75%.