assassinlike's picture
Upload folder using huggingface_hub
909636f verified
|
Raw History Blame Contribute Delete
24.9 kB

论文内容叙述

问题从哪里来

大模型预训练的数据并不是一个同质整体,而是由通用文本、数学、代码、推理等多个来源混合而成。设第 (k) 个来源的数据集为 (D_k),它在混合数据中的采样概率为 (\alpha_k)。论文把该来源的加权经验损失写成

Lk(θ)=−αk∑j=1∣Dk∣log⁡p(xj∣θ), L_k(\theta)=-\alpha_k\sum_{j=1}^{|D_k|}\log p(x_j\mid \theta),

总预训练目标则是

Ltrain(θ)=1K∑k=1KLk(θ). L_{\mathrm{train}}(\theta)=\frac1K\sum_{k=1}^K L_k(\theta).

通常只看这个总损失是否足够低,但总损失相同并不代表模型处在相同的几何位置。用两个来源说明:即使

Ltrain(θ)=12(L1(θ)+L2(θ)) L_{\mathrm{train}}(\theta)=\frac12\bigl(L_1(\theta)+L_2(\theta)\bigr)

取得相同最小值,也可能有两种完全不同的情形。一种情形中,总损失的极小点只是两个任务损失相互折中的“和的极小点”:它离 (L_1) 和 (L_2) 各自的极小点都很远。另一种情形中,各任务的极小点彼此集中,总损失极小点也靠近它们的交集。论文分别称它们为 Sum of Minima 和 Intersection of Minima。

作者要解决的核心问题是:当预训练损失相同时,第二种更接近各任务共同极小点的解,是否会在未知下游任务上泛化得更好?标准 AdamW 只最小化混合损失,没有机制主动选择这种解。论文因此把优化器的“隐式偏置”作为研究对象:不同优化器即使把训练损失降到同一水平,也可能系统地偏向不同几何性质的解。

这里没有独立的 related work 章节。论文与以往工作的关系主要体现在两条对比中。已有泛化理论大量研究极小点附近的平坦性,但很少研究训练解与不同任务极小点之间的距离;已有新优化器如 Muon 通常通过更快地降低预训练损失改善下游能力,而本文希望在预训练损失不变时,仅靠改变最终解的几何偏置来改善泛化。

“接近性”为什么是一个泛化因素

设预训练最终参数为

θtrain∗∈arg⁡min⁡θLtrain(θ). \theta_{\mathrm{train}}^*\in\arg\min_\theta L_{\mathrm{train}}(\theta).

对一个未知下游任务 (T),令其局部极小点集合为

ST={θ∣∃ϵ>0, ∀θ′∈Bϵ(θ), LT(θ)≤LT(θ′)}. S_T=\left\{\theta\mid \exists\epsilon>0,\ \forall\theta'\in B_\epsilon(\theta),\ L_T(\theta)\leq L_T(\theta')\right\}.

由于神经网络可能有许多局部极小点,论文不任意选择一个,而是定义离当前预训练解最近的下游极小点

θT∗=arg⁡min⁡θ∈ST∥θ−θtrain∗∥2. \theta_T^*=\arg\min_{\theta\in S_T}\|\theta-\theta_{\mathrm{train}}^*\|_2.

在 (\theta_T^) 处对下游损失作二阶 Taylor 展开。因为 (\theta_T^) 是局部极小点,(\nabla L_T(\theta_T^*)=0),所以一阶项消失:

LT(θtrain∗)=LT(θT∗)+12(θtrain∗−θT∗)⊤∇2LT(θT∗)(θtrain∗−θT∗)+O(∥θtrain∗−θT∗∥3)≤LT(θT∗)+12∥θtrain∗−θT∗∥22max⁡ξ∈[θT∗,θtrain∗]∥∇2LT(ξ)∥2. \begin{aligned} L_T(\theta_{\mathrm{train}}^*) &=L_T(\theta_T^*) +\frac12(\theta_{\mathrm{train}}^*-\theta_T^*)^\top \nabla^2L_T(\theta_T^*) (\theta_{\mathrm{train}}^*-\theta_T^*) +O(\|\theta_{\mathrm{train}}^*-\theta_T^*\|^3)\\ &\leq L_T(\theta_T^*) +\frac12\|\theta_{\mathrm{train}}^*-\theta_T^*\|_2^2 \max_{\xi\in[\theta_T^*,\theta_{\mathrm{train}}^*]} \|\nabla^2L_T(\xi)\|_2. \end{aligned}

