精读:TLM——深度感知的参数容量分配
挑战Transformer均匀层宽的默认设计范式,通过余弦递减MLP宽度实现零成本困惑度改善1.84点
1 论文概览
一句话总结: 在固定参数量和计算预算下,将MLP宽度沿深度方向余弦递减(前1.5×/后0.5×基线宽度),在四种架构×三种规模上零成本实现困惑度和下游性能的全面提升——一个被忽视了8年的”免费杠杆”。
核心问题: 自2017年Transformer诞生以来,所有语言模型都沿袭了”均匀层宽”的默认设计——每一层的参数量完全相同。然而大量证据表明层的贡献高度非均匀:早期层进行核心表征变换,后期层更多是精炼残差流。为什么参数容量要”一视同仁”地分配给贡献迥异的层?
主要贡献:
- 首次系统性验证深度感知容量分配的因果效应——通过严格的四组对照实验(前宽后窄/前窄后宽/中间宽/均匀),证明分配方向导致>1个ppl点的差异,确立”前置容量”策略的因果优势
- 提出TLM架构原则——参数承载组件的容量沿深度单调递减,总参数量和FLOPs严格不变,从”长方形”变”楔形”
- 跨架构×跨规模的全面验证——同一配置(余弦调度、1.5/0.5比例)无需再调,在Transformer/Gated Attention/Hope-attention/Titans四种架构和760M/1.3B两种规模上一致有效
- 机制性解释——通过余弦相似度分析揭示后期层MLP输出与残差流高度对齐(“重复已有内容”而非”计算新特征”),为前置容量的合理性提供了因果链条
- 确立深度感知容量分配为语言模型设计的”免费维度”——零额外参数、零额外计算、零架构改动,纯结构重排即可收益
2 方法详解
架构/方法概述
TLM的核心干预极为简洁:将常数MLP中间维度 替换为深度依赖函数 ,同时约束所有层的平均值等于基线宽度:
选择MLP宽度作为变化维度有三个自然优势:
- 参数主导:MLP在所有现代LM家族中占据绝大多数参数量
- 单维清晰:宽度是单一、干净的变化轴,不像注意力头数那样牵涉多维耦合
- FLOPs可控:MLP的FLOPs正比于宽度,便于严格保持计算预算
三种递减调度
作者测试了三种平滑衰减曲线:
| 调度 | 公式 | 特征 |
|---|---|---|
| Linear | 匀速递减,两端突变 | |
| Cosine | 两端平缓,中段渐变 | |
| Sigmoid | 两端保持,中间急缩 |
余弦调度全面胜出,其最优配置()在440M Transformer上实现ppl从16.28降至14.44(-1.84),而相同比例下Linear仅-0.32,Sigmoid甚至+0.16。
余弦调度的优势可以用一个直觉来理解:它既不”一刀切”地损失两端的灵活性(vs Linear),也不平均用力错过最该收缩的地方(vs Sigmoid),而是让容量在前后端都平滑过渡,在中段逐渐收紧。
关键创新点
-
从经验观察到设计原则的跃迁:此前层重要性分析(层剪枝、提前退出)仅停留在”层不等价”的观察层面;TLM首次将其转化为可操作的结构化设计原则
-
参数总量和FLOPs的双重严格约束:不是简单地”加宽前层”,而是在预算内重新分配。消融实验确认参数量、训练FLOPs、推理FLOPs、优化器状态、学习率全部固定,唯一变量是每层MLP中间维度
-
440M超参→跨架构跨规模零调迁移:所有后续实验都直接套用440M上搜索到的最优配置,没有做任何规模/架构特定的调优。这种”一次搜索、到处使用”的特性极大增强了工程实用性
与基线方法的区别
- vs 层剪枝:层剪枝是”砍掉低贡献层”,TLM是”重新分配容量”——不丢弃任何层,只是让每层按需获得资源
- vs NAS:神经架构搜索代价极高且搜索空间巨大;TLM搜索空间极简(一个调度函数+一个比例),440M上的5×3=15次扫描即可确定最优配置
- vs MoE:MoE通过稀疏激活增加总参数量但限制激活参数;TLM总参数量不变,只是重新分配
- vs 非均匀层宽的零散尝试:此前没有系统性将非均匀分配提升为通用架构原则,也没有跨架构验证
3 实验结果
主要指标
440M Transformer上的调度×比例扫描
| 锥度范围 | Cosine ppl | Linear ppl | Sigmoid ppl |
|---|---|---|---|
| 1.25→0.75 | 15.18 (-1.10) | 15.96 (-0.32) | 16.16 (+0.16) |
| 1.375→0.625 | 14.59 (-1.69) | 15.80 (-0.48) | 16.16 (+0.16) |
| 1.50→0.50 | 14.44 (-1.84) ⭐ | 15.96 (-0.32) | 16.12 (-0.16) |
| 1.625→0.375 | 14.59 (-1.69) | 15.64 (-0.64) | 15.96 (-0.32) |
| 1.75→0.25 | 15.49 (-0.79) | 16.28 (0.00) | 17.12 (+0.84) |
均匀基线ppl = 16.28。关键发现:Cosine调度最差配置(15.49)仍优于Linear最佳配置(15.64),说明调度形状比锥度比例更重要。
1.3B规模——困惑度
| 模型 | WikiText Uniform | WikiText Tapered | Δ | LAMBADA Uniform | LAMBADA Tapered | Δ |
|---|---|---|---|---|---|---|
| Transformer | 17.39 | 17.17 | -0.22 | 17.62 | 16.93 | -0.69 |
| Gated Attn | 16.03 | 15.92 | -0.11 | 14.26 | 14.11 | -0.15 |
| Hope-Attn | 15.91 | 15.94 | +0.03 | 15.48 | 14.92 | -0.56 |
| Titans | 16.05 | 15.76 | -0.29 | 14.19 | 14.04 | -0.15 |
7/8 WikiText配置改善,8/8 LAMBADA配置改善。Hope-attention 1.3B是唯一轻微退化(+0.03)的配置。
1.3B规模——8项常识推理基准平均准确率
| 模型 | Uniform | Tapered | Δ |
|---|---|---|---|
| Transformer | 56.05 | 56.38 | +0.33 |
| Gated Attention | 56.51 | 56.80 | +0.29 |
| Hope-attention | 56.95 | 57.05 | +0.10 |
| Titans | 56.73 | 57.08 | +0.35 |
760M规模——困惑度与准确率
| 模型 | WikiText Δ | LAMBADA Δ | 准确率 Δ |
|---|---|---|---|
| Transformer | -0.44 | -1.04 | +0.59 |
| Gated Attention | -0.76 | -0.41 | +0.27 |
| Hope-attention | -0.12 | -0.22 | +0.36 |
| Titans | -0.81 | -0.17 | +0.99 |
Titans在760M规模上获益最大,WikiText改善0.81 ppl、准确率提升0.99pp。
与 SOTA 的对比
值得注意的是,本论文并非在benchmark上追求SOTA,而是证明一个架构设计原则的普遍有效性。其对比对象是同等参数量、同等FLOPs的均匀宽度基线——这是最严格意义上的”苹果对苹果”比较。在440M模型上1.84个ppl点的改善,换算到更大规模可能意味着更显著的绝对差距。
消融实验关键发现
1. 四种分配方案的因果验证
| 方案 | 描述 | 验证ppl | Δ vs Uniform |
|---|---|---|---|
| Uniform | 均匀分配 | 16.28 | 0.00 |
| Wider-Early | 前段容量大 | 15.96 | -0.32 ✅ |
| Wider-Late | 后段容量大 | 17.29 | +1.01 ❌ |
| Wider-Middle | 中间容量大 | 16.61 | +0.33 ❌ |
分配方向导致>1.3个ppl点的差距(15.96 vs 17.29),这一结果强有力地证明了容量的空间分配方向是关键因素,而非简单的不均匀分配本身。
2. 长上下文检索(Needle-in-a-Haystack)
1.3B Transformer在4K/8K/16K三个上下文长度、四种检索变体上的测试表明:Tapered模型匹配或略优于Uniform对应物。例如MQ-NIAH在16K上下文下,Uniform 27.2% → Tapered 27.6%。这排除了”前置容量以牺牲长程能力为代价”的担忧。
3. 机制分析:后期层的”冗余性”
通过测量GPT-2系列模型中MLP输出与残差流的余弦相似度 ,发现:
两个关键量( 和 )都随深度单调上升。
- 表示正交写入新内容
- 表示强化已有方向(“重复”已有内容)
Pearson相关系数 在0.49-0.71之间(与层索引的相关性),MLP信号的单调趋势更清晰。这直接解释了为什么从冗余的后段层移出容量到创新的前段层是有益的。
4 个人分析
亮点
-
问题定义的范式价值:在”堆参数”和”换架构”成为主流的2026年,回到最基本的设计假设——“层为什么要一样宽?“——这种第一性原理的思考方式极具启发性。正如Sutton的”苦涩教训”所指出的,很多时候真正的问题不是”能力不够”,而是”已有能力没有被正确使用”
-
因果推断的严谨性:不是简单说”不均匀更好”,而是通过四组对照(前宽/后宽/中宽/均匀)确定方向性因果——只有”前置容量”才有效。Wider-Late导致+1.01 ppl退化这一反直觉结果,比Wider-Early的改善更能说明问题
-
跨架构泛化的深刻含义:在Transformer/Gated Attention/Hope-attention/Titans四种token-mixing机制完全不同的架构上均有效,说明容量分配不均是残差连接架构的内在属性,而非特定注意力机制的副产物
-
极简方法论的力量:5×3=15次扫描在440M模型上确定最优配置,然后零调迁移到所有其他设置。这种”小模型搜索、大模型部署”的范式对工程实践有直接价值
-
“免费杠杆”的经济学意义:零额外参数、零额外FLOPs、零架构改动,仅修改MLP宽度配置。在算力成本高企的当下,这种”重新分配而非堆料”的思路具有显著的现实意义
局限性
-
超参搜索规模有限:最优配置仅在440M参数模型上搜索。作者坦诚:“是否存在更适合不同规模、不同架构的’专属配方’,仍是开放问题”。从440M到100B+的迁移是否保真,需要进一步验证
-
评估基准偏窄:仅评估了8项常识推理基准(LAMBADA/PIQA/HellaSwag/WinoGrande/ARC-e/ARC-c/SIQA/BoolQ),未报告MMLU、GSM8K等核心推理基准。考虑到不同能力可能依赖不同深度的层,缺乏推理基准的结果是一个明显的信息缺口
-
仅针对解码器模型:研究专注于因果语言模型,在编码器(BERT系列)、编码器-解码器(T5系列)、视觉Transformer、扩散模型等其他架构上的适用性未验证
-
未探索注意力头/键值维度的tapering:论文指出MLP是”自然选择”,但注意力头的数量和维度也是参数分配的重要维度。是否可以在多维度同时tapering,可能产生协同效应也可能产生干扰
-
缺乏与现有优化技术的交互实验:TLM与量化、剪枝、蒸馏、MoE等主流优化技术的兼容性如何?能否叠加收益?这些关键工程问题未被触及
-
机制分析的因果方向未完全建立:虽然随深度递增与前置容量的有效性相关,但”后期层冗余→前置容量有益”的因果链还缺少一个关键环节——如果我们强制后期层不再冗余(例如通过训练目标),TLM的优势是否消失?
对大模型领域的影响
-
架构设计的”隐藏维度”被打开:如果MLP宽度可以非均匀分配,那注意力头数、KV维度、甚至层间dropout率是否也可以深度感知地分配?这可能开启一整个研究方向的探索
-
对训练效率的直接影响:在LLaMA等大规模训练中,零成本的ppl改善直接转化为更少的训练步数或更好的最终性能。考虑到LLaMA 405B的训练成本,即使0.1 ppl的改善也意味着可观的资金节省
-
与Scaling Law的交互:如果非均匀分配在更大规模上效果更显著(或更弱),这将直接影响Chinchilla式的最优模型大小预测。这与我们此前讨论的TPP Coverage论文(缩放定律外推的可靠性问题)形成呼应
-
对模型压缩的启示:如果后期层确实冗余,那么量化/剪枝应该更激进地作用于后段层。这为混合精度训练和推理提供了新的理论基础
5 实践启示
可以借鉴的技术
-
余弦MLP宽度递减作为默认配置:在新模型训练中,将 的余弦MLP宽度调度作为默认起点,几乎没有任何下行风险(即使最差的tapering配置也不劣于均匀基线)
-
“小模型搜索→大模型部署”范式:在440M模型上快速扫描最优tapering配置,然后直接应用于目标规模,避免大规模超参搜索的昂贵成本
-
余弦相似度作为层冗余度诊断工具: 可以用于诊断现有模型中哪些层存在冗余,指导剪枝/量化策略
-
深度感知的混合精度分配:根据的分布,对冗余度更高的后段层使用更低精度(如FP4/INT4),对创新度更高的前段层保持高精度(如BF16/FP8)
潜在应用场景
-
大规模预训练成本优化:在固定训练预算下,TLM可以零成本获得更好的ppl,等效于节省训练步数。对于LLaMA级别的训练,这直接转化为百万美元级的成本节省
-
边缘部署的效率提升:TLM后期层更窄,天然适合渐进式推理(early exit)——在简单输入时可以更早退出,因为前段层已经拥有了更多的处理容量
-
视觉/多模态模型架构优化:ViT和扩散模型同样继承”均匀层宽”的默认设定。如果容量分配不均是残差架构的内在属性,那TLM原则应该可以直接迁移
-
MoE的深度感知专家分配:当前MoE的专家数量在各层间通常是均匀的。如果结合TLM原则,前段层使用更多专家(更大的激活参数),后段层使用更少专家,可能实现MoE的进一步优化
-
模型蒸馏的新策略:在知识蒸馏中,根据源模型各层的信息新颖度(值)来决定蒸馏权重的分配——对高新颖度层投入更多蒸馏预算
“These findings establish depth-aware capacity allocation as a simple, architecture-agnostic axis of language model design, a free lever hidden in plain sight.”
—— Bayat, Behrouz & Courville, Tapered Language Models, arXiv:2606.23670
这篇论文的真正价值不在于1.84个ppl点的绝对改善,而在于它提醒我们:在追逐”更多参数”和”更新架构”之前,也许应该先检查——我们是否已经把已有的参数放对了位置? 在2026年行业全面转向MoE和稀疏架构的浪潮中,这个来自Mila的”免费午餐”提示了一个被长期忽视的事实:有时候,优化不是做更多的事,而是把同样的事做得更合理。