沙漏里的智慧:当Transformer学会收腰
> *一篇关于Variable-Width Transformers的深度解读* > > *论文:Wu et al., "Variable-Width Transformers", arXiv:2606.18246*
---
🏛️ 引子:帕特农神庙的秘密
你有没有站在帕特农神庙的柱子前,感到一种说不出的和谐?
古希腊的建筑师们早就发现了一个秘密:一根从头到尾一样粗的柱子,看起来会"胖"——它中间那段会因为视觉错觉而显得臃肿。所以他们在柱子的中间做了轻微的收缩,让整体看起来更加挺拔优雅。
这个技巧叫"收分"(entasis),是人类在建筑美学上的千年智慧。它不是装饰性的花哨,而是对视觉感知的深刻洞察:人眼会被绝对的均匀欺骗,而微妙的比例变化反而能唤醒更深的美感。
有趣的是,两千多年后的今天,一群AI研究员在做一件类似的事——他们让Transformer的中间层"收腰",结果发现,这座数字神庙不仅更优雅了,还更聪明了。
这就是我们今天要聊的Variable-Width Transformers,MIT和IBM的研究团队给它取了一个俏皮的名字:">
这篇论文的作者阵容堪称豪华: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个车道怎么办?填零?学一个投影矩阵来预测?
这就是>
让我用一个更生动的比喻来解释。
想象一条贯穿所有24层的传送带,宽度等于最宽的那一层。每一层工作时,只在这条传送带上的一段区域里"动手"——窄层只操作一小段,宽层操作一大段。如果某一层变窄了,它不再处理的区域怎么办?答案是:保持原样,直接传下去,就像没动过一样。
具体来说:
- 缩小时:如果第6层比第5层窄,那么多出来的维度就直接截断,不处理。这些维度上的信息原封不动地进入下一层。
- 扩大时:如果第7层比第6层宽,多出来的维度从哪里获取?从最近一次处理过这些维度的层复制过来。
作者还提出了一个等价的视角:可以把> > 小贴士:你可能会问,为什么"复制"比"学习"更好?这有点像搬家时整理箱子——与其试图预测每个物品在新家的位置(可能猜错,导致混乱),不如先把它们原封不动地搬过去,到了新家再重新排列。论文中的消融实验(表4)证实了这一点:carry-forward策略的loss是3.099,zero padding是3.124,训练投影层是3.150——甚至比基线还差。 ---📐 第三层数学:为什么收腰能省算力
好的,架构设计听起来 clever,但有没有理论保证呢?
有的,而且很漂亮。
假设基线模型有L层,每层宽度都是d,总参数量正比于Ld²。这里的平方来自于注意力机制:每个token需要与所有其他token计算相似度,复杂度是O(N²d),而参数数量正比于d²。
现在我们把模型改成沙漏形:早期层宽,中间窄,晚期再宽,但保持总参数量不变。也就是说:
$$\frac{1}{L}\sum_{\ell=1}^L d_\ell^2 = d^2$$
问题来了:平均宽度是多少?
根据Jensen不等式(一个关于凸函数的优雅定理),因为平方函数是严格凸的,我们有:
$$\left(\frac{1}{L}\sum_{\ell=1}^L d_\ell\right)^2 < \frac{1}{L}\sum_{\ell=1}^L d_\ell^2 = d^2$$
所以:
$$\frac{1}{L}\sum_{\ell=1}^L d_\ell < d$$
平均宽度严格小于基线!
这意味着什么?Transformer中最昂贵的操作之一是注意力机制的计算,它的复杂度与层宽成正比。平均宽度降低了,总的FLOPs(浮点运算次数)就减少了——论文显示能减少约22%。
这就是数学之美:我们没有任何额外的假设,仅凭凸函数的性质,就证明了沙漏形架构必然更省算力。这不是近似,不是启发式,这是严格的数学保证。
> 小贴士:Jensen不等式说的是,对于一个凸函数(比如平方),函数在平均点的值小于等于函数值的平均。想象一个碗的形状——碗底的点总是比碗沿的平均高度更低。这就是严格不等号的来源。如果你把弹珠放在碗沿,它会滚到碗底;如果你放在碗底,它会待在那里——这就是凸函数的"谷底"。
---
🧪 第四层实验:数字不说谎
理论再漂亮,也要看实验数据。作者训练了一系列模型来验证这个想法,从200M参数到3B参数的MoE模型。
形状对比实验(图2,500M规模)
他们尝试了四种宽度变化模式,结果出乎意料:
| 形状 | 描述 | 表现 |
|---|---|---|
| ∨形 | 从窄到宽(像一个漏斗) | 较差 |
| ∧形 | 从宽到窄(像一个倒漏斗) | 较差 |
| 3形 | 窄-宽-窄(中间最宽,像沙漏倒过来) | 中等 |
| ×形 | 宽-窄-宽(中间最窄,像真正的沙漏) | 最佳 ✓ |
为什么?论文后面的分析给出了答案:中间层过宽会导致"表示坍缩"——信息扩散到太多维度上,反而变得稀薄,就像把一杯浓咖啡倒进游泳池。收窄中间层反而像一个聚焦透镜,强迫信息在更紧凑的空间里保持密集和高效。
作者在论文里坦诚地写道:"我们最初的直觉是追求3形模型,增加中间层的计算...但经验结果促使我们采用×形。"这种对实验结果的尊重——即使它违背直觉——是优秀科学家的标志。
性能数据(表2)
他们在200M到3B参数的多个规模上进行了实验,结果一致:
| 规模 | 基线Loss | > | FLOPs减少 | 平均层宽减少 | |
|---|---|---|---|---|
| 200M | 3.452 | 3.430 | 3.2% | 10.0% |
| 500M | 3.138 | 3.099 | 3.7% | 11.0% |
| 1B | 2.926 | 2.890 | 2.6% | 10.5% |
| 2B | 2.751 | 2.726 | 2.5% | 10.9% |
| 3B MoE | 2.726 | 2.710 | 4.6% | 10.9% |
更惊人的是扩展定律分析(图4):当用拟合曲线比较时,> 在等效loss匹配条件下,达到2B基线模型同样质量的> 超参数搜索(图3)
在确定了×形是最优形状后,作者面临下一个问题:瓶颈应该放在哪里?应该多窄?
他们参数化了这个问题:
- $r_\ell$:瓶颈位置(相对于总层数的比例)
- $r_d$:瓶颈宽度(相对于基线宽度的比例)
最优配置是 $r_\ell = 0.75$,$r_d = 0.3$。
也就是说:瓶颈放在75%的深度处,宽度收缩到基线的30%。
这个发现有几个有趣的解读:
1. 瓶颈不在正中间:如果对称性是最优的,瓶颈应该在50%处。但最优位置是75%,这意味着模型需要更长的"扩张期"来准备输出。这符合直觉——生成答案需要丰富的表达能力,所以晚期层应该更宽。
2. 收缩幅度很大:30%意味着中间层只有基线宽度的不到三分之一。这是一个相当激进的收缩,但实验表明这是最优的。这说明中间层真的不需要那么多容量。
3. 跨规模一致性:200M、500M、1B三个规模都收敛到相似的配置,说明这不是偶然的过拟合,而是某种普适的结构规律。
扩展定律的深层含义(图4)
扩展定律(Scaling Laws)是AI领域最重要的发现之一。它告诉我们:模型性能随计算量、参数量、数据量的增长而可预测地提升。
> 拟合公式:
注意两个关键差异:
1. 截距更低(3.12 vs 3.16):同样的计算量,loss更低
2. 指数更陡(-0.0503 vs -0.0498):规模扩大时,优势可能进一步放大
第二点是更重要的。如果趋势持续,当模型规模达到100B或更大时,>
与MoE的协同效应
Mixture of Experts(MoE,混合专家模型)是另一种提升参数效率的方法——它用稀疏激活让模型在不增加计算量的情况下拥有更多参数。
有趣的是,> KV缓存优化的实际意义
在大语言模型的实际部署中,KV缓存往往是真正的瓶颈——不是计算,而是内存带宽。
每次生成新token时,模型需要读取之前所有token的key和value。对于长序列(比如4096 tokens),KV缓存可能占用几十GB的显存。
>
---
🔬 第五层解剖:为什么收腰让模型更聪明
论文用了一系列分析工具来"透视"模型的内部运作。这些发现不仅解释了为什么> 激活利用率(图5, 6)
在常宽基线模型中,作者发现MLP(多层感知机)的激活存在"死维度"——有些神经元几乎从不激活,白白占用参数预算。就像一个办公室里,有些座位永远空着,但房租照样付。
而在>
残差流矩阵熵(图7)
这里有一个特别漂亮的发现。作者测量了每一层残差流的信息熵(可以理解为"混乱程度"或"丰富程度")。
常宽基线模型在中间层出现了明显的"压缩谷"——归一化熵降到接近0,表示信息坍缩到了一个低维子空间里。换句话说,中间层虽然有那么多的维度,但实际上只在使用其中很小一部分。就像一个巨大的会议室,只坐了前排几个人。
而> 参与比(Participation Ratio, 图9)
参与比衡量的是"有多少维度真正有信息"。常宽基线的中间层参与比极低(<5%),意味着95%以上的维度是"死的"。而> Logit Lens预测动态(图8)
用一个叫Logit Lens的技术(可以"偷看"每一层对最终答案的预测),作者发现> ---🏗️ 第六层工程:理想丰满,现实骨感
当然,任何漂亮的研究都有现实的约束。作者在论文最后坦诚地列出了几个工程挑战:
1. 专用Kernel需求:不同宽度的层需要不同的GPU kernel优化,现有深度学习框架(PyTorch、JAX等)没有原生支持异构层宽度的计算图。
2. 并行性挑战:现有的张量并行(把一层切到多个GPU上)和流水线并行(把不同层放到不同GPU上)都假设所有层结构一致,异构层宽度会破坏这些假设。
3. 基础设施不匹配:当前的AI硬件(如NVIDIA GPU的Tensor Core)和软件栈(cuDNN、FlashAttention等)都是为统一宽度设计的。不同宽度的矩阵乘法无法充分利用硬件的向量化指令。
但作者强调:这些是工程限制,不是算法限制。就像卷积神经网络刚出现时也没有专用硬件,后来GPU和cuDNN让它飞起来。如果沙漏形架构的优势被持续验证,硬件和软件层面的适配只是时间问题。
实际上,论文已经给出了一个乐观的信号:即使是当前的通用框架,> ---🌌 尾声:从帕特农到硅基神庙
让我们回到帕特农神庙的柱子。
古希腊人没有流体力学的知识,也没有有限元分析。他们凭直觉和审美,在柱子上做了那个微妙的收分。两千年后,工程师们用计算机模拟才发现:这个设计不仅更美,还在结构上更稳定——它减少了风荷载下的涡振,让柱子在地震中更不容易断裂。
> 我们设计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架构 #效率优化 #小凯
🌟 智谱 GLM-5 已上线
我正在智谱大模型开放平台 BigModel.cn 上打造 AI 应用,智谱新一代旗舰模型 GLM-5 已上线,在推理、代码、智能体综合能力达到开源模型 SOTA 水平。
🎁 领取 2000万 Tokens