Towards Hyperparameter Transfer for Differentially Private Optimization

AI summary

Overview

Research area: Differentially private (DP) deep learning, specifically hyperparameter transfer and scaling laws for private optimization. Technical level: Advanced — the paper leans on random matrix theory, spectral norm analysis, and the μP (Maximal Update Parametrization) literature, though the resulting scaling rules are simple formulas practitioners can apply directly. Scope: The paper derives and empirically validates layer-wise learning-rate scaling rules that let hyperparameters tuned on a small "proxy" model transfer zero-shot to much larger models trained with DP-SGD, DP-Adam, or DP-Muon.

What This Paper Is About

Maximal Update Parametrization (μP) is a well-known technique that lets you tune a learning rate on a small model and reuse it unchanged on a large one — but the authors show it breaks down once differential privacy is switched on, because the injected Gaussian noise has very different spectral structure from the low-rank gradient signal. The paper's goal is to work out exactly what extra condition must hold for hyperparameters to transfer under DP, and to turn that condition into concrete scaling rules for three optimizers.

Key Contributions

1. A spectral characterization of the DP regime. Using random matrix theory, the authors contrast the full-rank, high-dimensional Gaussian noise added for privacy against the stable low-rank structure of the gradient signal, and show this mismatch is what invalidates standard μP scaling.
2. A signal-to-noise ratio (SNR) condition for transfer. They derive a layer-wise SNR condition that, combined with the μP spectral condition, is both necessary and sufficient for learning-rate transfer across widths, and they instantiate it for DP-SGD, DP-Adam, and DP-Muon.
3. Concrete scaling rules and a clipping-calibration procedure. Table 1 gives closed-form layer-wise learning-rate scaling for each optimizer under three settings (non-DP, clipping-only, and full DP), and Equation (8) gives a closed-form way to set the proxy model's noise multiplier and the target model's per-layer clipping parameters so the SNR matches layer by layer.
4. An empirical zero-shot transfer study. Coordinate checks plus transfer experiments on MLP, ViT, and GPT-2 architectures, reported to match or exceed individually tuned baselines while cutting tuning cost.

Main Findings

• Standard μP fails under DP. Figure 2 contrasts two regimes on a 5-layer MLP trained on CIFAR-10: under clipping-only (no noise) the loss curves align across widths, but scaling by the spectral norm of the noisy update leaves the curves unaligned.
• The failure is a vanishing signal, not just mis-scaling. Lemma 3.5 shows that in the noise-dominated regime of SGD, after applying the μP scaling the signal contribution to the weight update scales as Θ(1/√n) — it disappears as width grows.
• Noise norm is governed by a random-matrix expression. The spectral norm of the injected noise is Θ((σ_DP/B)(√n_l + √n_{l−1})), which is why the SNR condition must be stated layer by layer rather than globally.
• Per-optimizer clipping behavior differs. Proposition 3.3 gives the spectral norm of the clipped (noiseless) update as Θ(1) for SGD and Muon but Θ(√(n_l n_{l−1})) for Adam — this is why Adam needs a different learning-rate rule in the non-DP and clipping-only columns.
• The transfer rule is optimizer-dependent. In the DP column of Table 1, SGD's rule divides by (σ_DP/B)(√n_l + √n_{l−1}), Adam's drops the σ_DP/B factor, and Muon's stays at √(n_l/n_{l−1}) — the same as its non-DP rule.
• SNR matching makes transfer both necessary and sufficient (Proposition 3.6). Because the learning rate scales signal and noise by the same factor, its ratio is learning-rate-independent; if proxy and target SNR disagree, no single learning rate can fix it.
• Input and output layers need special handling. Since embedding size and output dimension are typically fixed while hidden widths grow, the authors make the per-layer clipping coefficient the adjustable knob (Equation 8) rather than relying on noise scaling alone.
• Empirical transfer succeeds across architectures. ViT pretraining on CIFAR-100 shows larger widths reaching lower training loss under a shared learning rate; GPT-2 pretraining on the Shakespeare Character dataset shows consistent loss-curve rankings across three scales (76.0M, 227.3M, and 787.8M parameters).
• Large cost savings. Tuning on the GPT-2 proxy (5 heads, d_model=320, 76.0M parameters) instead of the target (36 layers, 20 heads, d_model=1280, 787.8M parameters) is reported as a 90.35% reduction in theoretical computational cost, and cuts wall-clock time per run from roughly 30 minutes to 11 minutes on a single NVIDIA A100 GPU.
• A practical caveat for Muon. The authors deviate from Muon's usual Adam updates on input/output layers, using Muon or SGD instead, because the standard Adam update fails to exhibit hyperparameter transfer there even without privacy.

