Understanding Practical Optimization in Machine Learning

Egor Shulgin

cs.LGOct 26, 2026

Understanding Practical Optimization in Machine Learning

Authors: Egor Shulgin

Organizations: King Abdullah University of Science and Technology Thuwal, Kingdom of Saudi Arabia

Abstract

Modern machine learning relies on practical optimization methods that often differ from the algorithms studied in classical analyses: they train submodels, clip or normalize gradients, use layer-wise non-Euclidean directions, or approximate costly inner computations. This thesis analyzes the objective, assumptions, update, and computation defined by these procedures, rather than those of simplified proxies. For submodel training, we show that Independent Subnetwork Training can converge to a shifted fixed point because sketching before differentiation biases the update for the original objective. Model-Agnostic Sparsified Training instead incorporates the sketch distribution into its objective, which makes sketched gradients unbiased and admits convergence guarantees. For private optimization, we show that adaptive quantile clipping retains bias at any fixed quantile, while suitable quantile schedules recover near-optimal nonconvex rates under heavy-tailed noise. We then combine smoothed normalization with error compensation and obtain distributed private convergence guarantees under standard assumptions, without assuming bounded gradients. For optimizers based on linear minimization oracles, we develop and validate a layer-wise generalized smoothness model and build on it the Gluon framework for Muon- and Scion-style methods. We also show that the approximate orthogonalization in Muon matters: less precise orthogonalization calls for smaller stepsizes, so its precision should be tuned together with the learning rate. Overall, practical methods are better understood when theory retains the features that determine their optimization behavior.

Figures & tables

Figure 1.1: Representative trajectory-smoothness measurements from Chapter 6 . Left: the fitted layer-wise

(L_{i}^{0},L_{i}^{1})

model tracks observed trajectory smoothness better than a constant-smoothness proxy. Right: trajectory smoothness varies substantially across layers and matrix groups, motivating layer-wise constants rather than a single global smoothness parameter.

Ch.	Object	Modeled feature	Main message
2	IST / submodel-gradient training	Objective	Sketching before differentiation biases the update, and the method can converge to a shifted fixed point.
3	MAST / sketch-aware objective	Objective	Placing the sketch distribution in the objective makes sketched gradients unbiased for the induced target.
4	Adaptive quantile clipping	Update	Fixed quantiles retain clipping bias; a time-varying quantile schedule removes it and recovers optimal rates.
5	Smoothed normalization / $\alpha$ -NormEC	Update	Smoothed normalization with error compensation gives private distributed guarantees without bounded gradients.
6	Gluon / layer-wise generalized smoothness	Assumptions	Layer-wise generalized smoothness matches training trajectories and yields guarantees with per-layer constants.
7	Inexact Muon / LMO updates	Update	Orthogonalization precision changes the convergence bound and the best stepsize.

Table 1.1: Thesis map: main objects, modeled features, and messages.

Panel of Figure 2 — Figure 2.1: Performance of simplified IST on a quadratic problem for varying step size values.

Figure 3.1: Distributions of test accuracies of sparsified solutions for the ERM formulation ( 3.1 ) and MAST problem ( 3.2 ). “Sparsity” corresponds to the percentage of zeroed weights.

Panel of Figure 3 — Figure 3.2: Performance of Algorithm 3 with Bernoulli sketches ( 3.6 ) on standard loss ( 3.18 ) (for ${\bf S}_{i}\equiv\mathbf{I}$ ).

Figure 4.1: Adaptive clipping thresholds for five target quantiles on six federated learning tasks without privacy noise. The trajectories vary substantially across tasks and may increase, decrease, or change nonmonotonically. The curves were redrawn for this thesis from the vector trajectories in Figure 4 of Andrew et al. (2021) ; the original clipping norms were estimated using geometric updates with

\eta_{C}=0.2

and

C_{0}=0.1

Property	DP-SGD	Clip21	$\alpha$ -NormEC ( Ours )
Non-private Convergence	✗	✓	✓
Private Convergence	✗	✗	✓
Operators	Clipping / Normalization	Clipping	Normalization
Additional Assumptions	Bounded gradient or heterogeneity	No	No

Table 5.1: Comparison of distributed private optimization methods. Unlike prior work, which either lacks private convergence guarantees or requires restrictive assumptions like bounded gradients,

\alpha

-NormEC is the first to provide these guarantees under standard smoothness conditions alone.

