arXiv CS open access

Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining

2026 · arxiv_cs

arXiv CS · Papers · 2026

deep-learning machine-learning neural-networks

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining Shuchen Zhu1,2 , Yuxin Fang1 , Mingze Wang2,† , Kun Yuan2,†

arXiv:2608.28442v1 [cs.LG] 28 Aug 2026

ByteDance Seed, 2 Peking University †

Corresponding authors

Abstract Pretraining accounts for a large fraction of the total computational cost in LLM training. However, noise-dominant gradients and the highly ill-conditioned loss landscape bring severe challenges. Although modern adaptive optimizers such as AdamW and Muon have achieved great success in large-scale pretraining, their reliance on gradient normalization offers limited mitigation of the ill-conditioned curvature. The progress along flat directions (eigen-directions of small eigenvalues), which dominates the final loss reduction, remains relatively slow. To enhance training dynamics along flat directions, we propose a curvature-conditioned multiscale momentum method with sphere constraints, delivering steady acceleration in LLM pretraining. This multiscale momentum, applied only along flat directions, pairs a slow-decay component for noise reduction with a fast-decay component for rapid curvature adaptation, harnessing their complementary strengths. Crucially, we employ a sphere constraint technique to prevent parameter inflation and excessively rapid effective learning rate decay that would otherwise arise from a naive combination. Extensive experiments show that the proposed method significantly accelerates Muon across diverse architectures (dense, MoE) and model sizes (0.12B–2.3B parameters). Theoretically, we verify the acceleration effect and provide insight into the design principles underlying the flat-direction multiscale momentum. Date: August 31, 2026 Correspondence: Mingze Wang at mingzewang.math@gmail.com, Kun Yuan at kunyuan@pku.edu.cn.

Introduction

The optimizer is one of the most critical factors for training efficiency and continued scaling in the pretraining stage of large language models (LLMs). Several inherent challenges complicate LLM pretraining. Because the total number of training tokens vastly exceeds the per-step batch size, noise dominates the stochastic gradient: its magnitude far exceeds the true gradient signal. Moreover, the loss landscape is highly nonconvex and ill-conditioned, characterized by numerous flat directions alongside a few sharp ones. Notably, the training dynamics along the flat directions primarily drive the final loss reduction [33, 42]. Since the learning rate is constrained by the largest eigenvalue along the sharp directions and the noise level [3, 5], progress along these crucial flat directions remains slow, hampering overall training and potentially compromising stability. Adaptive optimizers such as AdamW [14, 22] and Muon [13, 20] partially alleviate ill-conditioning through coordinate- or block-wise normalization of gradient statistics. They have demonstrated stability and favorable scaling behavior in industrial-scale LLM training. However, this normalization (whitening) only partially compensates for the ill-conditioned curvature, and progress along the flat directions leaves room for improvement. 1

Recent works [37, 40, 54, 56] attempt to address this by estimating flat directions and manually amplifying the learning rate along them. Yet this approach also amplifies the stochastic noise along flat directions, yielding limited improvement and potentially destabilizing training. This motivates the question: How to accelerate training dynamics along flat directions by reducing gradient noise?

We address this with a momentum mechanism tailored to flat directions. Our contributions are summarized as follows: • Algorithm design. We propose a sphere-constrained multiscale momentum method that combines fastand slow-decay momentum components to enhance training dynamics along the flat directions. The multiscale momentum is curvature-conditioned : it activates in flat directions for robust noise reduction and deactivates in sharp directions to preserve stability. On its own, however, it induces norm inflation and effective learning rate collapse. To address this, we introduce sphere constraints to weights with parallel transport on each momentum component, which prevents norm inflation and ensures proper inheritance of momentum across varying tangent spaces. Together, they unlock the potential of multiscale momentum in LLM pretraining. • Empirical evaluation. We evaluate our proposed method (denoted by appending the suffix “M” to the base optimizer, e.g., MuonM) across a wide range of LLM pretraining settings, covering both dense and MoE architectures, model sizes from 0.12B to 2.3B, and diverse learning rate schedules (cosine decay and WSD). The results show that MuonM consistently outperforms Muon with lower terminal loss. Notably, these loss gains persist under extended token budgets, suggesting promising scalability to longer training horizons. • Theoretical analysis. Recent studies [16, 19, 25] show that appropriately designed linear regression models can serve as a useful proxy for LLM pretraining dynamics, capturing key trends in loss evolution. Within this framework, we analyze how slow-decay momentum injection along flat directions accelerates optimization. Our analysis reveals that it enables faster forgetting of accumulated noise while improving the signal learning rate, thereby leading to lower terminal loss. These findings support the design intuition behind the multiscale momentum for flat directions.

1.1

Notations

For a matrix M ∈ Rm×n , we define the matrix sign function as msign(M ) = M (M ⊤ M )−1/2 = (M M ⊤ )−1/2 M , where the matrix inverse square root is interpreted via the Moore-Penrose pseudo-inverse if M is not full rank. Let ∥ · ∥F denote the Frobenius norm. The operator RowScaler (M ) scales each row of M to have a Frobenius norm of r, and Normr (M ) = rM/∥M ∥F . For a positive semi-definite matrix H, the induced semi-norm of a √ ⊤ vector v is defined as ∥v∥H = v Hv. The symbol ⊙ represents the element-wise product. We use the term curvature in an intuitive sense to refer to Hessian-related information (e.g., eigenvalue distributions). ≲, ≳, and ≍ denote O, Ω, and Θ, respectively, to absorb absolute constants.

Preliminaries and Related Works

Empirical studies [10, 35, 48, 51] show that the loss landscape in deep learning is highly ill-conditioned. Specifically, the Hessian possesses a small proportion of large positive eigenvalues (the corresponding eigenspaces are referred to as the sharp directions) and a massive number of near-zero and negative eigenvalues (corresponding to the flat directions). Recent studies [6, 33, 42] show that this landscape structure drives the fast–slow training dynamics: oscillatory behavior along sharp directions without divergence, and slow but persistent progress along flat directions that ultimately dominates the final loss reduction. Thus, accelerating optimization along flat directions is critical, and this is precisely the focus of our work. Landscape geometry and training dynamics in deep learning.

Adaptive optimizers have become indispensable in large-scale LLM pretraining. AdamW [22] has become the de facto choice, owing to its simple and effective coordinate-wise preconditioning. However, its diagonal preconditioner inherently limits its capacity Efficient adaptive optimizer design for LLM pretraining.

to capture cross-parameter correlations within the full curvature. Muon [13, 20] instead applies matrix-level preconditioning to parameter blocks, exhibiting favorable convergence speed and scaling behavior. Recent methods, such as [18, 28, 36], further develop structured matrix preconditioners to better approximate curvature-related information. Our work is complementary to this line of research. Instead of improving the preconditioner structure, we focus on the design of momentum. Heavy Ball momentum [30] has been widely adopted in modern adaptive optimizers to facilitate training. Building upon it, Nesterov-type momentum [47, 50] offers additional steady performance improvements. More recently, multiscale momentum has been introduced to smooth and accelerate the dynamics of momentum methods. AggMo [23] linearly combines multiple momentum buffers with different decay rates to dampen oscillations. AdEMAMix [26] extends this by integrating multiscale momentum with Adam’s preconditioner in training neural networks. GPA [9] employs extra primal averaging to smooth the iterates of Nesterov momentum. SODA [29] and EMA-Nesterov [49] introduce an extra momentum via primal extrapolation and exponential moving average of look-ahead directions respectively to achieve acceleration in pretraining, while accommodating various base optimizers. Different from these approaches which are motivated primarily by deterministic optimization, our work focuses on its role in variance reduction under noise-dominated stochastic environments, as well as on distinct behaviors at different time scales across regions of varying curvature (sharp and flat). In addition, we provide a unified high-order ODE formulation that offers a continuous-time perspective for understanding these different algorithmic forms in Appendix D. Improved momentum mechanism for adaptive optimizers.

Normalization layers in modern LLMs introduce redundant degrees of freedom, making the learning rate alone insufficient to fully characterize the parameter update speed. For instance, an RMSNorm layer of the form W1 γ ⊙ Norm(W2 x) exhibits a radial redundancy structure W1 diag(γ) and norm-scale invariance with respect to W2 (when W2 is learnable). This naturally motivates treating radial and angular updates separately. Recent studies [15, 32, 43, 55] regard the angular update rate ∥∆W ∥F /∥W ∥F as the effective learning rate, which has been empirically shown to play a dominant role in shaping the loss trajectory [17, 21, 45]. The radial evolution determines the weight norm and, in turn, affects the effective learning rate. Weight decay can be interpreted as controlling the noise-induced inflation of the weight norm, thereby preventing the effective learning rate from decaying too rapidly. Reference [8] uses corrected weight decay coefficients to stabilize the weight norm. To achieve finer control, some recent approaches [2, 12, 43, 46] replace weight decay with direct norm constraints to weights, and they have demonstrated promising advantages over weight decay in LLM pretraining. Controlling the effective learning rate: from weight decay to sphere constraints.

Illustrative Examples Motivating Flat-Direction Multiscale Momentum

In this section, we illustrate how varying momentum decay rates affect stochastic optimization dynamics.

3.1

Momentum Accelerates Flat-Direction Optimization Through Variance Reduction

We begin with a simple case: at iteration k, the stochastic gradient follows gk ∼ N (µ, Σ2 ) with a constant true gradient µ, and the heavy-ball momentum updates as mk = (1 − α)mk−1 + αgk . The coefficient 1 − α measures the momentum decay rate: smaller α corresponds to slower decay, and larger α gives faster decay. α The estimator mk has the same mean µ and a reduced equilibrium covariance 2−α Σ2 . Thus, using a smaller α (slower momentum decay) allows mk to estimate the true gradient signal µ with lower variance. However, the above reasoning requires the true gradient to be fixed or nearly unchanged. Since gradient variation is directly governed by its derivative (i.e., the Hessian), slow momentum decay is only safe in flat directions, where small Hessian eigenvalues lead to slowly changing gradients. Applying overly slow momentum decay to sharp directions can introduce bias and cause instability. A simple example of minimizing f (x, y) = x2 − 0.1y with Gaussian gradient noise illustrates this clearly (Figure 1). Figure 1 (a,b) show that compared to a baseline with α = 0.1, using a smaller α (slower decay) for the flat direction y yields a faster and less noisy trajectory, whereas the same slow decay for the sharp direction x amplifies oscillations. While 3

(a) x-trajectory with varying momentum decay (b) y-trajectory with varying momentum decay (c) 80

αx = 0.001 αx = 0.01 αx = 0.1 mul. momentum

−5 0

2000

4000

6000

8000

10000

κ = 1.0 κ = 2.0 κ = 5.0

y-trajectory under different LR ratios

αy = 0.001 αy = 0.01 αy = 0.1 mul. momentum

20 10

0 0

2000

4000

step

6000

8000

10000

2000

step

4000

6000

8000

10000

step

Figure 1 Coordinate-wise trajectories of different momentum methods for minimizing f (x, y) = x2 − 0.1y with the step size η = 0.01 and sign-momentum preconditioning. The x- and y-axes correspond to the sharp and flat directions, respectively. (a, b) employ various momentum decay coefficients αx , αy for x, y respectively. “mul. momentum” refers to (2) with (αfast , αslow , c) = (0.1, 0.001, 0.75) and a biased flat-direction projection Qk (x, y) = √ 21 (0.2x + 0.8y). 2 0.2 +0.8

increasing the learning rate for y also promotes progress, it simultaneously amplifies noise and leads to an oscillatory and relatively more unstable acceleration (Figure 1 (c)). These observations motivate the use of anisotropic momentum decay to better suppress noise: mk = (I − αsharp Pk − αflat Qk )mk−1 + gk ,

wk+1 = wk − ηk mk ,

(1)

where 1 > αsharp ≫ αflat > 0 denote the decay coefficients for sharp and flat directions, and Pk , Qk denote projections onto the sharp and flat subspaces at the k-th step, respectively.

3.2 Flat-Direction Multiscale Momentum: Robust Variance Reduction against Projection Errors In practice, the acceleration provided by (1) is compromised by various imperfections, including errors in estimating the projection onto flat subspaces and unmodeled variations in the Hessian eigenspace. Projection estimation errors introduce a portion of the slow decay for momentum in sharp directions, thereby destabilizing training. At the same time, persistent rotation and alternating transitions of the Hessian eigenspaces inject large bursts of sharp-direction momentum into the flat subspaces. Due to the slow forgetting rate in flat −1 regions (requiring Θ(αflat ) steps), such biases can persist and cause oscillations. To address this, we use multiscale momentum in flat directions, which linearly interpolates a fast-decay component mfast for quick adaptation to gradient and curvature changes and a slow-decay component mslow for variance reduction. At the same time, we retain the fast-decay momentum mfast to help maintain stability. This leads to the following algorithm: # m# k = (1 − α# )mk−1 + α# gk ,

# ∈ {fast, slow},

wk+1 = wk − ηk cQk mslow + (1 − c)mfast k k fast fast = wk − ηk Qk (cmslow + (1 − c)m , k k ) + (1 − c)Pk mk

(2)

where αfast ≫ αslow > 0 denote the momentum decay rates, 0 < c < 1 is the interpolation coefficient, and Pk , Qk are the projections onto sharp and flat subspaces. The interpolation cmslow + (1 − c)mfast also yields k k lower noise in estimating the true gradient compared to mfast under the setting of Section 3.1 (Appendix E.1). k As illustrated in Figure 1 (a, b) and Figure 2, with projection errors and varying curvature, the flat-direction multiscale momentum still yields less noisy acceleration along the flat direction while not amplifying oscillations in the sharp direction.

(a)

αsharp,flat = (0.1, 0.1)

3.5

(b)

αsharp,flat = (0.1, 0.0005)

(c)

αfast,slow = (0.1, 0.0005), flat

(d)

αfast,slow = (0.1, 0.0005), full

50000

start end flat direction

3.0

40000

2.5

step

1.5

30000

2.0

20000

1.0 0.5

10000

0.0 −2

−1

−2

−1

−2

−1

−2

−1

Figure 2 Optimization trajectories on f (x, y) = (1 − x)2 + 1000(y − x2 )2 (global minima at (1, 1)) under Gaussian noise. (a) Standard momentum progresses slowly. (b) Anisotropic momentum (1) advances faster along flat directions but induces back-and-forth oscillations. (c) Flat-direction multiscale momentum (2) with c = 0.5 achieves stable acceleration. (d) Total-space multiscale momentum with c = 0.5 leads to instability, confirming that this technique should be restricted to flat directions.

The intuition behind the robust variance reduction of flat-direction multiscale momentum is as follows. Assume mfast = ∇f (wk ) + bfast + ξkfast , mslow = ∇f (wk ) + bslow + ξkslow with bias ∥bslow ∥ ≫ ∥bfast k k k k k k ∥ (due to slower slow fast forgetting over longer gradient histories), zero mean noise ∥ξk ∥ ≪ ∥ξk ∥ and noise dominating the fast momentum such that ∥ξkfast ∥ ≫ ∥bfast k ∥ + ∥∇f (wk )∥. Then, for an appropriate choice of c, the expected loss decrease conditioned on wk resulting from the Qk -subspace update can be approximated via a second-order Taylor expansion: Ek [f (wk+1 ) − f (wk )]|Qk

≈ − ηEk Qk ∇f (wk ), cQk mslow + (1 − c)Qk mfast k k

η2 ⊤ 2 slow Ek (cmslow + (1 − c)mfast + (1 − c)mfast k k ) Qk ∇ f (wk )Qk (cmk k ) 2 (1 − c)2 η 2 Ek (ξkfast )⊤ Qk ∇2 f (wk )Qk ξkfast . ≈ −η∥Qk ∇f (wk )∥2 + cη Qk ∇f (wk ), −Qk bslow + k {z } | | 2 {z } | {z } +

gradient descent

bias injection

noise variance

The noise-variance term is typically positive, since noise tends to accumulate along directions with larger Hessian eigenvalues [33, 39, 44]. Assuming the bias Qk bslow in Qk subspace is moderate relative to ∇f (wk ), k we can choose suitable c < 1 so that the bias injection is negligible compared to the gradient-descent term, while the noise variance term is simultaneously suppressed, yielding an improved local loss reduction. For the Pk subspace, we note that Ek [f (wk+1 ) − f (wk )]|Pk ≈ 0 as training typically lies in the Edge of Stability regime [3–5]. On the other hand, take c > 0 also reduces the actual learning rate in the Pk subspace, which helps decrease oscillations. We exclude slow momentum from Pk , as slow-momentum bias in sharp directions can be much larger than the gradient signal, which would force c to be vanishingly small and thereby restrict the acceleration in flat directions. Consequently, the overall effect is a net reduction in the local loss.

Obstacles with Naive Multiscale Momentum: Parameter Inflation and Effective Learning Rate Collapse

We now extend the idea of (2) to the Muon optimizer for LLM pretraining. To estimate flat directions, we consider distinguishing them within each parameter block. Following [56], we empirically identify the bottom 90% singular spaces of the momentum matrix as flat directions, based on empirical evidence that the top singular spaces align well with the top eigenspaces of the Hessian. See implementation details in Appendix A.1. For a transformer matrix W ∈ Rm×n , we extend the preconditioner-free algorithm (2) by equipping mfast with Muon preconditioning, while additionally incorporating the slow momentum mslow in flat directions. The

update rule 1 at the k-th iteration is # Mk# = (1 − α# )Mk−1 + Gk , # ∈ {fast, slow}, Wk+1 = (1 − ληk )Wk − ηk msign((1 − αfast )Mkfast + Gk ) + χQk (RowScalemin(1,√ n ) (Mkslow )) ,

(3)

where λ is the weight decay coefficient, χ > 0 is the slow-momentum coefficient, ηk is the scaled learning rate 2 , and Qk : Rm×n → Rm×n is the estimated linear projection to flat directions. Rather than designing a specialized preconditioner for the slow momentum, we simply apply row-wise normalization to match the norm of the msign term. This approach is inspired by the near-zero inter-row Hessian observed in [52], and it proves efficient in practice. Muon with vanilla multiscale momentum

300

Muon Alg. (3) (λ=0.1)

250

Alg. (3) (λ=0.2)

3.2

Muon Alg. (3) (λ=0.1) Alg. (3) (λ=0.2)

kWupkF

validation loss

3.3

Muon with vanilla multiscale momentum

200

3.1

150

100

tokens (billions)

Figure 3 Straightforward application of flat-direction multiscale momentum to Muon on a 0.12B dense model. (Left) Validation loss; (Right) Frobenius norm evolution of an up projection block during training.

