Inner Momentum for Differentially Private Muon

preprint

I NNER M OMENTUM FOR D IFFERENTIALLY P RIVATE M UON Bishnu Bhusal1,2 Minh Vu1 Ben Southworth1 Rohit Chadha2 Manish Bhattarai1 1

Los Alamos National Laboratory

2

University of Missouri

Geigh Zollicoffer1,3 3

Georgia Institute of Technology

arXiv:2610.02738v1 [cs.LG] 2 Oct 2026

A BSTRACT Differentially private training clips each per-example gradient before adding noise. This clipping is radial for each example, yet unequal clipping factors can distort the relative singular-vector geometry of their average. Muon is particularly exposed to this effect, since its update is an approximate polar factor U V ⊤ that depends only on the singular vectors that clipping can shift. To curb this degradation, we propose averaging each sampled example’s Muon gradient over the current model and a short history of recent models before clipping. The clipped batch matrix then separates into a common rescaling and a covariance residual R between sampled gradients and clipping values, with ∥R∥F ≤ σλ σG , bounding the clipping-induced distortion directly. We further show that a finite Newton–Schulz iteration preserves the polar factor of its input under these spectral conditions, confirming that our correction survives orthogonalization. In private GPT-2 fine-tuning on E2E and DART at ϵ ∈ {1, 2, 4, 8}, DP-Muon-IM improves BLEU and ROUGE-L over DP-Muon in every seed-matched comparison, and non-private diagnostics show 2–4% lower pre-noise polar error.

1

I NTRODUCTION

Training a neural network can reveal information about individual training records. Differential privacy (DP) bounds this leakage by ensuring that a model’s output distribution changes only slightly depending on whether any single record was used in training (Dwork and Roth, 2014). A standard approach in private deep learning is DP stochastic gradient descent (DPSGD). It clips each perexample gradient before averaging, which limits the influence of one record, and adds Gaussian noise calibrated to that limit (Abadi et al., 2016). Both operations can affect utility. Noise perturbs the averaged gradient, while clipping can change its direction by rescaling per-example gradients unequally (Chen et al., 2020). This work studies how clipping affects Muon (Jordan et al., 2024), an optimizer that approximately orthogonalizes matrix-valued momentum using a small number of Newton–Schulz iterations. Muon has recently been used alongside AdamW (Loshchilov and Hutter, 2019) to train modern large language models, including DeepSeek-V4 (DeepSeek-AI, 2026) and Kimi K2 via MuonClip (Kimi Team, 2025), with Muon-based training outperforming AdamW baselines on small transformers and Moonlight (Liu et al., 2025). Motivated by these gains, we ask whether Muon’s advantages persist under the per-example gradient clipping and Gaussian noise required for differentially private training. In fact, per-example clipping affects Muon by changing the relative weights of the gradients in the batch average. As shown in Fig. 1, when clipping factors differ across examples, the clipped average can have different singular-vector bases, and hence a different polar factor (Chen et al., 2020; Xiao et al., 2023). DP-Muon (Kim and Fan, 2026) adds Gaussian noise to this clipped average and then applies momentum and Newton–Schulz. Because the Newton–Schulz iteration preserves the polar factor of its input in exact arithmetic, it cannot correct clipping-induced changes to that polar factor. Instead, we propose to smooth each sampled example’s Muon gradient across the current and recent private model states before clipping. Under the stated variance and temporal-covariance assumptions, averaging can tighten a bound on the expected clipping-residual norm. We measure the realized residual and polar error separately. The resulting method, DP-Muon-IM, adapts the inner-momentum 1

preprint

Figure 1: Intuition for DP-Muon with inner momentum (DP-Muon-IM). (a) DP-Muon-IM averages each example’s Muon gradients at the current and recent private model states before flat clipping. (b-c) Unequal clipping factors introduce a residual R relative to a common rescaling of the batch mean. IM averages per-example gradients before this reweighting. (d) The residual can change the polar factor. We measure residual size and polar error separately; the utility experiments evaluate the complete method. approach of Xiao et al. (2023) to realize that intuition. We further formalize this idea by decomposing the clipped average into a common-rescaling term, which preserves the desired polar factor, and a residual governed by variation in the gradients and clipping factors. Privacy accounting remains unchanged: each example contributes a single query clipped to the same radius as in the baseline Gaussian mechanism, with the same sampling procedure and release schedule. In GPT-2 fine-tuning on E2E and DART, DP-Muon-IM improves BLEU and ROUGE-L over DP-Muon. Our contributions are as follows. • Section 3 explains how clipping can alter the matrix geometry relevant to Muon by decomposing the clipped average into a common rescaling and a residual R, with ∥R∥F ≤ σλ σG . We also show that the stated finite Newton–Schulz iteration preserves its immediate input’s polar factor in exact arithmetic, so it cannot remove polar error present at that stage. • Section 4 introduces DP-Muon-IM, which averages each sampled example’s gradients at current and recent private model states before clipping. The sampling rule, clipping radius, noise calibration, release schedule, and baseline per-run privacy accounting remain unchanged. • Section 5 presents private GPT-2 experiments showing that DP-Muon-IM improves both BLEU and ROUGE-L scores on E2E and DART compared with DP-Muon. We report E2E and DART results at target ϵ ∈ {1, 2, 4, 8}, and a single-seed history-length ablation quantifies the tradeoff between generation quality, training time, and memory. Organization. Section 2 introduces notation and background. Section 3 analyzes clipping-induced distortion, and Section 4 describes DP-Muon-IM. Section 5 reports the experiments. Sections 6–8 cover related work, discussion, and conclusions.

2

P RELIMINARIES

2.1 P ROBLEM SETUP AND NOTATION Let D = {xi }ni=1 be the private training set, let θ denote the model parameters, and let ℓ(θ; xi ) be the loss on example xi . Training seeks to minimize the empirical risk n

F (θ) =

1X ℓ(θ; xi ). n i=1

(1)

At training step t, the optimizer samples a minibatch St (also called a lot in DP training), and we write m = |St | for its realized size. The per-example gradient is (i)

Gt = ∇θ ℓ(θt ; xi ). 2

(2)

preprint

Throughout we treat a gradient as a vector over all trainable parameters. However, algorithmically Muon operates on matrix representations of specific subsets of parameters. Thus let W be the set of hidden-layer matrices operated on by Muon, with each matrix corresponding to a specific subset of (i) parameter indices. For any W ∈ W, the corresponding block of Gt is implicitly converted to its matrix form, W W (i) Gt,W = ∇W ℓ(θt ; xi ) ∈ Rnout ×nin . (3) Muon acts separately on each W ∈ W. Non-Muon parameters are omitted from the blockwise (i) notation. If present, their gradients remain in Gt , receive the same clipping-and-noise mechanism, (i) and are updated by the auxiliary optimizer. The norm ∥Gt ∥2 treats all parameter blocks as one vector. 2.2 D IFFERENTIAL PRIVACY AND PRIVATE OPTIMIZATION Two datasets are neighbors if they differ by the addition or removal of one record. A randomized algorithm M is (ϵ, δ)-differentially private (Dwork and Roth, 2014) if, for every neighboring pair D, D′ and every measurable output set E, Pr[M(D) ∈ E] ≤ eϵ Pr[M(D′ ) ∈ E] + δ,

(4)

where ϵ is the privacy budget and δ allows a small additive relaxation; smaller values give stronger privacy. To achieve this in optimization, given a clipping radius C > 0, flat clipping bounds each example’s contribution by rescaling its complete gradient: ( ) C (i) (i) (i) (i) clipC (Gt ) = λt Gt , λt = min 1, , (5) (i) ∥Gt ∥2 (i)

with λt

(i)

= 1 when Gt

