Add DpsgdConfig execution plan for Poisson and fixed-size DP-SGD by Ayush7614 · Pull Request #337 · google-deepmind/jax_privacy

Ayush7614 · 2026-07-28T21:41:21Z

Summary

Add a single DpsgdConfig Tier-2 DPExecutionPlan config for vanilla DP-SGD, with DpsgdBatchSelection (POISSON / FIXED) routing batch selection and matching accounting events.
Mirror the existing BandMFConfig calibrate() / make() API, including expected_participations (no public num_examples required for add/remove Poisson) and privacy-correctness tests.
Validation helpers live in _validate (not_none, instance_of, at_most); expected batch size is exposed as expected_batch_size.

AI usage disclosure

AI (Cursor) was used to draft and iterate on this PR under my direction. I specified the API shape (unify configs, expected_participations, enum routing, drop replace / use_zcdp, lean validation, remove the extra example), ran and checked the test suite locally (pytest tests/execution_plan_test.py), and reviewed the resulting diffs before each push. Several review rounds with @ryan112358 drove further edits; I verified each change against the comments and re-ran tests after updates.

Test plan

pytest tests/execution_plan_test.py (21 passed)
CI on this PR

Provide Tier-2 DP-SGD configs that couple Poisson or fixed-size batch selection, clipping, Gaussian noise, and matching accounting events so users can run vanilla DP-SGD without hand-wiring Tier-3 components.

github-actions

👋 Welcome, and thanks for opening your first PR!

We're excited to have you contribute. A few things to help you get started:

Join the community! Say hello on the Google Differential Privacy Slack — we'd love to hear who you are, what you're working on, and what brought you here.
A maintainer will review your PR shortly. If you have questions in the meantime, don't hesitate to ask here or on Slack.

Thanks for contributing to privacy-preserving machine learning! 🔐

ryan112358

Nice addition, thanks! Please address the comments and I'll take another look

ryan112358 · 2026-07-28T21:50:32Z

+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class DpsgdConfig:


Let's try to unify DpsgdConfig and FixedSizeDpsgdConfig under a single dataclass, where the batch selection / accounting is routed based on one of the dataclass fields

Done — merged into a single DpsgdConfig. Routing is now via an explicit batch_selection: DpsgdBatchSelection field (POISSON / FIXED), so additional strategies can be added later without a second config class.

ryan112358 · 2026-07-28T21:52:03Z

+
+  Attributes:
+    iterations: Number of training iterations / batch draws.
+    expected_batch_size: Expected Poisson batch size. Sampling probability is


Being able to specify this implies you know the number of examples, which rules out ADD_OR_REMOVE neighboring relation. Let's parameterize this in terms of "expected_participations" like BandMFConfig

Done — Poisson path now takes expected_participations (same idea as BandMFConfig). num_examples is only required for fixed-size (and truncated Poisson) where the dataset size is needed to derive the batch size / sampling probability.

ryan112358 · 2026-07-28T21:53:11Z

