ShellElastic一次训练出五个模型:弹性训练的三个坑怎么填
你有没有想过,一家做手机应用的公司,为什么总是要同时准备好几个版本的AI模型。旗舰手机用大模型,跑得动也效果好。入门手机用小模型,省电省内存。平板用中等模型,刚好够用。问题是,这些模型通常是分开训练出来的。训练一个十几亿参数的大模型要花多少钱,业内人都清楚。现在要五个不同大小的模型,难道要花五倍的钱,...
做手机应用的公司,为什么总得同时备着好几个版本的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%。Max层级(最大子网络)上它领先最多,差距达到0.023到0.050个单位。论文的解释是,其他弹性方法里只归最大子网络独享的那部分参数,平均每五步才训练一次,而ShellElastic的每个参数每一步都在更新。
再拿花了五倍步数、标准单独训练的模型来比:在最大规模(675M到1.69B参数)下,ShellElastic的平均验证损失只差0.014个单位,5-shot准确率差距不到0.8个百分点。五分之一的训练步数,换来的是接近五个独立模型各自训练的效果。
论文还对比了同期竞品StairFormer。它同样用前缀结构实现弹性子网络,但走等宽分组、累积前缀归一化,所有子网络共用一个读出层。ShellElastic的区别在于不等宽的"大核心小外壳"设计,每个子网络各有独立的归一化层和读出层。
架构搜索也做了一次:固定每个隐藏矩阵保留的参数块数量和子网络数量,只挪动那个比基础宽度更宽的分组的位置。结论相当一致——宽分组放核心处损失最低;放进外壳,包含它的最窄子网络损失会明显变差,差距最高到0.052个单位。宽核心、薄外壳这个默认设计由此得到直接验证。
这就像盖房子打地基:最厚实、承重最强的材料放在地基上,上面轻便的结构可以随便加盖;反过来把最厚实的部分留到中间某层去用,下面楼层的承重成了瓶颈,整体稳定性反而更差。
子网络集成:一个还没被完全解锁的加分项
消融实验里,论文还探了一个挺有意思的方向——子网络集成。
一次主干传播能同时产出五个子网络的输出分布,那么按某种权重把这五个预测组合起来,会不会比单用最大的子网络更好?
论文试了四种集成方式:平均权重、按置信度加权、只挑最有信心的那个读出层、事后找出让损失最小的固定权重组合。结果没有一种能超过单用最大子网络——置信度加权和平均权重差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:它是一种弹性Transformer架构,隐藏维度被组织成密集核心外加若干层外壳,借助块下三角权重矩阵,一次完整宽度的前向传播就天然包含所有更小子网络的精确隐藏状态,于是用接近单一模型的训练成本,同时训出多个不同体积的模型。
Q2:μ-ELASTIC解决了什么问题?A:这是一套超参数缩放方案,能把在小型非弹性代理模型上调好的学习率、初始化规模等超参数直接搬到弹性模型家族,并保证共享参数中每个子网络都能被同时稳定训练到位,不会出现某个宽度被忽略或过度更新的情况。
Q3:ShellElastic相比单独训练多个模型有什么优势?A:它只需要单独训练五倍步数的五分之一左右,验证损失和下游任务准确率就能接近分别训练五个独立模型的效果,而且每一步训练的计算负荷和成本都接近只训一个最大模型的开销。