摘要
包括Transformer、RNN和基于记忆机制的各类变体在内的现代语言模型,都采用了一种通用的架构:由若干结构完全相同的层堆叠而成,在深度方向上均匀分配参数 。但是现在的一些研究表明,不同层对最后输出的贡献是非均匀的:深层更多在于细化残差流(residual stream),而非对其进行转化。为此,作者提出参数容量是否反映出了这种不对称性。
作者在固定参数总量情况下,通过对照实验发现:与宽度均匀的基线相比,将更多容量分配给浅层、而将较少容量分配给深层可以降低困惑度;而反向分配则会损害性能。基于这一发现,作者提出了锥形语言模型(Tapered Language Models,TLM)架构设计原则,即在固定总预算的前提下,让承载参数的组件在深度上呈现单调递减(锥形化)的变化。
在所有现代语言模型家族中,MLP都占据了参数数量的大部分,并且提供"宽度"这一单一、清晰的可调维度。在三种模型规模和四种架构(Transformer、门控注意力、Hope-attention和Titans)上的实验结果表明,通过平滑余弦调度策略对MLP宽度进行锥形化处理,在无需增加任何额外参数或计算成本的情况下,始终能比均匀基线更有效地降低困惑度。这些发现确立了"深度感知的容量分配"作为语言模型设计中一条简单、与架构无关的语言模型设计维度,这是一个长期被忽视的直观易行的优化杠杆。
引言
现代的语言模型,尽管他们表面多样性,但是其实都共享同一种骨干。像Transformer、门控注意力变体、循环状态空间模型(recurrent and state-space models)以及基于记忆的架构模型,它们在tokens之间混合信息的方式不同,但是都是由L个相同的层堆叠而成。每个层包括了token混合模块和前馈网络(FFN)。在这个骨干网络中,参数在深度上均匀分布,每层都是相同的参数分配。
但是一些实验表明,不同层对最后的输出共享是不同的。很多模型在预训练到最后一层之前就已经定形了,最后几层只是在微调而已。如图一所示,作者在保持总参数不变的情况下,将MLP的参数分布分别沿着深度方向,做了逐步降低、线性降低和余弦降低,与均匀分布对比,其困惑度都得到了减低,其中余弦减少策略困惑度最低。

同时,针对如何非均匀分布参数,能够取得最佳结果这个问题。作者采用了A:浅层变宽、B:深层变宽和C:中间变宽三组实验说明了,随着层数的加深,参数量逐渐降低最好。因此在总参数预算固定时,采用非均匀分布参数设计网络架构,使得它们能够被正确分配到需要它们的地方,而不是均匀分配。
对于如何非均匀分配参数,将哪些地方这样设计这个问题,MLP是最自然的的选择。因为现代语言模型家族中,MLP是占据主导地位的参数储存单元,并且MLP的结构暴露了一个单一宽度的参数,即中间维度,这个参数独立于周围架构进行调整。此外,通过测量每一层向预训练后的Transformer的残差流中写入了多少新信息,发现了随着网络深度的增加,MLP的输出与现有的残差对齐成度逐渐提高,而不是计算新的特诊。这解释了为什么锥形化逐渐降低参数容量方法的有效性。
主要贡献:
- 提出"锥形语言模型"(Tapered Language Models,TLMs),这是一种架构原则:在总预算固定的情况下,带参数的组件会随深度单调递减;通过三种平滑递减方案(线性、余弦和 sigmoid)将该原则应用于多层感知机(MLP)的宽度,同时保持总参数数和浮点运算次数(FLOPs)不变。
- 证明了在三种模型规模和四种架构家族(Transformer、Gated Attention、Hope-attention 和 Titans)中,锥形化都能持续降低困惑度和提升下游基准测试性能,从而确立了深度感知型参数容量分配作为一种稳健且与架构无关的设计选择。
- 通过机制分析表明,随着深度增加,MLP 的输出与残差流的对齐程度逐渐提高,这直接解释了为何前置分配容量能提升性能。
锥形语言模型
Background
decoder-only的语言模型大多遵循由L个相同的块堆叠而成,每个都是由token-mixing模块M和多层感知机F组成。假设给定在l层的隐藏状态时,则每个block的组成方式为:
-
token-mixing:
(1),
和M的输出相加,形成中间状态
-
MLP变换:
(2),
和经过MLP变换后结果相加,得到下一层状态
其中不同架构的M是不同,如softmax attention (Transformer等)、gated attention、循环/记忆机制(RNN等),但是MLP组件F是一致的,采用深度可分离变换,由中间维度控制。
Tapering
在任意架构组件C中,第l层的维度为,其中
。均匀设计将所有l层固定为
,而对于锥形化的设计,需要满足
(3)
所以容量逐层递减,并且满足总容量固定。
Tapering MLP width