这个上界把二阶泛化差距拆成两个因子:

  • (\max_\xi|\nabla^2L_T(\xi)|_2) 描述从下游极小点到预训练解这条方向上的曲率,即平坦性;
  • (|\theta_{\mathrm{train}}^*-\theta_T^*|_2^2) 描述预训练解离下游极小点多远,论文称为 closeness。距离越小,closeness 越好。

如果这条一维方向上的损失严格是二次函数,上式成为精确等式而不只是上界。因此作者认为,在深度网络损失沿典型方向近似二次的经验前提下,平坦性和接近性共同概括了二阶层面的泛化信息。本文只研究过去较少受到关注的接近性。

但不能无代价地让所有任务极小点靠近。若在不增加别的代价时就能大幅减小

1K∑k=1K∥θtrain∗−θk∗∥22, \frac1K\sum_{k=1}^K\|\theta_{\mathrm{train}}^*-\theta_k^*\|_2^2,

通常也意味着训练误差更小、收敛更快,这与寻找显著更快优化器本身很困难的事实冲突。论文特意研究一个受控比较:不同解的总训练损失完全相同,接近性的改善以各任务所在盆地的“内禀损失”升高为代价,并暂时排除平坦性变化的影响。这样才能单独判断接近性的作用。

严格二次情形

为模拟非凸网络,作者设参数空间被分成多个互不相交的吸引盆地 (B)。在某个盆地内,从任务分布 (P) 抽到的任一任务都具有局部二次损失

L(θ)=a2∥θ−θL∗∥22+cB, L(\theta)=\frac a2\|\theta-\theta_L^*\|_2^2+c_B,

其中任务极小点满足

θL∗∼P(μB,σB2I), \theta_L^*\sim P(\mu_B,\sigma_B^2I),

(c_B) 是该盆地自身最低能达到的内禀损失。预训练的 (K) 个任务和下游任务独立同分布地来自 (P)。只比较训练损失都等于同一常数 (C_{\mathrm{train}}) 的候选盆地。

总训练损失的驻点是样本极小点的均值:

θtrain,B∗=1K∑k=1Kθk,B∗. \theta_{\mathrm{train},B}^*=\frac1K\sum_{k=1}^K\theta_{k,B}^*.

固定训练损失后,盆地深度必须满足

cB=Ctrain−a2K∑k=1K∥θtrain,B∗−θk,B∗∥22. c_B=C_{\mathrm{train}}- \frac{a}{2K}\sum_{k=1}^K \|\theta_{\mathrm{train},B}^*-\theta_{k,B}^*\|_2^2.

因此,任务极小点越集中,第二项越小,为维持同一个 (C_{\mathrm{train}}),(c_B) 就必须越高。这正是论文设定的取舍:用较差的盆地内禀深度换更小的任务间方差。

对新任务 (T),将这个 (c_B) 代回下游损失。训练样本均值与独立新样本之间的期望平方距离为 ((1+1/K)\sigma_B^2),训练样本围绕样本均值的平均平方距离为 ((K-1)\sigma_B^2/K)。两者相减后,内禀损失的影响恰好消掉,得到

ET∼P[LT(θtrain,B∗)]=Ctrain+aKσB2. \mathbb E_{T\sim P}\left[L_T(\theta_{\mathrm{train},B}^*)\right] =C_{\mathrm{train}}+\frac aK\sigma_B^2.

所以在这个模型中,只要训练任务和下游任务同分布,在训练损失相同的条件下,下游损失严格随任务极小点方差 (\sigma_B^2) 增长。任务极小点越接近,泛化越好。

超出二次函数的推广

一般情形中,作者先选总体期望损失 (\mathbb E_{L\sim P}[L(\theta)]) 的一个局部极小点 (\theta^) 作为盆地锚点。对每个任务 (L),从其局部极小点集合中选择离 (\theta^) 最近的 (\theta_L^*),并定义

μ=E[θL∗],σ2=E∥θL∗−μ∥22. \mu=\mathbb E[\theta_L^*],\qquad \sigma^2=\mathbb E\|\theta_L^*-\mu\|_2^2.

这里 (\sigma^2) 就是任务极小点的总体接近性指标。定理要求:沿每条 ([\theta_L^*,\theta^*]) 方向,损失局部且方向性地强凸,即对相关子空间中的单位向量 (u),