(i)

= 0. For matrix-level equations, clipC (Gt,W ) denotes the W block of

(i)

clipC (Gt ), with the clipping factor still determined by the full gradient norm. Letting B be the (expected) lot size, the private release is ! 1 X (i) e clipC (Gt ) + σDP CZt , Gt = Zt ∼ N (0, I), (6) B i∈St

e t,W denoting the matrix block of G e t for W ∈ W. Clipping controls sensitivity, Gaussian with G noise provides privacy, and all subsequent optimizer operations are post-processing. 2.3 M UON AND DP-M UON e t,W Muon updates each hidden-layer matrix separately (Jordan et al., 2024). Fix W ∈ W, and let G be its block of the private release in Eq. (6). DP-Muon (Kim and Fan, 2026) applies e t,W , Mt,W = µMt−1,W + G Ot,W = Newton-SchulzJ (Nt,W ),

e t,W + µMt,W , Nt,W = G Wt+1 = Wt − ηOt,W .

(7)

Here Mt,W is the outer-momentum buffer, Nt,W is the Nesterov direction, and Ot,W is the orthogonalized direction. The operator Newton-SchulzJ applies J Newton–Schulz steps; Section 3.2 gives the exact iteration. We call Mt,W outer momentum to distinguish it from the per-example inner momentum introduced next. Inner momentum acts before clipping, whereas outer momentum acts on the private release. Section 4 gives the complete training procedure. 2.4 M ETRICS For nonzero matrices X and Y , define the cosine error dcos (X, Y ) = 1 −

⟨X, Y ⟩F . ∥X∥F ∥Y ∥F

(8)

This is cosine distance after vectorizing the matrices. For a reduced singular value decomposition X = U ΣV ⊤ over the positive singular values, define polar(X) = U V ⊤ and polar(0) = 0. Individual singular vectors are not unique within a repeated-singular-value subspace, but U V ⊤ is 3

preprint

invariant to the corresponding joint change of basis. It is the partial isometry mapping the right singular subspace to the left singular subspace. The polar error is dpolar (X, Y ) = ∥polar(X) − polar(Y )∥F .

(9)

Cosine error compares all entries, whereas polar error compares this basis-invariant map used by Muon. We report both because two matrices can differ in entrywise cosine similarity while sharing a polar factor.

3

G EOMETRIC EFFECTS OF GRADIENT CLIPPING

Muon’s Newton–Schulz update preserves the input’s polar factor while transforming its singular values. This makes flat clipping, as used in standard DP-SGD, a poor fit for Muon. The reason is that different examples receive different clipping factors, clipping reweights each example’s contribution. In other words, the straightforward rescale-then-average mechanism that makes DP-SGD’s privacy guarantee simple actually corrupts the geometric object that drives Muon’s update. We make this concrete in three steps. First, in §3.1 we decompose the clipped batch average into a common-rescaling term, which would leave the polar factor unchanged, and a residual term R that captures exactly the distortion caused by unequal per-example clipping factors. Second, in §3.2 we show that Newton–Schulz orthogonalization, which computes the polar factor of a matrix, does not undo this distortion. Third, in §B we give a minimal two-example numerical example that makes the resulting change in update direction explicit and quantifiable. 3.1

T HE CLIPPING RESIDUAL

Fix a step t and a block W ∈ W. Write Gi for the matrix supplied by example i for this block before clipping, and let λi be the clipping factor computed from its complete query. For DP-Muon, (i) Gi = Gt,W . We suppress t and W below. For a batch St of size m, define Gavg =

1 X Gi , m

Gclip =

i∈St

1 X λi Gi , m i∈St

λ̄ =

1 X λi . m

(10)

i∈St

If every example has the same clipping factor, then Gclip = λ̄Gavg , so clipping cannot change the singular vectors or polar factor of the mean. Unequal factors add the residual below. Define the variation in clipping factors and matrix gradients by σλ2 =

1 X (λi − λ̄)2 , m

2 σG =

i∈St

1 X 2 ∥Gi − Gavg ∥F . m

(11)

i∈St

Proposition 1 (Residual identity and bound). The clipped mean can be written as Gclip = λ̄Gavg + R,

R=

1 X (λi − λ̄)(Gi − Gavg ), m

(12)

i∈St

where the residual satisfies ∥R∥F ≤ σλ σG . The proof is in Appendix A.1. The term λ̄Gavg preserves the singular vectors and polar factor, so any clipping-induced change to them must arise from R. Not every component of R changes the polar factor, however, and the bound in Proposition 1 is only a worst-case envelope; it need not predict the observed polar discrepancy tightly. The residual vanishes when the clipping factors are equal (σλ = 0) or the per-example matrices are equal (σG = 0). The same argument applies to per-matrix clipping using the block’s own clipping factor. Section 5.3 reports polar error, and Appendix D.1 reports the scale-normalized clipping residual for historical and current-only queries. Appendix B gives a worked two-example calculation of the clipping factors, residual, and polar error. 4

preprint

3.2

W HAT N EWTON –S CHULZ PRESERVES

Muon uses a small number of Newton–Schulz iterations, each a polynomial in XX ⊤ applied to X: X + = aX + b(XX ⊤ )X + c(XX ⊤ )2 X,

(a, b, c) = (3.4445, −4.7750, 2.0315).

(13)

Let Newton-SchulzJ (X) denote J such steps, starting from X/(∥X∥F + ε) with ε ≥ 0.1 Observe that each step acts only on singular values: if X = U ΣV ⊤ in reduced SVD form, then X + = U p(Σ) V ⊤ , where p(σ) = aσ + bσ 3 + cσ 5 acts entrywise on Σ. The singular vectors are therefore unchanged, and the polar factor is preserved as long as no singular value changes sign. Proposition 2 makes this precise. Proposition 2 (Newton–Schulz preserves the polar factor). Suppose c > 0 and b2 < 4ac. Then, in exact arithmetic, for
every nonzero matrix X, every ε ≥ 0, and every finite J ≥ 1, polar Newton-SchulzJ (X) = polar(X). Consequently, for nonzero X and Y ,
dpolar Newton-SchulzJ (X), Newton-SchulzJ (Y ) = dpolar (X, Y ). (14) Muon’s coefficients satisfy both conditions: 4ac − b2 = 5.189382 > 0. The proof is in Appendix A.2. For a constant clipping factor across samples, λi = λ̄, we have R = 0 in Eq. (12). Then Gclip is a positive rescaling of Gavg , so their polar factors, and hence their Muon updates, coincide. Any change in the polar factor due to non-constant clipping is therefore captured by the residual R. Note that Eq. (14) concerns only the matrices entering Newton–Schulz. In DP-Muon (Kim and Fan, 2026), Gaussian noise and outer momentum are applied first and may change the discrepancy. Even a common clipping rescaling can change the signal-to-noise ratio or the balance between the current gradient and the momentum buffer.

4

I NNER MOMENTUM BEFORE CLIPPING

Section 3 showed that unequal per-example clipping changes the clipped mean’s polar factor only through the residual R, with ∥R∥ ≤ σλ σG (Proposition 1). To reduce this distortion, we average each example’s gradients over the current and recent private model states before clipping, while leaving the rest of DP-Muon unchanged. We call this inner momentum, distinguishing it from Muon’s outer momentum, which averages noisy batch gradients after clipping. The section has three parts. §4.1 defines the per-example average. §4.2 is the central result: averaging shrinks the envelope of R by a factor ρ + (1 − ρ)SK , where ρ measures how correlated an example’s gradient fluctuations are across states and SK is the sum of squared averaging weights. The factor is below one when ρ < 1 and at least two weights are nonzero. §4.3 gives the resulting algorithm, its cost, and its privacy guarantee. The envelope result bounds R and does not prove the realized R is smaller, we measure the resulting polar error directly in Section 5. 4.1 AVERAGING ONE EXAMPLE ACROSS MODEL STATES Let Kt = min{t, Kmax }. For i ∈ St and each Muon matrix block W ∈ W, evaluate (i)

