PaperlayerView exact revision on arXiv

Source-linked AI summary

Linearizing Vision Transformer with Test-Time Training

Yining Li, Dongchen Han, Zeyu Liu, Hanyi Wang, Yulin Wang, Gao Huang

arXiv:2605.02772v2cs.CV

TL;DR

Converting pretrained Softmax Transformers to linear attention remains difficult because their representational spaces differ. This paper uses TTT and representation-alignment modules to linearize SD3.5, achieving comparable text-to-image quality while accelerating inference.

  • Problem

    Converting pretrained Softmax Transformers to linear-complexity attention remains challenging because their representational spaces differ, limiting effective weight inheritance.

  • Method

    The method uses structurally compatible TTT with a nonlinear inner model, key instance normalization, and lightweight locality enhancement for efficient adaptation.

  • Results

    1.32× and 1.47× speedups at 1024 and 2048 resolutions accompany comparable performance after 3,000 fine-tuning steps on Stable Diffusion 3.5.

  • Takeaways & Limitations

    The approach enables practical linearization of pretrained vision and diffusion Transformers without distillation or complex multi-stage training.

  • Takeaways & Limitations

    Direct fine-tuning after weight inheritance can be numerically unstable, with training diverging to NaN.

Abstract

from arXiv · show

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

1. Introduction

The paper addresses the costly conversion of pretrained Softmax Transformers to linear-complexity models through structural and representational alignment. It identifies TTT for weight inheritance and introduces key normalization and locality enhancement, achieving efficient adaptation with near-original performance.

  • Motivation: Softmax attention has quadratic time and memory complexity in sequence length N, limiting scalability for high-resolution images and long-context generation.
  • Approach: The work frames pretrained-Transformer linearization as a conversion problem requiring both structure alignment and representation alignment.
  • Structural alignment: TTT provides a linear-complexity architecture structurally compatible with Softmax attention, enabling full pretrained-weight inheritance without complex distillation.
  • Representational alignment: Key shift-invariance and strong locality bias motivate key instance normalization and a lightweight locality-enhancement module for data-efficient adaptation.
  • Validation: 0.57% of the original training steps yields comparable DiT-XL/2 performance after 8 epochs without distillation or complex operator activation.

2. Related Work

Related work addresses Vision Transformer scalability through local, sparse, circulant, and linear-complexity attention, while TTT offers an online-learning formulation for linear-complexity sequence modeling. This work focuses on TTT-based linear attention and the broader challenge of converting pretrained Softmax Transformers without prohibitively expensive retraining.

  • Efficient Vision Transformers: Vision Transformers face quadratic attention complexity O(N^2), motivating local, sparse, circulant, and linear-complexity alternatives for long visual sequences.The cited approaches target scalability limitations in efficient Vision Transformers.
  • Efficient Vision Transformers: This work focuses on TTT-based linear attention for efficient visual modeling.TTT is positioned as the paper’s chosen linear-attention approach for visual applications.
  • Test-Time Training (TTT): TTT reformulates attention as online learning, updating hidden states through gradient-based inference-time updates while enabling linear-complexity O(N) sequence modeling.The formulation is described as retaining expressive power and has been extended to language modeling and video generation.
  • Linearizing Pretrained Transformers: Converting pretrained Softmax Transformers into linear-complexity alternatives is motivated by the prohibitive cost of retraining large models from scratch.Prior work includes learnable query and key activations intended to approximate Softmax behavior.

3. Preliminary

The preliminary section contrasts Softmax and linear attention, then introduces Test-Time Training (TTT) as an inference-time framework with a dynamically updated inner model. TTT supports multiple inner-model architectures and combines learned parameters with sequence-specific fast-weight updates.

  • Attention Mechanisms: Softmax attention operates on queries, keys, and values, whereas linear attention replaces Softmax with a kernel function.
  • Test-Time Training: TTT updates an inner model during inference using learned parameters and fast weights computed per sequence through gradient descent on a self-supervised loss.The fast weights are sequence-dependent updates, while the inner-model parameters are optimized during training.
  • Test-Time Training: TTT permits inner models ranging from a single linear layer to a two-layer MLP or convolutional networks, creating capacity-efficiency trade-offs.
  • Test-Time Training: TTT applies query inputs to an inner model whose parameters combine learned weights with sequence-specific updates, compressing cached K, V information into a compact model.

4. Method

The method linearizes Softmax attention with a two-layer TTT inner model that preserves structural alignment and enables pretrained-weight inheritance. It further aligns key shift-invariance and locality through key normalization and depthwise query-key convolutions.

  • Structural alignment: Softmax attention acts as a nonlinear two-layer dynamic MLP, whereas standard linear attention is a single-layer dynamic linear transformation with limited representational capacity.This structural mismatch hinders effective transfer from pretrained Softmax models.
  • Structural alignment: Two-layer TTT matches Softmax attention’s structure, allowing pretrained-weight inheritance while its learnable inner parameters adapt to the original representation space.TTT dynamically updates a lightweight inner model on the KV cache at test time.
  • Empirical validation: 24.39 is the Freeze performance of ProjQK linear attention, improving over Vanilla linear attention’s 3.71 but remaining far below effective transfer levels.This result confirms that additional projections do not resolve the structural gap between standard linear and Softmax attention.
  • Empirical validation: 67.33% vs. 67.09%: TTT-SwiGLU slightly outperforms TTT-3Layer with equal parameters, supporting two-layer gated nonlinearity as the default inner model.Performance improves as the TTT inner model’s nonlinearity and depth increase.
  • Representation alignment: Key instance normalization recenters inherited keys to emulate Softmax’s shift-invariance, stabilizing training and enabling smoother weight transfer.Without this alignment, directly fine-tuning inherited TTT-SwiGLU weights can diverge to NaN.
  • Representation alignment: Depthwise convolutions applied to queries and keys enhance locality before TTT computation, with DWCQK achieving the best performance.TTT can also be combined with Neighborhood Attention for local-window processing.