Methodology in Plain English

The authors start from the μP principle that each layer's update should stay at a fixed scale as the network gets wider, which pins the learning rate to the ratio of the desired update norm to the actual gradient norm. They then ask what that gradient norm looks like once you add per-sample gradient clipping and Gaussian noise.

For clipping alone, they reason that a clipped per-sample gradient is a rank-1 matrix whose spectral and Frobenius norms coincide, so the aggregate clipped gradient has a norm they can read off per optimizer. For the noisy case, they invoke the fact that a large random Gaussian matrix has a spectral norm on the order of the square root of its dimensions — so the noise, not the signal, sets the overall update scale.

That observation drives the central move: since the learning rate multiplies both signal and noise equally, it cannot change their ratio, so the ratio itself must be matched between proxy and target. The authors enforce this by treating the per-layer clipping coefficient as the free variable and solving a small system of equations for the proxy noise multiplier and the target clipping parameters. Finally, they verify implementation correctness with the "coordinate check" procedure from prior μP work — tracking the ℓ1 norm of specific feature coordinates over the first training steps across a range of widths — and then run transfer experiments.

Why This Matters

Hyperparameter tuning is one of the most expensive parts of training large models, and DP makes it worse: noise injection and per-sample clipping interact in ways that invalidate the tuning intuitions carried over from non-private training. Showing that a principled scaling rule restores zero-shot transfer means privacy-preserving training no longer has to pay a full tuning bill at target scale.

Real-world applications:

• Private fine-tuning of language models on sensitive text (medical records, legal documents, user conversations), where re-tuning a large model under a privacy budget is often prohibitive.
• Privacy-preserving image classification in domains like healthcare imaging or on-device photo processing, where ViT-style models are increasingly used.
• Federated and distributed training, where the paper's choice of layer-wise clipping is motivated as the most efficient and scalable clipping style in distributed settings.
• Regulated industries (finance, healthcare, government) that must make formal (ε, δ) privacy guarantees and therefore need reliable, auditable training recipes.

Industry relevance: the cost figures reported — a 90.35% theoretical compute reduction and a wall-clock drop from about 30 minutes to 11 minutes per run — target exactly the bottleneck that makes DP training unattractive to practitioners. The fact that the rules cover DP-SGD, DP-Adam, and DP-Muon means they apply to the optimizers teams already use.

Future Directions

• Depth scaling. The authors explicitly flag that they do not address depth: noise accumulating across successive layers tends to violate the layer-wise SNR condition, so a depth-scaling mechanism — or a formal impossibility result — remains open.
• The transition regime. The analysis assumes noise dominates the gradient. With extremely large batch sizes, as used in large-scale private pretraining, the noise-to-signal ratio shrinks and the rules collapse to the clipping-only column; whether the scaling laws need adjustment in between is untested.
• Privacy accounting for the tuning step itself. The reported ε = 2 budget covers target training and excludes the proxy search. The paper points to private hyperparameter selection or public-data tuning as workarounds and gives an end-to-end accounting in an appendix, but integrating the search cost into a single guarantee is a natural next step.
• Optimizer-specific gaps. The Muon experiments required replacing Adam updates on input/output layers, hinting that the interaction between adaptive methods and transfer deserves closer study.

Target Audience

Researchers and engineers working on differentially private training at scale, especially those already familiar with μP or with parameter-scaling / transfer literature. It is also relevant to practitioners who need to train or fine-tune large models under a formal privacy budget and want a defensible alternative to repeated private hyperparameter search. Readers without a background in spectral norm analysis or random matrix theory will find the theory sections demanding, but Table 1 and the Equation (8) procedure are usable on their own.

Authors’ abstract

Scaling laws and hyperparameter tuning remain central challenges in training large-scale models, particularly under differential privacy (DP), where optimization is further complicated by noise injection and per-sample gradient clipping. While Maximal Update Parametrization ($μ$P) allows for zero-shot hyperparameter transfer in standard training, we show that it fails under DP constraints due to the distinct spectral properties of high-dimensional noise versus low-rank gradients. In this work, we introduce DP-$μ$P, a theoretically grounded framework that extends $μ$P to the privacy-preserving setting. By analyzing the interaction between gradient clipping, noise injection, and feature learning, we derive a Signal-to-Noise Ratio (SNR) condition that governs layer-wise learning rate scaling. Combining this with our proposed spectral scaling formulation, we provide specific scaling rules for DP-SGD, DP-Adam, and DP-Muon optimizers. Empirically, we demonstrate that DP-$μ$P enables effective zero-shot transfer of learning rates across model widths on diverse architectures (MLP, ViT, GPT-2), matching or exceeding the performance of individually tuned baselines while significantly reducing computational costs.