重新审视 Transformer 语言模型的形状约定 Revisiting the Shape Convention of Transformer Language Models
用宽-窄-宽（沙漏）FFN 替代传统窄-宽-窄 FFN，在 1B 以内参数规模达到更优或相当的性能
前置知识
Transformer FFN（前馈网络）
在 Transformer 中，每一层由注意力模块和前馈网络（FFN）组成。FFN 对每个 token 独立地进行非线性变换，通常采用窄-宽-窄的 MLP 结构：先将模型维度 dmodel 扩展到更宽的中间维度 dh（扩展比通常为 2-4），再投影回 dmodel。标准公式为 FFN(z) = z + W2 * sigma(W1 * norm(z))，其中 W1 属于 R^{dh x dmodel}，W2 属于 R^{dmodel x dh}。FFN 占据 Transformer 每层参数的大部分，通常超过注意力模块。
FFN 的形状直接决定了模型参数在深度、宽度和注意力之间的分配方式，是本文挑战的核心设计约定
Hourglass（沙漏）MLP
与传统窄-宽-窄 MLP 相反，沙漏 MLP 采用宽-窄-宽结构：先将输入投影到更窄的瓶颈维度 dh < dmodel，再投影回 dmodel。多个这样的子块可以通过残差连接堆叠，形成更深的网络。理论证明，具有残差连接的沙漏 MLP 保留了万能逼近能力（Lin & Jegelka, 2018），甚至可以实现最优逼近（Liu et al., 2024a）。
这是本文提出的核心替代方案，理解沙漏结构的理论优势是理解本文动机的关键
参数分配与模型形状
Transformer 的整体形状由三个关键维度决定：模型宽度 dmodel（隐藏维度）、深度 L（层数）以及 FFN 扩展比 dh/dmodel。这些参数共同决定了参数在注意力（上下文处理）和 FFN（逐点特征变换）之间的分配比例。传统 Transformer 通常是 FFN 主导的，FFN 参数约为注意力参数的 3 倍。
本文的核心贡献之一就是通过沙漏 FFN 实现了从 FFN 主导到注意力主导的参数再分配
SwiGLU 激活函数
SwiGLU 是 LLaMA 等现代 Transformer 广泛使用的激活函数，是 GLU（Gated Linear Unit）的变体。它使用 SiLU（Swish）激活函数对一个输入进行门控，其中 odot 表示逐元素乘法。相比简单的 ReLU，SwiGLU 引入了门控机制，提供了更丰富的非线性表达能力，但需要额外的参数。
本文的沙漏 FFN 实现基于 SwiGLU，每个子块包含两个下投影矩阵和一个上投影矩阵
研究动机
自 Transformer 架构提出以来，密集 Transformer 语言模型的 FFN 一直遵循一个固定的形状约定：窄-宽-窄 MLP，即先将模型维度 dmodel 扩展到更大的中间维度 dh（扩展比通常为 2-4），再投影回 dmodel。这一设计已被 GPT、LLaMA、Qwen、Gemma、OLMo 等几乎所有主流密集 LLM 采用。问题在于，FFN 占据了 Transformer 每层参数的绝大部分（通常是注意力模块的 3 倍左右），因此 MLP 的形状隐式地决定了模型容量如何在宽度、深度和注意力之间分配。这意味着一个未经充分检验的设计约定可能限制了模型架构的设计空间，导致参数分配不够高效。Kaplan 等人（2020）在 scaling law 研究中曾发现沙漏 MLP 效果更差，但这一结论从未在现代训练范式下被严格重新验证。
本文的目标是本文的目标是系统性地检验 Transformer FFN 中窄-宽-窄 MLP 形状约定的必要性。具体而言，作者希望验证：（1）用宽-窄-宽（沙漏）FFN 替代传统 FFN 是否能保持或超越性能；（2）沙漏 FFN 节省下来的参数是否能被更有效地利用，例如增大模型维度或增强注意力容量；（3）在 113M 到 1B 参数规模范围内，这种方法的可扩展性如何。最终目标是在固定的参数预算下找到更优的架构配置。
与已有工作不同的是，本文的独特视角在于，它不是在做新架构，而是在重新检验一个长期被视为理所当然的设计选择。作者抓住了一个被忽视的理论线索：沙漏 MLP 在理论上具有更强的函数逼近能力（万能逼近定理），在视觉生成任务中已经证明优于传统 MLP（Chen et al., 2025），但在语言模型中尚未被系统验证。此外，沙漏 FFN 通过缩小 FFN 宽度来释放参数，这些参数可以被再投资到注意力模块或增加 FFN 内部深度，这种参数再分配策略是传统 FFN 设计无法实现的。本文的切入点不是简单地比较两种 FFN，而是探索一个全新的参数分配范式：让注意力成为模型的主导组件，而非 FFN。
核心方法
本文的方法可以用一个直觉来理解：传统的窄-宽-窄 FFN 就像一个大力出奇迹的方案——用很宽的中间层来强行提升表达能力，但这种宽而浅的结构可能存在参数冗余。沙漏 FFN 则像是用精雕细琢的方式——通过多个窄而深的子块逐步精炼特征表示，就像反复用细砂纸打磨比用粗砂纸一次成型效果更好。具体技术路线是：将 Transformer 每层的传统 FFN 替换为一个由 K 个沙漏子 MLP 组成的堆栈，每个子 MLP 采用宽-窄-宽结构（先投影到瓶颈维度 dh < dmodel，再投影回 dmodel），子块之间通过残差连接。这使得 FFN 参数大幅减少，释放的参数预算可以重新分配给更大的模型维度 dmodel 或更多的注意力参数。
本文最核心的创新不在于沙漏 MLP 本身（这是已有的理论结果），而在于发现了一个在 Transformer 语言模型中被长期忽视的参数分配失衡：传统设计中 FFN 占据约 75% 的非嵌入参数，而注意力只占约 25%。通过沙漏 FFN，作者将这个比例颠倒——在 113M 模型中，注意力参数可以达到 66M 而 FFN 仅需 47M，注意力成为主导组件。这一发现的深刻含义是：对于较小规模的语言模型，上下文处理能力（注意力）比静态事实记忆（FFN）更宝贵。沙漏 FFN 不仅本身更高效，更重要的是它解锁了一种全新的参数再分配策略，使得模型可以在固定预算下同时优化宽度、深度和注意力容量之间的平衡。这种参数再分配是传统 FFN 设计根本无法实现的。
方法步骤详情
方法分为以下几个步骤：（1）定义沙漏 Transformer 层：每层包含一个标准注意力模块和一个沙漏 FFN，注意力模块的输入输出公式为 u^{(l)} = z^{(l)} + Attn(norm(z^{(l)}))。（2）构建沙漏 FFN：将注意力输出 h^{(l)}_0 = u^{(l)} 作为输入，通过 K 个沙漏子 MLP 依次处理，每个子块的输出为 h^{(l)}_{i+1} = h^{(l)}_i + MLP_i(h^{(l)}_i)，最终输出 z^{(l+1)} = h^{(l)}_K。（3）定义子 MLP 结构：采用 SwiGLU 激活，每个子块包含两个下投影矩阵和一个上投影矩阵，公式为 MLP_i(x) = W_u^{(i)} (SiLU(W_{d1}^{(i)} x_bar) odot W_{d2}^{(i)} x_bar))。（4）参数搜索：在 113M 规模下系统搜索最优的 dh/dmodel 比例、FFN 深度 K、Transformer 层数 L 和模型维度 dmodel 的组合，发现最优配置后外推到 403M、906M 和 1B 规模。（5）参数再分配：利用沙漏 FFN 释放的参数预算，增大 dmodel 使注意力参数占比提升。
技术新颖性
本文的技术新颖性体现在几个方面。首先，虽然沙漏 MLP 的理论性质已知，但这是首次将其系统性地应用于 Transformer 语言模型的 FFN 中进行大规模实验验证。其次，Kaplan 等人（2020）曾报告沙漏 MLP 效果更差，但本文在现代训练范式下推翻了这一结论，表明旧的负面结果可能受限于当时的实验条件。第三，作者发现的 U 形宽度-深度权衡曲线（dmodel/L 在 100-250 之间最优）是一个新的架构设计规律。第四，沙漏 FFN 允许 FFN 宽度和深度的独立控制（dh 和 K 是独立参数），这在传统 FFN 中是不可能的——传统 FFN 的宽度直接决定了参数量，无法单独调节深度。最后，参数从 FFN 向注意力的再分配范式是一种全新的设计理念，与传统的 FFN 主导设计形成鲜明对比。
实验结果
本文在 113M 到 1B 参数规模上进行了全面的实验验证，核心发现如下。首先，在固定 dmodel 和层数 L 的受控实验中（113M，dmodel=768，L=12），沙漏 FFN 变体（K=8，dh=384）的验证困惑度为 36.179，优于传统基线的 36.441，证明沙漏结构本身更高效。其次，在固定层数 L=12 的情况下搜索最优 dh/dmodel 比例，发现 K=4，dh/dmodel 约为 0.4 时效果最佳，表明 FFN 不需要统一的宽扩展就能有效工作。第三，最关键的发现是参数再分配实验：在 113M 模型中，(K=4, L=6) 配置的验证困惑度为 35.101，优于传统基线的 36.441，而且该配置将注意力参数提升到 45M、FFN 参数降低到 68M，颠倒了传统的分配比例。第四，宽度-深度权衡分析揭示了 U 形曲线，最优 dmodel/L 比例在 100-250 之间。第五，可扩展性验证表明：在 113M、403M、906M 规模上，沙漏 FFN 一致性地优于传统基线（验证困惑度分别从 36.441 降至 35.335、25.398 降至 24.909、22.473 降至 22.282）；在 1B 规模上达到相当水平（困惑度 20.082 vs 20.002），同时超越 OLMo2 配置（20.466）。在下游任务上，沙漏模型在 403M 以下的所有推理和 QA 任务中均表现更优，1B 规模下表现相当。消融实验还表明，增加 FFN 深度 K 从 1 到 8 持续改善性能（困惑度从 40.153 降至 32.832），dh/dmodel 从 0.4 到 0.8 之间性能差异不大，但低于 0.4 时性能显著下降。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 验证困惑度（113M） | Val PPL | 35.335 | 36.441 | 降低 3.0% |
| 验证困惑度（403M） | Val PPL | 24.909 | 25.398 | 降低 1.9% |
| 验证困惑度（906M） | Val PPL | 22.282 | 22.473 | 降低 0.8% |
| 验证困惑度（1074M） | Val PPL | 20.082 | 20.002 | 持平 |
| Arc Easy（113M） | Accuracy | 0.456 | 0.453 | +0.3pp |
| SciQ（1074M） | Accuracy | 0.825 | 0.806 | +1.9pp |
| HellaSwag（1074M） | Accuracy | 0.403 | 0.407 | 持平 |
| NaturalQS 困惑度（113M） | PPL | 1.537 | 1.604 | 降低 4.2% |
局限与改进
本文的局限性主要体现在以下几个方面。首先，作者承认全面的参数搜索仅在 113M 规模下进行，更大规模（403M-1B）的配置是从 113M 的发现外推而来，仅有有限的消融实验来确认这种外推的有效性。因此，沙漏 FFN 在数十亿甚至数百亿参数规模下的可扩展性仍是一个开放的实证问题。其次，本研究将注意力模块固定为标准多头注意力（MHA），未探索沙漏 FFN 与 Group-Query Attention（GQA）或 Multi-head Latent Attention（MLA）等高级注意力机制的交互作用，这些替代注意力机制可能会改变注意力-FFN 的参数分配格局。第三，1B 规模下沙漏 FFN 的相对增益明显缩小（困惑度仅改善 0.08），这可能暗示 FFN 有最低容量要求，进一步压缩空间有限。第四，训练数据规模较小（最大 21B tokens），未验证在大规模训练下的表现。第五，本文仅关注密集 Transformer，未涉及 MoE 架构中的沙漏 FFN 效果。
独立分析的弱点
本研究存在几个值得改进的弱点。第一，沙漏 FFN 的参数搜索策略主要基于穷举式网格搜索，缺少更高效的自动化架构搜索方法（如 NAS），在更大规模模型上这种搜索成本会急剧上升，建议未来可引入基于代理模型的贝叶斯优化或权重共享的超网络搜索来降低搜索成本。第二，所有实验使用相同的训练数据和 token 数量，未在 compute-optimal 设定下（如 Chinchilla scaling law）进行对比，这意味着某些架构配置可能在更充分的训练下展现出更大优势，建议在 compute-optimal 条件下重新评估。第三，沙漏 FFN 增加了网络深度，但未分析其对推理延迟和吞吐量的影响——更窄更深的结构可能在 GPU 上因为 kernel launch overhead 和内存访问模式不同而导致实际推理速度下降，建议增加推理效率的对比实验。第四，1B 规模的最优配置（K=1）实际上退化为传统 FFN 的变体，说明沙漏 FFN 的优势可能主要局限于中等规模，需要在 3B-7B 规模上进一步验证。第五，未与已有的高效 FFN 设计（如 LLaMA 的 SwiGLU、FlashAttention 等优化）进行联合对比。
未来方向
作者提出的未来工作方向包括：首先，验证沙漏 FFN 在更大模型规模（数十亿参数）下的有效性，特别是随着模型变大，深度的诅咒（信号传播衰减）可能对更深的沙漏 FFN 结构产生不利影响。其次，探索沙漏 FFN 与高级注意力机制（如 GQA、MLA）的组合，可能发现新的最优参数分配策略。第三，继续扩展模型维度 dmodel 和 FFN 深度 K，类似 manifold-constrained 架构（Xie et al., 2025）已表明扩展宽度的益处。基于本文成果可延伸的方向包括：将沙漏 FFN 应用于 MoE 架构中（Wang et al., 2024 已在 MoE 中验证了沙漏 MLP 的有效性）；研究沙漏 FFN 对模型可解释性的影响（更窄的瓶颈维度是否产生更清晰的特征表示）；以及将沙漏 FFN 的设计原则应用于视觉 Transformer 和多模态模型中。
复现评估
本文的复现条件较好。所有实验基于 OLMo 官方训练代码，使用 NVIDIA RTX 6000 Ada 和 B200 GPU，采用固定随机种子 6198。训练数据使用 OLMo-2 Stage 1 预训练语料的前 2.5B-21B tokens，并复用了 OLMo-2 1B 的精确数据排序以控制随机性。验证数据集使用 Dolma Common Crawl、Dolma The Stack、M2D2 和 WikiText 的标准划分。模型配置细节在论文的 Table 3 和 Table 7 中有完整记录。复现的主要障碍是：（1）需要多块高端 GPU 进行训练（1B 模型需要 21B tokens）；（2）超参数搜索在 113M 规模下进行，具体搜索空间和搜索次数未完全公开；（3）未提供预训练权重或完整代码仓库。总体而言，对于有足够算力的研究者，复现难度为中等。
论文图表