5. Experiments

Experiments evaluate TTT weight inheritance on ImageNet classification and image generation, including DiT-XL/2 and Stable Diffusion 3.5-Medium. TTT consistently improves efficient-attention classification results and preserves generation quality while reducing computational cost.

  • ImageNet Classification: On ImageNet-1K, T5 reaches 71.19 accuracy with 6.2M parameters and 1.34G computation, outperforming standard linear attention, LiT, CLEAR, and Hedgehog.All methods inherit DeiT weights under the same setting, except LiT does not inherit attention weights.
  • Class-Conditional Image Generation: DiT5 replaces attention with simplified TTT-2layer blocks using depthwise convolution on Q and K, L2 inner-model loss, and NAT with window size 5.The setup inherits all weights from a pretrained DiT-XL/2 checkpoint and requires neither knowledge distillation nor an operator-activation stage.
  • Stable Diffusion 3.5-Medium: SD3.5-T5 replaces approximately 50% of SD3.5-Medium transformer blocks and fine-tunes for 3,000 steps, requiring approximately 1 hour on 4 NVIDIA H20 GPUs.It is initialized from pretrained SD3.5-Medium weights and trained on Flux-generated 1024×1024 images.
  • Stable Diffusion 3.5-Medium: 84.43 vs. 83.83 on DPG-Bench and 0.69 vs. 0.66 on GenEval show that SD3.5-T5 improves over original SD3.5-Medium while retaining generation quality.Under the same fine-tuning budget, SD3.5-T5 remains comparable to SD3.5-FT, scoring 84.43 vs. 82.74 on DPG-Bench and 0.69 vs. 0.70 on GenEval.
  • Ablations: Without NAT, TTT achieves FID 72.98 versus 68.52 with NAT, indicating NAT is an optional enhancement rather than a core component.The text-to-image experiments therefore omit NAT and use pure TTT blocks.

6. Conclusion

The work presents T5 as a structurally aligned linear-complexity architecture that inherits Softmax attention weights, achieves competitive performance with minimal fine-tuning, and reduces computational cost. It requires no distillation or multi-stage training and is validated across image classification, class-conditional generation, and text-to-image generation.

  • 6. Conclusion: T5 enables direct inheritance of pretrained Softmax attention weights while providing linear-complexity Transformer attention.The framework addresses Transformer linearization from a structural perspective.
  • 6. Conclusion: The training procedure requires no distillation or multi-stage training strategies.T5 achieves competitive performance with only minimal fine-tuning while reducing computational complexity.
  • 6. Conclusion: 3,000 training steps yield performance comparable to the Softmax model on SD3.5-Medium while reducing computational FLOPs.The framework is validated across image classification, class-conditional generation, and text-to-image generation.

A. Image Classification

For ImageNet-1K classification, the models use Swin Transformer’s 30-epoch fine-tuning learning-rate schedule. Linear attention uses ELU + 1, while TTT uses SiLU with an inner-product loss for its inner-model objective.

  • A. Image Classification: The ImageNet-1K classification experiments adopt Swin Transformer’s 30-epoch fine-tuning learning-rate schedule.The detailed hyperparameters are summarized in Table 8.
  • A. Image Classification: Linear attention uses ELU + 1 as its activation function.
  • A. Image Classification: TTT uses SiLU activation and an inner-product loss for the inner-model objective.These settings follow ViT3.

B. Text-to-Image Generation

The method selectively replaces attention branches in SD3.5-Medium’s MMDiT architecture while preserving other attention components. Evaluation uses fixed inference settings and generates four images per prompt.

  • Architecture: SD3.5-Medium has 24 MMDiT blocks, with the first 13 designated as dual-attention blocks.These early blocks combine text-image joint attention with an additional image-token self-attention branch.
  • Architecture: In the first 13 blocks, only the additional image-token self-attention branch is replaced, while text-image joint attention remains unchanged.The proposed module substitutes for the additional image-token self-attention branch.
  • Training and Evaluation: 28 inference steps, guidance scale 3.5, and four generated images per prompt are used for DPGBench and GenEval experiments.These settings apply to all experiments described for text-to-image generation.

C. Ablation Study

The ablation study evaluates weight inheritance strategies, learning-rate scaling for TTT-specific parameters, and Neighborhood Attention across DiT-S/2 and DeiT-tiny settings.

  • Weight inheritance strategies: The study compares weight inheritance strategies on DiT-S/2 with batch size 256 for 50K training steps.This experiment is reported in Table 10.
  • Learning-rate scaling: The study examines learning-rate scaling for TTT-specific parameters using DiT-S/2 and DeiT-tiny.Results are separated into DiT-S/2 and DeiT-tiny experiments.
  • Neighborhood Attention: The study evaluates Neighborhood Attention on DiT-S/2 using ImageNet at 256×256 resolution.This experiment is reported in Table 12.
Loading 2605.02772v2…