diff --git a/research/slaclip/README.md b/research/slaclip/README.md new file mode 100644 index 000000000..4a5db3c8b --- /dev/null +++ b/research/slaclip/README.md @@ -0,0 +1,241 @@ +# SlaClip: Gradient Norm Slacks can be an Indicator for Adaptive Clipping in DP-SGD + +SlaClip is a privacy-preserving adaptive clipping method for differentially +private stochastic gradient descent (DP-SGD). It uses the norm budget left +unused by clipping—the *slack*—to release a noisy, binned estimate of the +gradient-norm cumulative distribution function (CDF). That Slack Indicator can +then drive the clipping-threshold update from the SlaClip paper or another +post-processing controller. + +This directory is a self-contained research prototype. It does not modify +`PrivacyEngine` or any other Opacus core API. + +## Method + +For a per-sample gradient `g_i`, clipping threshold `C`, and `K` CDF slots, +SlaClip defines the extended gradient from equations (6)-(8) as + +```text +g_i^+ = [Clip_C(g_i); s_i] +lambda = C / sqrt(K) +sqrt(K) * max(C - ||g_i||, 0) = a * lambda + b +s_i = [lambda * 1(a); b; 0], where 0 <= b < lambda. +``` + +This construction guarantees `||g_i^+||_2 <= C`. Under add/remove adjacency +and the fixed normalization constant `B`, the extended average query therefore +has the same `C / B` L2 sensitivity as the vanilla clipped-gradient query. With +the same sampling rule and noise multiplier, it can use the same per-step +privacy-accounting parameters as DP-SGD. + +The first `d` coordinates are exactly the native DP-SGD gradient release. The +implementation computes the gradient and `K` slack coordinates separately to +avoid materializing a `d + K` tensor. Adding independent `N(0, (sigma*C)^2)` +noise to both parts is distributionally identical to one isotropic Gaussian +draw in `d + K` dimensions. + +After aggregation and noise, the last `K` coordinates are divided by +`B * lambda`. The resulting `optimizer.slack_indicator` is the paper's noisy, +bin-averaged CDF estimate from equation (11). Its first coordinate describes +norms near `C`; its last coordinate describes norms near zero. + +> **Privacy boundary:** only the noisy `slack_indicator` property is a public +> output. Per-sample slack and the unnoised slack aggregate are internal to the +> joint DP-SGD mechanism and must not be released. Computing a separate CDF +> query outside this mechanism would require its own privacy analysis and +> accounting. + +### Selecting the slack dimension K + +When `num_slots=None` (the default), the optimizer selects `K` using the +99%-confidence SNR rule in equation (36): + +```text +K_max = (B / (2 * 2.576 * sigma))^(2/3) +K = max(1, floor(K_max)). +``` + +Here, `B` is the fixed normalization constant used by Opacus +(`expected_batch_size`), not a realized Poisson batch size or a physical +microbatch size. Table 3 gives the following illustrative practical choices +for representative batch sizes when `sigma=1`: + +| B | 128 | 256 | 512 | 1024 | 2048 | +|---:|---:|---:|---:|---:|---:| +| Illustrative practical K | 8 | 10 | 20 | 30 | 50 | + +These rounded values show the approximate scale of `K`; they are not treated +as special cases by the implementation. The automatic rule always evaluates +equation (36) directly, giving `K = {8, 13, 21, 34, 54}` for the batch sizes +above when `sigma=1`. Pass a positive integer as `num_slots` to reproduce a +particular practical choice or run an ablation. The selected value is +available as `optimizer.K`. + +## Usage + +The prototype supports two composable steps: + +1. `SlaClipDPOptimizer` jointly releases the DP gradient and private Slack + Indicator. With no controller, the clipping threshold remains unchanged. +2. `SlaClipController` consumes the released indicator and applies equations + (28)-(30) to update the threshold. Since this is post-processing of a DP + release, it adds no privacy cost. + +### Prepare private training + +`SlaClipPrivacyEngine` is the recommended entry point. It uses Opacus's native +model wrapping, data loader, secure RNG, and privacy accountant while selecting +`SlaClipDPOptimizer` automatically: + +```python +from research.slaclip import SlaClipPrivacyEngine + + +privacy_engine = SlaClipPrivacyEngine() +model, optimizer, train_loader = privacy_engine.make_private( + module=model, + optimizer=optimizer, + data_loader=train_loader, + noise_multiplier=1.0, + max_grad_norm=1.0, +) +``` + +This entry point also supports the inherited +`PrivacyEngine.make_private_with_epsilon()` API. It intentionally rejects +distributed training, non-flat clipping, and ghost clipping, which the current +research optimizer does not support. + +### Step 1 only: obtain private CDF information + +The default `SlaClipPrivacyEngine()` has no controller, so it keeps `C` fixed +while obtaining the noisy Slack Indicator: + +```python +for images, targets in train_loader: + optimizer.zero_grad() + loss = criterion(model(images), targets) + loss.backward() + optimizer.step() + + private_cdf = optimizer.slack_indicator + # Use private_cdf only in DP-safe post-processing. +``` + +Because Gaussian noise is unbounded, individual coordinates can fall outside +`[0, 1]` or fail to be monotone. A downstream method may project or smooth the +released vector as post-processing without additional privacy cost. + +### Steps 1 and 2: paper SlaClip + +Construct the entry point with the paper controller to adapt `C` after every +release, then call `make_private()` as above: + +```python +from research.slaclip import SlaClipController, SlaClipPrivacyEngine + + +privacy_engine = SlaClipPrivacyEngine( + clipping_controller=SlaClipController( + eta=0.5, + min_clipbound=0.1, + max_clipbound=50.0, + ), +) +model, optimizer, train_loader = privacy_engine.make_private( + module=model, + optimizer=optimizer, + data_loader=train_loader, + noise_multiplier=1.0, + max_grad_norm=1.0, +) + +for images, targets in train_loader: + optimizer.zero_grad() + loss = criterion(model(images), targets) + loss.backward() + optimizer.step() + + print(optimizer.current_clip) +``` + +The controller implements + +```text +r_t = clip_[0,1](slack_indicator[K] / C_t) +gamma_t = clip_[0,1](1 - (1 - r_t) / 2) +C_(t+1) = clip_[C_min,C_max]( + C_t * exp(eta * (gamma_t - slack_indicator[1])) + ). +``` + +The released indicator contains unbounded Gaussian noise. The controller +therefore projects `r_t` and `gamma_t` onto `[0, 1]`, as specified in the +paper, and bounds the next positive clipping threshold to +`[min_clipbound, max_clipbound]`. It intentionally does not force +`gamma_t - slack_indicator[1]` to be nonnegative: a negative value is the +feedback that decreases an overly large clipping threshold. + +A custom callable with signature `(current_clip, slack_indicator) -> next_clip` +can replace `SlaClipController`. Such a method reuses the SlaClip Slack +Indicator but is not the paper's threshold controller and should be described +accordingly. + +### Parameters + +- `num_slots`: number `K` of CDF bins. `None` automatically selects it from + equation (36). A positive integer overrides the automatic value. + Larger values increase resolution but also increase normalized indicator + noise. Pass it to `SlaClipPrivacyEngine`. +- `clipping_controller`: optional post-processing callable. `None` enables the + indicator-only mode. Pass it to `SlaClipPrivacyEngine`. +- `eta`: positive multiplicative update step size used by + `SlaClipController`. +- `min_clipbound`, `max_clipbound`: positive lower and upper bounds for the + clipping threshold. Their defaults, `0.1` and `50.0`, match Opacus + `AdaClipDPOptimizer` and the SlaClip experiment configuration. +- All other optimizer arguments have the same meaning as in Opacus + `DPOptimizer`. + +## Limitations + +- `SlaClipPrivacyEngine` supports the standard, non-distributed + `DPOptimizer` path. SlaClip is not registered as a clipping mode on the core + Opacus `PrivacyEngine`. +- Use it from the repository root through `research.slaclip`; research modules + are not part of the installed Opacus public API. +- The privacy argument assumes the same sampling rule, normalization constant, + clipping threshold, noise multiplier, and accountant parameters for the + gradient and slack parts of the joint release. +- As research code, it is not covered by Opacus public API compatibility + guarantees. + +## Tests + +From the repository root: + +```bash +python -m pytest research/slaclip -q +``` + +The tests cover the `SlaClipPrivacyEngine` entry point and accountant hook, +automatic `K` selection, equations (7)-(8), the extended-gradient norm bound, +exact agreement of the first `d` coordinates with native `DPOptimizer`, +indicator-only operation, the paper controller, and empty Poisson batches. + +## Citation + +```bibtex +@inproceedings{zou2026slaclip, + title={{SlaClip}: Gradient Norm Slacks Can Be an Indicator for Adaptive + Clipping in {DP-SGD}}, + author={Zou, Shuyan and Wang, Shaowei and Zhu, Zhanxing and Li, Jin and + Dong, Changyu and Sassone, Vladimiro and Wu, Han}, + booktitle={Proceedings of the 43rd International Conference on Machine + Learning}, + year={2026} +} +``` + +The authors' reference implementation is available at +[ZsyRock/SlaClip](https://github.com/ZsyRock/SlaClip). diff --git a/research/slaclip/__init__.py b/research/slaclip/__init__.py new file mode 100644 index 000000000..2acc0ad98 --- /dev/null +++ b/research/slaclip/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 SlaClip authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Research package for the SlaClip prototype.""" + +from .privacy_engine import SlaClipPrivacyEngine # noqa: F401 +from .slaclipoptimizer import SlaClipController, SlaClipDPOptimizer # noqa: F401 diff --git a/research/slaclip/privacy_engine.py b/research/slaclip/privacy_engine.py new file mode 100644 index 000000000..e422d3f41 --- /dev/null +++ b/research/slaclip/privacy_engine.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 SlaClip authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PrivacyEngine entry point for the SlaClip research prototype.""" + +from __future__ import annotations + +from typing import Callable, List, Optional, Union + +import torch +from opacus.optimizers import DPOptimizer +from opacus.privacy_engine import PrivacyEngine +from torch import optim + +from .slaclipoptimizer import SlaClipDPOptimizer + + +class SlaClipPrivacyEngine(PrivacyEngine): + """Prepare Opacus training objects with :class:`SlaClipDPOptimizer`. + + Passing no controller enables the indicator-only mode. Passing the paper's + ``SlaClipController`` enables the complete SlaClip adaptation rule. + """ + + def __init__( + self, + *, + accountant: str = "prv", + secure_mode: bool = False, + num_slots: Optional[int] = None, + clipping_controller: Optional[Callable[[float, torch.Tensor], float]] = None, + ): + super().__init__(accountant=accountant, secure_mode=secure_mode) + self.num_slots = num_slots + self.clipping_controller = clipping_controller + + def _prepare_optimizer( + self, + *, + optimizer: optim.Optimizer, + noise_multiplier: float, + max_grad_norm: Union[float, List[float]], + expected_batch_size: int, + loss_reduction: str = "mean", + distributed: bool = False, + clipping: str = "flat", + noise_generator=None, + grad_sample_mode: str = "hooks", + **kwargs, + ) -> SlaClipDPOptimizer: + if distributed: + raise ValueError("SlaClip does not currently support distributed training") + if clipping != "flat": + raise ValueError("SlaClip requires clipping='flat'") + if "ghost" in grad_sample_mode: + raise ValueError("SlaClip does not currently support ghost clipping") + if isinstance(max_grad_norm, list): + raise ValueError("SlaClip requires a scalar max_grad_norm") + + if isinstance(optimizer, DPOptimizer): + optimizer = optimizer.original_optimizer + + generator = None + if self.secure_mode: + generator = self.secure_rng + elif noise_generator is not None: + generator = noise_generator + + return SlaClipDPOptimizer( + optimizer=optimizer, + noise_multiplier=noise_multiplier, + max_grad_norm=float(max_grad_norm), + expected_batch_size=expected_batch_size, + loss_reduction=loss_reduction, + generator=generator, + secure_mode=self.secure_mode, + num_slots=self.num_slots, + clipping_controller=self.clipping_controller, + **kwargs, + ) diff --git a/research/slaclip/slaclipoptimizer.py b/research/slaclip/slaclipoptimizer.py new file mode 100644 index 000000000..ee51816a6 --- /dev/null +++ b/research/slaclip/slaclipoptimizer.py @@ -0,0 +1,293 @@ +"""Experimental SlaClip optimizer implementation for research use. + +This self-contained implementation is stored under ``research/slaclip`` and +does not modify Opacus core routing or public APIs. +""" + +# Copyright (c) 2026 SlaClip authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math +from typing import Callable, Optional + +import torch +from opacus.optimizers.optimizer import ( + DPOptimizer, + _check_processed_flag, + _generate_noise, + _mark_as_processed, +) +from torch.optim import Optimizer + + +def paper_recommended_k(batch_size: int, noise_multiplier: float = 1.0) -> int: + """Select the largest integer ``K`` satisfying SlaClip equation (36).""" + + z_0995 = 2.576 + + batch_size = int(batch_size) + noise_multiplier = float(noise_multiplier) + if batch_size <= 0: + raise ValueError("batch_size must be positive") + if not math.isfinite(noise_multiplier) or noise_multiplier <= 0: + raise ValueError("noise_multiplier must be finite and positive") + + k_max = (batch_size / (2.0 * z_0995 * noise_multiplier)) ** (2.0 / 3.0) + return max(1, math.floor(k_max)) + + +class SlaClipController: + """Paper-faithful clipping update from SlaClip equations (28)-(30).""" + + def __init__( + self, + *, + eta: float = 0.5, + min_clipbound: float = 0.1, + max_clipbound: float = 50.0, + ): + if not math.isfinite(eta) or eta <= 0: + raise ValueError("eta must be finite and positive") + if not math.isfinite(min_clipbound) or min_clipbound <= 0: + raise ValueError("min_clipbound must be finite and positive") + if not math.isfinite(max_clipbound) or max_clipbound <= min_clipbound: + raise ValueError( + "max_clipbound must be finite and larger than min_clipbound" + ) + self.eta = float(eta) + self.min_clipbound = float(min_clipbound) + self.max_clipbound = float(max_clipbound) + + def __call__(self, current_clip: float, slack_indicator: torch.Tensor) -> float: + if not math.isfinite(current_clip) or current_clip <= 0: + raise ValueError("current_clip must be finite and positive") + if slack_indicator.ndim != 1 or slack_indicator.numel() == 0: + raise ValueError("slack_indicator must be a non-empty vector") + if not bool(torch.isfinite(slack_indicator).all().item()): + raise ValueError("slack_indicator must contain only finite values") + + near_threshold = float(slack_indicator[0].item()) + near_zero = float(slack_indicator[-1].item()) + near_zero_adjusted = max(0.0, min(1.0, near_zero / current_clip)) + target = max(0.0, min(1.0, 1.0 - (1.0 - near_zero_adjusted) / 2.0)) + log_next_clip = math.log(current_clip) + self.eta * (target - near_threshold) + if log_next_clip <= math.log(self.min_clipbound): + return self.min_clipbound + if log_next_clip >= math.log(self.max_clipbound): + return self.max_clipbound + return float(math.exp(log_next_clip)) + + +class SlaClipDPOptimizer(DPOptimizer): + """DPOptimizer that jointly releases a private Slack Indicator. + + The optimizer implements SlaClip equations (6)-(11) without materializing a + ``d + K`` tensor. Clipped gradients and slack coordinates are accumulated + separately, then perturbed with independent coordinates of the same Gaussian + mechanism. Passing a ``clipping_controller`` additionally applies step 2 of + SlaClip; leaving it as ``None`` exposes only the private CDF information. + """ + + def __init__( + self, + optimizer: Optimizer, + *, + noise_multiplier: float, + max_grad_norm: float, + expected_batch_size: int, + loss_reduction: str = "mean", + generator=None, + secure_mode: bool = False, + num_slots: Optional[int] = None, + clipping_controller: Optional[Callable[[float, torch.Tensor], float]] = None, + **kwargs, + ): + if expected_batch_size is None or expected_batch_size <= 0: + raise ValueError("expected_batch_size must be positive") + if max_grad_norm <= 0: + raise ValueError("max_grad_norm must be positive") + super().__init__( + optimizer, + noise_multiplier=noise_multiplier, + max_grad_norm=max_grad_norm, + expected_batch_size=expected_batch_size, + loss_reduction=loss_reduction, + generator=generator, + secure_mode=secure_mode, + **kwargs, + ) + + self.K = ( + paper_recommended_k(expected_batch_size, noise_multiplier) + if num_slots is None + else int(num_slots) + ) + if self.K <= 0: + raise ValueError("num_slots must be a positive integer") + + self.clipping_controller = clipping_controller + self._lambda_t = 0.0 + self._slack_sum: Optional[torch.Tensor] = None + self._slack_indicator: Optional[torch.Tensor] = None + + @property + def current_clip(self) -> float: + """Current clipping threshold.""" + + return float(self.max_grad_norm) + + @property + def slack_indicator(self) -> torch.Tensor: + """Most recently released noisy, normalized Slack Indicator. + + This property never exposes the unnoised per-sample slack or aggregate. + """ + + if self._slack_indicator is None: + raise RuntimeError( + "Slack Indicator is available only after optimizer.step()" + ) + return self._slack_indicator.detach().clone() + + def zero_grad(self, set_to_none: bool = False): + super().zero_grad(set_to_none) + if not self._is_last_step_skipped: + self._slack_sum = None + self._lambda_t = 0.0 + + def _release_denom(self) -> float: + denom = float(self.expected_batch_size) * float(self.accumulated_iterations) + if denom <= 0: + raise ValueError("Expected release denominator must be positive") + return denom + + def _encode_slack( + self, per_sample_norms: torch.Tensor, current_clip: float + ) -> torch.Tensor: + """Encode SlaClip equations (7)-(8) for a batch of gradient norms.""" + + lambda_t = float(current_clip / math.sqrt(self.K)) + scaled_slack = torch.clamp( + current_clip - per_sample_norms, min=0.0 + ) * math.sqrt(self.K) + full_slots = torch.floor(scaled_slack / lambda_t).to(torch.int64) + full_slots = torch.clamp(full_slots, min=0, max=self.K) + residual = scaled_slack - full_slots.to(scaled_slack.dtype) * lambda_t + residual = torch.where( + full_slots >= self.K, torch.zeros_like(residual), residual + ) + + slot_index = torch.arange(self.K, device=per_sample_norms.device).view( + 1, self.K + ) + slack = (slot_index < full_slots.view(-1, 1)).to( + dtype=per_sample_norms.dtype + ) * float(lambda_t) + + has_residual = full_slots < self.K + if has_residual.any(): + residual_slot = torch.clamp(full_slots, max=self.K - 1) + slack[has_residual, residual_slot[has_residual]] = residual[has_residual] + return slack + + def clip_and_accumulate(self): + """Clip gradients as DPOptimizer and also aggregate encoded slack.""" + + grad_samples = self.grad_samples + if not grad_samples: + return + + if len(grad_samples[0]) == 0: + per_sample_norms = torch.zeros( + (0,), device=grad_samples[0].device, dtype=grad_samples[0].dtype + ) + per_sample_clip_factor = per_sample_norms + else: + per_param_norms = [ + grad_sample.reshape(len(grad_sample), -1).norm(2, dim=-1) + for grad_sample in grad_samples + ] + target_device = per_param_norms[0].device + per_param_norms = [norm.to(target_device) for norm in per_param_norms] + per_sample_norms = torch.stack(per_param_norms, dim=1).norm(2, dim=1) + per_sample_clip_factor = ( + self.current_clip / (per_sample_norms + 1e-6) + ).clamp(max=1.0) + + for parameter in self.params: + _check_processed_flag(parameter.grad_sample) + grad_sample = self._get_flat_grad_sample(parameter).to(parameter.dtype) + clip_factor = per_sample_clip_factor.to( + device=grad_sample.device, dtype=parameter.dtype + ) + grad = torch.einsum("i,i...", clip_factor, grad_sample) + if parameter.summed_grad is None: + parameter.summed_grad = grad + else: + parameter.summed_grad += grad + _mark_as_processed(parameter.grad_sample) + + lambda_t = float(self.current_clip / math.sqrt(self.K)) + if self._lambda_t and not math.isclose( + self._lambda_t, lambda_t, rel_tol=1e-12, abs_tol=0.0 + ): + raise ValueError( + "Clipping threshold changed while accumulating a logical batch" + ) + self._lambda_t = lambda_t + + batch_slack_sum = self._encode_slack(per_sample_norms, self.current_clip).sum( + dim=0 + ) + if self._slack_sum is None: + self._slack_sum = batch_slack_sum + else: + self._slack_sum += batch_slack_sum.to(self._slack_sum.device) + + def add_noise(self): + """Release gradient and Slack Indicator as one extended mechanism.""" + + current_clip = self.current_clip + + # Keep the first d coordinates on the native DPOptimizer path. Drawing + # the K slack coordinates afterwards is distributionally identical to + # one isotropic Gaussian draw in d + K dimensions. + super().add_noise() + + if self._slack_sum is None or self._lambda_t <= 0: + raise RuntimeError("Slack must be accumulated before adding noise") + slack_noise = _generate_noise( + std=self.noise_multiplier * current_clip, + reference=self._slack_sum, + generator=self.generator, + secure_mode=self.secure_mode, + ) + self._slack_indicator = (self._slack_sum + slack_noise) / ( + self._lambda_t * self._release_denom() + ) + + if self.clipping_controller is not None: + next_clip = float( + self.clipping_controller(current_clip, self._slack_indicator) + ) + if not math.isfinite(next_clip) or next_clip <= 0: + raise ValueError( + "clipping_controller must return a finite positive value" + ) + self.max_grad_norm = next_clip + + +SlaClipOptimizer = SlaClipDPOptimizer diff --git a/research/slaclip/test_slaclipoptimizer.py b/research/slaclip/test_slaclipoptimizer.py new file mode 100644 index 000000000..542c66aea --- /dev/null +++ b/research/slaclip/test_slaclipoptimizer.py @@ -0,0 +1,341 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 SlaClip authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import unittest + +import torch +from opacus.optimizers.optimizer import DPOptimizer +from torch.utils.data import DataLoader, TensorDataset + +from research.slaclip import SlaClipController, SlaClipDPOptimizer, SlaClipPrivacyEngine +from research.slaclip.slaclipoptimizer import paper_recommended_k + + +def make_optimizer( + optimizer_class=SlaClipDPOptimizer, + *, + noise_multiplier: float = 0.0, + max_grad_norm: float = 2.0, + expected_batch_size: int = 1, + num_slots: int = 2, + clipping_controller=None, + generator=None, + dtype=torch.float32, +): + model = torch.nn.Linear(2, 1, bias=False, dtype=dtype) + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + common_args = dict( + noise_multiplier=noise_multiplier, + max_grad_norm=max_grad_norm, + expected_batch_size=expected_batch_size, + generator=generator, + ) + if optimizer_class is SlaClipDPOptimizer: + common_args.update( + num_slots=num_slots, + clipping_controller=clipping_controller, + ) + return model, optimizer_class(optimizer, **common_args) + + +class SlaClipOptimizerResearchTest(unittest.TestCase): + def test_paper_recommended_k_uses_equation_36(self): + expected = { + 128: 8, + 256: 13, + 512: 21, + 1024: 34, + 2048: 54, + } + for batch_size, num_slots in expected.items(): + with self.subTest(batch_size=batch_size): + self.assertEqual(paper_recommended_k(batch_size, 1.0), num_slots) + + def test_paper_recommended_k_accounts_for_noise_multiplier(self): + batch_size = 2000 + noise_multiplier = 1.415787 + expected = math.floor( + (batch_size / (2.0 * 2.576 * noise_multiplier)) ** (2.0 / 3.0) + ) + + self.assertEqual(paper_recommended_k(batch_size, noise_multiplier), expected) + self.assertEqual(expected, 42) + + def test_optimizer_auto_selects_k_and_allows_explicit_override(self): + _, automatic = make_optimizer( + noise_multiplier=1.0, + expected_batch_size=512, + num_slots=None, + ) + _, explicit = make_optimizer( + noise_multiplier=1.0, + expected_batch_size=512, + num_slots=7, + ) + + self.assertEqual(automatic.K, 21) + self.assertEqual(explicit.K, 7) + + def test_paper_recommended_k_rejects_invalid_inputs(self): + invalid_inputs = ( + (0, 1.0), + (128, 0.0), + (128, float("inf")), + ) + for batch_size, noise_multiplier in invalid_inputs: + with self.subTest( + batch_size=batch_size, + noise_multiplier=noise_multiplier, + ), self.assertRaises(ValueError): + paper_recommended_k(batch_size, noise_multiplier) + + def test_paper_controller_equations_28_to_30(self): + controller = SlaClipController(eta=0.5) + slack_indicator = torch.tensor([0.2, 0.4]) + current_clip = 2.0 + adjusted_near_zero = 0.4 / current_clip + target = 1.0 - (1.0 - adjusted_near_zero) / 2.0 + expected = current_clip * math.exp(0.5 * (target - 0.2)) + self.assertAlmostEqual( + controller(current_clip, slack_indicator), expected, places=6 + ) + + def test_controller_projects_near_zero_signal(self): + controller = SlaClipController( + eta=0.5, + min_clipbound=0.1, + max_clipbound=50.0, + ) + current_clip = 2.0 + + projected_low = controller(current_clip, torch.tensor([0.2, -100.0])) + expected_low = current_clip * math.exp(0.5 * (0.5 - 0.2)) + self.assertAlmostEqual(projected_low, expected_low, places=6) + + projected_high = controller(current_clip, torch.tensor([0.2, 100.0])) + expected_high = current_clip * math.exp(0.5 * (1.0 - 0.2)) + self.assertAlmostEqual(projected_high, expected_high, places=6) + + def test_controller_clamps_next_threshold(self): + controller = SlaClipController( + eta=1.0, + min_clipbound=1.0, + max_clipbound=3.0, + ) + + self.assertEqual(controller(2.0, torch.tensor([-100.0, 1.0])), 3.0) + self.assertEqual(controller(2.0, torch.tensor([100.0, 1.0])), 1.0) + + def test_controller_allows_threshold_to_decrease(self): + controller = SlaClipController( + eta=0.5, + min_clipbound=0.1, + max_clipbound=50.0, + ) + next_clip = controller(2.0, torch.tensor([0.9, 0.0])) + + self.assertGreater(next_clip, 0.0) + self.assertLess(next_clip, 2.0) + + def test_controller_rejects_invalid_bounds_and_nonfinite_values(self): + invalid_kwargs = ( + {"eta": float("inf")}, + {"min_clipbound": 0.0}, + {"min_clipbound": 1.0, "max_clipbound": 1.0}, + {"max_clipbound": float("inf")}, + ) + for kwargs in invalid_kwargs: + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): + SlaClipController(**kwargs) + + controller = SlaClipController() + with self.assertRaises(ValueError): + controller(1.0, torch.tensor([float("nan"), 0.0])) + + def test_equations_7_and_8_preserve_extended_norm_bound(self): + _, optimizer = make_optimizer(max_grad_norm=2.0, num_slots=4) + norms = torch.tensor([0.0, 0.5, 1.25, 2.0, 3.0]) + slack = optimizer._encode_slack(norms, optimizer.current_clip) + + lambda_t = optimizer.current_clip / math.sqrt(optimizer.K) + encoded_amount = slack.sum(dim=1) + expected_amount = torch.clamp( + optimizer.current_clip - norms, min=0.0 + ) * math.sqrt(optimizer.K) + torch.testing.assert_close(encoded_amount, expected_amount) + self.assertTrue(torch.all(slack >= 0)) + self.assertTrue(torch.all(slack <= lambda_t)) + + clipped_norms = torch.clamp(norms, max=optimizer.current_clip) + extended_norms = torch.sqrt(clipped_norms.square() + slack.square().sum(dim=1)) + self.assertTrue(torch.all(extended_norms <= optimizer.current_clip + 1e-6)) + + def test_gradient_release_matches_native_dpoptimizer(self): + native_model, native = make_optimizer( + DPOptimizer, + noise_multiplier=1.0, + generator=torch.Generator().manual_seed(1234), + dtype=torch.float64, + ) + slaclip_model, slaclip = make_optimizer( + noise_multiplier=1.0, + generator=torch.Generator().manual_seed(1234), + dtype=torch.float64, + ) + grad_sample = torch.tensor([[[3.0, 4.0]], [[0.25, 0.5]]], dtype=torch.float64) + next(native_model.parameters()).grad_sample = grad_sample.clone() + next(slaclip_model.parameters()).grad_sample = grad_sample.clone() + + native.clip_and_accumulate() + slaclip.clip_and_accumulate() + torch.testing.assert_close( + next(slaclip_model.parameters()).summed_grad, + next(native_model.parameters()).summed_grad, + ) + + native.add_noise() + slaclip.add_noise() + torch.testing.assert_close( + next(slaclip_model.parameters()).grad, + next(native_model.parameters()).grad, + ) + + def test_indicator_only_mode_does_not_change_threshold(self): + model, optimizer = make_optimizer( + max_grad_norm=2.0, + expected_batch_size=2, + num_slots=4, + clipping_controller=None, + ) + parameter = next(model.parameters()) + parameter.grad_sample = torch.tensor( + [[[0.0, 0.0]], [[1.0, 0.0]]], dtype=parameter.dtype + ) + + optimizer.clip_and_accumulate() + optimizer.add_noise() + + self.assertEqual(optimizer.current_clip, 2.0) + expected = torch.tensor([1.0, 1.0, 0.5, 0.5]) + torch.testing.assert_close(optimizer.slack_indicator, expected) + + def test_controller_consumes_private_indicator(self): + model, optimizer = make_optimizer( + max_grad_norm=2.0, + expected_batch_size=2, + num_slots=4, + clipping_controller=SlaClipController(eta=0.5), + ) + parameter = next(model.parameters()) + parameter.grad_sample = torch.tensor( + [[[0.0, 0.0]], [[1.0, 0.0]]], dtype=parameter.dtype + ) + + optimizer.clip_and_accumulate() + optimizer.add_noise() + + expected = SlaClipController(eta=0.5)(2.0, torch.tensor([1.0, 1.0, 0.5, 0.5])) + self.assertAlmostEqual(optimizer.current_clip, expected, places=6) + + def test_empty_batch_releases_noisy_indicator(self): + model, optimizer = make_optimizer( + noise_multiplier=0.0, + expected_batch_size=2, + num_slots=2, + ) + parameter = next(model.parameters()) + parameter.grad_sample = torch.empty( + (0,) + tuple(parameter.shape), dtype=parameter.dtype + ) + + optimizer.clip_and_accumulate() + optimizer.add_noise() + + torch.testing.assert_close(optimizer.slack_indicator, torch.zeros(2)) + + def test_privacy_engine_entrypoint_records_one_joint_release(self): + model = torch.nn.Linear(2, 1) + base_optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + data_loader = DataLoader( + TensorDataset(torch.randn(8, 2), torch.randn(8, 1)), batch_size=4 + ) + privacy_engine = SlaClipPrivacyEngine( + num_slots=4, + clipping_controller=SlaClipController(eta=0.5), + ) + model, optimizer, data_loader = privacy_engine.make_private( + module=model, + optimizer=base_optimizer, + data_loader=data_loader, + noise_multiplier=1.0, + max_grad_norm=1.0, + poisson_sampling=False, + ) + + inputs, targets = next(iter(data_loader)) + optimizer.zero_grad() + torch.nn.functional.mse_loss(model(inputs), targets).backward() + optimizer.step() + + self.assertEqual(len(privacy_engine.accountant), 1) + self.assertIsInstance(optimizer, SlaClipDPOptimizer) + self.assertEqual(optimizer.slack_indicator.shape, (4,)) + + def test_privacy_engine_supports_epsilon_entrypoint(self): + model = torch.nn.Linear(2, 1) + base_optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + data_loader = DataLoader( + TensorDataset(torch.randn(8, 2), torch.randn(8, 1)), batch_size=4 + ) + privacy_engine = SlaClipPrivacyEngine(accountant="rdp", num_slots=2) + + _, optimizer, _ = privacy_engine.make_private_with_epsilon( + module=model, + optimizer=base_optimizer, + data_loader=data_loader, + target_epsilon=10.0, + target_delta=1e-5, + epochs=1, + max_grad_norm=1.0, + poisson_sampling=False, + ) + + self.assertIsInstance(optimizer, SlaClipDPOptimizer) + self.assertEqual(optimizer.K, 2) + self.assertIsNone(optimizer.clipping_controller) + + def test_privacy_engine_rejects_non_flat_clipping(self): + model = torch.nn.Linear(2, 1) + base_optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + data_loader = DataLoader( + TensorDataset(torch.randn(8, 2), torch.randn(8, 1)), batch_size=4 + ) + privacy_engine = SlaClipPrivacyEngine() + + with self.assertRaisesRegex(ValueError, "requires clipping='flat'"): + privacy_engine.make_private( + module=model, + optimizer=base_optimizer, + data_loader=data_loader, + noise_multiplier=1.0, + max_grad_norm=1.0, + poisson_sampling=False, + clipping="per_layer", + ) + + +if __name__ == "__main__": + unittest.main()