The results on 0.12B dense model are in Figure 3 (left), where both the Muon baseline and the value of χ and αslow have been thoroughly tuned (see details in Appendix B.1). The inclusion of slow momentum along flat directions yields a certain reduction in loss. However, the gain diminishes progressively, resulting in a terminal loss improvement of less than 0.01. This behavior is attributed to the amplified weight norm ∥Wk ∥F by slow momentum (e.g., the norm of an up projection block in Figure 3, right) and in turn accelerates k ∥F the decay of the effective learning rate ∥∆W ∥Wk ∥F . We provide an illustrative theoretical explanation to this phenomenon in Appendix E.2, which reveals its intrinsic link to the slow momentum. This mechanism also explains the evolution of the loss gain: a sharp initial loss drop driven by the rapid effective learning rate decay in the early phase, followed by the baseline gradually catching up in the later phase, though the flat-direction slow-momentum acceleration still retains a marginal edge. A natural remedy would be to use a larger weight decay coefficient λ (default 0.1) to restrain the norm inflation. However, even with λ = 0.2 which brings the norm back to the baseline level, the loss becomes even higher. This suggests that simply tuning existing hyper-parameters is insufficient to unleash the acceleration potential of multiscale momentum. We therefore turn to weight-decay-free techniques.

The Flat-Direction Multiscale Momentum Method with Sphere Constraints

We address the issue of weight norm inflation discussed in the previous section by replacing weight decay with sphere constraints. Before presenting the final algorithm, we first give a short review of manifold optimization.

5.1

Preliminaries on Manifold Optimization

Manifold optimization deals with optimization problems where the variable is constrained to lie on a manifold. In this work, we restrict ourselves to the Frobenius sphere constraint SR = {w ∈ Rp : ∥w∥F = R} with R > 0. 1 Here we set the coefficient of G to 1. This choice differs from the EMA format in (2) only by a constant factor and is thus k essentially equivalent. 2We use the RMS norm alignment rescaling in [20].

The tangent space at a point w ∈ SR is Tw SR = {v ∈ Rp : v ⊤ w = 0}. To define the optimization geometry, the tangent space is equipped with a metric (a positive definite quadratic form). Here we simply use the Euclidean metric. At the k-th step, the update direction ∆wk (e.g., gradient or momentum) is projected onto the tangent space via PTwk SR (∆wk ) =

mk ∆w⊤ w ∆wk − Rk2 k wk . A retraction wk+1 = NormR (wk − ηk ∆wk ) is then applied to keep the iterate on the sphere, where ηk is the γ : wk →wk + 1 mk step size. In this framework, momentum thus accumulates the wk Riemannian gradients PTw SR (g), with g denoting the Euclidean wk + 1 gradient. After moving from wk to wk+1 , the momentum vector in Twk SR is transferred to the new tangent space Twk+1 SR . This is T wk →wk + 1 mk accomplished by parallel transport Twk →wk+1 : Twk SR → Twk+1 SR , which moves tangent vectors along the manifold while preserving their norms. Intuitively, parallel transport ensures the momentum Figure 4 Illustration of parallel transport is adjusted to the new geometry, preventing a gradual distortion on the sphere. of the accumulated information. For the Frobenius sphere, parallel transport admits a simple closed form. Omitting it would amplify direction errors in the accumulated gradient history, especially when momentum decays slowly and the tangent spaces vary substantially along the optimization trajectory. We will elaborate on this point in Section 6.1. Given v ∈ Tw SR and a new point w′ ∈ SR (w′ ̸= −w), the transported vector is Tw→w′ (v) = v −

v ⊤ w′ R2 + w ⊤ w ′

(w + w′ ).

(4)

The above discussion can be directly extended to other tensor types such as matrices.

5.2

The Proposed Approach

Motivated by [43], we adopt a fixed Frobenius c to the same norm R, which translates norm constraint (denoted as R) and rescale the optimizer update ∆W the learning rate into an angular velocity on the sphere. To decouple radial and angular updates, we rec , where ∥W c ∥F ≡ R and r ∈ R is a learnable parameterize each transformer weight matrix W as W = rW scalar. Compared to a purely fixed-norm parameterization, the learnable r automatically adapts to the natural heterogeneity across blocks and layers [41]. Vector parameters (e.g., biases) are re-parameterized analogously. Given the token-level heterogeneity in the embedding and output matrices W ∈ RV ×d with vocabulary size c , where γ ∈ RV and V and model dimension d, we adopt a row-wise re-parameterization as W = diag (γ)W c maintains a fixed norm. We denote a base optimizer used on the re-parameterized model by each row of W appending the suffix "S" (sphere) to its name (e.g., MuonS). Directional Sphere Parameterization with Learnable Radius

We now apply the manifold optimization paradigm (parallel transport T and tangent space projection P) from Section 5.1 to the base multiscale momentum method (3) under the sphere constraint described above. We rescale both the fast-momentum and slow-momentum update norms to the sphere radius R, so that the learning rate acts as the angular velocity and simplifies tuning. The resulting algorithm (denoted as MuonM) follows the Nesterov-type momentum update of [13], using the Nesterov-type momentum ffast := (1 − αfast )M fast + PT S Gk in (5). The complete procedure for transformer blocks is summarized M Wk R k k in Algorithm 1, which degenerates to MuonS when χ = 0. We apply Adam to scalar parameters and sphere-constrained Adam to vector parameters. Please refer to Section A.2 for the complete algorithm details. Algorithm Design

Algorithm 1 MuonM for a transformer block matrix W ∈ Rm×n

1: Input: {ηk }, αfast > αslow > 0, χ ≥ 0. 2: for k = 0 to K do # 3: Update multiscale momenta Mk# = (1 − α# )TWk−1 →Wk (Mk−1 ) + PTW SR Gk , 4:

# ∈ {fast, slow}. k Estimate projection Qk onto the flat directions in the tangent space. Then compute the update direction

(5)

fkfast )) + χQk (RowScaleR/√m (Mkslow )). Uk = NormR (PTWk SR msign(M

5: Update parameters Wk+1 = NormR (Wk − ηk Uk ). 6: end for

The row-wise rescaling aligns the norm of the slow momentum with that of the msign term. The flat-direction projection Qk already includes the projection to the tangent space, and is computed by online power iterations ffast . Implementation details are provided in Section A.1. to estimate top singular subspaces of M k Our method incurs only marginal additional overhead in computation, memory, and communication compared to the base Muon optimizer. The extra computation stems mainly from manifold-related operations (tangent space projection, retraction, and parallel transport) and the flat-direction projection Qk . The former involves no dense matrix–matrix products while the latter’s main cost consists of two lightweight msign operations on matrices much smaller than W , which is nearly negligible as analyzed in Section A.1. The additional memory comes from the slow momentum buffer and the small projection matrices used by Qk , totaling roughly the size of one extra set of model parameters. In pretraining, memory is typically dominated by activations, and the extra optimizer state can be further reduced to a marginal amount per GPU via Fully Sharded Data Parallelism (FSDP) [53]. Regarding communication, in large-scale distributed training that employs parallelism strategies such as expert parallelism (EP) or pipeline parallelism (PP), the communication volume is similarly dominated by activations, and the extra communication from the sharded optimizer states is insignificant. Computation, Communication and Memory Overhead.

Experiments

In this section, we evaluate our proposed optimizer on dense and Mixture-of-Expert (MoE) architectures with scales from 0.12B to 2.3B. All models are trained using data from a 350Btoken high-quality pretraining corpus. The training budget is approximately 100 tokens per parameter (TPP) for dense models and approximately 500 tokens per activated parameter for MoE models, substantially exceeding the 20 TPP of the Chinchilla-optimal data [11] and aligning more closely with industrial pretraining practice. Unless otherwise specified, the token batch size is about 0.5M for dense models and 4M for MoE models. More experimental details are in Appendices A, B.

6.1 Validation of Critical Algorithm Components We evaluate the roles of the flat-direction projection Q and parallel transport for momentum on a 0.12B dense model, using MuonM as the baseline (Figure 5). We first replace the flat-direction projection Q with the sharp-direction projection (i.e., the orthogonal complement projection of Q), which leads to a large increase in loss, indicating that applying multiscale momentum along sharp directions is detrimental.

Ablation Studies on Dense-0.12B

relative loss

0.015

MuonM (baseline) MuonM Sharp

0.010

no T : mslow

0.005

no T : mfast, slow

no T : mfast

0.000 − 0.005 − 0.010 4

tokens (billions) Figure 5 Ablation studies on projection direction and parallel transport T . The label "Sharp" corresponds to replacing Q with the sharp-direction projection.

Next, we examine the effect of parallel transport on the fast and slow momentum separately. Removing 8

parallel transport for the fast momentum alone results in a negligible change in the loss trajectory. In contrast, removing it for the slow momentum causes a significant increase in loss, and additionally disabling parallel transport for the fast momentum has little further impact. This suggests that parallel transport is unnecessary for the fast momentum but indispensable for the slow momentum. A plausible explanation is that the fast momentum decays rapidly. Over its short decay horizon, relatively few parameter updates occur and the tangent space changes little. Thus the accumulated error from skipping parallel transport remains small. The slow momentum, on the other hand, decays slowly and spans many updates over a large parameter range, during which the tangent space can change substantially. Parallel transport is therefore required to keep the error accumulation under control. For instance, with a decay factor of 0.95 (half-life ≈ 14 steps) for the fast momentum and 0.999 (half-life ≈ 700 steps) for the slow momentum, this difference in update ranges becomes evident.

6.2

Results on Dense Models

We first compare our sphere constraint strategy (Section 5.2) with two recent methods of the same type: MuonH [43] (Frobenius norm sphere) and SSO [46] (spectral norm sphere). All sphere-constrained methods outperform the weight-decayed Muon baseline, with MuonS (i.e., MuonM without slow momentum) achieving the lowest loss (Figure 6, left). We attribute this advantage to the adaptivity of the learnable radius across heterogeneous blocks. Next, we compare MuonM with other enhanced momentum approaches, including AdEMAMix [26] (using AdamW as the base optimizer) and EMA-Nesterov [49] (using Muon as the base optimizer). MuonM achieves a lower terminal loss than these methods by a clear margin (Figure 6, right). Notably, the loss improvement of MuonM over MuonS (-0.019) and over Muon (-0.028) is substantially larger than the gain from directly applying flat-direction multiscale momentum to Muon (-0.008, see Figure 3). This indicates that the sphere constraint more effectively unlocks Comparison with Related Methods.

the acceleration potential of flat-direction multiscale momentum.

Dense-0.12B

Dense-0.12B 3.25

Muon MuonS MuonH SSO

3.10

validation loss

3.15

3.05 7

Muon MuonS MuonM (ours) AdEMAMix EMA-Nesterov

3.20 3.15 3.10 3.05 6

tokens (billions)

Figure 6 Comparison of MuonS and MuonM with other sphere-constrained methods (left) and momentum-enhancement approaches (right).

We systematically evaluate the acceleration benefits of MuonM against well-tuned MuonS and Muon baselines across dense models ranging from 0.12B to 1.4B parameters (Figures 7 and 6). Across all settings, MuonM consistently achieves the lowest terminal loss, outperforming MuonS and Muon by approximately 0.02 and 0.03, respectively. Main Results.

We evaluate Muon, MuonS, and MuonM on the 0.12B dense model across token budgets ranging from approximately 100 to 1000 TPP, with a peak learning rate sweep performed at each budget. The terminal losses are fitted using L(N ) = L∞ + AN −γ with the terminal loss L(N ), where N is the number of training tokens. As shown in Figure 8, MuonM consistently achieves the lowest terminal loss across all budgets, and the performance gap widens as the training budget increases, demonstrating its potential in long-horizon training. In addition, the searched optimal learning rate of MuonS (and thus MuonM) decays much more slowly (from at 1e-2 at 100 TPP to 7e-3 at 1000 TPP) compared to Scaling Behavior under Extended Training.

Muon MuonS MuonM

2.9 2.8

Muon MuonS MuonM

2.7

validation loss

3.0

Dense-1.4B

Dense-0.6B validation loss

Dense-0.25B

2.6

2.5

2.7 10

tokens (billions)

Muon MuonS MuonM

2.5

2.4

2.3 40

100

tokens (billions)

Figure 7 Validation losses on dense models across different optimizers. MuonM consistently achieves lower loss than all baselines with a notable margin.

Muon, which decays from 5e-3 to 1e-3 over the same range (Table 3). This mild shift in the optimal learning rate makes it more readily transferable across varying training horizons.

Scaling Behavior over Extended Training Horizons 3.25 Muon

3.20

MuonS MuonM

3.15

loss

3.10 3.05 3.00 2.95 2.90

Muon L∞ = 2.9157

MuonS L∞ = 2.9100

MuonM L∞ = 2.8884

100

training tokens (billions) Figure 8 Terminal validation loss on the 0.12B model under extended training. Dots indicate the terminal loss at each token budget with the corresponding tuned learning rate, and the solid curves are the fitted power-law trends.

6.3

Results on MoE Models

Figure 9 shows the comparison of Muon, MuonS, and MuonM on MoE models with 0.64B (A0.13B) and 2.3B (A0.36B) parameters. MuonM consistently achieves notable terminal loss reductions over the well-tuned Muon and MuonS baselines. We further compare these methods under the WSD schedule on the 0.64B (A0.13B) model, where MuonM also attains the lowest terminal loss (Figure 11). These results indicate the advantages of MuonM extend robustly from dense architectures to MoE models.

Theoretical Analysis

In this section, we analyze the training dynamics of flat-direction multiscale momentum and characterize its acceleration over vanilla momentum. Due to the inherent complexity of neural network loss landscapes, we resort to a proxy that captures the essential features of the true optimization trajectory. Following recent practices that demonstrate the surprising effectiveness of quadratic models (e.g., linear regression with a power-decayed spectrum) approximating actual LLM training dynamics [16, 19, 24, 25, 27], we adopt this setting to approximate the pretraining landscape. Specifically, we instantiate our analysis under the feature space linear regression framework established in [16, 38]. 10

MoE-2.3B (A-0.36B)

Muon MuonS MuonM

2.7

validation loss

MoE-0.6B (A-0.13B)

2.6

2.5 30

2.3

2.2

Muon MuonS MuonM

2.4

110

140

170

tokens (billions)

Figure 9 Validation losses on MoE models across different optimizers.

7.1

Problem Formulation

Let X and D denote the data domain and data distribution, respectively. Let ϕ : X → Rd be a feature map with ϕ(x) ∼ N (0, H) for x ∼ D, and let the target function be f ⋆ (x) = ⟨ϕ(x), w⋆ ⟩ with target parameter w⋆ ∈ Rd . Without loss of generality, we assume H = diag(λ1 , . . . , λd ) is diagonal with entries λ1 ≥ · · · ≥ λd > 0. Labels are generated according to y(x) = ⟨ϕ(x), w⋆ ⟩ + ϵ with ϵ ∼ N (0, σ 2 ). We learn f ⋆ using a student model f (x, w) = ⟨ϕ(x), w⟩ with w ∈ Rd by minimizing the population risk R(w) =

1 1 1 Ex∼D (f (x, w) − y(x))2 = u⊤ Hu + σ 2 , 2 2 2

(6)

where u = w − w⋆ . For brevity, we define the excess risk E(w) = 12 u⊤ Hu.

Assumption 7.1 (Power-law model). There exist constants ν > 1, s > 0, such that the Hessian spectrum and target parameter show a power-decayed structure with λj ≍ j −ν ,

|wj⋆ |2 ≍ j −1 λs−1 ≍ j −1−(s−1)ν , j

1 ≤ j ≤ d.

(7)

Here, the spectral exponent ν characterizes the model capacity of ϕ, while s measures the relative difficulty of the task with respect to the model capacity. Let Σ(w) denote the single-batch gradient noise covariance at w. Without loss of generality, we set the batch size to 1. The stochastic gradient at iteration k, computed on the sampled data xk , takes the form gk = ∇w 12 (f (xk , wk ) − y(xk ))2 = Huk + ξk , where the gradient noise satisfies Ek ξk = 0 and has covariance Ek [ξk ξk⊤ ] = Σ(uk ). We consider the multiscale momentum with Nesterov-type momentum # m# k = (1 − α# )mk−1 + gk ,

# ∈ {fast, slow}, ,

wk+1 = wk − ηk χQmslow + γgk + mfast k k

(8)

slow with initializations w0 = mfast −1 = m−1 = 0 and hyper-parameters 0 < αslow < αfast ≤ 1, χ ≥ 0, γ ≥ 0. For an Pd integer 0 ≤ r < d, we denote the flat-direction projection as Q = j=1 qj ej e⊤ j , where ej is the eigenvector of λj and qj = I{j>r} is the indicator for flat directions. For a given learning rate schedule, we define the Pk−1 intrinsic time as the accumulated learning rate Tk := i=0 ηi .

7.2

Theoretical Results

−1 Theorem 7.2 (Asymptotic Estimate of the Excess Risk). Define the constants atop = γ + αfast , afull = −1 −1 −1 γ + αfast + χαslow , p = 2 − ν . We further let aj = atop for j ≤ r and aj = afull for j > r. Suppose

Assumption 7.1 holds and d ≥ 2(r + 1). Then for any k and learning rate schedule {ηi } such that afull Tk ≲ dν and additional conditions in (148) hold, the iterates wk in (8) satisfy EE(wk ) ≍

+ σ2

S(Tk ) | {z }

signal learning

where S(t) = K(t) =

r X j=1 r X

k−1 X i=0

noise accumulation

λj |wj⋆ |2 e−2atop λj t + (r + 1)ν + afull t a2top λ2j e−2atop λj t + a2full

(9)

ηi2 K(Tk − Ti ),

}

−s

, (10)

(r + 1) + afull t

−p

j=1

Remark 7.3 (Scope of Tk ). The regime Tk ≲ dν /afull already encompasses the practically typical range k ≲ d. For Tk ≫ dν /afull , matching upper and lower bounds for EE(wk ) are not available within our current analysis. The same difficulty also arises in the SDE analysis of analogous scenarios of [16]. Under a constant learning rate, the asymptotic estimate admits a simplified form as follows. Corollary 7.4 (Constant Learning Rate). Let ηi ≡ η and r ≥ 4. Suppose the conditions of Theorem 7.2 hold. If we further have −1 ν ν a−1 (11) top r (1 + log(afull /atop )) ≲ ηk ≲ afull d , then

 EE(wk ) ≍ (afull ηk)−s + ησ 2 atop

r X

d X

j −ν + afull

j=1

j=r+1

 ∆

(12)

j −ν  = Eχ (η; k).

Building on this formulation, we demonstrate that for the same αfast , γ (while allowing η to be tuned) in (8), using χ > 0 (i.e., incorporating slow momentum) leads to accelerated convergence. A direct insight is that with a larger χ, taking a proportionally smaller η can reduce both the signal learning error and noise accumulation. We detail the mechanism underlying this improvement below. Corollary 7.5 (Acceleration by Slow Momentum). Suppose the conditions of Corollary 7.4 hold. Define the following quantities as functions of χ: aχ = γ + where ΛP =

j=1 j

−ν

1 αfast

, ΛQ =

χ αslow

, κχ =

j=r+1 j

−ν

aχ ΛP + ΛQ , Nχ = a0 (ΛP + κχ ΛQ ), Γχ = , a0 ΛQ + ΛP /κχ

(13)

. Then the following statements hold:

⋆ 1. (Proxy-optimal loss reduction). The optimal learning rate ηχ,k := argminη Eχ (η; k) satisfies ⋆ ηχ,k ≍

s σ 2 a0 Nχ (aχ k)s

1 s+1

⋆ , Eχ (ηχ,k ; k) ≍

σ2 k

ΛQ +

ΛP κχ

s s+1

(14)

⋆ where the signal learning term, noise accumulation term, and risk Eχ (ηχ,k ; k) are balanced at the same Eχ (η ⋆

;k)

−s/(s+1)

asymptotic order (up to constants). Consequently, the loss reduction ratio is E0 (ηχ,k ⋆ ;k) ≍ Γχ

0,k

⋆ 2. (Iteration complexity speedup). For any ϵ > 0, let kχ (ϵ) solve Eχ (ηχ,k ; kχ (ϵ)) = ϵ. Then we have

kχ (ϵ) ≍ Γ−1 χ . k0 (ϵ)

(15)

The proof follows by direct computation with Weighted AM-GM Inequality. According to the constraints ν ⋆ in Corollary 7.4 and Theorem 7.2, the optimal learning rate ηχ,k is admissible if k satisfies ra0 (1 + log κχ ) ≲ 1 s+1 ν k ≲ adχ and d is sufficiently large. The corollary explicitly quantifies the speedup induced by σ 2 Nχ (aχ )s slow momentum, which becomes more pronounced for larger values of χ/αslow .

Conclusion and Discussion

In this work, we propose a flat-direction multiscale momentum method with sphere constraints to accelerate LLM pretraining. The sphere constraint serves as a more effective foundation than standard weight decay for unleashing the acceleration potential of flat-direction multiscale momentum, by resolving the effective learning rate collapse inherent to the multiscale design. Extensive experiments show that our proposed method (MuonM) consistently accelerates Muon with substantial loss reductions across diverse model architectures and training setups. Our results also indicate that the sphere-constraint technique holds promise as a general building block for developing more efficient optimizers. Several limitations and open questions exist in this study. First, the optimal learning rate schedule for sphere-constrained methods remains an open question. Second, we have not investigated preconditioner design for the slow momentum, which we posit could further accelerate training.

Acknowledgment

We thank Prof. Lexing Ying and Prof. Weijie Su for helpful discussions.

References [1] Noah Amsel, David Persson, Christopher Musco, and Robert M. Gower. The polar express: Optimal matrix sign methods and their application to the muon algorithm. arXiv preprint arXiv:2505.16932, 2025. [2] Kang An, Jiaxiang Li, Donald Goldfarb, and Shiqian Ma. Demystifying manifold constraints in llm pre-training. arXiv preprint arXiv:2605.04418, 2026. [3] Arseniy Andreyev and Pierfrancesco Beneventano. Edge of stochastic stability: Revisiting the edge of stability for sgd. arXiv preprint arXiv:2412.20553, 2025. [4] Arseniy Andreyev, Advikar Ananthkumar, Marc Walden, Tomaso Poggio, and Pierfrancesco Beneventano. Momentum further constrains sharpness at the edge of stochastic stability. In Forty-third International Conference on Machine Learning, 2026. [5] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 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] DeepSeek-AI, Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, et al. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434, 2024. [8] Aaron Defazio. Why gradients rapidly increase near the end of training. arXiv preprint arXiv:2506.02285, 2025. [9] Aaron Defazio, Konstantin Mishchenko, Parameswaran Raman, Hao-Jun Michael Shi, and Lin Xiao. Smoothing diloco with primal averaging for faster training of llms. arXiv preprint arXiv:2512.17131, 2026. [10] Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2232–2241. PMLR, 09–15 Jun 2019. [11] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, et al. Training computeoptimal large language models. arXiv preprint arXiv:2203.15556, 2022. [12] Alexander Hägele, Alejandro Hernández-Cano, Atli Kosson, and Martin Jaggi. Improving neural network training by decoupling the magnitude and direction of weight vectors. arXiv preprint arXiv:2606.25971, 2026. [13] 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/. [14] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. [15] Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning across neural networks. In Forty-first International Conference on Machine Learning, 2024. [16] Binghui Li, Fengling Chen, Zixun Huang, Lean Wang, and Lei Wu. Functional scaling laws in kernel regression: Loss dynamics and learning rate schedules. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [17] Bingrui Li, Jiaxin Wen, Zhanpeng Zhou, Jun Zhu, and Jianfei Chen. Efficient hyperparameter tuning via trajectory invariance principle. arXiv preprint arXiv:2509.25049, 2025. [18] Zichong Li, Liming Liu, Chen Liang, Weizhu Chen, and Tuo Zhao. Normuon: Making muon more efficient and scalable. In Forty-third International Conference on Machine Learning, 2026. [19] Licong Lin, Jingfeng Wu, Sham M. Kakade, Peter Bartlett, and Jason D. Lee. Scaling laws in linear regression: Compute, parameters, and data. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. [20] Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for llm training. arXiv preprint arXiv:2502.16982, 2025. [21] Zihan Liu, Ruiheng Zheng, Shaobo Zhang, Changxin Tian, Kunlong Chen, Zhiqiang Zhang, and Lei Wu. Effective learning rate governs loss dynamics in language model pretraining. arXiv preprint arXiv:2608.24814, 2026.

