arXiv CSopen access

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

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

machine learning, deep learning, neural networks

Nexus: Same Pretraining Loss, Better Downstream Generalization via Common Minima Huanran Chen1,2,‡ , Huaqing Zhang1 , Xiao Li2,† , Yinpeng Dong1,∗ , Ke Shen2 , Jun Zhu1,∗ 1

arXiv:2604.09258v1 [cs.LG] 10 Apr 2026

‡

Tsinghua University, 2 ByteDance Seed

Work done at ByteDance Seed, † Project Lead, ∗ Corresponding authors

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., Fig. 2b), or merely a minimizer of the summed loss (e.g., Fig. 2a)? 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 Fig. 1). 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. Date: April 13, 2026 Email: Yingpeng Dong, Jun Zhu at {dongyinpeng, dcszj}@tsinghua.edu.cn;

Huanran Chen, Huaqing Zhang at {chenhr25, zhanghq22}@mails.tsinghua.edu.cn; Xiao Li, Ke Shen at {lixiao.20, shenke}@bytedance.com.

Introduction

Pretraining is the cornerstone of Large Language Models (LLMs). Accounting for 95% to over 99% of the total computational budget and data, it serves as the indispensable engine for their capabilities [23, 47]. 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 [10, 23, 41, 47]. To learn from such a heterogeneous corpus of K distinct sources, the standard practice PK 1 is to average the loss of each data source Lk (θ) and minimize the averaged loss Ltrain (θ) = K k=1 Lk (θ).

2.8

AdamW Nexus

2.6 2.4

AdamW Nexus

2.0

0.40

MMLU

OOD Loss

Eval Loss

0.35

1.6

2.0

0.30

1.8

1.4

0.25

1.6 0

20000

40000

60000

Step

80000

100000

(a) Pretraining Loss 2.6

20000

40000

60000

Step

80000

2.0

20000

AdamW Nexus

Math500 Loss

80000

100000

3.2

AdamW Nexus

3.0

MBPP Loss

GSM8k Loss

1.8

60000

Step

2.8

1.6

2.0

40000

1.8

2.2

2.6 2.4

1.4

1.6

2.2 1.2

1.4 0

20000

40000

60000

Step

80000

100000

1.0

2.0 0

(d) GSM8k Loss 0.6

20000

40000

60000

Step

80000

100000

1.8

(e) Math500 Loss

AdamW Nexus

0.5

100000

MBPP

Math500

GSM8k

80000

0.2

0.15 0.10

0.1

60000

Step

0.3

0.20

0.2

40000

AdamW Nexus

0.4

0.25

0.3

20000

(f) MBPP Loss

AdamW Nexus

0.30

0.4

0.0

100000

(b) OOD Loss AdamW Nexus

2.4

1.2

AdamW Nexus

0.45

1.8

2.2

0.50

0.1

0.05 0.0 0

20000

40000

60000

Step

(g) GSM8k

80000

100000

20000

40000

60000

Step

(h) Math500

80000

100000

20000

40000

60000

Step

80000

100000

(i) MBPP

Figure 1 Illustration of "same pretraining loss, better downstream task". The training loss of baseline and our Nexus are exactly the same. However, our methods achieves much better downstream generalization.

In this work, we investigate an interesting geometric question: Does the model converge to a common minimizer across all data sources Lk , or does it merely find a minimizer of the summed loss Ltrain ? To illustrate this, consider a simplified setting composed of two data sources (L1 and L2 ), yielding a training loss of Ltrain (θ) = 12 (L1 (θ) + L2 (θ)). As depicted in Fig. 2, there exist two distinct types of minimizers that achieve the exact same training loss Ltrain . The first type corresponds to the Sum of Minima (Fig. 2a), where ∗ the converged parameter θtrain successfully minimizes the total training loss Ltrain yet remains geometrically distant from the minimizers of individual tasks Lk . The second type approaches the Intersection of Minima ∗ (Fig. 2b), where θtrain is not only a minimizer of Ltrain , but is also geometrically close to the minimizer of each individual task Lk . We hypothesize that this geometric “ closeness”—the distance between task-specific minima—is strongly correlated with downstream generalization. Even when achieving the exact same pretraining loss, these two types of minimizers yield drastically different downstream losses LT (see the blue curve LT (θ) in Fig. 2). Intuitively, if the training losses Lk and downstream task LT are quadratic and i.i.d. distributed, the Intersection-type minimizer (Fig. 2b) will strictly outperform the Sum-type minimizer (Fig. 2a) on the downstream task LT , given the same pretraining loss (see Theorem 2.2). Therefore, we posit that this intuition may generalize beyond quadratics to LLM pretraining, and steering the optimization toward the Intersection-type minimizer would achieve the “same pretraining loss, better downstream task”. 2

1 2 train( )

()

3 2 1 0

1 2

Loss

train( )

()

3 2 1

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

(a) Distant (Sum of Minima)

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

(b) Close (Intersection of Minima)

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 LT .

However, directly optimizing for this geometric “closeness” is computationally intractable, as it requires knowing the exact minimizer of each Lk at every training step. To overcome this, we prove that the gradient ∇LT ∇L similarity between tasks, CosSim(∇Li , ∇Lj ) ≜ ∥∇Lii∥∥∇Ljj ∥ , upper bounds the geometric closeness. The rationale is straightforward: if the gradient directions of each loss ∇Lk are always exactly the same throughout optimization, their respective minimizers θk∗ must be exactly the same. Based on this insight, we propose the Nexus algorithm, which approximates the gradient of gradient similarity ∇CosSim(∇Li , ∇Lj ). Combining Nexus with pretraining optimizer [17, 20, 45] effectively maximizes CosSim(∇Li , ∇Lj ). In Sec. 5.1, we show that both gradient similarity and geometric closeness generalize to downstream tasks, thus leading to lower downstream loss and better downstream performance, even when achieving the same pretraining loss. We empirically validate Nexus across various settings, including model scales ranging from 130M to 3B parameters [42, 45, 47], diverse pretraining data and mixtures [2, 40], learning rate schedules [15, 26, 46], and training compute [18]. Experimental results demonstrate that, across nearly all settings, Nexus reduces the downstream loss by over 0.02 compared to the base optimizers—a substantial margin that typically requires doubling the pretraining compute [18]—while achieving the exact same pretraining loss. For instance, on the 3B model, Nexus improves GSM8K accuracy by 15%, MATH500 by 8% and HumanEval by 4%. These consistent and substantial downstream gains demonstrate the importance of implicit biases in unlocking downstream generalization [24], particularly as the current pretraining paradigm transitions from being compute-bound to data-bound [19, 29, 33, 37].

Closeness: A Second-Order Property Related to Generalization

2.1

Problem Formulation

Formally, let the pretraining corpus be the union of K distinct data sources, denoted as Dtrain = ∪K k=1 Dk . Let αk represent the sampling probability (data mixing ratio) for the k-th source. We define the weighted empirical loss function for the k-th source as: |Dk |

Lk (θ) = −αk

log p(xj |θ).

(1)

j=1

Consequently, the total pretraining objective is simply the average of these weighted losses: K

Ltrain (θ) =

1 X Lk (θ). K k=1

(2)

2.2

Flatness and Closeness are both Second Order Generalization Biases

∗ Our primary interest lies in how well our pretraining minimizer θtrain ∈ arg minθ Ltrain (θ) performs on the ∗ downstream task T , i.e., the downstream loss LT (θtrain ).

Let ST = {θ | ∃ϵ > 0, ∀θ ′ ∈ Bϵ (θ), LT (θ) ≤ LT (θ ′ )} be the set of local minimizers for the downstream task. We define θT∗ as the closest minimizer of downstream loss: ∗ θT∗ = arg min ∥θ − θtrain ∥2 . θ∈ST

(3)

By applying a second-order Taylor expansion of LT around the optimal point θT∗ , we can bound the downstream ∗ loss at the converged point θtrain : 1 ∗ ∗ ∗ ∗ ∗ LT (θtrain ) = LT (θT∗ ) + (θtrain − θT∗ )⊤ ∇LT (θT∗ ) + (θtrain − θT∗ )⊤ ∇2 LT (θT∗ )(θtrain − θT∗ ) + O(∥θtrain − θT∗ ∥3 ) 2 1 ∗ ∗ ∗ = LT (θT∗ ) + (θtrain − θT∗ )⊤ ∇2 LT (θT∗ )(θtrain − θT∗ ) + O(∥θtrain − θT∗ ∥3 ) 2 1 ∗ max∗ ∥∇2 LT (ξ)∥2 , ≤ LT (θT∗ ) + ∥θtrain − θ ∗ ∥2 · {z T }2 ξ∈[θT∗ ,θtrain 2| ] | {z } Closeness Flatness

(4) ∗ ∗ where the notation [θT∗ , θtrain ] denotes the line segment connecting θT∗ and θtrain . Note that the first-order term vanishes because θT∗ is a local minimizer (i.e., ∇LT (θT∗ ) = 0). The remaining term is controlled by two factors: the Flatness of the downstream loss landscape along the path, and crucially, the Closeness between our converged point and the task optimal. Remark 2.1. Prior literature extensively characterizes the local loss landscape as exhibiting high quadraticity, at least along most directions [4, 22, 44]. It is worth noting that the inequality in Eq. (4) becomes an exact equality when the loss function is strictly quadratic along the one-dimensional direction connecting θT∗ and ∗ θtrain . Therefore, if the standard assumption that the loss landscape is locally quadratic holds (which only needs to be true along typical directions), this bound would be extremely tight and serve as an accurate proxy for the generalization gap. ∗ Therefore, the flatter the local loss landscape of LT and the closer the converged parameter θtrain is to the task minimizer θT∗ , the better the generalization. Together, flatness and closeness encapsulate all secondorder information for downstream generalization. While flatness has been well-studied in prior literature [4, 12, 21, 38, 49], in this work, we focus solely on our new implicit bias: closeness.

2.3

Closeness Improves Out-of-Distribution Generalization

Eq. (4) reveals that the closeness between the trained parameters and the downstream task minimizers directly ∗ correlates with downstream generalization. In other words, if one could minimize ∥θtrain − θT∗ ∥22 without ∗ compromising the intrinsic loss LT (θT ) and the flatness, one would directly boost downstream generalization. However, in practice, minimizing closeness typically comes at a cost: either (1) an increase in intrinsic loss LT (θT∗ ) or (2) an increase in sharpness (see Sec. 5.1). This trade-off is expected; if one were to minimize the PK 1 ∗ ∗ 2 closeness even among training tasks (i.e., K k=1 ∥θtrain − θk ∥2 ) without penalty, it would imply achieving significantly smaller training error and faster optimization rates. This contradicts the prevailing assumption and empirical observations regarding the inherent hardness of discovering significantly faster optimizers [35, 45]. In this paper, we specifically focus on the "same training loss" regime. We demonstrate that a "close" minimizer (Fig. 2b) yields significantly better out-of-distribution generalization compared to a "distant" PK 1 minimizer (Fig. 2a), even at the same pretraining loss Ltrain (θ) = K k=1 Lk (θ). We analyze the specific scenario where improved closeness is achieved solely at the cost of increasing the intrinsic task loss Lk (θk∗ ). This assumption decouples our analysis from the flatness bias (thereby eliminating flatness as a confounding factor) and aligns with the actual behavior observed in our experiments (see Sec. 5.1).

The core intuition is illustrated in Fig. 2: as long as the loss landscape is quadratic-like along the directions of interest (i.e., locally and directionally strongly convex), and the pretraining and downstream tasks share a common task distribution, improved closeness will inherently lead to a lower generalization gap. We begin with a simplified analysis assuming strictly quadratic loss functions to mathematically substantiate this intuition. Theorem 2.2 (Generalization of Closeness in the Quadratic Case). To model the non-convex landscape, assume the parameter space Rd is partitioned into a set of disjoint basins of attraction {B}. Within any specific basin B, assume that any task L sampled from a distribution P is locally a quadratic function: ∗ 2 ∗ ∥2 + cB , where the local task minimizers are distributed as θL L(θ) = a2 ∥θ − θL ∼ P(µB , σB2 I) with mean µB 2 and variance σB , and cB is the intrinsic loss (depth) of basin B. ∗ Let the pretraining tasks {Lk }K k=1 and the downstream task LT be i.i.d. samples from P. Let Θ = {θtrain,B | ∗ Ltrain (θtrain,B ) = Ctrain } be the set of converged minimizers across different basins that achieve the exact same ∗ training loss Ctrain . For any candidate θtrain,B ∈ Θ, the expected downstream error on an unseen task T ∼ P 2 is strictly proportional to the task variance σB : ∗ ET ∼P [LT (θtrain,B )] = Ctrain +

a 2 σ . K B

(5)

