跳转至

Sparse Layers are Critical to Scaling Looped Language Models

会议: NeurIPS 2026
arXiv: 2605.09165
OpenReview: https://openreview.net/forum?id=zt77AglOvq
论文: https://arxiv.org/abs/2605.09165
领域: LLM 其他
关键词: 循环语言模型, 混合专家模型, 缩放定律, 早退推理, 路由分流

一句话总结

本文通过严格的等算力（IsoFLOP）对照实验证明，稀疏混合专家（MoE）层能通过跨循环路由分流打破权重绑定的表达力瓶颈，解决循环语言模型（Looped LM）的缩放赤字，并在循环边界实现天然且高效的免训练早退推理。

研究背景与动机

标准自回归语言模型依赖堆叠数十甚至上百层独占参数的 Transformer 计算层。这种架构存在显著的参数效率瓶颈：每个计算层消耗数十亿参数的显存占用，而在单步自回归推理中，每个参数仅被激活使用一次。循环语言模型（Looped Language Models，如 Universal Transformers、LoopLM 和 Ouro）通过在深度方向上将一组固定物理层重复循环执行多轮，显著压低了模型存储与权重显存占用，并提供了天然的分阶段计算边界。然而，现有研究普遍发现循环架构在训练算力对齐（IsoFLOP）的前提下，扩展表现严格劣于具有独立参数的标准 Transformer。

这种扩展性赤字的核心矛盾在于前馈网络（FFN）的表达力瓶颈。在经典的稠密循环模型中，每一轮循环都对所有 token 无差别地执行完全相同的非线性变换，这种逐层不变性在循环复合下急剧压缩了模型的特征空间与表征容量，导致参数绑定直接牺牲了 Scaling 潜力。此前虽有工作尝试将循环与稀疏专家（MoE）结合，但往往伴随注意力头或归一化层等多处架构改动，缺乏因果明确的对照分析，使得稀疏化是否能够拯救循环架构的扩展性悬而未决。

本文的切入角度是剥离无关架构干扰，通过极简且受控的对比研究揭示稀疏路由对权重绑定的本质补益。核心 idea：在循环模型中用 Top-k 稀疏 MoE 替换稠密 FFN，利用 token 在多轮循环中激活不同物理专家的路由分流机制恢复网络表达力；同时利用循环边界天然对齐输出层的结构特性，实现无须额外训练的高保真早退推理。

方法详解

整体框架

本文构建了具有相同解码器骨干的统一实验平台：基于 Multi-Head Attention、RoPE 旋转位置编码、Pre-RMSNorm（无增益偏置）与 SwiGLU 结构的 Transformer。对比矩阵严格控制两个自变量维度——FFN 是稠密还是稀疏（Top-k MoE），以及网络深度是常规独占层还是多轮循环结构（Looped）。输入序列 token 在进入模型后，历经 \(L\) 个物理层构成的基础模块，并循环执行 \(R\) 轮，构成等效深度为 \(L \times R\) 的前向计算路径。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    A["输入 Token 序列"] --> B["输入嵌入与位置编码"]
    B --> C["循环模块第 1 轮前向"]

    subgraph S1 ["跨循环专家路由分流"]
        direction TB
        C --> D["自注意力层 MHSA"]
        D --> E["稀疏 MoE 路由：激活专家集合 T1"]
        E --> F["循环边界状态评估"]
    end

    F -->|满足熵阈值条件| G["免训练循环边界早退"]
    F -->|未满足早退阈值| H["循环模块第 2 轮前向"]

    subgraph S2 ["深度特化表达恢复"]
        direction TB
        H --> I["共享自注意力层 MHSA"]
        I --> J["稀疏 MoE 路由：分流至专家集合 T2"]
    end

    J --> K["最终 RMSNorm 与 LM Head"]
    G --> K
    K --> L["输出 Token 分布预测"]

关键设计

1. 跨循环专家路由分流：动态计算恢复网络表达力

稠密循环 Transformer 的根本缺陷在于同一批 token 在每一轮循环经过物理 FFN 层时均受到同一套权重的固化变换。Looped-MoE 将固定的 SwiGLU FFN 替换为由 \(E=8\) 个专家组成的稀疏层，并采用可学习的门控路由器 \(W_{\text{router}} \in \mathbb{R}^{d_{\text{model}} \times E}\) 选出激活的 Top-\(k\)（\(k=2\)）个专家。由于序列隐状态经过注意力与残差累加后在不同循环轮次发生了表征演化，同一 token 在第二轮循环经过同一物理层时，门控网络会为其分派完全不同的专家子集：

\[\mathcal{T} = \text{top-}k(h(x)), \quad p_i(x) = \frac{e^{h(x)_i}}{\sum_{j \in \mathcal{T}} e^{h(x)_j}}, \quad y = \sum_{i \in \mathcal{T}} p_i(x) \, \text{FFN}_i(x)\]

机制追踪表明，在各物理层中，高达 25%–53% 的 token 在两轮循环之间分配到的专家集合完全不重叠，且绝大多数 token 至少切换了一个专家。这使得物理上共享的同一参数层在不同计算深度上动态演化出差异化的等效子网络，以近乎不增加激活参数的代价消除了权重共享的退化瓶颈。

