Nexus: Same Pretraining Loss, Better Downstream Generalization via Common Minima
2026-04-10 • Machine Learning
Machine Learning
AI summaryⓘ
The authors studied how large language models learn from different types of data during pretraining. They found that standard training methods often end up with solutions where the best points for different tasks are far apart, which may hurt general performance. To fix this, they created a new training method called Nexus that encourages shared learning by making gradients for different tasks more similar. Their experiments showed that Nexus improves performance on difficult reasoning tasks even though the overall training loss is the same. This suggests that looking only at training loss doesn’t fully capture how well models learn.
Large Language ModelsPretrainingOptimizationAdamWGradient SimilarityDownstream GeneralizationMulti-task LearningLoss MinimizationOut-of-distribution PerformanceReasoning Tasks
Authors
Huanran Chen, Huaqing Zhang, Xiao Li, Yinpeng Dong, Ke Shen, Jun Zhu
Abstract
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.