テスト時訓練によるビジョントランスフォーマーの線形化

要旨

線形複雑性のアテンション機構は、二次のボトルネックを克服するSoftmaxアテンションの有望な代替手段を提供するが、そのようなモデルをゼロから訓練することは非常に高コストである。事前学習済みTransformerから重みを継承することは魅力的な近道となるが、Softmaxアテンションと線形アテンションの間の基本的な表現のギャップが、効果的な重みの転送を妨げる。本研究では、この変換の課題にアーキテクチャの整合性と表現の整合性という二つの観点から取り組む。我々は、テスト時訓練（TTT）を線形複雑性のアーキテクチャとして特定し、その二層の動的定式化がSoftmaxアテンションと構造的に整合しているため、事前学習済みアテンション重みの直接的な継承を可能にする。さらに、キーシフト不変性や局所性を含む表現特性を整合させるために、キーインスタンス正規化と軽量な局所性強化モジュールを導入する。我々はStable Diffusion 3.5を線形化することで本手法を検証し、SD3.5-T^5（Transformer To Test Time Training）を導入する。4台のH20 GPU上でのわずか1時間のファインチューニングにより、SD3.5-T^5はファインチューニングされたSoftmaxモデルと同等のテキストから画像への品質を達成しつつ、1Kおよび2K解像度で推論をそれぞれ1.32倍、1.47倍高速化する。コードは https://github.com/LeapLabTHU/Transformer-to-TTT で公開されている。

English

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.