Linealizando el Transformador de Visión con Entrenamiento en Tiempo de Prueba

Resumen

Mientras que los mecanismos de atención de complejidad lineal ofrecen una alternativa prometedora a la atención Softmax para superar el cuello de botella cuadrático, entrenar estos modelos desde cero sigue siendo prohibitivamente costoso. Heredar los pesos de Transformers preentrenados proporciona un atractivo atajo, sin embargo, la brecha representacional fundamental entre la atención Softmax y la atención lineal impide una transferencia de pesos efectiva. En este trabajo, abordamos este desafío de conversión desde dos perspectivas: la alineación arquitectónica y la alineación representacional. Identificamos el Entrenamiento en Tiempo de Prueba (TTT) como una arquitectura de complejidad lineal cuya formulación dinámica de dos capas está estructuralmente alineada con la atención Softmax, permitiendo la herencia directa de pesos de atención preentrenados. Para alinear aún más las propiedades representacionales, incluyendo la invariancia al desplazamiento de claves y la localidad, introducimos la normalización de instancias de claves y un módulo ligero de mejora de localidad. Validamos nuestro enfoque linealizando Stable Diffusion 3.5 e introducimos SD3.5-T^5 (Transformer a Entrenamiento en Tiempo de Prueba). Con solo 1 hora de ajuste fino en 4×H20 GPUs, SD3.5-T^5 logra una calidad de texto a imagen comparable al modelo Softmax ajustado, mientras acelera la inferencia en 1,32× y 1,47× en resoluciones de 1K y 2K. El código está disponible en 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.