PRISM: Gauge-Invariant Tangent-Space Differentially Private LoRA

Summary

PRISM identifies a fundamental mismatch between DP-SGD's parameterization dependence and LoRA's non-identifiable factorization: since $Z = AB^\top$ admits infinitely many equivalent $(A, B)$ representatives, applying DP-SGD to the factors makes both the clipping threshold and the induced intrinsic noise depend on an arbitrary gauge. The paper works this out carefully, decomposing the failure into gauge-dependent clipping (Issue I), bilinear and gauge-amplified intrinsic noise (Issue II), and adaptive-preconditioning noise amplification (Issue III). PRISM's response is conceptually clean: project per-example gradients onto the tangent space of the rank-$r$ manifold at the current $Z$, clip with a global intrinsic Frobenius norm across all LoRA modules, inject isotropic Gaussian noise in the tangent space using an efficient $O((m+n)r)$ sampler, retract via truncated SVD, and use DP-aware floors to bound preconditioner gain. The main theoretical results — gauge invariance of the perturbation distribution, closed-form intrinsic noise energy $(\sigma C/b)\sqrt{r(m+n-r)}$, standard subsampled Gaussian accounting, a spectral bound on adaptive noise amplification — follow cleanly from this design. Empirically the paper reports DP wins on GLUE8 + Math-10K at $\varepsilon \in \{3, 6\}$, and demonstrates that the story scales to Gemma-2-9B / Gemma-3-12B-pt and to LoRA ranks $\{8, 16, 32\}$.

The conceptual move is genuinely useful: intrinsic mechanisms are the right abstraction for DP on quotient manifolds. However, the empirical evidence is more brittle than the write-up suggests. Reported gains rest on shared hyperparameters and a single seed, with no reported variance; execution of the released repository confirms that only the PRISM optimizer is implemented — none of the five baselines, no factor-space DP-SGD arm — so the diagnostic figures could not be regenerated end-to-end and the "PRISM has the best DP average" ranking is untestable from what was shipped. The one PRISM cell that did complete ($\varepsilon = 6$, 12-task mean) lands close to the paper (0.705 vs. 0.690), which supports the PRISM-side numbers but not the comparative claim on which the paper's headline rests.

Strengths

• Conceptual contribution. Framing factor-space DP-LoRA as a gauge-dependence problem on the rank-$r$ manifold is a genuinely useful shift in perspective; Corollary 2.3 crisply captures the worst case, and Issues I–III organize what breaks.
• Clean intrinsic construction. Tangent-space projection with a global intrinsic clip and isotropic Gaussian noise gives a closed-form effective noise energy $r(m+n-r)$ that is manifestly gauge invariant, with retraction shown first-order exact (Prop. 3.2).
• Efficient sampler. The horizontal-lift construction reduces Gaussian generation from $O(mn)$ to $O((m+n)r)$ per module, keeping PRISM at LoRA-scale FLOPs and — as the reproduced PRISM run also confirms — at essentially unchanged peak memory.
• DP-aware adaptive rule. Explicit floors on the rank-space preconditioner, combined with the tight spectral bound in Theorem 3.5, address a real practical failure mode (naive Adam absorbing DP noise into its second moment).
• Empirical breadth. Two multi-task suites, three Gemma backbones, three LoRA ranks, and ~20 diagnostic figures probing each of the three motivating issues.
• Thorough theory appendix. Quotient-geometry framing, tangent characterization, gauge-invariance proofs, isotropy of projected Gaussians, and PRV-based DP accounting are all laid out in detail.
• PRISM-side numbers are directionally consistent when they can be checked. In the one cell that fully reran, the reproduced 12-task $\varepsilon=6$ average (0.705) sits within the paper's own tolerance band around 0.690, and the reproduced PRISM step time (21.08 s) is in the same territory as Table 5's 18.64 s.

Weaknesses