Operator	Property
Contractive compressor $\mathcal{C}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$	$\left\\|\mathcal{C}(g)-g\right\\|^{2}\leq(1-\eta)^{2}\left\\|g\right\\|^{2}$
Clipping ${\rm Clip}_{\tau}\left(g\right):=\min\left(1,\frac{\tau}{\left\\|g\right\\|}\right)g$	$\left\\|{\rm Clip}_{\tau}\left(g\right)-g\right\\|^{2}\leq\max(0,\left\\|g\right\\|-\tau)^{2}$
Smoothed normalization ${\rm Norm}_{\alpha}\left(g\right):=\frac{1}{\alpha+\left\\|g\right\\|}g$	$\left\\|{\rm Norm}_{\alpha}\left(g\right)-g\right\\|^{2}\leq\left(1-\frac{1}{\alpha+\left\\|g\right\\|}\right)^{2}\left\\|g\right\\|^{2}$

Table 5.2: Comparison of contractive compressors, clipping, and smoothed normalization with their properties. Smoothed normalization, unlike clipping, satisfies a contractive-like property similar to that of compressors.

Figure 5.1: The highest test accuracy achieved by $\alpha$ -NormEC with different $\alpha$ and $\beta$ values.

Figure 5.4: Gradient norm

\|\nabla f(x^{k})\|

across methods on a logistic regression problem. Curves are averaged over

5

seeds.

M_{i}^{k}=\beta^{k}M^{k-1}_{i}+(1-\beta^{k})\nabla_{i}f_{\xi^{k}}(X^{k})

Algorithm 5 Gluon: Stochastic Adaptive Layer-Wise LMO-based Optimizer with Momentum

Panel of Figure 6 — Figure 6.1: Training NanoGPT on FineWeb validates our layer-wise $(L^{0},L^{1})$ -smoothness model.

Figure 6.2: Validation of 6.2.1 in NanoGPT-124M along training trajectories of unScion . The estimates

\hat{L}_{i}[k]

are closely matched by

\hat{L}_{i}^{\text{approx}}[k]

, while a constant

L_{0}

(corresponding to the classical smoothness assumption), results in a substantially poorer fit.

Figure 6.4: Validation of layer-wise

(L^{0},L^{1})

-smoothness for different groups of parameters in NanoGPT-124M along AdamW training trajectories.

\hat{d}^{k}\approx\underset{\|d\|\leq 1}{\operatorname{argmin}}\ \left\langle m^{k+1},d\right\rangle

Algorithm 6 Inexact Generalized Muon

Panel of Figure 7 — Figure 7.1: NanoGPT on FineWeb: convergence (left) and hyperparameter sensitivity (right) as PolarExpress (PolarExp) iterations vary.

Appendix figures & tables40 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 21 — Figure C.1: Distributions of accuracies of sparsified solutions for the ERM ( 3.1 ) and MAST ( 3.2 ) formulations.

Panel of Figure 22 — Figure C.2: Test accuracies of sparsified solutions for the ERM formulation ( 3.1 ) and MAST problem ( 3.2 ).

Panel of Figure 23 — Figure C.3: Finite-sum MAST loss ( 3.17 ) convergence for Algorithm 2 (II) with subsampling.

Figure C.4: MAST loss ( 3.2 ) convergence for Algorithm 2 (I) with varying step size and Rand-

1

sketches.

Panel of Figure 25 — Figure C.5: Performance of Algorithm 3 with Bernoulli sketches ( 3.6 ) on standard loss ( 3.18 ) (for ${\bf S}_{i}\equiv\mathbf{I}$ ).

Panel of Figure 26 — Figure E.1: Minimal train loss (left), final train loss (middle), and final test accuracy (right) achieved by non-private $\alpha$ -NormEC , after 300 communication rounds using a fine-tuned constant step size $\gamma$ .

Figure E.2: Training loss and test accuracy of non-private

\alpha

-NormEC with

\alpha=0.01

[solid],

0.1

[dashed], and

1.0

[dotted], and

\beta=0.01

[blue],

0.1

[green],

1.0

[orange], and

10.0

[red].

Figure E.5: Training loss and test accuracy of

\alpha

-NormEC with [solid] and without [dashed] server normalization.

Method: $\alpha$ -NormEC	$\beta$	$\gamma$	Final Accuracy
With server normalization	0.01	0.01	82.86%
	0.1	0.1	85.43%
	1.0	0.1	84.29%
	10.0	0.1	81.48%
Without server normalization	0.01	0.1	84.04%
	0.1	0.1	86.09 %

Appendix

Table E.1: Best configurations and final test accuracies of

\alpha

-NormEC with and without server normalization.

