← 返回 2026-02-03

SPARKLING:平衡信号保持与对称性打破的宽度渐进学习 SPARKLING: Balancing Signal Preservation and Symmetry Breaking for Width-Progressive Learning

Qifan Yu, Xinyu Ma, Zhijian Zhuo, Minrui Wang, Deyi Liu, Shiyi Zhan, Yiyuan Ma, Liang Xiang, Xingyan Bin, Di He 📅 2026-02-02 👍 48 2026-07-13 08:35
MoE 优化器状态管理 大语言模型 学习率调度 宽度扩展 渐进学习 预训练效率

通过RMS一致性保留与非对称梯度解耦实现中期宽度扩展,节省35%训练成本

前置知识

Progressive Learning(渐进学习)

渐进学习是一种资源高效的训练范式,在训练过程中从小模型逐步扩展到目标规模,而不是从一开始就训练完整规模的模型。其核心思想是先用较小的模型快速学习基础表示,然后通过深度(层堆叠/块插入)或宽度(维度扩展)增长逐步扩大参数规模。这种方法的理论依据是:小模型在训练早期已经能捕获有用的特征,扩展后可以利用这些已有知识加速收敛。现有方法已在深度扩展上取得显著成功,但宽度扩展仍是一个未被充分研究的方向。

SPARKLING是专注于宽度维度渐进学习的框架,理解渐进学习的整体范式和已有深度/宽度扩展方法的局限性,是理解本文动机和创新点的基础。

Mixture-of-Experts(MoE,混合专家模型)

MoE是一种稀疏激活架构,模型包含多个专家子网络(通常是MLP),每次输入仅激活其中的Top-k个专家。例如本文使用的OLMoE模型有64个专家,每次激活8个(Top-k=8)。这种设计使得总参数量可以很大(本文2.5B),但每次前向传播的活跃参数量较小(本文0.5B),从而在保持模型容量的同时控制计算成本。MoE模型中的宽度扩展有特殊含义:可以扩展专家的中间维度(inner dimension)或模型的隐藏维度(hidden dimension),两者都是论文研究的关键扩展轴。

本文的所有实验均基于MoE架构(OLMoE),论文提出了两种重要的宽度扩展轴:专家中间维度扩展和隐藏维度扩展。理解MoE的工作原理对于理解论文的实验设置和结论至关重要。

RMSNorm(均方根归一化)

RMSNorm是LayerNorm的简化变体,在现代LLM中广泛使用(如Qwen3、DeepSeek-V3、OLMoE)。与标准LayerNorm不同,RMSNorm省略了均值中心化步骤,直接用均方根(Root Mean Square)对隐藏状态进行缩放:对于输入向量 $\mathbf{x} \in \mathbb{R}^d$,输出为 $\mathbf{z} = \mathbf{x} \odot \gamma / \text{RMS}(\mathbf{x})$,其中 $\gamma$ 是可学习的缩放参数。RMSNorm计算效率更高,且在实践中表现与LayerNorm相当。在预归一化架构中,RMSNorm位于残差分支的输入端,控制着残差流和分支输出之间的信号平衡。

论文的核心技术贡献之一是提出了RMS尺度一致性原则,即在宽度扩展时保持激活的RMS统计量不变。理解RMSNorm的工作原理是理解这一关键思想的前提。

AdamW优化器

AdamW是目前大语言模型训练中最常用的优化器,结合了Adam的自适应学习率与权重衰减正则化。Adam维护每个参数的一阶矩估计(动量)和二阶矩估计(方差),通过 $m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t$ 和 $v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2$ 进行指数移动平均,然后用 $\theta_{t+1} = \theta_t - \eta \cdot \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon)$ 更新参数。AdamW在此基础上将权重衰减与梯度更新解耦。在本文中,AdamW的动量和方差状态(optimizer state)在宽度扩展时的处理方式是打破对称性的关键。

论文发现copy-based初始化会导致AdamW的动量状态也保持对称,从而使扩展参数持续耦合。理解AdamW的状态机制是理解对称性锁定问题和非对称状态重置方案的关键。

