Relative Generalization Invariance of LLM Pretraining
Fengzhuo Zhang, Shuche Wang, Shenggui Li, Tianyu Ruan, Jianliang He, Ivor Tsang, Tianyu Pang, Chao Du, Tianwei Zhang, Zhuoran Yang
cs.LG, cs.AI
2026-09-27
Yale与Sea AI Lab用124M到4.5B GPT证明:同数据下Adam与Muon等六种优化器2000步后rCCC超过0.95,token损失只差常数;换成ArXiv或代码数据流,相对泛化结构即破裂。
换一个优化器、改一下注意力、换一批预训练数据,验证损失都会动。问题是这三类改动动的是不是同一种东西。Soap、Muon、滑动窗口注意力、动态选数据都在压低预训练损失,缺的是把优化器、架构、数据流拆开,看各自改了泛化的哪一层结构。
Yale、NUS、NTU与Sea AI Lab把问题收成一件可测的事:两个模型对任意两个token的验证损失之差是否相同。如果相同,相对难易排序没变,整体只平移了一截。他们把这种现象叫做Relative Generalization Invariance(RGI,相对泛化不变性)。
RGI要求两个参数θ和θ′对验证集里任意token对d、d′满足损失差相等。等价说法是两套token级损失只差一个与token无关的常数C。训练侧用三元组T=(优化器,初始化及架构,数据流)生成模型;若对几乎所有初始化与数据流realization都成立,就说两条训练轨迹之间有RGI。足够晚之后一直成立,则称渐近RGI(ARGI)。
近似程度用两个归一化指标。rCCC把验证集切成K块,算两套中心化损失的Concordance Correlation Coefficient,越接近1越像纯平移。rdiff比较成对子集上的损失差,越接近0越接近精确RGI。
主干实验是124M Nano-GPT,在FineWeb上训练,验证集同分布但与训练集不相交。六种优化器里,坐标型是Adam、Lion、Adam-mini,矩阵型是Muon、Soap、Scion。同一设置共享初始化与数据流,学习率各自调到最好,调度器用Warmup-Stable-Decay。主文报124M,0.7B与4.5B结论一致。跨10个随机种子取最差的rCCC与rdiff。
对照逐项拧:学习率乘0.25到1.25、动量0.85/0.9/0.95、独立初始化、隐藏维从768收到384和64、注意力换成NoPE、ALiBi或带RoPE的SWA、FFN换成SReLU或SiLU、数据流换成C4、ArXiv或CodeParrot,验证集再做同一套分布偏移,外加按词表均匀抽样的Random。
同初始化、同FineWeb时,优化器几乎只做平移。10000步、20480个验证batch上,Muon对Adam的点拟合为y=x-0.07,R²=0.99。2000步之后,其余五种相对Adam都满足rCCC>0.95、rdiff<0.05。
| 对照 | 指标 | 结果 |
| Muon vs Adam,10000步,20480 batch | 验证损失直线 | y=x-0.07,R²=0.99 |
| 五种优化器 vs Adam,≥2000步 | rCCC / rdiff | >0.95 / <0.05 |
| 隐藏维768 vs 384 | RGI | 成立 |
| 隐藏维768 vs 64 | RGI | 失败 |
| FineWeb vs ArXiv或CodeParrot | rCCC | 全程低于0.95 |
独立初始化开头会偏,大约2000步后仍回到近似RGI。RoPE全注意力的Adam对照NoPE、ALiBi、SWA,以及GELU对照SReLU,rCCC同样能越过0.95。学习率与动量扫描也一样,只改变出现时间。1.5倍学习率未纳入,因为部分优化器会发散。
训练损失曲线在2000步后形状几乎一样。相邻步损失拆成一次参数更新α和相邻batch自身的损失差β,|β|大约是|α|的50倍。β是还没见过的下一批数据上的损失差,近似RGI会让它跨优化器几乎不变,所以曲线长得像。
数据流会拆掉这件事。Adam训FineWeb、对照优化器训C4时,rCCC先掉再升再掉,中间那截升可能因为两者都部分来自Common Crawl。ArXiv与CodeParrot上,即便对照也用Adam,rCCC全程低于0.95。验证集偏移温和(C4)时RGI还能转过去,ArXiv与代码上变弱但仍随训练上升,Random无结构token上看不到RGI。
NTK懒学习解释不了。5000步后参数和隐状态都远离初始化,跨优化器余弦相似度还在下降。Mean-field也撑不住:宽度收到192(基准768的四分之一)时,同宽度内优化器之间RGI仍在,但这条宽度的验证损失对学习率曲线已经和384、576、768不像同一套平均场画像。跨宽度比768与64时RGI失败,边界在「中等」和「砍太狠」之间。
理论给了一个过参数二次模型f=ABx。目标矩阵列正交且非零元等幅、协变量各向同性时,GD、去掉动量的Adam(按元素取符号)、去掉动量的Muon(谱归一化)对等范数输入精确满足RGI,同时|cos-sim(GD,Adam)|=O(1/t),|cos-sim(Adam,Muon)|=O(1/t²)。
对预训练来说,这篇把「换优化器」和「换数据」拆成两种干预。优化器和中等架构改动更像给所有token加减同一截损失,难的还是难。数据来源一换,相对结构会动。
想改哪些token学得好,优先动数据流,别指望换Muon就能重排难度。
不同优化器训练曲线长得很像,这里有一个机制:形状由β主导,β被RGI钉住。它解释不了谁的绝对损失更低,也没报下游任务。选优化器之前,先问相对泛化有没有被改过的那一轴动掉。
范围自己划过:只做自回归LLM,扩散模型留给以后。二次模型假设很窄,Adam与Muon还去掉了动量,放松假设是未来工作。隐藏维收到64时RGI失败,「中等架构变化」有硬边界,论文没给出这条边界在真实大模型上对应什么规模。
主结果是验证交叉熵,没有MMLU、代码或长上下文数字,不能把相对泛化不变读成能力结构不变。C4上rCCC先升后降,只给了分布相近的猜测,没有量化「相近到什么程度会保住RGI」。4.5B只在附录复现同结论,主文数字几乎全是124M。Random验证集失败也说明,RGI依赖验证分布仍有语义结构,不是任意输入上都成立。