Sparse Layers are Critical to Scaling Looped Language Models

TL;DR

本研究提出稀疏专家层在循环变换器中的关键作用,结合早期退出实现更优扩展和节省资源。

cs.LG 🔴 高级 2026-05-10 42 次浏览
Ryan Lee Jacob Biloki Edward J. Hu Jonathan May
深度学习 变换器 稀疏专家 模型扩展 早期退出

核心发现

方法论

通过对比标准变换器、循环变换器及其稀疏专家(MoE)变体,系统分析了参数复用、路由多样性及早期退出机制对模型扩展性的影响。采用μP参数化确保不同规模模型的训练一致性,利用IsoFLOP曲线评估不同架构在固定计算预算下的性能,结合专家路由差异分析和输出收敛性验证,揭示稀疏专家在循环模型中的表达能力恢复作用。

关键结果

  • 循环MoE模型在相同参数量下,测试损失优于标准模型,达到0.077的损失指数,显著优于非稀疏循环模型的0.085。路由差异分析显示,超过50%的tokens在不同循环中被分配到不同专家,有效实现层级特征多样性。
  • 在早期退出实验中,循环模型在节省10%的计算量时,保持perplexity在50以下,优于非循环模型的55,验证了循环结构与输出收敛的关系。多循环配置进一步提升早退出的性能,达成更优的计算-质量折衷。
  • 在AI2 OLMES基准测试中,循环MoE模型在参数较少的情况下,整体得分高于标准模型,平均提升1.0分,显示其在实际任务中的优越性。

研究意义

本研究突破了循环变换器在参数复用中的表达瓶颈,通过引入稀疏专家机制,有效提升模型扩展性和推理效率,解决了大规模模型存储与计算成本高的问题。结合早期退出策略,为实际部署提供了高效的推理方案,推动大规模语言模型的实用化进程。

技术贡献

提出路由差异驱动的稀疏专家在循环模型中的应用,系统分析了专家路由的多样性对模型表达能力的影响。引入μP参数化确保不同规模模型训练的稳定性,结合输出收敛性验证,揭示循环模型在边界处的输出优化机制。创新性地结合稀疏专家与早期退出,优化模型存储与推理效率,为未来模型设计提供新思路。

新颖性

首次系统性验证稀疏专家在循环变换器中的表达能力恢复作用,明确指出专家路由差异是提升扩展性关键。提出利用循环边界的输出收敛特性实现训练无关的早期退出,突破传统深层模型对中间层的依赖,为模型压缩与推理优化提供新途径。

局限性

  • μP参数化在深度变化时的迁移效果有限,未来需结合深度扩展的参数调优策略。
  • 实验主要集中在参数规模在数亿级别,尚未验证在更大规模(如百亿参数)模型中的表现。
  • 早期退出策略依赖输出分布的熵阈值,可能在不同任务或数据分布下表现不一致,需进一步泛化验证。

未来方向

未来将探索深度变化对μP参数化的适应性,结合更大规模预训练验证模型的扩展性。研究多循环层数与专家路由多样性对模型性能的影响,优化专家路由机制,提升模型在多任务环境中的泛化能力。同时,结合硬件优化,提升推理速度和能效,为大规模语言模型的实际部署提供更全面的解决方案。

AI 总览摘要

随着大规模语言模型的不断发展,模型存储成本与推理速度成为瓶颈。传统变换器架构采用逐层不同参数,虽效果优异,但难以在资源有限条件下扩展。循环变换器通过参数重用,降低存储需求,但其表达能力受限,难以匹配非循环模型的性能。本文提出在循环变换器中引入稀疏专家(MoE)层,利用专家路由差异实现每次循环的不同特征提取,从而恢复模型的表达能力。实验显示,稀疏专家循环模型在参数相同时,测试损失优于标准模型,且在早期退出机制中表现出更优的计算-质量折衷。通过分析专家路由的差异性,验证了不同专家在不同循环中的激活多样性,有效增强模型的表达丰富度。输出收敛性分析表明,循环边界的输出更接近最终分布,使得在这些点提前退出成为可能,极大节省推理成本。结合μP参数化,确保不同模型规模训练的稳定性,为未来大规模模型的扩展提供了理论基础。整体而言,稀疏专家与循环结构的结合,为大规模语言模型的存储、推理效率提供了新路径,推动其在实际应用中的落地。

深度分析

研究背景

近年来,深度学习中的变换器架构成为自然语言处理的主流,尤其是GPT、BERT等模型推动了预训练技术的发展。传统变换器通过堆叠多层参数密集的变换层实现强大表达能力,但其存储和推理成本极高。为缓解这一问题,循环变换器(如Universal Transformer)引入参数重用,减少存储,但表达能力受限。近年来,稀疏专家(MoE)技术被引入,利用路由机制激活不同子网络,显著提升模型能力。此前研究多集中在单一技术改进,缺乏系统性分析稀疏专家在循环结构中的作用。本研究结合两者,探索稀疏专家在循环变换器中的潜力,旨在突破参数重用的表达瓶颈,提升模型扩展性与推理效率。

核心问题

循环变换器虽在参数效率上优越,但其表达能力受限,尤其是在参数共享导致的表达瓶颈上。传统模型在参数复用时,所有循环共享相同参数,限制了模型的多样性和复杂性。如何在保持参数节省的同时,增强模型的表达能力,成为关键难题。此外,深层模型推理成本高,早期退出机制虽有潜力,但在循环结构中如何有效结合,尚缺乏系统性研究。解决这些问题,将极大推动大规模模型的实用化。

核心创新

本研究的核心创新包括:1)引入稀疏专家(MoE)层,利用不同专家的激活差异,实现每次循环的特征多样性,突破参数共享的限制;2)分析专家路由的差异性,验证其在模型扩展中的作用;3)结合输出收敛性,利用循环边界的特性实现训练无关的早期退出,显著降低推理成本。这些创新区别于以往仅在单一技术层面优化的工作,为模型扩展提供了全新思路。

方法详解

  • �� 构建基础变换器模型,采用多头自注意力(MHSA)和SwiGLU前馈网络。
  • �� 引入循环机制,将8层参数重复2次,形成16层等效深度模型。
  • �� 替换部分全连接层为Top-k稀疏专家(MoE),利用路由机制(h(x)=Wrouterx)激活不同专家。
  • �� 采用μP参数化,确保不同规模模型的训练稳定性,调节初始化和学习率。
  • �� 通过专家路由差异分析,统计不同循环中的专家激活差异,验证多样性。
  • �� 利用输出分布的Jensen-Shannon散度(JSD)分析模型在边界的输出收敛性。
  • �� 设计早期退出策略,在循环边界根据熵阈值提前退出,节省计算。
  • �� 在不同参数规模(如5亿到10亿)下训练模型,使用IsoFLOP曲线评估性能,验证扩展性。

实验设计

采用10B token的FineWeb-Edu数据集进行预训练,比较标准变换器、循环变换器及其MoE变体在参数相同条件下的性能。利用不同模型规模(128到1024维)进行训练,测量测试损失、输出收敛性和早退出效果。结合AI2 OLMES基准测试,评估模型在实际任务中的表现。通过专家路由差异分析,验证多样性贡献。采用熵阈值进行早期退出,比较不同循环次数和层数的效果。所有实验均在固定计算预算下进行,确保公平性。

结果分析

稀疏专家循环模型在参数相同时,测试损失降低至0.077,优于非稀疏模型的0.085。专家路由差异分析显示,超过50%的tokens在不同循环中被分配到不同专家,增强表达能力。早期退出实验表明,节省10%的计算量时,perplexity仍低于50,优于非循环模型的55。多循环配置进一步提升早退出性能,模型在保持较低perplexity的同时,显著减少推理成本。在实际任务中,循环MoE模型在AI2基准测试中得分高于标准模型,平均提升1分,验证其实际应用潜力。

应用场景

该架构适用于大规模预训练语言模型,特别是在存储和推理资源有限的场景。可部署于边缘设备或低功耗环境,提升推理速度和效率。结合早期退出机制,可实现动态计算调节,满足不同应用的性能需求。未来,结合硬件优化,有望推动大模型在实际产品中的落地,改善用户体验。

局限与展望

μP参数化在深度变化时效果有限,需结合深度扩展策略。实验主要在亿级参数规模,尚未验证在百亿参数级别的性能表现。早期退出策略依赖熵阈值,可能在不同任务或数据分布下表现不一。模型训练和推理仍需大量计算资源,未来需优化算法和硬件支持以实现更广泛应用。

通俗解读 非专业人士也能看懂

想象你在一家工厂工作,工厂里有很多工人(模型中的层),每个工人都能做不同的任务。传统工厂里,每个工人都一样,重复做同样的事情,效率高但缺乏变化。现在,为了让工厂更聪明,工人们被分成不同的小组(专家),每组专门做某一类任务。工厂还可以让工人们轮流工作(循环),每次轮到他们时,他们可以用不同的小组(专家)来完成任务,从而做出更复杂、更聪明的产品。而且,工厂在某些阶段可以提前结束工作(早期退出),不用等全部完成就可以交付。这种设计让工厂既省钱又快,还能做出更好的产品。

原文摘要

Looped language models repeat a set of transformer layers through depth, reducing memory costs and providing natural early-exit points at loop boundaries. However, looped models do not scale as favorably as standard transformers with unique layers. We compare standard and Mixture-of-Experts (MoE) transformers, with and without looping, and find two main results. First, we find Looped-MoE models scale better than the standard baseline while dense looped models do not. We trace this to routing divergence between loops: in Looped-MoE models, different experts are activated on each pass through the same shared layers, recovering expressivity without additional parameters. Our second finding is that looped models have better compute-quality trade-offs with early exits than standard models. Because each loop ends with the same layers that produce the final output, loop boundaries are superior exit points, as confirmed by earlier output convergence at these points. In sum, we provide a clear direction for scaling looped models: a Looped-MoE model with early exits can not only beat standard transformers at scale, but also enable significant memory and inference savings with minimal degradation in quality.

cs.LG cs.CL