沙漏里的智慧：当Transformer学会收腰

一篇关于Variable-Width Transformers的深度解读

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

🏛️ 引子：帕特农神庙的秘密

你有没有站在帕特农神庙的柱子前，感到一种说不出的和谐？

古希腊的建筑师们早就发现了一个秘密：一根从头到尾一样粗的柱子，看起来会"胖"——它中间那段会因为视觉错觉而显得臃肿。所以他们在柱子的中间做了轻微的收缩，让整体看起来更加挺拔优雅。

这个技巧叫"收分"（entasis），是人类在建筑美学上的千年智慧。它不是装饰性的花哨，而是对视觉感知的深刻洞察：人眼会被绝对的均匀欺骗，而微妙的比例变化反而能唤醒更深的美感。

有趣的是，两千多年后的今天，一群AI研究员在做一件类似的事——他们让Transformer的中间层"收腰"，结果发现，这座数字神庙不仅更优雅了，还更聪明了。

这就是我们今天要聊的Variable-Width Transformers，MIT和IBM的研究团队给它取了一个俏皮的名字："> <former"——你看，这个名字本身就画出了一个沙漏的形状。> 是左边开口，< 是右边开口，中间那段窄窄的，就是瓶颈层。

这篇论文的作者阵容堪称豪华：Zhaofeng Wu（RoPE位置编码的共同提出者）、Oliver Sieberling、Shawn Tan、Rameswar Panda（IBM Watson的资深研究员）、Yury Polyanskiy（MIT信息论大牛）、Yoon Kim（MIT NLP领域的明星教授）。当这样一群人聚在一起挑战一个被默认了七年的基本假设时，你就知道这件事不简单。

小贴士：Transformer是当今所有大语言模型（如GPT、Claude、Gemini）的基础架构。你可以把它想象成一个拥有很多层的"思考流水线"，每一层都在前一层的基础上进一步提炼信息。传统上，每一层的"宽度"（即同时处理多少信息，专业术语叫"隐藏维度"）都是一样的。这就好比你开了一家工厂，24个工位，每个工位的大小、设备、人员配置完全一样——不管这个工位是负责原材料分拣还是最终质检。

🤔 第一层追问：为什么所有层必须一样宽？

让我们从最基本的问题开始，像费曼那样，假设我们对神经网络一无所知。

假设你正在组装一台精密的机器。这台机器有24个工作站，排成一条直线。每个工作站都一样大，处理信息的能力也完全相同。你投入了24份完全相同的资源，然后期待着奇迹发生。

但等等——这真的合理吗？

想想看人类的大脑。大脑皮层的不同区域，神经元密度并不一样。视觉皮层有极其密集的六层结构，专门处理来自视网膜的海量信息；而前额叶皮层的结构就相对松散，但它负责的却是最高级的抽象推理。小脑的神经元数量甚至超过了大脑皮层其余部分的总和——因为运动协调需要极其精确的时间控制。

进化花了数亿年才微调出这种非均匀的资源分配，而我们在设计神经网络时，却从2017年Transformer诞生那天起，就默认了"一视同仁"。

这种"均匀假设"最初来自哪里？

2017年，Vaswani等人在《Attention Is All You Need》中提出了原始的Transformer架构。那是一个优雅的数学框架，用自注意力机制（Self-Attention）和全连接前馈网络（Feed-Forward Network）构建了序列到序列的映射。在那个开创性的设计中，每个Transformer层确实使用了相同的隐藏维度——但这更多是为了简洁和可解释性，而不是因为有什么理论证明"等宽是最优的"。

后来的研究（Kaplan et al., 2020的Scaling Laws）专注于"多大"的问题——模型的总参数量如何影响性能——但对"怎么分配"的问题，所有人都默契地避开了。偶尔有工作（如DeLighT、OpenELM）尝试了层间缩放，但通常只是改变前馈网络的中间维度，而不是完整层的隐藏维度。

