Token导航 LogoToken导航

AI 资讯 · 科技行者

ShellElastic一次训练产出五种体积模型

更新 2026-11-12科技行者正文 7633字阅读约 24分钟

你有没有想过,一家做手机应用的公司,为什么总是要同时准备好几个版本的AI模型。

旗舰手机用大模型,跑得动也效果好。入门手机用小模型,省电省内存。平板用中等模型,刚好够用。

问题是,这些模型通常是分开训练出来的。

训练一个十几亿参数的大模型要花多少钱,业内人都清楚。现在要五个不同大小的模型,难道要花五倍的钱,跑五次完整的训练流程?

这篇来自Rutgers大学和Adobe Research的论文,就是在解决这个让人头疼的问题。

弹性模型,一个老想法遇到的新麻烦

把多个不同大小的模型塞进一套共享参数里训练,这个思路其实不新鲜。

弹性模型*(Elastic Models):在一套共享参数里同时训练多个不同规模的子网络,训练一次就能得到多种推理成本的模型版本,不用分别训练。

早在2018年,slimmable networks和Once-for-All就在视觉模型上验证了这个想法:一个卷积网络可以同时适配多种宽度。后来MatFormer把这个思路搬到了Transformer架构上,专门调整FFN(前馈网络)层的宽度,谷歌的Gemma 3n模型就用了这套方案。

听起来很美好,对吧,一次训练,多种规格。但论文作者指出,这个美好想法在真正用于大规模预训练时,撞上了三个实实在在的坑。

第一个坑是训练负荷不稳定。

大多数弹性训练方案的做法是,每一步随机抽一个或几个子网络来训练,而不是每次都训练全部。这就意味着,这一步训练的是小模型,计算图很简单;下一步训练的是大模型,计算图又变了。

这种忽大忽小的工作量,对于依赖固定计算模式进行调度和内存管理的大规模分布式训练系统来说,是个麻烦事。

想象一下你在工厂里安排流水线。如果每个工位每天处理的产品规格都在随机变化,一会儿是螺丝,一会儿是齿轮,那整条流水线的物料准备、设备切换都会乱套。如果你坚持每一步都固定处理同一套标准化流程,效率反而更高。弹性训练的随机抽样策略,正是让计算系统每天都在"切换生产线"。

那如果每一步都训练全部子网络呢?负荷确实稳定了,但代价是每多一个子网络,就要多跑一次完整的前向和反向传播,训练成本随子网络数量线性增长,省下来的钱又变回去了。

第二个坑是超参数难以迁移。

学习率、初始化规模这些超参数*(Hyperparameters, HP):模型训练前需要人工设定的配置项,比如学习率决定了模型每次调整的步长大小。

在不同模型宽度下的最优值是不一样的。标准参数化方式下,一个在小模型上调好的学习率,搬到大模型上可能完全不管用。这就是说,即使你训练出了弹性模型家族,调参这件事本身也要在每个目标规模上重新来一遍,这又会把弹性训练节省下来的算力吃掉一部分。

第三个坑最隐蔽:没人能保证共享参数里的所有子网络能同时被"喂饱"。

论文里提到一个概念,叫最大化更新*(Maximally Updated):一个参数在初始化时输出规模是常数级别,训练一步后的变化幅度也保持在合理范围,既不会消失也不会爆炸。

这个概念听起来抽象,但背后的问题很实际。因为核心网络和外围网络共享着同一组权重,如果超参数设置得对核心网络刚好合适,很可能对外围的扩展部分就太大或太小了。调整一个全局的超参数值,往往是在不同宽度之间"拆东墙补西墙",总有某个宽度被牺牲掉。

这三个问题凑在一起,让弹性训练这个本该省钱的想法,在实际落地时处处是坑。论文作者给出的解法分两部分:一个叫ShellElastic的新架构解决第一个问题,一个叫μ-ELASTIC的参数化方法解决后两个问题。

ShellElastic:把隐藏维度切成核心和外壳

