jax_privacy.execution_plan

Module for defining DP Execution Plans.

API Stability: - DPExecutionPlan: 9/10 – Stable, backwards-compatible changes possible. - BandMFExecutionPlanConfig: 7/10 – Mostly stable, minor changes possible.

# Writing General-Purpose DP Training Loops via DPExecutionPlan

This module introduces the DPExecutionPlan, an object designed to encapsulate the core components of a differentially private (DP) mechanism. The primary aim is to simplify the process of constructing and applying DP mechanisms by packaging these components cohesively. A key benefit is the assurance that, when used correctly, the combination of these components will achieve the stated DP properties.

The design is framework-agnostic, specifying the essential pillars of a DP mechanism—such as batch selection and noise addition—without tightly coupling them to a specific training loop. Each component is exposed through a simple, well-documented API, allowing for flexible integration into various frameworks or direct use with JAX.

By programming against our DPExecutionPlan interface, it is easy to swap out different components or entire mechanisms without changing the core training loop logic. While the components are designed to work together, they can also be used selectively. For instance, a researcher might choose to use only the noise addition component if their dataset doesn’t support the efficient random access required by the batch selection strategy. Though this might invalidate the formal DP guarantee, it can still be valuable for research or when a heuristic quantification of privacy is acceptable. Ultimately, DPExecutionPlan aims to free users to concentrate on their training pipeline setup, rather than on the intricacies of correctly assembling DP components to achieve a desired privacy guarantee.

# Selecting and using a DPExecutionPlan

Constructors for these plans are highly configurable, offering access to the full capabilities of the underlying components while also providing sensible defaults. Our primary entry point is currently BandMFExecutionPlanConfig, although more will become available in the future when we feel the API has stabilized.

Classes

BandMFExecutionPlanConfig(*, iterations, ...)	Configuration for an Amplified BandMF-based DPExecutionPlan.
DPExecutionPlan(clipped_grad, ...)	Class for defining a DP execution plan.
PerformanceFlags([dtype, noise_seed, ...])	Performance-only flags that do not affect mechanism behavior or privacy.