MIT的Wu等人站出来说：这不对。不同层在计算中可能扮演不同角色，固定的参数和计算预算不必在深度上均匀分配。

他们的核心假设简单到近乎天真，却直击要害：

"不同层在计算中可能扮演不同角色，固定的参数和计算预算不必在深度上均匀分配。"

换句话说，也许早期层需要做宽一点——因为它们要接收原始输入，处理的信息维度高、噪音大，需要从混乱中捕捉模式；也许晚期层也需要做宽一点——因为它们要输出最终答案，需要丰富的表达能力来生成多样化的输出；而中间层呢？也许它们只是在传递和压缩信息，不需要那么多资源。

这就像一条高速公路：入口和出口需要多车道来处理汇入和分流的车流，但中间那段也许两车道就够了——因为那里的车流已经稳定，不需要频繁的加减速和变道。

🔧 第二层设计：沙漏是怎么炼成的

好，假设我们接受了"不同层应该有不同宽度"这个想法。怎么实现呢？

这里有一个技术陷阱，很多人第一次听到"可变宽度"时不会想到，但它实际上是最难的问题。

Transformer有一个核心机制叫残差连接（residual connection），这是He Kaiming在2015年提出的ResNet中首次引入的。你可以把它理解为信息的高速公路——每一层处理完信息后，会把结果和原始输入相加，然后传给下一层。这个设计的妙处在于：即使某一层"偷懒"了（比如权重随机初始化导致传递了零向量），信息也能通过高速公路顺利通过，不会丢失。

但如果相邻两层的宽度不一样呢？

假设第5层有1000个神经元（维度是1000），第6层只有300个。高速公路突然从十车道变成了三车道——信息怎么过去？

反过来更棘手：从300扩到1000，多出来的700个车道怎么办？填零？学一个投影矩阵来预测？

这就是> <former最核心的创新所在：固定残差流的参数自由调整机制（Parameter-Free Residual Resizing）。

让我用一个更生动的比喻来解释。

想象一条贯穿所有24层的传送带，宽度等于最宽的那一层。每一层工作时，只在这条传送带上的一段区域里"动手"——窄层只操作一小段，宽层操作一大段。如果某一层变窄了，它不再处理的区域怎么办？答案是：保持原样，直接传下去，就像没动过一样。

具体来说：

• 缩小时：如果第6层比第5层窄，那么多出来的维度就直接截断，不处理。这些维度上的信息原封不动地进入下一层。
• 扩大时：如果第7层比第6层宽，多出来的维度从哪里获取？从最近一次处理过这些维度的层复制过来。

这种"向前追溯"的策略被作者称为"carry-forward"，论文里通过消融实验证明这是最优策略——比填零好，比学一个投影矩阵更好（后者甚至会发散）。

作者还提出了一个等价的视角：可以把> <former理解为一个统一宽度的模型，其中每一层只读写残差流的一个子集。这个视角很关键，因为它意味着> <former的实现可以复用大量现有的深度学习基础设施——你只需要在每层的输入输出处做维度切片和复制即可。

小贴士：你可能会问，为什么"复制"比"学习"更好？这有点像搬家时整理箱子——与其试图预测每个物品在新家的位置（可能猜错，导致混乱），不如先把它们原封不动地搬过去，到了新家再重新排列。论文中的消融实验（表4）证实了这一点：carry-forward策略的loss是3.099，zero padding是3.124，训练投影层是3.150——甚至比基线还差。

📐 第三层数学：为什么收腰能省算力

好的，架构设计听起来 clever，但有没有理论保证呢？

有的，而且很漂亮。

假设基线模型有L层，每层宽度都是d，总参数量正比于Ld²。这里的平方来自于注意力机制：每个token需要与所有其他token计算相似度，复杂度是O(N²d)，而参数数量正比于d²。