Muon优化器

Muon是一种光谱式(spectral-style)优化器,通过Newton-Schulz正交化对矩阵值动量进行去相关处理。与AdamW的逐元素更新不同,Muon对整个权重矩阵施加正交化变换,形式为 $\mathbf{X}_{k+1} = \mathbf{X}_k \phi(\mathbf{X}_k^\top \mathbf{X}_k)$,其中 $\phi(\cdot)$ 是Gram矩阵的多项式映射。Muon试图通过正交化来打破更新中的相关性,从而加速学习。本文通过理论分析证明,即使Muon的正交化步骤也无法打破copy-based初始化引入的对称性,因为复制维度的Gram矩阵具有特殊的块交换对称结构。

论文的重要贡献之一是证明了SPARKLING框架在不同优化器家族(AdamW和Muon)下均有效,而copy-induced对称性在两种优化器下都存在。理解Muon的工作原理有助于理解论文的理论分析深度。

函数保持(Function Preservation, FP)

函数保持是宽度扩展中的一种初始化策略,要求扩展后的模型在扩展点的瞬时输出与扩展前完全相同。最常见的实现方式是复制初始化(copy initialization):将原始权重矩阵的行或列直接复制到扩展后的矩阵中。这保证了 $f_{\text{expanded}}(\mathbf{x}) = f_{\text{original}}(\mathbf{x})$,即前向映射无缝连续。然而,本文揭示了一个反直觉的发现:尽管复制初始化在前向连续性上最强,但它由于引入了后向梯度对称性,反而会导致扩展后收敛更慢、最终损失更高。

函数保持是现有宽度扩展方法的核心思想,而本文的关键创新之一是超越函数保持的视角,转而关注激活统计量的保持(RMS保持)和梯度动态的对称性打破。

研究动机

大语言模型的训练成本极高,渐进学习(Progressive Learning)通过在训练过程中逐步扩大模型规模来降低计算开销。现有的渐进学习方法已在深度扩展(层堆叠、块插入等)方面取得了显著成功,但宽度扩展(增加隐藏维度、MoE专家中间维度等)这一同样重要的参数增长维度仍严重欠缺研究。仅有的少数宽度扩展方法基本局限于训练的早期阶段(不到10%-30%的训练token),此时进行扩展所能节省的计算量相对于从头训练目标宽度模型几乎可以忽略不计,从根本上违背了渐进学习降低训练成本的核心动机。在训练中期进行宽度扩展虽然能最大化计算节省,却面临严重的训练不稳定性:朴素初始化会破坏激活统计量导致损失突增,而基于复制的初始化虽然保持前向连续性,却会在反向传播中引入梯度对称性,使扩展的参数无法有效分化、形成功能冗余。这两个问题在以往的工作中被孤立地处理,缺乏一个系统性的框架来同时解决。

本文的目标是本文的目标是建立一个系统性的中期宽度扩展框架,能够在预训练的中间阶段(例如50%的训练进度处)安全地扩展模型宽度,同时实现两个看似矛盾的目标:第一,保持扩展时刻的前向信号稳定性,确保激活统计量不发生剧烈波动;第二,打破由复制初始化引入的反向梯度对称性,使扩展后的新增参数能够分化为有意义的独立特征。最终,该框架需要在固定token预算下匹配甚至超越从头训练目标宽度模型的下游性能,同时显著减少训练计算量。

与已有工作不同的是,本文的独特切入角度在于同时从信号保持(Signal Preservation)和对称性打破(Symmetry Breaking)两个互补原则来审视宽度扩展问题。与以往仅关注函数保持(Function Preservation)的方法不同,作者首次指出:函数保持只是表面现象,其背后的真正机制是RMS尺度保持——即保持中间激活的均方根统计量不变。同时,作者系统地揭示了copy-based初始化的对称性锁定(Symmetry Lock)问题:不仅AdamW的逐元素更新会维持对称性,连Muon的正交化更新也无法打破这一对称性,因为复制维度的Gram矩阵具有特殊的块交换对称结构。这种双重视角使得SPARKLING能够设计出同时满足前向稳定性和后向多样性的扩展机制,这在以往的工作中从未被系统地实现过。