Figure E.6: Training loss and test accuracy of Clip21 [solid] and

\alpha

-NormEC [dashed] without server normalization in non-private training.

Figure E.9: The highest test accuracy of DP - $\alpha$ -NormEC with [left] and without [center] Server Normalization (SN), and their difference [right].

Figure E.10: Training loss and test accuracy of DP-Clip21 [solid], and DP-

\alpha

-NormEC with [dotted] and without [dashed] server normalization (SN) across different

\beta

values (with fine-tuned step sizes).

Figure E.11: Training loss and test accuracy of DP-Clip21 [solid] and DP-

\alpha

-NormEC [dashed] across different

\beta

values for

\mathchar 271\relax=1

. DP-

\alpha

-NormEC demonstrates both faster convergence and higher final accuracy for each

\beta

Figure E.12: Best test accuracy of DP- $\alpha$ -NormEC across different $\beta,\gamma$ values.

Figure E.13: Training loss and test accuracy of DP-

\alpha

-NormEC across different

\beta

values.

Figure E.14: Gradient norm and test accuracy of DP-Clip21 [solid], DP-

\alpha

-NormEC [dotted], and DP-SGD [dash-dotted] on Logistic Regression using the best hyperparameters

(\beta,\gamma)

selected by final test accuracy and then averaged over

5

seeds.

Method	Best $(\beta,\gamma)$	Final test acc (%)	Final grad norm
DP- $\alpha$ -NormEC	$(10^{-7},\,0.1)$	$\mathbf{93.68}\pm 0.656$	$\mathbf{0.074}\pm 0.023$
DP-Clip21	$(10^{-4},\,10)$	$93.509\pm 0.895$	$0.077\pm 0.027$
DP-SGD	$(10^{-1},\,0.1)$	$92.982\pm 0.961$	$0.176\pm 0.015$

Appendix

Table E.2: Final performance (mean

\pm

std over

5

seeds; selected by final test accuracy).

X_{i}^{k+1}=\mathrm{LMO}_{\mathcal{B}_{i}^{k}}\left(\nabla_{i}f(X^{k})\right):=\underset{X_{i}\in\mathcal{B}_{i}^{k}}{\arg\min}\ \langle\nabla_{i}f(X^{k}),X_{i}\rangle_{(i)}

(F.2)

Appendix

Algorithm 9 Deterministic Adaptive Layer-Wise LMO-based Optimizer

Figure F.1: Validation of layer-wise

(L^{0},L^{1})

-smoothness for the group of parameters from the embedding layer of NanoGPT-124M along unScion training trajectories. The group norm is

\|\cdot\|_{(p)}=n_{p}\|\cdot\|_{1\to\infty}

, with fitted values

L_{p}^{0}\approx 0

L_{p}^{1}\approx 1.3

. The same plot is shown twice with different

y

-axis limits.

Panel of Figure 42 — Figure F.2: Validation of layer-wise $(L^{0},L^{1})$ -smoothness for the group of parameters from the 4th transformer block of NanoGPT-124M along unScion training trajectories. The group norms are $\|\cdot\|_{(i)}=\sqrt{\nicefrac{{n_{i}}}{{m_{i}}}}\|\cdot\|_{2\to 2}$ , with fitted values $L_{i}^{0}\approx 0$ , $L_{i}^{1}\approx 70$ .

Panel of Figure 43 — Figure F.3: Validation of layer-wise $(L^{0},L^{1})$ -smoothness for the group of parameters from the 8th transformer block of NanoGPT-124M along unScion training trajectories. The group norms are $\|\cdot\|_{(i)}=\sqrt{\nicefrac{{n_{i}}}{{m_{i}}}}\|\cdot\|_{2\to 2}$ , with fitted values $L_{i}^{0}\approx 0$ , $L_{i}^{1}\approx 70$ .

Panel of Figure 44 — Figure F.4: Validation of layer-wise $(L^{0},L^{1})$ -smoothness for different groups of parameters in NanoGPT-124M along training trajectories of unScion using the specialized norm choices defined in ( F.4 ).

Panel of Figure 45 — Figure F.5: Validation of layer-wise $(L^{0},L^{1})$ -smoothness for different groups of parameters in NanoGPT-124M along training trajectories of unScion using the standard Euclidean norm $\|\cdot\|_{2}$ .

Panel of Figure 46 — Figure F.6: Validation of layer-wise $(L^{0},L^{1})$ -smoothness for different groups of parameters in NanoGPT-124M along AdamW training trajectories.

