![]()
这项由独立研究者完成的研究以预印本形式发布于2026年7月21日,论文编号为arXiv:2607.19058,有兴趣深入了解的读者可以通过该编号在arXiv平台查询完整论文。
**一个让AI研究者头疼已久的内存账单**
训练大型AI模型,就像经营一家大型餐厅。你不仅要备好食材(模型参数),还要备好厨师的工作记录本(优化器状态)——记录每道菜上次怎么做、下次应该怎么调整。问题是,这本"记录本"比食材本身还要厚得多。
在这项研究中,研究者聚焦的是一类叫做"稀疏混合专家模型"(MoE)的AI架构。这种架构的设计哲学非常聪明:模型里有许多"专家",处理每条信息时只调用其中一小部分,而不是所有专家都同时上阵。这样就做到了参数总量很大,但每次实际运算量很小。研究者测试的具体模型有67.84亿个参数,但每次处理一个词语时,实际激活的参数只有约4.4亿,节省了大量计算资源。
然而,训练时用来"教"模型学习的优化器——具体来说是业界最常用的AdamW优化器——并不享受这种折扣。AdamW需要给每个参数保存两份"学习记录"(分别叫一阶动量和二阶动量),而且是用高精度浮点数存储的。结果就是:67.84亿个参数对应的权重数据只占12.6GB,但AdamW光是这两份记录就要占50.6GB,整个训练峰值内存高达81.4GB——远超市面上主流的40GB专业显卡。
这个问题在学术界和工业界都是真实的痛点。很多研究团队为了训练大模型,要么租用昂贵的超大内存服务器,要么把模型拆分到多张显卡上联合训练,成本和复杂度都大幅上升。于是,这位独立研究者开始问一个没有人认真追究过的问题:**真的有必要对所有参数都保存同样多的学习记录吗?**
**二、混合专家模型里的三类"员工"**
为了理解这项研究的核心洞见,我们先要搞清楚这类AI模型的内部结构。可以把这个67.84亿参数的模型想象成一家大型公司,里面有三类截然不同的员工群体。
第一类是"骨干员工",也就是论文中所说的密集骨干网络(dense backbone),包括词嵌入层、注意力机制、密集前馈层等。这批人数量不多,只占总参数量的5%(约3.41亿),但他们是每处理一个词语都要上班的,工作量稳定、密集,工作记录的意义非常大。
第二类是人数最多的"专家团队",也就是模型里的128个专家网络。这批人占了总参数量的95%(约64.43亿),是绝对的大头。但他们的工作方式很特殊:每次处理一个词语,只有其中约2个专家被调用(top-2路由机制),也就是说每个专家平均只处理1/64的词语。他们的工作是稀疏的、零散的,每次上班都是"偶发任务"。
第三类是人数极少但地位关键的"调度员",也就是路由器(router)。这个小小的门控网络只有约52万个参数,占总量不到0.01%,但它的职责是决定每个词语该发给哪些专家处理——它的每一个决策都影响着整个公司的运转效率,尤其是各专家之间的工作量是否均衡。
以往所有的优化器,不管是AdamW、Lion还是Muon,都把这三类"员工"一视同仁,给每个人都保存同样格式、同样精度的工作记录。研究者认为这是一种巨大的浪费,而且可能并不合理。
**三、SkewAdam:一套"按需分配"的记账策略**
针对上述问题,研究者提出了名为SkewAdam的优化器。这个名字本身就揭示了设计思路:Skew的意思是"倾斜、偏向",也就是把记账资源向真正需要的地方倾斜。
对于骨干员工(5%的参数),SkewAdam保留了完整的float32精度动量缓冲区,外加一个"分解式二阶动量"。动量这个概念,可以用物理上的惯性来理解:一个滚动的球会保持方向继续滚,不会因为某一步地面稍微不平就立刻转向。骨干员工每步都有密集的梯度信号,动量能让学习方向更加稳定,1.27GB的存储代价在这里是值得的。
对于专家团队(95%的参数),SkewAdam完全不保存动量,只保存一个分解式二阶动量。所谓"分解式二阶动量",是借鉴自Adafactor优化器的技巧:对于一个n×m的权重矩阵,不直接存储n乘以m个数,而是只存储n个行统计量和m个列统计量,然后用它们的乘积来近似重建完整矩阵。对于一个4096×4096的专家矩阵,这相当于从存储1600万个数减少到只存储8192个数,节省了近2000倍空间。不保存动量的原因是:每个专家平均每步只接收约128个词语的梯度,信号本来就稀疏嘈杂,硬要用动量去平滑这样的噪声,效果有限,代价却高达24GB。
对于调度员路由器(不到0.01%的参数),SkewAdam反而给了最精确的待遇:完整的、不分解的float32二阶动量,并且不施加权重衰减。原因是路由器的工作性质特殊——它的128个输出数值之间的相对大小,直接决定了哪些专家被激活。如果用分解的方式估算,就会把不同输出位置的统计量混在一起,破坏这种精细的相对关系。2MB的精确存储,换来的是整个调度系统的稳定性,非常划算。
这套策略的整体结果是:总优化器状态从AdamW的50.55GB降低到1.29GB,只有后者的2.6%。训练峰值内存从81.4GB降到31.3GB,终于可以在一张40GB的显卡上运行。
**四、具体的数学机制:记账本是怎么写的**
为了让对技术细节有兴趣的读者有更清晰的认识,这里用比较直观的方式解释SkewAdam内部的运作逻辑。
分解式二阶动量的更新方式是这样的:每一步,先分别计算梯度矩阵在行方向和列方向的均方值,然后用指数移动平均的方式逐步积累历史信息。重建时,把行向量和列向量相乘,再除以行向量的均值,得到一个近似的完整二阶动量矩阵。这个近似在梯度矩阵的真实二阶动量接近"秩为1"(也就是行列之间高度相关)时是精确的,在实际训练中通常是足够好的近似。
更新步骤还包括一个"更新幅度裁剪"机制:计算出预调整方向后,如果它的均方根超过1,就把整个向量等比例缩小到均方根恰好等于1。这就好比规定每步最多走一定的距离,防止偶发的大梯度导致参数跳飞。
主权重是用bfloat16格式存储的(一种低精度浮点数,每个数只占2字节,而float32占4字节)。每一步计算完更新量后,在将其写回bfloat16权重时,会加入一个微小的随机噪声(范围是当前权重数值精度的一半),这个技巧叫做"抖动随机舍入",目的是避免系统性的舍入误差积累。值得一提的是,这项研究中所有对比的优化器都使用了完全相同的这套写回机制,确保比较的公平性。
**五、实验设计:一场严格的"同台竞技"**
为了让结果有说服力,研究者设计了一套极为严格的对照实验。
模型固定为同一个67.84亿参数的解码器结构,两个模块构成:第一个模块用密集SwiGLU前馈层,第二个用128专家的MoE层,宽度均为4096,使用分组查询注意力机制(32个查询头,8个KV头)。训练数据来自OpenWebText语料库,按文档哈希值划分为95%训练集和5%验证集,确保两侧没有重叠文档。每次运行固定10000步,批量大小为64条序列×128个词语,总计约8190万个词语,单轮不重复。
最关键的控制是:所有参与对比的优化器都从同一个随机初始化出发,按照完全相同的数据顺序进行训练,使用相同的bfloat16主权重和相同的写回路径。这就像让四位厨师从同一份食材开始,按同一份菜单顺序烹饪,只是各自使用不同的烹饪手法——这样最终菜品的差异就只能来自手法本身。
参与对比的优化器除SkewAdam外,还有AdamW(使用3×10??学习率)、Lion(使用1×10??学习率,按其论文建议降低3到10倍)、以及Muon(矩阵参数学习率0.02,其余参数用内置Adam,学习率10??)。训练在一块英伟达H200显卡(141GB显存)上进行,之所以选择这么大内存的卡,是因为要包含AdamW这个对照组——如果换成40GB卡,AdamW根本跑不起来。
**六、主要结果:内存大降的同时,收敛更快**
实验结果在内存和效果两个维度上都给出了清晰的答案。
内存方面,SkewAdam的峰值训练内存为31.3GB,是唯一能在40GB显卡上运行的方案。AdamW需要81.4GB,Lion和Muon虽然各自将状态降到25.27GB,但由于模型权重、梯度、激活值等其他开销,总峰值仍分别达到56.6GB和57.6GB,均超过40GB的门槛。
吞吐量方面,SkewAdam每秒处理约5000个词语,比AdamW快6.6%,仅比Lion慢1.5%。Muon因为每步都要对64.43亿专家参数运行牛顿-舒尔茨正交化迭代,吞吐量只有3409词语/秒,比SkewAdam慢了约32%。
收敛质量方面,训练前3000步,AdamW和Muon的验证困惑度(衡量语言模型好坏的指标,数字越低越好)领先于SkewAdam——第1000步时,SkewAdam比AdamW落后114个困惑度点。然而从第4000步开始,SkewAdam反超,最终以108.4的验证困惑度结束训练,领先Muon的120.2和AdamW的126.8。Lion的表现非常糟糕,最终停在393.7,训练过程中困惑度甚至在第9000步之后开始反弹恶化,始终未能有效学习。
路由均衡方面,理想状态下,128个专家应该被大致均等地使用,此时负载均衡损失函数的理论最低值是0.05。SkewAdam和AdamW从第4000步起就稳定在0.0505和0.0502,几乎贴着理论下限。Lion虽然路由相对均衡(0.0537),但语言模型完全没学好。Muon则在最后1000步突然跳升到0.0608,比理论下限高出22%,与此同时其验证困惑度也出现了唯一一段平台期,两者在时间上高度吻合。
**七、与Adafactor和GaLore的额外比较**
研究者后续在英伟达H100 NVL 47GB MIG切片上补充了两项对照实验:均匀Adafactor和GaLore风格基线。
Adafactor是SkewAdam最直接的"亲戚"——SkewAdam使用的分解式二阶动量估计器就是从Adafactor那里借鉴来的。Adafactor完全不保存动量,对所有参数统一使用分解式二阶动量,优化器状态只需12MB,比SkewAdam的1.29GB还少得多。然而,均匀Adafactor的最终验证困惑度是149.5,比SkewAdam的109.0足足高出40个点,而且在训练最后1000步里改善幅度不到0.1,已经近乎停滞。
这个对比非常有说服力,因为两个优化器共享完全相同的分解式二阶动量机制和更新裁剪策略,唯一的区别是SkewAdam给骨干参数保存了动量,而Adafactor没有。由此可以清晰地定位:SkewAdam之所以效果更好,核心原因是动量,而不是分层分配策略本身——分层策略带来的是内存节省,不是额外的精度提升。
GaLore风格基线(秩128的低秩梯度投影)在这个设置下彻底失败,最终困惑度为1839.9,路由虽然均衡,但语言模型完全没有学到有效表示。研究者明确指出,这不能作为对GaLore方法本身的否定——他们使用的是自己的实现版本,未经系统调参,与GaLore原始论文的设置有所不同。但这个失败案例至少提示,把稀疏专家梯度投影到低秩子空间这件事,需要非常谨慎。
**八、哪个"分层"决策真正起作用?消融实验的答案**
一个合理的质疑是:SkewAdam的效果是因为分层策略,还是因为其中某个特定的设计选择?研究者在AMD MI300X(192GB)上运行了消融实验,系统地把各个决策一一关闭,观察结果变化。
实验包含四个变体:完整的SkewAdam策略,在专家上恢复动量的版本,把路由器二阶动量改为分解式的版本,以及完全均匀分配(所有参数都用动量加分解式二阶动量)的版本。
结果是:四个变体的验证困惑度全部落在108.2到108.9之间,互相之间的差异完全在单次运行的随机噪声范围内。路由均衡损失也全部在0.050附近,几乎没有区别。真正发生显著变化的只有优化器状态大小:完整SkewAdam是1.29GB,恢复专家动量后变为25.29GB,完全均匀版本同样是25.29GB。
这个结果说明了两件事。其一,分层策略的贡献是内存,而不是额外的困惑度提升——完整策略和均匀策略达到完全相同的训练效果,前者只用了后者二十分之一的优化器状态。其二,专家层的动量是纯粹的浪费:加上它花费24GB,但困惑度只变动了0.2,完全在噪声范围内。
路由器使用精确二阶动量的选择(而非分解式)也被证明既无害也无益——改为分解式后,困惑度和路由均衡均无变化。研究者将其描述为"无害但非关键的选择",保留它只是出于对路由器特殊角色的谨慎考量。
**九、调参之后,基线能追上吗?**
为了排除"SkewAdam用了更合适的学习率"这一可能的解释,研究者系统地对AdamW和Adafactor进行了学习率扫描,而SkewAdam保持不调参的默认设置3×10??。
AdamW扫描了10??、3×10??、10??三个学习率,每个跑三个随机种子。最优结果在10??时出现,达到118.5±0.5,比默认的126.8改善了约8个点。Adafactor扫描了3×10??到3×10??共五个学习率,最优结果同样在10??,达到139.7(两个种子平均,差异0.005)。值得注意的是,3×10??时Adafactor出现了欠训练(困惑度257.5),而更高的学习率全部更差,说明10??的选择已经是有效区间内的最优点,上下都有了明确的边界。
调参之后,差距缩小了,但没有消弭:未经调参的SkewAdam(108.4到109.0之间)领先最优AdamW约10个困惑度点,领先最优Adafactor约30个点。10个点的差距约是AdamW种子间标准差的20倍,统计上非常显著。
研究者还指出,调过参的AdamW和调过参的Adafactor之间仍然差了约21个点(118.5对139.7),而这两者的唯一区别正是动量——AdamW有动量,Adafactor没有。这个跨优化器的对比,与消融实验的结论完全吻合:动量是关键,分层分配是内存的解法。
**十、零样本评估:诚实的"不知道"**
研究者还用标准的语言模型评估工具对训练结束的四个模型进行了零样本推理测试,任务包括PIQA(物理常识推理)、WinoGrande(代词消歧)、HellaSwag(故事续写)和ARC-Challenge(科学问答)。
结果很诚实:在仅8190万词语的训练量之后,所有四个模型在这些任务上的表现几乎全部接近随机猜测水平。PIQA上四者得分在53.5%到55.9%之间(随机猜测是50%),HellaSwag和ARC-Challenge接近随机猜测的25%。任何两个优化器之间的差异都在一到两倍标准误差之内,没有统计显著性。
8190万词语,对于现代大语言模型来说只是零头——正式训练通常需要数千亿甚至数万亿词语。研究者坦然地将这些零样本结果定性为"完整性检查"而非能力证明,并明确提醒读者:困惑度上的提升不等于下游任务的提升,不要对这份结果赋予超出它本身的意义。
**十一、局限与未竟之处**
研究者在论文中对这项工作的局限性进行了相当坦诚的讨论。
模型只有两个模块,这是出于成本考量的主动选择——它把95%的参数集中在一个专家库里,制造了一个极端的压力测试场景。但真实的MoE模型通常有几十乃至上百个交替的注意力层和MoE层,分层策略在多层叠加的复杂路由结构下是否同样有效,尚未得到验证。
大多数配置只跑了一个种子,消融实验的四个变体在0.6个困惑度点的范围内波动,而这恰好是单次运行间随机噪声的量级,因此消融结果应该被理解为"内存代价相当,效果相当",而不是对各变体之间存在任何具体顺序的声明。
训练只有8190万词语,上下文长度只有128个词语,按现代标准极为短小。权重衰减在bfloat16精度下是一个无效操作(变化量远小于浮点精度,被直接舍入为零),这意味着整个实验实际上是在没有权重衰减的条件下进行的,与真实的长时间生产训练有所不同。如果要将这套方案应用于生产规模,必须把权重衰减项整合进float32更新步骤,在精度截断之前完成计算。
GaLore风格基线使用的是研究者自己写的实现,未经系统调参,结果不代表GaLore方法的真实上限。Lion和Muon也只测试了单个学习率,没有进行系统调参,其表现可能还有提升空间。
**归根结底,这项研究告诉了我们什么**
说到底,这项工作的核心贡献是一个朴素却被忽视的观察:混合专家模型里的三类参数,工作方式截然不同,理应得到不同的记账待遇。把骨干网络的动量、专家层的分解式方差估计、和路由器的精确方差估计组合在一起,优化器状态从50.55GB降到1.29GB,训练峰值内存从81.4GB降到31.3GB,而最终模型效果不但没有变差,反而因为保留了骨干层的动量而优于AdamW。
这对于想在资源有限的条件下训练大模型的研究者和工程师来说,是一个实际可用的工程方案。它不依赖量化、不依赖梯度压缩、不依赖多卡分布式,只是重新思考了"每一分内存都值得花在哪里"这个问题。研究者自己也强调,这不是一个更好的优化算法,而是一个更好的内存分配策略——Adam家族的效果,花了2.6%的Adam状态存储。
当然,在更深的模型、更长的训练、更多的随机种子上验证这套方案,是这项工作留给未来的任务。但作为一个方向,它提出了一个颇具启发性的设计原则:在构建优化器时,先问清楚"这部分参数是谁,它们是怎么工作的",再决定给它们记多少账。有兴趣深入了解的读者,可以通过arXiv:2607.19058查阅完整论文,代码也已公开在GitHub上供复现使用。
Q&A
Q1:SkewAdam和AdamW的训练效果哪个更好?
A:在这项研究的实验设置下,SkewAdam的最终验证困惑度为108.4,优于AdamW的126.8。即使对AdamW进行学习率调参,其最优结果也只能达到118.5,仍然落后于未调参的SkewAdam约10个点。但需要注意,实验规模较小(约8190万词语),且模型结构特殊,结果能否推广到更大规模仍需验证。
Q2:SkewAdam为什么不给专家层保存动量?
A:因为每个专家平均每步只处理约128个词语的梯度,信号稀疏且噪声高,动量平滑的效果有限,却要消耗24GB内存。消融实验证实,给专家层加回动量后,验证困惑度变化仅0.2个点,完全在随机噪声范围内,属于纯粹的浪费。
Q3:分解式二阶动量是怎么节省内存的?
A:对于一个n×m的权重矩阵,传统方法要存储n×m个数,分解式方法只需存储n个行统计量和m个列统计量,通过它们的乘积来近似完整矩阵。以4096×4096的专家矩阵为例,存储从1600万个数减少到约8192个数,节省约2000倍,整个专家库的二阶动量只需12MB。
特别声明:以上内容(如有图片或视频亦包括在内)为自媒体平台“网易号”用户上传并发布,本平台仅提供信息存储服务。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.