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

Huanran Chen; Huaqing Zhang; Xiao Li; Yinpeng Dong; Ke Shen; Jun Zhu

cs.LGApr 10, 2026

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

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

Organizations: Tsinghua University · ByteDance Seed

Abstract

The foundational capabilities of large language models are acquired during pretraining on internet-scale, highly heterogeneous data mixtures. 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 significantly boosts downstream performance, despite 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.

Figures & tables

Panel of Figure 1 — Figure 1: Illustration of “same pretraining loss, better downstream task”. The training loss of baseline and our Nexus are nearly the same. However, our methods achieves better downstream generalization.

Panel of Figure 2 — Figure 2: Illustration of two types of minimizer. (a) Distant: Minimizers of each source are distant from each other. (b) Close: Minimizers are geometrically close to each other. Although both configurations achieve the same total training loss , they perform significantly differently on a new downstream task $\mathcal{L}_{\mathcal{T}}$ .

Figure 3: Intuitive illustration of Nexus Algorithm. $\overset{\mathbb{E}}{=}$ denotes equality in expectation over the task permutations.

Figure 4: Comparison of Standard Pretraining and Nexus Engineering Adaptation. Left: Standard training accumulates gradients over multiple mini-batches before performing a single optimizer update (at the micro-batch/accumulation step). Right: We adapt Nexus from Algorithm ˜ 1 to pretraining by keeping an auxiliary inner_model . It performs immediate updates on the inner_model at every mini-batch step. At the accumulation boundary, the total displacement (

\texttt{inner\_model}-\texttt{model}

) serves as the pseudo-gradient

\hat{\bm{g}}

for the outer optimizer, after which the inner model is re-synchronized.

Model	Optim.	Metric	Loss Metrics ( $\downarrow$ )		Gen.	Reasoning			Math		Code		Avg.
Model	Optim.	Metric	Pretrain.	OOD	MMLU	GPQA	GPQA-D	BBH	GSM8k	MATH	HumanEval	MBPP	All
1B	AdamW	Acc. ( $\uparrow$ )	1.826	1.433	32.1	25.0	21.8	29.6	18.0	13.0	19.0	17.0	21.9
	AdamW	Loss ( $\downarrow$ )	1.826	1.433	2.363	2.221	2.124	1.640	1.429	1.204	1.270	2.035	1.786
	Nexus	Acc. ( $\uparrow$ )	1.826	1.428	33.5	30.4	21.8	29.3	20.0	13.0	19.0	22.0	23.6
	Nexus	Loss ( $\downarrow$ )	1.826	1.428	2.316	2.201	2.102	1.638	1.396	1.176	1.261	1.977	1.758
	Improv.	Acc. ( $\uparrow$ )	-	-	+1.4	+5.4	0.0	-0.3	+2.0	0.0	0.0	+5.0	+1.7

Table 1: Main Results. Comparison of validation losses and downstream capabilities. Notably, Nexus consistently achieves nearly identical pretraining losses compared to the base optimizers, yet demonstrates superior performance across downstream losses and benchmarks .

Panel of Figure 5 — Figure 5: Benchmark Performance across Model Scales. We compare downstream capabilities for models ranging from 130M to 2.3B parameters. Notably, the relative gains of Nexus amplify as model capacity increases, with the average benchmark accuracy improvement growing from +0.8% on the 130M model to +3.2% on the 2.3B model.

Chinchila	Optim.	Metric	Loss Metrics ( $\downarrow$ )		Gen.	Reasoning			Math		Code		Avg.
Chinchila	Optim.	Metric	Pretrain.	OOD	MMLU	GPQA	GPQA-D	BBH	GSM8k	MATH	HumanEval	MBPP	All
2	AdamW	Acc. ( $\uparrow$ )	1.606	1.302	47.8	32.8	22.6	36.6	44.0	32.0	43.0	38.0	37.1
	AdamW	Loss ( $\downarrow$ )	1.606	1.302	2.265	2.005	1.910	1.534	1.259	1.054	1.116	1.922	1.633
	Nexus	Acc. ( $\uparrow$ )	1.602	1.290	48.9	29.6	23.4	36.6	59.0	40.0	47.0	38.0	40.3
	Nexus	Loss ( $\downarrow$ )	1.602	1.290	2.179	1.981	1.881	1.504	1.227	1.026	1.086	1.921	1.601
	Improv.	Loss ( $\uparrow$ )	+0.004	+0.012	+0.086	+0.024	+0.029	+0.030	+0.032	+0.028	+0.030	+0.001	+0.032