ShellElastic的设计思路,可以用一句话概括:把模型的隐藏维度组织成一个密集核心,加上一层一层往外扩展的外壳。

具体来说,论文把隐藏维度分成若干个连续的分组。第一个分组叫核心,之后依次是壳一、壳二……每个支持的子网络,都是核心加上若干层外壳的组合。

举个具体数字,论文的默认配置是基础宽度128,支持的子网络索引集合是{4, 5, 6, 7, 8}。这意味着核心是512维(4乘128),然后依次叠加四层各128维的外壳,最终得到五个子网络,宽度从512一路到1024。

这里有个关键设计:每一个隐藏层的权重矩阵,都被设计成块下三角形式*(Block Lower Triangular):一种矩阵结构,矩阵中某个输出分组只能读取它自己和它之前的所有分组的输入,绝不能读取排在它后面的分组。

说得直白一点,壳二可以看到核心和壳一的信息,但壳二绝对看不到壳三、壳四的信息。这个单向的信息流动规则,是整个架构能省钱的核心所在。

为什么这个规则重要?因为它保证了一件事:你跑一次完整宽度(也就是最大子网络)的前向传播,这次传播里天然就包含了所有更小子网络的精确隐藏状态,它们是作为前缀存在的。

这就好比你去自助餐厅打菜,一份满配的大份套餐里,天然就包含了小份套餐该有的所有菜品,只是多加了几样。你不需要单独为小份套餐重新炒一遍菜,大份套餐里切出来的那一部分,和单独做一份小份套餐,味道分毫不差。如果满配套餐里的食材摆放是乱序的,小份套餐想要的那几样菜混在后面加的菜里分不清楚,那你就必须单独重新准备,这正是MatFormer和完全弹性模型(Fully Elastic)面临的情况,它们的隐藏矩阵是稠密的,不同宽度子网络的隐藏状态互不包含,所以每种宽度都得单独跑一遍完整的前向传播。

论文把这个性质写成了命题一(Proposition 1),并且不仅在理论上证明了,还做了数值验证:用一个两层的小模型,对比一次共享主干传播和五次独立评估的输出,两边的梯度最大误差只有1.68乘以10的负8次方,基本可以认为完全一致。

那么训练全部子网络的成本到底是多少?论文给出了一个公式思路:联合训练的成本约等于主干网络跑一次完整宽度,再加上每个子网络各自终端归一化层和读出层的少量开销。而分别训练五个子网络的成本,则是五份完整主干加五份终端开销。随着支持的子网络数量增加,差距会越拉越大,但联合训练的成本始终接近单独训练最宽那个子网络的成本。

实测数据也验证了这一点。在119M到285M参数规模下,ShellElastic训练一步,耗时601毫秒(平均到五个层级),而五个独立训练的模型平均每步要737毫秒。这还只是算力层面的节省,没算上五次独立训练需要五倍的总步数这件事。

当然这个架构还有一个巧妙的细节,叫做分组局部归一化*(Group-local Normalization):模型里每个分组各自计算自己的均值方差来做归一化,而不是用整个网络宽度统一的统计量。

这个设计同样是为了保证前缀能被精确复用。如果归一化是跨越整个激活宽度计算的,那加一层壳进去,统计量就会变,连带前面核心部分的输出都会跟着变化,前缀复用的性质就被破坏了。论文还对比了另外两种归一化方式,块局部归一化和累积前缀归一化,发现三者的损失差异都在0.001个单位以内,但累积前缀归一化因为要为每个前缀计算统计量,训练时间多花6%,内存多花23%,所以分组局部是性价比最高的选择。

架构搭好了,接下来要解决的,是超参数怎么在这个共享架构里设置才合理。

μ-ELASTIC:让每一层都被喂到刚好的分量

超参数迁移这件事,本质上是一个物理尺度问题。