Gt−j,W = ∇W ℓ(θt−j ; xi ),

j = 0, . . . , Kt ,

(15)

where j = 0 is the current model and every term uses the same example xi . With decay γ ∈ [0, 1] PKt ℓ and normalized weights αj = γ j / ℓ=0 γ , form (i)

Ht,W =

Kt X

(i)

αj Gt−j,W .

(16)

j=0 (i)

(i)

If Kt = 0 or γ = 0, then Ht,W = Gt,W and we recover DP-Muon. For each example, the blocks (i)

Ht,W are combined with the non-Muon blocks, which are averaged in the same way with their own 1 These coefficients, from Jordan et al. (2024), are chosen so that small singular values grow quickly; the iterates approach an approximately orthogonal matrix rather than converging exactly to the polar factor.

5

preprint

history length and decay (KA , γA ), and clipped as one query, with a scalar factor λi computed from (i) all trainable coordinates. Proposition 1 then applies with Gi = Ht,W . P We normalize by j γ j , unlike the unnormalized sum of Xiao et al. (2023), for two reasons. First, (i)

Ht,W stays on the scale of a single gradient, so the clipping radius C keeps its meaning. An P unnormalized sum would inflate query norms by up to j γ j and force stronger clipping. Second, the weights sum to one, which is what the analysis in §4.2 uses. 4.2

AVERAGING SHRINKS THE RESIDUAL ENVELOPE

Proposition 1 bounds R by the product σλ σG , so the question is whether averaging lowers both factors. Let gj (z) = ∇θ ℓ(θt−j ; z) ∈ Rd be example z’s full-query gradient over all trainable coordinates, including Muon and non-Muon parameters. Define µj,u = E gj (z) and . With P ξj,u (z) = gj (z) − µj,uP the same weights for all coordinates, the averaged query is ui = j αj gj (zi ). Let SK = j αj2 , which is 1 without averaging and 1/(K + 1) for equal weights. We assume, for constants τu ≥ 0 and ρ ∈ [0, 1], that E∥ξj,u ∥22 ≤ τu2 and E⟨ξj,u , ξℓ,u ⟩ ≤ ρτu2 for j ̸= ℓ on the same example. The examples in a batch are independent. Set
s2K,u = τu2 ρ + (1 − ρ)SK . (17) Here τu bounds the root-mean-square fluctuation of a full-query gradient, and ρ bounds its cross-state covariance relative to τu2 . Proposition 3 (Residual envelope under averaging). Let λi = min{1, C/∥ui ∥2 } be the flat-clipping factor of example i’s averaged query ui , and let RW be the residual of Proposition 1 for a block W of the query. Then E ∥RW ∥F ≤

m − 1 s2K,u m − 1 τu2 = ρ + (1 − ρ)SK . m C m C

(18)

The current-only query (K = 0) satisfies the same bound with s0,u = τu . Averaging therefore shrinks the envelope by the factor ρ + (1 − ρ)SK , which is below one whenever ρ < 1 and at least two weights are nonzero. Why it holds. The proof has two steps, each answering one factor of Proposition 1. (i) σG shrinks. The weights sum to one, so the population means µj,u cancel when one example’s averaged gradient is compared with the batch mean. The remaining fluctuations combine with variance at most τu2 (ρ + (1 − ρ)SK ). (ii) σλ is controlled by the same dispersion. The map n 7→ min{1, C/n} is 1/CLipschitz, and | ∥ui ∥ − ∥ū∥ | ≤ ∥ui − ū∥, so σλ ≤ σu /C, where σu is the dispersion of the complete queries. A block is a subset of coordinates, so σG,W ≤ σu , and hence ∥RW ∥F ≤ σλ σG,W ≤ σu2 /C, which is quadratic in the dispersion. The full proof is in Appendix F. Separate histories. Proposition 3 uses one set of weights for all coordinates. Suppose instead that each parameter group b (for example, Muon and AdamW blocks) uses its
own weights, P with constants τb , ρb and squared-weight sum SK,b . Set s2b = τb2 ρb + (1 − ρb )SK,b and s2u = b s2b . The same argument, combined with Cauchy–Schwarz, gives E ∥RW ∥F ≤

m − 1 su sW , m C

(19)

where sW is the constant of the group containing W . A group without history (SK,b = 1) contributes an unreduced term to su . What it does and does not say. It bounds R from above and does not prove that the realized R is smaller. No lower bound on the current-only residual is available in general. The reduction also needs ρ < 1, meaning an example’s gradient fluctuations at nearby states must not be perfectly correlated. This is plausible with dropout or other evaluation randomness and with parameter drift. Staleness does not enter Eq. (18): averaging over older states affects the residual only through fluctuations, and affects the bias of the averaged gradient through a drift term DK that grows with K and γ. Appendix F (Theorem 6) shows how the two combine in a polar-error bound. This is the trade-off that Kmax and γ control. 6

preprint

4.3 T RAINING PROCEDURE Algorithm 1 gives the update for one Muon block W ∈ W, where W is the set of Muon matrices. We drop the block subscript and apply the same steps to every block. The only change from DP-Muon is (i) (i) line 4, which replaces the current gradient Gt by the averaged Ht of Eq. (16). Clipping, noise, outer momentum, and Newton–Schulz are unchanged. Because Newton–Schulz preserves the polar factor (Proposition 2), any correction must act before it, so we leave it unchanged. Algorithm 1 DP-Muon with inner momentum (DP-Muon-IM). Require: Iterations T , learning rate η, momentum µ, and NS steps J Require: History Kmax , γ and privacy parameters C, σDP , q, B 1: M−1 ← 0. 2: for t = 0, . . . , T − 1 do 3: Poisson-sample St ⊂ D with rate q; Kt ← min{t, Kmax }. PKt j j=0 γ ∇W ℓ(θt−j ; xi ) (i) 4: Ht ← , i ∈ St . PKt j j=0 γ ! 1 X (i) e 5: Gt ← clipC (Ht ) + σDP CZt , Zt ∼ N (0, Idim(W ) ). B i∈S t

et . 6: Mt ← µMt−1 + G e t + µMt . 7: Nt ← G 8: Ot ← Newton-SchulzJ (Nt ). 9: Wt+1 ← Wt − ηOt . 10: end for 11: return θT

Cost. The history holds at most Kmax pre-update model states, and each sampled example needs up to Kt + 1 gradient evaluations, so the per-step cost grows linearly in Kmax . Privacy. History averaging adds no private releases. To analyze privacy at each step, we hold the previous noisy releases fixed, which also fixes the current and stored model states. Each sampled example then contributes one query clipped at C, and Gaussian noise is added once to the batch sum. The baseline accounting applies whenever the sampling rule, clipping radius, noise calibration, release denominator, and number of releases are unchanged. Theorem 4 makes this precise (see Appendix E for the proof). Theorem 4 (Privacy of DP-Muon-IM). Algorithm 1 is (ϵ, δ)-DP under add/remove adjacency whenever a valid Poisson-sampled Gaussian accountant certifies this guarantee for T updates with sampling rate q and noise multiplier σDP . The guarantee holds for any fixed public history length Kmax and decay γ ∈ [0, 1].

5

E XPERIMENTS