0<λmin⁡≤u⊤∇2L(ξ)u≤λmax⁡,ξ∈[θL∗,θ∗]. 0<\lambda_{\min}\leq u^\top\nabla^2L(\xi)u \leq\lambda_{\max}, \qquad \xi\in[\theta_L^*,\theta^*].

此外,任务的 Hessian 与任务极小点位置在分布上要统计独立。令条件数 (\kappa=\lambda_{\max}/\lambda_{\min}),可以证明:

ET∼P[LT(θtrain∗)]−Ctrain≤λmax⁡2K[(λmax⁡λmin⁡)2+1]σ2. \mathbb E_{T\sim P}[L_T(\theta_{\mathrm{train}}^*)]-C_{\mathrm{train}} \leq \frac{\lambda_{\max}}{2K} \left[ \left(\frac{\lambda_{\max}}{\lambda_{\min}}\right)^2+1 \right]\sigma^2.

这说明在这些假设下,一般损失的泛化差距仍按 (O(\sigma^2/K)) 缩放。不过这不是无条件结论:同分布、局部方向强凸以及 Hessian 与极小点位置独立都是证明不可缺少的前提。

从极小点距离转为可计算的梯度相似度

直接优化 (|\theta-\theta_k^*|) 不可行,因为每一步都要先为每个来源完整求出一个极小点。论文改用任务梯度相似度。

设 (\theta) 已是总训练损失的驻点:

1K∑k=1K∇Lk(θ)=0. \frac1K\sum_{k=1}^K\nabla L_k(\theta)=0.

令 (\theta_k^) 是离 (\theta) 最近的第 (k) 个任务局部极小点,并设沿 ([\theta,\theta_k^]) 的位移方向,Hessian 曲率下界为 (\lambda_{\min}>0),同时 (G=\sup_k|\nabla L_k(\theta)|_2)。论文证明了

1K∑k=1K∥θ−θk∗∥22≤1Kλmin⁡2∑i≠j−∇Li(θ)⊤∇Lj(θ)≤G2Kλmin⁡2∑i≠j(1−CosSim⁡(∇Li,∇Lj)). \frac1K\sum_{k=1}^K\|\theta-\theta_k^*\|_2^2 \leq \frac{1}{K\lambda_{\min}^2} \sum_{i\ne j}-\nabla L_i(\theta)^\top\nabla L_j(\theta) \leq \frac{G^2}{K\lambda_{\min}^2} \sum_{i\ne j} \left(1-\operatorname{CosSim}(\nabla L_i,\nabla L_j)\right).

最大化不同任务的梯度余弦相似度,便会压低参数接近性的上界。

梯度相似度还有一个不依赖极小点距离的直接解释。若在训练集上做一步梯度下降,则对下游任务的一阶影响是

LT(θ)−LT(θ−γ∇Ltrain(θ))=γ∇Ltrain(θ)⊤∇LT(θ)+O(γ2). L_T(\theta)-L_T(\theta-\gamma\nabla L_{\mathrm{train}}(\theta)) =\gamma\nabla L_{\mathrm{train}}(\theta)^\top\nabla L_T(\theta) +O(\gamma^2).

也就是说,若训练梯度和下游梯度同向,训练一步会同时降低下游损失。因此论文把“参数极小点接近”和“梯度接近”合并在 closeness 这个词下使用:后者既是前者的可计算代理,又直接表示任务之间是否产生建设性干涉。

Nexus 如何近似二阶相似度梯度

若直接对梯度余弦相似度求参数梯度,必然出现 Hessian。记 (g_i=\nabla L_i(\theta))、(H_i=\nabla^2L_i(\theta)),则

∇θCosSim⁡(gi,gj)=(I−gigi⊤∥gi∥22)Higj∥gi∥2∥gj∥2+(I−gjgj⊤∥gj∥22)Hjgi∥gj∥2∥gi∥2. \begin{aligned} \nabla_\theta\operatorname{CosSim}(g_i,g_j) ={}&\left(I-\frac{g_ig_i^\top}{\|g_i\|_2^2}\right) \frac{H_i g_j}{\|g_i\|_2\|g_j\|_2}\\ &+\left(I-\frac{g_jg_j^\top}{\|g_j\|_2^2}\right) \frac{H_j g_i}{\|g_j\|_2\|g_i\|_2}. \end{aligned}

