Inner Momentum for Differentially Private Muon

1 INTRODUCTION

• Differentially private SGD clips per-example gradients before averaging and adds calibrated Gaussian noise, while unequal clipping can change the average’s direction.Clipping limits each record’s influence; noise perturbs the averaged gradient.Source: §1 · p. 1
• Muon’s update depends on singular vectors, which unequal per-example clipping can change by reweighting gradients in the batch average.Newton–Schulz preserves the polar factor of its input, so it cannot correct a clipping-induced change already present there.Source: §1 · p. 1
• DP-Muon-IM averages each sampled example’s Muon gradients at the current and recent private model states before clipping.The method adapts the inner-momentum approach of Xiao et al. (2023).Source: §1 · pp. 1–2, Figure 1 · p. 2
• DP-Muon-IM improves BLEU and ROUGE-L over DP-Muon on E2E and DART while keeping the baseline privacy accounting unchanged.The method uses the same clipping radius, sampling procedure, and release schedule as the baseline.Source: §1 · p. 2
• The paper decomposes the clipped average into a common-rescaling term and a residual governed by variation in gradients and clipping factors.The stated residual bound is ∥R∥F ≤ σλσG.Source: §1 · p. 2

2.1 PROBLEM SETUP AND NOTATION

• The training setup uses a private dataset, model parameters θ, example loss ℓ(θ; x_i), and minibatch S_t with realized size m.Muon operates on matrix blocks for hidden-layer parameters; the gradient norm for clipping treats all parameter blocks as one vector.Source: §2.1 · pp. 2–3

2.2 DIFFERENTIAL PRIVACY AND PRIVATE OPTIMIZATION

• The private optimization mechanism clips each example’s complete gradient to radius C, averages the clipped contributions, and adds calibrated Gaussian noise.Clipping controls sensitivity, noise provides privacy, and subsequent optimizer operations are post-processing.Source: §2.2 · p. 3

2.3 MUON AND DP-MUON

• DP-Muon forms a momentum buffer and Nesterov direction from the private release, then applies Newton–Schulz to produce the orthogonalized update.Outer momentum acts on the private release; inner momentum, introduced later, acts before clipping.Source: §2.3 · p. 3

2.4 METRICS

• Cosine error compares matrix entries, whereas polar error compares the basis-invariant partial isometry used by Muon.Matrices can differ in cosine similarity while sharing a polar factor.Source: §2.4 · pp. 3–4

3 GEOMETRIC EFFECTS OF GRADIENT CLIPPING

• Unequal clipping reweights per-example gradients, potentially changing the singular-vector bases and polar factor that determine Muon’s update.The section’s analysis separates this clipping effect from later noise and outer-momentum operations, which may also change the discrepancy.Source: §3 · p. 4, §3.2 · p. 5

3.1 THE CLIPPING RESIDUAL

• The clipped mean decomposes into a common-rescaling term and residual R, which captures distortion from unequal per-example clipping factors.The common-rescaling term preserves the mean’s singular vectors and polar factor; the residual vanishes when clipping factors or per-example matrices are equal, and its bound is a worst-case envelope rather than a tight prediction of polar error.Source: §3.1 · p. 4

3.2 WHAT NEWTON–SCHULZ PRESERVES

• Muon’s Newton–Schulz update preserves its input’s polar factor, so clipping-induced changes to that factor are not undone by orthogonalization.The result holds in exact arithmetic under the stated coefficient conditions and for finite iterations.Source: §3 · p. 4, §3.2 · p. 5

4.1 AVERAGING ONE EXAMPLE ACROSS MODEL STATES

• DP-Muon-IM averages each example’s gradients at the current and recent model states before clipping, distinguishing inner momentum from Muon’s outer momentum.Normalized weights keep the averaged query on the scale of a single gradient, preserving the clipping radius’s meaning.Source: §4 · p. 5, §4.1 · pp. 5–6

4.2 AVERAGING SHRINKS THE RESIDUAL ENVELOPE