+      ``expected_batch_size`` when ``None``.
+    truncated_batch_size: Optional truncation cap for Poisson sampling. When
+      set, accounting uses ``truncated_dpsgd_event``.
+    use_zcdp: If True, account with a ``ZCDpEvent`` (e.g. for discrete


Discrete gaussian is currently not fully compatable with this API or the noise_addition API, so let's drop that reference and this flag here.

Done — removed use_zcdp and the discrete-Gaussian references from the config and docs.

Merge FixedSizeDpsgdConfig into DpsgdConfig, route Poisson vs fixed-size via batch_size, parameterize Poisson like BandMF with expected_participations, and drop the use_zcdp / discrete-Gaussian references.

Ayush7614 · 2026-07-28T22:19:14Z

Thanks for the review @ryan112358 — addressed in the latest commit:

Unified DpsgdConfig and FixedSizeDpsgdConfig into a single DpsgdConfig, routed by batch_size (unset = Poisson, set = fixed-size).
Poisson path now uses expected_participations like BandMFConfig (no num_examples required under add/remove). num_examples is only required for fixed-size or truncated Poisson.
Removed use_zcdp and discrete-Gaussian references.

Happy to adjust further if needed.

ryan112358

Thx for the quick turnaround, please respond to each comment individually or mark it as resolved. (from prior review & this one)

ryan112358 · 2026-07-28T23:10:50Z

+  batch_size: int | None = None
+  num_examples: int | None = None
+  truncated_batch_size: int | None = None
+  replace: bool = False


Let's drop this kwarg

Done — dropped the replace kwarg.

ryan112358 · 2026-07-28T23:12:35Z

+      Poisson sampling.
+  """
+
+  iterations: int


In the future, we will want to support more than just Poisson / Fixed Batch. There's also Random allocation - support is being added for that in dp-accounting now. Whatever design we have here needs to be done with that future in mind

Agreed — switched the routing field to an explicit DpsgdBatchSelection enum (POISSON, FIXED) so we can extend with RANDOM_ALLOCATION (or others) once dp-accounting support lands, without another config class or a boolean/batch_size-as-sentinel design.

ryan112358 · 2026-07-28T23:13:42Z

+  l2_clip_norm: float = 1.0
+  rescale_to_unit_norm: bool = True
+  normalize_by: float = 1.0
+  batch_size: int | None = None


This can be inferred from expected_participations and num_examples so is probably not needed

Done — removed batch_size as a constructor argument. For FIXED, it is now derived as expected_participations * num_examples / iterations (exposed as a read-only property for callers that still need it).

Introduce DpsgdBatchSelection so more strategies (e.g. random allocation) can be added later, infer fixed batch size from expected_participations and num_examples, and remove the replace kwarg.

Ayush7614 · 2026-07-29T13:21:06Z

Thanks @ryan112358 — replied on each thread. Summary of the latest push:

DpsgdBatchSelection enum routes Poisson vs fixed-size (extensible for random allocation later)
Dropped replace and constructor batch_size (fixed batch size inferred from expected_participations + num_examples)
Earlier feedback already applied: unified config, expected_participations, no use_zcdp

Happy to adjust further if the enum / inference shape isn’t quite right.

Keep both Unreleased entries: DpsgdConfig (this PR) and the clipped_fun Formal Guarantees fix from google-deepmind#338.

ryan112358

Please update the PR description to disclose how you used AI for this PR. Things like how did you guide it, how did you check it's work, how much iteration did you do before you opened the PR, etc. AI usage is generally fine, but reviewing has a cost so it's helpful to know

ryan112358 · 2026-07-29T18:17:22Z

+  return clipped_grad_transform
+
+
+def _check_noise_multiplier_set(noise_multiplier: float | None) -> None:


Maybe add this to _validate.py as say "def notnull(**kwargs)" or similar

Done — added _validate.not_none(**kwargs) and use it from DpsgdConfig.make() (and BandMF’s calibrated check) instead of a one-off helper.

ryan112358 · 2026-07-29T18:21:23Z

+          f' {type(self.batch_selection)!r}.'
+      )
+
+    if (


This is adding a significant amount of bloat. Please think carefully about if it can be reduced or otherwise factored into _validate.py

Done — pulled common checks into _validate.not_none / instance_of / at_most and slimmed __post_init__ accordingly (net ~150 lines removed in this push with the example deletion).

ryan112358 · 2026-07-29T18:23:01Z

+        )
+
+  @property
+  def batch_size(self) -> int:


For poisson we have variable batch sizes, so this would have to be 'expected_batch_size'

Done — renamed the property to expected_batch_size (float). Fixed-size sampling still validates/uses an integer via a private _fixed_batch_size() helper.

ryan112358 · 2026-07-30T18:05:35Z

@@ -0,0 +1,138 @@
+# coding=utf-8


Generally we prefer smaller focused PRs and we pay close attention to "diff hygiene" and whether a change warrants the maintenance burden. Given that users can (or should be able to) just take an existing example and swap out the execution plan, I think a totally separate example is unnecessary here, so I'd remove this here align with the change to examples_guide.md

Done — removed examples/dp_sgd_execution_plan.py and the corresponding examples_guide.md / CHANGELOG references. Existing examples can swap in DpsgdConfig as needed.

Move not_none/instance_of/at_most into _validate, rename batch_size to expected_batch_size, and remove the standalone execution-plan example.

Ayush7614 · 2026-07-30T19:30:35Z

Thanks @ryan112358 — addressed the latest review in 4d8c3df:

_validate.not_none / instance_of / at_most; slimmer DpsgdConfig validation
renamed property to expected_batch_size
removed the standalone example + docs/CHANGELOG mentions
updated the PR description with an AI-usage note

Happy to adjust further if anything still feels heavy.

ryan112358

Mostly looks good, will send for an internal review next. Thanks!

ryan112358 · 2026-08-03T15:36:40Z

+    )
+    if self.noise_multiplier is not None:
+      _validate.non_negative(noise_multiplier=self.noise_multiplier)
+    _validate.instance_of(


The type checker should catch this at "compile" time, no need for runtime validation

Done — removed the runtime instance_of check (and the unused _validate.instance_of helper).

ryan112358 · 2026-08-03T15:38:38Z

+      ValueError: If ``num_examples`` is unset.
+    """
+    _validate.not_none(num_examples=self.num_examples)
+    return self.expected_participations * self.num_examples / self.iterations


This might not be an integer under FIXED

Done — clarified that expected_batch_size is always a float and need not be integral under FIXED. The integer passed to FixedBatchSampling is derived separately via _fixed_batch_size(), which still requires an exact positive integer for accounting consistency.

Probably for fixed_batch_size this should still be an integer (rounded) since by definition it has to be fixed across all steps. I think you should fold in the logic from the _fixed_batch_size helper here, and delete that method.

ryan112358 · 2026-08-03T15:45:49Z

+          f'Unsupported batch_selection={self.batch_selection!r}. Additional'
+          ' strategies will be wired as dp-accounting support lands.'
+      )
+    query_sensitivity = clipped_grad_transform(lambda: None).sensitivity(


So I think dp_accounting treats "noise_multiplier" as relative to "l2_norm_bound" which is 2x smaller than "sensitivity" under REPLACE_ONE DP. It's worth changign this to l2_norm_bound and adding an inline comment explaining why, along with a test that noise get's calibrated correctly in this setting.

Done — make() now scales noise by l2_norm_bound (with an inline comment explaining why sensitivity() would over-noise under REPLACE_ONE). Added test_fixed_size_noise_matches_l2_norm_bound_not_sensitivity to lock this in.

Remove instance_of validation, clarify expected_batch_size may be non-integral under FIXED, and calibrate Gaussian noise using l2_norm_bound so REPLACE_ONE does not over-noise vs dp_accounting.

Ayush7614 · 2026-08-03T20:05:48Z

Thanks @ryan112358 — addressed the latest review:

Dropped runtime instance_of validation
Clarified expected_batch_size may be non-integral under FIXED
Noise now uses l2_norm_bound (not sensitivity()), with a REPLACE_ONE calibration test

pytest tests/execution_plan_test.py: 22 passed.

arung54 · 2026-08-04T18:31:28Z

+    return self.expected_participations * self.num_examples / self.iterations
+
+  def _fixed_batch_size(self) -> int:
+    """Integer batch size required by ``DpsgdBatchSelection.FIXED``.


Our linter is blocking approval since this multi-line docstring doesn't have a "Returns:"

arung54 · 2026-08-04T18:39:36Z

+          batch_size=self._fixed_batch_size(),
+          replace=False,
+      )
+    if self.batch_selection is not DpsgdBatchSelection.POISSON:


Style nitpick: For readability/consistency with similar statement in other parts of this library/related libraries could you reformat this as:

if self.batch_selection is DpsgdBatchSelection.FIXED: ... elif self.batch_selection is DpsgdBatchSelection.POISSON: ... else: raise NotImplementedError(...)

ryan112358 · 2026-08-05T13:04:46Z

+      ValueError: If ``num_examples`` is unset.
+    """
+    _validate.not_none(num_examples=self.num_examples)
+    return self.expected_participations * self.num_examples / self.iterations


Probably for fixed_batch_size this should still be an integer (rounded) since by definition it has to be fixed across all steps. I think you should fold in the logic from the _fixed_batch_size helper here, and delete that method.

ryan112358 · 2026-08-05T13:06:08Z

+          batch_size=self._fixed_batch_size(),
+          replace=False,
+      )
+    if self.batch_selection is not DpsgdBatchSelection.POISSON:


