Linearizando o Vision Transformer com Treinamento em Tempo de Teste

Resumo

Embora mecanismos de atenção com complexidade linear ofereçam uma alternativa promissora à atenção Softmax para superar o gargalo quadrático, treinar tais modelos do zero ainda é proibitivamente caro. Herdar pesos de Transformers pré-treinados oferece um atalho atraente, porém a lacuna representacional fundamental entre a atenção Softmax e a atenção linear impede a transferência eficaz de pesos. Neste trabalho, abordamos este desafio de conversão sob duas perspectivas: alinhamento arquitetural e alinhamento representacional. Identificamos o Test-Time Training (TTT) como uma arquitetura de complexidade linear cuja formulação dinâmica de duas camadas está estruturalmente alinhada com a atenção Softmax, permitindo a herança direta de pesos de atenção pré-treinados. Para alinhar ainda mais as propriedades representacionais, incluindo invariância ao deslocamento de chave e localidade, introduzimos a normalização de instância de chave e um módulo leve de aprimoramento de localidade. Validamos nossa abordagem linearizando o Stable Diffusion 3.5 e introduzimos o SD3.5-T^5 (Transformer To Test Time Training). Com apenas 1 hora de ajuste fino em GPUs 4×H20, o SD3.5-T^5 alcança qualidade de texto para imagem comparável ao modelo Softmax ajustado, enquanto acelera a inferência em 1,32× e 1,47× nas resoluções de 1K e 2K. O código está disponível em 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.