• The residual envelope shrinks by ρ + (1 − ρ)SK relative to the current-only bound when ρ < 1 and at least two averaging weights are nonzero.This bounds R from above; it does not establish that the realized residual is smaller.Source: §4 · p. 5, §4.2 · p. 6
• The reduction requires imperfectly correlated gradient fluctuations across states; older states can increase averaged-gradient bias through a drift term that grows with K and γ.Kmax and γ control this trade-off.Source: §4.2 · p. 6

4.3 TRAINING PROCEDURE

• Per-step cost grows linearly with Kmax because each sampled example needs up to Kt + 1 gradient evaluations.The history stores at most Kmax pre-update model states.Source: §4.3 · p. 7
• History averaging adds no private releases, and the baseline privacy accounting applies when the sampling, clipping, noise, denominator, and release count remain unchanged.Each sampled example contributes one query clipped at C, followed by one Gaussian-noised batch sum.Source: §4.3 · p. 7

5 EXPERIMENTS

• The utility runs fine-tune GPT-2-small on E2E and DART for ten epochs, targeting ϵ ∈{1, 2, 4, 8} and evaluating seeds 37–39.Each per-example query is clipped at C = 1 before noise; configurations are selected by validation loss on seed 37.Source: §5 · p. 7

5.1 COMPARISON WITH EXISTING METHODS

• Table 1 compares Muon-based optimizers on E2E at target ϵ = 8, including DP-Muon, DP-Muon-BC, and Kalman-filter variants.DP-Muon-BC corrects privacy-noise bias after Newton–Schulz; the KF variants apply DiSK’s simplified Kalman filter before optimizer updates.Source: Table 1 · p. 8, §5.1 · p. 8

5.2 UTILITY ACROSS PRIVACY BUDGETS

• DP-Muon-IM improves BLEU and ROUGE-L over DP-Muon in all 24 seed-matched comparisons across two datasets, four privacy budgets, and three seeds.Mean BLEU gains are 0.60–2.57 points on E2E and 0.50–1.06 on DART.Source: Table 2 · p. 8, §5.2 · p. 8

5.3 EFFECT OF HISTORY LENGTH

• Validation BLEU reaches 67.76 at K = 1 in 59.0 minutes, versus 67.82 at K = 4 in 123.7 minutes; gains across K are nonmonotonic.Peak GPU allocation rises from 26.70 GiB without history to 49.85 GiB with history.Source: §5.3 · p. 8, Figure 2 · p. 9
• Before optimizer update step 10, history reduces polar error by 1.76% at K = 1 and 3.71% at K = 4 relative to the current-only estimate.The Muon-only-history comparison uses a shared unclipped current-gradient reference before noise and outer momentum.Source: §5.3 · pp. 8–9, Figure 2 · p. 9
• With K = 4 and γ = 0.9, historical averaging lowers normalized clipping residual in all 21 seed–checkpoint comparisons, with mean paired reductions of 25.0%–46.1%.These raw-query diagnostics precede noise and outer momentum.Source: §5.3 · p. 9

• Prior analyses connect clipping bias to changes in aggregate gradients and per-example sampling noise, while inner–outer momentum uses pre-clipping averaging.Other approaches include adaptive clipping, error feedback, and per-sample momentum with post-release filtering.Source: §6 · p. 9
• This work adapts pre-clipping averaging to Muon’s matrix updates and studies the resulting clipping residual.Source: §6 · p. 9
• Private Muon research studies Gaussian perturbations under the transform, while separate work analyzes optimization with finite Newton–Schulz steps.The paper’s related-work discussion positions its geometry analysis alongside these studies.Source: §6 · p. 9

• Unequal clipping can change the batch mean’s polar factor, while the stated Newton–Schulz iteration preserves its immediate input’s polar factor in exact arithmetic.Source: §7 · p. 9
• With K = 4 and γ = 0.9, IM reduces polar error by 2.4–3.7% across measured checkpoints, before Gaussian noise and outer momentum.The probe does not measure the final Muon update.Source: §7 · p. 9
• On E2E, K = 1 comes within 0.06 validation BLEU points of K = 4 in less than half the training time.The paper notes that extra gradient evaluations and stored model states remain costs to address in larger-scale experiments.Source: §7 · p. 9