如上图所示,给出了三种递减策略:
Linear:以一个固定的比例去逐层递减。第l层维度为:
(4)
Cosine: 遵循半余弦曲线,在两个端点处(l=0和l=L-1)衰减缓慢,在中间衰减最快
(5)
Sigmoid: 由k>0控制陡峭程度(实验中k=10),在中间层附近快速下降,形成平滑的阶梯状过渡
(6)
实施细节 :第一层和最后一层的维度分别固定为,中间每层的宽度四舍五入至最接近16的倍数,并且满足公式(3),单调递减,只改变MLP的中间维度
,其它不变。
实验

表1给出了440M的Transformer中,使用三种递减策略和二u中宽度比例下的实验结果。标准的均匀设计的模型困惑度为16.28,列除了与标准的差值,绿色说明困惑度降低,红色反之,并且颜色的深浅表示相差的程度。可以看出,cosine策略下,维度从1.50到0.50时,模型效果最好。
逐层新颖性(Layer-wise Novelty)
前面的实验说明了锥形化参数量的有效性,为了探究其机制,在给定的参数容量前提下,测量每个层向预训练Transformer的残差流中写入了多少新信息。发现随着深度的增加,MLP的输出与现有残差流的对齐度逐渐提高,更多的是强化已有特征,而不是计算新的特征。锥形化使结构和这一模式对齐。

每层测量两个余弦量,根据公式(1)和(2),定义:
(7)
(8)
并在token上平均。表示写入与残差正交的内容,
对应强化残差已编码的方向。 随深度上升的趋势表明新颖性降低。
在GPT-2家族的WikiText-2的2048个token上计算和
,该家族提供了在受控架构下跨多种规模的公开预训练检查点,使我们能够在没有训练变异混淆的情况下隔离深度模式。首层和末层作为边界情况被排除。图4的模式在家族中一致:两个量在早中期层下降到低值,并在每个模型的网络后半部分上升。与层索引的皮尔逊相关性在每次测量中均为正,且MLP的相关性比完整块更紧密:
的相关性范围从r = 0.49到r = 0.71,
从r = 0.27到r = 0.71。两个量都上升说明了两点:表明MLP本身变得冗余,而非在注意力工作时闲置;并且显示深度模式延伸到整个层,而非仅限于MLP。
与锥形化的联系是直接的:若均匀容量下的后期MLP输出与残差对齐(而非正交),额外隐藏维度未被利用。锥形化减少后期高对齐度的隐藏维度,并将节省的参数重新分配到早期层,使架构与"新颖性随深度降低"的模式对齐。
局限性
在更大的模型和其它架构上的实验改进,表明所选配置是鲁棒的默认选项,但并未证明它对每个模型都是最优的。偏好的锥形剖面可能依赖于诸如模型深度、隐藏维度、分配给MLP的参数比例、token混合架构或训练预算等属性。特别是,随着模型规模的扩大,最佳调度或端点比可能会变化,从而可能带来超出此处报告的收益。
讨论和总结
TLMs提出了一种架构设计原则:在固定总预算下,承载参数的组件沿深度单调"锥形化"。通过在MLP中间宽度上实施平滑余弦锥形化,锥形化在三个模型规模和四种架构家族(涵盖softmax注意力、门控注意力、循环自修改记忆、神经长期记忆)上,相对均匀宽度基线降低了困惑度并提升了下游基准性能,并提供了机制证据------随着深度增加,MLP输出与残差流的对齐度逐步提高,直接为锥形化方向提供了动机。这些收益未增加参数数量或训练计算量。在动机实验、调度消融实验和宽度扫描中,一致的模式浮现:模型容量在早期层最具价值,而平滑衰减的分配是有效且鲁棒的分配方式。
除MLP宽度外,许多决定参数分配的逐层维度(如注意力头数、键值维度、循环状态大小、内存槽数量、混合专家模型中的专家数量)都是锥形化的自然候选。此外,除了语言模型外,视觉Transformer、扩散Transformer和多模态模型都继承了类似的均匀参数容量的默认模式,后续可以考虑探索其它维度的模型领域上的作用。