🤖 machine learning

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

This paper introduces DP-Muon, a differentially private optimizer that integrates matrix-valued momentum and Newton-Schulz orthogonalization with rigorous privacy guarantees, and proposes a bias-corrected variant (DP-MuonBC) that significantly improves utility in private fine-tuning without increasing the privacy budget.

Original authors: Jihwan Kim, Chenglin Fan

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jihwan Kim, Chenglin Fan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to write stories based on sensitive personal diaries. You want the robot to learn the style of the writing without ever memorizing or revealing the specific secrets inside any single diary. This is the challenge of Differential Privacy (DP): teaching a model while protecting individual data points.

This paper introduces a new, smarter way to do this teaching, called DP-Muon, and an even better version called DP-MuonBC. Here is the breakdown using simple analogies.

1. The Problem: The "Noisy Classroom"

Standard methods for teaching AI with privacy (like DP-SGD) work like a noisy classroom.

The Teacher: The AI model.
The Students: The data (diaries).
The Lesson: The teacher asks every student for a hint (a gradient) on how to improve.
The Privacy Trick: To protect privacy, the teacher clips the hints so no single student can shout too loudly (clipping) and then adds static noise to the room (Gaussian noise) so no one can tell exactly what any single student said.

However, standard methods treat the AI's "brain" like a long list of numbers (a vector). The paper argues that for the hidden layers of modern AI, the brain is actually a grid of numbers (a matrix). Treating a grid like a simple list is inefficient and clumsy.

2. The Solution: The "Muon" Optimizer

The paper uses an existing tool called Muon.

The Analogy: Imagine the AI is a hiker trying to find the bottom of a valley.
- Standard Optimizers are like a hiker who only looks at the slope directly under their feet.
- Muon is like a hiker who carries a compass and a map. It looks at the "momentum" (the direction they were already moving) and then performs a special geometric trick (called Newton-Schulz orthogonalization) to ensure they are moving in the most efficient, straight-line direction possible, without getting stuck in local loops.
The Innovation: The paper figures out how to add the "privacy noise" to this specific Muon process without breaking the math.

3. The Big Discovery: Privacy is Free (Mostly)

The authors were worried that adding the complex "compass and map" steps (momentum and orthogonalization) after adding the privacy noise might accidentally leak more secrets.

The Finding: They proved mathematically that these extra steps are just post-processing.
The Metaphor: Imagine you take a photo of a crowd, blur the faces to protect privacy, and then crop the photo to focus on the sky. Cropping the photo doesn't reveal the faces you already blurred. Similarly, the complex math Muon does after the noise is added doesn't cost any extra privacy. The privacy guarantee is determined entirely by the initial "blurring" step.

4. The Hidden Flaw: The "Heat Smoothing" Bias

While the privacy is safe, the authors found a subtle problem with the quality of the learning.

The Analogy: Imagine you are trying to guess the shape of a sculpture, but you are looking at it through a thick, warm fog (the privacy noise).
- The momentum part of the optimizer is like a straight line; the fog just makes the line wobble a bit, but it stays straight.
- The Newton-Schulz part (the geometric trick) is like a curved mirror. When you look at a curved mirror through fog, the reflection doesn't just wobble; it gets distorted and smoothed out. The fog makes the sharp edges of the sculpture look rounder than they really are.
The Consequence: This "smoothing" creates a bias. The AI thinks the sculpture is rounder than it is, so it learns the wrong shape. This happens specifically because the privacy noise interacts with the non-linear math of the Muon optimizer.

5. The Fix: DP-MuonBC (The "De-Fogger")

To fix this, the authors created DP-MuonBC (Bias-Corrected).

The Analogy: To fix the foggy mirror, the AI doesn't just look once. It performs a clever trick:
1. It looks at the sculpture through the fog (the standard noisy step).
2. It then imagines looking at the sculpture through double the fog (by adding a tiny, random, fake "probe" noise).
3. It compares the two views. By mathematically subtracting the "double fog" effect from the "single fog" effect, it can cancel out the distortion and recover the true, sharp shape of the sculpture.
The Result: This correction removes the error caused by the privacy noise without requiring any more privacy protection. It gets a better result for the same amount of privacy budget.

6. The Results: Better Stories, Same Privacy

The team tested this on two real-world tasks (E2E and DART), which involve turning structured data into natural language stories.

The Outcome:
- DP-Muon (the standard version) was already much better at writing stories than the old standard methods (DP-SGD and DP-Adam).
- DP-MuonBC (the bias-corrected version) wrote even better stories, with clearer and more accurate language.
The Catch: They achieved this improvement without spending any extra "privacy budget." They got more utility for free by fixing the math.

Summary

This paper takes a powerful, matrix-based optimizer (Muon), wraps it in privacy protection, and proves that the wrapping doesn't leak secrets. Then, it notices that the wrapping causes a slight distortion (bias) in the learning process and invents a mathematical "eraser" (DP-MuonBC) to remove that distortion, resulting in a smarter, more private AI that learns faster and better.

Technical Summary: DP-Muon

Problem Statement