8 CONCLUSION

• 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, it achieves higher BLEU and ROUGE-L than DP-Muon on E2E and DART at target ϵ ∈{1, 2, 4, 8}.Source: §8 · p. 10

AI USE STATEMENT

• The authors report using language models for language editing, LaTeX support, bibliographic checks, and mathematical consistency, and taking responsibility for the paper.Source: AI USE STATEMENT · p. 10

A.1 RESIDUAL IDENTITY AND BOUND

• The residual-bound proof uses the means and dispersions defined in Eqs. (10) and (11), with the triangle inequality and Cauchy–Schwarz.Source: §A.1 · p. 12

A.2 NEWTON–SCHULZ POLAR-FACTOR PRESERVATION

• The appendix states that the Newton–Schulz iteration preserves its input’s polar factor in exact arithmetic.Its proof keeps the singular-vector factors fixed while the positive diagonal singular values change.Source: §A.2 · p. 12

B ILLUSTRATIVE EXAMPLE OF CLIPPING DISTORTION

• In the two-example construction, the clipping bound is attained, although it need not be tight in general.The example also shows that clipping can change the aggregate polar factor even when it preserves each example’s polar factor.Source: B ILLUSTRATIVE EXAMPLE OF CLIPPING… · p. 12

C UTILITY EXPERIMENT DETAILS

• At ε∈{1, 2, 4}, E2E uses Muon/AdamW histories (4, 0.5)/(3, 0.5), while DART uses (3, 0.7)/(3, 0.5); E2E ε=8 uses (4, 2) with both decays 0.5.The ε=8 learning rates remain 0.002/0.0035 from the earlier validation search without retuning.Source: C UTILITY EXPERIMENT DETAILS · p. 13
• At E2E ε=8, DP-Muon-IM averages 65.34 ± 0.39 BLEU and 67.60 ± 0.17 ROUGE-L versus DP-Muon’s 64.74 ± 0.05 and 67.31 ± 0.25.The per-metric maxima are 65.65/67.80 for IM and 64.79/67.58 for the baseline, all from seed 39.Source: C UTILITY EXPERIMENT DETAILS · p. 13
• Utility runs use Poisson sampling at q=1/41 for 410 E2E releases and q=1/61 for 610 DART releases, with one Gaussian release per logical update.Clipping uses the full concatenated per-example query; the implementation uses release denominators 1025 and 1027 rather than the nominal 1024.Source: C UTILITY EXPERIMENT DETAILS · p. 13
• DiSK combines base and extrapolated queries before clipping and noise, then filters the private release; its single-seed KF joint ablation uses κ=0.5 and extrapolation coefficient 0.8.The IM setting was selected by validation loss from 30 new trials and one reused anchor.Source: C UTILITY EXPERIMENT DETAILS · pp. 13–14

C.1 OPTIMIZER-GROUP ABLATION

• The E2E ε=4 optimizer-group ablation switches Muon and AdamW histories independently while holding optimizer moments and privacy calibration fixed.The table measures NLL in nats per target token.Source: §C.1 · p. 14, Table 5 · p. 14

D GEOMETRIC ERROR ACROSS TRAINING

• The corrected diagnostic compares IM with a current-only query at the same model state against a shared finite raw-current reference, not independently trained IM and baseline models.It measures per-record mean target-token loss before noise and outer momentum.Source: D GEOMETRIC ERROR ACROSS TRAINING · p. 14
• Polar-error reduction is 3.67% before update 10 and 2.44% before update 2000; both cosine and polar errors are lower in all 21 seed–checkpoint pairs.Figure 3 reports the corrected three-seed polar diagnostic at (K, γ) = (4, 0.9).Source: D GEOMETRIC ERROR ACROSS TRAINING · p. 14, Figure 3 · p. 14
• Cosine-error reductions range from 3.20% to 5.98%; both distances use a finite reference, not a population gradient or the noisy Newton–Schulz update.The diagnostic compares queries before noise and outer momentum.Source: D GEOMETRIC ERROR ACROSS TRAINING · p. 15
• The Figure 2 K-sweep points for K=0–3 are completed early probes, not final-training results, and its panels differ in diagnostic settings.The K=4 point reuses the corrected seed-42 run.Source: D GEOMETRIC ERROR ACROSS TRAINING · p. 15

