Differential privacy training improves muon accuracy with inner momentum averaging
Inner Momentum for Differentially Private Muon
Machine LearningCryptography and Security
Summary
Training AI models with privacy guarantees involves adding noise and limiting how much each example can influence the model. The authors found that a specific training method called Muon suffers when each gradient (direction of learning) is clipped differently, causing distortions. To fix this, they average gradients over recent models before clipping, which reduces distortion and helps keep the model’s updates more accurate. Their approach, tested on language models, showed better performance under privacy constraints than previous methods.
What this means in practice
- •For machine learning engineers: Enhance privacy-preserving language model fine-tuning by reducing gradient clipping distortion using inner momentum averaging.
- •For data privacy teams: Improve accuracy of differentially private model training pipelines that rely on clipping per-example gradients without compromising privacy guarantees.
Authors
Bishnu Bhusal, Minh Vu, Ben Southworth, Geigh Zollicoffer, Rohit Chadha, Manish Bhattarai
Abstract
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 UV^T 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 <= sigma_lambda sigma_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 epsilon in {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.