现在我们把模型改成沙漏形：早期层宽，中间窄，晚期再宽，但保持总参数量不变。也就是说：

问题来了：平均宽度是多少？

根据Jensen不等式（一个关于凸函数的优雅定理），因为平方函数是严格凸的，我们有：

所以：

平均宽度严格小于基线！

这意味着什么？Transformer中最昂贵的操作之一是注意力机制的计算，它的复杂度与层宽成正比。平均宽度降低了，总的FLOPs（浮点运算次数）就减少了——论文显示能减少约22%。

这就是数学之美：我们没有任何额外的假设，仅凭凸函数的性质，就证明了沙漏形架构必然更省算力。这不是近似，不是启发式，这是严格的数学保证。

小贴士：Jensen不等式说的是，对于一个凸函数（比如平方），函数在平均点的值小于等于函数值的平均。想象一个碗的形状——碗底的点总是比碗沿的平均高度更低。这就是严格不等号的来源。如果你把弹珠放在碗沿，它会滚到碗底；如果你放在碗底，它会待在那里——这就是凸函数的"谷底"。

🧪 第四层实验：数字不说谎

理论再漂亮，也要看实验数据。作者训练了一系列模型来验证这个想法，从200M参数到3B参数的MoE模型。

形状对比实验（图2，500M规模）

他们尝试了四种宽度变化模式，结果出乎意料：

注意这里有个反直觉的发现：作者最初以为中间层应该更宽（3形），因为中间层往往在做最复杂的特征提取和模式整合。但实验结果恰恰相反——沙漏形（×形）完胜。

为什么？论文后面的分析给出了答案：中间层过宽会导致"表示坍缩"——信息扩散到太多维度上，反而变得稀薄，就像把一杯浓咖啡倒进游泳池。收窄中间层反而像一个聚焦透镜，强迫信息在更紧凑的空间里保持密集和高效。

作者在论文里坦诚地写道："我们最初的直觉是追求3形模型，增加中间层的计算...但经验结果促使我们采用×形。"这种对实验结果的尊重——即使它违背直觉——是优秀科学家的标志。

性能数据（表2）

他们在200M到3B参数的多个规模上进行了实验，结果一致：

每一行都是更低的loss + 更少的计算。这不是妥协，这是双赢。

更惊人的是扩展定律分析（图4）：当用拟合曲线比较时，> <former不仅起点更低（同样计算量下loss更低），斜率也更陡——这意味着随着模型规模扩大，优势可能会进一步放大。

在等效loss匹配条件下，达到2B基线模型同样质量的> <former，只需要77.8%的FLOPs和85.1%的平均层宽。

超参数搜索（图3）

在确定了×形是最优形状后，作者面临下一个问题：瓶颈应该放在哪里？应该多窄？

他们参数化了这个问题：

• \(r_\ell\)：瓶颈位置（相对于总层数的比例）
• \(r_d\)：瓶颈宽度（相对于基线宽度的比例）

在200M、500M、1B三个规模上进行网格搜索后，发现了一个惊人的一致模式：

最优配置是 \(r_\ell = 0.75\)，\(r_d = 0.3\)。

也就是说：瓶颈放在75%的深度处，宽度收缩到基线的30%。

这个发现有几个有趣的解读：

1. 瓶颈不在正中间：如果对称性是最优的，瓶颈应该在50%处。但最优位置是75%，这意味着模型需要更长的"扩张期"来准备输出。这符合直觉——生成答案需要丰富的表达能力，所以晚期层应该更宽。

2. 收缩幅度很大：30%意味着中间层只有基线宽度的不到三分之一。这是一个相当激进的收缩，但实验表明这是最优的。这说明中间层真的不需要那么多容量。

3. 跨规模一致性：200M、500M、1B三个规模都收敛到相似的配置，说明这不是偶然的过拟合，而是某种普适的结构规律。

扩展定律的深层含义（图4）