论文提到的最大更新参数化*(Maximal Update Parameterization, 简称μP):一种让神经网络在不同宽度下保持特征更新规模一致的参数化方法,由Yang & Hu在2021年提出,使得在小模型上调好的超参数能直接用在大模型上。

μP的核心贡献是,针对单一的、不共享参数的模型,证明了只要按照特定规则缩放初始化方差、学习率和权重衰减,就能让网络在变宽时保持每一层的特征更新幅度稳定,既不会随着宽度增加而消失,也不会炸掉。这给了业界一个重要的工具:在小模型上调好的超参数,可以直接搬到大模型上用,不用重新搜索。

但μP解决的是单个独立模型的问题。弹性模型的麻烦在于,一套参数里塞了好几个不同宽度的子网络,它们共享着同一批权重矩阵。这时候该怎么设置超参数?

论文给出的答案叫μ-ELASTIC,核心是定义了一个组合比例q,它由两个因子相乘得到:一个叫μ-Transfer因子(记作m),负责把从非弹性的小代理模型上调好的超参数,迁移到弹性模型家族;另一个叫μ-Internal因子(记作子网络宽度s),负责处理同一个弹性模型内部不同宽度之间的差异。

这个设计巧妙的地方在于,对于ShellElastic架构里的某个输出分组,无论它被包含在哪个更宽的子网络里跑,它实际读取的输入维度数永远是固定的(因为块下三角结构保证了这一点)。既然读取的信息量不变,那针对这个分组设置的初始化方差、学习率、权重衰减,自然可以是同一套,不用因为它身处哪个子网络而变来变去。

论文把这个性质证明成了定理一(Theorem 1),叫做μ-Internal:简单来说,只要按照这套规则设置,无论某个分组被哪个包含它的子网络调用,它产生的激活值在初始化时都保持常数级别大小,更新一步后的变化幅度也都稳定在同一个水平,不会因为所处的子网络不同而忽大忽小。

这就好比一个工厂里的某个零件生产线,不管最终组装成的是小型玩具车还是大型卡车模型,这条生产线每次加工零件消耗的原材料和加工强度始终是固定的,不会因为组装的整车大小而被要求加大马力或者偷工减料。如果不这样设计,而是用同一套超参数笼统地套用到核心和最外层的壳上,论文的命题二(Proposition 2)证明了一件更残酷的事:核心部分和最外层壳的激活值规模会相差根号M除以s1倍(M是最大宽度索引,s1是最小宽度索引),更新幅度相差M除以s1倍。换句话说,不管你怎么调这个共享的超参数值,总有一端被喂得过饱,另一端被饿着。

论文还证明了μ-ELASTIC能同时实现两层迁移,这叫做推论一(Corollary 1):第一层,从非弹性的小代理模型迁移到整个弹性模型家族;第二层,当弹性模型本身被加宽或者加壳变得更大时,之前调好的超参数依然适用。

实验数据印证了这套理论。论文在119M到285M参数规模下,对基础学习率和初始化规模分别做了扫描实验。代理模型(宽度128)调出来的最优学习率是0.002,这个值直接搬到ShellElastic的全部五个子网络上,损失和各自最优点的差距都在0.005个单位以内。换到更大规模(375M到906M参数)也依然成立,证明这套规则确实能同时支持两种迁移。

而标准参数化方式(也就是不做任何按宽度缩放调整的做法)在同一组实验里就没这么稳定了,它的最优超参数会随着规模变化而漂移,论文测出标准参数化的最佳表现,比μ-ELASTIC还差了0.03到0.11个单位。

论文还单独测了μ-Internal这条性质带不带来实际收益。做法是,先给标准参数化方式一个"作弊"的机会,让它在12到13种配置里挑出验证损失最低的那一个,再和μ-ELASTIC不经任何微调直接用默认值去对比。结果是,即便标准参数化挑了自己最好的那一版,μ-ELASTIC依然在训练损失和验证损失上全线领先,证明光靠"事后调好"是弥补不了结构性的宽度失衡问题的,必须从根本上让每个分组按照自己真实的信息吞吐量来设置参数。