• Motivating "unbounded amplification" is a worst-case limit. Corollary 2.3 requires the pathological gauge $c \to 0$ or $c \to \infty$; the paper's own Figure 3 sweeps $S_t$ across only ~3.6% of its range from natural optimizer drift, over which amplification is a bounded ~1.5–2×. The Section-2 prose reads as if the unbounded regime describes ordinary optimization; a calibrated statement would fit the evidence better.
• Adaptive-preconditioning motivation is imprecise. The identity "$P \approx \Sigma_\xi$" tacitly assumes a full-matrix preconditioner. Adam uses a diagonal second moment, and under the isotropic DP noise of the paper's own Gaussian mechanism the alleged $P^{-1/2}\Sigma_\xi P^{-1/2} \approx I$ reshaping vanishes trivially. The DP-aware floors that Section 3.3 proposes remain sound, but the motivating argument conflates two distinct regimes.
• No ablation. PRISM bundles three fixes for three issues, but nothing decomposes their contributions. A table separating intrinsic clipping, isotropic tangent noise, and DP-aware floors would let the reader attribute the DP gain to a mechanism rather than to the composite.
• Non-DP PRISM underperforms every alternative except FFA. PRISM sits at 0.737 non-DP vs. LoRA-RITE's 0.782, a ~4.5-point gap dismissed in one sentence. Since the intrinsic geometry is present in both regimes, the utility cost of the retraction / adaptive rule / tangent construction deserves discussion.
• Runtime doubling under-discussed. Table 5 reports a ~2× step-time cost, which is nontrivial for LoRA-scale training. Only a single sentence addresses it, and the main-text limitations section leaves it implicit; the reproduced 21.08 s / step corroborates the order-of-magnitude cost, but the LoRA+/AdamW baseline denominators of Table 5 cannot themselves be regenerated (see below).
• Retraction is nonlinear post-processing. Prop. 3.2's first-order exactness is what the paper needs for DP additivity, but for practical step sizes the $O(\eta^2)$ residual couples PRISM noise into a nonlinear correction on $Z$ that is not controlled by the intrinsic clip. When this matters is not discussed.
• Rank sensitivity is non-monotone at $r=32$. PRISM's total Avg is 0.684 → 0.690 → 0.682 across $r \in \{8,16,32\}$; the dip is directly predicted by Theorem 3.1's $r(m+n-r)$ noise scaling, but the paper does not interpret it.
• Terminology drift in diagnostics. Figures 1–4, 14–20 label the comparator as "Baseline," while the prose alternates between "factor-space DP," "factor-space DP-AdamW," and "DP-AdamW." Without a single consistent label a reader cannot verify that all figures concern the same method.
• Undefined "aggressiveness" axis. Figures 16 and 19 use "preconditioner aggressiveness" as an x-axis without a formal definition, yet the Issue-III interpretation depends on this proxy tracking actual update-noise amplification, and the paper uses a different proxy for Adam-like diagonals than for low-rank preconditioners.
• Post-processing immunity depends on a boundary condition that is never stated. Post-processing DP requires that the moments $m_{A,\ell}, m_{B,\ell}, V_{A,\ell}, V_{B,\ell}$ be updated only from the sanitized $(\Delta A^{\rm dp}_\ell, \Delta B^{\rm dp}_\ell)$. Algorithm 1 implies this, but a reimplementer could easily accumulate over raw per-example gradients and silently violate DP.
• Local math / notation issues. The rank-space noise moments carry a $m_\ell/(m_\ell - r)$ prefactor on the $A$-side but not on the $B$-side, without a main-text justification; the Sylvester optimality condition in Prop. A.11 elides the differentiation step so the RHS sign cannot be checked; the bilinear noise "scales as $mnr\tau^4$" describes $\mathbb{E}\|\cdot\|_F^2$ but is stated in $E_Z$-language despite Definition 2.1 defining $E_Z$ as a square-root quantity; Figure 3's linear fit to a full-order measurement using a first-order predictor implies the bilinear term is empirically negligible in that regime, but this is never stated. The $(m+n)r$-dim sampler landing on an $r(m+n-r)$-dim subspace also deserves a one-sentence note.
• "Amplification can even decrease" (Figure 20) is stronger than the theory. Theorem 3.5 gives a saturating $1/\sqrt{\lambda}$ bound; a directional "decreases" claim needs a scatter that actually trends downward, not one that merely stays flat while the baseline rises.

Reproducibility & code

Executed replication of github.com/osu-srml/PRISM-DP-LoRA distinguishes what reproduces, what is available-but-incomplete, and what is unrunnable from what was shipped.

