PETRecipe: one shared DP-SGD training path

[fazelehh]

Stacked on #438 (base is feature/pet-optimization; GitHub retargets this to main automatically once #438 merges).

What

Adds leakpro/optimization/training.py: the structured training contract (PETRecipe) and the single shared Opacus path (train_with_dpsgd, build_campaign_fns) that replaces the DP-SGD loop previously duplicated in every campaign example. All three examples (LOS-LR, CIFAR, LOS GRU-D) are ported onto it and shrink to declaring a recipe; the standalone GRU-D script is superseded by the recipe-ported one.

A PETRecipe is the plan's structured contract: instead of a monolithic train(data) -> model, the user supplies factories the optimizer can intervene in (make_model, make_optimizer, make_loader, criterion, epochs). Each factory receives the sampled knob config and reads what it needs, which is what makes the next PET cheap — a regularization bundle is just make_model reading dropout and make_optimizer reading weight_decay, with no interface change. Exotic loops that do not fit keep the escape hatch of writing the three campaign callables directly.

Hardening beyond the refactor

• make_opacus_compatible() — ModuleValidator pass (BatchNorm → GroupNorm), disables in-place activations, and rewrites torchvision residual blocks whose in-place out += identity ModuleValidator cannot see because it lives in forward rather than a submodule. Without this, any BatchNorm architecture — including both webapp image presets — fails at make_private().
• The privacy accountant is a parameter, defaulting to "prv" (was hard-coded "rdp"). RDP is the looser bound, so identical noise reported a larger ε here than via the webapp's DP-SGD path. This changes reported ε for existing campaigns (downward, for identical noise); empirical attack numbers are unaffected.
• The accuracy metric thresholds single-logit binary heads at 0 instead of argmax — argmax over one column is always 0, silently yielding 0% or 100%.

Review fixes applied

An adversarial review of this PR produced 10 findings; all are addressed in 18db5a2a. The theme was silent behaviour changes versus the per-example code this PR replaced.

• train_with_dpsgd fails closed. A missing or misspelled noise_multiplier silently trained a fully non-private model, from a function whose only purpose is DP-SGD. The key is now required; non-private runs pass an explicit 0.0.
• make_opacus_compatible runs on both branches. Rewriting BatchNorm → GroupNorm only for private configs left the noise=0 anchor a structurally different model, so its utility gap mixed the cost of DP noise with the cost of an architecture change — the frontier's utility axis was not comparable end to end. Documented consequence: the submitted architecture is not necessarily what trains.
• Full mimicry is enforced, not silently downgraded. n_ref_train used min(target, ref_pool), so a short reference pool quietly trained weaker references and biased the audit optimistic. A short pool is now an error — which is what the old per-example code did loudly.
• _patch_residual_blocks only rewrites blocks still using the stock forward. Subclasses with their own forward (SE, attention) were being silently replaced by the vanilla residual path, changing the model rather than making it Opacus-safe. Such blocks are now left alone with a warning.
• The auc utility metric refuses multiclass output instead of reshape(-1)-ing it into a meaningless number.
• The accountant travels with the epsilon in campaign_extras. The prv default is deliberate, but PRV and RDP epsilons are not comparable — measured 4.22 vs 4.87 for identical noise — so a record that does not name its accountant cannot safely be compared with another run's.

Examples: GRU-D regains the max_physical_batch=128 cap the pre-PR code set deliberately (it is the memory-heaviest target in the repo; the shared 256 default would OOM mid-campaign) and gains the small-dataset split guard its LR sibling already had. The LOS campaign imports target_models.LR instead of redeclaring an identical LRSigmoid, so it trains the same class the audited target used.

One earlier test asserted the now-incorrect behaviour (BatchNorm surviving the non-private path) and has been inverted to assert architecture consistency instead.

Verification

• pytest leakpro/tests/test_optimization/ — 47 tests pass.
• ruff check . clean.
• The DP-SGD compatibility tests fail with the fix disabled (checked), so they are real regression tests rather than vacuous ones.

The webapp's near-verbatim copy of the Opacus-compat block is removed in #448, which now calls this shared helper.

🤖 Generated with Claude Code

https://claude.ai/code/session_01QBPG7MprEqhqG95fFURk8h

[fazelehh]

Addressed the code-review findings on this PR. The theme was silent behaviour changes versus the per-example code this PR replaced.

Correctness

• train_with_dpsgd fails closed: a missing or misspelled noise_multiplier silently trained a fully non-private model. The key is now required; non-private runs pass an explicit 0.0.
• make_opacus_compatible runs on both branches. Rewriting BatchNorm→GroupNorm only for private configs left the noise=0 anchor a structurally different model, so its utility gap mixed the cost of DP noise with the cost of an architecture change — the frontier's utility axis was not comparable end to end.
• Full mimicry is enforced, not silently downgraded: n_ref_train used min(target, ref_pool), so a short reference pool quietly trained weaker references and biased the audit optimistic. Now an error, as the old per-example code was.
• _patch_residual_blocks only rewrites blocks still using the stock forward; subclasses with a custom forward (SE/attention) were being silently replaced by the vanilla residual path.
• The auc utility metric refuses multiclass output instead of reshape(-1)-ing it into a meaningless number.
• The accountant travels with the epsilon in campaign_extras. The prv default is deliberate (it matches the webapp's DP-SGD path), but PRV and RDP epsilons are not comparable — measured 4.22 vs 4.87 for identical noise — so a record that does not name its accountant cannot safely be compared with another run's.

Examples

• GRU-D regains max_physical_batch=128 (it is the memory-heaviest target here; the shared 256 default would OOM mid-campaign) and gains the small-dataset split guard its LR sibling already had.
• The LOS campaign imports target_models.LR instead of redeclaring an identical LRSigmoid, so it trains the same class the audited target used.

Also: one of my earlier tests asserted the now-incorrect behaviour (BatchNorm surviving the non-private path) and has been inverted to assert architecture consistency. The webapp's duplicate Opacus-compat block is removed in #448, which now calls this shared helper.

47 tests pass on this branch, ruff clean.