Table 2: Scaling Analysis on Training Tokens. We extend the pretraining duration of the 3B model from 2

\times

to 4

\times

Chinchilla optimal tokens. The results demonstrate that the downstream performance advantage of Nexus over the AdamW baseline persists strictly.

Panel of Figure 6 — Figure 6: Results on varying data mixtures (3B models). As the proportion of math data increases (10% $\to$ 70%), the relative performance gains of Nexus on math benchmarks gradually diminish, whereas its advantages on General domain progressively expand. This suggests Nexus boosts the sample-sparse or harder-to-learn domains in the mixture.

Schedule	Optim.	Metric	Loss Metrics ( $\downarrow$ )		Gen.	Reasoning			Math		Code		Avg.
Schedule	Optim.	Metric	Eval	OOD	MMLU	GPQA	GPQA-D	BBH	GSM8k	MATH	HumanEval	MBPP	All
WSD	AdamW	Acc. ( $\uparrow$ )	1.606	1.302	47.8	32.8	22.6	36.6	44.0	32.0	43.0	38.0	37.1
	AdamW	Loss ( $\downarrow$ )	1.606	1.302	2.265	2.005	1.910	1.534	1.259	1.054	1.116	1.922	1.633
	Nexus	Acc. ( $\uparrow$ )	1.602	1.290	48.9	29.6	23.4	36.6	59.0	40.0	47.0	38.0	40.3
	Nexus	Loss ( $\downarrow$ )	1.602	1.290	2.179	1.981	1.881	1.504	1.227	1.026	1.086	1.921	1.601
	Improv.	Loss ( $\uparrow$ )	+0.004	+0.012	+0.086	+0.024	+0.029	+0.030	+0.032	+0.028	+0.030	+0.001	+0.032

Table 3: Results under different learning rate schedulers. We evaluate the 3B model trained with AdamW and Nexus using both WSD and Cosine schedulers. The results demonstrate that the “same pretraining loss, better downstream performance” phenomenon is highly robust regardless of the scheduler.

Metric	Optim.	Pretrain Set	OOD Set	GPQA-D	GSM8k	Math500	HumanEval	MBPP
Grad Sim. ( $\uparrow$ )	AdamW	0.4499	0.2228	0.0824	0.0374	0.0422	0.0367	0.0091
Grad Sim. ( $\uparrow$ )	Nexus	0.4661	0.2464	0.0924	0.0325	0.0427	0.0382	0.0092
Param. Closeness. ( $\downarrow$ )	AdamW	1.452	2.812	4.500	3.418	3.775	4.472	3.645
Param. Closeness. ( $\downarrow$ )	Nexus	1.441	2.806	4.482	3.326	3.766	4.444	3.648
Loss ( $\downarrow$ )	AdamW	1.606	1.302	1.910	1.259	1.054	1.116	1.922
Loss ( $\downarrow$ )	Nexus	1.602	1.290	1.881	1.227	1.026	1.086	1.921

Table 4: Analysis of Gradient Similarity, Loss, and Benchmarks. By optimizing gradient similarity within the pretraining corpus, Nexus achieves higher gradient similarity between the pretraining corpus and downstream corpus. Consistent with Eq. ˜ 8 , this first order gradient similarity directly translates into lower zero-th order downstream losses, ultimately yielding better benchmark performance.

Panel of Figure 7 — Figure 7: Ablation Studies. (a) Implicit biases of various optimizers, illustrated by the correlation between pretraining and downstream losses. (b) Pretraining trajectory of Nexus-Dot, demonstrating that optimizing the unnormalized dot product disrupts pretraining loss minimization. (c) Loss landscape visualization of Adam and Nexus.

Data and Loss Functions
Notation	Description
$K$	Total number of distinct pretraining data sources (tasks).
$\alpha_{k}$	Sampling probability (data mixing ratio) for the $k$ -th data source.
$\mathcal{L}_{k}(\bm{\theta})$	The expected / empirical loss on the $k$ -th source task.
$\mathcal{L}_{\text{train}}(\bm{\theta})$	The averaged pretraining loss: $\frac{1}{K}\sum_{k=1}^{K}\mathcal{L}_{k}(\bm{\theta})$ .
$\mathcal{L}_{\mathcal{T}}(\bm{\theta})$	The loss on an unseen downstream evaluation task $\mathcal{T}$ .

Table 5: Summary of key notations used in this paper.