2. 免训练循环边界早退：输出层结构性天然对齐

常规 Transformer 执行中间层早退（Early Exit）通常需要专门设计辅助预测头、引入蒸馏损失，或者采用 LayerDrop 等破坏性正则化训练。而在循环模型中，由于第 \(r\) 轮循环的终点层恰好就是产生最终输出分布的物理终点层，该层在预训练过程中天然被优化为对接最终 LayerNorm 与 LM Head 的输出空间。因此，模型在循环边界处的特征分布对词表预测具有结构性的自适应成熟度。

测试推理阶段直接通过词表输出分布的信息熵判定是否早退：当输出熵小于预设阈值 \(\tau\) 时立即退出，将后续循环的理论计算量直接折算为节省的 FLOPs。通过 Jensen-Shannon 散度（JSD）对比各层输出分布与最终输出分布的逼近过程可以发现，循环模型在首轮循环结束即跳跃式收敛至最终预测，不仅完全免去微调训练，其 Pareto 前沿还严格超越了允许在任意层早退的标准非循环模型。

3. 循环与稀疏 μP 缩放对齐：消除 IsoFLOP 超参偏置

为了精确测算四类架构的缩放定律，必须杜绝不同宽度下学习率未最优化导致的系统性偏差。本文将最大更新参数化（Maximal Update Parameterization, \(\mu\)P）无缝推广至循环稀疏架构中。对于输入宽度扩展比率 \(w_{\text{ratio}} = d_{\text{model}} / d_{\text{base}}\)，非嵌入权重初始化方差按 \(\sigma^2_{\text{base}} / w_{\text{ratio}}\) 缩放，学习率按 \(\eta_{\text{base}} / w_{\text{ratio}}\) 缩放，并将该机制统一应用至稀疏路由权重 \(W_{\text{router}}\)、各专家参数 \(W_{\text{experts}}\) 以及输出解嵌入矩阵 \(W_{\text{unembed}}\)。

对于循环绑定的物理层，其梯度天然是多轮反向传播的累加，在标准 \(\mu\)P 体系下无需引入额外的人工缩放因子。通过在小宽度（\(d_{\text{base}}=128\)）下网格搜索最优学习率，可直接无损零样本迁移至数倍宽度的模型训练，验证各尺寸模型损失漂移低于 0.8%，保证了 IsoFLOP 缩放曲线拟合的严谨性。

损失函数 / 训练策略

模型预训练在 100 亿 Token 的 FineWeb-Edu 子集上展开，采用 AdamW 优化器配合 Warmup-Stable-Decay（WSD）学习率调度，并在训练末期 10% 步数内通过平方根衰减（sqrt-decay cooldown）平滑降至峰值学习率的 5%。总体优化损失函数由交叉熵语言建模损失 \(\mathcal{L}_{\text{LM}}\) 与两项防止 MoE 路由崩溃的辅助损失构成：

\[\mathcal{L} = \mathcal{L}_{\text{LM}} + \lambda_{\text{LB}} \mathcal{L}_{\text{LB}} + \lambda_{\text{RZ}} \mathcal{L}_{\text{RZ}}\]
\[\mathcal{L}_{\text{LB}} = E \sum_{i=1}^E f_i \cdot \bar{p}_i, \qquad \mathcal{L}_{\text{RZ}} = \frac{1}{B} \sum_{j=1}^B \left( \log \sum_{i=1}^E \exp(h(x_j)_i) \right)^2\]

其中 \(f_i\) 表示路由给专家 \(i\) 的 token 比例，\(\bar{p}_i\) 为其平均路由概率；负载均衡损失 \(\mathcal{L}_{\text{LB}}\) 保证各专家负载均匀，路由 Z-loss \(\mathcal{L}_{\text{RZ}}\) 惩罚过大 logits 以维持训练数值稳定性。

实验关键数据

主实验

在等算力预算为 \(10^{18}\) 训练 FLOPs 的严苛设置下，本文对比了四类计算最优架构在 AI2 OLMES 常识推理基准套件上的表现（Table 4 & Table 6）。循环 MoE 架构（Looped-MoE）凭借参数复用与稀疏路由的结合，以显著更少的存储参数超越了稠密非循环基线（Base）。

模型架构 激活参数 (M) 存储参数 (M) ARC-Easy ARC-Challenge BoolQ CSQA PIQA WinoGrande Core 9 均分
Looped (8×2) 246 168 39.1 24.6 49.6 27.1 57.1 51.2 37.4
Base (16层稠密) 246 246 39.3 25.3 51.4 29.3 57.8 50.5 38.7
Looped-MoE (8×2) 246 216 40.4 24.3 63.9 30.9 55.2 51.7 39.6
MoE (16层稀疏) 246 366 38.9 23.0 39.0 30.0 56.1 49.7 36.4

消融实验

