Линеаризация трансформера зрения с помощью обучения во время тестирования

Аннотация

В то время как механизмы внимания с линейной сложностью предлагают многообещающую альтернативу Softmax-вниманию для преодоления квадратичного узкого места, обучение таких моделей с нуля остается непомерно дорогим. Наследование весов от предварительно обученных трансформеров представляет собой привлекательный путь, однако фундаментальный разрыв в представлениях между Softmax и линейным вниманием препятствует эффективному переносу весов. В данной работе мы решаем эту задачу преобразования с двух точек зрения: архитектурного согласования и согласования представлений. Мы определяем Test-Time Training (TTT) как архитектуру с линейной сложностью, чья двухслойная динамическая формулировка структурно согласована с Softmax-вниманием, что позволяет напрямую наследовать предварительно обученные веса внимания. Для дальнейшего согласования свойств представлений, включая инвариантность к сдвигу ключей и локальность, мы вводим нормализацию экземпляров ключей и легковесный модуль улучшения локальности. Мы проверяем наш подход на примере линеаризации Stable Diffusion 3.5 и представляем SD3.5-T^5 (Transformer To Test Time Training). Всего за 1 час дообучения на 4×H20 GPU SD3.5-T^5 достигает качества преобразования текста в изображение, сопоставимого с дообученной Softmax-моделью, при этом ускоряя инференс в 1,32 и 1,47 раза при разрешениях 1K и 2K. Код доступен по адресу 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.