We fine-tune GPT-2 on two table-to-text generation tasks and compare DP-Muon-IM with DP-Muon across privacy budgets. We also compare with published Muon-based private optimizers, and study how the history length affects generation quality, training cost, and polar error. Data and model. We fully fine-tune GPT-2-small on E2E (Novikova et al., 2017) and DART (Nan et al., 2021). This is a standard benchmark for private fine-tuning (Li et al., 2022; Zhang et al., 2025) and the setting used to evaluate DP-Muon (Kim and Fan, 2026). Training/validation/test splits contain 42,061/4,672/4,693 records on E2E and 62,659/6,980/12,552 on DART. Optimizers and training. Our baseline, DP-Muon, disables both inner histories (K = 0). DP-MuonIM adds separate Muon and AdamW histories, with lengths KM , KA and decays γM , γA , extending Section 4. Each per-example query is clipped at C = 1 before noise is added. We train for ten epochs on one H100 80GB, with nominal logical/physical batches of 1024/32. Muon uses momentum 0.95, Nesterov updates, and five Newton–Schulz steps. Learning rates decay linearly without warm-up or weight decay. Hyperparameter details for our runs are provided in Appendix C. Privacy and selection. Our utility runs target ϵ ∈ {1, 2, 4, 8} with δ = 1/(2N ), where N is the number of training records. We calibrate noise using PRV accounting (Gopi et al., 2021). We select configurations by validation loss on seed 37 and evaluate seeds 37–39. 7

preprint

Table 1: E2E test results for GPT-2 fine-tuning at target ϵ = 8. BLEU ↑

ROUGE-L ↑

Kim and Fan (2026) Kim and Fan (2026)

63.78 63.66

66.96 66.71

This work This work This work This work

64.77 64.95 64.79 65.65

67.37 67.36 67.58 67.80

Method

Source

DP-Muon DP-Muon-BC KF-DP-Muon-IM KF-DP-Muon DP-Muon DP-Muon-IM

Table 2: Test generation after ten epochs: mean ± standard deviation over seeds 37–39. DP-Muon

DP-Muon-IM

BLEU ↑

ROUGE-L ↑

BLEU ↑

ROUGE-L ↑

1 2 4 8

57.64 ± 1.97 63.05 ± 0.34 64.10 ± 0.29 64.74 ± 0.05

62.71 ± 0.21 65.89 ± 0.62 66.72 ± 0.18 67.31 ± 0.25

60.21 ± 2.65 63.76 ± 0.26 64.83 ± 0.57 65.34 ± 0.39

64.41 ± 1.02 66.43 ± 0.60 67.03 ± 0.39 67.60 ± 0.17

1 2 4 8

32.34 ± 0.59 34.88 ± 0.37 36.48 ± 0.35 37.44 ± 0.33

50.20 ± 0.23 52.27 ± 0.12 53.52 ± 0.27 54.35 ± 0.30

33.19 ± 0.42 35.94 ± 0.29 37.16 ± 0.33 37.95 ± 0.19

50.91 ± 0.16 53.09 ± 0.13 54.24 ± 0.17 54.79 ± 0.12

Task

Target ϵ

E2E

DART

Evaluation. We compute BLEU (Papineni et al., 2002) and multi-reference ROUGE-L (Lin, 2004). Both scores are reported on a 0–100 scale; higher is better. E2E has 630 test prompts and 4,693 references; DART has 5,097 prompt groups and 12,551 nonempty references. 5.1

C OMPARISON WITH EXISTING METHODS

Table 1 compares Muon-based optimizers on E2E at target ϵ = 8. We take the DP-Muon and DPMuon-BC scores from Kim and Fan (2026); DP-Muon-BC adds a correction for privacy-noise bias after Newton–Schulz. Our KF variants apply DiSK’s simplified Kalman filter to the private release before the Muon and AdamW updates (Zhang et al., 2025). KF-DP-Muon disables inner history; KF-DP-Muon-IM adds history averaging. DP-Muon and DP-Muon-IM use per-metric maxima over seeds 37–39. 5.2

U TILITY ACROSS PRIVACY BUDGETS

DP-Muon-IM improves BLEU and ROUGE-L over DP-Muon in all 24 seed-matched comparisons (two datasets, four privacy budgets, three seeds; Table 2). Mean BLEU gains are 0.60–2.57 points on E2E and 0.50–1.06 on DART. Selected E2E configurations use KM = 4, γM = γA = 0.5, and KA = 3 for ϵ ∈ {1, 2, 4} or KA = 2 for ϵ = 8. DART uses (KM , KA ) = (3, 3) and (γM , γA ) = (0.7, 0.5) throughout. Appendix C.1 reports a fixed-learning-rate ablation with Muon-only history, AdamW-only history, both, and neither. 5.3

E FFECT OF HISTORY LENGTH

Utility and cost. We vary KM = KA = K ∈ {0, . . . , 4} on E2E at ϵ ≈ 4 for ten epochs, with Muon/AdamW learning rates 0.002/0.0035 and decays γM = γA = 0.5 for K > 0; K = 0 is DP-Muon without history. One historical state provides most of the BLEU gain: K = 0, 1, 4 yield validation BLEU scores of 66.97, 67.76, and 67.82 in 36.9, 59.0, and 123.7 minutes, respectively, with nonmonotonic gains across K (Figure 2, left). Peak GPU allocation rises from 26.70 GiB without history to 49.85 GiB with history. Timings exclude setup, evaluation, and checkpoint I/O. Polar error. We measure polar error with Muon-only history, γ = 0.9 (Figure 2, right). Each K follows a separate trajectory. Within each trajectory, we compare IM and current-only clipped means 8

Validation BLEU ↑

68.0

Polar-error reduction (%) ↑

preprint

K=4

K=1

K=3

67.5 K=2

67.0

K=0 (no history)

40

80

120

Training time (min)

K=4

4 K=2

3 2 1

K=3

K=1

K=0 (no history)

No reduction

0 0

1

2

3

4

Muon history length K

Figure 2: Effect of history length in separate utility and geometry experiments. Left: validation BLEU versus training time after ten epochs. Right: pre-noise polar-error reduction before optimizer update step 10, with Muon-only history.

at the same model state against a shared unclipped current-gradient reference, before noise and outer momentum. Error is the root-sum-square polar-factor distance across 48 Muon matrices, using 512 query records and 2,048 disjoint held-out reference records. Just before the 10th optimizer update, IM reduces polar error by 1.76% at K = 1 and 3.71% at K = 4 relative to the current-only estimate. Clipping residual. With K = 4 and γ = 0.9, historical averaging lowers the normalized clipping residual in all 21 seed–checkpoint comparisons. Mean paired reductions range from 25.0% to 46.1% across checkpoints (Appendix D.1). These raw-query diagnostics precede noise and outer momentum.

6

R ELATED W ORK

Clipping bias and pre-clipping averaging. Geometric analyses show that clipping can bias private optimization by changing the aggregate gradient (Chen et al., 2020). Xiao et al. (2023) relate this bias to per-example sampling noise, introduce inner–outer momentum, and report cosine similarity between unclipped and averaged-clipped gradients. Adaptive clipping and error feedback offer other treatments of clipping bias (Bu et al., 2023; Zhang et al., 2024); DP-PMLF combines per-sample momentum with post-release filtering (Xu et al., 2026). We adapt pre-clipping averaging to Muon’s matrix updates and study the resulting clipping residual (Section 3.1). Muon under differential privacy. Muon applies momentum and a finite Newton–Schulz polynomial to hidden-layer matrix gradients (Jordan et al., 2024). DP-Muon studies private matrix optimization and Gaussian perturbations under this transform (Kim and Fan, 2026); Kim and Oh (2026) analyze optimization with finitely many Newton–Schulz steps. Section 3.2 characterizes the polar geometry preserved by the iteration used here.

7

D ISCUSSION

