Variable-Width Transformers：当Transformer不再「等宽」——×形架构的缩放定律启示

论文：Variable-Width Transformers
作者：Zhaofeng Wu, Oliver Sieberling, Shawn Tan, Rameswar Panda, Yury Polyanskiy, Yoon Kim (MIT, MIT-IBM Watson AI Lab)
论文链接：https://arxiv.org/abs/2606.18246
代码库：https://github.com/ZhaofengWu/variable-width-transformers

核心发现速览

一句话总结

在Transformer中，所有层保持相同宽度是工程方便，不是数学必然。让中间层「瘦身」、头尾层「增宽」，不仅性能更好，还更省算力和显存。×形架构证明了：瓶颈不是缺陷，而是结构正则化。

问题：为什么Transformer每一层必须一样宽？

Transformer架构自2017年以来，有一个从未被真正质疑的默认设定：

一旦选定了隐藏维度 d，所有 L 层的宽度都是 d。

这个设定极其方便——所有层共享相同的张量形状，kernel可以一次编译、流水线并行天然对齐、分布式训练代码简洁优雅。但方便不等于最优。

论文提出一个根本性问题：

不同层在Transformer中承担不同计算角色——早期层做词法/句法，中间层做语义，晚期层做输出准备。既然如此，为什么每一层必须分配相同的参数和计算预算？

这是一个合理的质疑。在计算机视觉中，U-Net早就用编码器-解码器结构，中间层分辨率更低（空间压缩）。ResNet使用bottleneck残差块来减少计算。但在NLP的Transformer中，这个「层间不均匀」的自由度从未被系统探索。

核心设计：> <former 的 × 形架构

四种形状探索

论文首先对比了4种宽度变化模式（在500M参数上搜索最优形状）：

×形最优，这个结果有点反直觉。论文团队最初的直觉是◇形（中间层更宽，因为中间层做语义计算），但实验结果显示×形更好。

残差流的「参数自由」缩放机制

改变层宽度的最大挑战是：残差连接（skip connection）怎么跨维度？

如果第ℓ层输出dℓ维，第ℓ+1层输入dℓ+1维，维度不匹配。常见做法：

• 训练一个投影层 f: ℝᵈℓ → ℝᵈℓ₊₁
• 零填充（padding with zeros）
• 直接截断

论文发现投影层增加参数，零填充损害性能。它们提出一个参数自由的方法：

保持全局残差流宽度等于最宽层的宽度，每个层只读取和写入残差流的一个子集。

具体规则：

• 收缩时（dℓ < dℓ₋₁）：简单截断多余的维度，f(𝐱) = 𝐱[:dℓ]
• 扩张时（dℓ > dℓ₋₁）：从最近一层处理过该维度的层复制该维度。即每个坐标维度 i 在残差流中一直存在，如果当前层不处理它，它就「bypass」当前层，从上一层直接复制。

数学上，这等价于：整个模型是统一宽度的，但每个层只与残差流的一个子集交互。inactive的维度通过残差连接「跳过」该层。

消融实验证实：copy-forward（复制）方法显著优于零填充和投影层。

参数化公式

给定4个超参数：

• ℓ*：瓶颈层索引
• dℓ*：瓶颈层维度
• α⁻：早期层递减率（ℓ ≤ ℓ*）
• α⁺：晚期层递增率（ℓ > ℓ*）

层宽度按几何级数变化：

早期（ℓ ≤ ℓ*）：dℓ = α⁻ · dℓ₋₁（递减，α⁻ < 1）
晚期（ℓ > ℓ*）：dℓ = α⁺ · dℓ₋₁（递增，α⁺ > 1）

约束条件：总参数量 = 均匀宽度基线（匹配参数量做公平对比）。

论文发现，使用比例参数化（rℓ = ℓ*/L，rd = dℓ*/d）能在不同模型规模上大致复现最优趋势。推荐值：ℓ* ≈ 0.75L（瓶颈在3/4深度），dℓ* ≈ 0.3d（瓶颈为原始宽度的30%）。

关键发现1：×形在相同参数量下性能更好，同时更省资源

模型规模覆盖

缩放定律分析

论文拟合了Loss vs FLOPs和Loss vs 平均层宽的缩放曲线，发现：

1. > <former 的缩放曲线截距更低（同样FLOPs下Loss更低）
2. 斜率更陡（规模越大，优势可能越明显）
3. 达到2B基线Loss（2.751）只需77.8% FLOPs和85.1%平均层宽

下游任务表现

2B模型在11个下游任务上的对比：

<former在perplexity-based任务（LAMBADA, WikiText）上优势最显著，NLU任务整体也更好。

MoE扩展

3B MoE模型（1B active参数）同样受益：

• active参数还少了3%（因为瓶颈层更窄），但Loss更低
• 说明变宽度设计不仅适用于dense模型，也兼容MoE

关键发现2：压缩山谷——中间层的表示坍塌

论文最重要的机制分析，是揭示了常宽Transformer的一个深层问题：中间层的表示坍塌（compression valleys）。

什么是压缩山谷？

近期可解释性研究发现，Transformer中间层会把token表示「压缩」到一个低秩子空间——很多维度被「浪费」，表示熵急剧下降。这被称为压缩山谷或attention sinks的对偶面。

论文用归一化矩阵熵（normalized matrix entropy）来衡量每层残差流表示的「展开程度」：

H = -1/log(r) · Σⱼ pⱼ log pⱼ
其中 pⱼ = σⱼ² / ‖X‖²_F，σⱼ是X的奇异值

熵越高 = 表示越「均匀展开」= 利用的维度越多。

基线 vs > <former 的对比

常宽基线：

• 早期层：熵高，表示展开
• 中间层：熵暴跌到接近0，表示坍塌到一个极低秩子空间
• 晚期层：熵部分恢复

