Research Papers
EN ZH

高效 AI · Transformer · 文生图

用 Test-Time Training 线性化 ViT(SD3.5-T5)

SD3.5-T5 把 Softmax 换成继承预训练 Q/K/V/O 权重的 TTT 层,3000 步微调后 GenEval 0.69,1K 提速 1.32 倍、2K 提速 1.47 倍。

Tsinghua University · 论文发表 2026-05-28 · 3 分钟读完

用 Test-Time Training 线性化 ViT(SD3.5-T5)

快速答案

SD3.5-T5 把 Stable Diffusion 3.5 里的 Softmax 注意力换成 Test-Time Training(TTT)层,只用 3000 步微调(4 张 H20 上约 1 小时)就达到接近原模型的文生图质量,然后在单张 H20 上 1024x1024 提速 1.32 倍、2048x2048 提速 1.47 倍。GenEval 上它拿 0.69,微调后的 Softmax 模型是 0.70,原始 SD3.5 是 0.66;DPG-Bench 上 84.43,略高于原始的 83.83。为什么这么便宜:TTT 的两层动态结构跟 Softmax 注意力是同一个形状,预训练的 Q/K/V/O 投影能直接搬过去,不用丢掉重练。这篇讲的是扩散图像生成上的转换性价比,不是说线性注意力从此追平了 Softmax。

把 Softmax 转线性,为什么通常会崩

线性注意力去掉了二次方开销,高分辨率扩散因此变慢的根源就在这里。但从头训一个线性模型,代价跟训原模型一样大。大家都想要的捷径是继承预训练 Transformer 的权重再微调几步。这条路通常走不通:Softmax 注意力层学到的表示塞不进线性核。之前的转换方法,比如 LiT 那一类嫁接,只保留 MLP 权重,把注意力投影全扔了,而做混合的恰恰是这部分。论文在 ImageNet 上做了对照:从 DeiT 骨干微调出来的普通线性注意力 Top-1 只有 63.30%,LiT 69.52%,CLEAR 68.66%,Hedgehog 直接发散成 NaN,对照的 Softmax DeiT-Tiny 是 72.05%。这个差距就是丢掉注意力知识的代价。

TTT 怎么把 Softmax 的注意力权重继承下来

真正的点子是结构对齐,不是发明新核。论文把 Softmax 注意力改写成一个中间带非线性的两层变换:query 先过一个由 key 构成的矩阵,做一次非线性压缩,再过一个由 value 构成的矩阵。Test-Time Training 恰好是这个形状。它的内层模型是个小的两层 MLP,第一层在推理时会被一步梯度更新,所以也读作 query、非线性中间、输出。两套结构逐项对应,预训练的 Q、K、V、O 投影就能原样插进 TTT 层,含义不变。跟固定线性核不同,这个非线性内层 MLP 有足够容量去拟合 Softmax 层学到的函数,所以继承下来不掉点。

消融把这点单独验证了:在 DiT-S/2 生成器上,完整继承权重 FID 68.52;只继承 MLP 权重是 89.71;只继承注意力权重是 93.18。两边都留住才稳得住质量,注意力权重不能省。

让继承能训起来的两个修补

光对齐结构还不够。Softmax 和 TTT 在两个表示习惯上不一样,论文各打了一个补丁。

一是 key 的平移不变性。给所有 key 加一个常数偏移,Softmax 会把它抵消掉,偏移后的 key 给出同样的注意力。TTT 没有这个性质,而预训练的 key 偏移很大:在 5K 张 ViT 激活上测出的 key 偏移比约 0.5,随机初始化只有 0.07。把这些没对中的 key 喂进 TTT 内层优化,梯度会炸,这就是直接移植会出 NaN 的原因。key 实例归一化在进层之前把 key 沿序列方向对中,就这一个改动,把发散变成了 DeiT-Tiny 上 71.19% 的 Top-1。

二是局部性。ViT 里的 Softmax 注意力其实悄悄是局部的,TTT 把这个偏置丢了。一个轻量模块补回来:在 query 和 key 上做深度可分卷积,把 DeiT-Tiny 从 69.25% 抬到 71.19%;再加一小项邻域注意力到 72.06%,反超了 72.05% 的 Softmax 基线。这个局部模块是挂在核心架构洞见上的便宜工程,论文也讲清楚了它做的是补归纳偏置,不是变魔术。

关键结果

  • GenEval(SD3.5): SD3.5-T5 总分 0.69,微调 Softmax 模型 0.70,原始 SD3.5 0.66,转换基本保住了指令遵循质量。分项:Single 1.00、Two 0.88、Counting 0.58、Color 0.88、Position 0.25、Color-Attr 0.55。
  • DPG-Bench(SD3.5): 84.43,略高于原始 SD3.5-Medium 的 83.83。
  • 提速: 单张 H20 上 1024x1024 延迟从 25 秒降到 19 秒,1.32 倍;2048x2048 升到 1.47 倍。序列越长二次项越大,线性层省掉的正是这块。
  • 转换成本: 3000 步,batch 64,数据是 Flux 生成的 1024x1024 图,4 张 H20 约 1 小时。类条件 DiT 上,线性化的 DiT5-XL/2 微调 8 epoch(原从头训 1400 epoch)拿到 FID 2.48,原模型 2.27,约 0.57% 的训练成本。
  • ImageNet 转换(DeiT-Tiny,30 epoch): T5-T 71.19%、T5+-T 72.06%,对照 Softmax DeiT-Tiny 72.05%,普通线性注意力只有 63.30%,Hedgehog 发散。
  • 权重继承消融(DiT-S/2): 完整继承 FID 68.52,只 MLP 89.71,只注意力 93.18,说明注意力投影里有能搬的知识。