[22] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. [23] James Lucas, Shengyang Sun, Richard Zemel, and Roger Grosse. Aggregated momentum: Stability through passive damping. arXiv preprint arXiv:1804.00325, 2019. [24] Alexandru Meterez, Pranav Ajit Nair, Depen Morwani, Cengiz Pehlevan, and Sham Kakade. Anytime pretraining: Horizon-free learning-rate schedules with weight averaging. arXiv preprint arXiv:2602.03702, 2026. [25] Alexandru Meterez, Pranav Ajit Nair, Depen Morwani, Cengiz Pehlevan, Sham Kakade, and Alex Damian. A defense of the quadratic model. arXiv preprint arXiv:2607.21716, 2026. [26] Matteo Pagliardini, Pierre Ablin, and David Grangier. The adEMAMix optimizer: Better, faster, older. In The Thirteenth International Conference on Learning Representations, 2025. [27] Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+3 phases of compute-optimal neural scaling laws. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. [28] Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with norm-constrained LMOs. In Forty-second International Conference on Machine Learning, 2025. [29] Thomas Pethick, Wanyun Xie, Roman Machacek, and Volkan Cevher. Optimistic dual averaging unifies modern optimizers. arXiv preprint arXiv:2605.11172, 2026. [30] B.T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5):1–17, 1964. ISSN 0041-5553. [31] Thomas Robert, Mher Safaryan, Ionut-Vlad Modoranu, and Dan Alistarh. LDAdam: Adaptive optimization from low-dimensional gradient statistics. In The Thirteenth International Conference on Learning Representations, 2025. [32] Simon Roburin, Yann de Mont-Marin, Andrei Bursuc, Renaud Marlet, Patrick Pérez, and Mathieu Aubry. Spherical perspective on learning with normalization layers. arXiv preprint arXiv:2006.13382, 2022. [33] Minhak Song, Kwangjun Ahn, and Chulhee Yun. Does SGD really happen in tiny subspaces? In The Thirteenth International Conference on Learning Representations, 2025. [34] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024. [35] Qian-Yuan Tang, Yufei Gu, Yunfeng Cai, Mingming Sun, Ping Li, Zhou Xun, and Zeke Xie. Investigating the overlooked hessian structure: From CNNs to LLMs. In Aarti Singh, Maryam Fazel, Daniel Hsu, Simon Lacoste-Julien, Felix Berkenkamp, Tegan Maharaj, Kiri Wagstaff, and Jerry Zhu, editors, Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 58805–58831. PMLR, 13–19 Jul 2025. [36] Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham M. Kakade. SOAP: Improving and stabilizing shampoo using adam for language modeling. In The Thirteenth International Conference on Learning Representations, 2025. [37] Jinbo Wang, Mingze Wang, Zhanpeng Zhou, Junchi Yan, Weinan E, and Lei Wu. The sharpness disparity principle in transformers for accelerating language model pre-training. In Forty-second International Conference on Machine Learning, 2025. [38] Jinbo Wang, Binghui Li, Zhanpeng Zhou, Mingze Wang, Yuxuan Sun, Jiaqi Zhang, Xunliang Cai, and Lei Wu. Fast catch-up, late switching: Optimal batch size scheduling via functional scaling laws. In The Fourteenth International Conference on Learning Representations, 2026. [39] Jinbo Wang, Mingze Wang, Jiaqi Zhang, Wei Wang, Peng Pei, Xunliang Cai, Weinan E, and Lei Wu. Gradpower: Powering gradients for faster language model pre-training. In Forty-third International Conference on Machine Learning, 2026. [40] Mingze Wang, Jinbo Wang, Haotian He, Zilin Wang, Guanhua Huang, Feiyu Xiong, Zhiyu li, Weinan E, and Lei Wu. Improving generalization and convergence by enhancing implicit regularization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.

[41] Mingze Wang, Shuchen Zhu, Yuxin Fang, Binghui Li, Kai Shen, and Shu Zhong. Negligible in size, significant in effect: On scale vectors in large language models. arXiv preprint arXiv:2605.26895, 2026. [42] 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. [43] Kaiyue Wen, Xingyu Dang, Kaifeng Lyu, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them ii: Hyperball optimization. arXiv preprint arXiv:2606.16899, 2026. [44] Lei Wu, Mingze Wang, and Weijie J Su. The alignment property of SGD noise and how it helps select flat minima: A stability analysis. In Advances in Neural Information Processing Systems, 2022. [45] Yihao Xiao, Jialong Sun, Zitian Gao, Zeming Wei, Chutian Wang, Ran Tao, Jiaye Teng, and Bryan Dai. Hyperball may not be a free lunch. arXiv preprint arXiv:2607.22444, 2026. [46] Tian Xie, Haoming Luo, Haoyu Tang, Hu Yiwen, Jason Klein Liu, Qingnan Ren, Yang Wang, Xin Zhao, Rui Yan, Bing Su, Chong Luo, and Baining Guo. Controlled LLM training on spectral sphere. In Forty-third International Conference on Machine Learning, 2026. [47] Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. arXiv preprint arXiv:2208.06677, 2024. [48] Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pages 581–590. IEEE, 2020. [49] Chung-Yiu Yau, Dawei Li, Athanasios Glentis, Valentyn Boreiko, Hoi-To Wai, and Mingyi Hong. Ema-nesterov: Stabilizing nesterov’s lookahead for accelerated deep learning optimization. arXiv preprint arXiv:2605.25395, 2026. [50] Huizhuo Yuan, Yifeng Liu, Shuang Wu, Zhou Xun, and Quanquan Gu. MARS: Unleashing the power of variance reduction for training large models. In Forty-second International Conference on Machine Learning, 2025. [51] Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why transformers need adam: A hessian perspective. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. [52] Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P Kingma, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. In The Thirteenth International Conference on Learning Representations, 2025. [53] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. Pytorch fsdp: Experiences on scaling fully sharded data parallel. arXiv preprint arXiv:2304.11277, 2023. [54] Wenjie Zhou, Bohan Wang, Wei Chen, and Xueqi Cheng. Bsfa: Leveraging the subspace dichotomy to accelerate neural network training. arXiv preprint arXiv:2510.25244, 2025. [55] Zhanpeng Zhou, Wenjie Zhou, Yufei Gu, Juqiu Wang, Guanwei Zhang, Ruobing Xie, Xianbiao Qi, Yushi Bai, Kaifeng Lyu, and Shunyu Yao. From lr to elr: A better heuristic for pretraining dynamics, 2026. URL https://hy.tencent.ai/research/elr. [56] Shuchen Zhu, Rizhen Hu, Mingze Wang, Mou Sun, Xue Wang, Kun Yuan, and Zaiwen Wen. Accelerating llm pre-training through flat-direction dynamics enhancement. arXiv preprint arXiv:2602.22681, 2026.

Appendix A Algorithm Details

. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

A.1 Estimating Flat Directions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

A.2 Update Rules of MuonM for Different Parameter Types . . . . . . . . . . . . . . . . . . . . .

B Experimental Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

B.1 Experimental Details of Section 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

B.2 Experimental Details of Section 6.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

C Additional Experimental Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

C.1 Results with WSD Learning Rate Schedules . . . . . . . . . . . . . . . . . . . . . . . . . . . .

D A Continuous-Time Perspective for Unifying Enhanced Momentum Methods . . . . . . .

D.1 Unified ODE Representation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

E Illustrative Theoretical Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

E.1 Variance Reduction of Multiscale Momentum . . . . . . . . . . . . . . . . . . . . . . . . . . .

E.2 Weight Norm Inflation by Slow Momentum . . . . . . . . . . . . . . . . . . . . . . . . . . . .

F Proof . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

F.1 Useful Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Algorithm Details

A.1

Estimating Flat Directions

Coverage Score

In this study, we focus on separating sharp and Row Hessian (Top-4) Covered by G TG (Top-6) Column Hessian (Top-4) Covered by GG T (Top-6) flat directions within each block. The construction Mean Value Mean Value 1.0 1.0 of the flat-direction projection is motivated by the 0.8 0.8 empirical observation in [56] that, for a transfor0.6 0.6 mation block W ∈ Rm×n , the top eigenspaces of 0.4 0.4 G⊤ G and GG⊤ are highly aligned with the per-row 0.2 0.2 0.0 0.0 and per-column Hessians, respectively (Figure 10). 0 5 10 15 20 25 30 0 2 4 6 8 10 12 14 16 Row Index Column Index Therefore, at the k-th iteration, we estimate the top singular space of the Nesterov-type momentum Figure 10 (Figure 3 in [56]) Coverage of the top eigenspaces ffast := (1 − αfast )M fast + Gk , which serves as a M k k of row (column) Hessians by those of G⊤ G (GG⊤ ) for the stabilized proxy for the gradient Gk , as the sharp up_proj block of an FFN layer in a toy dense language directions. Specifically, we maintain a projection model, where G ∈ R32×16 . A higher score indicates a greater degree of containment. matrix P right ∈ Rn×d (d ≪ n, e.g., d = 0.1n in our experiments) that captures the sharp directions of the row-wise Hessian. At iteration k, P right is updated via online power iteration to track the top d right singular spaces: ⊤ right fast fast right f f P = msign Mk Mk P , k−1

and the corresponding projection for the column-wise Hessian is then computed as ffast P right ∈ Rm×d . Pkleft = msign M k k Finally, for any U ∈ Rm×n , the flat-direction projection is given by Qk (U ) = PTWk SR Im − Pkleft (Pkleft )⊤ U In − Pkright (Pkright )⊤ , where I denotes the identity matrix, PTWk SR is the projection to the tangent space at Wk . Similarly, [31] applies the power method to estimate the top eigenspaces. Here we employ the more efficient msign operator to estimate the subspace projection, replacing the computationally expensive QR decomposition. For the embedding and output matrices W ∈ RV ×dmodel , where V and dmodel denote the vocabulary size and model dimension, respectively, we simply select the d rows with the largest norm of the second moment in Adam. This choice is motivated by [52], which shows that the token-wise Hessian exhibits an approximately block-diagonal structure across rows, each row (token) corresponding to a dense sub-block. In practice we set d ≪ V (e.g., d = 0.1V in our experiments). For vector and scalar parameters, we do not apply flat-direction acceleration. The computational cost (in FLOPs) for computing M ∈ Rm×n in computing msign(M ) via Newton-Schulz (NS) iterations is approximately 6sr2 c, where c = max{m, n}, r = min{m, n}, and s denotes the number of iterations. Consequently, compared to performing NS iterations on the full-dimensional matrix of size m × n, employing the right and left projection matrices of sizes m × d and n × d reduces the computation to approximately (d/n)2 of the original cost (e.g., 0.1 × 0.1 = 0.01) , rendering it nearly negligible. The additional memory overhead primarily stems from storing the right projection matrix of size n × d, which is significantly smaller than that required for the momentum. Computation and Memory Overhead of Flat-Direction Projection.

A.2

Update Rules of MuonM for Different Parameter Types

Notations.

For a matrix A, we use A[i, :] to denote its i-th row. For a vector v, v[i] denotes its i-th element.

We categorize the parameters into the following three types based on their structure:

1. Matrix parameters: Transformer weight matrices, such as Q, K, V, O, and FFN blocks. We apply MuonM (Algorithm 1) to these. 2. Vector parameters: biases, learnable vectors in RMSNorm layers, each row of embedding/output layers (corresponding to each token), and the learnable grouped radii introduced by the sphere constraint. For these, we use Adam with sphere constraints, with slow momentum applied to partial vector blocks. 3. Scalar parameters: native architectural scalars and the learnable radii from the sphere constraint. We directly apply Adam to these. The specific update rules are as follows. • Scalar. It is initialized as 1 and trained directly by Adam without weight decay. The learning rate is the same as other vector and matrix parameters. In practice, if several scalars are interchangeable (e.g., the scalar multipliers of a weight matrix and its preceding RMSNorm, of the query and key matrices in attention, or of the up and down matrices in the FFN), we group them into a single learnable parameter. • Vector. For a vector γ ∈ Rd , we parameterize it as γ = r γ b, where r ∈ R and γ b ∈ Rd has a fixed norm R. At the k-th iteration, it is updated by adam mk = (1 − αfast ) Tγbk−1 →bγk (mk−1 ) + PTγbk SR gk ,

vk = (1 − βadam ) vk−1 + βadam gk ⊙ gk , m k γ bk+1 = NormR γ bk − ηk PTγbk SR √ , vk + ϵ

(16)

