Nexus: Same Pretraining Loss, Better Downstream Generalization via Common Minima¶

会议: NeurIPS 2026
arXiv: 2604.09258
OpenReview: https://openreview.net/forum?id=gEYdmlnR5G
领域: 预训练
关键词: 预训练优化器、隐式偏置、公共极小值、梯度相似度、下游泛化

一句话总结¶

针对传统优化器收敛于仅最小化加权总损失的“平均型极小值”而损害跨任务泛化的问题，本文提出即插即用的 Nexus 二阶梯度近似器，通过最大化不同数据源间的梯度余弦相似度将模型引导至各任务极小值紧密聚集的“公共极小值”，在预训练损失几乎一致的前提下大幅提升下游泛化性能。

研究背景与动机¶

在大语言模型（LLM）的预训练阶段，模型需要从涵盖自然语言、代码、数学与复杂推理等多源异构语料中汲取通用能力。行业标准实践通常是将各数据源的加权经验损失进行算术平均并求和最小化。在此范式下，从业者普遍将预训练验证集损失（Pretraining Loss）视作评价模型能力与优化进展的唯一代理指标，默认更低的训练损失必然对应更优秀的下游能力。

然而，从高维损失曲面的几何视角审视，相同的预训练总损失可能对应着拓扑性质截然不同的收敛状态。一方面，优化轨迹可能落入各数据源极小值彼此分散的“平均型极小值”（Average-type Minimizer）——此时加权总损失虽然达到局部极小，但当前参数距离任意单个数据源的真实最优解均存在较大欧氏距离；另一方面，优化轨迹也可能收敛于各数据源极小值紧密聚集的“公共极小值”（Intersection-type Minimizer）。理论推导表明，当预训练与下游任务服从相同任务分布且局部具有强凸性时，下游泛化误差与各任务极小值之间的几何紧密度（Closeness）直接相关，极小值距离越远，未见下游任务的期望泛化误差就越高。标准优化器（如 AdamW、Muon）缺乏促进极小值对齐的隐式偏置，导致下游表现受限。

解决该问题的直接瓶颈在于显式最小化任务极小值间欧氏距离在计算上是不可行的，因为在预训练的每一个迭代步追踪各个子任务的真实最优点需要极为昂贵的嵌套优化。核心 idea：本文证明了多任务间的梯度余弦相似度是参数空间几何紧密度的紧致上界，进而提出 Nexus 双循环优化框架，通过在微步积累阶段使用无动量归一化 SGD 进行内循环前向扰动计算伪梯度，在不增加反向传播开销的前提下显式最大化源间梯度相似度，从而在保持相同预训练损失的同时显著提升下游下游任务性能。

方法详解¶

整体框架¶

Nexus 的核心思想是通过二阶近似将各数据源之间的梯度余弦相似度作为正则项融入优化轨迹，使优化方向兼顾预训练经验损失下降与源间梯度方向一致性。整体流水线被设计为双层循环（Dual-loop）结构：内循环作为无显式 Hessian 计算的二阶梯度近似器（Gradient Approximator），通过序列化的归一化梯度下降（Normalized SGD, NSGD）探测损失曲面的曲率变化并积累源间交互；外循环则作为标准的优化器载体（如 AdamW），消费内循环生成的伪梯度进行全局模型参数更新。工程实现上，Nexus 巧妙复用了现代分布式预训练中的梯度累积（Gradient Accumulation）边界，引入辅助镜像权重实现无额外反向传播的高效计算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：多源异构预训练 Mini-batch 流"] --> B["1. 几何紧密度与二阶泛化偏置分析"]
    B --> C["2. 梯度相似度对偶上界推导"]
    C --> D["3. Nexus 双循环梯度近似机制"]
    D --> E["4. 分布式梯度累积工程适配"]
    E --> F["输出：消费伪梯度更新主模型，收敛至公共极小值"]

关键设计¶

1. 几何紧密度与二阶泛化偏置分析：揭示同损失下泛化差异的几何根源