扩展定律（Scaling Laws）是AI领域最重要的发现之一。它告诉我们：模型性能随计算量、参数量、数据量的增长而可预测地提升。

<former在这个框架下展现了一个令人兴奋的模式：

拟合公式：

• 常宽基线：\(L = 3.16 \times C^{-0.0498}\) (R²=0.999)

• <former：\(L = 3.12 \times C^{-0.0503}\) (R²=0.997)

注意两个关键差异：

1. 截距更低（3.12 vs 3.16）：同样的计算量，loss更低
2. 指数更陡（-0.0503 vs -0.0498）：规模扩大时，优势可能进一步放大

第二点是更重要的。如果趋势持续，当模型规模达到100B或更大时，> <former的优势可能从目前的几个百分点扩大到更显著的差距。这意味着沙漏形架构可能不只是"小修小补"，而是一条新的扩展路径。

与MoE的协同效应

Mixture of Experts（MoE，混合专家模型）是另一种提升参数效率的方法——它用稀疏激活让模型在不增加计算量的情况下拥有更多参数。

有趣的是，> <former与MoE的结合效果特别好：

• 3B/1B MoE模型上，> <former的相对优势（4.6% FLOPs减少）比dense模型更大
• 困惑度改善尤为显著

这表明两种效率提升机制可能是互补的：MoE在宽度维度上做稀疏化，> <former在深度维度上做非均匀分配。两者结合，可能开辟出一条"双稀疏"的高效架构路线。

KV缓存优化的实际意义

在大语言模型的实际部署中，KV缓存往往是真正的瓶颈——不是计算，而是内存带宽。

每次生成新token时，模型需要读取之前所有token的key和value。对于长序列（比如4096 tokens），KV缓存可能占用几十GB的显存。

<former的平均层宽减少10%直接意味着：

• KV缓存减少10%
• 内存带宽需求减少10%
• 长序列推理的延迟降低10%

在长上下文（100K+ tokens）成为标配的今天，这个优化可能带来比训练FLOPs减少更实际的收益。因为训练是一次性的，而推理是持续发生的——每一次用户提问都在消耗资源。

🔬 第五层解剖：为什么收腰让模型更聪明

论文用了一系列分析工具来"透视"模型的内部运作。这些发现不仅解释了为什么> <former有效，还揭示了Transformer内部表示的一些深层规律。

激活利用率（图5, 6）

在常宽基线模型中，作者发现MLP（多层感知机）的激活存在"死维度"——有些神经元几乎从不激活，白白占用参数预算。就像一个办公室里，有些座位永远空着，但房租照样付。

而在> <former中，激活密度更高，维度利用率更均匀。瓶颈层像一个结构正则化器，强迫网络在有限的容量里打包更密集、更有意义的表示。这类似于Dropout或L1正则化的效果——限制容量反而促使网络学习更鲁棒的特征。

残差流矩阵熵（图7）

这里有一个特别漂亮的发现。作者测量了每一层残差流的信息熵（可以理解为"混乱程度"或"丰富程度"）。

常宽基线模型在中间层出现了明显的"压缩谷"——归一化熵降到接近0，表示信息坍缩到了一个低维子空间里。换句话说，中间层虽然有那么多的维度，但实际上只在使用其中很小一部分。就像一个巨大的会议室，只坐了前排几个人。

而> <former：

• 早期层主动降低熵（为即将到来的"收腰"做准备，主动压缩信息）
• 瓶颈及后续层维持更高的归一化熵
• 避免了中间层的表示坍缩

这就像一条河：常宽模型是一条平坦但淤泥沉积的河道，> <former是先收窄加速、再放宽舒展的河道——水流更湍急，泥沙更少沉积。

参与比（Participation Ratio, 图9）

参与比衡量的是"有多少维度真正有信息"。常宽基线的中间层参与比极低（<5%），意味着95%以上的维度是"死的"。而> <former维持约1000个有效维度——即使总宽度减少了，有效维度反而增加了。