D.1 NORMALIZED CLIPPING RESIDUAL

• Normalized clipping-residual reduction ranges from 24.99% to 46.10% across checkpoints, with historical averaging yielding a smaller residual in all 21 seed–checkpoint pairs.Figure 4 covers E2E with K=4 and γ=0.9, using three seeds and all seven pre-update checkpoints.Source: §D.1 · p. 15, Figure 4 · p. 15

E PRIVACY AND IMPLEMENTATION CONDITIONS

• The mechanism fixes public q, C, σDP, release denominator, and update count; states must depend only on preceding noisy releases and public settings, with independent Poisson sampling and add/remove adjacency.The sampling decisions remain internal, and the accountant must bound adaptive composition.Source: E PRIVACY AND IMPLEMENTATION CONDITIONS · p. 15
• Each record’s complete query may combine Muon and AdamW histories or per-record preprocessing, but cannot use unaccounted private statistics from other records.Fresh data-independent randomness is drawn for each record and update; historical evaluations within one record’s query may share randomness.Source: E PRIVACY AND IMPLEMENTATION CONDITIONS · p. 16
• The conditional per-step update is a Poisson-sampled Gaussian mechanism with add/remove sensitivity at most C, regardless of how many historical gradients each query uses.Adaptive composition gives the release transcript the accountant-certified (ε, δ) guarantee; optimizer updates and final models derived from it are post-processing.Source: E PRIVACY AND IMPLEMENTATION CONDITIONS · p. 16
• The guarantee is per run with fixed public settings, not for the full experimental search; private validation, additional runs, and data-dependent diagnostics require separate accounting.It uses add/remove adjacency; replace-one adjacency has an unsampled sensitivity bound of 2C and requires corresponding accounting.Source: E PRIVACY AND IMPLEMENTATION CONDITIONS · p. 16

F AVERAGING AND POLAR ERROR

• The polar-error analysis applies to the clipped mean before Gaussian noise and outer momentum, not to the final Muon update.The reference population gradient is required to have rank min(r, c).Source: F AVERAGING AND POLAR ERROR · pp. 16–17
• The analysis reuses each of m ≥ 2 independent examples across fixed current and historical model states, with states and weights independent of the sampled tuples.Expectations include local evaluation randomness with the same conditional law given each example at every state.Source: F AVERAGING AND POLAR ERROR · p. 17
• The bound assumes blockwise fluctuation and cross-state covariance controls; ρ = 1 follows from the fluctuation bound, while independent centered fluctuations satisfy ρ = 0.Nearby states control staleness but do not by themselves imply temporal decorrelation.Source: F AVERAGING AND POLAR ERROR · pp. 17–18
• Theorem 6 bounds polar error with terms for staleness, sampling error, and unequal clipping, under its assumptions and a full-rank current population-gradient reference.No rank condition is imposed on N.Source: F AVERAGING AND POLAR ERROR · pp. 17–18
• Averaging makes the upper bound no larger when ρ < 1 and at least two weights are nonzero; strict inequality yields a strictly smaller bound.This compares upper bounds, not actual errors; when ρ = 1, the variance terms in the bound do not decrease.Source: F AVERAGING AND POLAR ERROR · p. 18
• The expected empirical-variance ratio for averaged versus current-only queries is ρ + (1 − ρ)SK, which is below one when ρ < 1 and at least two weights are nonzero.With equal weights, SK = 1/(K + 1).Source: F AVERAGING AND POLAR ERROR · pp. 17–18