ryan112358 · 2026-08-05T13:06:23Z

+    if accountant_fn is None:
+      if self.batch_selection is DpsgdBatchSelection.FIXED:
+
+        def _rdp_accountant(


Not sure what this does, can it be deleted?

ryan112358 · 2026-08-05T13:07:20Z

+          iterations=self.iterations,
+          cycle_length=1,
+          truncated_batch_size=self.truncated_batch_size,
+          partition_type=(


Since we have cycle_length=1, partition_type is a no-op and can be removed.

ryan112358 · 2026-08-11T16:01:58Z

Hi @Ayush7614 , I have mulled over this change a bit and have a few thoughts.

DpSGDConfig is a special case of BandMFConfig in most cases, with the exception of fixed batch sampling (which is not SOTA). For that reason, the maintenance burden of this extra code is likely not worth the delta in capabilities it brings to Jax privacy currently.

With that said, if you are able to land a couple of other PRs before this one, then this class can feature some nice additional capabilities. Specifically, if you can generalize DpSgd accounting to support a list of sampling probabilities and a list of batch sizes respectively, along with a list of noise multipliers, this would provide enough of a difference to potentially warrant bringing this in.

If you want to work on this, please open an issue first with your proposed change including what files it affects and how.

Add DpsgdConfig and FixedSizeDpsgdConfig execution plans.

a60fa01

Provide Tier-2 DP-SGD configs that couple Poisson or fixed-size batch selection, clipping, Gaussian noise, and matching accounting events so users can run vanilla DP-SGD without hand-wiring Tier-3 components.

github-actions Bot reviewed Jul 28, 2026

View reviewed changes

ryan112358 reviewed Jul 28, 2026

View reviewed changes

Ayush7614 added 2 commits July 29, 2026 03:27

Fix pyink formatting and unused import in execution_plan tests.

d0d953f

Address review: unify DpsgdConfig and use expected_participations.

30f7a49

Merge FixedSizeDpsgdConfig into DpsgdConfig, route Poisson vs fixed-size via batch_size, parameterize Poisson like BandMF with expected_participations, and drop the use_zcdp / discrete-Gaussian references.

Ayush7614 force-pushed the feat/dpsgd-execution-plan-configs branch from 46c3978 to 30f7a49 Compare July 28, 2026 22:16

ryan112358 reviewed Jul 28, 2026

View reviewed changes

Ayush7614 added 2 commits July 29, 2026 18:49

Address second review: enum routing, drop batch_size and replace.

732acde

Introduce DpsgdBatchSelection so more strategies (e.g. random allocation) can be added later, infer fixed batch size from expected_participations and num_examples, and remove the replace kwarg.

Apply pyink formatting to execution_plan.py.

7ce25dd

Ayush7614 changed the title ~~Add DpsgdConfig and FixedSizeDpsgdConfig execution plans~~ Add DpsgdConfig execution plan for Poisson and fixed-size DP-SGD Jul 29, 2026

Merge upstream/main; resolve CHANGELOG.md conflict.

91e534b

Keep both Unreleased entries: DpsgdConfig (this PR) and the clipped_fun Formal Guarantees fix from google-deepmind#338.

ryan112358 reviewed Jul 30, 2026

View reviewed changes

Address review: slim DpsgdConfig validation and drop extra example.

4d8c3df

Move not_none/instance_of/at_most into _validate, rename batch_size to expected_batch_size, and remove the standalone execution-plan example.

ryan112358 requested changes Aug 3, 2026

View reviewed changes

Address review: drop runtime type check; scale noise by l2_norm_bound.

da75ebe

Remove instance_of validation, clarify expected_batch_size may be non-integral under FIXED, and calibrate Gaussian noise using l2_norm_bound so REPLACE_ONE does not over-noise vs dp_accounting.

ryan112358 approved these changes Aug 4, 2026

View reviewed changes

arung54 reviewed Aug 4, 2026

View reviewed changes

ryan112358 reviewed Aug 5, 2026

View reviewed changes



		@dataclasses.dataclass(frozen=True, kw_only=True)
		class DpsgdConfig:

		return clipped_grad_transform


		def _check_noise_multiplier_set(noise_multiplier: float \| None) -> None:

Uh oh!

Conversation

Ayush7614 commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

AI usage disclosure

Test plan

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

ryan112358 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Ayush7614 commented Jul 28, 2026

Uh oh!

ryan112358 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Ayush7614 commented Jul 29, 2026

Uh oh!

ryan112358 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Ayush7614 commented Jul 30, 2026

Uh oh!

ryan112358 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Ayush7614 commented Jul 28, 2026 •

edited

Loading

arung54 Aug 4, 2026 •

edited

Loading