Logit Lens预测动态（图8）

用一个叫Logit Lens的技术（可以"偷看"每一层对最终答案的预测），作者发现> <former的预测动态更"从容"：

• 早期到中期，目标token的概率更高（模型更早"有了方向"）
• 中间层的分布更稳定、变化更渐进（不像基线那样剧烈波动）
• 最终层快速集中概率质量（一旦到达深层，模型迅速"拍板"）

换句话说，> <former就像一个提前做足了功课的学生——不是到最后才匆忙凑答案，而是层层推进，胸有成竹。

🏗️ 第六层工程：理想丰满，现实骨感

当然，任何漂亮的研究都有现实的约束。作者在论文最后坦诚地列出了几个工程挑战：

1. 专用Kernel需求：不同宽度的层需要不同的GPU kernel优化，现有深度学习框架（PyTorch、JAX等）没有原生支持异构层宽度的计算图。

2. 并行性挑战：现有的张量并行（把一层切到多个GPU上）和流水线并行（把不同层放到不同GPU上）都假设所有层结构一致，异构层宽度会破坏这些假设。

3. 基础设施不匹配：当前的AI硬件（如NVIDIA GPU的Tensor Core）和软件栈（cuDNN、FlashAttention等）都是为统一宽度设计的。不同宽度的矩阵乘法无法充分利用硬件的向量化指令。

但作者强调：这些是工程限制，不是算法限制。就像卷积神经网络刚出现时也没有专用硬件，后来GPU和cuDNN让它飞起来。如果沙漏形架构的优势被持续验证，硬件和软件层面的适配只是时间问题。

实际上，论文已经给出了一个乐观的信号：即使是当前的通用框架，> <former在训练FLOPs和KV缓存上都已经展示了明确的效率优势。KV缓存是推理阶段的关键瓶颈——每次生成新token，都需要读取之前所有token的key和value。平均层宽减少10%意味着KV缓存也减少10%，这在长序列推理中意义重大。

🌌 尾声：从帕特农到硅基神庙

让我们回到帕特农神庙的柱子。

古希腊人没有流体力学的知识，也没有有限元分析。他们凭直觉和审美，在柱子上做了那个微妙的收分。两千年后，工程师们用计算机模拟才发现：这个设计不仅更美，还在结构上更稳定——它减少了风荷载下的涡振，让柱子在地震中更不容易断裂。

<former的故事某种程度上是相似的。

我们设计Transformer时，默认了所有层等宽，因为这样做最简单、最容易实现。但当有人跳出这个假设，问了一句"为什么"时，一个被隐藏了多年的设计自由度突然显现了出来。

不是每一个"为什么"都能带来突破。但每一个突破，往往都始于一个"为什么"。

Wu等人用严格的数学证明和扎实的实验验证告诉我们：神经网络的深度维度上，存在着一个被遗忘的旋钮——宽度分配。把它拧到"沙漏"的位置，模型不仅更省资源，还更聪明。

这让人想起费曼说过的一句话：

"自然界总是以最经济的方式运作。"

也许，一个聪明的AI架构，也应该学会像自然界那样——不平均用力，而是把资源放在最需要的地方。

收腰，是为了更好地发力。

📚 参考文献

1. Wu, Z., Sieberling, O., Tan, S., Panda, R., Polyanskiy, Y., & Kim, Y. (2026). Variable-Width Transformers. arXiv preprint arXiv:2606.18246.
2. Kaplan, J., et al. (2020). Scaling Laws for Neural Language Models. arXiv:2001.08361.
3. Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017.
4. Levine, Y., et al. (2020). Depth-Width Trade-offs for Neural Networks. arXiv:2001.10648.
5. Mehta, S., et al. (2021). DeLighT: Deep and Light-weight Transformer. ICLR 2021.

#论文解读 #arXiv #Transformer #AI架构 #效率优化 #小凯