Model Size	Optimizer	Outer LR	Inner LR ( $\gamma$ )	Chinchilla	Tokens (B)	Weight Decay	Reference
1B	Adam	0.002	-	$4\times$	50	0.2	Sec. ˜ 4.3
1B	Nexus	0.002	0.04	$4\times$	50	0.2	Sec. ˜ 4.3
3B	Adam	0.001	-	$2\times$	110	0.2	Secs. ˜ 4.2 and 4.5 Secs. ˜ 4.4 and 14.3
	Muon	0.001	-	$2\times$	110	0.1
	Nexus	0.001	0.1	$2\times$	110	0.2

Table 6: Summary of key hyperparameters for the main pretraining experiments.

Data	Optim.	Metric	Loss Metrics ( $\downarrow$ )		Gen.	Reasoning			Math		Code		Avg.
Data	Optim.	Metric	Pretrain.	OOD	MMLU	GPQA	GPQA-D	BBH	GSM8k	MATH	HumanEval	MBPP	All
Math10	AdamW	Acc. ( $\uparrow$ )	1.606	1.302	47.8	32.8	22.6	36.6	44.0	32.0	43.0	38.0	37.1
	AdamW	Loss ( $\downarrow$ )	1.606	1.302	2.265	2.005	1.910	1.534	1.259	1.054	1.116	1.922	1.633
	Nexus	Acc. ( $\uparrow$ )	1.602	1.290	48.9	29.6	23.4	36.6	59.0	40.0	47.0	38.0	40.3
	Nexus	Loss ( $\downarrow$ )	1.602	1.290	2.179	1.981	1.881	1.504	1.227	1.026	1.086	1.921	1.601
	Improv.	Loss ( $\uparrow$ )	+0.004	+0.012	+0.086	+0.024	+0.029	+0.030	+0.032	+0.028	+0.030	+0.001	+0.032

Table 7: Results on varying data mixtures (3B models). Hyperparameters follow Secs. ˜ 4.1 and 4.2 . As the proportion of math data increases (10%

\to

70%), the relative performance gains of Nexus on math benchmarks gradually diminish, whereas its advantages on other domains (General, Reasoning) progressively expand. This suggests Nexus boosts the sample-sparse or harder-to-learn domains in the mixture.

Model	Optim.	Metric	Loss Metrics ( $\downarrow$ )		Gen.	Reasoning			Math		Code		Avg.
Model	Optim.	Metric	Pretrain.	OOD	MMLU	GPQA	GPQA-D	BBH	GSM8k	MATH	HumanEval	MBPP	All
1B	Adam	Acc. ( $\uparrow$ )	1.331	1.863	34.2	25.8	18.8	24.8	18.0	14.0	38.0	1.0	21.8
	Adam	Loss ( $\downarrow$ )	1.331	1.863	2.552	2.280	2.191	1.700	1.708	1.346	1.324	2.991	2.011
	Nexus	Acc. ( $\uparrow$ )	1.338	1.835	31.4	25.0	18.8	23.0	22.0	12.0	41.0	15.0	23.5
	Nexus	Loss ( $\downarrow$ )	1.338	1.835	2.446	2.261	2.172	1.689	1.749	1.325	1.325	2.908	1.984
	Improv.	Loss ( $\uparrow$ )	-0.007	+0.028	+0.106	+0.019	+0.019	+0.011	-0.041	+0.021	-0.001	+0.083	+0.027

Table 8: Results on a public pretraining dataset [ 5 ] . The Adam baseline exhibits artificial performance inflation on leaked benchmarks. In contrast, Nexus effectively resists shortcut over-memorization, successfully reallocating model capacity to uncontaminated tasks and achieving superior overall OOD generalization.

Size	Optim.	Metric	Loss Metrics ( $\downarrow$ )		Gen.	Reasoning			Math		Code		Avg.
Size	Optim.	Metric	Pretrain.	OOD	MMLU	GPQA	GPQA-D	BBH	GSM8k	MATH	HumanEval	MBPP	All
130M	AdamW	Acc. ( $\uparrow$ )	2.038	1.559	28.0	22.6	24.2	25.1	7.0	6.0	0.0	10.0	15.4
	AdamW	Loss ( $\downarrow$ )	2.038	1.559	2.555	2.438	2.338	1.793	1.612	1.360	1.407	2.230	1.967
	Nexus	Acc. ( $\uparrow$ )	2.031	1.549	27.0	25.7	21.8	28.8	5.0	11.0	0.0	10.0	16.2
	Nexus	Loss ( $\downarrow$ )	2.031	1.549	2.523	2.414	2.312	1.792	1.601	1.330	1.384	2.183	1.942
	Improv.	Acc. ( $\uparrow$ )	-	-	-1.0	+3.1	-2.4	+3.7	-2.0	+5.0	0.0	0.0	+0.8