实际效果:五合一的训练到底能打多少分

理论讲完了,接下来看看这套方案在真实预训练里表现如何。

论文在三个规模上做了对比:119M到285M参数、375M到906M参数、675M到1.69B参数。对比的对象包括MatFormer、完全弹性模型、单独训练的标准模型(Standalone),以及把标准模型塞进同样步数训练的采样标准模型(Sampled Standalone)。

数据集是DCLM-Baseline,一个经过过滤的Common Crawl网页文本语料库,默认训练步数是25000步,相当于喂进去65.5亿个token。

在119M到285M规模下,ShellElastic相比MatFormer和完全弹性模型,平均验证损失低了0.028个单位,五个层级(从Low到Max)的5-shot准确率平均值也是四种方法里最高的,达到36.5%。更关键的是,ShellElastic在Max层级(最大子网络)相比其他基线方法领先得最多,差距可以达到0.023到0.050个单位,论文给出的解释是,其他弹性方法里只有最大子网络独享的那部分参数,平均每五步才被训练一次,而ShellElastic的每个参数每一步都被更新。

拿标准单独训练的模型(花了五倍步数)来对比,ShellElastic在最大规模(675M到1.69B参数)下,平均验证损失只差了0.014个单位,5-shot准确率差距不到0.8个百分点。这意味着花五分之一的训练步数,拿到了接近单独训练五个独立模型的效果。

论文还和一个同期的竞品StairFormer做了对比。StairFormer也是用前缀结构来实现弹性子网络,但它用等宽分组,用累积前缀归一化,而且所有子网络共享同一个读出层。两者的区别在于ShellElastic用的是不等宽的"大核心小外壳"设计,并且每个子网络都有自己独立的归一化层和读出层。

论文还专门做了一次架构搜索,固定每个隐藏矩阵保留的参数块数量和子网络数量,只改变那个比基础宽度更宽的分组放在哪个位置。结果很统一:把宽的那个分组放在核心位置,损失最低;放在外壳里,会让包含这个宽分组的最窄子网络损失明显变差,差距最高到0.052个单位。这个发现直接验证了ShellElastic把宽核心、薄外壳作为默认设计的合理性。

这就像盖房子打地基,如果你把最厚实、承重能力最强的那部分材料放在地基上,上面轻便的结构可以随意加盖楼层。但如果把最厚实的部分留到某一层中间去用,下面楼层的承重反而成了瓶颈,整体稳定性反而更差。

子网络集成:一个还没被完全解锁的加分项

论文在消融实验里探索了一个很有意思的方向,叫子网络集成。

既然一次主干传播同时产生了五个子网络的输出分布,那把这五个预测结果按某种权重组合起来,是不是能比单独用最大的那个子网络效果更好?

论文测了四种集成方式:平均权重、按置信度加权、只挑最有信心的那个读出层、以及事后找出能让损失最小的固定权重组合。结果是,这些组合方式都没能超过单用最大子网络,置信度加权和平均权重比最大子网络差了0.03到0.065个单位,只挑最有信心的读出层更差,差了0.068个单位。

但论文接着做了一个"开天眼"的实验:假设有一个oracle(神谕),每次都能提前知道哪个子网络对当前token预测得最准,直接选那个最准的。这个oracle的损失,比单用最大子网络低了0.19个单位,而且它选中最窄子网络的比例高达22.7%,意味着并非永远是最大的模型表现最好,有超过七成的token上,某个更小的子网络其实表现得比最大子网络更好。

这个发现挺有意思的。你可能会以为,模型越大,预测能力就该是全方位压制小模型的,但实验数据戳破了这个假设。不同宽度的子网络其实在不同类型的token上各有所长,只是目前没有办法在不知道答案的前提下,提前判断该信哪个子网络。