Differential Privacy (DP) is the standard framework for training models on sensitive data, with DP-SGD serving as the dominant baseline. However, the interaction between per-example gradient clipping, Gaussian noise injection, momentum, and nonlinear orthogonalization in matrix-valued optimizers has not been systematically analyzed. Specifically, Muon, a recent optimizer that updates hidden-layer weights using momentum followed by Newton–Schulz orthogonalization to produce an orthogonalized momentum matrix, lacks a formal privacy analysis. It is unclear how the nonlinear Newton–Schulz map interacts with privacy noise and whether standard DP-SGD accounting templates apply to such matrix-valued updates. Furthermore, existing private optimizers often suffer from bias issues (e.g., DP-Adam collapsing to a rescaled DP-SGD form) when privacy noise is introduced without explicit correction.

Methodology

The paper formulates DP-Muon, a differentially private variant of the Muon optimizer. The core procedure involves:

Per-Matrix Clipping: For each hidden-layer matrix block $W$ , per-example gradients are clipped to a Frobenius norm threshold $C_W$ .
Joint Release: The clipped gradients are averaged over the sampled minibatch (lot), and Gaussian noise is added to this average. Crucially, the noise is calibrated to the sensitivity of the joint release of all matrix blocks in the lot, rather than treating blocks as independently subsampled mechanisms.
Post-Processing: The noisy average undergoes momentum accumulation and a finite-step Newton–Schulz iteration to produce the update direction. These steps are treated as post-processing, meaning they do not incur additional privacy cost.

To address the bias introduced by the nonlinear Newton–Schulz map acting on noisy inputs, the authors propose DP-MuonBC. This variant applies a bias-correcting post-processing step. By analyzing the "frozen clipped-path" model, the authors identify that the privacy noise induces a "heat-smoothing" bias at the output of the Newton–Schulz map (a matrix-valued curvature term). DP-MuonBC cancels the leading $O(\rho^2)$ bias term using an antithetic probe average: it evaluates the Newton–Schulz map at the noisy momentum plus and minus independent Gaussian probes, then extrapolates back to the unsmoothed direction.

Key Contributions

1. Privacy Guarantees

The paper proves that DP-Muon inherits the privacy guarantee of the corresponding same-lot subsampled Gaussian accountant.

Joint Accounting: The privacy analysis accounts for the joint product Gaussian mechanism of all matrix blocks released in a single step, rather than summing independent privacy costs.
Post-Processing Invariance: The momentum update and Newton–Schulz orthogonalization are proven to be differentially private post-processing operations, adding no additional privacy cost to the noise-added gradient release.

2. Optimization Theory

The authors establish a finite-horizon stationarity bound for DP-Muon under per-matrix clipping. The bound separates four distinct error sources:

Optimization Error: Standard terms related to the learning rate and horizon.
Clipping Residual: The bias introduced by clipping gradients (captured by $B_T(C)$ ).
Privacy Noise: Stochastic error from minibatch sampling and Gaussian noise, scaled by momentum averaging.
Newton–Schulz Approximation Error: The error from using a finite number of iterations ( $q$ ) instead of an exact SVD/polar factor, captured by a quality factor $\chi_q$ .

Under a light-tail condition on the gradient norms, the paper shows that the average stationarity measure vanishes as $O(T^{-1/4}\sqrt{\log T})$ with appropriate tuning of the learning rate, momentum, and clipping threshold.

3. Bias Correction (DP-MuonBC)

The paper identifies that the primary bias in DP-Muon arises not in the linear momentum buffer but after the nonlinear Newton–Schulz map, where Gaussian noise induces a heat-smoothing effect.

Mechanism: DP-MuonBC uses antithetic Gaussian probes to estimate and subtract the leading second-order heat-smoothing bias term.
Privacy: Since the correction uses only data-independent probes applied to already privatized gradients, DP-MuonBC satisfies the exact same $(\epsilon, \delta)$ -DP guarantee as DP-Muon.

Experimental Results

The authors evaluated the methods on two table-to-text generation benchmarks, E2E and DART, fine-tuning a GPT-2 model under a target privacy budget of $\epsilon \approx 8$ .

Muon vs. Baselines: DP-Muon significantly outperformed both DP-SGD and DP-Adam. On E2E, DP-Muon improved the evaluation Negative Log-Likelihood (NLL) from 0.4961 (DP-Adam) to 0.4553 and increased BLEU from 18.49 to 32.91.
Bias Correction Impact: DP-MuonBC further improved upon the DP-Muon baseline without increasing the privacy budget. On E2E, it achieved the best NLL (0.4490) and BLEU (34.92). On DART, DP-MuonBC achieved the best results across all three metrics (NLL: 0.4460, BLEU: 7.50, ROUGE-L: 34.42).
Training Dynamics: Muon-based methods consistently reached lower NLL than DP-SGD and DP-Adam throughout training, with DP-MuonBC showing particular gains in the later stages of optimization.

Significance and Claims

The paper claims that Muon-style matrix updates are effective for private fine-tuning, offering a promising alternative to conventional vector-based optimizers in the differentially private regime. The work provides the first systematic analysis of the interaction between matrix orthogonalization and differential privacy.

The significance of DP-MuonBC lies in its ability to improve utility by correcting the specific bias induced by the nonlinear orthogonalization step, all while maintaining the exact same privacy budget as the uncorrected version. The authors frame this as a step toward understanding how to leverage advanced optimizer structures (like matrix-valued updates) in privacy-preserving settings without sacrificing theoretical guarantees.

The paper concludes by noting that extending the theory to fully adaptive multi-block trajectories and evaluating larger models remain natural future directions.