where gk is the stochastic gradient. c , and we impose row-wise sphere • Embedding and Output. They are re-parameterized as W = diag(γ) W V ×d c constraints on W ∈ R , where the i-th row has a fixed Frobenius norm Ri . The vector γ is treated as a vector parameter and follows the vector update rules above. At the k-th iteration, for each row i we update the momenta as # adam Mk# [i, :] = (1 − α# ) TWk−1 [i,:]→Wk [i,:] (Mk−1 [i, :]) + PTWk [i,:] SRi Gk [i, :],

# ∈ {fast, slow},

(17)

where Gk is the stochastic gradient, and parallel transport is performed row-wise. Then the flat-direction projection Qk is computed as described in Section A.1, yielding a vector qk ∈ [0, 1]V , where a value of 1 indicates that the corresponding row is selected as a flat direction. The update direction for each row is then M ffast [i, :] Uk [i, :] = NormRi PTWk [i,:] SRi p k + χ qk [i] NormRi Mkslow [i, :] , (18) Vk [i, :] + ϵ where the second moment Vk is maintained as Vk = (1 − βadam ) Vk−1 + βadam Gk ⊙ Gk . Finally, the parameters are updated row-wise via Wk+1 [i, :] = NormRi Wk [i, :] − ηk Uk [i, :] . (19)

Experimental Details

We use 6 Newton-Schulz (NS) iterations with coefficients in Polar Express [1] to implement the msign operation, which achieves higher precision than the method in [13]. The peak learning rate is searching in {5e-4, 7e-4, 1e-3, 2e-3, 3e-3, 5e-3, 7e-3, 1e-2, 1.5e-2, 2e-2} for all optimizers.

B.1

Experimental Details of Section 4

In our hyperparameter search, the optimal learning rate for Muon was found to be 0.005. The sharp subspace ratio was set to 10%, following the procedure described in Appendix A.1. Beyond the learning rate, we set αfast = 0.05, searched over αslow ∈ {0.005, 0.002, 0.001, 0.0005} and χ ∈ {0.1, 0.2, 0.5}. The final best configuration is lr = 0.005, αslow = 0.002, and χ = 0.1. 19

B.2

Experimental Details of Section 6.2

We use the cosine decay learning rate schedules for all optimizers. For MuonH and SSO, the optimal learning rates found by grid search are 1 × 10−2 and 5 × 10−3 , respectively. We also tuned the radius for SSO over {1.0, 2.0, 5.0} and found that 2.0 yields the lowest loss. For AdEMAMix, we searched β3 ∈ {0.01, 0.001, 0.0001} and α ∈ {1, 3, 5, 8}, obtaining the best configuration (lr, β3 , α) = (3 × 10−3 , 1 × 10−3 , 5.0). Following [26], the warmup iterations Tα and Tβ3 are set to the total number of iterations. For EMA-Nesterov, we searched γ ∈ {0.99, 0.995, 0.999} and βmax ∈ {0.3, 0.5, 0.7, 0.9}, and selected (lr, βmax , γ) = (3 × 10−3 , 0.9, 0.99). The coefficient schedule of βt involves two transition points Tw and Tr , and we adopted the recommended values Tw = 0.3 T and Tr = 0.8 T from [49], where T denotes the total number of iterations. Comparing Different Optimizers on the 0.12B Dense Model.

Muon Configuration. We follow the standard Muon setup described in [20]. Muon is applied exclusively to all 2D transformer weight matrices, while AdamW handles all other parameters (RMSNorm, embeddings, and p the output layer). For a Muon-updated matrix of shape m × n, the learning rate is scaled by 0.2 max{m, n} to align the RMS of the updates with that of AdamW. We employ Nesterov momentum with decay rate θmuon = 0.95, weight decay λ = 0.1, and a gradient clipping threshold of 1.0. MuonM Configuration. Unless otherwise specified, we uniformly adopt χ = 0.2, αfast = 0.05, αslow = 0.001,

and set the top 10% subspace as sharp directions for each block using MuonM. Given the rapid gradient variations during the early training phase, we apply a warmup mechanism to the slow momentum coefficient αslow as  αfast , t ≤ Twarmup ,      αfast , Twarmup < t ≤ 2Twarmup , αslow,t = αfast /αslow −1  1 + Twarmup (t − Twarmup )     αslow , t > 2Twarmup , where Twarmup is the warmup iterations for the learning rate. Models.

We utilize two popular classes of LLM architectures in the experiments:

• Dense. We adopt a dense decoder-only Transformer architecture that uses Rotary Positional Encoding (RoPE) [34], SwiGLU feed-forward layers, RMSNorm, and a Pre-Norm design. Multi-head attention (MHA) is employed in self attention. Detailed configurations are provided in Table 1. • MoE. The Mixture-of-Experts (MoE) architecture we use is built on a decoder-only Transformer with Multi-head Latent Attention (MLA) [7]. Each MoE layer employs one shared expert together with partially activated experts per token. The first layer is configured as dense. We use an auxiliary-loss-free load balancing strategy. RMSNorm with a Pre-Norm design and RoPE (applied within the decoupled rotary components of MLA) are used throughout. Detailed configurations are provided in Table 2. Table 1 Dense LLaMA model configurations. All models are trained with a default sequence length of 1024.

Acronym

Size

dmodel

dFFN

nhead

nlayer

Batch size

lr_max Muon

MuonS

MuonM

Dense (0.12B)

0.12B

768

2048

480

5e-3

1e-2

7e-3

Dense (0.25B)

0.25B

768

2048

512

3e-3

7e-3

4.5e-3

Dense (0.6B)

0.6B

1280

3412

512

1e-3

5e-3

3.5e-3

Dense (1.4B)

1.4B

2048

5461

512

7e-4

3e-3

2e-3

We evaluate two common learning rate schedules with min{1000, T /50} warm-up steps in all experiments (T is the total steps): Learning Rate Schedulers.

• cos (cosine decay): a linear warm-up to peak lr_max, followed by cosine decay to a terminal learning rate lr_min = 0.05 × lr_max. 20

Table 2 MoE model configurations. All models are trained with a default sequence length of 4096 and batch size of 1024. Acronym

Size

Activated Size

dmodel

nhead

nlayer

MoE-0.6B

0.64B

0.13B

704

MoE-2.3B

2.27B

0.36B

1280

dFFN Dense

Routed Expert

Shared Expert

3584

672

7168

704

1536

nexpert

MLA Dimensions

lr_max

rkv

dqk

drope qk

Muon

MuonS

4/32

704

240

128

7e-3

1e-2

7e-3

4/64

1536

512

128

3e-3

7e-3

4.5e-3

MuonM

• wsd (warmup-stable-decay): a linear warm-up to lr_max, followed by√ a stable phase with constant learning rate of lr_max (up to 70% of total steps), and then a minus-sqrt (1− t) decay to lr_min = 0.05×lr_max over the final 30% steps. The searched peak learning rates for Muon and MuonS on the 0.64B MoE model are 5e-3 and 7e-3 respectively. Then we directly set the peak learning rate of MuonM to 4.5e-3 ≈ 7e-3×2/3. Table 3 Optimal peak learning rates for Muon and MuonS at various training steps on the 0.12B dense model. The peak learning rate of MuonM is set to approximately 2/3 of that of MuonS without search.

Iterations

20k

35k

50k

75k

100k

150k

200k

Muon MuonS MuonM

5e-3 1e-2 7e-3

3e-3 1e-2 7e-3

2e-3 7e-3 4.5e-3

1e-3 7e-3 4.5e-3

Additional Experimental Results

C.1

Results with WSD Learning Rate Schedules

Figure 11 shows the loss trajectories of Muon, MuonS, and MuonM under the WSD learning rate schedule. MuonM achieves the lowest terminal loss. The peak learning rates for Muon, MuonS, and MuonM are 0.005, 0.007 and 0.0045, respectively. MoE-0.6B (A-0.13B), WSD Muon MuonS MuonM

validation Loss

2.7

2.6

2.5 30

tokens (billions)

Figure 11 Validation losses of different optimizers on the 0.6B MoE model under the WSD schedule.

A Continuous-Time Perspective for Unifying Enhanced Momentum Methods

In this section, we establish the connections among AdEMAMix, GPA, EMA-Nesterov, SODA, and our method: all share the same continuous-time ODE form when preconditioning is ignored. For GPA and EMA-Nesterov, we assume the underlying base optimizer is Heavy Ball momentum. For brevity, our ODE derivation is conducted under gradient-noise-free conditions and with constant hyper-parameters. Let gk = ∇f (wk ) denote the gradient and Pk the preconditioner at the k-th iteration. Let h > 0 denote the discretization interval and tk = kh. For the EMA coefficient ρ, we use the scaling ρ = e−h/τ ≈ 1 − hτ where τ > 0 is the corresponding time constant. The algorithmic details and continuous limits are as follows.

We illustrate our approach using MuonM as an example. Omitting the sphere constraint and Muon preconditioning yields the following discrete-time recursion (constants have been absorbed into step sizes for notational simplicity): Our Approach

# m# k = (1 − α# )mk−1 + α# gk ,

# ∈ {fast, slow}, .

wk+1 = wk − ηk χmslow + γgk + mfast k k

(20)

Under a minor change of notation, the continuous-time counterpart takes the form ṁf = −τf−1 mf + ∇f (w),

ṁs = −τs−1 ms + ∇f (w),

(21)

ẇ = −af mf − as ms − ag ∇f (w).

d d Let gt = ∇f (wt ). The momentum equations can be written as 1 + τf dt mf = τf gt , 1 + τs dt ms = τs gt . d d d wt = −af mf − as ms − ag gt , we obtain Applying the operator 1 + τf dt 1 + τs dt to dt d d d d 0 = 1 + τf 1 + τs wt + af τf 1 + τs gt dt dt dt dt d d d + a s τs 1 + τ f gt + ag 1 + τf 1 + τs gt . dt dt dt Expanding the operator products gives the third-order ODE d2 d d3 wt + (τf + τs ) 2 wt + wt + (af τf + as τs + ag ) ∇f (wt ) 3 dt dt dt d d d d2 2 3 2 + [τf τs (af + as ) + ag (τf + τs )] ∇ f (wt ) wt + ag τf τs ∇ f (wt ) wt , wt + ∇ f (wt ) 2 wt , dt dt dt dt

0 = τf τs

where we also use

d d d2 d2 3 2 ∇f (w ) = ∇ f (w ) w , w wt . + ∇ f (w ) t t t t t dt2 dt dt dt2

GPA maintains a base sequence zk , an exponentially averaged sequence xk , and an interpolation point yk at which the stochastic gradient is evaluated [9]. The generic GPA framework permits an arbitrary base optimizer. We instantiate it with preconditioned Heavy Ball: GPA

yk = µy xk + (1 − µy )zk ,

bk = ρH bk−1 + (1 − ρH )gk ,

dk = Pk bk ,

(22)

zk+1 = (1 − ηk λ)zk − ηk dk ,

xk+1 = µx xk + (1 − µx )zk+1 . Here bk is the Heavy-Ball state, µx controls primal averaging, and µy controls the interpolation between the averaged and base iterates. The continuous time limit of (22) without preconditioner and weight decay can be written as ḃ = −τf−1 b + τf−1 ∇f (y),

ż = −b,

ẋ = −τs−1 x + τs−1 z,

(23)

y = µy x + (1 − µy )z.

The last two equations imply (1 + τs

d d )y = 1 + (1 − µy )τs z. dt dt 22

(24)

The first two equations give

d d ) z + ∇f (y) = 0. dt dt Eliminating z yields the third-order GPA (Heavy-Ball) ODE

(25)

(1 + τf

d d d d )(1 + τs ) y + 1 + (1 − µy )τs ∇f (y) = 0. dt dt dt dt

(26)

d2 d d d3 y + τ + τ yt + yt + ∇f (yt ) + (1 − µy )τs ∇2 f (yt ) yt = 0. t f s 3 2 dt dt dt dt

(27)

(1 + τf Namely τf τs

AdEMAMix augments AdamW with a second, substantially slower first-moment EMA [26]. Let mfk and msk denote the fast and slow first moments, and let vk denote the second moment. Its complete update is mfk , mfk = β1 mfk−1 + (1 − β1 )gk , m b fk = 1 − β1k AdEMAMix

msk = β3,k msk−1 + (1 − β3,k )gk , vk = β2 vk−1 + (1 − β2 )gk⊙2 ,

vbk =

m b fk + ak msk . wk+1 = (1 − ηk λ)wk − ηk √ vbk + ϵ

vk , 1 − β2k

(28)

After removing the second-moment preconditioner, bias-correction and weight decay, the continuous state equations are ṁf = −τf−1 mf + τf−1 ∇f (w), ṁs = −τs−1 ms + τs−1 ∇f (w),

(29)

ẇ = −af mf − as ms .

d d Applying (1 + τf dt )(1 + τs dt ) to the parameter equation and eliminating the two momentum states gives

d d d d d )(1 + τs ) w + af (1 + τs )∇f (w) + as (1 + τf )∇f (w) = 0. dt dt dt dt dt

(30)

d2 d3 d d w + τ + τ wt + wt + (af + as )∇f (wt ) + (af τs + as τf )∇2 f (wt ) (wt ) = 0. t f s dt3 dt2 dt dt

(31)

(1 + τf Equivalently, τf τs

EMA-Nesterov is a wrapper that replaces the one-step Nesterov look-ahead direction by an EMA of parameter increments [49]. Let zk denote the EMA look-ahead state. When the wrapped base optimizer is preconditioned Heavy Ball momentum, the update rule is EMA-Nesterov

yk = xk + βk zk , bk = ρH bk−1 + (1 − ρH )gk ,

xk+1 = (1 − ηk λ)yk − ηk Pk bk ,

(32)

zk+1 = ρz zk + (1 − ρz )(xk+1 − xk ).

If Ak denotes the base-optimizer map, the original wrapper is xk+1 = Ak (xk + βk zk ) followed by the final line of (32). The continuous time limit of (32) without preconditioning and weight decay is ḃ = −τf−1 b + τf−1 ∇f (x),

ż = −τs−1 z + τs−1 ẋ,

ẋ = βz − b.

(33)

Eliminating b and z gives d d d d d ) (1 + τs ) − β x + (1 + τs )∇f (x) = 0. dt dt dt dt dt

(34)

d d2 d3 d x + τ + τ − βτ xt + 1 − β xt + ∇f (xt ) + τs ∇2 f (xt ) xt = 0. t f s f 3 2 dt dt dt dt

(35)

(1 + τf Equivalently, τf τs

SODA combines gradient averaging, dual optimism, primal averaging, and primal extrapolation [29]. Its full update is mk+1 = (1 − αk )mk + αk ∇f (yk ), SODA

m̄k+1 = (1 − ᾱk )mk+1 + ᾱk ∇f (yk ),

zk+1 ∈ ∂h∗k (−γk m̄k+1 ) = arg min {γk ⟨m̄k+1 , z⟩ + hk (z)} , z

(36)

xk+1 = (1 − λk )xk + λk zk+1 ,

yk+1 = (1 − λ̄k )xk+1 + λ̄k zk+1 .

We consider the Euclidean mirror map h(z) = 12 ∥z∥2 , and freeze these coefficients with new notations. The continuous limit is ṁ = −τ −1 m + τ −1 ∇f (y), (37) ẋ = −λ1 x − η(m + ω∇f (y)), y = (1 − λ2 )x − ν(m + ω∇f (y)).

Define the quantities s = m + ωg(y), δ = ηa + νλ1 . The state equations then become ẋ = −λ1 x − ηs,

y = (1 − λ2 )x − νs.

(38)

Moreover, ṡ = ṁ + ω∇2 f (y)ẏ = −τ −1 s +

1+ω ∇f (y) + ω∇2 f (y)ẏ. τ

(39)

Differentiating y = ax − νs and using (38)–(39) yield ν(1 + ω) I + νω∇2 f (y) ẏ = −λ1 y + qs − g(y), τ

(40)

where q = ν τ −1 − λ1 − η(1 − λ2 ).

Differentiating (40) and eliminating s and ṡ give a closed equation for y: " # 1 λ1 (1 + ω)δ 2 3 2 I + νω∇ f (y) ÿ + νω∇ f (y)[ẏ, ẏ] + λ1 + I + ℓ∇ f (y) ẏ + y + ∇f (y) = 0, τ τ τ

(41)

where ℓ = ν(1+ω) + ωδ, and the λτ1 y term denotes the weight decay. Optimistic momentum and primal τ −1 extrapolation together yield an implicit preconditioner I + νω∇2 f (y) .

D.1

Unified ODE Representation

The preceding derivation indicates that these continuous-time limits can be cast in a unified form d d A wt + B ∇f (wt ) = 0, dt dt 24

(42)

d where A, B are matrix-valued polynomials in the differential operator dt : X X i i n m d d d d A = , B = . Ai Bi dt dt dt dt i=0 i=0

(43)

For two-timescale momentum, A is of degree three. This polynomial framework naturally generalizes to the n-timescale case, where the degree of A becomes n + 1.

Illustrative Theoretical Analysis

E.1

Variance Reduction of Multiscale Momentum

In this section, we analyze the variance of multiscale momentum through an illustrative Gaussian example, demonstrating the variance-reduction effect of introducing a slow momentum component. Consider mslow = (1 − αslow ) mslow k k−1 + αslow gk ,

(44)

mfast = (1 − αfast ) mfast k k−1 + αfast gk , mk = c mslow + (1 − c) mfast k k ,

i.i.d.

where 0 < αslow < αfast ≤ 1 and gk ∼ N (µ, Σ2 ). We are interested in the stationary covariance of the update mk which preserves the stationary mean, i.e., E[m∞ ] = µ. Let ϵk = gk − µ denote the zero-mean noise. Unrolling the recursions, the noise components of the two momenta are mslow −µ= k

k X (1 − αslow )j αslow ϵk−j ,

mfast −µ= k

j=0

k X j=0

(1 − αfast )j αfast ϵk−j .

(45)

Therefore, in stationarity, Cov(mslow ∞ )=

∞ X 2 (1 − αslow )2j αslow Σ2 ,

Cov(mfast ∞ )=

j=0

∞ X j=0

2 (1 − αfast )2j αfast Σ2 ,

∞ X fast Cov(mslow (1 − αslow )j (1 − αfast )j αslow αfast Σ2 . ∞ , m∞ ) =

(46)

j=0

Summing the geometric series, it follows that 2 2 c2 αslow (1 − c)2 αfast 2c(1 − c) αslow αfast + + Σ2 . Cov(m∞ ) = 1 − (1 − αslow )2 1 − (1 − αfast )2 1 − (1 − αslow )(1 − αfast )

(47)

Simplifying yields

c2 αslow (1 − c)2 αfast 2c(1 − c) αslow αfast + + Σ2 . 2 − αslow 2 − αfast αslow + αfast − αslow αfast {z } |

Cov(m∞ ) =

(48)

:= f (c)

Differentiating with respect to c gives f ′ (c) = −

2(αfast − αslow ) αfast (2 − αslow ) − 2(αfast − αslow ) c (αslow + αfast − αslow αfast )(2 − αslow )(2 − αfast )

< 0,

(49)

since every factor is positive for 0 < αslow < αfast ≤ 1 and c ∈ [0, 1]. Hence f is strictly decreasing on [0, 1]. Increasing the weight c of the slow momentum monotonically reduces the stationary variance. In particular, the two limiting regimes are αfast αslow f (0) = (fast momentum only), f (1) = (slow momentum only). (50) 2 − αfast 2 − αslow 25

αslow (2 − αfast ) Thus incorporating slow momentum reduces the stationary variance by a factor of up to ≈ αfast (2 − αslow ) αslow relative to fast momentum alone. αfast

E.2

Weight Norm Inflation by Slow Momentum

In this section, we provide a theoretical understanding of the role of slow momentum in norm inflation. Motivated by the fact that noise dominates in LLM pretraining, the intuition behind norm inflation for a scale-invariant loss L (i.e., L(cW ) = L(W ) for any c > 0) is that a Brownian-motion-like movement drives the norm to blow up, while weight decay can suppress it. We therefore approximately assume that the stochastic i.i.d. gradients gk ∼ N (0, σ 2 Id ). For brevity, we assume that the projection onto the flat directions Q is constant across iterations and has rank q ≈ d. We then consider the update with weight decay coefficientλ > 0 # m# k = (1 − α# )mk−1 + gk ,

# ∈ {fast, slow},

(51)

χ wk+1 = (1 − ηλ)wkχ − η(χQmslow + mfast k k ),

slow where mfast −1 = m−1 = 0 and the superscript χ denotes the slow momentum coefficient. We have the following conclusion.

Proposition E.1 (Slow-momentum-induced norm inflation). Suppose ρ := 1 − ηλ > 0 and both processes start from the same deterministic w0 . Then for every k ≥ 0, the iterates in (51) satisfy χ 0 E∥wk+1 ∥2 − E∥wk+1 ∥2 = qη 2 σ 2

k X

(52)

2χAfast,n Aslow,n + χ2 A2slow,n ,

n=0

where βfast = 1 − αfast , βslow = 1 − αslow and A#,n :=

n X

ℓ , ρ n−ℓ β#

# ∈ {fast, slow}.

ℓ=0

Moreover, the weight norm converges to lim E∥wkχ ∥2 = Vχ := dη 2 σ 2

k→∞

∞ X

A2fast,n + qη 2 σ 2

n=0

∞ X

2χAfast,n Aslow,n + χ2 A2slow,n ,

(53)

n=0

inflated norm by slow momentum

where Vχ < ∞ if and only if ρ < 1 (i.e., λ > 0).

}

