Back to Blog 动态短卷积:给 Transformer 加一个

动态短卷积:给 Transformer 加一个"会看情况的"局部感受野

Paper
Dynamic Short Convolutions Improve Transformers
Transformer 动态卷积 MIT cs.LG
TL;DR: 将 Transformer 中静态短卷积替换为输入依赖的动态卷积,在 150M–2B 参数的语言模型上带来 1.33×(QKV 版)到 1.60×(全线性层版)的计算优势。配合定制 Triton kernel,端到端训练仅增加 8% 开销。效果泛化至 MoE、Mamba-2、Gated DeltaNet。

背景与动机

Transformer 统治大模型时代,核心是 attention,但卷积从未真正退出舞台。Mamba、DeltaNet 等线性 RNN 都保留了 depthwise separable 短卷积(kernel width 3-5)作为局部序列混合组件。但这些卷积的权重是静态的——同一个滤波器在序列每个位置完全相同。

MIT 的这篇论文提出了一个直觉上很自然的改进:让卷积权重随输入动态变化。每个 token 根据自己的上下文生成一组专属的局部滤波器,然后从邻域中按需提取信息。

论文用一个精妙的语言例子说明动机:考虑 "the old can opener" 和 "the old can swim"。前三个词完全相同,但局部组合方式完全不同——前者是 [the [old [can opener]]],后者是 [[the old] [can swim]]。静态卷积用同一个滤波器处理所有位置,无法区分这种上下文依赖的局部组合;动态卷积让每个位置自己决定"如何看邻居",自然就能处理这种变化。

核心方法

参数化:标准短卷积计算 y_t = Σ w_k ⊙ x_{t-k},其中 w 是固定的。动态短卷积将 w 替换为 w(t),由当前隐藏状态通过线性投影生成:

y_t := Σ w_k(t) ⊙ x_{t-k}

直接生成 D → W×D 的权重会使参数量翻倍(W=4 时),因此论文提出两种高效参数化:

低秩分解:D → R → W×D,用一个小中间维度 R 控制参数量。实验中 R=16-52,表现最佳。

按头共享:D → W×(D/H),每个注意力头共享一个滤波器,天然适合 GPU kernel 优化。

放置位置:主要实验将动态卷积放在 QKV 投影之后、RoPE 之前,以残差方式连接:X = X + DynamicShortConv(X)。也实验了放在所有线性层之后的变体。

Figure 1
Figure 1:自定义 Triton kernel 在 H100 上的延迟对比。橙色为 Triton 实现,灰色为 PyTorch eager / torch.compile 基线,虚线为 CUDA 优化的静态卷积 kernel。head-wise 变体在 H≥16 时甚至比静态卷积更快。

合成实验:动态组合能力验证

论文设计了一个变长 key 的关联回忆任务(modified MQAR)。每个 key 由 1-3 个 token 组成,由一个特殊分隔符标记长度。由于 key 之间存在共享 token 且长度不同,没有任何静态滤波器能区分它们

在 Fuzzy Recall 上,动态卷积将准确率从 0.505(静态)拉到 0.726——这类任务要求模型根据上下文动态决定如何聚合变长局部信息,正是动态卷积的甜区。

Figure 2
Figure 2:左图——变长 key MQAR 任务上,动态卷积一致优于静态卷积和普通 Transformer。右图——MAD benchmark 上各任务的表现,动态卷积在 Fuzzy Recall 上优势尤为突出。

语言建模 Scaling Law

在 150M 到 2B 参数的 dense Transformer 上,使用 Nemotron-CC 语料、Llama-style 架构进行训练。Scaling law 拟合表明:仅加在 QKV 上就有 1.33× 计算优势;加在所有线性层上达到 1.60×。这意味着:达到相同 loss,动态卷积模型只需普通 Transformer 约 60-75% 的计算量。

Figure 3
Figure 3:Scaling law 拟合结果。左图——QKV 版动态卷积的 1.33× 计算优势。右图——全线性层版本的 1.60× 优势,提升更显著。

下游任务评估

动态卷积在各规模上一致优于标准 Transformer 和静态卷积变体。一个值得注意的结果:Mamba-2 换上动态卷积后,性能接近 Gated DeltaNet(带静态卷积),说明改进局部混合器可能比重新设计序列混合器更有效。

训练效率

QKV 版动态卷积仅比 Transformer 基线慢约 8%,静态卷积约 6%。全线性层版本慢 22-25%,但 1.60× 的计算优势完全弥补了这一开销。

Figure 4
Figure 4:端到端训练吞吐量。QKV 版动态卷积仅比 Transformer 基线慢约 8%,静态卷积约 6%。全线性层版本慢 22-25%,但 1.60× 的计算优势完全弥补了这一开销。
Figure 5
Figure 5:RULER benchmark 上 MoE 模型的各子任务表现。动态卷积在 multi-key (MK)、multi-query (MQ)、multi-value (MV) 子任务上优势显著——这些任务需要输入依赖的局部聚合。

消融实验要点

卷积宽度:W=3-4 是最佳甜区,更大不会带来额外增益。这与静态卷积的发现一致。

放置位置:单投影放置时,Value 投影增益最大;Q+K+V 三者都加效果最好。

QK-Norm 兼容:在带 QK-Norm 的 Transformer 上,动态卷积仍然显著优于静态卷积——静态卷积在 QK-Norm 下几乎无收益。

局限性与未来方向

Scaling 实验仅到 2B dense + 7B MoE(1B active),frontier 级别尚未验证。训练只用了单一数据混合(Nemotron-CC)。Triton kernel 在 H100 上效率不错,但推理优化和其他硬件(TPU)的支持还需额外工程。参数化、放置方式、kernel width 的设计空间远未穷尽。

总结与评价

动态短卷积的思路清晰、实现简洁、效果稳定。它不是架构层面的革命,而是一个被忽视的"免费午餐"式改进——用极少的工程代价换取显著的性能提升。对 Mamba-2 和 DeltaNet 的泛化能力说明这是一个通用的架构增强手段,不局限于 Transformer。

对于 2B 以下模型,这是一个几乎无脑值得加的改进。对于更大规模模型,scaling law 的趋势令人乐观,但最终结论需要更多实验。论文配套的 Triton kernel 已开源。

Tags: #y_t := Σ w_k(t) ⊙ x_{t-k}