Table 9: Benchmark Performance across Model Scales. We compare downstream capabilities for models ranging from 130M to 2.3B parameters. Notably, the relative gains of Nexus over the base optimizer amplify as model capacity increases, with the average benchmark accuracy improvement growing from +0.8% on the 130M model to +3.2% on the 2.3B model.

Optim.	Metric	Loss Metrics ( $\downarrow$ )		Gen.	Reasoning			Math		Code		Avg.
Optim.	Metric	Pretrain.	OOD	MMLU	GPQA	GPQA-D	BBH	GSM8k	MATH	HumanEval	MBPP	All
AdamW	Acc. ( $\uparrow$ )	1.606	1.302	47.8	32.8	22.6	36.6	44.0	32.0	43.0	38.0	37.1
AdamW	Loss ( $\downarrow$ )	1.606	1.302	2.265	2.005	1.910	1.534	1.259	1.054	1.116	1.922	1.633
Adam+Nexus	Acc. ( $\uparrow$ )	1.602	1.290	48.9	29.6	23.4	36.6	59.0	40.0	47.0	38.0	40.3
Adam+Nexus	Loss ( $\downarrow$ )	1.602	1.290	2.179	1.981	1.881	1.504	1.227	1.026	1.086	1.921	1.601
(- AdamW)	Acc. ( $\uparrow$ )	-	-	+1.1	-3.2	+0.8	0.0	+15.0	+8.0	+4.0	0.0	+3.2

Table 10: Comparison with Muon Optimizer on 3B Models. As shown, Muon improves downstream performance by decreasing the pretraining loss. While Nexus achieves nearly the same pretraining loss as AdamW, it achieves comparable performance to Muon on downstream tasks.

Phase	Optim.	Metric	Loss Metrics ( $\downarrow$ )		Gen.	Reasoning			Math		Code		Avg.
Phase	Optim.	Metric	SFT	OOD	MMLU	GPQA	GPQA-D	BBH	GSM8k	MATH	HumanEval	MBPP	All
Pre-SFT	AdamW	Acc. ( $\uparrow$ )	1.655	1.263	52.7	28.9	25.8	41.1	54.0	37.0	50.0	42.0	41.4
	AdamW	Loss ( $\downarrow$ )	1.655	1.263	2.221	1.938	1.842	1.484	1.233	1.031	1.053	1.859	1.583
	Nexus	Acc. ( $\uparrow$ )	1.647	1.258	50.9	22.7	22.7	35.9	57.0	40.0	48.0	42.0	39.9
	Nexus	Loss ( $\downarrow$ )	1.647	1.258	2.138	1.929	1.838	1.489	1.179	1.006	1.030	1.803	1.552
	Improv. Loss ( $\uparrow$ )		+0.008	+0.005	+0.083	+0.009	+0.004	-0.005	+0.054	+0.025	+0.023	+0.056	+0.031

Table 11: Downstream SFT Results. As shown, Nexus does not prematurely compromise the model’s SFT capabilities; on the contrary, it continues to outperform AdamW after SFT.

Algorithm Newton-Schulz Iteration	Algorithm RMS-Aligned Update
def zeropower_via_ns(G, steps=5): a, b, c = 3.4445, -4.7750, 2.0315 X = G / (torch.norm(G) + 1e-7) for _ in range(steps): A = X @ X.mT X = aX + (bA + c*A@A) @ X return X	def sgd_alignRMS_update(G): K = min(G.shape[-2:]) # Normalize to match Muon scale norm = torch.norm(G) + 1e-7 return G * (K**0.5) / norm

Figure 8: Muon-inspired SGDM. Left: Quintic Newton-Schulz iterations for orthogonalization. Right: Per-tensor RMS normalization ensuring consistent update scale as Muon.

Panel of Figure 9 — Figure 9: Experimental Results using SGDM. Nexus both gives acceleration and implicit bias.

Panel of Figure 10 — Figure 10: Experimental Results including the BS-128 Baseline. The standard AdamW and Nexus trajectories are retained from the main trials in Sec. ˜ 4.2 for direct comparison.