The E2E and DART results support averaging per-example gradients before clipping. The residual decomposition explains why this stage matters: unequal clipping can change the batch mean’s polar factor, while the Newton–Schulz iteration used here preserves its immediate input’s polar factor in exact arithmetic. The geometry probes compare clipped IM and current-only means against the same current-gradient reference. With history length K = 4 and decay γ = 0.9, IM reduces polar error by 2.4–3.7% across measured checkpoints (Appendix D). These measurements precede Gaussian noise and outer momentum; they do not measure the final Muon update. The E2E ablation favors a short history when compute is limited: K = 1 came within 0.06 validation BLEU points of K = 4 in less than half the training time (Figure 2). This result favors a short history in the tested setting. The extra gradient evaluations and stored model states remain costs to address in larger-scale experiments. 9

preprint

8

C ONCLUSION

We identify a covariance residual caused by unequal clipping and show that the Newton–Schulz iteration used here preserves its input’s polar factor in exact arithmetic. DP-Muon-IM averages per-example gradients at current and recent private model states before clipping, retaining baseline per-run privacy accounting under the stated conditions. In GPT-2 fine-tuning, DP-Muon-IM achieves higher BLEU and ROUGE-L than DP-Muon on E2E and DART at target ϵ ∈ {1, 2, 4, 8}. AI U SE S TATEMENT We used large language models for language editing, LaTeX support, bibliographic checks, and mathematical consistency. We reviewed all AI-assisted work and take full responsibility for the paper.

R EFERENCES Martín Abadi, Andy Chu, Ian Goodfellow, H. Brendan McMahan, Ilya Mironov, Kunal Talwar, and Li Zhang. Deep learning with differential privacy. In Proceedings of the 2016 ACM SIGSAC Conference on Computer and Communications Security, CCS ’16, pages 308–318. Association for Computing Machinery, 2016. doi: 10.1145/2976749.2978318. URL https://doi.org/10. 1145/2976749.2978318. Zhiqi Bu, Yu-Xiang Wang, Sheng Zha, and George Karypis. Automatic clipping: Differentially private deep learning made easier and stronger. In Advances in Neural Information Processing Systems, volume 36, pages 41727–41764, 2023. doi: 10.52202/075280-1808. First posted as arXiv:2206.07136 in 2022. Xiangyi Chen, Zhiwei Steven Wu, and Mingyi Hong. Understanding gradient clipping in private SGD: A geometric perspective. In Advances in Neural Information Processing Systems, volume 33, pages 13773–13782, Red Hook, NY, USA, 2020. Curran Associates Inc. URL https://proceedings.neurips.cc/paper/2020/hash/ 9ecff5455677b38d19f49ce658ef0608-Abstract.html. DeepSeek-AI. DeepSeek-V4: Towards highly efficient million-token context intelligence. arXiv preprint arXiv:2606.19348, 2026. URL https://arxiv.org/abs/2606.19348. Cynthia Dwork and Aaron Roth. The algorithmic foundations of differential privacy. Foundations and Trends in Theoretical Computer Science, 9(3–4):211–487, August 2014. doi: 10.1561/0400000042. URL https://doi.org/10.1561/0400000042. Sivakanth Gopi, Yin Tat Lee, and Lukas Wutschitz. Numerical composition of differential privacy. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 11631–11642. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/ 2021/file/6097d8f3714205740f30debe1166744e-Paper.pdf. 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/. Gyu Yeol Kim and Min-hwan Oh. Convergence of Muon with Newton– Schulz. In International Conference on Learning Representations, 2026. URL https://proceedings.iclr.cc/paper_files/paper/2026/hash/ 4c0129f1f7a036af5ed2409f1f9f8131-Abstract-Conference.html. Also available as arXiv:2601.19156. Jihwan Kim and Chenglin Fan. DP-Muon: Differentially private optimization via matrixorthogonalized momentum. arXiv preprint arXiv:2605.12994, 2026. URL https://arxiv. org/abs/2605.12994. Kimi Team. Kimi K2: Open agentic intelligence. arXiv preprint arXiv:2507.20534, 2025. URL https://arxiv.org/abs/2507.20534v2. Revised February 2026. 10

preprint

Ren-Cang Li. New perturbation bounds for the unitary polar factor. SIAM Journal on Matrix Analysis and Applications, 16(1):327–332, 1995. doi: 10.1137/S0895479893256359. Xuechen Li, Florian Tramèr, Percy Liang, and Tatsunori Hashimoto. Large language models can be strong differentially private learners. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=bVuP3ltATMz. Chin-Yew Lin. ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Out, pages 74–81, Barcelona, Spain, July 2004. Association for Computational Linguistics. URL https://aclanthology.org/W04-1013/. Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang, Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982, 2025. URL https://arxiv. org/abs/2502.16982. Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id= Bkg6RiCqY7. Linyong Nan, Dragomir Radev, Rui Zhang, Amrit Rau, Abhinand Sivaprasad, Chiachun Hsieh, Xiangru Tang, Aadit Vyas, Neha Verma, Pranav Krishna, Yangxiaokang Liu, Nadia Irwanto, Jessica Pan, Faiaz Rahman, Ahmad Zaidi, Mutethia Mutuma, Yasin Tarabar, Ankit Gupta, Tao Yu, Yi Chern Tan, Xi Victoria Lin, Caiming Xiong, Richard Socher, and Nazneen Fatema Rajani. DART: Opendomain structured data record to text generation. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tur, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou, editors, Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 432–447, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.naacl-main. 37. URL https://aclanthology.org/2021.naacl-main.37/. Jekaterina Novikova, Ondřej Dušek, and Verena Rieser. The E2E dataset: New challenges for end-to-end generation. In Kristiina Jokinen, Manfred Stede, David DeVault, and Annie Louis, editors, Proceedings of the 18th Annual SIGdial Meeting on Discourse and Dialogue, pages 201–206, Saarbrücken, Germany, August 2017. Association for Computational Linguistics. doi: 10.18653/v1/W17-5525. URL https://aclanthology.org/W17-5525/. Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Pierre Isabelle, Eugene Charniak, and Dekang Lin, editors, Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics, pages 311–318, Philadelphia, Pennsylvania, USA, July 2002. Association for Computational Linguistics. doi: 10.3115/1073083.1073135. URL https://aclanthology.org/P02-1040/. Hanshen Xiao, Zihang Xiang, Di Wang, and Srinivas Devadas. A theory to instruct differentiallyprivate learning via clipping bias reduction. In 2023 IEEE Symposium on Security and Privacy (SP), pages 2170–2189, 2023. doi: 10.1109/SP46215.2023.10179409. Xincheng Xu, Thilina Ranbaduge, Qing Wang, Thierry Rakotoarivelo, and David Smith. Enhancing DPSGD via per-sample momentum and low-pass filtering. Proceedings of the AAAI Conference on Artificial Intelligence, 40(32):27341–27349, 2026. doi: 10.1609/aaai.v40i32.39951. Xinwei Zhang, Zhiqi Bu, Zhiwei Steven Wu, and Mingyi Hong. Differentially private SGD without clipping bias: An error-feedback approach. In International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=uFbWHyTlPn. Xinwei Zhang, Zhiqi Bu, Borja Balle, Mingyi Hong, Meisam Razaviyayn, and Vahab Mirrokni. DiSK: Differentially private optimizer with simplified kalman filter for noise reduction. In International Conference on Learning Representations, 2025. URL https://proceedings.iclr.cc/paper_files/paper/2025/hash/ c7793beb7f32b559df55d48370f2b8ae-Abstract-Conference.html. 11

preprint

A

P ROOFS

A.1

R ESIDUAL IDENTITY AND BOUND

We use the means and dispersions defined in Eqs. (10) and (11). Proof of Proposition 1. Expanding R and using R=