主流泛化理论主要关注局部极小值的平坦度（Flatness，由 Hessian 最大特征值或迹表征），然而在多任务预训练环境下，平坦度无法解释相同训练损失下截然不同的下游表现。本文通过对未见下游任务 $\mathcal{T}$ 在最优点 $\boldsymbol{\theta}^*_{\mathcal{T}}$ 处进行带拉格朗日余项的二阶泰勒展开，将期望下游损失分解为内生损失、方向锐度（Directional Sharpness）与几何紧密度（Closeness）的乘积： $$\mathcal{L}_{\mathcal{T}}(\boldsymbol{\theta}_{\text{train}}^*) = \mathcal{L}_{\mathcal{T}}(\boldsymbol{\theta}_{\mathcal{T}}^*) + \frac{1}{2} \|\boldsymbol{\theta}_{\text{train}}^* - \boldsymbol{\theta}_{\mathcal{T}}^*\|_2^2 \cdot \boldsymbol{u}^\top \nabla^2 \mathcal{L}_{\mathcal{T}}(\boldsymbol{\xi}) \boldsymbol{u}$$ 在局部强凸（$\lambda_{\min} \mathbf{I} \preceq \nabla^2 \mathcal{L} \preceq \lambda_{\max} \mathbf{I}$）且任务极小值服从方差为 $\sigma^2$ 的同分布假设下，定理证明了在固定预训练损失 $C_{\text{train}}$ 的约束下，下游期望泛化误差上界严格受控于极小值分布方差： $$\mathbb{E}_{\mathcal{T} \sim \mathcal{P}}[\mathcal{L}_{\mathcal{T}}(\boldsymbol{\theta}_{\text{train}}^*)] - C_{\text{train}} \le \frac{\lambda_{\max}\left(\left(\frac{\lambda_{\max}}{\lambda_{\min}}\right)^2 + 1\right)}{2K}\sigma^2$$ 这一理论从本质上确立了几何紧密度 $\sigma^2$ 是与平坦度地位等同且相互解耦的二阶泛化隐式偏置：通过以轻微提高各任务本征极小值深度为代价换取更紧凑的极小值聚合，可在保持宏观训练损失不变的情况下实质性降低下游泛化间隙。

2. 梯度相似度对偶上界推导：将不可求解的极值距离转化为梯度对齐

由于显式追踪各任务最优解 $\boldsymbol{\theta}_k^*$ 属于 NP-hard 级的嵌套极值问题，本文建立了梯度余弦相似度与极小值欧氏距离之间的对偶不等式。由柯西-施瓦茨不等式及强凸性下界条件可得： $$\frac{1}{K}\sum_{k=1}^K \|\boldsymbol{\theta} - \boldsymbol{\theta}_k^*\|_2^2 \le \frac{1}{K\lambda_{\min}^2}\sum_{i \neq j} \left(-\nabla \mathcal{L}_i(\boldsymbol{\theta})^\top \nabla \mathcal{L}_j(\boldsymbol{\theta})\right) \le \frac{G^2}{K\lambda_{\min}^2}\sum_{i \neq j}\left(1 - \text{CosSim}(\nabla \mathcal{L}_i(\boldsymbol{\theta}), \nabla \mathcal{L}_j(\boldsymbol{\theta}))\right)$$ 该上界表明，最大化任务梯度的两两余弦相似度等价于最小化极小值之间的离散程度。同时从一阶动力学来看，单步梯度下降在任务 $j$ 上对任务 $i$ 带来的损失下降一阶近似为 $\gamma \nabla \mathcal{L}_i(\boldsymbol{\theta})^\top \nabla \mathcal{L}_j(\boldsymbol{\theta}) + \mathcal{O}(\gamma^2)$。因此，高梯度相似度保证了在任一数据源上的梯度更新均能同向赋能其他源及下游任务，避免负迁移振荡。

3. Nexus 双循环梯度近似机制：无显式 Hessian 的二阶梯度重构

直接对余弦相似度求导会引入昂贵的雅可比-向量积（Jacobian-vector product, JVP）： $$\nabla_{\boldsymbol{\theta}} \text{CosSim}(\nabla \mathcal{L}_i, \nabla \mathcal{L}_j) = \left(\nabla_{\boldsymbol{\theta}} \frac{\nabla \mathcal{L}_i}{\|\nabla \mathcal{L}_i\|_2}\right)^\top \frac{\nabla \mathcal{L}_j}{\|\nabla \mathcal{L}_j\|_2} + \left(\nabla_{\boldsymbol{\theta}} \frac{\nabla \mathcal{L}_j}{\|\nabla \mathcal{L}_j\|_2}\right)^\top \frac{\nabla \mathcal{L}_i}{\|\nabla \mathcal{L}_i\|_2}$$ Nexus 通过设计内循环序列更新来隐式构造这一二阶导数。在每个外循环步中，内循环从当前位置 $\boldsymbol{\theta}_{t,0}$ 出发，依次采样任务并使用无动量的归一化梯度下降迭代 $K$ 步：$\boldsymbol{\theta}_{t,m} \leftarrow \boldsymbol{\theta}_{t,m-1} - \gamma \frac{\boldsymbol{g}}{\|\boldsymbol{g}\|_2}$。内循环结束后的累积位移差量定义为伪梯度 $\hat{\boldsymbol{g}}_t = \boldsymbol{\theta}_{t,0} - \boldsymbol{\theta}_{t,K}$。理论证明，在期望意义下，该伪梯度精确逼近一阶平均梯度与余弦相似度梯度的加权组合： $$\mathbb{E}[\hat{\boldsymbol{g}}_t] = \gamma \sum_{i=1}^K \frac{\nabla \mathcal{L}_i(\boldsymbol{\theta}_t)}{\|\nabla \mathcal{L}_i(\boldsymbol{\theta}_t)\|_2} - \gamma^2 \frac{K-1}{4K}\left(\nabla_{\boldsymbol{\theta}}\sum_{i\neq j}\text{CosSim}(\nabla \mathcal{L}_i, \nabla \mathcal{L}_j) + \boldsymbol{\mathcal{E}}_{\text{sym},i,j}\right) + \mathcal{O}(\gamma^3)$$ 其中非对称残差 $\boldsymbol{\mathcal{E}}_{\text{sym},i,j}$ 正交于目标梯度，且在神经网络沿 Hessian 主特征向量对齐的稳定边缘（Edge of Stability）机制下自然趋近于零。这使得 Nexus 在数学上等效于以 $\gamma^2 \frac{K-1}{4K}$ 为惩罚强度隐式最大化梯度相似度。

4. 分布式梯度累积工程适配：零通信开销的生产级即插即用部署

将理想双循环算法直接拓展到大规模分布式预训练时，若每步外更新都强求遍历所有 $K$ 个数据源，会导致有效批量剧增并破坏学习率等既定超参体系。为此，Nexus 提出了与梯度累积无缝集成的工程方案：在累积步内保留一份镜像权重副本 inner_model；在每个 mini-batch 步骤，前向与反向传播直接在 inner_model 上执行，随后调用 per-parameter 的 NSGD 立即更新 inner_model；到达累积边界时，计算主模型与内模型的位移差 $\hat{\boldsymbol{g}} = \text{inner\_model} - \text{model}$ 作为伪梯度传给外层 AdamW 优化器，最后将主模型同步回 inner_model。该机制巧妙借助异构语料在微批次间的自然交替，以极低显存开销（可通过 CPU Offloading 完全消除）与完全零额外反向传播算力消耗实现了对批次间梯度余弦相似度的实时正则化。

损失函数 / 训练策略¶

在外循环全局视角下，Nexus 隐式最小化的二阶复合优化目标函数为： $$\mathcal{J}_{\text{2nd}}(\boldsymbol{\theta}) = \gamma \sum_{i=1}^K \mathcal{L}_i(\boldsymbol{\theta}) - \gamma^2 \frac{K-1}{4K} \sum_{i \neq j} \text{CosSim}\left(\nabla \mathcal{L}_i(\boldsymbol{\theta}), \nabla \mathcal{L}_j(\boldsymbol{\theta})\right)$$ 内循环采用逐层（Per-parameter/Per-matrix）而非全局拉平向量的 L2 归一化，与 $\mu\text{P}$ 理论及 RMSNorm 保持一致，从而确保了超参数跨尺度的稳定迁移。外循环兼容任意基优化器（实验主要采用 AdamW），且严格复用其最优学习率和权重衰减配置。

实验关键数据¶

主实验¶

在严格清洗脱敏的预训练语料上，针对 1B（4× Chinchilla tokens）和 3B（2× Chinchilla tokens）Llama 架构模型，对比标准 AdamW 基线与引入 Nexus 的优化表现。在预训练验证损失几乎完全一致的约束下，下游各项基准指标显著提升：

