技术摘要：重新审视 Transformer 语言模型的形状约定

问题陈述

稠密 Transformer 语言模型在历史上一直遵循一种僵化的架构约定：前馈网络（FFN）采用“窄–宽–窄”的多层感知器（MLP）形状。在这种设计中，模型维度（$d_{model}$）被扩展到一个更宽的中间维度（$d_h$），其扩展比例通常在 2 到 4 之间，然后再投影回原维度。这种结构主导了参数分配，通常占据了模型大部分的参数和浮点运算量。

作者质疑这种特定的形状是必不可少的，还是仅仅是一种便捷的实例化。最近的理论工作表明，残差型的“宽–窄–宽”（沙漏型）MLP 可以作为最优函数近似器，并可能以更少的参数提供更优越的表达能力。此外，在生成式视觉模型和混合专家（MoE）Transformer 中的经验证据表明，沙漏结构可以优于标准设计。核心问题在于：传统的“窄–宽–窄”FFN 是一种必要的约束，还是一个更深的、更窄的沙漏型 FFN 可以实现注意力模块与 FFN 模块之间更高效的容量分配？

方法论

本文提出了一种 Transformer 变体，用**沙漏型 FFN（Hourglass FFN）**取代了传统的 FFN。其方法论包括：

1. 架构重新设计：

   • 沙漏型 FFN 用由 $K$ 个通过残差路径连接的沙漏型子 MLP 组成的堆叠，取代了单个宽扩展层。
   • 每个子 MLP（$MLP_i$）遵循压缩-扩展结构：它将输入从 $d_{model}$ 投影到瓶颈维度 $d_h$（其中 $d_h < d_{model}$），应用非线性激活（SwiGLU），然后投影回 $d_{model}$。
   • 层输出是输入通过 $K$ 个此类残差块的结果。

2. 参数重新分配策略：

   • 通过降低中间维度 $d_h$（使 FFN 变得“更窄”），作者节省了通常分配给 FFN 宽度的参数。
   • 这些节省下来的参数被重新分配到以下方向：
     • 增加模型维度（$d_{model}$）和注意力容量。
     • 增加 FFN 堆叠的深度（$K$）。
     • 增加 Transformer 层数（$L$）。

3. 实验设置：

   • 基准测试： 与标准的 LLaMA 式 Transformer 以及最先进的 OLMo-2 架构进行对比。
   • 规模： 实验涵盖了从 113M 到 1B 参数的模型规模。
   • 约束： 实验在匹配的总参数预算内（差异在 0.001% 以内）且固定注意力大小的情况下进行，以隔离 FFN 形状的影响。
   • 搜索空间： 作者进行了架构搜索，改变了 $K$（沙漏堆叠深度）、$L$（Transformer 层数）、$d_{model}$ 以及 $d_h/d_{model}$ 的比例。

核心贡献

本文做出了四个主要贡献：

1. 挑战标准： 它挑战了事实上的“窄–宽–窄”FFN 形状，提出了一种高效的“宽–窄–宽”沙漏型 FFN，该结构通过堆叠的、具有残差连接的子 MLP 来保持表达能力。
2. 参数重分配： 它证明了一种策略，即通过将参数从 FFN 宽度转向注意力宽度或内部 FFN 深度（$K > 1$），可以获得更优的性能，尤其是在较小规模（113M）时。
3. 架构洞察： 它识别出了一个 U 型的宽度-深度权衡，找到了 $d_{model}/L$ 和 $d_h/d_{model}$ 的最优比例，并证实了更深的沙漏型 FFN 结构具有鲁棒性。
4. 可扩展性验证： 它在 113M 到 1B 参数范围内验证了沙漏型 FFN，证明其是竞争力的替代方案，在较小规模下表现优于传统模型，并在较大规模下与之持平。

结果

实证验证得出以下发现：

• 固定注意力下的性能： 在 113M 规模且固定注意力参数的情况下，沙漏型 FFN 比传统 FFN 实现了更低的验证损失和困惑度（PPL）。例如，在 $K=8$ 时，PPL 从 36.441（传统型）降至 36.179。
• 最优配置：
  • 深度 vs. 宽度： $d_h/d_{model} \approx 0.4$ 且 $K=4$ 层的配置产生了最低的验证损失。
  • 参数转移： 在 113M 规模下，通过转移参数使注意力模块占据主导地位（例如，66M 注意力 vs. 47M FFN）获得了最佳性能（PPL 35.101），这表明对于小型模型，上下文处理比静态事实检索更有价值。
  • 宽度-深度权衡： 观察到了 $d_{model}/L$ 比例明显的 U 型曲线，其最优甜点区位于 100 到 250 之间。
• 可扩展性：
  • 中小规模（113M – 906M）： 沙漏型变体在验证损失、困惑度和下游推理任务（如 Arc Easy、HellaSwag、PIQA）方面始终优于传统基准模型。
  • 大规模（1B）： 在 1B 参数规模下，沙漏型变体达到了与强力传统基准相当的性能（Val PPL 20.082 vs. 20.002），并超越了 OLMo-2 配置（20.466）。
• 消融研究： 增加沙漏深度 $K$ 持续提升性能，但在超过 $K=6$ 后收益递减。将 $d_h/d_{model}$ 降低至 0.4 以下会导致性能显著下降，这表明 FFN 压缩存在一个下限。

重要性与主张

本文声称，Transformer 中长期存在的“窄–宽–窄”MLP 形状并非必不可少，且可能存在冗余。通过采用沙漏型 FFN，作者证明了：

1. 效率： 模型可以利用更高效的参数分布实现具有竞争力或更优越的性能。
2. 资源重分配： 沙漏型结构实现了一种根本性的转变，即注意力机制可以成为主要的参数消耗者，挑战了传统的 FFN 主导型预算。
3. 重新评估设计： 这些发现促使人们重新思考注意力与 FFN 之间的平衡，表明对于高效且具表达能力的现代语言模型，更深、更窄的 FFN 结合更宽的注意力机制可能是一种更优的设计策略。

作者对向百亿或千亿级参数规模的可扩展性保持谨慎，指出其发现是从 113M 规模外推而来的，并且“深度诅咒”（信号传播问题）在更深的架构中仍需进一步研究。他们还指出，沙漏型 FFN 与高级注意力机制（如 Group-Query Attention）的相互作用仍有待探索。