PK 1 ∗ ∗ = K Proof. By stationarity, the converged parameter is the mean of local minimizers: θtrain,B k=1 θk,B . Constraining the training loss to Ctrain and closeness to σB2 explicitly determines the basin’s intrinsic depth: PK a ∗ ∗ 2 cB = Ctrain − 2K k=1 ∥θtrain,B − θk,B ∥2 . This enforces the core trade-off: to achieve the identical Ctrain , a basin with tightly clustered minimizers inherently requires a higher intrinsic loss cB to compensate. ∗ − θT∗ ,B ∥22 ] + cB . Substituting For an unseen task T ∼ P, the expected downstream loss is E[LT ] = E[ a2 ∥θtrain,B cB perfectly cancels out the intrinsic depth, leaving the generalization dependent on the variance K−1 2 gap aentirely 1 ∗ of the distributions: E[LT (θtrain,B )] − Ctrain = a2 1 + K − K σB = K σB2 .

Consequently, as long as downstream tasks and pretraining tasks follow the same distribution, by trading intrinsic loss cB for improved closeness (i.e., smaller σB2 ), one obtains better out-of-distribution generalization due to the reduction in task variance. We can also extend Theorem 2.2 beyond purely quadratic loss functions to the broader class of general loss landscapes exhibiting local and directional strong convexity, as demonstrated in the following theorem. Theorem 2.3 (Generalization of Closeness beyond Quadratics, Proof in Sec. C.2). Let θ ∗ be a specific local ∗ minimizer of the population loss EL∼P [L(θ)]. For any task L sampled from P, let θL = arg minθ∈SL ∥θ ∗ − θ∥2 be its corresponding local minimizer. Assume that for any task L ∼ P, the loss function is locally and ∗ directionally strongly convex along the segments [θL , θ ∗ ], i.e., λmax ≥ u⊤ ∇2 L(ξ)u ≥ λmin > 0 for any ∗ ∗ ∗ ∗ ∗ ∗ ξ ∈ [θL , θ ] and any unit vector u ∈ span{θ − θL | L ∼ P}. Let µ = E[θL ] and σ 2 = E[∥θL − µ∥22 ]. 2 ∗ Assuming the statistical independence between the task flatness ∇ LL (ξ) and the task closeness θL across the distribution P. Conditioned on achieving a fixed training loss Ctrain , the expected out-of-distribution ∗ generalization error of the converged training parameter θtrain is bounded by: 2 λmax λmax + 1 λmin ∗ ET ∼P [LT (θtrain )] − Ctrain ≤ σ2 . (6) 2K Therefore, as long as the loss landscape exhibits quadratic-like behavior (i.e., local and directional strong convexity) along these typical directions [θk∗ , θ ∗ ], explicitly optimizing for closeness σ 2 = E[∥θk∗ − µ∥22 ] would be beneficial for a lower downstream loss.

Nexus Optimizer: Enhancing Closeness via Second-Order Approximation

Both geometric intuition and our analysis of quadratic functions support the conclusion that a “close” minimizer (Fig. 2b) generalizes to out-of-distribution data significantly better than a “distant” minimizer (Fig. 2a), even 5

when achieving the same training loss. Consequently, we aim to explicitly optimize this closeness during LLM pretraining. In this section, we introduce a second-order gradient approximator named “Nexus”, which PK 1 ∗ 2 effectively optimizes parameter closeness on the training tasks (i.e., K ∥θ − θ k ∥2 ), and successfully k=1 generalizes to the closeness of unseen downstream tasks (i.e., ∥θ − θT∗ ∥22 ).

3.1

Gradient Similarity Upper Bounds Closeness

Directly optimizing the closeness metric ∥θ − θk∗ ∥ involves finding the specific minimizer θk∗ for each task, which is itself a minimization problem and computationally P prohibitive. Fortunately, we observe that the gradient similarity between different source tasks, given by i̸=j −∇Li (θ)⊤ ∇Lj (θ), provides a tractable upper bound for closeness. Intuitively, if the gradients of distinct tasks Lk consistently align in direction, their respective minimizers be exactly the same. Theoretically, both the gradient dot product and cosine similarity serve as tight bounds for closeness: Theorem 3.1 (Gradient Similarity Upper Bounds Closeness). Let θ be the converged parameter satisfying PK 1 ′ ′ ∇Ltrain (θ) = K k=1 ∇Lk (θ) = 0. Let Sk = {ϑ | ∃ϵ > 0, ∀ϑ ∈ Bϵ (ϑ), Lk (ϑ) ≤ Lk (ϑ )} be the set of local

minimizers for task k, and θk∗ = arg minϑ∈Sk ∥ϑ−θ∥2 . Let λmin = mink inf ξ∈[θ,θk∗ ]

∗ ⊤ ∗ (θ−θk ) θ−θk 2 ∗ ∥ ∇ Lk (ξ) ∥θ−θ ∗ ∥ ∥θ−θk 2 k 2

0, and G = supk ∥∇Lk (θ)∥2 . Then, the closeness between the minimizers is bounded by: K 1 X 1 X G2 X ∥θ − θk∗ ∥22 ≤ −∇Li (θ)⊤ ∇Lj (θ) ≤ (1 − CosSim(∇Li (θ), ∇Lj (θ))) . 2 K Kλmin Kλ2min k=1

i̸=j

(7)

i̸=j

In other words, optimizing the training trajectory towards a regime where CosSim(∇Li (θ), ∇Lj (θ)) remains consistently high guarantees high closeness (i.e., a small distance ∥θ − θk∗ ∥2 ). This "gradient similarity upper bound" also provides a more intuitive understanding of why closeness improves downstream generalization. Suppose that the high gradient similarity achieved among training tasks (i.e., high Sim(∇Li (θ), ∇Lj (θ))) successfully generalizes to the similarity between the training objective and the downstream task (i.e., high Sim(∇Ltrain (θ), ∇LT (θ))). This similarity directly represents the reduction in downstream loss after a single Gradient Descent (GD) step on the training set (in the first-order sense): = γ∇Ltrain (θ)⊤ ∇LT (θ) + O(γ 2 ).

LT (θ) − LT (θ − γ∇Ltrain (θ)) {z } |

(8)

decrease of downstream loss after one GD step on training set

Therefore, we view gradient similarity CosSim(∇Li (θ), ∇Lj (θ)) as a strong proxy for parameter closeness: it not only provides a tight upper bound on parameter distance (thereby enforcing closeness), but also leads to the same beneficial effects on downstream generalization. Given this strong connection, in the remainder of this paper, we use the term "closeness" to refer to both parameter closeness and gradient closeness.

3.2

Optimizing Gradient Similarity via Nexus

Therefore, to encourage parameter closeness, it suffices to maximize the gradient similarity. However, directly optimizing this objective is computationally intractable because the gradient of the cosine similarity involves the Hessian matrix: ! ∇Lj ∇L⊤ ∇2 Li ∇Lj ∇2 Lj ∇Li ∇Li ∇L⊤ j i ∇θ CosSim(∇Li , ∇Lj ) = I − + I− . (9) 2 2 ∥∇Li ∥2 ∥∇Li ∥2 ∥∇Lj ∥2 ∥∇Lj ∥2 ∥∇Lj ∥2 ∥∇Li ∥2 To address this challenge, we propose the Nexus optimizer, which approximates the gradient in Eq. (9) through a dual-loop mechanism. The complete procedure is outlined in Algorithm 1. Conceptually, one should view each step in the outer loop as a standard parameter update, while the K steps in the inner loop serve as a gradient approximator for Eq. (9). Specifically, for each outer iteration, we perform K normalized SGD steps (inner loop) to accumulate the approximated gradient ĝt . This ĝt is then passed to the outer optimizer (e.g., AdamW [20, 25], Muon [17]) to perform the actual update. The following theorem demonstrates that Nexus algorithm effectively maximizes gradient similarity. 6

Algorithm 1 Standard Nexus Algorithm Require: Initial params θ0 , losses {Li }ki=1 , total iterations T . Require: Optimizers: Optinner (Normalized SGD), Optouter (e.g., AdamW). Inner learning rate γ. 1: for t = 1 to T do 2: 3: 4: 5: 6: 7: 8:

θt,0 ← θt−1 {Initialize inner loop} for m = 1 to K do

Sample task index sm ∼ Uniform({1, . . . , K}) g ← ∇Lsm (θt,m−1 ) g {Update inner trajectory} θt,m ← θt,m−1 − γ · ∥g∥ 2 ĝt ← θt,0 − θt,K {Compute Nexus pseudo-gradient} θt ← Optouter (θt−1 , ĝt ) {Outer-update} Return: θT Nexus Update:

Nexus Regularizer

Conventional Update + Nexus Regularizer: Conventional Update

Nexus Regularizer

Nexus Regularizer:

Figure 3 Intuitive illustration of Nexus Algorithm.

Theorem 3.2 (Nexus Maximizes Gradient Similarity). Assume there exist constants Gmin , L, ρ > 0 such that for any t ∈ [1, T ] and m ∈ [1, K]: ∥∇Li (θt,m )∥2 ≥ Gmin ;

∥∇2 Li (θ)∥2 ≤ L;

∥∇2 Li (x) − ∇2 Li (y)∥2 ≤ ρ∥x − y∥2 .

(10)

Then, the sequence {θt } generated by Algorithm 1 effectively minimizes the following second-order objective: J2nd (θ) = γ

K X i=1

∥Li (θ)∥2 − γ 2

K −1X CosSim ∇Li (θ), ∇Lj (θ) . 4K

(11)

i̸=j

This holds because the expected update direction satisfies: E[ĝt ] = γ

K X

X ∇Li (θt ) K −1 − γ2 ∇θ CosSim(∇Li , ∇Lj ) + E 2nd , ∥∇Li (θt )∥2 4K i=1

(12)

i̸=j

where the approximation error is bounded by ∥E 2nd ∥2 ≤ 16

4L2 +ρGmin G2min

K 3 γ 3 = O(γ 3 ).

Intuitive Understanding of the Inner Loop. To intuitively understand why Nexus’s inner loop optimizes

Eq. (9), consider a simplified scenario with two loss functions, L1 and L2 , as illustrated in Fig. 3. At the current parameter state θt,0 , a conventional optimizer (e.g., AdamW, Muon) would simply aggregate the gradients as ∇L1 (θt,0 ) + ∇L2 (θt,0 ) for the update. In contrast, the Nexus inner loop operates sequentially: it first takes a step using ∇L1 (θt,0 ) to reach an intermediate point θt,1 , and subsequently evaluates the next gradient ∇L2 (θt,1 ) at this displaced location. As shown in the figure’s equations, this sequential trajectory is 7

Algorithm 2 Standard Pretraining Require: model, loader Require: opt_outer (e.g., AdamW) Require: accum_steps 1: for i, batch in loader do 2: {Mini-batch Step} 3: L ← model(batch) 4: L.backward() 5: if i % accum_steps == 0 then 6: {Accumulation Step} 7: opt_outer.step() 8: opt_outer.zero_grad()

Algorithm 3 Nexus (Engineering Adaptation) Require: model, loader, opt_outer, accum_steps 1: inner_model ← model.clone() 2: opt_inner ← NSGD(inner_model) 3: for i, batch in loader do 4: L ← inner_model(batch) 5: L.backward() 6: opt_inner.step() 7: if i % accum_steps == 0 then 8: 9: 10:

ĝ ← inner_model − model opt_outer.step(grad=−ĝ) inner_model ← model.clone()

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 microbatch/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 (inner_model − model) serves as the pseudo-gradient ĝ for the outer optimizer, after which the inner model is re-synchronized.

mathematically equivalent to a conventional update plus a "Nexus regularizer." Crucially, this regularizer naturally yields a Hessian-gradient product, which equals the gradient of the gradient similarity (in the first-order sense). Consequently, the pseudo-gradient ĝt produced by the inner loop effectively serves as the sum of the gradient of the pretraining loss and the gradient of the gradient similarity defined in Eq. (9). Thus, Nexus serves as an effective mechanism for maximizing parameter closeness. Strictly speaking, Nexus should be conceptualized as a gradient approximator rather than a standalone optimizer, for two reasons: (1) the inner optimization step must be exactly vanilla SGD without any momentum (otherwise, as shown in Sec. 5.1, it fails to maximize gradient similarity), whereas (2) the outer optimizer can be any standard optimization algorithm. Consequently, Nexus is fully orthogonal to the choice of the outer base optimizer (e.g., it can be combined with AdamW, Muon, etc.).

3.3

Adapting Nexus to Practical Pretraining

We establish that Nexus effectively maximizes gradient similarity with controllable higher-order errors in Theorem B.3 and 3.2. However, directly applying Algorithm 1 to pretraining is still difficult. This is because Algorithm 1 requires computing gradients for every data source to perform a single effective outer update. In pretraining, the number of data sources is typically large (e.g., K > 50), which would result in an effective batch size that differs significantly from standard settings [18, 45]. This prevents us from leveraging established hyperparameters, thereby increasing tuning costs and preventing the wide application of Nexus. To address this, we propose an engineering adaptation to better adapt Nexus to practical pretraining. As shown in Algorithm 2, standard pretraining can be viewed as a gradient accumulation workflow: it computes gradients in every mini-batch and performs an optimizer update in every accumulation step. Leveraging this structure, we adapt Nexus as illustrated in Algorithm 3. Specifically, we introduce an auxiliary inner_model. For each mini-batch, we perform an immediate Normalized SGD (NSGD) update on this inner model to approximate the hessian-gradient product. Upon completing the accumulation steps, we compute the displacement between the inner model and the frozen main model, using this displacement as the pseudo-gradient ĝ for the outer optimizer. Therefore, our adapted Nexus actually maximizes the cosine similarity between mini-batches within a single accumulation step. Since the pretraining corpus is typically vast and the mixing ratio for each source is typically low, two consecutive mini-batches are highly likely to be sampled from different sources. Thus, this approach effectively achieves the objective of Algorithm 1. 8

Remark. It is worth noting that our adapted Nexus incurs almost no extra computational cost. The

total number of forward and backward passes remains exactly the same as standard pretraining. The only computational overhead comes from the copy and update of the inner model, but this is negligible compared to the forward-backward pass (considering the classical 6N BS approximation [18]). The only memory overhead comes from the inner model, but this can be reduced to nearly zero through techniques like CPU offloading and asynchronous processing. We employ Algorithm 3 for all experiments, with the exception of specific ablation studies. Readers may proceed directly to Sec. 4. In Sec. B, we also provide a theoretical analysis of Nexus’s convergence speed and discuss its implications for standard Normalized SGD.

Experiments

In this section, we validate that Nexus achieves nearly the same pretraining loss while delivering better downstream performance through comprehensive experiments across various datasets, learning rate schedules, model scales and token scales.

4.1

Experimental Settings

Our experimental setup largely follows the protocols established in Wen et al. [45] and OLMo et al. [30]. Pretraining Datasets. We utilize an in-house pretraining dataset similar to [40]. This corpus is: (1) strictly

cleaned to ensure no data contamination regarding the evaluated benchmarks or distillation data; and (2) of higher quality and stability than typical open-source datasets, allowing us to observe smooth and clear optimization trends. We also conduct experiments on public datasets [2] in Sec. H.3. However, these public datasets are not strictly decontaminated and contain training samples from our benchmarks. This leads to artificially inflated performance on certain tasks while underperforming on others. Consequently, we primarily rely on the strictly cleaned dataset for more stable analysis. Model Architecture. Following Wen et al. [45], we train Llama-architecture models of 130M, 300M, 520M, 1.2B, and 2.3B parameters (excluding embeddings). We primarily analyze the 520M (1B total parameters) and 2.3B (3B total parameters) models, hereafter referred to by their total parameter counts for brevity, except in the scaling law analysis (Sec. 4.3) as required by Kaplan et al. [18]. Hyperparameters. Wen et al. [45] have already conducted extensive parameter searches using grid search,

coordinate descent, and fine-grained tuning. To ensure fairness, we always apply exact the same hyperparameters to both Nexus and its corresponding base optimizers. For the base optimizers, we adopt the optimal hyperparameters identified in Wen et al. [45]. We further verified these settings by sweeping the learning rate with a multiplier of 2 (i.e., verifying 0.5× and 2.0×), confirming that their configurations remain optimal for our dataset. See Sec. H.1 for the detailed hyperparameters in each experiment. Benchmarks. We evaluate on diverse benchmarks encompassing general knowledge (MMLU [13]), reasoning (GPQA, GPQA Diamond [34], BBH [39]), math (GSM8k [5], MATH500 [14]), and coding (HumanEval [11], MBPP [1]). Beyond discrete accuracies, we also track downstream task losses and out-of-distribution (OOD) loss. The OOD loss is evaluated on a strictly cleaned proprietary in-house corpus, which exhibits a strong correlation with downstream benchmark capabilities. Highlighting Strategy. We use bold to highlight non-trivial performance gaps, defined as a loss difference

> 0.01 or a benchmark improvement > 2%, following Wen et al. [45].

4.2

Main Experimental Results

Settings. We train 1B models by 4× Chinchilla and 3B models for 2× Chinchilla tokens using two optimizer

configurations: the standard AdamW baseline, AdamW equipped with our Nexus regularizer (Nexus).

Nexus achieves Same Pretraining Loss, Better Downstream Task. As detailed in Tab. 1, Nexus strictly

satisfies the “same pretraining loss” condition, showing an immaterial difference of 0.004 compared to the baseline. Despite this parity in pretraining loss, Nexus demonstrates substantial improvements across nearly 9

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. Model Optim.

Metric

Loss Metrics (↓)

Gen.

Reasoning

Math

Code

Avg.

Pretrain. OOD MMLU GPQA GPQA-D BBH GSM8k MATH HumanEval MBPP

AdamW

Acc. (↑) Loss (↓)

1.826

1.433

Nexus

Acc. (↑) Loss (↓)

1.826

1.428

Improv.

Acc. (↑) Loss (↑)

AdamW

32.1 2.363

25.0 2.221

21.8 2.124

29.6 1.640

18.0 1.429

33.5

30.4 2.201

21.8 2.102

29.3 1.638

All

13.0 1.204

19.0 1.270

17.0 21.9 2.035 1.786

20.0

13.0

1.396

1.176

19.0 1.261

22.0 1.977

0.000

+1.4 +5.4 0.0 -0.3 +2.0 0.0 +0.005 +0.047 +0.020 +0.022 +0.002 +0.033 +0.028

0.0 +0.009

+5.0 +1.7 +0.058 +0.027

Acc. (↑) Loss (↓)

1.606

1.302

Nexus

Acc. (↑) Loss (↓)

1.602

1.290

Improv.

2.316

47.8 2.265

23.6 1.758

22.6 1.910

36.6 1.534

44.0 1.259

32.0 1.054

43.0 1.116

38.0 37.1 1.922 1.633

23.4 1.881

36.6 1.504

59.0 1.227

40.0 1.026

47.0 1.086

38.0 1.921

Acc. (↑) +1.1 -3.2 +0.8 0.0 +15.0 +8.0 Loss (↑) +0.004 +0.012 +0.086 +0.024 +0.029 +0.030 +0.032 +0.028

+4.0 +0.030

0.0 +3.2 +0.001 +0.032

32.8

2.005

48.9

29.6

2.179

1.981

40.3 1.601

all evaluated out-of-distribution and downstream metrics. Specifically, it reduces the OOD validation loss by 0.012 and yields significant accuracy gains on complex reasoning benchmarks, including a +15.0% improvement on GSM8k, +8.0% on MATH, and +4.0% on HumanEval. These consistent gains across diverse domains validate our core hypothesis: steering optimization toward the intersection of task minima effectively unlocks downstream generalization in the same pretraining loss regime. Comparison of Muon and Nexus. Compared to the standard AdamW baseline, Muon reduces the pretraining

loss by 0.029 and improves the average downstream accuracy by 2.3%. In contrast, Nexus achieves a negligible 0.004 reduction in pretraining loss yet yields a 3.2% improvement in average downstream accuracy, reaching a downstream performance level comparable to Muon (see Tab. 10). This observation indicates a fundamental divergence in their optimization pathways: while Muon’s downstream improvements rely primarily on achieving a lower pretraining loss, the gains from Nexus stem from its implicit bias despite maintaining a nearly identical pretraining loss as AdamW. Output Analysis. Compared to the AdamW baseline, Nexus improves accuracy by 15.0% on GSM8k, 8.0% on

MATH, and 4.0% on HumanEval. To investigate the source of these improvements, we analyze the model outputs on these benchmarks. We observe that the set of correctly answered questions by Nexus is almost a strict superset of those answered correctly by AdamW. Specifically, on GSM8k and HumanEval, Nexus retains a >95% retention rate on the questions already solved by AdamW, while the 15.0% net improvement stems entirely from exclusively solving previously failed questions. This additive behavior indicates that the performance gain provided by Nexus over the base optimizer is highly stable, expanding the capability boundaries without regressing on previously learned knowledge.

4.3

Scaling Analysis on Model Size

Motivation. In the following two subsections, we investigate the scalability of Nexus across model size and

training duration (tokens). Prevailing literature on implicit bias suggests that the role of implicit regularization becomes increasingly prominent with greater overparameterization and extended computational budgets, since sufficient expressive power and optimization steps grant the model the flexibility to satisfy the geometric implicit bias without compromising the minimization of the pretraining loss [3, 27, 28, 32, 36, 44, 48]. Since Nexus operates via such implicit bias, we hypothesize that its downstream generalization benefits will also amplify at larger compute and model scales. Settings. We evaluate models across five distinct sizes as outlined in Sec. 4.1. Please refer to Sec. H.1 for the

detailed hyperparameters of each experiment. The results are shown in Tab. 9 and Fig. 5. 10

Loss ( )

1.4 1.3

35 30

2.0

130M

300M

520M

Model Size

1.2B

2.3B

(a) Downstream Loss

+1.7%

+1.5%

1.5

+0.8%

1.0

1.1

+2.6%

2.5

1.2

+3.2%

3.0

Nexus over AdamW ( %)

1.5

3.5

AdamW Nexus

Avg. Benchmark Accuracy (%)

AdamW (OOD) Nexus (OOD) AdamW (GSM8k) Nexus (GSM8k) AdamW (HumanEval) Nexus (HumanEval)

1.6

0.5

130M

300M

520M

Model Size

1.2B

0.0

2.3B

(b) Downstream Benchmark

130M

300M

520M

Model Size

1.2B

2.3B

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. Table 2 Scaling Analysis on Training Tokens. We extend the pretraining duration of the 3B model from 2× to 4× Chinchilla optimal tokens. The results demonstrate that the downstream performance advantage of Nexus over the AdamW baseline persists strictly. Chinchila Optim.

Metric

Loss Metrics (↓)

Gen.

Reasoning

Math

Code

Avg.

Pretrain. OOD MMLU GPQA GPQA-D BBH GSM8k MATH HumanEval MBPP

Acc. (↑) AdamW Loss (↓)

1.606

1.302

Acc. (↑) Loss (↓)

1.602

1.290

22.6 1.910

All

36.6 1.534

44.0 1.259

32.0 1.054

43.0 1.116

38.0 1.922

37.1 1.633

59.0 1.227

40.0 1.026

47.0 1.086

38.0 1.921

40.3 1.601

Improv. Loss (↑) +0.004 +0.012 +0.086 +0.024 +0.029 +0.030 +0.032 +0.028

+0.030

+0.001 +0.032

Nexus

47.8 2.265

32.8

2.005

48.9

29.6

23.4

36.6

2.179

1.981

1.881

1.504

AdamW

Acc. (↑) Loss (↓)

1.591

1.293

48.3 2.240

1.975

21.9 1.880

35.2 1.513

54.0 1.245

33.0 1.038

45.0 1.119

43.0 1.976

38.0 1.623

Nexus

Acc. (↑) Loss (↓)

1.588

1.281

52.8 2.216

25.0 1.863

44.1 1.501

62.0 1.229

33.0

1.957

1.008

49.0 1.087

47.0 1.885

41.7 1.593

Improv. Loss (↑) +0.003 +0.012 +0.024 +0.018

+0.017

+0.012 +0.016 +0.030

+0.032

+0.091 +0.030

23.4

20.3

Universal ‘‘Same Pretraining Loss, Better Downstream’’. Across all model sizes ranging from 130M to 2.3B,

Nexus consistently maintains the pretraining validation loss within a negligible margin (defined as ∆ < 0.01 in Sec. 4.1) compared to the baseline, satisfying "same pretraining loss." Despite this parity in pretraining loss, Nexus achieves non-trivial loss reduction on nearly all downstream tasks. For instance, at the 1.2B scale, while the validation loss difference is merely 0.007, Nexus reduces MMLU loss by 0.086, and both BBH and HumanEval losses by 0.03, more than 7 times larger than the pretraining loss gap. Performance Gains Amplify with Scale. We observe that the relative advantage of Nexus over the AdamW

baseline expands monotonically as model capacity increases. Specifically, the average benchmark accuracy improvements across the five evaluated scales are +0.8% (130M), +1.5% (300M), +1.7% (520M), +2.6% (1.2B), and +3.2% (2.3B). This amplification is particularly pronounced in complex reasoning tasks: the accuracy gap on GSM8k widens from negligible levels at the 130M scale to +15.0% (59.0 vs. 44.0) at the 2.3B scale, accompanied by a 0.032 reduction in downstream loss. These results demonstrate that Nexus scales favorably with model capacity, effectively leveraging the increased expressive power to enforce the geometric closeness bias.

Nexus over AdamW ( %)

+3.4%

4 3

+11.5%

+1.1%

6 4

+3.0%

+2.5%

40%

Math Data Proportion

70%

(a) General Domain (MMLU)

+2.0%

+1.5%

2 0

10%

+5.2%

Nexus over AdamW ( %)

+5.8%

0 10%

40%

Math Data Proportion

(b) Math Domain (Avg.)

70%

10%

40%

Math Data Proportion

70%

Figure 6 Results on varying data mixtures (3B models). As the proportion of math data increases (10% → 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.

4.4

Scaling Analysis on Training Tokens

Settings. To evaluate scalability with respect to compute, we extend the training duration of the 3B model

from the standard 2× Chinchilla optimal token count to 4× Chinchilla optimal (i.e., doubling the original training time). All other configurations, including the data mixture, model architecture, and base optimizer hyperparameters, remain strictly identical to those in the main experiments in Sec. 4.1. The advantage of Nexus does not diminish with more training tokens. As shown in Table 1, while the AdamW

baseline naturally improves with extended training (average accuracy increasing from 37.1 to 38.0), it still fundamentally lags behind Nexus. Notably, the overall performance gap between Nexus and AdamW does not shrink with more tokens; Nexus at 4× Chinchilla achieves an average accuracy of 41.7, effectively maintaining and even slightly widening its substantial lead over the baseline. This confirms that the current implicit bias of standard SGD is insufficient to naturally reach optimal geometric closeness, making Nexus’s explicit regularization strictly necessary even under extended compute budgets.

4.5

Robustness to Data Mixing

Motivation. In Sec. 3.2 and Eq. (8), we show that the gradient similarity implies the marginal gains on task i when optimizing on task j (in the first-order sense):

= γ∇Li (θ)⊤ ∇Lj (θ) + O(γ 2 ).

Li (θ) − Li (θ − γ∇Lj (θ)) {z } |

(13)

decrease of task i after one GD step on task j

Since Nexus encourages gradient similarity across the training set, optimizing a sample-dense domain implicitly optimizes sample-sparse domains. Therefore, we conjecture that Nexus acts like a dynamic data mixture, which boosts the sample-sparse or harder-to-learn domains within the mixture without manual re-weighting. Setup. To validate our hypothesis, we construct three distinct data mixtures by explicitly anchoring the

sampling weight of the mathematics domain to 10%, 40%, and 70% (denoted as Math10, Math40, and Math70). Accordingly, we downsample the remaining data sources to fulfill the complementary proportion (e.g., Math70 consists of 70% math and 30% downsampled other data). We train 3B models on each mixture using both AdamW and Nexus, strictly adhering to the hyperparameter settings detailed in Sec. 4.1. Results. As shown in Fig. 6 and Tab. 7, increasing the proportion of math data from 10% to 70% gradually

reduces Nexus’s relative gain on math reasoning. Conversely, as general data becomes the relative minority, Nexus yields a larger improvement in this domain, increasing its gain from +1.1% to +5.8%. Interestingly, the gain on coding tasks exhibits a non-monotonic trend, which we hypothesize is because code generation is a composite capability requiring a complex balance of both logical reasoning and domain knowledge. Furthermore, Nexus mitigates the performance fluctuations observed in the baseline across these mixture

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. Schedule Optim.

WSD

Cosine

Metric

Loss Metrics (↓)

Eval

Gen.

Reasoning

Math

Code

Avg.

OOD MMLU GPQA GPQA-D BBH GSM8k MATH HumanEval MBPP

AdamW

Acc. (↑) 1.606 Loss (↓)

1.302

Nexus

Acc. (↑) 1.602 Loss (↓)

1.290

22.6 1.910

36.6 1.534

44.0 1.259

32.0 1.054

43.0 1.116

38.0 1.922

37.1 1.633

23.4 1.881

36.6 1.504

59.0 1.227

40.0 1.026

47.0 1.086

38.0 1.921

40.3 1.601

Improv. Loss (↑) +0.004 +0.012 +0.086 +0.024 +0.029 +0.030 +0.032 +0.028

+0.030

+0.001 +0.032

AdamW

Acc. (↑) 1.526 Loss (↓)

1.255

Nexus

Acc. (↑) 1.528 Loss (↓)

1.250

47.8 2.265

All

32.8

2.005

48.9

29.6

2.179

1.981

53.2 2.195

1.480

60.0 1.212

32.0 1.022

56.0 1.045

39.0 1.867

41.0 1.572

34.8 1.479

59.0

41.0 0.994

54.0

1.169

1.025

46.0 1.805

43.4 1.541

Improv. Loss (↑) -0.002 +0.005 +0.080 +0.007 +0.003 +0.001 +0.043 +0.028

+0.020

+0.062 +0.030

54.9 2.115

26.6 1.924

19.5 1.829

30.5

27.3

1.917

1.826

41.5

shifts. These results support our conjecture that Nexus acts as an implicit balancer, dynamically prioritizing under-optimized tasks without manual mixture tuning.

4.6

Robustness to Learning Rate Schedule

Motivation. While the Warmup-Stable-Decay (WSD) scheduler [15] has become increasingly popular in recent

LLM pretraining, the Cosine annealing schedule remains a widely adopted standard [45, 46]. To ensure that our observed generalization benefits are not merely an artifact of a specific learning rate dynamic, we evaluate the robustness of Nexus across different schedulers. Settings. We conduct an ablation study by replacing the default WSD scheduler with a standard Cosine

learning rate scheduler. All other training configurations, including the 3B model architecture, data mixture, and base optimizer hyperparameters, remain strictly identical to the main setup detailed in Sec. 4.2. Results. As demonstrated in Tab. 3, the "same pretraining loss, better downstream performance" phenomenon

persists consistently across both schedulers. Under the Cosine schedule, Nexus maintains a negligible pretraining loss difference compared to the AdamW baseline (1.528 vs. 1.526) while delivering substantial improvements on downstream metrics, such as a +0.03 loss gain on downstream benchmarks. This confirms that the implicit bias introduced by Nexus is highly robust and orthogonal to the choice of learning rate trajectory.

Discussions

In this section, we conduct several interesting ablation studies of Nexus.

5.1

Experimental Validation of Our Theory

Settings. To validate our theory, we analyze the training trajectories of the 3B AdamW and 3B Nexus

models from Sec. 4.2. During pretraining, we record the gradient cosine similarity between test set and each downstream corpus every 1,000 steps and compute the average to approximate the averaged gradient similarity during training. Upon the completion of pretraining, we perform full batch Gradient Descent using AdamW with learning rate 2 × 10−5 and weight decay 0 on each downstream task LT to locate the respective task-specific minimizer θT∗ for subsequent visualization and distance evaluation. Nexus Encourages Training Set Closeness. As demonstrated in Tab. 4, Nexus effectively increases the gradient

similarity across the pretraining set, Ei̸=j [CosSim(∇Li , ∇Lj )], compared to the base optimizer, as analyzed in Theorem B.3 and 3.2. 13

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. Pretrain Set

OOD Set

AdamW Nexus

0.4499

0.2228

0.0824

0.4661

0.2464

0.0924

0.0325

Param. Closeness. (↓)

AdamW Nexus

1.452

2.812

4.500

3.418

3.775

4.472

1.441

2.806

4.482

3.326

3.766

4.444

3.648

Loss (↓)

AdamW Nexus

1.606 1.602

1.302

1.910

1.259

1.054

1.116

1.290

1.881

1.227

1.026

1.086

1.922 1.921

Benchmark (↑)

AdamW Nexus

22.6 23.4

44.0

32.0

43.0

59.0

40.0

47.0

Metric

Optim.

Grad Sim. (↑)

GPQA-D

GSM8k

Math500

HumanEval

MBPP

0.0374

0.0422

0.0367

0.0091

0.0427

0.0382

0.0092 3.645

38.0 38.0

Training Set Closeness Generalizes to Downstream Closeness. Fortunately, this gradient closeness generalizes

beyond the pretraining corpus to unseen downstream tasks T , effectively increasing the similarity between the training objective and the downstream task, CosSim(∇Ltrain , ∇LT ).

Downstream Closeness Yields Smaller Downstream Loss and Better Performance. Since this gradient

similarity CosSim(∇Ltrain , ∇LT ) generalizes, optimizing the pretraining objective inherently optimizes the downstream tasks, as indicated by the first-order approximation in Eq. (8). This gradient closeness translates into lower downstream losses and better benchmark performance.

∗ Empirical Landscapes Align with Fig. 2. As shown in Fig. 7c, Nexus reduces the downstream loss LT (θtrain )

∗ by decreasing the geometric distance ∥θT∗ − θtrain ∥2 between the converged parameter and the task-specific minimizer. This observation matches the analyses in Theorem 2.2 and 2.3. While Nexus reduces this distance, it does not cause all minima to nearly intersect as depicted in Fig. 2b—which would theoretically yield a nearly 0% OOD generalization error. Instead, it achieves a moderate reduction in geometric distance, leading to a proportionally lower downstream loss. We hope future work can design stronger Nexus variants capable of approaching this extreme closeness without introducing significant computational overhead.

5.2

Implicit Biases of Other Optimizers

Motivation. To explicitly demonstrate the "same pretraining loss, better downstream task" phenomenon

and analyze the implicit biases of different optimizers, we visualize the correlation between pretraining and downstream losses, using the results of AdamW, Muon, and AdamW-Nexus from Sec. 4.2. We plot the averaged downstream loss (y-axis) against the pretraining validation loss (x-axis) at corresponding checkpoints. The results are presented in Fig. 7a. Muon does not possess a superior implicit bias. As illustrated in Fig. 7a, the curves for Adam and pure Muon

almost completely overlap. This indicates that despite its orthogonalization mechanism, Muon seems not to introduce a favorable implicit bias for downstream generalization beyond what is explained by the pretraining loss itself. This observation aligns with the findings in Wen et al. [45], which suggest that for Muon-like optimizers, achieving the same pretraining loss typically translates to the same downstream performance. While recent work [43] demonstrates that Muon tends to optimize towards representations with a higher weight rank than Adam, empirical results suggest that this structural difference in weight matrices seems not inherently translate into observable generalization benefits on downstream tasks. Implicit bias does not stem from gradient normalization. To further isolate the source of Nexus’s generalization

benefits, we conduct an ablation study where the normalized gradient g/∥g∥2 is directly fed into the Adam optimizer instead of the raw gradient g. The results are shown as the NSGD curve in Fig. 7a. We observe that this variant still fails to introduce any favorable implicit bias, which closely overlaps with the standard Adam baseline. This indicates that the downstream gains of Nexus do not originate from the mere act of 14

AdamW Muon Nexus NSGD

Averaged Downstream Loss

1.75

1.775 1.750 1.725

1.65

1.50

1.650

1.45

1.625

1.40

1.600

1.60

1.65

1.70

1.75

1.80

1.85

Pretraining Loss

1.90

1.95

2.00

(a) Implicit Biases

2.50

1.75

1.675

1.55

HumanEval (AdamW) HumanEval (Nexus) GSM8K (AdamW) GSM8K (Nexus)

2.75

2.00

1.700

1.60

3.00

2.25

Pretraining Loss

1.70

Adam Nexus Nexus-Dot

1.800

Downstream Loss

1.80

1.50 1.25 1.00

102000

104000

106000

Step

108000

110000

(b) Nexus-Dot

1.0

0.5

0.0

* 2 Distance to train

0.5

1.0

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.

normalizing gradients. Mathematically, this ablation is strictly equivalent to executing the Nexus algorithm with an inner loop step count of K = 1. This empirical observation perfectly aligns with Theorem 3.2: when K = 1, the coefficient of the gradient similarity regularizer K−1 4K becomes strictly zero, stripping the optimizer of its consensus-seeking property and reducing it to a purely first-order method.

5.3

Cosine Similarity Instead of Dot Product Similarity

Although as discussed in Sec. 3, the dot product similarity of gradients offers a more direct theoretical connection—yielding a tighter bound for parameter closeness (Theorem 3.1) and a more straightforward interpretation for downstream generalization (Eq. (8))—it proves practically challenging to optimize. This difficulty primarily arises because the dot product objective introduces a pathological optimization shortcut. Specifically, the dot product is highly scale-dependent: if the overall loss magnitude scales by a factor of k, the gradient norm scales proportionally by k, causing the dot product similarity to artificially inflate by a factor of k 2 . Consequently, directly maximizing the dot product severely disrupts the primary minimization of the pretraining loss, as the optimizer may exploit this shortcut by inadvertently increasing the gradient norms rather than discovering genuine task consensus. As demonstrated in Fig. 7b, the optimization trajectory of Nexus-Dot lags significantly behind the standard Adam baseline. The resulting degradation in pretraining loss heavily outweighs any potential generalization benefits conferred by its implicit bias. Therefore, we adopt cosine similarity (via normalized gradients) as our primary regularization objective in this work. Note that the progressive deceleration of Nexus-Dot observed in Fig. 7b is a persistent geometric phenomenon, occurring consistently regardless of the learning rate scheduler or the choice of base optimizer (e.g., AdamW or Muon). Due to space constraints, we selectively present the ablation results for the 3B model with Adam, corresponding to the main setup in Sec. 4.2.

Conclusion and Limitation

In this work, we investigate the geometric closeness of minimizers of different losses in LLM pretraining. We show that this closeness strongly correlates with downstream generalization. To optimize this closeness, we propose the Nexus algorithm, which encourages gradient similarity across different tasks. We show that both gradient closeness and geometric closeness generalize to downstream tasks, thus leading to lower downstream loss and better downstream performance. Experimental results across various settings validate our claims. We reckon that as the LLM scaling paradigm transitions from being compute-bound to data-bound, explicitly engineering the implicit biases of optimizers to unlock generalization may serve as a critical frontier for developing more capable language models. Limitations. Despite its empirical success and theoretical consistency on AdamW, Nexus currently remains

incompatible with the Muon optimizer. Specifically, Muon combined with Nexus even underperforms the

AdamW-Nexus configuration on downstream tasks, due to its deceleration on Muon (in contrast to the slight acceleration observed with AdamW as demonstrated in Sec. B.1). We hypothesize this may be due to several subtle factors, such as numerical sensitivities involving the pseudo-gradient coefficient γ (see Eq. (11)) or potential interactions arising from the Newton-Schulz iterations. We are currently investigating these challenges and aim to resolve this incompatibility in future work.

Acknowledgement This work was conducted for research and validation purposes only. The algorithms and methodologies described herein are experimental prototypes and have not been integrated into any commercial products or services of the affiliated organizations. We gratefully acknowledge the support of the National Science Foundation (Grant 625B2104). We also thank Kaiyue Wen, Haodong Wen, Yan Wu, Jianhui Duan, Chengyin Xu, Kaiyuan Chen for their insightful comments and helpful discussions.

References [1] Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732, 2021. [2] Aarti Basant, Abhijit Khairnar, Abhijit Paithankar, Abhinav Khattar, Adithya Renduchintala, Aditya Malte, Akhiad Bercovich, Akshay Hazare, Alejandra Rico, Aleksander Ficek, et al. Nvidia nemotron nano 2: An accurate and efficient hybrid mamba-transformer reasoning model. arXiv preprint arXiv:2508.14444, 2025. [3] Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine-learning practice and the classical bias–variance trade-off. Proceedings of the National Academy of Sciences, pages 15849–15854, 2019. [4] Huanran Chen, Yinpeng Dong, Zeming Wei, Yao Huang, Yichi Zhang, Hang Su, and Jun Zhu. Understanding pre-training and fine-tuning from loss landscape perspectives. arXiv preprint arXiv:2505.17646, 2025. [5] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021. [6] Jeremy Cohen, Alex Damian, Ameet Talwalkar, J Zico Kolter, and Jason D Lee. Understanding optimization in deep learning with central flows. In The Thirteenth International Conference on Learning Representations, 2025. [7] Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, pages 27449–27461, 2021. [8] Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024. [9] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022. [10] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. [11] Mark Chen et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021. [12] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020. [13] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300, 2020. [14] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874, 2021. [15] Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06395, 2024. [16] Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. Advances in Neural Information Processing Systems, 2013. [17] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https://kellerjordan.github.io/posts/muon/. [18] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. [19] Konwoo Kim, Suhas Kotha, Percy Liang, and Tatsunori Hashimoto. Pre-training under infinite compute. arXiv preprint arXiv:2509.14786, 2025. [20] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.

[21] Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning, pages 5905–5914, 2021. [22] Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in Neural Information Processing Systems, 31, 2018. [23] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024. [24] Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. In International Conference on Machine Learning, pages 22188–22214. PMLR, 2023. [25] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. [26] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017. [27] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2019. [28] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614, 2014. [29] Jinjie Ni, Qian Liu, Longxu Dou, Chao Du, Zili Wang, Hang Yan, Tianyu Pang, and Michael Qizhe Shieh. Diffusion language models are super data learners. arXiv preprint arXiv:2511.03276, 2025. [30] Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia, Yuling Gu, Shengyi Huang, Matt Jordan, Nathan Lambert, Dustin Schwenk, Oyvind Tafjord, Taira Anderson, David Atkinson, Faeze Brahman, Christopher Clark, Pradeep Dasigi, Nouha Dziri, Allyson Ettinger, Michal Guerquin, David Heineman, Hamish Ivison, Pang Wei Koh, Jiacheng Liu, Saumya Malik, William Merrill, Lester James V. Miranda, Jacob Morrison, Tyler Murray, Crystal Nam, Jake Poznanski, Valentina Pyatkin, Aman Rangapur, Michael Schmitz, Sam Skjonsberg, David Wadden, Christopher Wilhelm, Michael Wilson, Luke Zettlemoyer, Ali Farhadi, Noah A. Smith, and Hannaneh Hajishirzi. 2 olmo 2 furious, 2025. [31] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017. [32] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022. [33] Mihir Prabhudesai, Mengning Wu, Amir Zadeh, Katerina Fragkiadaki, and Deepak Pathak. Diffusion beats autoregressive in data-constrained settings. arXiv preprint arXiv:2507.15857, 2025. [34] David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R Bowman. Gpqa: A graduate-level google-proof q&a benchmark. In First Conference on Language Modeling, 2024. [35] Andrei Semenov, Matteo Pagliardini, and Martin Jaggi. Benchmarking optimizers for large language model pretraining. arXiv preprint arXiv:2509.01440, 2025. [36] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 2018. [37] Jacob Mitchell Springer, Sachin Goyal, Kaiyue Wen, Tanishq Kumar, Xiang Yue, Sadhika Malladi, Graham Neubig, and Aditi Raghunathan. Overtrained language models are harder to fine-tune. arXiv preprint arXiv:2503.19206, 2025. [38] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929–1958, 2014. [39] Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, , and Jason Wei. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261, 2022. [40] ByteDance Seed Team. Seed-oss open-source models. https://github.com/ByteDance-Seed/seed-oss, 2025.

[41] Qwen Team. Qwen3 technical report, 2025. URL https://arxiv.org/abs/2505.09388. [42] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. [43] Shuche Wang, Fengzhuo Zhang, Jiaxiang Li, Cunxiao Du, Chao Du, Tianyu Pang, Zhuoran Yang, Mingyi Hong, and Vincent YF Tan. Muon outperforms adam in tail-end associative memory learning. arXiv preprint arXiv:2509.26030, 2025. [44] Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint arXiv:2211.05729, 2022. [45] Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them. arXiv preprint arXiv:2509.02046, 2025. [46] Kaiyue Wen, Zhiyuan Li, Jason S Wang, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Understanding warmup-stable-decay learning rates: A river valley loss landscape view. In The Thirteenth International Conference on Learning Representations, 2025. [47] An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2.5 technical report. arXiv preprint arXiv:2412.15115, 2024. [48] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:1611.03530, 2016. [49] Yihao Zhang, Hangzhou He, Jingyu Zhu, Huanran Chen, Yifei Wang, and Zeming Wei. On the duality between sharpness-aware minimization and adversarial training. arXiv preprint arXiv:2402.15152, 2024.

Appendix A

Notations and Assumptions

To facilitate the theoretical analysis in the subsequent sections, we summarize the key mathematical notations and fundamental optimization assumptions used throughout this paper.

A.1

Notations

The primary mathematical notations for data mixtures, loss functions, geometries, and optimization dynamics are summarized in Tab. 5. Table 5 Summary of key notations used in this paper. Notation

Description

Data and Loss Functions K αk Lk (θ) Ltrain (θ) LT (θ)

Total number of distinct pretraining data sources (tasks). Sampling probability (data mixing ratio) for the k-th data source. The expected / empirical loss on the k-th source task. PK 1 The averaged pretraining loss: K k=1 Lk (θ). The loss on an unseen downstream evaluation task T . Geometric and Statistical Variables

∗ θtrain

Sk , ST θk∗ , θT∗ µ σ2

The converged parameter state that minimizes Ltrain . The set of local minimizers for task k and downstream task T , respectively. The specific local minimizer in Sk or ST closest to the current parameter. The statistical center of task-specific minimizers: E[θk∗ ]. The intrinsic variance (Closeness) of task-specific minimizers: E[∥θk∗ − µ∥22 ]. Optimization and Nexus Variables

γ ĝt CosSim(x, y) Sij (θ)

A.2

The inner learning rate (step size) used in the Nexus gradient approximator. The Nexus pseudo-gradient (displacement) passed to the outer optimizer. ⊤ y The cosine similarity between two vectors: ∥x∥x2 ∥y∥ . 2 Shorthand for gradient similarity: CosSim(∇Li (θ), ∇Lj (θ)).

Assumptions

The main assumptions used in our analysis are outlined below [6, 46]. Additional assumptions required for specific analyses will be stated in the respective theorems. Assumption A.1 (Bounded Gradients). For all tasks i ∈ [1, K] and parameters θ along the optimization trajectory, the gradient norm is strictly bounded from below and above: 0 < Gmin ≤ ∥∇Li (θ)∥2 ≤ G This ensures that the Normalized SGD step in Nexus is always well-defined and numerically stable. Assumption A.2 (Smoothness and Bounded Curvature). The loss function Li is L-smooth, meaning its Hessian spectral norm is bounded from above. Furthermore, within the local basin of attraction [θ, θk∗ ], the curvature is strictly lower-bounded by λmin > 0: λmin ≤ inf ∗ u⊤ ∇2 Li (ξ)u ≤ ∥∇2 Li (θ)∥2 ≤ L ξ∈[θ,θk ]

where u is any unit vector. Prior literature extensively characterizes the local loss landscape of deep neural networks as exhibiting high quadraticity, particularly along meaningful optimization trajectories [4, 22, 44]. 20

Consequently, under the standard premise that the loss landscape can be locally and directionally approximated by a quadratic function, this bounded curvature condition should not be viewed as a restrictive assumption. Assumption A.3 (Hessian Lipschitz Continuous). The Hessian matrix is ρ-Lipschitz continuous. For any parameters x, y: ∥∇2 Li (x) − ∇2 Li (y)∥2 ≤ ρ∥x − y∥2 This assumption is necessary to bound the Jacobian of the normalized gradient during the second-order Taylor expansion in Nexus’s inner loop.

Additional Discussions

B.1

Convergence Rate of Nexus

All of our analyses are based on the assumption that Nexus should not be slower than its base optimizer. This ensures that both can achieve the "same training loss," allowing the implicit bias of Nexus to subsequently achieve "better downstream performance." One might concern that since Nexus optimizes two joint objectives (see Theorem B.3 and 3.2), it may be slower than its base optimizer. Consequently, the downstream gains might not offset the speed loss, potentially leading to worse overall downstream performance. Fortunately, this concern does not hold in practice. Empirically, across all experiments, Nexus is not slower, and sometimes even slightly faster, than its base optimizer (see Sec. 4). Intuitively, Nexus makes the gradients of each Li similar; thus, optimizing Li effectively optimizes Lj simultaneously, as analyzed in Sec. 3.2. This "constructive interference" can lead to slightly faster convergence. We can also adopt the framework of Johnson and Zhang [16] (assuming each Li is L-smooth and µ-strongly convex) to obtain further theoretical intuition. In this setting, standard SGD typically achieves only an O(1/T ) convergence rate. However, if Nexus succeeds in finding a region where these tasks share common minimizers, it can achieve exponential convergence: Theorem B.1. Suppose each Li is L-smooth and µ-strongly convex. That is, for any θ1 , θ2 , we have: L ∥θ1 − θ2 ∥22 , 2 µ Li (θ1 ) ≥ Li (θ2 ) + ∇Li (θ2 )⊤ (θ1 − θ2 ) + ∥θ1 − θ2 ∥22 . 2

Li (θ1 ) ≤ Li (θ2 ) + ∇Li (θ2 )⊤ (θ1 − θ2 ) +

(14)

Additionally, assume there exists a common minimizer θ ∗ such that ∇Li (θ ∗ ) = 0 for all i ∈ [K]. Then, for 2 the sequence {θ0 , θ1 , . . . , θT } generated by Nexus with step size γ ∈ (0, L+µ ), we have: T 2γµL E[∥θT − θ ∥ ] ≤ 1 − ∥θ0 − θ ∗ ∥2 . L+µ ∗ 2

(15)

2 Specifically, setting γ = L+µ and defining the condition number κ = L/µ, we obtain the convergence rate:

E[∥θT − θ ∗ ∥2 ] ≤

κ−1 κ+1

∥θ0 − θ ∗ ∥2 .

(16)

Therefore, if Nexus guides the parameters into a locally convex and smooth regime where a common minimizer exists, it guarantees exponential convergence.

B.2

Implicit Bias of Normalized SGD

Interestingly, Nexus also offers a novel perspective on the success of Normalized SGD (NSGD). We observe that NSGD can be mathematically interpreted as a special case of Nexus, revealing that NSGD does not merely minimize the scalar loss but also implicitly optimizes gradient closeness. This implicit regularization provides a geometric explanation for why NSGD often generalizes better than standard Gradient Descent. 21

Theorem B.2 (Implicit Bias of NSGD). Let {θt } be the sequence generated by Normalized SGD with learning rate γ. This sequence implicitly minimizes the following expected joint objective: γ ′ E i.i.d [CosSim(∇L(x; θ), ∇L(x ; θ))], 8 x,x′ ∼ D 2 min γ3. subject to a discretization error bounded by 34 4L G+ρG 2

(17)

JNSGD (θ) = Ex∼D [L(x; θ)] −

min

Proof. A sequence of n updates of NSGD is algebraically equivalent to performing Nexus with k inner steps (using NSGD) and n/k outer steps (using SGD with step size 1). By Theorem 3.2, the magnitude of the gradient alignment signal scales as S(k) ≈ γ 2 k(k−1) , while the residual error is bounded by N (k) ≤ 16 4

4L2 +ρGmin G2min

k3 γ 3 .

Defining the signal-to-noise ratio as ρ(k) ≜ S(k)/N (k) and maximizing it with respect to k yields k = 2. Thus, viewing the NSGD updates through the lens of Nexus with k = 2 yields the stated results.

B.3

Other Approximators for Hessian Gradient Product

While the Hessian-vector product can theoretically be implemented via the Jacobian-vector product (JVP) in PyTorch [31] with only a constant factor of computational overhead, implementing exact Hessian-gradient products in practical LLM pretraining remains prohibitive. First, standard Hessian-vector product implementations are often incompatible with memory-efficient kernels like FlashAttention [8, 9] (which typically do not support second-order differentiation efficiently), leading to significantly higher memory usage and computational costs. Second, the constant margin of memory overhead poses significant infrastructure challenges for large-scale distributed training. Moreover, the Nexus algorithm exhibits a beneficial third-order effect. It actively seeks regions where gradients are not only aligned but also locally flat along the gradient dimension. This ensures that the gradient alignment property remains stable across a larger regime. Theorem B.3 (Nexus Maximizes Stability of Closeness). Assume the existence of constants Gmin , L, ρ > 0 as in Theorem 3.2, and let M3 be a constant such that ∇3 Li (θ)[u, v, w] ≤ M3 for any unit vectors u, v, w. Then, the sequence {θt } generated by Algorithm 1 effectively minimizes the following third-order objective: J3rd (θ) = J2nd (θ) + γ 3

(K − 1)(2K − 1) X ∇Li (θ)⊤ ∇2 Lj (θ)∇Lp (θ). 12K 2 i,j,p

(18)

The approximation error is bounded by: E 3rd ≜ ∥E[ĝt ] − ∇J3rd ∥ ≤

M3 M3 L + 24 8Gmin

K 4γ4 +

M3 L2 5 5 K γ = O(γ 4 ). 40G2min

(19)

Therefore, the third-order effect of Nexus works like a kind of "Multi-Task SAM": it minimize the directional sharpness along different tasks, leading to flatter landscape.

Proofs for Closeness Improving Generalization

C.1

Proof for Theorem 2.2

Proof. First, solving the stationarity condition ∇ PK 1 ∗ ∗ converged parameter: θtrain =K k=1 θk .

Lk (θ) = 0, we obtain the closed-form solution for the

The training loss at this optimum is given by: K

∗ Ltrain (θtrain )=

1 X a ∗ ∥θtrain − θk∗ ∥22 + c = Ctrain . K 2

(20)

k=1

From this, we can express the intrinsic loss constant c (which represents the "depth" of the minima) in terms of the fixed training loss Ctrain : K a X ∗ c = Ctrain − ∥θtrain − θk∗ ∥22 . (21) 2K k=1

Now, consider the loss on a new downstream task T with minimizer θT∗ ∼ P: ∗ LT (θtrain )=

a ∗ ∥θ − θT∗ ∥22 + c. 2 train

(22)

Substituting c, the generalization gap becomes: a ∗ LT (θtrain ) − Ctrain = 2

∗ ∥θtrain − θT∗ ∥22 −

1 X ∗ ∥θtrain − θk∗ ∥22 K

! .

(23)

k=1

Taking the expectation over the task distribution the property of variance for i.i.d. samples PP, ∗and utilizing 1 1 ∗ 2 (where E[∥θtrain − θT∗ ∥2 ] = (1 + K )σ 2 and E[ K ∥θtrain − θk∗ ∥2 ] = K−1 K σ ): a 1 K −1 2 a ∗ 2 E[LT (θtrain )] − Ctrain = 1+ σ − σ = σ2 . (24) 2 K K K This concludes the proof. It explicitly shows that for a fixed training loss budget Ctrain , the generalization error scales linearly with the task variance σ 2 .

C.2

Proof for Theorem 2.3

We now generalize the previous result to the general case. Assume that the pretraining tasks {Lk }K k=1 and the downstream task LT are sampled independently from a latent task distribution P. Due to the over-parameterized nature of LLMs, the minimizers are not unique. To rigorously analyze the closeness, we first define the set of local minimizers for the expected population loss: SP = {ϑ | ∃ϵ > 0, ∀ϑ′ ∈ Bϵ (ϑ), ET ∼P [LT (ϑ)] ≤ ET ∼P [LT (ϑ′ )]} .

(25)

Let θ ∗ ∈ SP be one specific local minimizer of the population loss. This serves as the anchor point for the basin of attraction. We then define the task-specific minimizer θk∗ as the projection of this population minimizer θ ∗ onto the set of local minimizers of task k: θk∗ = arg min ∥ϑ − θ ∗ ∥2 , ϑ∈Sk

where Sk denotes the set of local minimizers of Lk .

(26)

Given the distribution of these task-specific minimizers {θk∗ }, we define their statistical center µ and intrinsic covariance Σ as: µ := ET ∼P [θT∗ ], Σ := E[(θT∗ − µ)(θT∗ − µ)⊤ ]. (27) We also define the scalar intrinsic variance σ 2 = Tr(Σ) = E[∥θk∗ − µ∥22 ]. From this point forward, our analysis focuses on the closeness to the statistical center µ, as E[θT∗ − µ] = 0 holds by definition. 23

Step 1: Estimation Error.

∗ The converged parameter θtrain satisfies the stationarity condition:

∗ ∇Ltrain (θtrain ) = 0 ⇐⇒

K X

∗ ∇Lk (θtrain ) = 0.

(28)

k=1 ∗ ∗ ∗ Applying the Mean Value Theorem, there exists ξk ∈ [θtrain , θk∗ ] such that ∇Lk (θtrain ) = ∇2 Lk (ξk )(θtrain −θk∗ ). Thus: K K X X ∗ ∇2 Lk (ξk )(θtrain − µ) = ∇2 Lk (ξk )(θk∗ − µ). (29) k=1

k=1

We assume the local curvature is bounded: for any k and vector u, λmin ∥u∥2 ≤ u⊤ ∇2 Lk (ξk )u ≤ λmax ∥u∥2 . Bounding the estimation error norm: K

∗ ∥θtrain − µ∥2 ≤

1 X λmax ∥θk∗ − µ∥2 . Kλmin

(30)

k=1

Taking the expectation (noting cross-terms vanish because E[θk∗ − µ] = 0) and defining κ = λmax /λmin : ∗ E[∥θtrain − µ∥22 ] ≤

κ2 2 σ . K

(31)

Step 2: The Intrinsic Loss Trade-off. We condition on the training loss achieving a fixed value Ctrain . By exact Taylor expansion around the task minimizers, the training loss is: K K 1 X 1 ∗ 1 X ∗ ∗ ∗ ⊤ 2 ∗ ∗ Lk (θtrain ) = Lk (θk ) + (θtrain − θk ) ∇ Lk (ξk )(θtrain − θk ) . Ctrain = K K 2 k=1

(32)

k=1

Taking the expectation over the task distribution, we can express the expected intrinsic loss exactly as: " # K 1 1 X ∗ ∗ ∗ ⊤ 2 ∗ ∗ E[Lk (θk )] = Ctrain − E (θtrain − θk ) ∇ Lk (ξk )(θtrain − θk ) . (33) 2 K k=1 {z } | Qtrain (Expected Empirical Closeness Penalty)

We retain the term Qtrain explicitly without approximation. This term represents the curvature-weighted variance of the minimizers around the converged point. Step 3: Downstream Generalization (Rigorous Matrix Derivation). Finally, we analyze the expected performance on a downstream task T sampled from the same distribution P. We perform a Taylor expansion of the test loss around the task-specific minimizer θT∗ . Since ∇LT (θT∗ ) = 0, the first-order term vanishes: 1 ∗ ∗ ∗ LT (θtrain ) = LT (θT∗ ) + (θtrain − θT∗ )⊤ ∇2 LT (ξT )(θtrain − θT∗ ). 2

(34)

Taking the expectation over the task distribution, we define the expected test closeness penalty Qtest : 1 ∗ ∗ ∗ ET [LT (θtrain )] = E[LT (θT∗ )] + E (θtrain − θT∗ )⊤ ∇2 LT (ξT )(θtrain − θT∗ ) . 2 | {z }

(35)

Qtest

Recalling the intrinsic loss trade-off from Eq. (33), we have E[LT (θT∗ )] = Ctrain − Qtrain . Substituting this into the equation above yields the generalization gap decomposition: ∗ ET [LT (θtrain )] = Ctrain + (Qtest − Qtrain ).

(36)

Let H̄ = EP [∇2 L(ξ)] denote the expected Hessian matrix over the task distribution. Since tasks are i.i.d., both training and test tasks share this expected geometry. For the test term Qtest , we use the identity x⊤ Ax = Tr(Axx⊤ ). Replacing the specific task Hessian with the expected Hessian H̄ = EP [∇2 L(ξ)]: Qtest =

∗ 1 ∗ Tr H̄ · E (θtrain − θT∗ )(θtrain − θT∗ )⊤ . 2

(37)

We expand the covariance term fully around the statistical center µ: ∗ ∗ ∗ ∗ E (θtrain − θT∗ )(θtrain − θT∗ )⊤ = E ((θtrain − µ) − (θT∗ − µ))((θtrain − µ) − (θT∗ − µ))⊤ (38)

∗ ∗ = E[(θtrain − µ)(θtrain − µ)⊤ ] + E[(θT∗ − µ)(θT∗ − µ)⊤ ] ∗ ∗ − E[(θtrain − µ)(θT∗ − µ)⊤ ] − E[(θT∗ − µ)(θtrain − µ)⊤ ].

∗ The cross-terms vanish strictly because θT∗ is independent of θtrain and is centered at µ (i.e., E[θT∗ − µ] = 0 by definition of µ). Substituting E[(θT∗ − µ)(θT∗ − µ)⊤ ] = Σ back:

Qtest =

1 1 ∗ ∗ Tr H̄ · E[(θtrain − µ)(θtrain − µ)⊤ ] + Tr(H̄Σ). 2 2

(39)

For the training term Qtrain , we consider the expected quadratic penalty averaged over the training tasks. By linearity of expectation, we replace ∇2 Lk with H̄ exactly: K 1 X ∗ ∗ Qtrain = E (θtrain − θk∗ )⊤ H̄(θtrain − θk∗ ) . 2K

(40)

k=1

We apply the Generalized Centroid Property. For any positive semi-definite matrix H̄, the weighted sum of squared errors is minimized by the mean θ̄. Thus, we have the rigorous lower bound: K X

∗ ∗ (θtrain − θk∗ )⊤ H̄(θtrain − θk∗ ) ≥

k=1

K X

(41)

(θ̄ − θk∗ )⊤ H̄(θ̄ − θk∗ ).

k=1

We perform the matrix variance decomposition on the RHS by inserting µ: K X

(θ̄ − θk∗ )⊤ H̄(θ̄ − θk∗ ) =

k=1

K X

((θ̄ − µ) − (θk∗ − µ))⊤ H̄((θ̄ − µ) − (θk∗ − µ))

k=1

K X

(θ̄ − µ)⊤ H̄(θ̄ − µ) +

k=1

K X

(θk∗ − µ)⊤ H̄(θk∗ − µ) − 2(θ̄ − µ)⊤ H̄

k=1

K X

(θk∗ − µ) .

k=1

K(θ̄−µ)

Simplifying the cross-term and combining with the first term: K X

(θ̄ − θk∗ )⊤ H̄(θ̄ − θk∗ ) = K(θ̄ − µ)⊤ H̄(θ̄ − µ) +

k=1

K X

(θk∗ − µ)⊤ H̄(θk∗ − µ) − K(θ̄ − µ)⊤ H̄(θ̄ − µ).

k=1

Taking expectations and using the trace identity E[x⊤ Ax] = Tr(AE[xx⊤ ]): PK • The first term: k=1 Tr(H̄E[(θk∗ − µ)(θk∗ − µ)⊤ ]) = KTr(H̄Σ). 25

(42)

(θk∗ − µ)⊤ H̄(θk∗ − µ) − 2K(θ̄ − µ)⊤ H̄(θ̄ − µ)

k=1 K X

}

(43)

1 1 • The second term (variance of the mean): E[(θ̄ − µ)(θ̄ − µ)⊤ ] = K Σ. Thus, KTr(H̄ · K Σ) = Tr(H̄Σ).

Combining these, the expected training penalty is bounded by: 1 1 Qtrain ≥ KTr(H̄Σ) − Tr(H̄Σ) = 2K 2

1 1− K

Tr(H̄Σ).

(44)

Subtracting the two terms (Qtest − Qtrain ), the dominant term 12 Tr(H̄Σ) cancels out exactly. We then bound the remaining terms using the spectral norm λmax and the estimation error bound derived in Eq. (31): 1 1 1 1 ∗ ∗ ∗ ET [LT (θtrain )] − Ctrain ≤ Tr(H̄E[(θtrain − µ)(θtrain − µ)⊤ ]) + Tr(H̄Σ) − 1− Tr(H̄Σ) 2 2 2 K 1 1 ∗ ∗ − µ)(θtrain − µ)⊤ ] + Tr(H̄Σ) = Tr H̄ · E[(θtrain 2 2K λmax λmax ≤ E[∥θ ∗ − µ∥22 ] + Tr(Σ) (45) 2 train 2K λmax κ2 2 λmax 2 σ + σ ≤ 2 K 2K λmax (κ2 + 1) 2 = σ . 2K 2

This confirms that the generalization gap scales with O( σK ), driven by the intrinsic task variance and the number of pretraining tasks.

Proof of Theorem 3.1

In this section, we provide the detailed proof for Theorem 3.1, which bounds the closeness between minimizers using gradient similarity. Proof. The proof proceeds in three main steps: (1) relating the closeness to the gradient norm via the Mean Value Theorem; (2) exploiting the stationarity condition of the total loss to decompose the gradient norms; and (3) bounding the cross-terms using the gradient upper bound and cosine similarity. Step 1: Relating Closeness to Gradient Norm. Recall that θk∗ is the projection of θ onto the global optimal set Sk . Since θk∗ is a minimizer, we have ∇Lk (θk∗ ) = 0. Applying the Mean Value Theorem to the vector-valued function ϑ 7→ ∇Lk (ϑ), there exists a point ξk on the line segment connecting θk∗ and θ such that: ∇Lk (θ) − ∇Lk (θk∗ ) = ∇2 Lk (ξk )(θ − θk∗ ).

(46)

Substituting ∇Lk (θk∗ ) = 0 and taking the norm: ∥∇Lk (θ)∥2 = ∥∇2 Lk (ξk )(θ − θk∗ )∥2 .

(47)

We assume the curvature condition where the smallest eigenvalue of the Hessian along the displacement vector is bounded below by λ > 0. Specifically: where uk =

2 u⊤ k ∇ Lk (ξk )uk ≥ λ,

θ − θk∗ . ∥θ − θk∗ ∥2

(48)

This implies ∥∇2 Lk (ξk )(θ − θk∗ )∥2 ≥ λ∥θ − θk∗ ∥2 . Rearranging this inequality gives an upper bound on the closeness: 1 ∥θ − θk∗ ∥2 ≤ ∥∇Lk (θ)∥2 . (49) λ Squaring and averaging over all K tasks yields: K

k=1

1 X 1 X ∥θ − θk∗ ∥22 ≤ ∥∇Lk (θ)∥22 . K Kλ2

(50)

Step 2: Force Balance Decomposition. Since θ is the converged parameter for the total loss, it satisfies the stationarity condition: K X ∇Lk (θ) = 0. (51) k=1

We analyze the squared norm of this sum, which must equal zero: K X

∇Lk (θ)

k=1

= 2

K X

∥∇Lk (θ)∥22 +

k=1

∇Li (θ)⊤ ∇Lj (θ) = 0.

(52)

i̸=j

By rearranging terms, we obtain an exact identity relating the sum of squared gradient norms to the negative sum of cross-task inner products: K X k=1

∥∇Lk (θ)∥22 =

−∇Li (θ)⊤ ∇Lj (θ) .

(53)

i̸=j

Substituting Eq. (53) into Eq. (50), we obtain the first inequality of the theorem: K 1 X 1 X ∥θ − θk∗ ∥22 ≤ −∇Li (θ)⊤ ∇Lj (θ) . 2 K Kλ k=1

i̸=j

(54)

Step 3: Bounding via Cosine Similarity. Finally, we bound the inner product term using the gradient magnitude upper bound G = supk ∥∇Lk (θ)∥2 . Recall that: ∇Li (θ)⊤ ∇Lj (θ) = ∥∇Li (θ)∥2 ∥∇Lj (θ)∥2 CosSim(∇Li (θ), ∇Lj (θ)).

(55)

We use the property that for any i, j, the following term is non-negative: (G2 − ∥∇Li (θ)∥2 ∥∇Lj (θ)∥2 )(1 − CosSim(∇Li (θ), ∇Lj (θ))) ≥ 0,

(56)

since ∥∇Lk (θ)∥2 ≤ G and CosSim ≤ 1. Adding this non-negative term to the negative inner product allows us to derive the bound directly: −∇Li (θ)⊤ ∇Lj (θ) = −∥∇Li (θ)∥2 ∥∇Lj (θ)∥2 CosSim(∇Li (θ), ∇Lj (θ)) ≤ −∥∇Li (θ)∥2 ∥∇Lj (θ)∥2 CosSim(∇Li (θ), ∇Lj (θ)) + (G2 − ∥∇Li (θ)∥2 ∥∇Lj (θ)∥2 )(1 − CosSim(∇Li (θ), ∇Lj (θ)))

(57)

= G (1 − CosSim(∇Li (θ), ∇Lj (θ))) − ∥∇Li (θ)∥2 ∥∇Lj (θ)∥2 2

≤ G2 (1 − CosSim(∇Li (θ), ∇Lj (θ))). Summing this inequality over all i ̸= j yields: X X −∇Li (θ)⊤ ∇Lj (θ) ≤ G2 (1 − CosSim(∇Li (θ), ∇Lj (θ))) . i̸=j

i̸=j

Combining this with the result from Step 2 completes the proof.

(58)

Implicit Bias of Nexus Optimizer

In this appendix, we provide the detailed proofs for Theorem 3.2. We rigorously analyze the update dynamics of Algorithm 1 using second-order Taylor expansions and derive the precise form of the implicit optimization objective with explicit non-asymptotic error bounds.

E.1

Preliminaries and Notation

Let Li : Rd → R denote the loss function for the i-th task, where i ∈ {1, . . . , k}. We denote the gradient and Hessian at parameters θ as ∇Li (θ) and ∇2 Li (θ), respectively. The cosine similarity between the gradients of task i and task j is defined as: Sij (θ) ≜ CosSim(∇Li (θ), ∇Lj (θ)) =

∇Li (θ)⊤ ∇Lj (θ) . ∥∇Li (θ)∥2 ∥∇Lj (θ)∥2

(59)

Algorithm 1 performs k inner updates in each outer iteration t. Let θt,0 be the parameters at the start of the inner loop (i.e., θt,0 = θt−1 ). At each inner step m ∈ {1, . . . , k}, a task index sm is sampled uniformly from {1, . . . , k}. The update rule is: ∇Lsm (θt,m−1 ) θt,m = θt,m−1 − γ . (60) ∥∇Lsm (θt,m−1 )∥2 Pk The Nexus pseudo-gradient passed to the outer optimizer is ĝt = θt,0 − θt,k = m=1 (θt,m−1 − θt,m ).

E.2

Assumptions and Derived Constants

To derive explicit non-asymptotic bounds, we utilize the following standard assumptions regarding the loss landscape. • Assumption 1 (Bounded Gradients): For all tasks i and parameters θ, the gradient norm is bounded from below: 0 < Gmin ≤ ∥∇Li (θ)∥2 . • Assumption 2 (Smoothness): The loss Li is L-smooth, i.e., ∥∇2 Li (θ)∥2 ≤ L. • Assumption 3 (Hessian Lipschitz): The Hessian is ρ-Lipschitz continuous, i.e., ∥∇2 Li (x) − ∇2 Li (y)∥2 ≤ ρ∥x − y∥2 . Based on the properties above, we further denote L1 and L2 as the Lipschitz constants for the normalized gradient and its Jacobian, respectively: 1. The normalized gradient is L1 -Lipschitz continuous: ∇Li (x) ∇Li (y) − ∥∇Li (x)∥2 ∥∇Li (y)∥2

≤ L1 ∥x − y∥2 .

(61)

2. The Jacobian of the normalized gradient is L2 -Lipschitz continuous: ∥Ji (x) − Ji (y)∥2 ≤ L2 ∥x − y∥2 , ∂ where Ji (θ) = ∂θ

Derivation of Constants.

∇Li (θ) ∥∇Li (θ)∥2

(62)

Here, we provide the detailed derivation of L1 and L2 based on Assumptions 1-3.

1. Derivation of L1 : By the Mean Value Theorem, L1 is bounded by the supremum of the spectral norm of the Jacobian Ji (θ). The Jacobian is explicitly given by: 1 ∇Li ∇L⊤ i Ji (θ) = I− ∇2 Li (θ). (63) ∥∇Li ∥2 ∥∇Li ∥22

The middle term is an orthogonal projection matrix with spectral norm 1. Using the bounds from Assumptions 1 and 2: 1 L L1 ≤ sup ∥Ji (θ)∥2 ≤ ·1·L= . (64) Gmin Gmin θ 2. Derivation of L2 : We decompose the Jacobian Ji (θ) into three components: a scalar term u(θ), a projection

term Π(θ), and the Hessian Hi (θ):

∇Li ∇L⊤ i Ji (θ) = ∥∇Li (θ)∥−1 · I − · ∇2 Li (θ) . | {z } ∥∇Li ∥22 | {z 2 } | {z } Hi (θ) u(θ)

(65)

Π(θ)

We apply the product Lipschitz rule. For a product of three functions f = abc, the Lipschitz constant satisfies Lf ≤ La Mb Mc + Ma Lb Mc + Ma Mb Lc , where M(·) denotes the upper bound of the magnitude and L(·) denotes the Lipschitz constant. • Part 1: Scalar u(θ) = ∥∇Li ∥−1 2 . Magnitude (Mu ): By Assumption 1, |u| ≤ G 1 . min

∇ Li ∇Li Hi ∇Li Lipschitz (Lu ): The gradient of u is ∇u = −∥∇Li ∥−2 = − ∥∇L . Taking the norm, we have 2 ∥∇L ∥ ∥3 i 2

i 2

∥Hi ∥2 L i ∥2 ∥∇Li ∥2 ∥∇u∥2 ≤ ∥H∥∇L = ∥∇L . 3 2 . Using the bounds L and Gmin , we get Lu = G2 i∥ i∥ 2

•

min

Part 2: Projection Π(θ) = I − hi h⊤ i .

Magnitude (MΠ ): The spectral norm is ∥Π∥2 = 1. Lipschitz (LΠ ): Π depends on the normalized gradient hi , which is L1 -Lipschitz. For any unit vectors

x, y, we have ∥xx⊤ − yy ⊤ ∥2 ≤ ∥x(x − y)⊤ ∥2 + ∥(x − y)y ⊤ ∥2 = 2∥x − y∥2 . By the chain rule, LΠ = 2L1 = G2L . min

• Part 3: Hessian Hi (θ) = ∇2 Li . Magnitude (MH ): By Assumption 2, ∥Hi ∥2 ≤ L. Lipschitz (LH ): By Assumption 3, LH = ρ. Substituting these values into the product rule formula: L2 ≤ Lu MΠ MH + Mu LΠ MH + Mu MΠ LH L 1 2L 1 ≤ · 1 · L + · · L + · 1 · ρ G2min Gmin Gmin Gmin 2

(66)

2L ρ L + 2 + . G2min Gmin Gmin

Combining terms yields the final constant: L2 =

E.3

3L2 + ρGmin . G2min

(67)

Derivation of the Update Direction

We now derive the expansion of the total pseudo-gradient ĝt and bound the error terms. E.3.1

Step 1: Expansion of the Normalized Gradient

We aim to expand the normalized gradient at the shifted parameters θt,m−1 around the initial point θt,0 . Let ∆θm−1 = θt,m−1 − θt,0 . The Jacobian of the normalized gradient is given explicitly by the projection of the Hessian: 1 ∇Lsm (θ)∇Lsm (θ)⊤ Jsm (θ) = I− ∇2 Lsm (θ). ∥∇Lsm (θ)∥2 ∥∇Lsm (θ)∥22 30

(68)

Applying Taylor’s theorem with the Lagrange remainder form: ∇Lsm (θt,m−1 ) ∇Lsm (θt,0 ) = + Jsm (θt,0 )∆θm−1 + rm . ∥∇Lsm (θt,m−1 )∥2 ∥∇Lsm (θt,0 )∥2

(69)

Using the L2 -Lipschitz property of the Jacobian, the residual vector rm is bounded by: ∥rm ∥2 ≤ E.3.2

L2 ∥∆θm−1 ∥22 . 2

(70)

Step 2: Recursive Substitution

The displacement ∆θm−1 is the sum of previous updates. Using the zeroth-order approximation: ∆θm−1 =

m−1 X

l=1

(θt,l − θt,l−1 ) = −γ

∇Lsl (θt,l−1 ) . ∥∇Lsl (θt,l−1 )∥2

(71)

We approximate the terms in the sum using the zeroth-order expansion around θt,0 . Using the L1 -Lipschitz property of the normalized gradient: ∇Lsl (θt,l−1 ) ∇Lsl (θt,0 ) − ∥∇Lsl (θt,l−1 )∥2 ∥∇Lsl (θt,0 )∥2

≤ L1 ∥θt,l−1 − θt,0 ∥2 = L1 ∥ − 2

l−1 X j=1

∇Lsl (θt,j ) ∥ ≤ L1 (l − 1)γ. (72) ∥∇Lsl (θt,j )∥2

Thus, we can write: ∆θm−1 = −γ

m−1 X l=1

∇Lsl (θt,0 ) + δm−1 , ∥∇Lsl (θt,0 )∥2

(73)

where the accumulated error δm−1 is bounded by summing the individual errors: ∥δm−1 ∥2 ≤ γ

m−1 X

L1 (l − 1)γ = L1 γ 2

l=1

L1 (m − 1)(m − 2) ≤ (m − 1)2 γ 2 . 2 2

(74)

Substituting this expression for ∆θm−1 back into Eq. (69): m−1

X ∇Lsm (θt,m−1 ) ∇Lsm (θt,0 ) ∇Lsl (θt,0 ) + E m. = −γ Jsm (θt,0 ) ∥∇Lsm (θt,m−1 )∥2 ∥∇Lsm (θt,0 )∥2 ∥∇Lsl (θt,0 )∥2

(75)

l=1

Here, the total error at step m, denoted E m , consists of the Taylor residual rm and the propagation error from δm−1 scaled by the Jacobian. Using ∥Jsm ∥2 ≤ L1 and ∥∆θm−1 ∥2 ≤ (m − 1)γ: L2 L1 L2 + L21 2 2 2 2 ∥E m ∥2 ≤ (m − 1) γ + L1 (m − 1) γ = (m − 1)2 γ 2 . (76) 2 2 2 E.3.3

Step 3: Aggregation of the Pseudo-Gradient Pk ∇L (θt,m−1 ) The total pseudo-gradient is ĝt = γ m=1 ∥∇Lssm(θt,m−1 )∥2 . Substituting the result from Step 2: m

ĝt = γ

k X

k m−1 X X ∇Lsl ∇Lsm − γ2 Jsm + E total . ∥∇L ∥ ∥∇L s 2 sl ∥2 m m=1 m=1

(77)

l=1

(We omit the argument θt,0 for brevity; all terms are evaluated at θt,0 ). The total error vector E total = Pk γ m=1 E m is explicitly bounded by summing the bounds from Step 2: ∥E total ∥2 ≤ γ

k X L2 + L2 1

m=1

(m − 1)2 γ 2 ≤

L2 + L21 3 3 k γ . 6

(78)

E.3.4

Step 4: Expectation Analysis and Connection to Cosine Similarity

We now compute the expectation of ĝt over the independent uniform sampling of indices s1 , . . . , sk and relate the second-order term to the gradient of the cosine similarity. Linear Term.

Let Tlinear =

∇Lsm m=1 ∥∇Lsm ∥2 . By linearity of expectation:

E[Tlinear ] = k ·

k k X 1 X ∇Li ∇Li = . k i=1 ∥∇Li ∥2 ∥∇L i ∥2 i=1

(79)

Pk Pm−1 ∇L Interaction Term. Let Tinteract = m=1 l=1 Jsm ∥∇Lssl∥2 . The double summation contains k(k−1) terms. 2 l Since m > l, sm and sl are independent. Thus: k k ∇Lsl 1 XX ∇Lj E Jsm = 2 . Ji ∥∇Lsl ∥2 k i=1 j=1 ∥∇Lj ∥2

(80)

k k ∇Lj k(k − 1) 1 X X k−1X ∇Lj Ji · 2 Ji = . 2 k i=1 j=1 ∥∇Lj ∥2 2k i,j ∥∇Lj ∥2

(81)

Summing over all pairs yields: E[Tinteract ] =

We define ∇Sij as the gradient of the cosine similarity between task i and j. Explicitly: ∇Sij = Ji Observing that the summation interaction term expectation as: E[Tinteract ] =

E.4

∇Lj i,j Ji ∥∇Lj ∥2

∇Li ∇Lj + Jj . ∥∇Lj ∥2 ∥∇Li ∥2

(82)

is symmetric with respect to i and j, we can rewrite the

k−1X ∇Lj k−1X ∇Li Ji ∇Sij (θt,0 ). + Jj = 4k i,j ∥∇Lj ∥2 ∥∇Li ∥2 4k i,j

(83)

Proof Conclusion

Combining the linear term and the interaction term, the expected Nexus update direction is: E[ĝt ] = γ

k X

∇Li (θt,0 ) k−1X − γ2 ∇Sij (θt,0 ) + E total . ∥∇Li (θt,0 )∥2 4k i,j i=1

Substituting the constants derived in Assumption E.2, the residual is bounded by: 1 ρGmin + 4L2 ∥E total ∥2 ≤ k3 γ 3 . 6 G2min

(84)

(85)

This confirms that the update direction follows the gradient of the loss plus the similarity alignment term, subject to a bounded cubic error. □

Proof of Convergence Rate (Theorem B.1)

In this section, we provide the detailed proof for Theorem B.1. Let θ ∗ be the common minimizer such that ∇Li (θ ∗ ) = 0 for all i ∈ [k]. Consider the update at step m: θm = θm−1 − γ∇Lsm (θm−1 ), where sm is the task index sampled uniformly at random. First, we expand the squared distance to the optimum for a specific realization of sm : ∥θm − θ ∗ ∥2 = ∥θm−1 − γ∇Lsm (θm−1 ) − θ ∗ ∥2 = ∥θm−1 − θ ∗ ∥2 − 2γ⟨∇Lsm (θm−1 ), θm−1 − θ ∗ ⟩ + γ 2 ∥∇Lsm (θm−1 )∥2 .

(86)

To bound the inner product term, we utilize the property of smooth and strongly convex functions. Define the auxiliary function ϕi (θ) = Li (θ) − µ2 ∥θ∥2 . Since each Li is L-smooth and µ-strongly convex, ϕi (θ) is convex and (L − µ)-smooth. By the co-coercivity property of convex smooth functions, for any θ, we have: ⟨∇ϕi (θ) − ∇ϕi (θ ∗ ), θ − θ ∗ ⟩ ≥

1 ∥∇ϕi (θ) − ∇ϕi (θ ∗ )∥2 . L−µ

(87)

Substituting ∇ϕi (θ) = ∇Li (θ) − µθ and noting that ∇Li (θ ∗ ) = 0, we substitute back: 1 ∥∇Li (θ) − µ(θ − θ ∗ )∥2 L−µ 1 = ∥∇Li (θ)∥2 − 2µ⟨∇Li (θ), θ − θ ∗ ⟩ + µ2 ∥θ − θ ∗ ∥2 . L−µ

⟨∇Li (θ) − µ(θ − θ ∗ ), θ − θ ∗ ⟩ ≥

(88)

Rearranging the terms, we obtain the following inequality which holds for any task index i, and thus specifically for the sampled index sm : ⟨∇Lsm (θ), θ − θ ∗ ⟩ ≥

1 µL ∥∇Lsm (θ)∥2 + ∥θ − θ ∗ ∥2 . L+µ L+µ

(89)

Substituting Eq. (89) back into Eq. (86) with θ = θm−1 : 1 µL ∗ 2 ∗ 2 2 ∗ 2 ∥θm − θ ∥ ≤ ∥θm−1 − θ ∥ − 2γ ∥∇Lsm (θm−1 )∥ + ∥θm−1 − θ ∥ + γ 2 ∥∇Lsm (θm−1 )∥2 L+µ L+µ 2γµL 2γ ∗ 2 2 = 1− ∥θm−1 − θ ∥ + γ − ∥∇Lsm (θm−1 )∥2 . L+µ L+µ (90) 2γ 2 Provided that the step size satisfies γ ∈ (0, L+µ ], the coefficient γ 2 − L+µ is non-positive. Since ∥∇Lsm (θm−1 )∥2 ≥ 0, we can drop the gradient norm term to obtain an upper bound: 2γµL ∥θm − θ ∗ ∥2 ≤ 1 − ∥θm−1 − θ ∗ ∥2 . L+µ

(91)

Since this inequality holds for any realization of the random sample sm , we take the expectation over the sampling distribution. Let E[·] denote the total expectation over the sequence of random indices {s1 , . . . , sm }. We have: 2γµL E[∥θm − θ ∗ ∥2 ] ≤ 1 − E[∥θm−1 − θ ∗ ∥2 ]. (92) L+µ Applying this recurrence relation recursively for T steps yields: T 2γµL E[∥θT − θ ∥ ] ≤ 1 − ∥θ0 − θ ∗ ∥2 . L+µ ∗ 2

(93)

2 Specifically, when choosing the step size γ = L+µ :

(L − µ)2 2γµL 4µL = = 1− =1− L+µ (L + µ)2 (L + µ)2

κ−1 κ+1

2 ,

(94)

where κ = L/µ is the condition number. Thus, we obtain the convergence rate: E[∥θT − θ ∗ ∥2 ] ≤

κ−1 κ+1

∥θ0 − θ ∗ ∥2 .

(95) □

Third-Order Implicit Bias Analysis

In this section, we analyze the third-order implicit bias of Nexus, inspired by recent works [6, 7, 46]. While the second-order analysis reveals how Nexus aligns gradients, it does not fully explain the stability of this alignment in complex landscapes. Here, we demonstrate that the Nexus update direction implicitly minimizes a "Generalized Directional Sharpness" metric. This implies that Nexus actively seeks regions where the loss landscape is not only aligned but also locally flat along the alignment direction, thereby preventing the "de-alignment" caused by sharp curvature.

G.1

Setup and Definitions

To perform this analysis, we verify the behavior of the third-order terms in the Taylor expansion. We introduce a standard assumption regarding the smoothness of the Hessian. Assumption 4 (Bounded Third Derivative). Assume the third-order derivative tensor is bounded, i.e., for any

unit vectors u, v, w and any task i, there exists a constant M3 such that ∥∇3 Li (θ)[u, v, w]∥2 ≤ M3 . This implies that the third-order Taylor remainder satisfies ∥rT aylor (δ)∥2 ≤ M63 ∥δ∥32 . Definition (Generalized Directional Sharpness). We define the generalized sharpness term involving the Hessian

of task j and the gradient directions of tasks i and p as: Ri,j,p (θ) ≜

1 ∇Li (θ)⊤ ∇2 Lj (θ)∇Lp (θ). 2

(96)

This term measures the curvature of task j along the plane spanned by the gradients of tasks i and p. When i = p, this reduces to the standard directional sharpness, quantifying how fast the gradient changes along the update direction.

G.2

Proof of Theorem B.3

1. Exact Expansion of the Gradient. Consider the m-th inner step with sampled task sm . Let θm−1 = θ0 +∆m−1 . The exact third-order Taylor expansion is:

1 (m) ∇Lsm (θm−1 ) = ∇Lsm (θ0 ) + ∇2 Lsm (θ0 )∆m−1 + ∇3 Lsm (θ0 )[∆m−1 , ∆m−1 ] + rT aylor . 2

(97)

(m)

The remainder is bounded by ∥rT aylor ∥2 ≤ M63 ∥∆m−1 ∥32 . Using the bound on displacement magnitude ∥∆m−1 ∥2 ≤ (m − 1)γ: M3 (m) (m − 1)3 γ 3 . (98) ∥rT aylor ∥2 ≤ 6 ˜ m−1 = 2. Displacement Decomposition. We define the ideal displacement using initial gradients as ∆ Pm−1

˜ m−1 + δm−1 . Using the Lipschitz constant L1 = L/Gmin −γ dˆsl . The true displacement is ∆m−1 = ∆ for the normalized gradient, the accumulated error is bounded by: l=1

∥δm−1 ∥2 ≤ γ

m−1 X

L1 (l − 1)γ ≤

l=1

L1 (m − 1)2 γ 2 . 2

(99)

3. Substitution into Quadratic Term. We substitute ∆m−1 into the third-order term. By multilinearity of the

tensor:

1 3 1 ˜ m−1 , ∆ ˜ m−1 ] + r (m) . ∇ Lsm [∆m−1 , ∆m−1 ] = ∇3 Lsm [∆ sub 2 2

(100)

(m)

The residual rsub accounts for the cross-terms and quadratic error terms. Its norm is strictly bounded by: (m)

∥rsub ∥2 ≤

1 ˜ m−1 ∥∥δm−1 ∥ + ∥∇3 ∥∥δm−1 ∥2 . 2∥∇3 ∥∥∆ 2

(101)

˜ m−1 ∥ and ∥δm−1 ∥: Substituting the bounds for ∥∆ 2 L1 M3 L1 (m) 2 2 2 2 ∥rsub ∥2 ≤ M3 ((m − 1)γ) (m − 1) γ + (m − 1) γ 2 2 2 M3 L 1 M3 L21 = (m − 1)3 γ 3 + (m − 1)4 γ 4 . 2 8

(102)

4. Derivation of the Expected Update Direction. The explicit third-order component of the update (excluding

residuals) is:

v3 =

k X

1 3 ˜ ˜ ∇ Lsm [∆m−1 , ∆m−1 ] . 2

m−1 m−1

−γ

m=1

˜ m−1 = Substituting ∆

Pm−1 l=1

(103)

−γ dˆsl : v3 = −

γ3 X X X 3 ∇ Lsm [dˆsl , dˆsp ]. 2 m=1 p=1

(104)

l=1

Taking the expectation over uniform sampling of indices sm , sl , sp , each triplet (i, j, p) appears with probability 1/k 3 : k m−1 m−1 h i γ3 X X X Esm ,sl ,sp ∇3 Lsm [dˆsl , dˆsp ] E[v3 ] = − 2 m=1 l=1 p=1  ! k m−1 m−1 3 X X X X γ 1 =− 1  3 ∇3 Lj [dˆi , dˆp ] (105) 2 m=1 k j,i,p l=1 p=1  ! k 3 X X 1 γ (m − 1)2  3 =− ∇3 Lj [dˆi , dˆp ] . 2 m=1 k i,j,p Using the summation formula

E[v3 ] = −γ 3

(k − 1)(2k − 1) X 3 ∇ Lj [dˆi , dˆp ]. 12k 2 i,j,p

m=1 (m − 1)

= k(k−1)(2k−1) : 6 (106)

Recognizing that ∇θ Ri,j,p = 12 ∇3 Lj [dˆi , dˆp ] (treating the direction vectors as locally constant for the gradient of the surrogate), we can rewrite the update as a gradient descent step on the sharpness metric:   X (k − 1)(2k − 1) E[v3 ] = −γ∇ γ 2 Ri,j,p  . (107) 6k 2 i,j,p This confirms that Nexus implicitly minimizes the generalized directional sharpness. Pk (m) (m) 5. Bounding the Total Residual. The total error vector is E 3rd = m=1 −γ(rT aylor + rsub ). Taking the norm: ∥E 3rd ∥2 ≤ γ

k X M3 m=1

Using summation bounds

M3 L 1 + 6 2

M3 L21 4 4 (m − 1) γ + (m − 1) γ . 8 3 3

(108)

Pk−1 3 Pk−1 4 k4 k5 j=1 j ≤ 4 and j=1 j ≤ 5 , and substituting L1 = L/Gmin : ∥E 3rd ∥2 ≤

M3 M3 L + 24 8Gmin

k4 γ 4 +

M3 L2 5 5 k γ . 40G2min

(109) □

More Experiments Details

H.1

Detailed Hyper-parameters

Our hyperparameter configurations strictly follow the baseline established in Wen et al. [45]. To ensure optimality for our specific pretraining corpus, we conducted a grid search over the learning rate with a multiplier of 2 (i.e., verifying 0.5× and 2.0×). The empirical results confirmed that the original learning rate settings remain optimal for our setup. For clarity and reproducibility, we summarize the key hyperparameters in Tab. 6. For the learning rate schedule, all experiments utilizing the Warmup-Stable-Decay (WSD) scheduler employ 1,000 warmup steps and 10,000 decay steps. Across all experiments, we maintain a global batch size of 256, an Adam β of (0.9, 0.95), an Adam ϵ of 10−10 , and a gradient clipping norm of 1.0. Table 6 Summary of key hyperparameters for the main pretraining experiments. Model Size Optimizer Outer LR Inner LR (γ ) Chinchilla Tokens (B) Weight Decay

H.2

Reference

Adam Nexus

0.002 0.002

0.01

4× 4×

50 50

0.2 0.2

Sec. 4.3

Adam Muon Nexus

0.001 0.001 0.001

0.01

2× 2× 2×

110 110 110

0.2 0.1 0.2

Secs. 4.2 and 4.5 Secs. 4.4 and H.3

Detailed Results for Data Mixture

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% → 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. Data

Math10

Math40

Math70

Optim.

Metric

Loss Metrics (↓)

Gen.

Reasoning

Math

Code

Avg.

Pretrain. OOD MMLU GPQA GPQA-D BBH GSM8k MATH HumanEval MBPP

Acc. (↑) AdamW Loss (↓)

1.606

1.302

44.0 1.259

32.0 1.054

43.0 1.116

38.0 1.922

37.1 1.633

Acc. (↑) Loss (↓)

1.602

1.290

59.0 1.227

40.0 1.026

47.0 1.086

38.0 1.921

40.3 1.601

Improv. Loss (↑) +0.004 +0.012 +0.086 +0.024 +0.029 +0.030 +0.032 +0.028

+0.030

+0.001 +0.032

32.8

2.005

22.6 1.910

All

36.6 1.534

Nexus

47.8 2.265 48.9

29.6

23.4

36.6

2.179

1.981

1.881

1.504

AdamW

Acc. (↑) Loss (↓)

1.336

1.330

47.8 2.210

29.6 1.989

22.6 1.891

38.1 1.522

64.0 1.171

44.0 0.969

41.0 1.144

38.0 1.976

40.6 1.609

Nexus

Acc. (↑) Loss (↓)

1.339

1.331

51.2 2.182

33.5

27.3

1.990

1.889

41.1 1.511

70.0 1.117

43.0 0.929

45.0 1.132

44.5 1.876

44.4 1.578

Improv. Loss (↑)

-0.003

-0.001 +0.028 -0.001

+0.002 +0.011 +0.054 +0.040

+0.012

+0.100 +0.031

AdamW

Acc. (↑) Loss (↓)

1.033

1.399

44.0 2.252

27.3 2.025

23.4 1.923

41.1 1.541

77.0 1.111

45.0 0.923

38.0 1.178

38.0 1.897

Nexus

Acc. (↑) Loss (↓)

1.040

1.409

49.8 2.221

30.4

2.037

23.4 1.936

42.5 1.548

76.0 1.082

52.0 0.911

1.176

1.872

Improv. Loss (↑)

-0.007

-0.010 +0.031 -0.012

-0.013

-0.007 +0.029 +0.012

+0.002

+0.025 +0.008

41.0

38.0

41.7 1.606 44.1

1.598

As shown in Tab. 7, we observe a dynamic trade-off mechanism: • In the sample-sparse regime (Math10): Where math data is scarce, the baseline optimizer struggles to generalize on reasoning tasks. Nexus provides the most significant gains here (e.g., +15.0 on GSM8k), 37

effectively "mining" the rare training signals to build robust reasoning capabilities. • In the sample-dense regime (Math70): As math data becomes abundant, the baseline catches up on math benchmarks. However, Nexus automatically shifts its advantage to the now-relative-minority domains. It significantly boosts General Knowledge (MMLU: +5.8) and broad Reasoning (GPQA: +3.1) compared to the baseline, which begins to suffer from domain dominance. • Lower sensitivity to mixture shifts: Nexus also demonstrates higher stability against drastic changes in data mixture. When shifting from a math-heavy (Math70) to a math-sparse (Math10) mixture, the performance variance of Nexus is significantly smaller than that of the baseline. For instance, while the baseline’s GSM8k score drops precipitously by 33.0 points (from 77.0 to 44.0), Nexus mitigates this degradation, dropping only 17.0 points (from 76.0 to 59.0). Similarly, on MMLU, while the baseline fluctuates by 3.8 points, Nexus remains highly stable with a variation of less than 1.0 point (49.8 vs. 48.9), demonstrating its stability against data mixture changes. This suggests that Nexus reduces sensitivity to manual data mixing ratios, acting as an automatic balancer that prioritizes representations for the most under-optimized tasks in the mixture.

H.3

Experiments on a Public Dataset

Motivation. While our primary analyses utilize strictly cleaned data to avoid confounding factors, many

popular open-source pretraining datasets inevitably suffer from data contamination, inadvertently including benchmark training sets (e.g., GSM8k). We evaluate Nexus on a public dataset from Basant et al. [2] to investigate whether its consensus-seeking mechanism remains robust and mitigates shortcut over-memorization in the presence of such noisy, contaminated signals. Settings. We train the 1B and 3B models on a public dataset [2]. All other training configurations, including

model architectures and base optimizer hyperparameters, are kept strictly identical to the main experiments detailed in Sec. 4.2. Table 8 Results on a public pretraining dataset [2]. 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. Model Optim. Metric

Gen.

Reasoning

Math

Code

Avg.

Pretrain. OOD MMLU GPQA GPQA-D BBH GSM8k MATH HumanEval MBPP

Adam

Acc. (↑) Loss (↓)

1.331

1.863

Nexus

Acc. (↑) Loss (↓)

1.338

1.835

Improv. Loss (↑)

Loss Metrics (↓)

34.2

2.552

25.8 2.280

18.8 2.191

24.8 1.700

18.0 1.708

14.0 1.346

38.0 1.324

1.0 21.8 2.991 2.011

41.0

31.4

25.0

18.8

23.0

2.446

2.261

2.172

1.689

1.325

15.0 2.908

-0.007 +0.028 +0.106 +0.019

+0.019

+0.011 -0.041 +0.021

-0.001

+0.083 +0.027

27.3

47.4 1.529

44.0

59.0 1.205

23.0 39.0 2.732 1.812

63.0

Adam

Acc. (↑) Loss (↓)

1.330

1.623

2.356

Nexus

Acc. (↑) Loss (↓)

1.338

1.606

Improv. Loss (↑)

-0.008

55.1

25.0 2.062

56.2 2.380

25.8

23.4

2.047

+0.017 -0.024 +0.015

22.0

1.749

12.0

All

23.5 1.984

1.519

31.0 1.121

1.957

44.4 47.0 1.540 1.533

1.106

1.199

38.0 2.530

+0.018

-0.011 -0.014 +0.015

+0.006

+0.202 +0.026

1.975

32.0

41.2 1.786

Results. As shown in Tab. 8, the Adam baseline exhibits artificial performance inflation on potentially

contaminated benchmarks like GSM8k. In contrast, Nexus resists overfitting to these leaked signals and effectively reallocates the model’s capacity to uncontaminated, sparse domains. This dynamic balancing is evidenced by the striking improvements on coding tasks—such as MBPP accuracy increasing from 1.0% to 15.0% (1B) and 23.0% to 38.0% (3B)—ultimately leading to a consistently lower OOD loss across both scales.

H.4

Detailed Results for Model Size Scaling

This section provides the detailed experimental results corresponding to the model size scaling analysis discussed in Sec. 4.3. 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. Size

Loss Metrics (↓)

Optim.

Metric

AdamW

Acc. (↑) Loss (↓)

2.038

1.559

Acc. (↑) Loss (↓)

2.031

1.549

Gen.

Reasoning

Math

Code

Avg.

Pretrain. OOD MMLU GPQA GPQA-D BBH GSM8k MATH HumanEval MBPP 7.0 1.612

6.0 1.360

0.0 1.407

5.0

11.0 1.330

0.0

10.0

16.2

1.384

2.183

1.942

Improv.

Acc. (↑) -1.0 +3.1 -2.4 +3.7 -2.0 +5.0 Loss (↑) +0.007 +0.010 +0.032 +0.024 +0.026 +0.001 +0.011 +0.030

0.0

+0.8

+0.023

+0.047 +0.024

AdamW

Acc. (↑) Loss (↓)

1.909

1.474

Acc. (↑) Loss (↓)

1.901

1.469

130M Nexus

28.0 2.555

22.6 2.438

2.338

27.0

25.7 2.414

2.312

2.523

24.2

21.8

25.1 1.793 28.8

1.792

1.601

0.0

10.0 2.230

All 15.4 1.967

26.5 2.296

21.8 2.196

31.8 1.704

15.0 1.471

15.0 1.255

6.0 1.322

16.0 2.141

20.7 1.860

30.3

27.3 2.278

25.0 2.177

30.7 1.707

14.0 1.470

16.0

2.381

1.237

13.0 1.298

21.0 2.046

1.824

Improv.

Acc. (↑) -3.0 +0.8 Loss (↑) +0.008 +0.005 +0.114 +0.018

+3.2 +0.019

-1.1 -1.0 +1.0 -0.003 +0.001 +0.018

+7.0 +0.024

+5.0 +1.5 +0.095 +0.036

AdamW

Acc. (↑) Loss (↓)

1.826

1.433

13.0 1.204

19.0 1.270

17.0 2.035

21.9 1.786

520M Nexus

Acc. (↑) Loss (↓)

1.826

1.428

1.758

Improv.

Acc. (↑) Loss (↑)

AdamW

300M Nexus

33.3

2.495

32.1 2.363

25.0 2.221

21.8 2.124

29.6 1.640

18.0 1.429

33.5

30.4 2.201

21.8

29.3 1.638

20.0

13.0

2.102

1.396

1.176

19.0 1.261

22.0 1.977

0.000

+1.4 +5.4 0.0 -0.3 +2.0 0.0 +0.005 +0.047 +0.020 +0.022 +0.002 +0.033 +0.028

0.0 +0.009

+5.0 +1.7 +0.058 +0.027

Acc. (↑) Loss (↓)

1.714

1.364

Acc. (↑) Loss (↓)

1.707

1.358

2.316

22.6 2.410

24.2 2.000

25.7 2.373

1.984

Improv.

Acc. (↑) -2.7 +3.1 Loss (↑) +0.007 +0.006 +0.160 +0.037

+0.016

AdamW

Acc. (↑) Loss (↓)

1.606

1.302

Acc. (↑) Loss (↓)

1.602

1.290

1.2B Nexus

2.3B Nexus Improv.

22.2

23.6

27.4 1.799

30.0 1.338

23.0 1.112

30.0 1.199

31.0 2.023

29.1 1.813

32.9

1.792

37.0 1.325

28.0 1.109

35.0 1.179

30.0 1.987

31.7 1.777

+5.5

+7.0

+5.0

+0.007 +0.013 +0.003

+5.0 +0.020

-1.0 +2.6 +0.036 +0.036

22.6 1.910

36.6 1.534

44.0 1.259

32.0 1.054

43.0 1.116

38.0 1.922

37.1 1.633

23.4 1.881

36.6 1.504

59.0 1.227

40.0 1.026

47.0 1.086

38.0 1.921

40.3 1.601

Acc. (↑) +1.1 -3.2 +0.8 0.0 +15.0 +8.0 Loss (↑) +0.004 +0.012 +0.086 +0.024 +0.029 +0.030 +0.032 +0.028

+4.0 +0.030

0.0 +3.2 +0.001 +0.032

44.4

2.626 41.7 2.466

47.8 2.265

32.8

2.005

48.9

29.6

2.179

1.981

23.4 -0.8

As demonstrated above and analyzed in Sec. 4.3, Nexus consistently outperforms the base optimizer across all evaluated model scales, with average benchmark accuracy improvements of +0.8% (130M), +1.5% (300M), +1.7% (520M), +2.6% (1.2B), and +3.2% (2.3B).

H.5

Experiments on Muon Optimizers

This section provides the detailed experimental results discussed in Secs. 4.2 and 5.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. Loss Metrics (↓)

Optim.

Metric

AdamW

Acc. (↑) Loss (↓)

1.606

1.302

Adam+Nexus

Acc. (↑) Loss (↓)

1.602

1.290

Gen.

Reasoning

Math

Code

Avg.

Pretrain. OOD MMLU GPQA GPQA-D BBH GSM8k MATH HumanEval MBPP 36.6 1.534

44.0 1.259

32.0 1.054

43.0 1.116

38.0 1.922

37.1 1.633

59.0 1.227

40.0 1.026

47.0 1.086

38.0 1.921

40.3 1.601

(- AdamW)

Acc. (↑) +1.1 -3.2 +0.8 0.0 +15.0 +8.0 Loss (↑) +0.004 +0.012 +0.086 +0.024 +0.029 +0.030 +0.032 +0.028

+4.0 +0.030

0.0 +3.2 +0.001 +0.032

Muon

Acc. (↑) Loss (↓)

38.0 1.035

40.0 1.091

43.0 1.951

(- AdamW)

Acc. (↑) +2.0 -0.8 +1.6 +5.3 +2.0 +6.0 Loss (↑) +0.029 +0.017 +0.077 +0.037 +0.036 +0.032 +0.023 +0.019

-3.0 +0.025

+5.0 +2.3 -0.029 +0.027

1.577

1.285

47.8 2.265

32.8 2.005

22.6 1.910

All

48.9

29.6

23.4

36.6

2.179

1.981

1.881

1.504

49.8 2.188

32.0 1.968

24.2 1.874

41.9 1.502

46.0 1.236

39.4 1.606

As demonstrated above, Nexus achieves comparable downstream performance to Muon, despite maintaining a pretraining loss that is nearly identical to the AdamW baseline. These results explicitly indicate that while Muon improves downstream performance primarily by reaching a significantly lower pretraining loss, the gains from Nexus stem directly from its favorable implicit bias.

H.6

Downstream SFT

Motivation and Settings. To verify whether the performance gains of Nexus are merely a result of "pre-

consuming" the potential improvements of the SFT phase in advance, we evaluate the supervised fine-tuning (SFT) performance of our checkpoints. We use an SFT dataset similar to [40] and branch off from the 100,000-step checkpoints of the experiments in Sec. 4.6. Training is conducted on the SFT data with a learning rate of 2 × 10−5 and a global batch size of 256, which matches the pretraining learning rate and batch size at the 100,000-step mark. This setup can be viewed as continuing the learning rate decay on the SFT dataset, consistent with standard practices [40, 41, 47]. 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. Phase

Pre-SFT

Post-SFT

Optim.

Metric

Loss Metrics (↓)

SFT

Gen.

Reasoning

Math

Code

Avg.

OOD MMLU GPQA GPQA-D BBH GSM8k MATH HumanEval MBPP

Acc. (↑) AdamW 1.655 Loss (↓)

1.263

Acc. (↑) 1.647 Loss (↓)

1.258

1.484

37.0 1.031

1.053

35.9 1.489

57.0 1.179

40.0 1.006

48.0

42.0

39.9

1.030

1.803

1.552

Improv. Loss (↑) +0.008 +0.005 +0.083 +0.009 +0.004 -0.005 +0.054 +0.025

+0.023

+0.056 +0.031

50.9 2.138

1.938

25.8 1.842

22.7 1.929

22.7 1.838

All

54.0 1.233

Nexus

52.7 2.221

28.9

41.1

AdamW

Acc. (↑) 1.035 Loss (↓)

1.278

51.4 2.244

28.9 2.006

Nexus

Acc. (↑) 1.028 Loss (↓)

1.274

54.7 2.220

28.9

29.7

39.3

1.990

1.901

1.551

Improv.

Acc. (↑) +3.3 0.0 Loss (↑) +0.007 +0.004 +0.024 +0.016

-3.9 +0.014

50.0

42.0 41.4 1.859 1.583

58.0 1.377

28.0 1.111

40.0 1.077

40.0 40.7 1.957 1.658

62.0 1.299

35.0 1.076

46.0 1.060

1.952

42.7 1.631

-6.3 +4.0 +7.0 +0.024 +0.078 +0.035

+6.0 +0.017

+6.0

+2.0

+0.005 +0.027

33.6

45.6

1.915

1.575

46.0

Nexus still outperforms AdamW after SFT. As shown in Tab. 11, after supervised fine-tuning, Nexus achieves

an average accuracy of 42.7%, surpassing the AdamW baseline by 2.0%. Specifically, Nexus outperforms AdamW by 7.0% on MATH, 6.0% on HumanEval, and 6.0% on MBPP. These results indicate that Nexus does not prematurely compromise the model’s capacity for downstream alignment. Nexus maintains lower SFT loss than AdamW throughout training. We observe that for the pre-SFT

checkpoints, Nexus already yields a lower loss on the SFT dataset compared to AdamW (1.647 vs. 1.655). This result demonstrates that the geometric properties optimized by Nexus during pretraining translate into better generalization even before any explicit fine-tuning. Furthermore, this lower SFT loss is consistently maintained throughout the entire training process, as evidenced by the post-SFT loss (1.028 for Nexus vs. 1.035 for AdamW). These observations indicate the potential of Nexus for continual training and extended optimization phases.

Record · ID 5982 · SHA-256 7859d11fea4ec888
Conceptio Open Knowledge Archive — every document is proof-bundled with source, license, and retrieval metadata.