用测试时训练让线性注意力模型快速继承预训练权重，提升推理速度。
Linearizing Vision Transformer with Test-Time Training
- 提出测试时训练架构，与Softmax注意力结构对齐，可直接复用预训练权重。
- 仅用1小时微调即在1K/2K分辨率下实现接近原始模型的图像生成质量。
- 适合追求高效推理的视觉生成研究者和工业部署场景使用。
虽然线性复杂度注意力机制有望克服Softmax注意力的二次计算瓶颈，但从头训练仍成本高昂。直接继承预训练Transformer权重看似便捷，但Softmax与线性注意力间存在根本表征差异，阻碍有效迁移。本文从结构对齐与表征对齐两方面解决该问题：识别出测试时训练（TTT）是一种线性复杂度架构，其双层动态形式与Softmax注意力结构一致，支持直接继承预训练权重；为进一步对齐键值平移不变性与局部性，引入键实例归一化与轻量局部增强模块。通过将Stable Diffusion 3.5线性化，构建SD3.5-T$^5$模型。仅需在4×H20 GPU上进行1小时微调，该模型在1K与2K分辨率下分别实现1.32×与1.47×的推理加速，且文本到图像生成质量与微调后的Softmax模型相当。代码已开源。
原文摘要 · Abstract (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 4$\times$H20 GPUs, SD3.5-T$^5$ achieves comparable text-to-image quality to the fine-tuned Softmax model, while accelerating inference by 1.32$\times$ and 1.47$\times$ at 1K and 2K resolutions. Code is available at https://github.com/LeapLabTHU/Transformer-to-TTT.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。