精确计算这些 Hessian-vector product 在大模型训练中代价高,而且常与 FlashAttention 等只高效支持一阶反向传播的内核不兼容。Nexus 的关键做法是不显式构造 Hessian,而用多个顺序的归一化 SGD 小步所产生的总位移,自动形成同样的 Hessian-gradient 交互项。

标准算法

每个外层迭代 (t) 开始时令

θt,0=θt−1. \theta_{t,0}=\theta_{t-1}.

然后做 (K) 个内层步骤。第 (m) 步均匀抽取一个任务 (s_m\in{1,\ldots,K}),计算当前位置的梯度并归一化:

gm=∇Lsm(θt,m−1),θt,m=θt,m−1−γgm∥gm∥2. g_m=\nabla L_{s_m}(\theta_{t,m-1}),\qquad \theta_{t,m}=\theta_{t,m-1}-\gamma\frac{g_m}{\|g_m\|_2}.

内层轨迹完成后,不把 (\theta_{t,K}) 直接当作模型参数,而是把总位移

g^t=θt,0−θt,K \widehat g_t=\theta_{t,0}-\theta_{t,K}

作为“伪梯度”交给任意外层优化器:

θt=Opt⁡outer(θt−1,g^t). \theta_t=\operatorname{Opt}_{\mathrm{outer}} (\theta_{t-1},\widehat g_t).

内层必须是无动量的原始归一化 SGD,因为理论所需的二阶交互来自“先沿任务 (i) 移动,再在移动后的位置测任务 (j) 的梯度”。外层则可以是 AdamW、Muon 等任意常规优化器。因此 Nexus 更准确地说是一个伪梯度构造器,而不是独立优化器。

可以证明,使用这种更新方式能够优化 $\nabla_\theta\operatorname{CosSim}(g_i,g_j)$.

预训练中的工程实现

理论算法每次外层更新要遍历 (K) 个来源;真实预训练可能有五十多个来源,这会改变有效 batch size,使已有学习率和批量超参数不能直接复用。作者因此把标准梯度累积过程改造成 Nexus 内循环。

实现中同时保留主模型 model 和辅助模型 inner_model。一个梯度累积周期开始时二者参数相同。每个 micro-batch 到来后:

  1. 只在 inner_model 上前向和反向;
  2. 立刻用无动量的归一化 SGD 更新 inner_model;
  3. 主模型在整个累积周期内保持冻结;
  4. 到累积边界时,计算辅助模型相对主模型的位移;
  5. 取与位移相反的方向作为外层伪梯度,交给 AdamW 等外层优化器更新主模型;
  6. 再把更新后的主模型复制给辅助模型,开始下一周期。

如果记辅助模型最终参数为 (\theta_{\mathrm{inner}})、主模型为 (\theta_{\mathrm{main}}),理论中的伪梯度是

g^=θmain−θinner. \widehat g=\theta_{\mathrm{main}}-\theta_{\mathrm{inner}}.

论文伪代码把位移写作 inner_model - model,随后向外层传入其负值,两种写法等价。

工程版不再显式知道每个 micro-batch 属于哪个来源,而是最大化同一梯度累积周期内不同 micro-batch 的梯度相似度。作者依赖一个数据假设:语料极大、单个来源的混合比例较低,因此相邻批次很可能来自不同来源。它于是近似理论算法的跨来源任务采样。

总前向和反向次数与标准梯度累积相同;新增计算主要是辅助参数的更新与同步。内存中原则上要多存一份模型。论文称可借助 CPU offload 和异步处理把额外显存压到接近零,但没有给出这些工程技术的具体实现或吞吐量数据,所以“几乎无额外成本”主要是按前反向次数和理论的 (6NBS) 训练成本近似得出的判断。

在强凸、光滑并且所有任务确实共享同一个极小点的理想区域内,附录还证明了收敛率。若每个 (L_i) 都是 (L)-smooth、(\mu)-strongly convex,步长 (\gamma\in(0,2/(L+\mu))),则

E∥θT−θ∗∥22≤(1−2γμLL+μ)T∥θ0−θ∗∥22. \mathbb E\|\theta_T-\theta^*\|_2^2 \leq \left(1-\frac{2\gamma\mu L}{L+\mu}\right)^T \|\theta_0-\theta^*\|_2^2.

取 (\gamma=2/(L+\mu))、条件数 (\kappa=L/\mu) 时,变成