Figure F.7: Estimated trajectory smoothness for NanoGPT-124M trained with AdamW . The same cross-layer heterogeneity pattern persists, indicating that layer-wise

(L^{0}_{i},L^{1}_{i})

-smoothness is not specific to unScion .

Panel of Figure 48 — Figure F.8: Illustrative best (top row) and worst (bottom row) per‑layer fits of Assumption 1 (measured $L_{i}[k]$ vs. $L_{i}^{\text{approx}}[k]$ ) along NanoGPT-124M training.

Figure F.9: Relative fit

\mathrm{MSE}_{\mathrm{rel}}

error of the layer-wise

(L^{0}_{i},L^{1}_{i})

model across transformer blocks in NanoGPT-124M for each matrix type. Dashed lines indicate the mean and median over blocks (embedding matrix:

\mathrm{MSE}_{\mathrm{rel}}=0.0005

Figure F.10: Estimated trajectory smoothness for NanoGPT-124M trained with unScion .

Figure F.11: Estimated trajectory smoothness for GPT-2 Medium (

\sim

355M parameters) trained with unScion .

Figure F.12: Estimated trajectory smoothness for GPT-2 Large (

\sim

774M parameters) trained with unScion .

Panel of Figure 53 — Figure F.13: Validation of layer-wise $(L^{0},L^{1})$ -smoothness for different groups of parameters of a CNN model along the training trajectories of unScion with full-batch gradients . The norms used for each group are as follows: $\|\cdot\|_{(i)}=\sqrt{\nicefrac{{1}}{{C_{i}^{out}}}}\|\cdot\|_{2}$ for biases, $\|\cdot\|_{(i)}=k^{2}\sqrt{\nicefrac{{C_{i}^{in}}}{{C_{i}^{out}}}}\|\cdot\|_{2\to 2}$ for conv, and $\|\cdot\|_{(p)}=n_{p}\|\cdot\|_{1\to\infty}$ for the last group $X_{p}$ , associated with classification head weights.

Panel of Figure 54 — Figure F.14: Validation of layer-wise $(L^{0},L^{1})$ -smoothness for different groups of parameters of a CNN model along the training trajectories of unScion with stochastic gradients . The norms used for each group are as follows: $\|\cdot\|_{(i)}=\sqrt{\nicefrac{{1}}{{C_{i}^{out}}}}\|\cdot\|_{2}$ for biases, $\|\cdot\|_{(i)}=k^{2}\sqrt{\nicefrac{{C_{i}^{in}}}{{C_{i}^{out}}}}\|\cdot\|_{2\to 2}$ for conv, and $\|\cdot\|_{(p)}=n_{p}\|\cdot\|_{1\to\infty}$ for the last group $X_{p}$ , associated with classification head weights.

Panel of Figure 55 — Figure G.1: Supplementary results for NanoGPT on FineWeb: convergence behavior (left) and runtime scaling (right) as the number of PolarExpress (PolarExp) iterations varies.

PolarExpress iterations	1	2	3	4	5	6	7	8
Validation loss	3.0675	3.0316	3.0109	3.0058	3.0036	3.0037	3.0029	3.0023
Training time (minutes)	652.2	668.8	660.4	665.8	655.7	663.1	668.8	675.0

Appendix

Table G.1: Final validation loss and total training time for NanoGPT on FineWeb across different numbers of PolarExpress iterations used for the LMO approximation.

Figure G.2: Aggregate PolarExpress error over training, defined as the sum over tracked layers of

\|\mathrm{PE}_{p}(M)-\mathrm{PE}_{100}(M)\|_{2}

on the same normalized matrix

M

. Higher iteration counts maintain lower aggregate orthogonalization error throughout training.

Figure G.3: Layer-wise PolarExpress convergence at validation step 6000 for NanoGPT. For each tracked matrix family in each transformer block, we plot the spectral-norm discrepancy

\|\mathrm{PE}_{k}(M)-\mathrm{PE}_{100}(M)\|_{2}

as a function of the inner-iteration count

k

, where

M

is the normalized pre-orthogonalization momentum matrix and

\mathrm{PE}_{100}

is used as a high-accuracy reference. MLP matrices converge in only a few iterations, while several attention matrices remain difficult until roughly

k\approx 10

–

20

, revealing strong heterogeneity across parameter families and blocks.

PE iters	Mean $e_{k}$	Std. $e_{k}$	Median $e_{k}$	Max $e_{k}$	Val. loss
1	0.9826	0.0311	1.0011	1.0200	3.0675
2	0.9879	0.0220	1.0016	1.0079	3.0316
3	0.9554	0.0644	1.0012	1.0143	3.0109
4	0.8480	0.2260	0.9964	1.0055	3.0058
5	0.6952	0.4047	0.9895	1.0032	3.0036
6	0.6507	0.4321	0.9797	1.0030	3.0037

Appendix

Table G.2: Effect of the number of PolarExpress iterations on the aggregate orthogonalization discrepancy

e_{k}(M)=\|\mathrm{PE}_{k}(M)-\mathrm{PE}_{100}(M)\|_{2}

, and the final NanoGPT validation loss on FineWeb. Discrepancy statistics are aggregated over all completed checkpoints and saved tensors. This is a solver-to-solver diagnostic relative to a high-accuracy PE 100 reference on the same normalized matrix, not the error to the exact polar factor.

Figure G.4: Full final test accuracy heatmaps on CIFAR-10 across a grid of step sizes and momentum values for different numbers of Newton-Schulz (“Newton”) iterations. Brighter colors indicate higher final test accuracy. The region of high performance is visibly larger and more stable for higher numbers of Newton-Schulz iterations.

Nonconvex Stochastic Optimization

Nonconvex stochastic optimization develops methods for minimizing nonconvex objectives using noisy gradients and analyzes their convergence.

See all 96 papers

DP-Muon: Differentially Private Optimization via Matrix-Orthogonalized Momentum
May 13, 2026Jihwan Kim, Chenglin FanDeep Learning Optimization Differentially Private Stochastic Gradient Descent
Seoul National University
How Does Muon Work in Private Deep Learning? Scalable Differentially Private Training via Orthogonalization
Oct 17, 2026Alexander Gaponov, Grigory Malinovsky, Egor Shulgin +1Muon Optimizer Deep Learning Optimization
King Abdullah University of Science and Technology (KAUST) KAUST Center of Excellence for Generative AI Thuwal, Saudi Arabia
OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality
Jun 7, 2026Ganzhao YuanMomentum Methods Nonconvex Stochastic Optimization
Faculty of Computer Science and Artificial Intelligence Shenzhen University of Advanced Technology (SUAT), China

May 13, 2026J/K move · Enter open · S save

DP-Muon: Differentially Private Optimization via Matrix-Orthogonalized Momentum

Jihwan Kim, Chenglin Fan

Seoul National University

Deep Learning Optimization Differentially Private Stochastic Gradient Descent Muon Optimizer

Open paper AskPDF arXiv

We study differentially private optimization with matrix-orthogonalized momentum. DP-Muon uses conventional global per-example clipping and one Gaussian gradient release per step; matrix updates and auxiliary updates are post-processing. Our main contribution concerns the additional mean distortion created when fresh Gaussian noise passes through a nonlinear matrix map. Conditioning on the actual adaptive history immediately before the current noise yields an exact Gaussian heat identity. For a smooth Newton-Schulz map, first-order DP-MuonBC reduces this conditional output bias from second to fourth order in the fresh noise scale, and an arbitrary-order extension has bias of order

2K+2

. We prove matrix-block stationarity bounds under global clipping, retain finite-step orthogonalization error explicitly, and give an exact criterion for improvement of the resulting upper bound. A separate inequality exposes the effect of auxiliary Adam updates. GPT-2 experiments on E2E at four privacy targets favor the reported Muon configurations over Adam baselines in test NLL.

Operator	Property
Contractive compressor $\mathcal{C}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$	$\left\\|\mathcal{C}(g)-g\right\\|^{2}\leq(1-\eta)^{2}\left\\|g\right\\|^{2}$
Clipping ${\rm Clip}_{\tau}\left(g\right):=\min\left(1,\frac{\tau}{\left\\|g\right\\|}\right)g$	$\left\\|{\rm Clip}_{\tau}\left(g\right)-g\right\\|^{2}\leq\max(0,\left\\|g\right\\|-\tau)^{2}$
Smoothed normalization ${\rm Norm}_{\alpha}\left(g\right):=\frac{1}{\alpha+\left\\|g\right\\|}g$	$\left\\|{\rm Norm}_{\alpha}\left(g\right)-g\right\\|^{2}\leq\left(1-\frac{1}{\alpha+\left\\|g\right\\|}\right)^{2}\left\\|g\right\\|^{2}$

Abstract

Nonconvex Stochastic Optimization

Explore similar work

DP-Muon: Differentially Private Optimization via Matrix-Orthogonalized Momentum

How Does Muon Work in Private Deep Learning? Scalable Differentially Private Training via Orthogonalization

OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality