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

📄 arXiv: 2604.09258v1 📥 PDF

作者: Huanran Chen, Huaqing Zhang, Xiao Li, Yinpeng Dong, Ke Shen, Jun Zhu

分类: cs.LG

发布日期: 2026-04-10

💡 一句话要点

Nexus优化器：通过寻找共同极小值提升大语言模型下游泛化能力

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 预训练 优化器 泛化能力 共同极小值

📋 核心要点

  1. 现有大语言模型预训练后，不同数据源对应的损失极小值分散，影响下游任务的泛化性能。
  2. Nexus优化器通过最大化训练过程中梯度的相似性，鼓励模型收敛到不同数据源损失的共同极小值。
  3. 实验表明，Nexus在保持预训练损失不变的情况下，显著提升了下游任务的性能，尤其是在复杂推理任务上。

📝 摘要（中文）

预训练是大语言模型（LLMs）的基石，消耗了绝大部分计算资源和数据，是模型能力的主要引擎。本文研究了预训练收敛状态的一个有趣的几何问题：模型是否收敛到所有数据源的共同极小值，还是仅仅是总损失的极小值？我们假设任务特定极小值的几何“接近程度”与下游泛化能力内在相关。研究表明，标准优化器（如AdamW）通常收敛到任务特定极小值彼此远离的点。为此，我们提出了Nexus优化器，通过最大化优化过程中的梯度相似性来鼓励这些极小值的接近。在1.3亿到30亿参数的模型、各种数据混合和超参数调度上的实验表明，Nexus在达到相同预训练损失的情况下，显著提高了下游性能。值得注意的是，在30亿参数的模型上，Nexus将超出分布损失降低了0.012，并在复杂推理任务（如GSM8k）上产生了高达15.0%的准确率提升。这一发现挑战了将预训练损失作为模型评估的唯一代理的观点，并证明了隐式偏差在解锁下游泛化中的重要性。

🔬 方法详解

问题定义：现有的大语言模型预训练方法，通常使用AdamW等优化器，虽然可以降低预训练损失，但模型最终收敛到的解，其不同数据源（例如，不同类型的文本、代码、数学题等）对应的损失极小值在参数空间中是分散的。这种分散性导致模型在下游任务上的泛化能力受限，尤其是在分布外数据和复杂推理任务上表现不佳。现有方法仅仅关注降低预训练损失，而忽略了模型收敛解的几何特性，即不同数据源损失极小值之间的关系。

核心思路：论文的核心思路是，如果模型能够收敛到一个所有数据源损失的“共同极小值”（Common Minima），即不同数据源对应的损失函数在该点都接近最小值，那么模型就能更好地泛化到下游任务。为了实现这一目标，论文提出了Nexus优化器，其核心思想是在优化过程中，鼓励不同数据源的梯度方向尽可能一致，从而引导模型向共同极小值收敛。

技术框架：Nexus优化器的整体框架与标准的优化器（如AdamW）类似，都是在每个训练迭代中，根据梯度更新模型参数。不同之处在于，Nexus在计算梯度更新时，会考虑不同数据源梯度之间的相似性。具体来说，对于每个训练批次，Nexus首先计算每个数据源的梯度，然后计算这些梯度之间的余弦相似度，并将其作为优化目标的一部分。通过最大化梯度相似度，Nexus鼓励模型学习到对所有数据源都通用的特征表示。

关键创新：Nexus优化器的关键创新在于，它将梯度相似性引入到优化过程中，从而改变了模型的隐式偏差。传统的优化器只关注降低预训练损失，而Nexus则同时关注降低预训练损失和提高梯度相似性。这种新的优化目标使得模型更容易收敛到共同极小值，从而提升下游泛化能力。与现有方法相比，Nexus的本质区别在于，它不仅仅关注损失函数的值，还关注损失函数的几何形状。

关键设计：Nexus优化器的关键设计在于梯度相似度损失函数的选择和超参数的设置。论文中使用余弦相似度作为梯度相似度的度量，并将其乘以一个超参数λ，作为正则化项添加到总损失函数中。λ控制了梯度相似度对优化过程的影响程度。论文通过实验发现，合适的λ值可以显著提升下游性能。此外，论文还探索了不同的超参数调度策略，例如，逐渐增加λ的值，以在训练初期允许模型探索更广阔的参数空间，而在训练后期鼓励模型收敛到共同极小值。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Nexus优化器在1.3亿到30亿参数的模型上，显著提高了下游性能，同时保持了预训练损失不变。在30亿参数的模型上，Nexus将超出分布损失降低了0.012，并在GSM8k等复杂推理任务上产生了高达15.0%的准确率提升。这些结果表明，Nexus优化器能够有效地提升大语言模型的泛化能力，并挑战了将预训练损失作为模型评估的唯一标准的观点。

🎯 应用场景

Nexus优化器可应用于各种大语言模型的预训练，尤其是在需要模型具备良好下游泛化能力的场景中。例如，在医疗、金融等领域，模型需要处理来自不同数据源的复杂信息，并进行准确的推理和预测。Nexus优化器可以帮助模型更好地适应这些场景，提高模型的可靠性和实用性。此外，Nexus的思路也可以推广到其他机器学习任务中，例如，多任务学习、领域自适应等。

📄 摘要（原文）

Pretraining is the cornerstone of Large Language Models (LLMs), dominating the vast majority of computational budget and data to serve as the primary engine for their capabilities. During pretraining, LLMs acquire foundational knowledge from an unprecedentedly massive and diverse data sources, encompassing a vast array of domains such as general language, mathematics, code, and complex reasoning. In this work, we investigate an interesting geometric question regarding the converged state of pretraining: Does the model converge to a common minimizer across all data sources (e.g., \cref{fig:cwa_illustration:close}), or merely a minimizer of the summed loss (e.g., \cref{fig:cwa_illustration:distant})? We hypothesize that the geometric "closeness" of task-specific minima is intrinsically linked to downstream generalization. We reveal that standard optimizers (e.g., AdamW) often converge to points where task-specific minima are distant from each other. To address this, we propose the Nexus optimizer, which encourages the closeness of these minima by maximizing gradient similarity during optimization. Experiments across models ranging from 130M to 3B parameters, various data mixtures and hyperparameter schedules, show that Nexus \textit{significantly boosts downstream performance}, despite \textit{achieving the same pretraining loss} (see \cref{fig:demo:benchmark}). Notably, on the 3B model, Nexus reduces the out-of-distribution loss by 0.012 and yields up to a 15.0\% accuracy improvement on complex reasoning tasks (e.g., GSM8k). This finding challenges the reliance on pretraining loss as the sole proxy for model evaluation and demonstrates the importance of implicit biases in unlocking downstream generalization.