Model	Configuration	Metric	Loss Metrics ( $\downarrow$ )		Gen.	Reasoning			Math		Code		Avg.
Model	Configuration	Metric	Pretrain.	OOD	MMLU	GPQA	GPQA-D	BBH	GSM8k	MATH	HumanEval	MBPP	All
1B	AdamW (BS=128)	Acc. ( $\uparrow$ )	1.937	1.506	29.3	25.0	20.3	23.3	8.0	5.0	11.0	14.0	17.0
1B	AdamW (BS=128)	Loss ( $\downarrow$ )	1.937	1.506	2.396	2.333	2.231	1.728	1.510	1.286	1.340	2.091	1.864
3B	AdamW (BS=128)	Acc. ( $\uparrow$ )	1.627	1.326	45.4	29.7	21.9	40.7	39.0	26.0	37.0	39.0	34.8
3B	AdamW (BS=128)	Loss ( $\downarrow$ )	1.627	1.326	2.262	2.033	1.937	1.556	1.281	1.067	1.148	1.951	1.654

Table 12: Results of Small Batch Size Baseline. Validation losses and downstream capabilities for AdamW trained with a reduced batch size of 128. Compared to the optimal BS-256 baseline reported in the main text ( Tab. ˜ 1 ), reducing the batch size results in significant performance degradation.

Panel of Figure 11 — Figure 11: Ablation Results on Regularization Strength $\lambda$ . As the coefficient $\lambda$ increases, there is a monotonic increase in gradient closeness across both training and OOD sets.

Language Model Pretraining

Language model pretraining teaches models language patterns and knowledge from large text corpora before they are adapted to specific tasks.

See all 440 papers

Mar 12, 2026

Neural Thickets: Diverse Task Experts Are Dense Around Pretrained Weights

Pretraining produces a learned parameter vector that is typically treated as a starting point for further iterative adaptation. In this work, we instead view the outcome of pretraining as a distribution over parameter vectors, whose support already contains task-specific experts. We show that in small models such expert solutions occupy a negligible fraction of the volume of this distribution, making their discovery reliant on structured optimization methods such as gradient descent. In contrast, in large, well-pretrained models the density of task-experts increases dramatically, so that diverse, task-improving specialists populate a substantial fraction of the neighborhood around the pretrained weights. Motivated by this perspective, we explore a simple, fully parallel post-training method that samples

N

parameter perturbations at random, selects the top

K

, and ensembles predictions via majority vote. Despite its simplicity, this approach is competitive with standard post-training methods such as PPO, GRPO, and ES for contemporary large-scale models.

Yulu Gan, Phillip Isola

MIT CSAIL

LLM Post-Training Weight-Space Analysis

Ask PDF

Sep 26, 2026

Relative Generalization Invariance of LLM Pretraining

Large Language Model (LLM) pretraining performance is jointly shaped by three components of the training triplet: the optimizer, model architecture, and training data stream. However, how these components influence performance in distinct ways remains unclear. We take a first step toward isolating their effects by studying relative generalization. We introduce Relative Generalization Invariance (RGI), the invariance of the validation-loss difference between any two tokens across models. We show that RGI approximately holds across a wide range of optimizers and moderate architectural variations, suggesting that these choices induce an approximately uniform shift in token-wise losses. In contrast, changing the training data stream can substantially alter relative generalization. We further show that RGI cannot be explained by the neural tangent kernel or mean-field regimes alone and prove that it can emerge in an overparameterized quadratic model. Overall, our work identifies RGI as a new phenomenon in LLM pretraining that helps distinguish the effects of optimizers and architectures from those of training data.

Fengzhuo Zhang, Shuche Wang, Shenggui Li +7

Yale University · National University of Singapore · Nanyang Technological University +2

Language Model Pretraining Language Modeling

Ask PDF

May 4, 2026

Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting

Pretraining optimizers are tuned to produce the strongest possible base model, on the assumption that a stronger starting point yields a stronger model after subsequent changes like post-training and quantization. This overlooks the geometry of the base model which controls how much of the base model's capabilities survive subsequent parameter updates. We study three pretraining optimization approaches that bias optimization toward flatter minima: Sharpness-Aware Minimization (SAM), large learning rates, and shortened learning rate annealing periods. Across model sizes ranging from 20M to 150M parameters, we find that these interventions consistently improve downstream performance after post-training on five common datasets with up to 80% less forgetting. These principles hold at scale: a short SAM mid-training phase applied to an existing OLMo-2-1B checkpoint reduces forgetting by 31% after MetaMath post-training and by 40% after 4-bit quantization.

Ishaan Watts, Catherine Li, Sachin Goyal +2

Carnegie Mellon University

Language Model Pretraining Sharpness-Aware Minimization LLM Post-Training Catastrophic Forgetting

Ask PDF

Abstract

Language Model Pretraining

Explore similar work

Neural Thickets: Diverse Task Experts Are Dense Around Pretrained Weights

Relative Generalization Invariance of LLM Pretraining

Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting