Linearisierung des Vision Transformers mit Testzeit-Training

Zusammenfassung

Während Aufmerksamkeitsmechanismen mit linearer Komplexität eine vielversprechende Alternative zur Softmax-Aufmerksamkeit darstellen, um den quadratischen Engpass zu überwinden, bleibt das Training solcher Modelle von Grund auf unerschwinglich teuer. Die Übernahme von Gewichten aus vortrainierten Transformatoren bietet einen attraktiven Abkürzungsweg, jedoch verhindert die grundlegende Repräsentationslücke zwischen Softmax- und linearer Aufmerksamkeit eine effektive Gewichtsübertragung. In dieser Arbeit adressieren wir diese Konvertierungsherausforderung aus zwei Perspektiven: strukturelle Angleichung und Repräsentationsangleichung. Wir identifizieren Test-Time Training (TTT) als eine Architektur mit linearer Komplexität, deren zweischichtige dynamische Formulierung strukturell mit der Softmax-Aufmerksamkeit übereinstimmt, was eine direkte Übernahme von vortrainierten Aufmerksamkeitsgewichten ermöglicht. Um die Repräsentationseigenschaften, einschließlich Key-Shift-Invarianz und Lokalität, weiter anzugleichen, führen wir eine Key-Instanz-Normalisierung und ein leichtgewichtiges Lokalitätsverbesserungsmodul ein. Wir validieren unseren Ansatz, indem wir Stable Diffusion 3.5 linearisieren, und stellen SD3.5-T³ (Transformer To Test Time Training) vor. Mit nur einer Stunde Feintuning auf 4 H20-GPUs erreicht SD3.5-T³ eine mit dem feingetunten Softmax-Modell vergleichbare Text-zu-Bild-Qualität, während die Inferenz bei 1K- und 2K-Auflösungen um das 1,32- bzw. 1,47-Fache beschleunigt wird. Der Code ist verfügbar unter 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.