怎么读这个标题

提速是真的,但有条件,把它读成”线性注意力追平 Softmax”就过头了。1.32 和 1.47 倍出现在 1K 和 2K,是因为二次方开销在那里才占大头;低分辨率、短序列时线性层没多少可省,甚至可能更慢,所以论文专门把图像生成的提速报在 1K 和 2K。质量也只是”接近”,锚在 GenEval 0.69 对 0.70、DPG-Bench 84.43 对 83.83,不是干净的赢,而且 SD3.5 上没给直接的 FID。老实的说法:如果你手头已经有一个高分辨率扩散 Transformer,这是一条便宜路子,花约一小时换来接近不掉点的真实推理提速;能搬走的洞见是,挑一个数学形状跟 Softmax 对得上的线性架构,胜过去发明一个巧妙的核。把预训练 Softmax 知识嫁接进更便宜注意力的同一个思路,也是 Full Attention Strikes Back 在做的;让扩散 Transformer 在高分辨率下更省的大方向,还能看 LoomVideo。

局限与存疑

没有专门的局限小节,对一个旋钮这么多的转换方法,这本身就值得记一笔。质量是接近不是更好:GenEval 0.69 对 0.70 在噪声内,Position 0.25 仍然弱,所以它保住了原模型的短板,没修。提速跟分辨率挂钩,1K 以下帮不上忙,论文也没给出拐点在哪。TTT 那套新参数要 20 倍学习率加 key 实例归一化才不炸,配方敏感,换骨干可能要重调。SD3.5 的结果靠 GenEval 和 DPG-Bench,没有直接 FID,微调数据又是 Flux 生成的合成图,这就留了个口子:这份接近,有多少是在蒸馏 Flux,而不是保住 SD3.5。最后,所有结果都是图像生成和 ImageNet 分类,视频、长序列、文本这些最想用线性注意力的场景一个都没测。

常见问题

SD3.5-T5 跟微调后的 Softmax 模型在文生图质量上怎么对比?

接近,没超过。SD3.5-T5 GenEval 0.69,微调 Softmax 是 0.70,原始 0.66;DPG-Bench 84.43,原始 83.83。论文用的词是”comparable”,而且没给 SD3.5 的直接 FID,所以当作噪声内的持平,不是质量提升。真正的回报是 1K、2K 上 1.32 和 1.47 倍的推理提速。

为什么 TTT 能继承预训练注意力权重,线性注意力却不行?

Softmax 注意力能写成两层变换:query 过一个 key 构成的矩阵,过一次非线性,再过一个 value 构成的矩阵。TTT 的内层模型也是带非线性中间的两层 MLP,两者逐项对齐,预训练的 Q/K/V/O 投影含义不变地搬过去。普通线性核没有非线性中间,表示不出 Softmax 学到的函数,所以线性注意力转换在 DeiT-Tiny 上只有 63.30% Top-1,而 T5 到 71.19%。

1.32 和 1.47 倍提速到底如何只在高分辨率出现?

在高分辨率。1.32 倍是 1024x1024(单张 H20 延迟 25 秒到 19 秒),1.47 倍是 2048x2048。线性层去掉了序列长度上的二次项,所以省的随分辨率增长,低分辨率几乎没有。把提速当成普适收益是最常见的误读。

SD3.5-T5 不做 key 实例归一化会怎样(对比 Hedgehog)?

训练发散成 NaN。预训练 key 带很大的常数偏移(实测偏移比约 0.5,随机初始化是 0.07),Softmax 能抵消,TTT 不能,内层优化梯度因此爆炸。用实例归一化把 key 沿序列对中,就把 NaN 变成 DeiT-Tiny 上能跑的 71.19% Top-1。Hedgehog 没这个修补,在同一个探针里发散。

转换比从头训便宜多少?

SD3.5 上是 3000 步、batch 64,用 Flux 生成的 1024x1024 图,4 张 H20 约 1 小时。类条件 DiT 上是 8 epoch,对照从头训的 1400 epoch,约 0.57% 的训练成本,FID 2.48 对 2.27。这个成本差是论文真正的论点,相关的省算力工作还可以看 MiniMax 稀疏注意力。

一句话:SD3.5-T5 说明,挑一个两层数学跟 Softmax 对齐的线性架构,就能继承预训练注意力权重,花约一小时微调换来高分辨率下 1.32-1.47 倍的提速,质量接近持平。阅读 arXiv 原文。