P

i∈St (λi − λ̄) = 0 gives

1 X λi Gi − λ̄Gavg = Gclip − λ̄Gavg . m i∈St

The triangle inequality and Cauchy–Schwarz yield 1 X ∥R∥F ≤ |λi − λ̄| ∥Gi − Gavg ∥F ≤ σλ σG . m i∈St

A.2

N EWTON –S CHULZ POLAR - FACTOR PRESERVATION

Recall p(σ) = aσ + bσ 3 + cσ 5 , the scalar polynomial associated with Eq. (13).   Proof of Proposition 2. Completing the square gives p(σ) = σ c(σ 2 + b/2c)2 + (4ac − b2 )/4c > 0 for all σ > 0. Starting from a reduced SVD X = U ΣV ⊤ , the normalization and each step replace Σ by another positive diagonal matrix while keeping U and V fixed. The new diagonal entries need not be in decreasing order, but sorting them only permutes the columns of U and V jointly, which leaves U V ⊤ unchanged. Hence every iterate has a reduced SVD with polar factor U V ⊤ , and by the definition in Section 2.4, polar(Newton-SchulzJ (X)) = polar(X).

B

I LLUSTRATIVE EXAMPLE OF CLIPPING DISTORTION

Consider two per-example gradient matrices:     10 0 0 −1 G1 = , G2 = . 0 10 1 0 √ With clipping radius C = 2, their clipping factors are 0.1 and 1. The means before and after clipping are     5 −0.5 0.5 −0.5 Gavg = , Gclip = . 0.5 5 0.5 0.5 The mean clipping factor is λ̄ = 0.55, so the residual is   −2.25 −0.225 R = Gclip − 0.55Gavg = . 0.225 −2.25 √ Here σλ = 0.45 and σG = 202/2, giving √ ∥R∥F = 0.225 202 = σλ σG ≈ 3.198. The bound is attained in this two-example construction, but it need not be tight in general. More importantly, Gclip is not a scalar multiple of Gavg : clipping changes the aggregate polar factor even though it preserves the polar factor of either example. In the metrics of Section 2.4, the resulting errors are 11 dcos (Gclip , Gavg ) = 1 − √ ≈ 0.226, 202

s

11 dpolar (Gclip , Gavg ) = 2 1 − √ ≈ 0.951. 202

In this two-dimensional special case, the polar factors happen to be planar rotations by 5.7◦ and 45◦ . In higher dimensions, the polar factor instead describes a partial isometry between the right and left singular subspaces. 12

preprint

C

U TILITY EXPERIMENT DETAILS

Parameter groups and updates. Muon updates tensors with at least two dimensions, excluding token and position embeddings and the language-model head. AdamW updates those excluded parameters, vectors, p and biases. For an r × c Muon matrix, the learning rate is the selected base rate times 0.2 max(r, c) (Liu et al., 2025). The five Newton–Schulz steps use coefficients (3.4445, −4.7750, 2.0315), with no extra output-level bias correction. AdamW retains standard moment bias correction and numerical ε = 10−8 . Training dropout is active; historical evaluations use a separate stochastic stream, and evaluation disables dropout. Selected histories and learning rates. Table 3 gives the base learning rates for the utility runs. At ϵ ∈ {1, 2, 4}, E2E uses (KM , γM ) = (4, 0.5) and (KA , γA ) = (3, 0.5). DART uses (KM , γM ) = (3, 0.7) and (KA , γA ) = (3, 0.5) at every budget. DP-Muon (No-IM) disables both histories. Both methods in this table use AdamW β1 = 0.9. The E2E ϵ = 8 campaign fixes (KM , KA ) = (4, 2) and γM = γA = 0.5 for IM; its baseline disables both histories. Both use Muon/AdamW learning rates 0.002/0.0035, retained from the earlier validation search without retuning. At E2E ϵ = 8, both methods are trained afresh at δ = 1/84122. Across seeds 37–39, IM gives 65.34 ± 0.39 BLEU and 67.60 ± 0.17 ROUGE-L, versus 64.74 ± 0.05 and 67.31 ± 0.25 for the baseline. Table 1 reports per-metric maxima: 65.65/67.80 for IM and 64.79/67.58 for the baseline, all from seed 39. Seed 37 informed the earlier search; seeds 38, 39 retain the previously used confirmation identities. Table 3: Validation-selected base learning rates for DP-Muon (No-IM) and DP-Muon-IM; each pair is (Muon, AdamW). Both methods retain the same standard optimizer moments. Task

Target ϵ

DP-Muon (No-IM)

DP-Muon-IM

E2E

1 2 4 8

(0.0016, 0.0028) (0.0016, 0.0035) (0.0016, 0.0035) (0.0020, 0.0035)

(0.0016, 0.0035) (0.0016, 0.0035) (0.0020, 0.0035) (0.0020, 0.0035)

DART

1 2 4 8

(0.00105, 0.00168) (0.0013125, 0.00168) (0.0013125, 0.0021) (0.0013125, 0.0021)

(0.00105, 0.00168) (0.0013125, 0.0021) (0.0013125, 0.0021) (0.0013125, 0.002688)

Private release. E2E uses Poisson sampling with q = 1/41 for 410 releases; DART uses q = 1/61 for 610 releases. The implementation uses ⌊B⌋ rather than B in Eq. (6), giving release denominators 1025 and 1027, respectively. These differ from the nominal logical batch size 1024. Clipping uses the full concatenated per-example query, and each logical update makes one Gaussian release, not one release per microbatch. Table 4 lists the verified per-run calibration. Table 4: Per-run PRV calibration for our utility runs, with δ = 1/(2N ) at every budget. Achieved ϵ is rounded. Task

Target ϵ

Noise multiplier σDP

Achieved ϵ

δ

E2E

1 2 4 8

2.05078125 1.2646484375 0.889892578125 0.67626953125

0.996290 1.999911 3.996703 7.998976

1/84122 1/84122 1/84122 1/84122

DART

1 2 4 8

1.767578125 1.123046875 0.8197021484375 0.63720703125

0.993732 1.998849 3.993658 7.995195

1/125318 1/125318 1/125318 1/125318

DiSK settings. Each DiSK closure forms its temporal query. DiSK combines the base and extrapolated queries before clipping and noise, then filters the private release before Muon/AdamW. The filter uses κ times the previous state plus (1 − κ) times the new input. 13

preprint

The single-seed KF rows in Table 1 form a separate joint ablation using learning rates 0.002/0.0021, κ = 0.5, and DiSK extrapolation coefficient 0.8. The IM row retains the same histories. The IM setting was selected by validation loss from 30 new trials and one reused anchor. Both KF studies use σDP = 0.67626953125, δ = 1/84122, and achieved ϵ ≈ 7.998976. DiSK’s extrapolation coefficient is distinct from either history decay. C.1

O PTIMIZER - GROUP ABLATION

We switch history on or off separately for Muon and AdamW, fixing their base learning rates at 0.002 and 0.0035. Enabled histories use (KM , γM ) = (4, 0.5) and (KA , γA ) = (3, 0.5). All four arms retain Muon momentum 0.95, AdamW β1 = 0.9 and the E2E ϵ = 4 privacy calibration in Table 4. Table 5: Optimizer-group ablation on E2E at target ϵ = 4 and δ = 1/84122. NLL is measured in nats per target token. History enabled

Test NLL ↓

BLEU ↑

ROUGE-L ↑

1.2842 1.2779 1.2686 1.2586

63.66 63.76 64.37 65.30

66.90 66.95 67.12 67.47

Neither Muon only AdamW only Both

D

G EOMETRIC ERROR ACROSS TRAINING