这就像一个团队里,经验最丰富的老员工平均水平确实最高,但遇到某些特定的客户难题,反而是刚入职不久、思路更灵活的新人能一针命中。问题是,领导不可能提前知道哪个难题该交给谁,只能靠事后复盘才发现"原来这单该让小张来"。

论文坦诚地说,这个oracle揭示的0.19个单位的提升空间,目前还没有办法在没有答案的情况下兑现,但这给未来的研究留下了一个明确的方向:怎么设计一个不依赖标签、仅凭输入特征就能判断哪个子网络更可信的路由机制。

一个能"长大"的架构

论文还提到了ShellElastic的一个额外性质,这个性质MatFormer和完全弹性模型都不具备:训练完成之后,还可以继续往外加壳。

因为块下三角结构保证了没有任何一个分组会读取在它之后的分组,而每个分组又是独立归一化的,这意味着你可以在一个已经训练好的ShellElastic模型上,新增一层壳,用全新的数据(比如下游任务数据)单独训练这层新壳,同时把之前训练好的所有权重冻结不动。这样一来,所有之前训练好的子网络,输出结果分毫不变,而你又额外获得了一个更大、更强的新子网络。

这个性质在MatFormer和完全弹性模型上是做不到的,因为它们的新增单元会直接写入已经训练好的隐藏状态里,或者改变已训练权重所接收的输入分布,相当于往已经建好的房子里强行拆墙加梁,原本的结构承受力都会受影响。而ShellElastic的加壳过程,更像是在已经封顶的建筑旁边新盖一栋独立的侧楼,两者共享地基但互不干扰主体结构。

写在后面

读完这篇论文,印象最深的其实不是ShellElastic本身的架构设计,而是论文把"弹性训练"这个听起来很简单的想法,拆解成了三个独立的工程问题,然后分别用严谨的数学证明去回应。

很多论文讲到参数共享,容易止步于"效果好"这种描述,但这篇论文花了大量篇幅去证明为什么效果好,甚至给出了反例证明别的做法为什么不行。命题二那个关于标准参数化方式必然导致核心和外壳不匹配的证明,看完让我意识到,很多工程上"调一调超参数就好了"的经验法则,背后其实藏着根号关系和线性关系的数学必然性,不是靠多试几组参数就能绕开的。

另一个让我反复琢磨的地方,是oracle实验里那个22.7%的数字。它说明了一件事:参数更多、算力更强的模型,不代表它在每一个具体任务上都比小模型强。这个事实如果放在更大的语境里想一想,其实对现在"越大越好"的大模型叙事是个提醒。模型的能力分布可能本身就是不均匀的,某些场景下,一个精简的小模型反而更懂行。

至于这个发现能不能真正转化成实用的路由机制,论文自己也没给出答案,只是诚实地把这块留白写了出来。这种诚实感,比强行给出一个"完美解决方案"的论文更让人信服。

Q&A

Q1:ShellElastic是什么?

A:ShellElastic是一种弹性Transformer架构,它把隐藏维度组织成一个密集核心加若干层外壳,通过块下三角权重矩阵,让一次完整宽度的前向传播天然包含所有更小子网络的精确隐藏状态,从而用接近单一模型的训练成本同时训练出多个不同体积的模型。

Q2:μ-ELASTIC解决了什么问题?

A:μ-ELASTIC是一种超参数缩放方案,它能把在小型非弹性代理模型上调好的学习率、初始化规模等超参数,直接迁移到弹性模型家族,并保证共享参数里的每个子网络都能被同时稳定地训练到位,不会出现某个宽度被忽略或过度更新的情况。

Q3:ShellElastic相比单独训练多个模型有什么优势?

A:ShellElastic只需单独训练五倍步数的五分之一左右,就能在验证损失和下游任务准确率上接近分别训练五个独立模型的效果,同时每一步训练的计算负荷和成本都接近只训练一个最大模型的开销。

文章标签AI资讯

资讯来源:由AI资讯编辑整理自互联网公开内容,版权归原作者所有,未经许可,不得转载。