Tapered Language Models:MLP 宽度单调递减,免费的性能提升

2026-06-23 arXiv 2606.23670
模型架构 MLP Tapering 深度感知容量分配 Mila · Cornell
MLP 中间层宽度沿深度 cosine 递减,总参数量和 FLOPs 完全不变。跨四种架构(Transformer、Gated Attention、Hope-attention、Titans)、三种规模(440M-1.3B)全部提升。来自 Aaron Courville 组,典型的"为什么之前没人做过"。

问题:所有层都一样宽,合理吗?

从原始 Transformer 到线性 RNN、记忆增强架构,一个共同默认设置从未被质疑:每一层的参数量相同。但早退法证明残差流在到达最后一层前就收敛了;层剪枝发现深层冗余度高;可解释性工作显示底层捕获语法,上层编码语义。

既然各层重要性不同,参数分配为什么还是均匀的?

控制实验:前宽后窄赢了

把 440M Transformer 的层分成三组,固定总参数量下分配不同 MLP 宽度:

Capacity allocation experiment
图2:三种容量分配方案。前宽后窄(A)ppl 15.96 最优;前窄后宽(B)ppl 17.29 恶化超过 1 点;中间宽(C)也劣于均匀 baseline(16.28)。方向比分配本身更重要。

前宽后窄改善 0.32 点,反向分配恶化超过 1 点。容量不是被动资源,方向决定了一点多。

cosine 衰减是最佳 schedule

把分块改为平滑递减,测试三种衰减(线性、cosine、sigmoid)和五种宽度比:

Tapering schedules
图1:三种递减 schedule 在固定参数量下的 MLP 宽度轮廓与验证 perplexity。Cosine(14.44)> 线性(15.18)> 阶梯(15.64)> 均匀(16.28)。右图:cosine 和线性在不同宽度比下均呈 U 形,cosine 全程优于线性。

Cosine 在所有宽度比上严格最优。最优配置 1.5/0.5:首层 MLP 宽度为 baseline 的 1.5 倍,末层 0.5 倍,ppl 从 16.28 降到 14.44,改善 1.84 点。

为什么有效?层的新颖度递减

测量每层 MLP 输出与输入残差流的余弦相似度。随深度增加,MLP 输出越来越"对齐"已有残差——后面的层在做的事情越来越冗余,宽度却一样,参数浪费了。Tapering 把容量从冗余层转移到仍在产出新信息的浅层。

Layer novelty analysis
图4:GPT-2 系列各层 MLP 输出与残差流的余弦相似度。随深度增加一致上升(r = 0.49-0.71),说明深层 MLP 越来越冗余——输出的方向已在残差中存在。

跨架构验证

在四种架构、两种规模上,cosine 1.5/0.5 全部提升平均常识推理准确率,LAMBADA perplexity 八组对比全部改善:

架构规模ppl ↓基准 pplAvg. ↑
Transformer760M21.4221.8652.84 / 52.25
Gated Attention760M19.9820.7452.88 / 52.61
Hope-attention760M20.5020.6254.05 / 53.69
Titans760M20.7721.5853.29 / 52.30
Transformer1.3B17.1717.3956.38 / 56.05
Gated Attention1.3B15.9216.0356.80 / 56.51
Hope-attention1.3B15.9415.9157.05 / 56.95
Titans1.3B15.7616.0557.08 / 56.73

长上下文检索(Needle-in-a-Haystack)也不降反升。

与 Variable-Width Transformers 对比

VWT 提出 × 形轮廓(两头宽、中间窄),TLM 提出单调递减(前宽后窄)。TLM 控制更干净——单一变量(只改 MLP 宽度)、固定参数预算、有机制分析支撑。控制实验明确证明反向分配有害,说明存在一个清晰的单调方向。

思考

工程价值极高:零成本(不增参数、不增 FLOPs)、零风险(任何架构都能用)、改动极小。作者提到 attention head 数、KV 维度、RNN 状态大小、MoE expert 数都可以 taper。如果所有轴都可以单调递减,那"均匀层"这个默认值确实该被终结了。