模型规模	优化器	预训练损失 (↓)	OOD损失 (↓)	MMLU Acc (↑)	GSM8k Acc (↑)	MATH Acc (↑)	HumanEval (↑)	全任务均分 (↑)
1B	AdamW	1.826	1.433	32.1%	18.0%	13.0%	19.0%	21.9%
1B	Nexus	1.826	1.428	33.5%	20.0%	13.0%	19.0%	23.6% (+1.7%)
3B	AdamW	1.606	1.302	47.8%	44.0%	32.0%	43.0%	37.1%
3B	Nexus	1.602	1.290	48.9%	59.0%	40.0%	47.0%	40.3% (+3.2%)

消融与扩展实验¶

在 3B 模型规模下，对优化器隐式偏置机制、训练计算量扩展以及相似度度量方式进行系统性消融验证：

实验维度 / 配置	预训练损失 (↓)	OOD损失 (↓)	GSM8k Acc (↑)	MATH Acc (↑)	机制说明与结论
AdamW 基线 (2× Chinchilla)	1.606	1.302	44.0%	32.0%	标准一阶累积，缺乏极小值几何约束
Muon 优化器	1.577	1.288	51.0%	36.0%	依赖降低训练损失提升泛化，无同损失偏置优势
NSGD 纯归一化 ($K=1$)	1.605	1.301	44.5%	32.2%	正则强度项 $\frac{K-1}{4K}=0$，证明收益非来自梯度归一化
Nexus-Dot (未归一化点积)	1.684	1.385	31.0%	21.0%	模长尺度敏感导致病态优化捷径，训练严重恶化
Nexus (4× Chinchilla 充分训练)	1.588	1.281	62.0%	33.0%	较 4× AdamW (54.0%) 优势不降反升，算力充裕下不可或缺

关键发现¶

泛化增益随模型尺度单调放大：从 130M、300M、520M、1.2B 到 2.3B（即 3B 总参数），Nexus 相对 AdamW 的平均下游准确率提升分别为 +0.8%、+1.5%、+1.7%、+2.6% 和 +3.2%。高度过参数化赋予了网络在不破坏经验损失的同时收敛至几何紧凑公共极小值的结构自由度。
能力边界扩展而非灾难遗忘：对 3B 模型输出的微观分析表明，Nexus 在 GSM8k 与 HumanEval 上对 AdamW 已正确解出的题目的保留率超过 95%，所带来的 +15.0% 净增益完全来源于攻克 AdamW 此前答错的复杂推理边界样本。
隐式数据混合自平衡效应：在将数学数据采样比例由 10% 调至 70% 的偏斜配比实验中，Nexus 在样本稀缺域上的相对增益最为显著（通用域提升随通用数据占比降低由 +1.1% 扩大至 +5.8%），证明梯度相似度优化促使密集域的梯度有效反哺稀疏域，充当了动态自适应重加权机制。

亮点与洞察¶

颠覆“预训练损失唯一论”的经验教条：严格从理论与实验双重视角揭示了“相同预训练损失、更强下游泛化”的客观存在。在预训练范式从算力受限（Compute-bound）向高质量数据受限（Data-bound）演进的背景下，调节优化器的隐式几何偏置成为了压榨模型能力的新维度。
极具工程美感的零通信开销实现：避开了传统元学习与二阶优化需计算 Hessian 矩阵或引入额外反向传播的致命缺陷，通过巧妙对齐微批次梯度累积流程，以极其轻量的内模型扰动完成了高维梯度余弦相似度的二阶解析。

局限与展望¶

Muon 优化器集成遇阻：作者承认当前 Nexus 机制与正交化优化器 Muon 存在不兼容现象，组合使用时出现训练降速，下游表现不及 AdamW-Nexus，推测原因涉及 Newton-Schulz 迭代与伪梯度系数的数值冲突，仍需进一步理论解析。
局部极小值重合度仍有提升空间：损失曲面距离度量显示，Nexus 目前仅实现了极小值几何间距的温和收缩，尚未真正达成完全重合的理想公共解，未来探索更强且保持零额外计算开销的收敛加速器具有重大价值。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统形式化了预训练多任务公共极小值的几何紧密度概念，并推导出余弦相似度上界与二阶近似算法。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 130M 至 3B 规模、2×至 4× Chinchilla 计算量、多样数据配比、学习率调度器及严谨微观消融。
写作质量: ⭐⭐⭐⭐⭐ 数学论证环环相扣，从二阶泰勒展开、几何上界到算法构造逻辑极为严密。
价值: ⭐⭐⭐⭐⭐ 对 LLM 预训练评估体系和下一代优化器设计具有深刻的指导意义，工程落地门槛极低。