Remark E.2. The norm inflation due to slow momentum has two contributing factors. The first is the explicit positive coefficient χ > 0, and the second is the asymptotic dominance Aslow,n ≫ Afast,n for large n.

Proof. Unrolling the momentum recursions gives m# k =

k X

β#k−j gj ,

# ∈ {fast, slow}.

j=0

Substituting it into the update rule yields χ wk+1 = ρk+1 w0 − η

k X j=0

(Afast,k−j Id + χAslow,k−j Q) gj .

i.i.d.

Since gk ∼ N (0, σ 2 Id ), χ E∥wk+1 ∥2 = ρ2(k+1) ∥w0 ∥2 + η 2 σ 2

k X j=0

∥Afast,k−j Id + χAslow,k−j Q∥F .

We obtain (52) by 2 2 ∥Afast,n Id + χAslow,n Q∥F − ∥Afast,n Id ∥F = q 2χAfast,n Aslow,n + χ2 A2s,n . Moreover,

 n+1 n+1  − β# ρ , ρ ̸= β# , A#,n = ρ − β#  (n + 1)ρn , ρ = β# ,

P∞ gives n=0 A2#,n < ∞ iff λ > 0. Therefore, the stationary mean-square parameter norms are finite iff λ > 0, and taking the trace of the stationary covariance matrix gives Vχ in (53).

Proof

Define ρfast = 1 − αfast , ρslow = 1 − αslow , ηmax,k = maxi<k ηi . For sharp/flat subspace indicator q ∈ {0, 1}, write dq = αρfast + χq αρslow , a0 = atop and a1 = afull . We take e = 2.71828... as the base of the 2 2 fast slow natural logarithm, unless explicitly defined otherwise. Notations.

We first introduce explicit constant coefficients in Assumption 7.1 to facilitate cleaner mathematical derivations. Assumption F.1 (Detailed Version of Assumption 7.1). There exist constants ℓ− , ℓ+ , ϑ− , ϑ+ > 0, a capacity exponent ν > 1 and a source exponent s > 0 such that, for every 1 ≤ j ≤ d, ℓ− j −ν ≤ λj ≤ ℓ+ j −ν ,

ϑ− j −1 λs−1 ≤ |wj⋆ |2 ≤ ϑ+ j −1 λs−1 . j j

Lemma F.2 (Hessian-aligned anisotropic noise structure). For every u ∈ Rd , the gradient noise covariance is Σ(u) = Huu⊤ H + (u⊤ Hu + σ 2 )H,

(54)

(2E(u) + σ 2 )H ⪯ Σ(u) ⪯ (4E(u) + σ 2 )H.

(55)

and it satisfies Proof. This lemma is Lemma A.1 in [38], and we restate its proof here for completeness. For one sample, the stochastic gradient satisfies g = ϕϕ⊤ u − ϕϵ, Eg = Hu.

Wick’s formula for the centered Gaussian vector ϕ gives

E[ϕϕ⊤ uu⊤ ϕϕ⊤ ] = 2Huu⊤ H + (u⊤ Hu)H. Further computation based on it proves (54). Finally, we have Huu⊤ H ⪯ (u⊤ Hu)H by Cauchy–Schwarz to the H-inner product. Since u⊤ Hu = 2E(u), this proves (55). Remark F.3. Lemma F.2 shows the alignment between the noise and the Hessian. In the early phase of training, E is large and the noise is dominated by the loss scale. In the later phase, where E ≲ σ 2 , the noise aligns well with the Hessian, and the stochastic gradient is noise-dominated. The following lemma gives the aggregate parameter-state update rule. Lemma F.4 (Discrete parameter-momentum recursion). For the j-th eigen-direction 1 ≤ j ≤ d, define slow ⊤ Xj,k := (uj,k , mfast j,k−1 , mj,k−1 ) .

Then the concrete algorithm (8) is equivalent with Xj,k+1 = Mj,k Xj,k + dj,k ξj,k ,

(56)

where

 1 − ηk cj λj λj Mj,k =  λj

 −ηk χqj ρslow , 0 ρslow

−ηk ρfast ρfast 0



 −ηk cj dj,k =  1  , 1

(57)

cj := γ + 1 + χqj , and ξj,k ∈ R denotes the shared noise in all three entries. Proof. Substituting mfast = ρfast mfast k k−1 + λj uj,k + ξj,k ,

mslow = ρslow mslow k k−1 + λj uj,k + ξj,k

into (8) gives the results. We next analyze the roots of the characteristic polynomial det(zI − Mj,k ) of the transition matrix Mj,k , which determine the final convergence rate. Consider any fixed coordinate j, and hide the subscript j for brevity. Define the constants 1 − αslow 1 − αfast −1 −1 +χ e , a := γ + αfast +χ eαslow , d := 2 2 αfast αslow where χ e ≥ 0 denotes the corresponding coefficient for the slow momentum. The matrix with zero learning rate has eigenvalues 1, ρfast = 1 − αfast , ρslow = 1 − αslow . The next lemma shows that for a sufficiently small positive learning rate, the root issuing from 1 remains separated from the two roots issuing from ρfast and ρslow .

Lemma F.5 (Spectrum distribution of the transition matrix). Let Mχe,λ,η be the matrix obtained from (57) with λj = λ, χqj = χ e and ηk = η. Define x = ηλ. If aηλ ≤

αslow , 8

(58)

then we have: • Mχe,λ,η has a real simple eigenvalue z0 in 1 − αslow 4 , 1 with a 2 log z0 + ax + a d + x ≤ clog x3 , 2

(59)

where clog = 4a(d + a)2 + 6a2 e. • The other two eigenvalues, counted with algebraic multiplicity, satisfy max

z∈σ(Mχ,λ,η )\{z0 } e

|z| ≤ 1 −

3αslow . 4

(60)

Proof. Expanding det(zI − Mχe,λ,η ) along its first row gives P (z) = (z − ρfast )(z − ρslow )(z − 1 + ηcλ) + ηλρfast (z − ρslow ) + η χ eλρslow (z − ρfast ),

(61)

fast slow where c = γ + 1 + χ e. Define the constants e = 2 1−α + 2e χ 1−α , and clog = 4a(d + a)2 + 6a2 e. α3 α3 fast

slow

• Bound the largest real root. Note that F (z) > 0 for z ≥ 1. Thus all real roots are less than 1. Let δ := 1 − z1 . The characteristic equation gives ρfast ρslow δ =x c+ +χ e . (62) αfast − δ αslow − δ | {z } :=Fx (δ)

On the interval I = [0, αslow /2], we have 0 ≤ Fx (δ) ≤ 2ax ≤ αslow /4. Thus Fx (I) ⊂ I. On the other hand, for δ ∈ I, ρfast ρslow 1 ′ Fx (δ) = x +χ e ≤ 4xd ≤ . 2 2 (αfast − δ) (αslow − δ) 2 Since Fx is a contraction mapping on I, it admits a unique real fixed point δx ∈ I, with 0 ≤ δx = Fx (δx ) ≤ 2ax ≤

αslow , 4

(63)

lim δx = 0.

x→0

Combining it with P (−∞) < 0, P (ρslow ) > 0, P has at least 2 real roots. Thus P has 3 real roots. Next we consider finer estimate to the largest root. Expanding the two (1 − δ)−1 terms in (62) via their geometric series (with remainder) in powers of δ gives ρfast ρslow 2e e +χ e ≤ a + dδ + δ 2 a + dδ + δ 2 ≤ c + 2 αfast − δ αslow − δ 3 by using δ ≤ 2ax ≤ αslow /4. Hence 8 0 ≤ δ − ax ≤ 2adx2 + a2 ex3 . 3

(64)

Further applying dx ≤ ax/αslow ≤ 1/8 gives 0 ≤ δ − ax − adx2 ≤ 2ad2 + 3a2 e x3 . Combining them with

n≥3 δ

/n ≤ δ 3 /[3(1 − δ)] gives |δ 2 − a2 x2 | X δ n + 2 n n≥3 1 ≤ 2ad2 + 3a2 e x3 + 3a2 d + a2 e x3 + 4a3 x3 . 2

| log(1 − δ) + ax + a(d + a/2)x2 | ≤|δ − ax − adx2 | +

(65)

Their sum is at most clog x3 , proving (59). • Bound the remaining two roots. Define 

0 0



1 T = −1 −1

αfast /λ 0

−xbfast ρfast + xbfast xbfast

 −xbslow , xbslow ρslow + xbslow

. αslow /λ

Then M is similar to  1 − ax Nq (x) := T M T −1 =  ax ax

bfast :=

ρfast , αfast

bslow := χ e

ρslow . αslow

(66) Thus ∥Nq (x) − diag(1, ρfast , ρslow )∥∞ ≤ 2ax ≤ αslow /4. The Gershgorin disk associated with the first row lies in Re z ≥ 1 − 2ax ≥ 1 − αslow /4. Each of the other two disks lies in |z| ≤ ρslow + 2ax ≤ 1 − 3αslow /4. The first disk is therefore disjoint from the union of the other two. By the Gershgorin disk theorem (Theorem F.13), the first component contains exactly one eigenvalue, while the complementary component contains exactly two eigenvalues counted with algebraic multiplicity, proving (60).

Lemma F.6 (Properties of spectral projections). We use the notation Nq (x) in (66), where q ∈ {0, 1} is the sharp subspace indicator. For x ≥ 0 satisfying aq x ≤ αslow /8, let Γ0 := {z ∈ C : |z − 1| = αslow /2}, Γ1 := {z ∈ C : |z| = 1 − αslow /2}, Z Z 1 1 P0 (x) := (zI − Nq (x))−1 dz, P1 (x) := (zI − Nq (x))−1 dz. 2πi Γ0 2πi Γ1

(67)

Both contours are positively oriented. The following properties hold. (i) The operators in (67) are complementary oblique subspace projections with P0 (x) + P1 (x) = I and P0 (x)P1 (x) = P1 (x)P0 (x) = 0. P0 (x) is associated with the largest eigenvalue and P1 (x) is associated with the complementary two-dimensional subspace3 . If x, y ≥ 0 and aq max{x, y} ≤ αslow /8, then ∥P0 (x) − P0 (y)∥∞ ≤

16aq |x − y|, αslow

(68)

max{∥P1 (y)P0 (x)∥∞ , ∥P0 (y)P1 (x)∥∞ } ≤

32aq |x − y|. αslow

(69)

∥P0 (x)∥∞ ≤ 2, and

(ii) At x = 0, one has P0 (0) = diag(1, 0, 0). Besides, the projector P0 is differentiable at x = 0 with   2 2 0 −ρfast /αfast −χqρslow /αslow , 0 0 P0′ (0) =  aq /αfast aq /αslow 0 0

(70)

and its quadratic remainder satisfies ∥P0 (x) − P0 (0) − xP0′ (0)∥∞ ≤

32a2q x2 . 2 αslow

(71)

(iii) Let δ = 1 − z0,q (x). The range of P0 (x) is spanned by the eigenvector rq (x) =

αfast − δ αslow − δ

⊤ with ∥rq (x)∥∞ = 1.

The complementary projection P1 satisfies for any v ∈ Range P1 (x), 17αslow ∥Nq (x)v∥∞ ≤ 1 − ∥v∥∞ . 20

(72)

(73)

(iv) Let a ≤ b, xm ≥ 0, and suppose aq xm ≤ αslow /8 for a ≤ m ≤ b. Define Lm := ∥P0 (xm ) − P0 (xm−1 )∥∞ .

With the ordered-product notation

← Q−b−−−

m=a Am = Ab · · · Aa , we have

←−− b Y Nq (xm )P1 (xm ) v m=a

≤e

m=a+1 Lm

b−a+1 17αslow ∥P1 (xa )v∥∞ . 1− 20

(74)

∞

Proof. Proof of (i). Lemma F.5 shows that Γ0 encloses exactly the largest root z0,q (x) (i.e, the largest eigenvalue of Nq (x)). The disk bounded by Γ1 contains the other two eigenvalues, counted with algebraic multiplicity, and excludes z0,q (x). Hence P0 is the (rank 1) oblique projection to the eigenspace corresponding 3 In general, P , P are not orthogonal to each other. 0 1

to z0,q (x), and P1 = I − P0 . On each contour, the Neumann resolvent identity and ∥Nq (x) − Nq (0)∥∞ ≤ 2aq x give 4 ∥(zI − Nq (x))−1 ∥∞ ≤ . αslow The Riesz formula and the second resolvent identity now give (68). Complementarity gives P1 (y)P0 (x) = [P0 (x) − P0 (y)]P0 (x),

P0 (y)P1 (x) = P0 (y)[P0 (y) − P0 (x)],

so (69) follows from (68). Proof of (ii). Write Nq (x) = Nq (0) + xNq′ (0) (which is linear). Differentiating the Riesz formula at x = 0 gives   2 2 Z 0 −ρfast /αfast −χqρslow /αslow 1 , 0 0 P0′ (0) = R0 (z)Nq′ (0)R0 (z) dz =  aq /αfast (75) 2πi Γ0 aq /αslow 0 0 where Rx (z) := (zI − Nq (x))−1 and we use R0 (z) = diag((z − 1)−1 , (z − ρfast )−1 , (z − ρslow )−1 ). Applying the second-order resolvent identity Rx = R0 + xR0 Nq′ (0)R0 + x2 R0 Nq′ (0)Rx Nq′ (0)R0 , we have

1 2πi

x2 = 2πi

P0 (x) − P0 (0) − xP0′ (0) =

Γ0

Rx (z) − R0 (z) − xR0 (z)Nq′ (0)R0 (z)x dz R0 (z)Nq′ (0)Rx (z)Nq′ (0)R0 (z)

(76)

Γ0

4 2 , ∥Rx ∥∞ ≤ αslow , ∥Nq′ (0)∥∞ ≤ 2aq , we get Using the bounds ∥R0 ∥∞ ≤ αslow

x2 |Γ0 |∥R0 ∥2∞ ∥Nq′ (0)∥2∞ ∥Rx ∥∞ 2π 2 32a2q x2 x2 2 4 ≤ (παslow ) . (2aq )2 ≤ 2 2π αslow αslow αslow

∥P0 (x) − P0 (0) − xP0′ (0)∥∞ ≤

(77)

Proof of (iii). Direct computation verifies that Nq (x)rq (x) = z0,q (x)rq (x). Since 0 ≤ δ ≤ αslow /4 < αfast , the last two coordinates of rq (x) belong to [0, 1/3], proving (72). ρfast χqρslow For the complementary subspace, set θq (x) := x αfast (α + αslow (αslow −δ) . Then by computation, an fast −δ) eigenvector for z0,q (x) with respect to Nq (x)⊤ is ⊤ xρfast xχqρslow ℓq (x) = 1, − ,− . αfast (αfast − δ)(1 − θq (x)) αslow (αslow − δ)(1 − θq (x)) The bounds δ ≤ αslow /4, dq ≤ aq /αslow , and aq x ≤ αslow /8 imply 0 ≤ θq (x) ≤

4 1 xdq ≤ . 3 6

Because Range P1 (x) = ker ℓq (x)⊤ , every vector in the fast range satisfies |v1 | ≤

θq (x) 1 max{|v2 |, |v3 |} ≤ max{|v2 |, |v3 |}. 1 − θq (x) 5

The last two coordinates of Nq (x)v are therefore bounded by 6 17αslow ∥v∥∞ . 1 − αslow + aq x ∥v∥∞ ≤ 1 − 5 20 31