• Headline aggregates (partial). Of 17 planned jobs the campaign completed 2, both the PRISM arm at $\varepsilon = 6$ on Gemma-3-4B-pt. Recomputed 12-task mean is 0.705, within the paper's own ~0.05 tolerance of 0.690, so the PRISM-side headline of Table 2 is directionally supported. The $\varepsilon = 3$ arm never finished (three retries; both eps=3 output directories empty), and the "best average" ranking cannot be checked in either budget because no baseline optimizer exists in the codebase.
• Baseline optimizers are absent. src/prism_cli/trainers.py:64 hard-codes self.method='prism' and train_eval.py's parser has no --method flag, so FFA, LoRA-RITE, LoRA+, LAMB and a factor-space DP-AdamW baseline are not selectable and are not implemented anywhere in src/prism_cli/optim/. LLM-Adapters/finetune.py's lora_rite optimizer additionally imports a module absent from the tree. Every baseline number the paper reports is therefore its own printed value; the "best DP average across 12 tasks" claim (Tables 2, 3, 4) is untestable from the release. README.md documents python train_eval.py --dataset math10k --method adam as a working baseline command; that command fails on the missing flag and, once patched, silently reruns PRISM.
• Diagnostic figures (unrunnable). Figures 1–4 and 12–20, plus the Appendix C.1–C.3 sweeps, need a factor-space DP baseline, a gauge-rescaled-init driver ($c \in \{0.25, 0.5, 1, 2, 4\}$ and the $\log_{10} c \in \mathrm{linspace}(-3, 3, 61)$ / 64-MC grid), a $\sigma$-sweep for $\varepsilon \in \{1.5, 3, 6, 12\}$, per-step $S_t$ / $\|\Delta Z_t\|_F$ / preconditioned-noise trackers, and plotting code. None of these are in the repository. The replication agent had to hand-author analysis/diagnostics.py, analysis/aggregate.py, analysis/profile_step.py, and the --gauge_c mechanism from scratch; even then only 1 of 5 gauge arms started (c=0.25, stopped at step 41 of 300), and not a single figure or panel of Figures 1–4 or 12–20 was produced anywhere.
• Runtime / memory profiling (runnable-but-fragile). No torch.cuda.Event, time.perf_counter, or torch.cuda.max_memory_allocated in the training loop, no profiling harness. Using the agent-authored profiler, PRISM measured 21.08 s / step and 20 742 MB peak memory on Math-10K at $r=16$, $\varepsilon=6$; the PRISM row of Table 5 is thus in the same ballpark as the paper's 18.64 s and 20 964 MB. But the LoRA+ / AdamW comparison rows in Table 5 cannot be produced from the release, so the "roughly doubles step time" ratio is one-sided, and the hardware differs (agent used A100-80GB, paper reports A100-40GB), so peak-memory headroom is not directly comparable.
• Single-seed everywhere. Table 7 pins seed 42 for every method and DP budget; no multi-seed orchestration script exists. The Math-10K PRISM branch additionally bypasses the loader RNG generator (trainers.py:154), risking further nondeterminism at a fixed seed. Table 2's DP-average margins are as small as 0.015, well within seed-to-seed variance typical of DP-LoRA at Gemma-3-4B scale.
• Silent cross-contamination risk in evaluation. LLM-Adapters/evaluate.py names its generation dump by (model, adapter, dataset) only, dropping $\varepsilon$, seed, rank, and privacy mode; src/prism_cli/eval_math.py reads accuracies back from that same path. Two runs differing only in $\varepsilon$ (or seed, or rank) silently overwrite and re-read each other's generations with no error — a defect caught only by code inspection. This is exactly the multi-configuration sweep Tables 2–4 and the appendix require, and a reproducer running the sweep in parallel or in-place could publish cross-contaminated numbers.
• Configuration divergence from Eq. (26). src/prism_cli/trainers.py:49 defaults prism_floor_mode='scalar' while the optim module src/prism_cli/optim/prism.py:253 defaults to 'geometry', and the paper's Eq. (26) prescribes a per-module, noise-and-geometry-aware floor. The reproduced numbers were produced with the 'scalar' floor.
• STS-B disagreement. The reproduced PRISM $\varepsilon=6$ run gave STS-B = 0.828 vs. the paper's 0.718 (a +0.110 gap the verifier could not explain after checking metric convention). This is the main reason the replicated mean (0.705) sits above the paper's (0.690) rather than falling short — worth a closer look before claiming per-task wins.
• dp_clip_frac saturated at 1.0. In both completed PRISM runs, per-example clipping was active at every logged step, whereas the appendix figures (Figs. 8–11) treat clipping fraction as a quantity that varies non-trivially across steps and gauges. The training regime under which the paper's non-saturated diagnostic curves arise is not stated.

None of the above changes my judgment of the theoretical contribution, but it means the empirical portion of the paper — including every diagnostic figure the mechanism story rests on and the comparative claim in every headline table — is currently a black box from a reproducibility standpoint.

Recommended Changes

Essential

