jax_privacy.execution_plan.BandMFExecutionPlanConfig

class jax_privacy.execution_plan.BandMFExecutionPlanConfig(*, iterations, strategy, noise_multiplier=None, l2_clip_norm=1.0, rescale_to_unit_norm=True, normalize_by=1.0, sampling_prob=1.0, truncated_batch_size=None, num_examples=None, column_normalize=False)[source]

Bases: object

Configuration for an Amplified BandMF-based DPExecutionPlan.

This config is designed to be fully serializable, defined in terms of simple types. The config can be created with or without a noise_multiplier. If created without one, call calibrate() to obtain a new config with a noise_multiplier calibrated to a target (epsilon, delta) guarantee.

The expected batch size of the batch selection strategy is (num_examples / num_bands) * sampling_prob. num_examples may or may not be passed to this config. It should only be passed if it is considered a public, non-sensitive quantity (i.e., when using zero-out adjacency rather than add-remove).

Example Usage (Calibrate from epsilon/delta):

>>> config = BandMFExecutionPlanConfig.default(
...   num_bands=1, iterations=1000, sampling_prob=0.1,
... ).calibrate(epsilon=1.0, delta=1e-5)

Example Usage (Direct noise_multiplier):

>>> config = BandMFExecutionPlanConfig.default(
...   num_bands=1, iterations=1000, sampling_prob=0.1,
...   noise_multiplier=1.0,
... )

Example Usage (BandMF with custom strategy):

>>> config = BandMFExecutionPlanConfig(
...   strategy=np.array([1.0, 0.5, 0.2]),
...   iterations=1000, sampling_prob=0.4,
... ).calibrate(epsilon=1.0, delta=1e-5)

References: https://arxiv.org/abs/2306.08153 and https://arxiv.org/abs/2405.15913

Variables:

noise_multiplier – The ratio of noise standard deviation to the query sensitivity. The actual noise stddev is determined by this value, the query sensitivity, and the strategy matrix column norm. If not set, use calibrate() to automatically determine it from target (epsilon, delta) privacy parameters.
iterations – The number of iterations the mechanism is defined for. Tip: Set this to be a multiple of num_bands for the best utility.
strategy – The toeplitz coefficeints of the BandMF strategy matrix.
l2_clip_norm – The maximum L2 norm of the per-example gradients.
rescale_to_unit_norm – Divide the clipped gradient by the l2_clip_norm.
normalize_by – Divide the sum-of-clipped gradients by this value.
sampling_prob – The Poisson sampling probability for each example in a group.
truncated_batch_size – If using truncated Poisson sampling, the maximum batch size to truncate to. If set, the plan.batch_selection_strategy will always return batches of size at most truncated_batch_size, and accounting will be based on truncated Poisson sampling (http://arxiv.org/html/2508.15089).
num_examples – The number of examples in the dataset. Required when truncated_batch_size is set. Only set when the dataset size is considered public, non-sensitive information (e.g., when using zero-out adjacency rather than add-remove). If specified, the dataset will be partitioned using batch_selection.PartitionType.EQUAL_SPLIT, and otherwise it will be partitioned using batch_selection.PartitionType.INDEPENDENT.
column_normalize – Whether to column-normalize the strategy matrix.

Methods

`__init__`
`calibrate`	Returns a new config with a calibrated noise_multiplier.
`default`	Returns a BandMFExecutionPlanConfig with an RMSE-optimized strategy.
`make`	Returns a DP execution plan for the given BandMF mechanism config.

Attributes

`column_normalize`
`l2_clip_norm`
`noise_multiplier`
`normalize_by`
`num_examples`
`rescale_to_unit_norm`
`sampling_prob`
`truncated_batch_size`
`iterations`
`strategy`

iterations: int

strategy: Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]

noise_multiplier: float | None = None

l2_clip_norm: float = 1.0

rescale_to_unit_norm: bool = True

normalize_by: float = 1.0

sampling_prob: float = 1.0

truncated_batch_size: int | None = None

num_examples: int | None = None

column_normalize: bool = False

calibrate(*, epsilon, delta, tol=None, accountant_fn=<class 'dp_accounting.pld.pld_privacy_accountant.PLDAccountant'>)[source]

Returns a new config with a calibrated noise_multiplier.

Parameters:

epsilon (float) – The target privacy budget.
delta (float) – The target privacy failure probability.
tol (float | None) – The tolerance in noise_multiplier space for the calibration binary search. Defaults to 1e-6 if not specified.
accountant_fn (Callable[[NeighboringRelation], PrivacyAccountant]) – A function that returns a fresh privacy accountant used for calibration given a neighboring relation. Defaults to PLDAccountant.

Return type:

BandMFExecutionPlanConfig

Returns:

A new BandMFExecutionPlanConfig with calibrated noise_multiplier.

classmethod default(num_bands, iterations, strategy_optimization_steps=500, **kwargs)[source]

Returns a BandMFExecutionPlanConfig with an RMSE-optimized strategy.

See BandMFExecutionPlanConfig for the full list of keyword arguments.

Parameters:

num_bands (int) – The number of bands in the strategy matrix.
iterations (int) – The number of iterations the mechanism is defined for.
strategy_optimization_steps (int) – The number of optimization steps to use for the strategy matrix.
**kwargs – Keyword arguments to pass to BandMFExecutionPlanConfig.

Return type:

BandMFExecutionPlanConfig

Returns:

A BandMFExecutionPlanConfig with an RMSE-optimized strategy.

make(performance_flags=None)[source]

Returns a DP execution plan for the given BandMF mechanism config.

Parameters:: performance_flags (PerformanceFlags | None) – Optional performance flags that control implementation details such as dtype, sharding, and microbatching. If None, default values are used for all performance flags.
Return type:: DPExecutionPlan
Returns:: A DPExecutionPlan configured from this config and the given performance flags.
Raises:: ValueError – If noise_multiplier has not been set.