通过测试时训练线性化视觉 Transformer

虽然线性复杂度注意力机制为克服二次瓶颈提供了 Softmax 注意力的有前途的替代方案，但从头开始训练此类模型仍然非常昂贵。从预训练的 Transformer 继承权重提供了一条有吸引力的捷径，但 Softmax 和线性注意力之间的基本表征差距阻碍了有效的权重转移。在这项工作中，我们从两个角度解决这一转换挑战：架构对齐和表征对齐。我们将测试时训练（TTT）确定为一种线性复杂度架构，其两层动态公式在结构上与 Softmax 注意力一致，从而能够直接继承预训练的注意力权重。为了进一步调整表征属性，包括关键移位不变性和局部性，我们引入了关键实例归一化和轻量级局部性增强模块。我们通过线性化稳定扩散 3.5 来验证我们的方法，并引入 SD3.5-T$^5$（Transformer 测试时间训练）。只需在 4$\times$H20 GPU 上运行 1 小时的 微调，SD3.5-T$^5$ 就可以实现与微调 Softmax 模型相当的文本到图像质量，同时在 1K 和 2K 分辨率下将推理速度提高 1.32$\times$ 和 1.47$\times$。代码可在 https://github.com/LeapLabTHU/Transformer-to-TTT 获取。