这证实了压缩山谷的存在——中间层有宽度但不好好利用。

> <former：

• 早期层：主动降低熵，压缩表示（因为知道后续层会变窄，提前准备）
• 中间层（瓶颈层）：熵反而更高，避免坍塌
• 晚期层：熵保持较高，直到最终层快速聚焦到目标token

关键洞察：物理上的瓶颈（更少的维度）迫使网络更有效地利用每个维度。常宽模型有「冗余的安全感」，导致中间层懒惰。> <former的瓶颈是「约束产生创造力」的工程实例。**

MLP激活维度负载均衡

论文还分析了MLP中间层的激活密度：

• 常宽基线：中间层大量MLP维度「死亡」（激活频率极低）
• > <former：各维度激活更均匀，负载均衡更好

这验证了瓶颈的「正则化效应」：物理约束消除了「冗余维度被浪费」的可能性。

关键发现3：Logit Lens揭示的预测动态

通过Logit Lens（将每层隐藏状态投影到词表空间，看模型「在想什么」）：

> <former vs 基线：

1. 目标token概率更高：早期到中期层就更快锁定正确答案
2. 解码分布熵更低：模型更「确信」自己在预测什么
3. 层间KL散度更低：分布变化更平缓，不是剧烈震荡

这意味着：> <former的信息流更平滑，从输入到输出的「推理路径」更稳定。常宽模型在中间层经历了剧烈的表示压缩和重建，而> <former通过瓶颈结构让这个过程更自然。

效率优势的数学根源

论文用数学证明，在匹配参数量的前提下，变宽度模型必然更省计算和内存：

参数量匹配

每层参数 ≈ K · dℓ²（线性投影矩阵占主导，与维度平方相关）。匹配总参数量：

Σℓ dℓ² = L · d²  （等于常宽基线）

由均方 ≥ 均值平方（Jensen不等式）：

(Σdℓ / L)² ≤ Σdℓ² / L = d²
⇒ 平均层宽 < d

所以平均层宽必然小于常宽基线。

FLOPs减少

• Dense FLOPs（线性投影）：与参数量线性相关，匹配参数量时相同
• Attention FLOPs：与序列长度² × 层宽度 线性相关，所以总attention FLOPs ∝ Σdℓ < L·d

KV cache减少

KV cache大小与层宽度线性相关，所以总KV cache ∝ Σdℓ < L·d。实际约减少10-15%。

局限与工程挑战

论文非常坦诚地列出了实际落地的挑战：

1. 高效训练Kernel复杂

不同层有不同形状，需要为每种形状优化kernel（GEMM尺寸不同）。当前训练框架（PyTorch/FSDP/DeepSpeed）高度优化了统一宽度的场景，变宽度需要大量定制kernel开发。

2. 分布式并行挑战

• 张量并行：层内不同部分需要不同并行策略
• 流水线并行：层间宽度不同，激活张量形状变化，通信模式复杂

3. 残差流额外开销

全局残差流比最宽层还宽，涉及切片、复制、零填充等操作，引入额外的kernel启动开销。但论文指出，通过kernel融合可以大部分消除。

关键论点

这些不是算法限制，而是实现限制。当前基础设施被重度优化为统一宽度，不是架构本身的内在属性。专用Kernel可以弥合理论和实际效率的差距。

更深层思考：均匀性的偏见

这篇论文挑战的是一个比「Transformer层等宽」更广泛的偏见：

在深度学习中，我们过度偏好「均匀性」和「对称性」：

• 所有层相同宽度
• 所有层相同深度（虽然这个已经被打破，比如有early exit）
• 所有头相同维度
• 所有token被相同对待（虽然patching/mixing已经不同）

但生物神经网络从不均匀——不同脑区有不同密度、不同连接模式。视觉皮层的V1、V2、V4、IT区有不同的神经元密度和感受野大小。均匀性是人类工程方便的选择，不是自然优化的结果。

> <former告诉我们：不对称可以是优势，约束可以是正则化。 瓶颈强迫网络更有效地利用资源，避免「冗余安全感」导致的浪费。

这也呼应了另一篇近期论文（Ikeda et al., 2025）的发现：在MLP中集中更多计算到中间层有好处。但> <former走得更远——不是只调整MLP中间维度，而是全块宽度都变化，且发现最优形状是×形（与直觉的◇形相反）。

对工程实践的启示

1. 架构设计新维度：层宽度调度（width schedule）成为与深度调度、学习率调度并列的新超参数。

2. 推理效率直接收益：KV cache减少10-15%意味着：

   • 同样显存可以服务更长序列
   • 同样序列长度可以支持更大batch
   • 边缘部署的内存瓶颈得到缓解

3. 训练成本优化：FLOPs减少22%（在匹配Loss的前提下）意味着相同预算可以训练更大模型或更多数据。

4. MoE兼容性：变宽度与MoE结合，可以在active参数减少的情况下保持性能，进一步降低推理成本。

5. Kernel融合是必经之路：要让理论效率变为实际效率，需要在编译器/Kernel层面支持变宽度的GEMM和通信。

一句话总结（再强调一次）

Transformer每一层保持相同宽度是工程方便，不是数学必然。让中间层「瘦身」、头尾层「增宽」的×形架构，不仅性能更好，还更省算力和显存。瓶颈不是缺陷，而是结构正则化——它迫使网络避免中间层的表示坍塌，更均匀地利用每个维度。这是一个被忽视的自由度，可能改变未来语言模型的设计范式。

参考论文：Wu et al., "Variable-Width Transformers", arXiv:2606.18246, 2026.

#AI #大语言模型 #Transformer #架构设计 #深度学习 #高效计算 #MIT #MoE #缩放定律 #可解释性