E∥θT−θ∗∥22≤(κ−1κ+1)2T∥θ0−θ∗∥22. \mathbb E\|\theta_T-\theta^*\|_2^2 \leq \left(\frac{\kappa-1}{\kappa+1}\right)^{2T} \|\theta_0-\theta^*\|_2^2.

这只是解释为什么进入共同极小点区域后可能指数收敛,并不证明一般非凸大模型训练一定进入该区域。

方法的整体图景

把前面的推导合在一起,Nexus 的逻辑链是:混合预训练存在许多总损失相同的解;靠近各来源共同极小点的解在同分布任务模型中有更小的下游泛化差距;极小点距离不能直接计算,但不同来源的梯度余弦不相似度给出了距离上界;顺序执行若干归一化 SGD 内步,会通过在位移后位置重新测量梯度而自然产生 Hessian-gradient 项;这些项在期望上最大化任务间梯度余弦相似度;最后再把内循环总位移交给 AdamW 等成熟外层优化器,从而保留其一阶优化能力,同时加入共同极小点偏置。

实验如何进行

主要实验使用一套内部预训练语料。作者称该语料经过严格去污染,不含评测基准或蒸馏数据,因而比包含泄漏的公开语料更适合观察平滑的优化趋势。模型采用 Llama 架构,规模为不含 embedding 的 130M、300M、520M、1.2B 和 2.3B;按包含 embedding 的总参数量,重点模型称作 1B 和 3B。

基线超参数沿用已有细致调参结果,并额外检查了基准学习率的 (0.5\times) 和 (2\times)。AdamW 与 Nexus 使用完全相同的外层超参数。主要复现参数是:全局 batch size 256,Adam 的 ((\beta_1,\beta_2)=(0.9,0.95)),(\epsilon=10^{-10}),梯度裁剪阈值 1.0,权重衰减 0.2;WSD 调度使用 1000 步 warmup 和 10000 步 decay。1B 模型外层学习率 0.002,3B 模型为 0.001;Nexus 内层学习率均为 (\gamma=0.01)。3B 的 Muon 对照使用权重衰减 0.1。

评测覆盖通识 MMLU,推理 GPQA、GPQA-Diamond、BBH,数学 GSM8K、MATH500,代码 HumanEval、MBPP。除离散准确率外还计算每项任务损失,以及一套严格清洗的内部 OOD 语料损失。作者认为 OOD 损失与下游能力高度相关。论文把损失差大于 0.01 或准确率差大于 2 个百分点视为非平凡差异。

主要结果:训练损失相同,下游不同

1B 模型训练到 (4\times) Chinchilla token 数,AdamW 与 Nexus 的预训练损失都为 1.826,几乎完全相同。Nexus 的 OOD 损失由 1.433 降到 1.428,下游任务平均损失由 1.786 降到 1.758;平均准确率从 21.9% 升到 23.6%。提升较明显的项目包括 GPQA 的 (+5.4) 个百分点和 MBPP 的 (+5.0) 个百分点。

3B 模型训练到 (2\times) Chinchilla token 数,AdamW 与 Nexus 的预训练损失分别为 1.606 和 1.602,差 0.004。与此同时,OOD 损失从 1.302 降到 1.290,八项下游任务的平均损失从 1.633 降到 1.601,平均准确率从 37.1% 升到 40.3%。最显著的是 GSM8K 从 44% 到 59%,MATH500 从 32% 到 40%,HumanEval 从 43% 到 47%。并非每项准确率都提高:GPQA 从 32.8% 降到 29.6%,BBH 和 MBPP 持平;但相应任务损失整体上仍几乎都更低。

输出重叠分析表明,在 GSM8K 和 HumanEval 上,Nexus 保留了 AdamW 已经答对样本中的 95% 以上;净增益主要来自额外解决 AdamW 失败的问题,而不是用一批新正确答案交换掉大量旧正确答案。作者据此认为这种能力扩展较稳定。

Muon 与 Nexus 的对比用来区分“降低训练损失”和“改变隐式偏置”。3B Muon 把预训练损失从 AdamW 的 1.606 降到 1.577,并将平均准确率提高 2.3 个百分点;Nexus 只降低 0.004,却提高 3.2 个百分点,最终下游水平与 Muon 接近。因此 Muon 的收益主要能由更低预训练损失解释,Nexus 的收益则不能。

模型和训练时长扩展

