테스트 시간 훈련을 통한 비전 트랜스포머의 선형화

초록

선형 복잡도 어텐션 메커니즘은 이차 병목 문제를 극복하기 위한 소프트맥스 어텐션의 유망한 대안을 제공하지만, 이러한 모델을 처음부터 훈련하는 것은 여전히 엄청난 비용이 든다. 사전 학습된 트랜스포머의 가중치를 상속하는 것은 매력적인 지름길을 제공하지만, 소프트맥스 어텐션과 선형 어텐션 간의 근본적인 표현적 차이로 인해 효과적인 가중치 전이가 불가능하다. 본 연구에서는 이러한 변환 과제를 구조적 정렬과 표현적 정렬이라는 두 가지 관점에서 다룬다. 우리는 테스트-시간 훈련(TTT)을 선형 복잡도 아키텍처로 식별하였으며, 이 아키텍처의 2계층 동적 공식은 소프트맥스 어텐션과 구조적으로 정렬되어 사전 학습된 어텐션 가중치의 직접 상속을 가능하게 한다. 키 이동 불변성 및 지역성을 포함한 표현적 특성을 더욱 정렬하기 위해, 키 인스턴스 정규화와 경량 지역성 강화 모듈을 도입한다. 우리는 스테이블 디퓨전 3.5를 선형화하여 SD3.5-T^5(Transformer To Test Time Training)를 소개함으로써 접근 방식을 검증한다. 4×H20 GPU에서 단 1시간의 미세 조정만으로 SD3.5-T^5는 미세 조정된 소프트맥스 모델과 유사한 텍스트-이미지 생성 품질을 달성하면서, 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.