The corrected diagnostic uses consistent per-record mean target-token loss for current and historical gradients. It compares IM with a current-only query at the same state of each IM-trained trajectory, before noise and outer momentum. Both are compared with the shared finite raw-current reference described in Section 5.3. This is not a comparison of independently trained IM and baseline models.

Polar-error reduction (%) ↑

Runs use E2E, GPT-2, batch size 8, learning rate 5 × 10−5 , Muon (K, γ) = (4, 0.9), and no AdamW history. Full-query clipping uses C = 1. The three training seeds are 42, 43, 44; probes occur before updates 10, 50, 100, 250, 500, 1000, 2000. Thus index 10 means nine completed updates. Means accumulate in FP32 and polar factors use CPU FP64 SVD.

4

3

IM: mean ± SD 3.67%

2.44%

2

1 Current-only baseline

0 10

100

1000

Recorded training step (log scale)

Figure 3: Corrected three-seed polar diagnostic at (K, γ) = (4, 0.9). Relative reduction is computed within each seed as 100(1 − EIM /Ecurrent ), then averaged. Faint lines show individual seeds; bars and shading are pointwise sample SD, not confidence intervals. Both cosine and polar error are lower with IM in all 21 seed–checkpoint pairs. Mean polar-error reduction decreases from 3.67% before update 10 to 2.44% before update 2000 (Figure 3); mean 14

preprint

cosine-error reductions range from 3.20% to 5.98%. These distances are relative to a finite reference, not a population gradient or the noisy Newton–Schulz update. The one-seed K sweep in Figure 2 uses the same raw-query diagnostic design at index 10. It pairs IM and current-only queries within each trajectory; the K = 4 point reuses the corrected seed-42 run above. The K = 0–3 points are completed early probes, not final-training results. The two panels of that figure therefore differ in history groups, normalization, seed, batch size, and training horizon. D.1

N ORMALIZED CLIPPING RESIDUAL

We measure the clipping residual on the same m = 512 fixed held-out query records at each (i) checkpoint. Let HW be record i’s query for Muon block W : a history average for IM and a P (i) single current-state evaluation for the control. Write Gavg,W = m−1 i HW and Gclip,W = P (i) m−1 i λi HW . Clipping factors come from the complete raw query at C = 1. We report P 2
1/2 W ∈W ∥RW ∥F RW = Gclip,W − λ̄Gavg,W , ηR = P . (20) 2
1/2 λ̄ W ∈W ∥Gavg,W ∥F The norms aggregate the same 48 Muon matrices. Each arm uses its own mean clipping factor λ̄ and unclipped query mean in the denominator. These measurements precede Gaussian noise and outer momentum.

Normalized residual ηR (dimensionless)

Historical averaging

Current-only

0.20 0.15 0.10 0.05 0.00

10

50

100

250

500

1000

2000

Pre-update index (log scale)

Figure 4: Normalized clipping residual ηR on E2E for historical averaging (K = 4, γ = 0.9) and a single current-state evaluation at the same model states. Points show means over training seeds 42–44; bars show sample standard deviation, not confidence intervals. All seven pre-update checkpoints are shown. Clipping uses the complete raw query at C = 1; the residual metric covers Muon coordinates. Historical averaging gives a smaller ηR in all 21 seed–checkpoint pairs (Figure 4). We compute the percentage reduction within each seed as 100(1 − ηR,IM /ηR,current ), then average over seeds. These means range from 24.99% to 46.10% across checkpoints.

E

P RIVACY AND IMPLEMENTATION CONDITIONS

Fix a public initialization and public values of the sampling rate q ∈ (0, 1], clipping radius C > 0, noise multiplier σDP > 0, release denominator b > 0, and update count T . Current and stored model states and optimizer states must be computed only from preceding noisy releases and public settings. At each update, sample every record independently with probability q, independently of previous sampling draws. Sampling decisions remain internal. The accountant must bound adaptive composition under add/remove adjacency. 15

preprint