• Release the missing baselines. Add implementations of factor-space DP-AdamW (the "Baseline" throughout the diagnostics), FFA, LoRA-RITE, LoRA+, and LAMB — or provide reference forks — so Tables 2, 3, 4 are regenerable from the shipped code. Wire the --method CLI argument to route to the correct optimizer, and update the README so its baseline commands actually run.
• Release the gauge-sweep and $\sigma$-sweep harnesses. Add the multi-gauge training driver ($c \in \{0.25, 0.5, 1, 2, 4\}$ and the finer Appendix C.2 grid with 64 Monte-Carlo samples), the $\sigma$-sweep for $\varepsilon \in \{1.5, 3, 6, 12\}$, the per-step $S_t$ and $\|\Delta Z_t\|_F$ trackers, and the plotting code so that Figures 1–4 and 12–20 are end-to-end reproducible. This is the infrastructure the entire Issue I–III mechanism story rests on.
• Report per-seed variance. Rerun Table 2 (and ideally Tables 3, 4) with at least three seeds, add standard deviations, and ship a multi-seed orchestration script. Current wins (as small as 0.015 in average) are within one-seed noise for DP-LoRA at this scale.
• Fix the evaluation-path collision. Include $\varepsilon$, seed, rank, and privacy mode in LLM-Adapters/evaluate.py's generation-dump naming so multi-configuration sweeps do not silently overwrite one another.
• Add an ablation. Report DP utility for at least three variants that turn off one of {intrinsic tangent projection, isotropic tangent noise / retraction, DP-aware floors} to decompose PRISM's gain across mechanisms.
• Reconcile prism_floor_mode with Eq. (26). Align the trainer default with the optim default and with the paper's prescription, and state which mode produced the reported numbers.
• Explain the STS-B disagreement. Confirm the exact metric / split / decoding used for STS-B in Table 2, and either report a reproducer-friendly protocol or a corrected number.

Suggested

• Calibrate the "unbounded amplification" framing. Reword Section 1 / Section 2 to state that factor-space DP-LoRA drifts to a gauge-dependent noise level during training, with the worst case being unbounded, rather than implying the unbounded regime describes ordinary optimization.
• Add a profiling script for Table 5. Include the 10-warmup / 30-measured protocol as a runnable script with torch.cuda.Event timing and torch.cuda.max_memory_allocated, and briefly discuss the ~2× step-time cost in the main text (not only in the limitations paragraph); the reproduced 21.08 s / step supports the direction of the claim but not the ratio.
• Tighten the adaptive-preconditioner argument. Clarify whether "$P \approx \Sigma_\xi$" is about diagonal Adam under anisotropic true-gradient signal or a full-matrix preconditioner; note that under isotropic DP noise alone the reshaping vanishes.
• Address the non-DP utility gap. Explain why PRISM underperforms LoRA-RITE, LoRA+, LAMB, and AdamW without DP (Table 2, Non-DP row), and identify which mechanism (retraction, canonical lift, floors) is responsible.
• Discuss non-monotonicity at $r = 32$. Interpret the dip in Table 4 in terms of Theorem 3.1's $r(m+n-r)$ noise scaling.
• Unify baseline terminology. Pick one label (e.g., "factor-space DP-AdamW") and use it consistently across Figures 1–4, 14–20 and the accompanying prose; explicitly define "preconditioner aggressiveness" where it first appears on an axis.
• State the moment-input boundary condition. In Section 3.3, add one sentence noting that the moments $m_{A,\ell}, m_{B,\ell}, V_{A,\ell}, V_{B,\ell}$ are updated from the sanitized $(\Delta A^{\rm dp}_\ell, \Delta B^{\rm dp}_\ell)$ only, so post-processing immunity applies.
• Verify the rank-space moment prefactors and Sylvester sign. Either explain why the $A$- and $B$-sides in the Eq. (26) context are asymmetric or fix the mismatch; include the differentiation step for Prop. A.11.
• Clarify Figure 3 and the sampler dimension. Note that Figure 3's $S_t$ range reflects natural drift (not a deliberate gauge sweep) and that the near-linear fit implies the bilinear noise is empirically negligible in that regime; add one sentence in Section 3.2 explaining why an $(m+n)r$-dim sampler produces an isotropic distribution on the $r(m+n-r)$-dim tangent space.
• Soften the "amplification can even decrease" claim. Replace with "saturates rather than grows" unless the PRISM points in Figure 20 show a genuinely decreasing trend.
• Discuss the retraction residual. Add one paragraph on when $\|\Delta A \Delta B^\top\|_F$ is small enough that Prop. 3.2's $O(\eta^2)$ correction is negligible under DP noise, and when it might resurface.
• Address dp_clip_frac saturation. Report the intrinsic clip threshold and the clipping trajectory that produced Figures 8–11, so a reader can reconcile them with the always-saturated regime observed in the reproduced runs.