This proves (73). Proof of (iv). If v ∈ Range P1 (xm−1 ), then ∥P1 (xm )v∥∞ ≤ (1 + Lm )∥v∥∞ . Iteration with (73) proves (74). Before proceeding with the proof, we recall some previous notations and introduce a few new ones. ( atop , j ≤ r, aj := A := atop (Id − Q) + afull Q, afull , j > r, SA (t) :=

d X j=1

λj |wj⋆ |2 e−2aj λj t ,

KA (t) :=

d X

(78)

a2j λ2j e−2aj λj t .

j=1

Lemma F.7 (Uniform kernel bounds). Suppose that Assumption F.1 and d ≥ 2(r + 1) hold. Then: (i) There are positive constants cS,− , cS,+ , cK,− , cK,+ , given in the proof and depending only on the power-law constants, s, ν, and Cd , such that, for 0 ≤ afull t ≤ Cd dν , min{1, cS,− }S(t) ≤ SA (t) ≤ max{1, cS,+ }S(t),

(79)

min{1, cK,− }K(t) ≤ KA (t) ≤ max{1, cK,+ }K(t). (ii) The kernel KA satisfies

Z ∞ KA (t) dt = 0

1 tr(AH). 2

(80)

(iii) If (r + 1)ν ≤ afull t ≤ Cd dν , then S(t) ≍

r X j=1

1/ν

λj |wj⋆ |2 e−2atop λj t + (afull t)−s ,

K(t) ≍ afull t−p .

(81)

(iv) There is a constant Ch depending only on the power-law constants, s, and ν, such that, if r ≥ 4 and atop t ≥ Ch rν (1 + log κ), then r X λj |wj⋆ |2 e−2atop λj t ≲ (afull t)−s , (82) j=1

where κ := afull /atop ≥ 1. Proof. Throughout the proof, we abbreviate βS := 1 + sν, βK := 2ν, R := r + 1. Proof of (i). For β > 1, define ν+1 1 L1 (β) := min 4−β e−2ℓ+ , 41−β e−2 ℓ+ Cd , 2 and set (83)

cS,− := ϑ− ℓs− L1 (βS ), ( cS,+ := ϑ+ ℓs+ 2s max 1 + cK,− := ℓ2− L1 (βK ), cK,+ := ℓ2+ 2p max 1 +

1 Γ(s) , +2 sν ν(2ℓ− )s

−1 ) −1 s+ν

s+ν 2eℓ−

(84) (85)

−2

1 Γ(p) 8e , + p 2ν − 1 ν(2ℓ− ) (2ℓ− )2

(86)

R∞ where Γ(a) := 0 y a−1 e−y dy denotes the Gamma function for a > 0. The following proof uses that c− A ≤ D ≤ c+ A implies min{1, c− }(B + A) ≤ B + D ≤ max{1, c+ }(B + A) for positive A, B, D. 32

• Lower bounds. Set u = afull t, and for any β > 1, Fβ,d (u) :=

d X

j −β e−2ℓ+ uj

−ν

δβ :=

j=R

We prove that the tail estimate

β−1 . ν

(87)

Fβ,d (u) ≥ L1 (β)(Rν + u)−δβ

holds for any 0 ≤ u ≤ Cd d . ν

– Case 1: u ≤ Rν . We have Fβ,d (u) ≥

2R−1 X

j −β e−2ℓ+ uj

−ν

≥ 2−β e−2ℓ+ R1−β ≥ 4−β e−2ℓ+ (Rν + u)−δβ .

j=R

– Case 2: u > Rν . Define J := u1/ν > R. If J ≤ d/4,

Fβ,d (u) ≥

If J > d/4, u ≤ Cd d gives uj ν

Fβ,d (u) ≥

−ν

j −β e−2ℓ+ uj

−ν

J≤j≤4J

≥ 4−β e−2ℓ+ J 1−β .

≤ u(d/2)−ν ≤ Cd dν (d/2)−ν = 2ν Cd . It follows that

j −β e−2ℓ+ uj

−ν

d 2 ≤j≤d

≥

ν+1 d −β −2ν+1 ℓ+ Cd 1 d e ≥ 41−β e−2 ℓ+ Cd J 1−β , 2 2

where the last inequality uses d < 4J. In both cases, (Rν + u)−δβ ≤ u−δβ = J 1−β , proving (87). For the signal lower bound, since λj |wj⋆ |2 e−2afull λj t ≥ ϑ− ℓs− j −βS e−2ℓ+ uj d X j=R

−ν

λj |wj⋆ |2 e−2afull λj t ≥ ϑ− ℓs− FβS ,d (u) ≥ ϑ− ℓs− L1 (βS )(Rν + u)−(βS −1)/ν ≥ cS,− (Rν + u)−s .

For the kernel lower bound, we use a2full λ2j e−2afull λj t ≥ a2full ℓ2− j −2ν e−2ℓ+ uj d X j=R

−ν

and get

a2full λ2j e−2afull λj t ≥ a2full ℓ2− FβK ,d (u) ≥ a2full ℓ2− L1 (βK )(Rν + u)−(βK −1)/ν ≥ cK,− a2full (Rν + u)−p .

• Upper bounds. For the upper bounds, it suffices to control the sum

Gβ (u) :=

∞ X

j −β e−2ℓ− uj

−ν

β > 1.

1 1+ β−1

j=R

– Case 1: u ≤ Rν . We have Gβ (u) ≤

∞ X

−β

j=R

≤R

1−β

– Case 2: u > Rν . Define

fβ,u (x) := x−β e−2ℓ− ux

−ν

We use the elementary inequality X j≥1

f (j) ≤

Z ∞ f (x) dx + 2 sup f (x). x>0

−ν For the integral term, set δβ = β−1 . This yields ν and make the change of variables y = 2ℓ− ux Z ∞ Z ∞ Z −ν (2ℓ− u)−δβ ∞ δβ −1 −y Γ(δβ ) f (x) dx = x−β e−2ℓ− ux dx = y e dy = . δβ ν ν(2ℓ − u) 0 0 0 d To bound the supremum term, we compute the derivative dx log fβ,u (x) = − βx +2ℓ− uνx−ν−1 . Hence β/ν β/ν the maximum is attained at the point satisfying 2ℓ− uνx−ν = β. Thus supx>0 fβ,u (x) = 2eℓ . −u

Since u > Rν ≥ 1 and β/ν = δβ + ν −1 > δβ , it follows that " β/ν # Γ(δβ ) β/ν Gβ (u) ≤ +2 u−δβ . 2eℓ− ν(2ℓ− )δβ In either case, Rν + u ≤ 2 max{Rν , u}. Consequently, ( β/ν ) Γ(δ ) β/ν 1 β , Gβ (u) ≤ 2δβ max 1 + +2 (Rν + u)−δβ . β − 1 ν(2ℓ− )δβ 2eℓ−

(88)

Substituting the value of β into it yields the upper bounds in (79). Proof of (ii). Direct computation gives Z ∞ KA (t) dt = 0

d X

a2j λ2j

j=1

d 1 1X 1 = aj λj = tr(AH). 2aj λj 2 j=1 2

Proof of (iii). Put u = afull t. Since y 2 e−2y ≤ e−2 for y ≥ 0, and r ≤ u1/ν , we have r X

a2top λ2j e−2atop λj t = a2full u−2

j=1

r X (κ−1 λj u)2 e−2bλj u ≤ e−2 a2full u−2 r ≤ e−2 a2full u−p . j=1

This proves the forgetting kernel part in (81). The signal rate part is obtained by an analogous argument, where the sharp subspace summation is retained. Proof of (iv). Define the constants c = 2ℓ− , x = atop t, y = rxν for notational brevity. We suppose y ≥ 1 and will verify this condition later. For β > 1, set the constant   X mν c 1 2 top Cβ = 2β max 1, + 2m(β−1) e− 2 (2 −1)  . (89) cν ec(2ν − 1) m≥1

We first consider the sum on r/2 ≤ j ≤ r. With n = r − j, Bernoulli’s inequality gives −β 1 + νn ≤ 2β r−β , we have r . Since j

X r/2≤j≤r

j −β e−cy(r/j) ≤ 2β r−β e−cy

X n≥0

ν r j

= 1 − nr

1 e−cνyn/r ≤ 2β max 1, r1−β (y −1 + r−1 )e−cy , cν 34

−ν

≥

(90)

where the last inequality uses (1 − e−z )−1 ≤ 1 + z −1 for z > 0. For the remaining summation, we define n Bm := j : On Bm , j −β ≤ Consequently, X j∈Bm

r o < j ≤ , 2m+1 2m

m+1 β 2 , r

m ≥ 1.

ν r ≥ 2mν . j i h mν = 2β r1−β y −1 e−cy 2m(β−1) ye−c(2 −1)y .

mν

j −β e−cy(r/j) ≤ 2β r1−β 2m(β−1) e−cy2

(91)

For am := c(2mν − 1), one has ye−am y = ye−am y/2 e−am y/2 ≤ Thus X j≤r/2

j −β e−cy(r/j) ≤ 2β

ec(2ν − 1)

2 −am /2 2 e ≤ e−am /2 . eam ec(2ν − 1)

r1−β y −1 e−cy

2m(β−1) e− 2 (2

mν

−1)

m≥1

Combining it with (90) gives r X

j −β e−cxj

j=1

−ν

≤ Cβtop r1−β (y −1 + r−1 )e−cy .

(92)

Applying (92) with β = βS = 1 + sν gives r X j=1

λj |wj⋆ |2 e−2atop λj t ≤ ϑ+ ℓs+

r X

j −βS e−cxj

−ν

j=1

≤ CStop r−sν (y −1 + r−1 )e−2ℓ− y ,

(93)

where the constant CStop := ϑ+ ℓs+ Cβtop depends only on the power-law constants, s, and ν. S Finally, set

s Ch := max 1, ℓ−

Then y ≥ Ch (1 + log κ) ≥ 1. Since κx = afull t, division by (κx)−s = κ−s r−sν y −s gives Pr ⋆ 2 −2atop λj t j=1 λj |wj | e ≤ CStop y s−1 + r−1 y s e−ℓ− y ≤ CStop sup(z s−1 + z s )e−ℓ− z < ∞, −s (afull t) z≥1 where the right hand side is a finite constant. This proves (82).

We now consider iterating the parameter-momentum update as specified in Lemma F.4. For the j-th coordinate and integers k > i, define Φj;k,i := Mj,k−1 Mj,k−2 · · · Mj,i , Φj;i,i := I3 , and

(94)

Rj;k,i := e⊤ 1 Φj;k,i+1 dj,i ,

where e1 = (1, 0, 0) . Let ⊤

1 esignal := k

d X

2 j=1

2 λ j e⊤ 1 Φj;k,0 Xj,0 ,

noise Kk,i :=

d X j=1

2 λ2j Rj;k,i .

(95)

Lemma F.8 (Two-sided convolution-type bounds on E). For the iterates wk generated by (8), X X σ2 σ2 signal signal noise noise ek + Kk,i EE(wi ) + ≤ EE(wk ) ≤ ek + Kk,i 2EE(wi ) + . 2 2 i<k

(96)

i<k

Proof. Iterating (56) yields uj,k = e⊤ 1 Φj;k,0 Xj,0 +

Rj;k,i ξj,i .

i<k

Since the noise is independent and has zero mean, it follows that d

EE(wk ) =

1X 1 XX 2 λj u2j,k = esignal λj Rj;k,i + E[e⊤ j Σ(ui )ej ]. k 2 j=1 2 j=1 i<k

Applying Lemma F.2 proves (96). noise Next we replace esignal , Kk,i with quantities related to SA , KA , which have closed forms defined in (78). k The following lemma characterizes the approximation error of this replacement. We first introduce some definitions for notational brevity. Let

Vk (η) :=

k−2 X i=0

|ηi+1 − ηi |,

δstep (k) :=

max

0≤i≤k−2

|ηi+1 − ηi | , αslow ηi

δstep (1) := 0.

Define δstep (k) afull λ1 ηmax,k + Vk (η) + , 2 αslow αslow X X afull εexp (k) := afull d1 + λ21 ηi2 + clog,1 λ31 ηi3 , 2

(97)

εproj (k) :=

i<k

(98)

i<k

(99)

εep (k) := 64εproj (k) + 8εexp (k), and for q ∈ {0, 1}, Gq :=

ρ2fast 2χqρfast ρslow χ2 q 2 ρ2slow + + 2 (1 − ρ2 ) 2 αfast αfast αslow (1 − ρfast ρslow ) αslow (1 − ρ2slow ) fast

21 .

(100)

For any nonnegative sequence {si }i , define Lj,k := ηk−1 +

k−2 X

k−1 Y

(1 − aqj λj ηm ) ηt+1 − ηt + aqj λj ηt ηt+1 ,

t=0 m=t+2

b δ(k) := 4

d X

λ2j G2qj L2j,k ,

(101)

b b k := δ(k) (Bs) max si . 0≤i<k

j=1

Lemma F.9 (Approximation errors). For any nonnegative sequence {si }, define the operator T that maps s to a sequence {T s}i as X (T s)k := ηi2 KA (Tk − Ti )si , (T s)0 = 0. (102) i<k

Suppose afull λ1 ηmax,k ≤

αslow , 8

εep (k) ≤

1 . 16

(103)

Then

1 esignal − SA (Tk ) ≤ 4εep (k)SA (Tk ), k 2

(104)

and k−1 X i=0

noise Kk,i si − (T s)k ≤ 2εep (k) + εep (k)2 (T s)k + 2 1 + εep (k)

q b k + (Bs) b k. (T s)k (Bs)

(105)

Further using Cauchy–Schwarz Inequality to (105) gives k−1 X 225 289 noise b k≤ b k. (T s)k − (Bs) (T s)k + 2(Bs) Kk,i si ≤ 512 128 i=0

(106)

Proof. Fix a coordinate j with eigenvalue λ := λj and sharp subspace indicator q := qj ∈ {0, 1}. Set xm := ηm λ, and consider the projectors from Lemma F.6. The first condition in (103) ensures aq xm ≤ αslow /8 at every step. Suppress q from the projector notation and define Lm := ∥P0 (xm ) − P0 (xm−1 )∥∞ ,

m ≥ 1.

The estimate in (68) gives k−1 X m=1

Lm ≤

16aq λ 1 . Vk (η) ≤ 16εproj (k) ≤ αslow 64

(107)

We considner the matrix Nq (x) in Lemma F.6, where q ∈ {0, 1} is the sharp subspace indicator. To be specific, let   1 0 0 0 . Tλ := −1 αfast /λ −1 0 αslow /λ Then Tλj Mj,ℓ Tλ−1 = Nq (xℓ,j ) and Tλ e1 = e1 , where xℓ,j := ηℓ λj . When no ambiguity arises, we omit the j ← Q−k−− subscript j for brevity. We define the ordered matrix product ℓ=1 Ai := Ak Ak−1 ...A1 .

For the deterministic part X, since Xj,0 = (−wj⋆ , 0, 0)⊤ , the T transformed coordinates is Tλ Xj,0 = −wj⋆ h for h = (1, −1, −1)⊤ . Define s Dj,k := −wj⋆ e⊤ 1

←−− k−1 Y

Nq (xm )P0 (xm ) h,

m=0

U j,k := −wj⋆

k−1 Y

z0,q (xm ),

m=0

0 Uj,k := −wj⋆ e−aj λj Tk .

Next we bound each of relevant approximation errors to prove the corresponding conclusions. s 0 Part 1: approximation errors related to Dj,k − Uj,k . fast slow + χq 1−α . Set Recall that cq = γ + 1 + χq, dq = 1−α α2 α2 fast

p0 := P0 (x0 )h,

slow

pm := P0 (xm )pm−1 ,

1 ≤ m < k.

′ Note that e⊤ 1 P0 (0)h = dq by (70) and P0 (0)h = e1 . Hence using (71), dq ≤ aq /αslow and aq x0 ≤ αslow /8, we obtain ⊤ ′ ⊤ ′ |e⊤ 1 p0 − 1| =|x0 e1 P0 (0)h + e1 P0 (x0 ) − P0 (0) − x0 P0 (0) h| (108) 32a2q x20 ≤ dq x0 + 2 ≤ 5εproj (k). αslow

Since pm ∈ Range P0 (xm ), (72) gives ∥pm ∥∞ = |e⊤ 1 pm |. Moreover, pm − pm−1 = [P0 (xm ) − P0 (xm−1 )]pm−1 , 37

∥pm ∥∞ ≤ (1 + Lm )∥pm−1 ∥∞ .

(109)

Thus

Y (109) k−1

max ∥pm ∥∞ ≤ m<k

(108)

(1 + Lm )∥p0 ∥∞ ≤ e

m=1 Lm

m=1

X (109) 17 k−1

∥pk−1 − p0 ∥∞ ≤

Pk−1

16 m=1

(107) 17

(1 + 5εproj (k)) ≤

, (110)

(107)

Lm ≤ 17εproj (k).

s Since Nq (x)P0 (x) = z0,q (x)P0 (x), we have Dj,k = −wj⋆ (110) imply

Qk−1

⊤ ⊤ m=0 z0,q (xm )e1 pk−1 = e1 pk−1 U j,k . Then (108) and

s ⊤ |Dj,k − U j,k | = |U j,k | |e⊤ 1 pk−1 − 1| ≤ |U j,k | ∥pk−1 − p0 ∥∞ + |e1 p0 − 1| ≤ 22εproj (k)|U j,k |.

(111)

By (64), z0,q (x) ≤ 1 − aq x ≤ e−aq x . Combining it with (59) gives k−1 X

k−1 k−1 X aq 2 X 2 3 log z0,q (xm ) + aq λTk ≤ aq dq + ηm + clog,q λ3 ηm λ ≤ εexp (k), 2 m=0 m=0 m=0

where we also use aq ≤ afull , dq ≤ d1 , and clog,q ≤ clog,1 . It follows that (112)

0 0 |U j,k − Uj,k | ≤ εexp (k)|Uj,k | 0 Combining this estimate with (111) and |U j,k | ≤ |Uj,k | yields

0 s 0 0 |Dj,k − Uj,k | ≤ 22εproj (k) + εexp (k) |Uj,k | ≤ εep (k)|Uj,k |.

(113)

s Part 2: approximation errors related to e⊤ 1 Φj;k,0 Xj,0 − Dj,k .

Set

y−1 := h,

ym := Nq (xm )ym−1 ,

and for m ≥ 1 consider the decomposition

0 ≤ m < k,

sm := Nq (xm )P0 (xm )ym−1 ,

fm := Nq (xm )P1 (xm )ym−1 ,

s0m := Nq (xm )P0 (xm )s0m−1 ,

am := sm − s0m

s . For m ≥ 1, we have Note that ym = sm + fm , and {s0m }m generates Dj,k

am = Nq (xm )P0 (xm )(am−1 + fm−1 ), fm = Nq (xm )P1 (xm )(s0m−1 + am−1 + fm−1 ). For any u ∈ Range P0 (xm−1 ), v ∈ Range P1 (xm−1 ), we have P0 (xm−1 )v = 0 and P0 (xm )u = u + (P0 (xm ) − P0 (xm−1 ))u, P0 (xm )v = (P0 (xm ) − P0 (xm−1 ))v,

Thus

∥P0 (xm )u∥∞ ≤ (1 + Lm )∥u∥∞ ,

Qm :=

∥s0m ∥∞ Pm , exp(−aq ℓ=0 xℓ )

Am :=

P1 (xm )v = (I − P0 (xm ))v.

∥P1 (xm )u∥∞ ≤ Lm ∥u∥∞ ,

∥P0 (xm )v∥∞ ≤ Lm ∥v∥∞ ,

Set

P1 (xm )u = (P0 (xm−1 ) − P0 (xm ))u,

∥P1 (xm )v∥∞ ≤ (1 + Lm )∥v∥∞ .

∥am ∥∞ Pm , exp(−aq ℓ=0 xℓ )

Fm :=

∥fm ∥∞ Pm . exp(−aq ℓ=0 xℓ )

(114)

(115)

−aq xm The inequality aq xm ≤ αslow /8 < 17αslow /20 implies 1 − 17 . Hence combining (115) and (73) 20 αslow ≤ e yields Qm ≤ (1 + Lm )Qm−1 ,

(116)

Am ≤ (1 + Lm )Am−1 + Lm Fm−1 ,

Fm ≤ Lm (Qm−1 + Am−1 ) + (1 + Lm )Fm−1 .

At m = 0, the identities s00 = z0,q (x0 )P0 (x0 )h, f0 = Nq (x0 )P1 (x0 )h together with z0,q (x0 ) ≤ e−aq x0 , (73), and −aq x0 1 − 17 give 20 αslow ≤ e Q0 ≤ ∥P0 (x0 )h∥∞ , F0 ≤ ∥P1 (x0 )h∥∞ . Using ∥h∥∞ = 1, (70) and (71), we get ∥P0 (x0 )h∥∞ = |e⊤ 1 P0 (x0 )h|

⊤ ′ ⊤ ′ ≤ |x0 e⊤ 1 P0 (0)h| + |x0 e1 P0 (0)h| + |e1 (P0 (x0 ) − P0 (0) − x0 P0 (0))h|

32a2q x20 65 ≤ 1 + 5εproj (k) ≤ , ≤ 1 + d q x0 + 2 αslow 64

(117)

⊤ where the third inequality also uses dq ≤ aq /αslow and aq x0 ≤ αslow 8 . On the other hand, P1 (0)h = (0, −1, −1) implies 16aq x0 65 ∥P1 (x0 )h∥∞ ≤ 1 + ∥P0 (x0 ) − P0 (0)∥∞ ≤ 1 + . ≤ αslow 64

Consequently max{Q0 , F0 } ≤ 65/64, and hence Qm ≤ 17/16. Define Cm := Am + Fm . Then (116) yields Cm ≤ (1 + 2Lm )Cm−1 + Lm Qm−1 . Using S0,k ≤ 1/64 and iterating gives 2S0,k

max Cm ≤ e m<k

65 17 + S0,k 64 16

≤

7 . 6

For k ≥ 2, iteration of the second line of (116) gives 7 Ak−1 ≤ eS0,k max Cm ≤ eS0,k S0,k ≤ 19εproj (k). m<k 6

(118)

It also holds for k = 1 as A0 = 0. Finally, e⊤ 1 P1 (0) = 0, so ∥e⊤ 1 P1 (xk−1 )∥1 ≤ ∥P0 (xk−1 ) − P0 (0)∥∞ ≤ Note that and

16aq xk−1 . αslow

(119)

s ⋆ ⊤ 0 ⋆ ⊤ e⊤ 1 Φj;k,0 Xj,0 − Dj,k = −wj e1 (yk−1 − sk−1 ) = −wj e1 (ak−1 + fk−1 ),

|e⊤ 1 ak−1 | ≤ ∥ak−1 ∥∞ ,

(119) 16a x q k−1

⊤ ⊤ |e⊤ 1 fk−1 | = |e1 P1 (xk−1 )fk−1 | ≤ ∥e1 P1 (xk−1 )∥1 ∥fk−1 ∥∞ ≤

Thus Pk−1 s ⋆ i=0 xi |e⊤ 1 Φj;k,0 Xj,0 − Dj,k | ≤ |wj |e (118)

αslow

16aq xk−1 Ak−1 + Fk−1 αslow

0 0 ≤ 38εproj (k)|Uj,k | ≤ εep (k)|Uj,k |,

∥fk−1 ∥∞ .

(120)

2 where the second inequality also uses aq xk−1 /αslow ≤ εproj (k). Combining this estimate with (113) proves

(121)

0 0 |e⊤ 1 Φj;k,0 Xj,0 − Uj,k | ≤ 2εep (k)|Uj,k |.

0 2 Pd 2 0 2 0 2 ≤ 4εep (k) 1 + εep (k) |Uj,k | = SA (Tk ), and εep (k) ≤ Thus |e⊤ | . Since j=1 λj |Uj,k 1 Φj;k,0 Xj,0 | − |Uj,k |

1/16, we have

1 esignal − SA (Tk ) ≤ 2εep (k) 1 + εep (k) SA (Tk ) ≤ 4εep (k)SA (Tk ). k 2 This proves (104). Part 3: approximation errors related to the noise decay. Here we focus on a fixed coordinate (i.e., fix λ, q), and omit the subscript for notational brevity. For 0 ≤ i ≤ k − 1, define (k)

0k := (0, . . . , 0)⊤ ∈ Rk ,

:= 0, . . . , 0, 1, | {z } i entries

u0 = 0k ,

⊤

0, . . . , 0 | {z }

∈ Rk ,

k−i−1 entries

(122)

(k) gi = λui + ei

mf−1 = ms−1 = 0k ,

If [x]i denotes coordinate i of x ∈ Rk , then the preceding definition implies [gt ]i = λ[ut ]i + I{i=t} for 0 ≤ i, t ≤ k − 1. We further let mft = ρfast mft−1 + gt , mst = ρslow mst−1 + gt ,

(123)

ut+1 − ut = −ηt γgt + mft + χqmst Then ut = Rj;t,0 , . . . , Rj;t,t−1 , 0, . . . , 0 | {z }

⊤

t−1 X

0 ≤ t ≤ k − 1.

i=0

k−t entries

and gt = λRj;t,0 , . . . , λRj;t,t−1 , 1,

(k)

Rj;t,i ei ,

0, . . . , 0 | {z }

⊤

0 ≤ t ≤ k,

(124)

0 ≤ t ≤ k − 1.

k−t−1 entries

We also define gr = 0k for r < 0. Since mft = 0 ≤ t ≤ k − 1, ut+1 − ut = −ηt

γgt +

t X

s ℓ ℓ=0 ρfast gt−ℓ , mt

ρℓfast gt−ℓ + χq

ℓ=0

t X

ℓ ℓ=0 ρslow gt−ℓ , we get that for

! ρℓslow gt−ℓ

ℓ=0

= −ηt

t X

Cℓ (q)gt−ℓ ,

(125)

ℓ=0

where Cℓ (q) := γI{ℓ=0} + ρℓfast + χqρℓslow for every integer ℓ ≥ 0.

where the series is absolutely convergent because 0 ≤ ρfast , ρslow < 1. For a Rk -valued sequence b = (bt )t≥0 , P∞ ρn ρn define the convolution with Dn (q) := ℓ=n Cℓ (q) = αfast + χq αslow (n ≥ 1) by fast slow (Dq ∗ b)t :=

 t+1 X   D (q)b n

 n=1 0,

t−n+1 ,

t ≥ 0,

(126)

t < 0.

It follows that for every t ≥ 0, t t X X (Dq ∗ b)t − (Dq ∗ b)t−1 = D1 (q)bt + Dℓ+1 (q) − Dℓ (q) bt−ℓ = aq bt − Cℓ (q)bt−ℓ . ℓ=1

ℓ=0

(127)

Applying (127) with b = g turns (125) into ut+1 − ut = −aq ηt gt + ηt (Dq ∗ g)t − (Dq ∗ g)t−1 ,

0 ≤ t ≤ k − 1.

(128)

Thus the effective gain aq is separated exactly from a discrete boundary difference; no spectral approximation has been made. Define the sequence uE 0 = 0, (k)

E uE t+1 = (1 − aq ληt )ut − aq ηt et ,

It satisfies uE t = −aq Define

t−1 X

t−1 Y

ηi

i=0

(k)

(1 − aq ληℓ )ei ,

0 ≤ t ≤ k.

ℓ=i+1

(k)

gtE := λuE t + et ,

(129)

0 ≤ t ≤ k − 1.

0 ≤ t ≤ k − 1,

vt := ut − uE t ,

0 ≤ t ≤ k,

and set gkE := 0. Let Xk := (Rk ){0,...,k} . Define the map Bq : Xk → Xk as (Bq b)0 := 0, (Bq b)t+1 := (1 − aq ληt )(Bq b)t + ηt (Dq ∗ b)t − (Dq ∗ b)t−1 , 0 ≤ t ≤ k − 1.

(130)

Since gt − gtE = λvt for 0 ≤ t ≤ k − 1, subtracting (129) from (128) gives that for 0 ≤ t ≤ k − 1, vt+1 = (1 − aq ληt )vt + ηt (Dq ∗ gE )t − (Dq ∗ gE )t−1 + ληt (Dq ∗ v)t − (Dq ∗ v)t−1 . Thus v ∈ Xk satisfies

(131)

v = Bq gE + λBq v.

Thus we consider the decomposition vt = We set v0 = 0.

(r) (r+1) := λBq v(r) (1 ≤ r < k) and v(1) := Bq gE . r=1 vt , where v

We first bound Bq for a general input. Set γm,t := ηt (Bq b)m =

m−1 X

Qm−1

ℓ=t+1 (1 − aq ληℓ ) for 0 ≤ t < m ≤ k. Then (130) gives

γm,t (Dq ∗ b)t − (Dq ∗ b)t−1 .

t=0

Since (Dq ∗ b)−1 = 0, summation by parts yields (Bq b)m = ηm−1 (Dq ∗ b)m−1 + Then by

m−2 X t=0

(132)

(γm,t − γm,t+1 )(Dq ∗ b)t .

P∞

n=1 Dn (q) = dq , we get

∥(Bq b)m ∥ ≤ (ηm−1 +

m−2 X t=0

|γm,t − γm,t+1 |)dq

sup 0≤t≤m−1

(133)

∥bt ∥.

Substituting m−2 X t=0

|γm,t+1 − γm,t | =

m−2 X

m−1 Y

t=0

ℓ=t+2

(1 − aq ληℓ ) (ηt+1 − ηt ) + aq ληt ηt+1

≤Vk (η) + ηmax,k

m−2 X m−1 Y

(

(1 − aq ληℓ ) −

t=0 ℓ=t+2

≤Vk (η) + ηmax,k . 41

m−1 Y

(1 − aq ληℓ )

ℓ=t+2

(1 − aq ληℓ ))

ℓ=t+1

(134)

into (133) yields

sup ∥(Bq b)m ∥ ≤ 2 ηmax,k + Vk (η) dq

0≤m≤k

sup 0≤t≤k−1

∥bt ∥.

(135)

Set yi := aq ληi and note that G2q =

∞ X

Dn (q)2 =

n=1

2χqρfast ρslow χ2 q 2 ρ2slow ρ2fast + + ≤ G21 . 2 (1 − ρ2 ) 2 αfast αfast αslow (1 − ρfast ρslow ) αslow (1 − ρ2slow ) fast

Direct computation gives ∞ X ρn

fast

αfast

≤

implying that dq =

P∞

n=1

2 αslow

∞ X ρ2n

!1/2

fast 2 α n=1 fast

n=1 Dn (q) ≤

∞ X χqρn

slow

≤

αslow

n=1

2 αslow

∞ X χ2 q 2 ρ2n

!1/2

slow

n=1

2 αslow

2G √ q . By (129), we have that for 1 ≤ t ≤ k, αslow

λuE t =−

t−1 X

t−1 Y

i=0

(k)

(1 − yℓ )ei .

ℓ=i+1

Then 0 ≤ yi ≤ αslow /8 gives that for 1 ≤ t ≤ k ∥λuE t ∥=

t−1 X

yi2

i=0

t−1 Y

(1 − yℓ )2 ≤

ℓ=i+1

max yi

0≤i≤t−1

t−1 Y

1−

For 0 ≤ t ≤ k − 1, E

(Dq ∗ g )t =

max yi

0≤i≤t−1

# (1 − yℓ )

ℓ=0

t+1 X

X t−1 i=0

Y t−1 1 − (1 − yi )2 (1 − yℓ )2 ℓ=i+1

(136)

αslow . ≤ 8 t+1 X

(k) Dn (q)et−n+1 +

n=1

Dn (q)λuE t−n+1 .

n=1

Hence for any 0 ≤ t ≤ k − 1, ∥(Dq ∗ gE )t ∥ ≤

t+1 X

(k)

Dn (q)et−n+1 +

n=1

t+1 X

!1/2 Dn (q)2

n=1

t+1 X n=1

Dn (q)∥λuE t−n+1 ∥

(137)

≤ Gq + dq sup ∥λuE s ∥ 0≤s≤t

1 1+ √ Gq . 2 Pm−2 Since (v(1) )m = (Bq gE )m = ηm−1 (Dq ∗ gE )m−1 + t=0 (γm,t − γm,t+1 )(Dq ∗ gE )t , combining (132), (134), and (137) yields ! m−2 X √ (1) ∥vm ∥ ≤ ηm−1 + |γm,t − γm,t+1 | sup ∥(Dq ∗ gE )t ∥ ≤ (2 + 2) ηmax,k + Vk (η) Gq . (138) (136)

≤

0≤t<m

t=0

1 Set θk := 2λ ηmax,k + Vk (η) dq . Then θk ≤ 2αslow εproj (k) ≤ 2εproj (k) ≤ 512 by dq ≤ aq /αslow and aq λ ≤ afull λ1 . (135) gives that for 1 ≤ r < k, (r+1) (r) sup ∥vm ∥ ≤ θk sup ∥vm ∥.

0≤m≤k

Combining this with (138) yields sup ∥vm ∥ ≤

0≤m≤k

≤

k X

(r) sup ∥vm ∥ ≤ (2 +

√

r=1 0≤m≤k

√

k X 2) ηmax,k + Vk (η) Gq θkr−1 r=1

2 ηmax,k + Vk (η) Gq .

2+ 1 − θk

Substituting it and (133) into (131) yields a finer estimate: ∥vk ∥ ≤

k−2 X 1 (1 + √ )Gq + λdq sup ∥vm ∥ (ηk−1 + |γk,t − γk,t+1 |) 2 0≤m≤k t=0

≤ 2Gq (ηk−1 +

k−2 X t=0

|γk,t − γk,t+1 |).

Qk−1 0 E Define Rj;k,i := −ηi aq e−aq λ(Tk −Ti ) and Rj;k,i := −aq ηi ℓ=i+1 (1 − aq ληℓ ) for 0 ≤ i ≤ k − 1. Note that (124) gives k−1 X (k) E vk = Rj;k,i − Rj;k,i ei . i=0

Then

k−1 X i=0

E (Rj;k,i − Rj;k,i )2 = ∥vk ∥2 ≤ 4(ηk−1 +

k−2 X t=0

|γk,t − γk,t+1 |)2 G21 .

(139)

E 0 Next we bound |Rj;k,i − Rj;k,i |. Note that 2 yi ≤ aq ληmax,k ≤ αslow εproj (k) ≤ εproj (k),

Pk−1 and ℓ=i+1 yℓ2 ≤ 2εexp (k) as a2q λ2 ≤ a2full λ21 ≤ 2afull (d1 + afull /2)λ21 . Combining them with 0 ≤ − log(1 − y) − y ≤ 47 y 2 for 0 ≤ y ≤ 1/8, we get log

k−1 E X |Rj;k,i | 8 1 log(1 − yℓ ) + yℓ ≤ εproj (k) + εexp (k) ≤ εep (k). = y + i 0 |Rj;k,i | 7 7 ℓ=i+1

Using |ex − 1| ≤ 2|x| for |x| ≤ 1/2, we obtain (140)

E 0 0 |Rj;k,i − Rj;k,i | ≤ εep (k)|Rj;k,i |.

Taking summation of (139), (140) over the coordinate j gives k−1 d XX i=0 j=1

and

E b b k, λ2j (Rj;k,i − Rj;k,i )2 si ≤ δ(k) max si = (Bs)

k−1 d XX i=0 j=1

0≤i<k

(141)

(142)

E 0 λ2j (Rj;k,i − Rj;k,i )2 si ≤ εep (k)2 (T s)k ,

respectively. Note that k−1 d XX i=0 j=1

0 λ2j (Rj;k,i )2 si = (T s)k ,

k−1 X i=0

noise Kk,i si =

k−1 d XX i=0 j=1

λ2j (Rj;k,i )2 si .

(143)

Thus (T s)k −

k−1 X

2 noise Kk,i si

≤

i=0

k−1 d XX i=0 j=1

0 0 λ2j (Rj;k,i + Rj;k,i )|Rj;k,i − Rj;k,i |si

v uk−1 d uX X ≤t λ2 (Rj;k,i + R0 j

j;k,i

i=0 j=1

v uk−1 d uX X 2 ) si t λ2 (Rj;k,i − R0

2 j;k,i ) si

(144)

i=0 j=1

q q p b k εep (k) (T s)k + (Bs) b k (T s)k + (Bs) q b k + (Bs) b k, = 2εep (k) + εep (k)2 (T s)k + 2(1 + εep (k)) (T s)k (Bs) ≤ (2 + εep (k))

where the second inequality uses Cauchy–Schwarz Inequality, the third inequality uses v v v uk−1 d uk−1 d uk−1 d uX X uX X uX X 0 E E 0 t )2 si ≤t )2 si + t − Rj;k,i )2 si λ2j (Rj;k,i − Rj;k,i λ2j (Rj;k,i − Rj;k,i λ2j (Rj;k,i i=0 j=1

i=0 j=1

(145)

q (141),(142) p b k + εep (k) (T s)k , ≤ (Bs) and v v uk−1 d uk−1 d uX X uX X 2 0 0 t 2 2 λj (Rj;k,i − Rj;k,i ) si + 2t λ2j (Rj;k,i )2 si j;k,i ) si ≤

v uk−1 d uX X t λ2 (Rj;k,i + R0 j

i=0 j=1

(145),(143)

≤

i=0 j=1

(146)

p b k + (2 + εep (k)) (T s)k . (Bs)

This proves (105) as εep (k) ≤ 1/16. We define the following notations for proving Theorem 7.2: F(k) := S(k) + σ 2

k−1 X i=0

ηi2 K(Tk − Ti ),

csk,− := min{1, cS,− , cK,− }, b δbmax (k) := max δ(n), 1≤n<k

(147)

csk,+ := max{1, cS,+ , cK,+ }, Kηmax (k) := max

1≤n<k

n−1 X i=0

ηi2 KA (Tn − Ti ),

εb(k) :=

b σ 2 δ(k) , F(k)

where we assume EE(w0 ) ≤ C0 σ with C0 treated as a constant. At k = 1, both maxima over 1 ≤ n < k in (147) are defined to be zero, and Kηmax (1) := 0. 2

Assumption F.10 (Admissible hyper-parameters). Suppose that the horizon k and the learning rate schedule {ηi } satisfy εep (k) ≤

1 , 16

Kηmax (k) ≤

1 , 8

1 δbmax (k) ≤ , 25

εb(k) ≤

csk,− , 16(6C0 + 1)

afull λ1 ηmax,k ≤

αslow . 8 (148)

Remark F.11 (Simplified sufficient conditions for admissible hyperparameters). Here we provide a cleaner version of the hyperparameter constraints in (148). The constraints on εep (k) and δbmax (k) reduce to Pk−2 upper bounds on ηmax,k , i=1 |ηi+1 − ηi |, and maxi<k−1 |ηi+1 /ηi − 1|. By (152), a sufficient condition for P −1 Pd r Kηmax (k) ≤ 18 is ηmax,k ≤ 14 a0 j=1 j −ν + a1 j=r+1 j −ν , which also becomes an upper bound constraint to ηmax,k . 44

We next turn to the upper bound constraint on εb(k). Deriving a simplified sufficient condition for general learning rate schedules is challenging. However, for some common choices such as constant and powerν law decay, it can be satisfied. Excluding the very early stage, we consider the regime Tk ≳ artop so that K(Tk − s) ≲ K(Tk − Ti ) for s ∈ [Ti , Ti+1 ]. Hence, Z Tk K(z) dz = 0

This yields

k−1 X Z Ti+1 Ti

i=0

K(Tk − s) ds ≲

k−1 X i=0

ηi K(Tk − Ti ) ≤

k−1 X

ηk−1 i=0

ηi2 K(Tk − Ti ).

(149)

Pk−1

2 ηi = η0 (1 + i/τ0 )−γ , where τ0 ∈ R+ ∪ {+∞} and i=0 ηi K(Tk − Ti ) ≳ ηk−1. Now consider γ t+1 = 1 + τ01+t − 1 ≤ τ0γ+t gives 1/2 < γ ≤ 1. The inequality ηtη−η t+1

#2 k−2 k−1 X ηt − ηt+1 2 k−2 X p 2γ 2 X 2 2 (ηt − ηt+1 ) K(Tk − Tt+2 ) ≲ η K(Tk − Ti ). (150) ηt+1 K(Tk − Tt+1 ) ≤ ηt+1 τ0 i=0 i t=0 t=0 t=0

"k−2 X

Combining this with

we obtain

2 ηk−1 K(0) ≲ ηk−1 , Pk−1 2 i=0 ηi K(Tk − Ti )

i2 h p p Pk−2 ηk−1 K(0) + t=0 (ηt − ηt+1 ) K(Tk − Tt+2 ) k−1 X i=0

≲ ηk−1 +

ηi2 K(Tk − Ti )

γ2 . τ0

Note that Vk (η) ≤ η0 , δstep (k) ≤ αslow τ0 . By Lemma F.12, for sufficiently large d, τ0 and sufficiently small η0 , the admissibility condition (148) holds. γ

Proof of Theorem 7.2. Suppose the conditions in (148) hold. We structure the proof as follows. Step 1: the uniform upper bound for E. 2

+ σ σ For 0 ≤ i ≤ k − 1, we define the sequence Ei := EE(wi ), s− i := Ei + 2 , si := 2Ei + 2 , and we denote by 1 the sequence consisting entirely of ones.

Since KA is decreasing, (T 1)n =

n−1 X i=0

ηi2 KA (Tn − Ti ) ≤ ηmax,k Z Tn

= ηmax,k 0

(80) η

KA (u) du ≤

n−1 X Z Ti+1 i=0

max,k

KA (Tn − s) ds

(151)

tr(AH).

Taking the maximum over 0 ≤ n < k gives max (T 1)n ≤ Kηmax (k) ≤

0≤n<k

ηmax,k tr(AH). 2

(152)

The upper bounds in (96) and (106) give that for 1 ≤ n ≤ k, En ≤esignal + n

n−1 X i=0

noise + Kn,i si ≤ esignal + n

289 b + )n (T s+ )n + 2(Bs 128

3 289 289 2 2 b b ≤ C0 σ 2 + (T E)n + σ (T 1)n + 4δ(n) max Ei + δ(n)σ . 0≤i≤n−1 2 64 256

(153)

2 2 b b + )n = δ(n) where the last inequality uses (T s+ )n = 2(T E)n + σ2 (T 1)n , (Bs 2 max0≤i≤n−1 Ei + σ2 and (104)

esignal ≤ n

1 3 3 + 4εep (n) SA (Tn ) ≤ SA (0) ≤ C0 σ 2 . 2 4 2

(154)

Then using max0≤n<k (T E)n ≤ Kηmax (k) max0≤i≤k−1 Ei and (152) to (153), we get max Ei ≤

0≤i≤k−1

3 289 max σ2 C0 σ 2 + Kη (k) + 4δbmax (k) max Ei + . 0≤i≤k−1 2 64 4

(155)

max (k) + 4δbmax (k) ≤ 43 in (148), we get With the condition 289 64 Kη

max Ei ≤

0≤i<k

6C0 +

3 4

σ2 .

(156)

Step 2: the asymptotic bounds for Ek .

2 + 5 2 The estimate in (156) gives max0≤i<k s− i ≤ 6C0 + 4 σ , max0≤i<k si ≤ 2 (6C0 + 1) σ .

Thus

k−1 X

σ2 σ2 (T 1)k , ηi2 KA (Tk − Ti ) Ei + ≥ 2 2 i=0 k−1 X σ2 + 2 (T s )k = ηi KA (Tk − Ti ) 2Ei + ≤ 2 (6C0 + 1) σ 2 (T 1)k , 2 i=0 5 − − b b b σ 2 δ(k), (Bs )k = δ(k) max si ≤ 6C0 + 0≤i<k 4 b b b + )k = δ(k) (Bs max s+ ≤ 2 (6C0 + 1) σ 2 δ(k)

(T s− )k =

0≤i<k

(157)

Besides, (104) gives 1 SA (Tk ) ≤ 4

1 1 3 − 4εep (k) SA (Tk ) ≤ esignal ≤ + 4ε (k) SA (Tk ) ≤ SA (Tk ). ep k 2 2 4

Define Gk := SA (Tk ) + σ 2 (T 1)k = SA (Tk ) + σ 2

Pk−1

2 i=0 ηi KA (Tk − Ti ). The estimates in (79) give

Gk ≥ min{1, cS,− }S(Tk ) + σ 2 min{1, cK,− } 2

k−1 X i=0

Gk ≤ max{1, cS,+ }S(Tk ) + σ max{1, cK,+ } we obtain (96)

k−1 X

(106) 1

ηi2 K(Tk − Ti ) ≥ csk,− F(k),

k−1 X i=0

noise − Kk,i si ≥

(159)

ηi2 K(Tk − Ti ) ≤ csk,+ F(k).

225 b − )k (T s− )k − (Bs 512 i=0 (157) 1 225 2 5 b SA (Tk ) + σ (T 1)k − 6C0 + σ 2 δ(k) ≥ 4 1024 4 (148) 225 24C0 + 5 966C0 + 145 ≥ − csk,− Gk = csk,− Gk . 1024 64(6C0 + 1) 1024(6C0 + 1)

Ek ≥ esignal + k

(158)

SA (Tk ) +

(160)

and (96)

k−1 X

(106) 3

289 b + )k (T s+ )k + 2(Bs 4 128 i=0 (161) (157) 3 (148) 289 c 289 sk,− b ≤ SA (Tk ) + (6C0 + 1)σ 2 (T 1)k + 4(6C0 + 1)σ 2 δ(k) (6C0 + 1)csk,+ + Gk . ≤ 4 64 64 4

Ek ≤ esignal + k

noise + Kk,i si ≤

SA (Tk ) +

Combining (160), (161) and (159), we get Ek ≍ F(k).

This proves the conclusion. Proof of Corollary 7.4. Set tk := kη, R := r + 1.

Part 1: estimate the noise accumulation term. Enlarge the constant Ch supplied by Lemma F.7 (iv) so that ν n o 1 4 2ℓ− Ch ≥ log 2, Ch ≥ max 1, 21/(1− ν ) − 1 . (162) 5 Ti = iη gives j ≤ r,

Pk−1

2 2 i=0 η K(Tk − Ti ) = η

m=1 K(mη). Decompose K(z) =

Kjtop (z) := a2top λ2j e−2atop λj z ,

top tail (z), where for j=1 Kj (z) + K

K tail (z) := a2full (Rν + afull z)−p .

Let r i atop X afull h 1−ν ν −(1−ν −1 ) JK (t) := K(z) dz = λj 1 − e−2atop λj t + R − (R + a t) . full 2 j=1 1 − ν −1 0

Z t

For 1 ≤ m ≤ k,

Kjtop (mη) Kjtop ((m − 1)η)

(163)

= e−2atop λj η ≥ e−2atop λ1 η ,

(164) −p afull η afull η −p K tail (mη) = 1+ ν ≥ 1+ . K tail ((m − 1)η) R + afull (m − 1)η Rν n o η −p Set the constant cm := min e−2atop λ1 η , 1 + afull . Then (164) implies that for z ∈ [(m − 1)η, mη], Rν K(mη) ≤ K(z) ≤ K((m − 1)η) ≤ c−1 m K(mη). Thus Z mη Z mη cm K(z) dz ≤ ηK(mη) ≤ (m−1)η

It follows that cm ηJK (tk ) ≤ η 2

K(z) dz.

(m−1)η

k X m=1

K(mη) ≤ ηJK (tk ).

(165)

By the hyper-parameter constraints in Theorem 7.2, we have afull η afull λ1 η αslow 1 1 = ≤ ≤ ≤ . ν ν ν ν R λ1 R 8λ1 R 8ℓ− R 8ℓ− −p 1 −1/4 Consequently, cm ≥ cm,0 := min e , 1 + 8ℓ− ≳ 1. 2atop λ1 η ≤ 2afull λ1 η ≤

αslow 1 ≤ , 4 4

(166)

Since d ≥ 2R, we have d X j=R

λj ≥

Then defining ctr = min we get

2R−1 X j=R

1 1 2 , ℓ+

d X

ℓ− j −ν ≥ ℓ− 2−ν R1−ν , , Ctr = max

j=R

1 2ν 2 , (1−ν −1 )ℓ−

λj ≤ ℓ+

∞ X j=R

j −ν ≤

ℓ+ 1−ν R . q

and using tr(AH) = atop

j=1 λj +afull

(167) Pd

j=R λj ,

ctr tr(AH) ≤=

atop X afull R1−ν ≤ Ctr tr(AH). λj + 2 j=1 1 − ν −1 | {z }

(168)

=JK (∞)

The constraint (11) implies that atop λj tk ≥ ℓ− Ch that for every j ≤ r,

ν r j

(1 + log κ) ≥ ℓ− Ch . Combining this with (162) implies

1 − e−2atop λj tk ≥ 1 − e−2ℓ− Ch ≥

1 . 2

(169)

ν 4 . 5

(170)

Moreover, κ ≥ 1 and r/R = r/(r + 1) ≥ 4/5 give r ν afull tk ≥ κC (1 + log κ) ≥ Ch h Rν R This implies that R1−ν − (Rν + afull tk )−(1−ν R1−ν

−1

)

afull tk =1− 1+ Rν

−(1−ν −1 )

ν −(1−ν −1 ) 1 4 ≥ 1 − 1 + Ch ≥ . 5 2

(171)

Applying (169) and (171) to the closed form of JK (tk ) yields 1 JK (∞) ≤ JK (tk ) ≤ JK (∞). 2

(172)

k X cm,0 ctr 2 η tr(AH) ≤ η K(mη) ≤ Ctr η tr(AH). 2 m=1

(173)

Combining (165), (168), and (172) gives

Part 2: estimate the signal term. By (82), there exists a fixed Ctop > 0 such that the lower inequality in (11) implies r X λj |wj⋆ |2 e−2atop λj tk ≤ Ctop (afull tk )−s . (174) j=1

(162) and (170) imply afull tk ≥ Rν , and therefore (175)

2−s (afull tk )−s ≤ (Rν + afull tk )−s ≤ (afull tk )−s . Using (174) and (175) to S(t) =

⋆ 2 −2atop λj t + j=1 λj |wj | e

(r + 1)ν + afull t

−s

2−s (afull ηk)−s ≤ S(ηk) ≤ (Ctop + 1)(afull ηk)−s . Finally, combining (173) and (176) gives (12).

, we have (176)

Lemma F.12 (Admissibility for a decreasing schedule). Suppose Assumption F.1 and d ≥ 2(r + 1) hold. Let k ≥ 2, and suppose afull Tk ≤ Cd dν ,

η0 ≥ η1 ≥ · · · ≥ ηk−1 > 0,

afull λ1 η0 ≤

αslow . 8

Then δbmax (k) ≤ 16η02

d X

λ2j G2qj ,

Kηmax (k) ≤

j=1

η0 tr(AH). 2

(177)

If we further assume max

G20 G21 , a2top a2full

" ηk−1

K(0) +

k−2 X t=0

#2 p (ηt − ηt+1 ) K(Tk − Tt+2 )

k−1 X csk,− ≤ η 2 K(Tk − Ti ), 256(6C0 + 1) max{1, cK,+ } i=0 i

it holds that εb(k) ≤

(178)

csk,− . 16(6C0 + 1)

Proof. For every 1 ≤ j ≤ d, 0 ≤ t ≤ k − 1, and 1 ≤ ℓ ≤ k, we have 0 ≤ aqj λj ηt ≤ afull λ1 η0 ≤

αslow < 1, 8

0 ≤ Πj;k,ℓ ≤ 1.

(179)

Using Πj;k,t+2 − Πj;k,t+1 = aqj λj ηt+1 Πj;k,t+2 , we get k−2 X

aqj λj ηt ηt+1 Πj;k,t+2 = ηk−1 +

t=0

k−2 X t=0

Πj;k,t+2 (ηt − ηt+1 ) − η0 Πj;k,1 .

(180)

Combining this with Πj;k,t+2 ≤ e−aqj λj (Tk −Tt+2 ) yields "

k−2 X

η0 Lj,k ≤ 2 ηk−1 + Πj;k,t+2 (ηt − ηt+1 ) − Πj;k,1 2 t=0 # " k−2 X −aqj λj (Tk −Tt+2 ) ≤ 2 ηk−1 + (ηt − ηt+1 )e .

# (181)

t=0

For 0 ≤ u ≤ Tk , (79) implies d X j=1

λ2j G2qj e−2aqj λj u ≤ max

G20 G21 , a2top a2full

max{1, cK,+ }K(u).

Combing it with (181) gives b δ(k) ≤ 16 max

G20 G21 , a2top a2full

#2 k−2 X p p (ηt − ηt+1 ) K(Tk − Tt+2 ) . max{1, cK,+ } ηk−1 K(0) + "

(182)

t=0

Pd Note that Lj,n ≤ 2η0 by (181) for n ≤ k. Therefore δbmax (k) ≤ 16η02 j=1 λ2j G2qj , and then we prove (177) by P k−1 using (152). Finally, using F(k) ≥ σ 2 i=0 ηi2 K(Tk − Ti ), (182) and (178) proves the conclusion.

F.1

Useful Lemmas

Theorem F.13 (Gershgorin disk theorem). For M = (mij ) ∈ Cn×n , define the row disks     X Di := z ∈ C : |z − mii | ≤ |mij | .   j̸=i

Every eigenvalue of M belongs to i=1 Di . More precisely, if a connected component of this union consists of exactly k disks and is disjoint from all remaining disks, then it contains exactly k eigenvalues of M , counted with algebraic multiplicity. Theorem F.14 (Neumann series and matrix resolvent identities). Let E ∈ Cn×n , and the norm ∥ · ∥ ∈ {∥ · ∥2 , ∥ · ∥∞ }. If ∥E∥ < 1, then (I − E)−1 =

∞ X

∥(I − E)−1 ∥ ≤

En,

n=0

1 . 1 − ∥E∥

(183)

For A ∈ Cn×n , write RA (z) := (zI − A)−1 whenever zI − A is invertible. If ∥RA (z)E∥ < 1, then RA+E (z) = (I − RA (z)E)−1 RA (z),

∥RA+E (z)∥ ≤

∥RA (z)∥ . 1 − ∥RA (z)E∥

(184)

Whenever zI − A and wI − A are invertible, RA (z) − RA (w) = (w − z)RA (z)RA (w).

(185)

For A, B ∈ Cn×n , whenever both zI − A and zI − B are invertible, RA (z) − RB (z) = RA (z)(A − B)RB (z).

(186)

Theorem F.15 (Riesz spectral projection formula). Let A ∈ Cn×n , and Γ be a positively oriented, smooth, simple closed contour with no eigenvalue of A on the contour. The matrix Z 1 PΓ (A) := (zI − A)−1 dz (187) 2πi Γ is the projection (i.e., PΓ (A)2 = PΓ (A)) onto the generalized eigenspaces corresponding to the eigenvalues inside Γ. It commutes with A, and for every integer n ≥ 0, Z 1 An PΓ (A) = z n (zI − A)−1 dz. (188) 2πi Γ If no eigenvalue of B ∈ Cn×n lies on Γ, then 1 PΓ (A) − PΓ (B) = 2πi

Z Γ

RA (z)(A − B)RB (z) dz.

(189)

Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining

Analyze & explore related research in Conceptio

Related documents