从 130M 到 2.3B,不同规模上 Nexus 与 AdamW 的预训练损失差始终小于 0.01,而平均准确率增益依次为 (+0.8,+1.5,+1.7,+2.6,+3.2) 个百分点。增益随模型规模单调扩大。GSM8K 的差距在 130M 时不明显,到 2.3B 时扩大为 15 个百分点。作者的解释是:模型的过参数化程度越高,越有自由度在不牺牲训练损失的前提下满足额外几何偏置。

把 3B 训练从 (2\times) 延长到 (4\times) Chinchilla token 后,AdamW 平均准确率从 37.1% 升到 38.0%,说明增加训练本身有效;但 Nexus 在相同 (4\times) 设置下达到 41.7%,对 AdamW 的优势没有消失,反而略有扩大。此时两者预训练损失仍几乎相同,分别为 1.591 和 1.588,OOD 损失差仍为 0.012。作者据此否定“只要训练更久,AdamW 自然会到达同样接近的解”这一解释。

数据混合比例

作者构造 Math10、Math40、Math70 三种混合,把数学数据比例固定为 10%、40%、70%,其余来源按比例下采样。动机来自一阶关系

Li(θ)−Li(θ−γ∇Lj(θ))=γ∇Li(θ)⊤∇Lj(θ)+O(γ2). L_i(\theta)-L_i(\theta-\gamma\nabla L_j(\theta)) =\gamma\nabla L_i(\theta)^\top\nabla L_j(\theta)+O(\gamma^2).

若 Nexus 让任务 (i,j) 的梯度更同向,那么在样本充足的任务 (j) 上更新,也会顺带降低样本稀少任务 (i) 的损失。因此它可能像一个不显式改采样权重的动态数据平衡器。

实验与这个解释基本一致。数学只占 10% 时,Nexus 在 GSM8K 上比 AdamW 高 15 个百分点;数学占到 70% 时,数学基线已经追上,Nexus 的优势转向成为少数的通识和推理数据,MMLU 提升 5.8 个百分点,GPQA 提升 3.1 个百分点。代码增益不单调,作者推测代码同时依赖逻辑推理和领域知识,不能由单一数据比例解释。

Nexus 对混合比例变化也更稳定。Math70 变为 Math10 时,AdamW 的 GSM8K 从 77% 降到 44%,下降 33 点;Nexus 从 76% 降到 59%,下降 17 点。AdamW 的 MMLU 跨设置波动 3.8 点,Nexus 的 Math70 与 Math10 结果分别为 49.8% 和 48.9%,波动不足 1 点。需要注意,Math70 下 Nexus 的预训练和 OOD 损失都略差于 AdamW,所以“自动平衡”表现为跨域能力分配更稳,并不保证每个混合下所有总体损失都改善。

学习率调度

默认使用 Warmup-Stable-Decay,作者又替换为标准 cosine annealing,其他设置不变。Cosine 下 AdamW 与 Nexus 的预训练损失为 1.526 和 1.528,Nexus 反而高 0.002,但下游平均损失仍由 1.572 降到 1.541,平均准确率从 41.0% 升到 43.4%。因此现象不依赖某一种学习率轨迹。

理论机制的实验核对

作者在 3B 训练过程中每 1000 步记录一次预训练测试集与各下游语料梯度的余弦相似度并取时间平均。训练结束后,又在每个下游任务上用学习率 (2\times10^{-5})、权重衰减 0 的 AdamW 做 full-batch 优化,以找到近似任务极小点 (\theta_T^*),再测参数距离。

训练来源之间的平均梯度相似度从 AdamW 的 0.4499 升到 Nexus 的 0.4661,说明 Nexus 确实完成了它直接优化的目标。预训练目标与 OOD 语料的梯度相似度从 0.2228 升到 0.2464;GPQA-Diamond、MATH500、HumanEval、MBPP 也略升。不过并不是所有任务都升,例如 GSM8K 从 0.0374 降到 0.0325。因此“训练梯度接近性泛化到所有下游任务”是整体趋势,不是逐任务无例外的事实。

参数距离同样总体减小。训练集近似极小点距离从 1.452 降到 1.441;OOD 从 2.812 降到 2.806;GSM8K 从 3.418 明显降到 3.326;GPQA-Diamond、MATH500、HumanEval 也下降。MBPP 则从 3.645 微增到 3.648。与这些总体变化同时,下游损失大多降低,准确率整体提高。损失地形插值图还显示,Nexus 的预训练解在 GSM8K 和 HumanEval 的任务特定极小点方向上更靠近低损失区域。

这些极小点并没有像示意图中的理想“完全交集”那样重合,距离只得到中等幅度缩短;下游损失也相应是中等幅度下降,而不是趋近于零。作者把更强但仍低成本的接近性优化留作未来方向。

消融揭示了收益来自哪里

把不同 checkpoint 的平均下游损失画成预训练损失的函数后,AdamW 和纯 Muon 的曲线几乎重合:在预训练损失相同处,它们的下游损失也相同。Muon 的权重正交化可能产生更高秩表示,但这项结构变化没有表现出独立于训练损失的额外泛化收益。

直接把归一化梯度送入 Adam 的 NSGD 变体也与 AdamW 曲线接近,说明单纯归一化不是 Nexus 收益来源。这个变体等价于 Nexus 内循环只有一步,即 (K=1)。此时二阶相似度项系数

K−14K \frac{K-1}{4K}

严格为零,没有任何跨任务交互;这与实验结果吻合。真正关键的是至少两个顺序内步之间的 Hessian-gradient 交互。

虽然梯度点积与下游一步改进关系更直接,并且在接近性定理中给出更紧的中间界,实践中却不能直接最大化点积。若整体损失放大 (k) 倍,梯度范数放大 (k) 倍,点积会人为放大 (k^2) 倍。优化器可通过增大梯度范数而不是真正让方向一致来走捷径。Nexus-Dot 的训练曲线因此明显落后于 AdamW,训练损失恶化超过潜在泛化收益;这种减速在不同调度器和 AdamW、Muon 外层上都存在。余弦相似度去除了尺度,因而是最终采用的目标。

补充实验提供的新增信息

在一套未严格去污染的公开预训练数据上,作者训练 1B 和 3B 模型。Nexus 的预训练损失都比 Adam 略高约 0.007--0.008,但 OOD 损失分别降低 0.028 和 0.017,平均下游损失也分别降低 0.027 和 0.026。最突出的变化在稀缺代码任务:1B 的 MBPP 从 1% 到 15%,3B 从 23% 到 38%。作者认为,面对可能含有 GSM8K 等基准泄漏的噪声数据,Adam 更容易把容量用于记忆捷径,Nexus 的跨任务共识偏置会把一部分容量重新分配到未污染、样本稀少的域。不过这是作者根据任务分布和结果作出的解释,实验没有直接定位或删除具体泄漏样本来验证因果机制。

作者还从 100000 步 checkpoint 继续做监督微调,以检查 Nexus 是否只是提前“用掉”了 SFT 本来能带来的提升。SFT 学习率为 (2\times10^{-5}),全局 batch size 256,与该预训练时点的学习率和 batch size 对齐。微调后,Nexus 的平均准确率为 42.7%,AdamW 为 40.7%;MATH、HumanEval、MBPP 分别高 7、6、6 个百分点。SFT 损失从训练前的 1.647 对 1.655,一直保持到训练后的 1.028 对 1.035,均是 Nexus 更低。因此它没有损害后续对齐容量,优势能延续到继续训练阶段。

适用边界与尚未解决的问题

理论结论描述的是一个受控机制,而不是对任意大模型损失的无条件保证。它要求训练与下游任务来自同一潜在分布、关注方向上损失局部近似二次或方向强凸、曲率有上下界,并在一般定理中要求任务曲率和任务极小点位置统计独立。真实 LLM 是否严格满足这些条件没有被证明,实验只能说明它们给出的预测与观察到的总体趋势一致。

Nexus 当前与 Muon 的组合并不成功。虽然理论上外层优化器可以任意选择,实际 Muon+Nexus 的下游表现甚至差于 AdamW+Nexus。作者观察到 Nexus 在 AdamW 上略有加速,却会拖慢 Muon,推测原因可能是伪梯度系数 (\gamma) 的数值敏感性,或与 Muon 内部 Newton--Schulz 正交化迭代发生了相互作用,但论文没有给出确定诊断。

此外,工程版实际优化的是同一累积周期内 micro-batch 的相似度,而不是显式按数据来源分组的相似度;其有效性依赖相邻批次经常来自不同来源。完整辅助模型带来的内存、同步和分布式通信代价也只作了理论估计,没有报告大规模训练吞吐量。因而论文已经证明并验证了“顺序归一化内步产生梯度对齐偏置”,但把它无摩擦地扩展到任意训练系统,仍需要额外工程验证。