针对推理时基于输出熵门控的免训练早退机制，在 \(10^{18}\) FLOPs 预算训练的最优模型上评估不同计算量节省比例下的测试困惑度（Perplexity，Table 7）。循环架构在早退下的表现显著优于非循环模型，且更深度的循环显著拉宽了效率优势。

模型配置 循环与层数结构 全深度 PPL 节省 5% FLOPs 节省 10% FLOPs 节省 20% FLOPs 节省 30% FLOPs
Base 16 层稠密 (无循环) 34.8 43.1 55.4 112.6 272.6
MoE 16 层稀疏 (无循环) 34.8 51.9 75.7 167.5 369.2
Looped 8 层 × 2 轮循环 36.2 40.9 50.2 84.3 156.0
Looped-MoE 8 层 × 2 轮循环 35.9 41.1 51.0 89.7 182.6
Looped-MoE 4 层 × 4 轮循环 35.4 38.2 44.3 71.0 160.1
Looped-MoE 2 层 × 8 轮循环 37.3 38.8 42.0 57.6 153.5

关键发现

  • 稀疏专家彻底扭转循环模型缩放赤字：在相同训练 FLOPs 下，稠密 Looped 模型测试损失全程落后于 Base 模型；而 Looped-MoE 凭借幂律缩放指数 \(\alpha = 0.077\) 全面反超 Base 模型（\(\alpha = 0.076\)），证明了稀疏化是循环架构 Scaling 的关键催化剂。
  • 早退韧性完全来源于循环架构而非 MoE：实验显示非循环 MoE 在早退时性能恶化最快（节省 10% 算力时 PPL 恶化至 75.7），而 Looped 稠密与 Looped-MoE 均能在早退时维持优异困惑度。这证实了早退优势并非源于专家稀疏性，而是因为循环模型的末层本身就是输出层，中间循环边界天然输出了结构健全的表征。
  • 更密集的循环进一步抬升早退 Pareto 前沿：随着循环次数增加（\(8 \times 2 \to 4 \times 4 \to 2 \times 8\)），循环边界早退点由 1 个扩充到 7 个。Looped-MoE \(2 \times 8\) 在节省 20% FLOPs 时仍能取得 57.6 的超高保真 PPL（相比 Base 模型的 112.6 优势极其巨大）。

亮点与洞察

  • 路由分流突破权重绑定上限：通过简单的 Top-k 动态路由，共享参数层在不同循环阶段自动激活不同的专家子网络，打破了递归神经网络长期以来的静态特征收缩瓶颈。
  • 内生性早退对齐：传统早退算法必须依赖额外的 Head 结构或复杂正则化训练，而循环模型天生让每一轮循环的结束点复用最终输出预测层，以零训练成本实现了表征的极早收敛。
  • 高知识容量与低显存占用的优雅兼顾：Looped-MoE 仅需保存极少的物理层权重，同时通过 MoE 专家池保留了极高的知识存储容量，成为受限显存端侧部署大模型极具潜力的架构范式。

局限与展望

  • 受限于深度缩放的 \(\mu\)P 适用性：当前 \(\mu\)P 体系主要针对宽度缩放成立，无法跨越物理层数变化。探索如 CompleteP 等联合深度与宽度的参数化方案是进一步完善循环模型缩放定律的必要方向。
  • 理论算力节省尚未转化为工程加速吞吐：当前早退实验基于理论层跳过统计 FLOPs，而在真实硬件批量推理中，变长序列早退导致的计算碎片化与专家动态加载开销仍需定制化推理引擎（如 Flash-Decoding / vLLM 调度层）进行实机验证。
  • 规模验证受限：由于算力约束，本研究验证的最大模型尺寸在 305M 激活参数与 711M 存储参数区间，未来仍需在 1B 至 7B 乃至更大规模上开展扩展性终验。

相关工作与启发

  • vs Ouro / LoopLM: Ouro 等在超大模型上尝试了循环架构，但因采用稠密 FFN，其等算力下的 Scaling 效率落后于标准 Transformer；本文证明引入稀疏 MoE 是解决该扩展瓶颈的核心钥匙。
  • vs MoEUT: MoEUT 尝试将 Universal Transformer 与 \(\sigma\)-MoE、SwitchHead 结合，但引入了过多并发改动；本文进行了干净受控的单变量对比，确立了 Top-k 稀疏路由在循环模型中的主导地位。
  • vs LayerSkip / Mixture of Depths: 现有自适应计算方案需改变训练目标或训练专用跳层路由；本文证明循环架构天然具备边界早退特性，无须任何修改即拥有更优的 Pareto 前沿。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从实验和机制上明确证明了稀疏路由对循环模型 Scaling 瓶颈的决定性修复作用。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖了基于 \(\mu\)P 的严格 IsoFLOP 缩放分析、OLMES 9 大常识评测、JSD 激活收敛追踪与早退 Pareto 扫描。
  • 写作质量: ⭐⭐⭐⭐⭐ 实验对照干净无冗余，动机明确，图表逻辑极其连贯。
  • 价值: ⭐⭐⭐⭐⭐ 为轻显存、长推理、计算自适应的下一代端侧与高效语言模型架构提供了极其清晰坚实的设计指南。

相关论文