e0 , . . . , G e t−1 ) for the preceding noisy releases, with Complete-query mechanism. Write Tt−1 = (G T−1 empty. Given these releases, let ut (x; Tt−1 , ω) ∈ Rd be record x’s query over all trainable coordinates. Here ω is fresh data-independent randomness, drawn independently for each record and update and independently of sampling, Gaussian noise, and preceding releases. Historical evaluations within one record’s query may share randomness. The query may combine different Muon and AdamW histories and apply per-record preprocessing. It must not use unaccounted private statistics from other records. Define
vt (x; Tt−1 , ω) = clipC ut (x; Tt−1 , ω) , ∥vt (x; Tt−1 , ω)∥2 ≤ C. (21) The complete release at step t is et = 1 G b

! X

vt (x; Tt−1 , ωt,x ) + σDP CZt

,

Zt ∼ N (0, Id ),

(22)

x∈St

where b > 0 is fixed and public. Gaussian draws are independent across updates and independent of sampling and query randomness. Algorithm 1 displays one Muon block of this release with b = B; the utility implementation uses b = ⌊B⌋. Neither choice changes the privacy bound when the denominator is held fixed across neighboring datasets. Independent Gaussian draws for individual coordinates or blocks implement the same joint release; they are not separately clipped mechanisms. Privacy argument. Consider neighboring datasets D and D′ = D ∪ {x⋆ }. Fix a common value of Tt−1 . With the public initialization and settings fixed, the current model and stored histories are the same in both conditional executions. This does not assert that their unconditional training trajectories coincide. For every common record, the conditional query law in Eq. (21) is the same on both datasets. Couple the sampling and local randomness of these common records. The additional record is selected independently with probability q and contributes a vector of norm at most C. Thus the conditional update is a Poisson-sampled Gaussian mechanism with noise multiplier σDP : the add/remove sensitivity of the unsampled sum is at most C, irrespective of how many historical gradients were used within each query. Division by public b is post-processing. This per-step privacy bound holds uniformly over preceding transcripts. Adaptive composition gives e0 , . . . , G e T −1 ) the (ϵ, δ) guarantee certified by the accountant (Abadi et al., the complete transcript (G 2016; Gopi et al., 2021). Updating the history, applying Muon momentum and Newton–Schulz, and applying AdamW are post-processing of these releases. The final model and any checkpoints computed solely from the transcript inherit the same guarantee (Dwork and Roth, 2014). Calibration and output scope. The privacy analysis holds calibration fixed across neighboring datasets. Any reference dataset size used to choose q, T , δ, or b must therefore be a public design value, not recomputed from the private dataset’s cardinality. The guarantee uses add/remove adjacency; replace-one adjacency has an unsampled sensitivity bound of 2C and requires the corresponding accounting. This is a per-run guarantee for fixed public settings, not a privacy bound for the full experimental search. Private validation, hyperparameter selection, additional training runs, or a privately trained initialization require separate accounting. Releasing raw gradients, sampling decisions, or additional data-dependent diagnostics is not covered by this guarantee. In particular, the geometric probes’ diagnostic releases and non-Poisson training protocol are not covered by the utility runs’ reported budgets. This privacy argument requires no smoothness, gradient-decorrelation, or polar-rank assumptions.

F

AVERAGING AND POLAR ERROR

Proposition 1 in Section 3 separates the clipped mean into a common rescaling and a residual caused by unequal clipping. Here we use this decomposition to bound polar error relative to the current population gradient, accounting for sampling error and the use of historical gradients. The analysis concerns the clipped mean before Gaussian noise and outer momentum, not the final Muon update. 16

preprint

Setup. Fix model states θ0 , . . . , θK and one Muon block W ∈ Rr×c , where θ0 is the current state. For an example z, let Gj (z) be its raw gradient for this block at state j, and define µj = EGj (z) and ξj (z) = Gj (z) − µj . Expectations include any local evaluation randomness, with the same conditional law given z at every state. Draw m ≥ 2 examples independently and reuse each example across states. Each example and all its evaluation draws form one independent, identically distributed tuple; the states and weights are fixed independently of these tuples. PK With fixed weights αj ≥ 0, j=0 αj = 1, shared across the block’s coordinates, set H (i) =

K X

m

αj Gj (zi ),

N=

j=0

1 X λi H (i) , m i=1

SK =

K X

αj2 .

(23)

j=0

As in Section 3, λi > 0 is computed by flat P clipping of example i’s Pcomplete raw query, not by clipping this block separately. Write λ̄ = m−1 i λi and σλ2 = m−1 i (λi − λ̄)2 . The current-only query has K = 0 and S0 = 1; equal weights give SK = 1/(K + 1). We assume, for fixed constants τW ≥ 0, ρ ∈ [0, 1], and cλ ≥ 0, 2

2 1. E ∥ξj ∥F ≤ τW for every j; 2 2. E ⟨ξj , ξℓ ⟩F ≤ ρτW for j ̸= ℓ, evaluated on the same example;

3. σλ ≤ cλ λ̄ almost surely. Here τW bounds fluctuations of block W , whereas τu in Section 4.2 bounds fluctuations of the full query. The blockwise covariance bound ρ here need not equal the full-query value used there. Assumption (1) is the blockwise version of Assumption 3.1 in Xiao et al. (2023). Assumption (2) allows temporal correlation: independent centered fluctuations satisfy ρ = 0, but independence is not required. The choice ρ = 1 follows from (1) by Cauchy–Schwarz and does not require identical fluctuations. Assumption (3) and the full-rank reference condition below are additional assumptions, not results of Xiao et al. (2023). Polar perturbation. For square nonsingular inputs, the following inequality follows from classical polar-factor perturbation bounds (Li, 1995). The proof below also covers rectangular matrices and rank-deficient X using the partial-polar convention in Section 2.4. Lemma 5 (Polar perturbation with a full-rank reference). Let T, X ∈ Rr×c with rank(T ) = min(r, c). Then 2 ∥X − T ∥F dpolar (X, T ) ≤ . (24) σmin (T ) Here σmin (T ) > 0 is the smallest of the min(r, c) singular values of T . Proof. First suppose r ≥ c. Let Q = polar(T ), U = polar(X), and write T = QP with P = (T ⊤ T )1/2 ⪰ sIc , where s = σmin (T ) > 0. Since Q⊤ Q = Ic and U ⊤ U ⪯ Ic , A = Ic − 12 (Q⊤ U + U ⊤ Q) ⪰ 0,

2

∥Q − U ∥F ≤ 2 tr(A).

The polar factor U , including when X is rank deficient, maximizes ⟨X, V ⟩F over ∥V ∥op ≤ 1. As Q is feasible, ⟨X, Q − U ⟩F ≤ 0. Hence s 2 ∥Q − U ∥F ≤ tr(P A) = ⟨T, Q − U ⟩F ≤ ⟨T − X, Q − U ⟩F ≤ ∥T − X∥F ∥Q − U ∥F . 2 Dividing by ∥Q − U ∥F proves the claim for r ≥ c when it is nonzero; the zero case is immediate. For r < c, apply this result to T ⊤ and X ⊤ . Transposition preserves Frobenius norms and singular values, and polar(X ⊤ ) = polar(X)⊤ . Theorem 6 (Polar-error bound under averaging). Under Assumptions (1)–(3), suppose µ0 ∈ Rr×c has rank min(r, c). Define sK,W = τW

p

ρ + (1 − ρ)SK ,

DK =

K X j=0

17

αj µj − µ0

. F

(25)

preprint

Then 2 E dpolar (N, µ0 ) ≤ σmin (µ0 )

sK,W DK + √ + cλ m

r

m−1 sK,W m

! .

(26)

The three terms account for staleness, sampling error, and unequal clipping, respectively. No rank condition is imposed on N . Proof. Let Yi =

P

P = m−1 i Yi . Applying Proposition 1 with Gi = H (i) gives   X N R − µ0 =  αj µj − µ0  + Ȳ + . (27) λ̄ λ̄ j

j αj ξj (zi ) and Ȳ

The first term has norm DK . For the second, Assumptions (1)–(2) and the nonnegative weights imply X X
2 2 2 2 E ∥Yi ∥F ≤ τW αj2 + ρτW αj αℓ = τW SK + ρ(1 − SK ) = s2K,W . j

j̸=ℓ

√ 2 The Yi are independent with mean zero, so E Ȳ F ≤ sK,W / m. For the last term, write σH = P 2 −1 (i) m − H̄ F . Proposition 1 and (3) give i H r r q ∥R∥F m−1 m−1 2 2 E ∥Yi ∥F ≤ sK,W . ≤ cλ σH , EσH ≤ EσH = m m λ̄ Finally, polar(N ) = polar(N/λ̄) because λ̄ > 0. Taking norms and expectations in Eq. (27) and applying Lemma 5 proves Eq. (26). When the bound improves. Without averaging, D0 = 0 and s0,W = τW . For τW > 0, averaging gives sK,W < τW when ρ < 1 and at least two weights are nonzero (SK < 1). At the same current state and batch size, use one p constant cλ satisfying (3) for both the averaged and current-only queries. With am = m−1/2 + cλ (m − 1)/m, the right-hand side of Eq. (26) is no larger than the current-only bound whenever DK ≤ am (τW − sK,W ). (28) Strict inequality gives a strictly smaller bound. This compares upper bounds, not actual errors. When ρ = 1, the variance terms in this bound do not decrease. If each per-example loss is β-smooth in the full parameter vector, uniformly over the local randomness, then the block mean is β-Lipschitz and DK ≤ β

K X

αj ∥θj − θ0 ∥2 .

(29)

j=0

This is the smoothness condition used by Xiao et al. (2023), with the common conditional evaluationrandomness law specified above. Nearby states control staleness, but do not by themselves imply temporal decorrelation. Proof of PropositionP3. Use the full-query with τu from P gradients and moment assumptions P Section 4.2. Let ui = j αj gj (zi ), Yi,u = j αj ξj,u (zi ), and σu2 = m−1 i ∥ui − ū∥22 . The map n 7→ min{1, C/n} is 1/C-Lipschitz: its derivative is 0 for n < C and −C/n2 ∈ [−1/C, 0) for n > C. The variance of the λi is at most their second moment about λ(∥ū∥), and | ∥ui ∥ − ∥ū∥ | ≤ ∥ui − ū∥, so σλ ≤ σu /C. Proposition 1 gives ∥RW ∥F ≤ σλ σH,W , and σH,W ≤ σu because a block is a sub-collection of coordinates. Hence ∥RW ∥F ≤ σu2 /C. The full-query moment bounds give P 2 2 E∥Yi,u ∥2 ≤ sK,u . Since the means j αj µj,u cancel in ui − ū, the same sample-variance identity m−1 2 2 used in Theorem 6 yields E σu2 = m−1 m E∥Y1,u ∥2 ≤ m sK,u . 2 If the full-query moment bounds in Section 4.2 hold with equality and τu > 0, then E σu2 = m−1 m sK,u exactly. The ratio of the expected empirical variances for averaged and current-only queries is then ρ + (1 − ρ)SK .

18

Related documents

• Inner Momentum for Differentially Private Muon#1187052
• Muon Under Gradient Noise and the Limits of Orthogonalization Near Optima#1117574
• On the Convergence of Muon and Beyond#1143622
• Muon as a Residual Connection#329079
• DMuon: Efficient Distributed Muon Training with Near-Adam Overhead#310782