核心方法

SPARKLING的核心思路是将中期宽度扩展分解为两个互补的技术挑战:信号保持和对称性打破。从直觉上看,扩展一个已经在训练中的模型就像在行驶中给汽车换轮胎——既要保持行驶的平稳性(信号保持),又要让新轮胎能够独立发挥作用(对称性打破)。技术路线分为四步:第一步,对新引入的参数使用基于复制的初始化,保持前向映射的函数连续性;第二步,对所有初始化策略(包括复制、随机、零初始化)应用RMS尺度一致性缩放,确保扩展后的激活统计量与扩展前保持一致;第三步,对新增参数的优化器状态进行非对称重置(保留原始参数的状态,重置新参数的状态),打破复制引入的梯度耦合;第四步,为新增参数引入非对称学习率重预热,提供额外的优化多样性。这四步机制协同工作,既保持了扩展时刻的前向稳定性,又确保了后续训练中新旧参数能够独立演化。

SPARKLING的核心创新在于同时解决了宽度扩展的两个根本矛盾。第一个创新点是RMS尺度一致性(RMS-scale Consistency):论文首次论证了现有函数保持方法的真正作用机制不是损失连续性,而是激活RMS统计量的保持。基于这一洞察,作者为不同的初始化策略(随机、零、复制)推导了精确的缩放因子 $\alpha$,使得扩展后每个隐藏坐标的方差满足 $\text{Var}(y'_i) = \text{Var}(y_i)$。第二个创新点是对称性锁定的识别与打破:作者证明了copy-based初始化不仅使原始参数和复制参数在前向传播中等价,更关键的是它们在反向传播中接收完全相同的梯度($\nabla_{\tilde{W}} L = \nabla_W L$),导致即使在后续训练中也无法分化。作者进一步证明,这种对称性对AdamW和Muon两种优化器都成立——Muon的Newton-Schulz正交化步骤由于Gram矩阵的块交换对称结构,同样无法自发打破对称。打破这一锁定的方法是非对称优化器状态重置:保留原始参数的动量/方差状态,将新参数的状态初始化为零。

方法步骤详情

SPARKLING方法的完整步骤如下。步骤1:确定扩展时刻和扩展轴。在预训练中期(如100B/200B tokens处)选择要扩展的宽度维度,论文研究了三种轴:MoE专家中间维度(Inner 2x,512到1024)、隐藏维度(Hidden 2x,1024到2048)、以及两者联合扩展。步骤2:参数初始化。对新引入的参数使用复制初始化(copy initialization),即将原始权重矩阵的行或列直接复制到扩展矩阵中。对于fan-out扩展(输出维度增长),$\tilde{W} = W$,扩展权重 $W' = [W^\top, W^\top]^\top$;对于fan-in扩展(输入维度增长),$W' = \alpha [W, \tilde{W}]$,输入 $x' = [x, \tilde{x}]$。步骤3:RMS尺度一致性缩放。根据初始化类型和扩展类型计算精确缩放因子。对于fan-in扩展中的两侧复制初始化,缩放因子为 $w'_{ij} = \frac{1}{\sqrt{1+3c}} w_{ij}$(当 $0 < c \leq 1$),其中 $c$ 是复制比例。对于RMSNorm参数 $\gamma$,同分布随机采样或复制自然保持RMS不变。对于零初始化侧,按随机初始化的缩放因子处理。步骤4:非对称优化器状态重置。保留原始参数 $W$ 的优化器状态 $S$(包括AdamW的动量和方差),将新参数 $\tilde{W}$ 的状态初始化为零:$S' = [S, \mathbf{0}]$。步骤5:非对称学习率重预热。原始参数保持原有的余弦学习率调度 $\eta(t)$;新参数从当前学习率 $\eta_e = \eta(t_e)$ 开始,在 $\tau_w$ 步内重预热到更高的峰值 $\hat{\eta}_{\max} = \rho \cdot \eta_e$($\rho = 1.3, \tau_w = 250$),然后跟随后续的余弦衰减。

技术新颖性

SPARKLING的技术新颖性体现在以下几个层面。理论层面:论文首次建立了函数保持到RMS保持的理论联系,论证了函数保持方法的真正作用机制是激活RMS统计量的保持,而非表面的损失连续性。论文还首次系统分析了copy-based初始化在不同优化器下的对称性问题:不仅证明了AdamW下梯度完全相同($\nabla_{\tilde{W}} L = \nabla_W L$),还证明了Muon的Newton-Schulz正交化步骤由于Gram矩阵的块交换对称结构($\mathbf{X}_k^\top \mathbf{X}_k$ 具有特殊块结构)同样无法打破对称。方法层面:论文为多种初始化组合(random-random、random-zero、random-copy、zero-copy、copy-copy)和扩展类型(fan-in、fan-out、RMSNorm)推导了统一的RMS保持缩放公式。非对称优化器状态重置方案简单而有效,避免了丢弃原始模型训练信号(Drop All)或保持对称(Copy All)的两个极端。非对称学习率重预热方案仅对新参数施加额外的优化激励,不影响原始参数的训练动态。实验层面:论文首次在MoE模型上系统验证了中期宽度扩展的可行性,覆盖了两种宽度轴和两种优化器家族,证明了SPARKLING的通用性。

RMS-preserving rescaling consistently improves late-stage convergence under MoE expert inner-dimension expansion
Figure 1: RMS-preserving rescaling consistently improves late-stage convergence under MoE expert inner-dimension expansion
Optimizer-state handling under copy-based expansion
Figure 2: Optimizer-state handling under copy-based expansion
Asymmetric re-warmup consistently improves convergence under mid-stage width expansion
Figure 3: Asymmetric re-warmup consistently improves convergence under mid-stage width expansion
Under zero initialization, RMS-preserved scaling enables the post-expansion activation RMS scale to quickly recover toward the original baseline
Figure 4: Under zero initialization, RMS-preserved scaling enables the post-expansion activation RMS scale to quickly recover toward the original baseline
Hidden-dimension expansion mirrors expert-inner growth
Figure 5: Hidden-dimension expansion mirrors expert-inner growth
A sample asymmetric learning rate re-warmup curve
Figure 6: A sample asymmetric learning rate re-warmup curve
Hyperparameter search for asymmetric re-warmup
Figure 7: Hyperparameter search for asymmetric re-warmup

实验结果

SPARKLING在OLMoE模型(0.5B活跃参数,2.5B总参数)上进行了全面验证。核心发现1:RMS尺度一致性缩放显著改善后期收敛。在所有测试的初始化组合中(random-random、random-zero、random-copy、zero-copy、copy-copy),RMS保持缩放虽然在扩展时刻可能产生略大的瞬时损失跳变,但后续训练中始终收敛到更低的最终损失,验证了RMS一致性是中期宽度扩展成功的关键。核心发现2:copy-based初始化存在对称性锁定。尽管复制初始化在前向连续性上最强,但其最终性能显著低于其他RMS保持策略。作者通过理论分析和实验验证了这是因为复制参数接收完全相同的梯度,形成对称性锁定,使扩展容量形同虚设。核心发现3:非对称优化器状态重置有效打破对称性。对比Drop All(全局重置)、Copy All(复制状态)、非对称重置(仅重置新参数状态)三种策略,非对称重置在恢复速度和最终损失上都显著优于对称基线,表明仅对新参数进行状态重置即可充分打破对称性。核心发现4:非对称学习率重预热进一步改善收敛。重新预热在所有宽度轴和初始化策略下都一致性地降低了最终损失,对copy-copy初始化的改善尤为显著——重预热使copy-copy达到了所有变体中的最低最终损失,有效弥补了对称性锁定的影响。核心发现5:下游任务匹配或超越从头训练。在12个下游基准上,SPARKLING的平均准确率在Inner 2x扩展中为60.30(从头训练59.64),Hidden 2x扩展中为61.13(从头训练60.46),联合扩展中为62.55(从头训练60.89)。核心发现6:显著节省计算量。SPARKLING节省20%-35%的训练FLOPs,最高实现1.49x的墙钟加速。

Downstream performance under 2x mid-stage width growth
Table 1: Downstream performance under 2x mid-stage width growth
Compute-cost comparison under a fixed token budget
Table 2: Compute-cost comparison under a fixed token budget
Baseline model Configuration
Table 3: Baseline model Configuration
Training hyperparameter configuration
Table 4: Training hyperparameter configuration
Step-law optimal learning rates across model sizes and their batch-size-scaled learning rates
Table 5: Step-law optimal learning rates across model sizes and their batch-size-scaled learning rates
Effectiveness under Muon optimizer
Figure 8: Effectiveness under Muon optimizer
Comparison to prior function-preserving symmetry-breaking heuristics
Figure 9: Comparison to prior function-preserving symmetry-breaking heuristics
查看结构化数据
任务指标本文基线提升
Inner 2x扩展(专家中间维度512到1024)下游平均 12个基准平均准确率 60.30 59.64(从头训练) SPARKLING超过从头训练基线1.1%,同时相比朴素扩展(58.96)提升2.3%
Hidden 2x扩展(隐藏维度1024到2048)下游平均 12个基准平均准确率 61.13 60.46(从头训练) SPARKLING超过从头训练基线1.1%,同时相比朴素扩展(59.25)提升3.2%
Hidden 2x与Inner 2x联合扩展下游平均 12个基准平均准确率 62.55 60.89(从头训练) SPARKLING超过从头训练基线2.7%,相比朴素扩展(60.89)提升2.7%
Arithmetic(Inner 2x设置) 准确率 60.50 55.67(从头训练) 提升8.7%,在所有基准中提升幅度最大
Arithmetic(Hidden 2x设置) 准确率 61.90 54.47(从头训练) 提升13.6%,提升幅度最大
Arithmetic(Hidden 2x与Inner 2x设置) 准确率 66.00 55.03(从头训练) 提升19.9%,提升幅度最大
Inner 2x扩展计算节省 FLOPs节省与加速比 7.21乘以10的20次方 FLOPs 9.01乘以10的20次方 FLOPs(从头训练) 节省20%训练FLOPs,1.27x墙钟加速
Hidden 2x扩展计算节省 FLOPs节省与加速比 8.10乘以10的20次方 FLOPs 10.80乘以10的20次方 FLOPs(从头训练) 节省25%训练FLOPs,1.29x墙钟加速
Hidden 2x与Inner 2x联合扩展计算节省 FLOPs节省与加速比 11.70乘以10的20次方 FLOPs 18.00乘以10的20次方 FLOPs(从头训练) 节省35%训练FLOPs,1.49x墙钟加速

局限与改进

论文承认的第一个局限是:尽管SPARKLING在下游任务上匹配或超越了从头训练的基线,但最终预训练损失仍略高于从头训练的目标宽度模型。例如在Inner 2x设置中,SPARKLING的损失为2.3153,而从头训练为2.3096。这表明宽度扩展本身不可避免地带来一定的性能代价,虽然这种代价在经过下游微调后可以被弥补甚至逆转。第二个局限是实验规模相对有限:所有实验均基于OLMoE(0.5B活跃/2.5B总参数)模型,在更大规模模型上的效果有待验证。第三个局限是目前仅研究了2x的宽度扩展倍率,对更大倍率(如3x或4x)的适用性尚不清楚。第四个局限是实验仅覆盖了MoE架构,未验证SPARKLING在Dense模型上的效果。从独立观察的角度来看,论文还缺少对扩展时机的系统研究——目前固定在50%训练进度处扩展,不同扩展时机(25%、33%、50%、66%)的效果差异未被探讨。此外,RMS保持缩放的推导依赖于高维各向同性假设(即宽网络中特征向量趋于渐近正交),这一假设在窄网络或训练早期是否成立需要更多验证。

独立分析的弱点

论文存在几个值得深入分析的弱点。弱点1:预训练损失仍有差距。SPARKLING的最终预训练损失始终略高于从头训练的目标宽度模型(差距约0.006到0.019),虽然下游任务表现可弥补,但这一差距的原因尚不完全清楚。改进方向:可以探索更精细的RMS保持策略,例如不仅保持一阶统计量(均值、方差),还保持二阶统计量(协方差结构),或者引入更渐进的宽度增长策略而非一次性扩展。弱点2:缺乏对Dense模型的验证。所有实验均基于MoE模型,Dense模型(如标准Transformer)没有MoE的稀疏激活特性,扩展策略可能需要调整。改进方向:在标准Dense模型上验证SPARKLING,特别是研究Dense MLP中fan-in/fan-out扩展的RMS保持公式的适用性。弱点3:扩展倍率和时机的敏感性分析不足。论文仅研究了2x扩展倍率和50%扩展时机,对更大倍率和不同时机的效果缺乏系统分析。改进方向:进行网格搜索,研究扩展倍率、扩展时机与最终性能之间的关系,为实践者提供更具体的指导。弱点4:非对称学习率重预热引入额外超参数。重预热比率和重预热步数虽然论文给出了推荐值(比率1.3、步数250),但最优值可能因模型规模、数据集、扩展轴而异。改进方向:研究自适应的重预热策略,例如根据扩展后损失的变化动态调整重预热强度。

未来方向

论文作者明确提出了三个未来研究方向。方向1:宽度和深度同时扩展的统一原则。目前宽度扩展和深度扩展是独立研究的,如何在同一次渐进学习过程中同时进行两种扩展,以及两种扩展的最优顺序和组合方式,是一个重要的开放问题。方向2:RMS保持与Maximal Update Parameterization的兼容性。MuP条件确保最优超参数在不同模型规模间可迁移。如果SPARKLING的RMS保持策略能满足MuP条件,那么扩展后的模型可以直接使用扩展前调好的超参数,实现免调优扩展。方向3:更全面的免调优渐进学习框架。结合超参数迁移和自动扩展策略,减少渐进学习中的人工干预。基于论文成果,还可以延伸以下方向:(1)探索SPARKLING在更大规模模型(如7B、13B、70B参数量级)上的效果,验证方法的可扩展性;(2)研究SPARKLING在多阶段渐进学习(连续多次宽度扩展)中的累积效应;(3)将RMS保持思想推广到深度扩展中,分析层插入时的激活统计量保持问题;(4)探索SPARKLING与其他效率技术(如知识蒸馏、数据课程)的组合效果。

复现评估

论文的复现可行性总体较好。开源情况:论文基于OLMoE架构,该架构的代码和预训练权重已开源。但论文本身的代码和训练配置未在arxiv页面明确公布。数据:论文使用的训练数据为OLMoE的原始预训练数据集(200B tokens),训练细节在附录Table 3和Table 4中详细列出,包括模型配置(24层、1024隐藏维度、64个MoE专家、Top-k=8、RoPE位置编码等)和训练超参数(AdamW、峰值学习率1.9556乘以10的负3次方、余弦调度等)。算力:实验使用64张NVIDIA A100 80GB GPU,对于200B tokens的训练来说算力需求较高。从头训练的对照实验需要84到209 GPU小时,SPARKLING需要66到140 GPU小时。难度:方法本身实现相对直接——RMS缩放是简单的数值缩放,非对称状态重置只需几行代码修改,非对称学习率重预热需要修改训练循环中的学习率调度器。最复杂的是RMS缩放因子的计算,特别是both-sides copied情况下的公式推导。总体而言,论文提供的数学推导和实验细节足够复现核心结果,但完整复现需要较大算力投入。