Noise Injection

Opake's noise mechanisms live next to the training paradigm they support: opake.dpsgd.noise for independent Gaussian-family noise, and opake.dpftrl.noise for matrix-factorization (correlated) noise. The base NoiseState type that both build on lives in opake.types.

Overview

After clipping gradients, DP-SGD requires adding noise proportional to the clip norm and noise multiplier. The noise obscures individual contributions, providing the actual privacy guarantee.

Opake provides several noise mechanisms:

Independent Noise (DP-SGD)

• gaussian_noise() — Gaussian noise.

Correlated Noise (DP-FTRL / Matrix Factorization)

• mf_gaussian_noise() — Unified correlated noise dispatcher. Takes a strategy object and creates the corresponding noise mechanism.

Strategy factories (passed to mf_gaussian_noise()):

• band_mf_strategy() — BandMF banded Toeplitz correlated noise
• blt_strategy() — Buffered Linear Toeplitz (BLT) correlated noise
• lambda_cgd_strategy() — DP-λCGD correlated noise (PRNG replay, zero extra memory)
• bisr_strategy() — BISR (Banded Inverse Square Root) correlated noise
• identity_strategy() — Identity (DP-SGD via MF API, easy to swap)

All noise functions return (noise_fn, state), where noise_fn(grads, state) -> (noisy_grads, new_state). When private second moments are enabled, the noisy value is a SecondMomentNoiseOutput carrying both noisy_grads and noisy_squared_grads.

State Classes

• NoiseState — Abstract base class for all noise state types. Defines _step_counter and _rng_key.
• SecondMomentNoiseOutput — Paired first/squared-gradient output for private second moments.
• GaussianNoiseState — State for gaussian_noise(). Holds step counter and RNG key.
• MFNoiseState — State for mf_gaussian_noise(). Holds internal correlation state, step counter, and RNG key.
• SecondMomentMFNoiseState — Paired-stream MF state used by mf_gaussian_noise(..., second_moment_strategy=...).

Distributed Sync Helpers

Use sync() from opake.distributed to validate noise state consistency across ranks. It auto-dispatches based on type:

• sync(GaussianNoiseState) — Validate RNG key and step counter match across ranks.
• sync(MFNoiseState) — Validate MF noise state matches across ranks.

See also: Noise Addition User Guide

Paired second-moment release

When clipped_grad(..., second_moment=True) produces a SecondMomentClippingOutput, both gaussian_noise (DP-SGD) and mf_gaussian_noise(..., second_moment_strategy=...) (DP-FTRL) consume it and emit a SecondMomentNoiseOutput with paired noise on both streams.

The runtime σ allocation is sensitivity-proportional and works identically for scalar max_norm and per-group PerGroup max_norm. For each group \(g\) (single-group K=1 for scalar clipping), let \(\Delta^{(1)}_g\) be the first-stream per-record bound (C_g / n for averaged clipping, or ζ · ‖C₁‖ strategy-amplified for DP-FTRL) and \(\Delta^{(2)}_g\) the second-stream per-record bound (C_g² / n derived from the same clip; ζ² · ‖C₂‖ for DP-FTRL). Then with noise multiplier \(\text{nm}\), set

\[ S := \sum_h \bigl(\Delta^{(1)}_h + \Delta^{(2)}_h\bigr),\qquad \sigma^{(1)}_g := \text{nm}\sqrt{\Delta^{(1)}_g\,S},\qquad \sigma^{(2)}_g := \text{nm}\sqrt{\Delta^{(2)}_g\,S}. \]

This satisfies the joint Mahalanobis budget with equality:

\[ \sum_{g,i} \Bigl(\frac{\Delta^{(i)}_g}{\sigma^{(i)}_g}\Bigr)^2 = \frac{1}{\text{nm}^2}. \]

So the paired release has the same PLD as a single sensitivity-1 Gaussian release at the joint effective multiplier passed to paired_noise_stddevs (see its docstring: no-op for DP-SGD; divide by ‖C₁‖ for DP-FTRL). Accounting is plain gaussian(nm) for DP-SGD and the underlying mf_gaussian(nm, …) for DP-FTRL; there is no separate transformation wrapper and no ρ knob.

mf_gaussian_noise accepts scalar or PerGroup max_norm on ClippedPytree inputs. Single-stream IID stddevs for PerGroup bounds match the MSE-optimal allocation from ClippedPytree.noise_stddev_for (the same Mahalanobis allocation as gaussian_noise). Leaf→group keys are optree ParamPath tuples (build them with per_group(params, …)); nested parameter trees are supported. Trainer/examples keep flat named_parameters by choice.

Gaussian

gaussian_noise(*, noise_multiplier, key, compute_dtype=torch.float32)

Create a Gaussian noise function with immutable state.

Returns (noise_fn, state) where noise_fn adds calibrated Gaussian noise to :class:opake.types.ClippedPytree inputs and returns updated state. The realized standard deviation is noise_multiplier * clipped.max_norm. The output is a :class:opake.types.NoisedPytree carrying that realized noise_stddev metadata.

The noise function uses exactly the key you provide — no auto-detection of distributed state. For synchronized noise in DDP, pass the same key on every rank. For independent noise, derive a per-rank key::

from opake.random import key, fold_in
my_key = fold_in(key(42), rank)  # different noise per rank
noise_fn, state = gaussian_noise(noise_multiplier=1.1, key=my_key)

Parameters:

Name	Type	Description	Default
noise_multiplier	float	Gaussian noise multiplier. The realized standard deviation is noise_multiplier * clipped.max_norm.	required
key	RngKey	Explicit RNG key for deterministic, functional randomness. Same key on all ranks → same noise (synchronized). fold_in(key, rank) → independent noise per rank.	required
compute_dtype	dtype	One of torch.float16, torch.bfloat16, torch.float32, or torch.float64 for sampling and addition. Defaults to torch.float32. Output is cast back to the input dtype.	float32

Returns:

Type	Description
GaussianNoiseFn	A tuple (noise_fn, state) where:
GaussianNoiseState	noise_fn(clipped_grads, state) -> (noisy_grads, new_state)
tuple[GaussianNoiseFn, GaussianNoiseState]	state is a :class:GaussianNoiseState

Example

import torch from opake.types import clipped from opake.dpsgd.noise import gaussian_noise from opake.random import key

noise_fn, state = gaussian_noise(noise_multiplier=1.1, key=key(42)) grads = torch.zeros(10) noisy_grads, state = noise_fn(clipped(grads, max_norm=1.0), state)

Source code in packages/opake-dpsgd/src/opake/api/dpsgd/noise/_gaussian.py

def gaussian_noise(
    *,
    noise_multiplier: float,
    key: RngKey,
    compute_dtype: torch.dtype = torch.float32,
) -> tuple[
    GaussianNoiseFn,
    GaussianNoiseState,
]:
    """Create a Gaussian noise function with immutable state.

    Returns ``(noise_fn, state)`` where ``noise_fn`` adds calibrated Gaussian
    noise to :class:`opake.types.ClippedPytree` inputs and returns updated
    state.  The realized standard deviation is
    ``noise_multiplier * clipped.max_norm``.  The output is a
    :class:`opake.types.NoisedPytree` carrying that realized
    ``noise_stddev`` metadata.

    The noise function uses exactly the ``key`` you provide — no auto-detection
    of distributed state. For synchronized noise in DDP, pass the same key on
    every rank. For independent noise, derive a per-rank key::

        from opake.random import key, fold_in
        my_key = fold_in(key(42), rank)  # different noise per rank
        noise_fn, state = gaussian_noise(noise_multiplier=1.1, key=my_key)

    Args:
        noise_multiplier: Gaussian noise multiplier.  The realized standard
            deviation is ``noise_multiplier * clipped.max_norm``.
        key: Explicit RNG key for deterministic, functional randomness.
            Same key on all ranks → same noise (synchronized).
            ``fold_in(key, rank)`` → independent noise per rank.
        compute_dtype: One of ``torch.float16``, ``torch.bfloat16``,
            ``torch.float32``, or ``torch.float64`` for sampling and addition.
            Defaults to ``torch.float32``. Output is cast back to the input dtype.

    Returns:
        A tuple ``(noise_fn, state)`` where:

        - ``noise_fn(clipped_grads, state) -> (noisy_grads, new_state)``
        - ``state`` is a :class:`GaussianNoiseState`

    Example:
        >>> import torch
        >>> from opake.types import clipped
        >>> from opake.dpsgd.noise import gaussian_noise
        >>> from opake.random import key
        >>>
        >>> noise_fn, state = gaussian_noise(noise_multiplier=1.1, key=key(42))
        >>> grads = torch.zeros(10)
        >>> noisy_grads, state = noise_fn(clipped(grads, max_norm=1.0), state)
    """
    resolved_noise_multiplier = _resolve_noise_multiplier(noise_multiplier)
    _validate_compute_dtype(compute_dtype)

    if not isinstance(key, RngKey):
        raise InputTypeError(*(f"key must be RngKey, got {type(key)}",))

    state = GaussianNoiseState(
        _step_counter=0,
        _rng_key=key,
    )

    def _sample(
        center: torch.Tensor, std: float, generator: torch.Generator
    ) -> torch.Tensor:
        """Sample unbounded N(center, std²) with the supplied CPU generator."""
        if std == 0:
            return center

        in_dtype = center.dtype
        noise = torch.randn(
            center.shape,
            dtype=compute_dtype,
            generator=generator,
        ).to(device=center.device)
        center_c = center.to(compute_dtype) if in_dtype != compute_dtype else center
        sample = center_c + std * noise

        return sample.to(dtype=in_dtype) if in_dtype != compute_dtype else sample

    def _add_noise_tree(grads, effective_stddev, generator):
        _validate_noise_stddev(effective_stddev)

        # Per-group σ: look up each leaf by optree ParamPath.
        if isinstance(effective_stddev, PerGroup):
            import optree

            from opake.api.engine.pytree import tree_flatten_with_paths

            paths, leaves, treedef = tree_flatten_with_paths(grads)
            noised_leaves = []
            for path, tensor in zip(paths, leaves, strict=True):
                if not isinstance(tensor, torch.Tensor):
                    raise InputTypeError(
                        *(
                            "gaussian_noise with PerGroup stddev expects tensor "
                            f"leaves; got {type(tensor).__name__} at path {path!r}.",
                        )
                    )
                group_std = effective_stddev.for_path(path)
                noised_leaves.append(_sample(tensor, group_std, generator))
            return optree.tree_unflatten(treedef, noised_leaves)

        return tree_map(lambda t: _sample(t, effective_stddev, generator), grads)

    def _clipped_stddev(clipped: ClippedPytree) -> float | PerGroup:
        if isinstance(clipped.max_norm, PerGroup):
            return per_group_noise_stddev(clipped.max_norm, resolved_noise_multiplier)
        # Non-private run: force std = 0 before the ``nm * max_norm`` product so
        # a disabled clip (max_norm = +inf) gives ``0`` rather than ``0 * inf =
        # NaN`` (the ``std == 0`` short-circuit in ``_sample`` then fires).
        effective = (
            0.0
            if resolved_noise_multiplier == 0.0
            else resolved_noise_multiplier * clipped.max_norm
        )
        _validate_noise_stddev(effective)
        return effective

    def _add_paired(
        clipped_input: SecondMomentClippingOutput, st: GaussianNoiseState
    ) -> tuple[SecondMomentNoiseOutput, GaussianNoiseState]:
        first_clipped, second_clipped, first_stddev, second_stddev = (
            resolve_paired_clipped(
                clipped_input,
                noise_multiplier=resolved_noise_multiplier,
            )
        )
        # Two independent noise streams, both beneath this mechanism's root
        # so they cannot collide with each other or with the single-stream
        # derivation (``fold_in(_rng_key, GAUSSIAN_STREAM_FOLD, step)``).
        first_step_key = rng_fold_in(
            st._rng_key,
            GAUSSIAN_STREAM_FOLD,
            PAIRED_FIRST_STREAM_FOLD,
            st._step_counter,
        )
        second_step_key = rng_fold_in(
            st._rng_key,
            GAUSSIAN_STREAM_FOLD,
            PAIRED_SECOND_STREAM_FOLD,
            st._step_counter,
        )
        noisy_grads = _add_noise_tree(
            first_clipped.pytree,
            first_stddev,
            generator_from_key(first_step_key),
        )
        noisy_squared = _add_noise_tree(
            second_clipped.pytree,
            second_stddev,
            generator_from_key(second_step_key),
        )
        next_state = GaussianNoiseState(
            _step_counter=st._step_counter + 1,
            _rng_key=st._rng_key,
        )
        return (
            SecondMomentNoiseOutput(
                NoisedPytree(
                    pytree=noisy_grads,
                    max_norm=first_clipped.max_norm,
                    noise_stddev=first_stddev,
                ),
                NoisedPytree(
                    pytree=noisy_squared,
                    max_norm=second_clipped.max_norm,
                    noise_stddev=second_stddev,
                ),
            ),
            next_state,
        )

    def noise_fn(grads, st):
        """Add Gaussian noise to a clipped pytree (or paired stream)."""
        with record_function("opake::gaussian_noise"):
            return _noise_fn_impl(grads, st)

    def _noise_fn_impl(grads, st):
        if isinstance(grads, SecondMomentClippingOutput):
            return _add_paired(grads, st)

        next_state = GaussianNoiseState(
            _step_counter=st._step_counter + 1,
            _rng_key=st._rng_key,
        )

        if isinstance(grads, NoisedPytree):
            raise InputTypeError(
                *(
                    "gaussian_noise expects ClippedPytree inputs, not NoisedPytree "
                    "values that have already passed through a noise mechanism.",
                )
            )

        if not isinstance(grads, ClippedPytree):
            raise InputTypeError(
                *(
                    "gaussian_noise expects ClippedPytree inputs. Wrap manual "
                    "values with opake.types.clipped(...).",
                )
            )

        effective_stddev = _clipped_stddev(grads)
        step_key = rng_fold_in(st._rng_key, GAUSSIAN_STREAM_FOLD, st._step_counter)
        noisy_tree = _add_noise_tree(
            grads.pytree,
            effective_stddev,
            generator_from_key(step_key),
        )
        return (
            NoisedPytree(
                pytree=noisy_tree,
                max_norm=grads.max_norm,
                noise_stddev=effective_stddev,
            ),
            next_state,
        )

    return noise_fn, state

Matrix Factorization Noise

Dispatcher

mf_gaussian_noise(grad_template, strategy, *, n_steps, min_sep=1, max_participations=None, noise_multiplier, key, compute_dtype=torch.float32, second_moment_strategy=None)

Create a correlated noise mechanism for the given MF strategy.

The polymorphic :meth:MfStrategy.streaming_matrix query selects the appropriate streaming primitive at the supplied amplification context; :class:LambdaCgdStrategy uses a PRNG-replay path (no streaming matrix is materialized).

The returned noise_fn dispatches on its input type:

• ClippedPytree → NoisedPytree (single-stream noise).
• SecondMomentClippingOutput → SecondMomentNoiseOutput (paired-stream noise; only when second_moment_strategy was supplied at construction).

The paired-stream release uses the sensitivity-proportional joint Mahalanobis allocation (:func:~opake.api.engine.noise_allocation.paired_noise_stddevs) with the MF translation nm / ‖C₁‖ as the joint effective multiplier so the joint PLD matches the single-stream MF Gaussian accountant at (noise_multiplier, ‖C₁‖). PerGroup max_norm is supported on both streams.

Parameters:

Name	Type	Description	Default
grad_template	Any	Pytree with same structure/shapes as gradients.	required
strategy	MfStrategy	MF strategy recipe from one of the factory functions.	required
n_steps	int	Total training steps (the amplifier's horizon).	required
min_sep	int	Minimum separation between participations (default 1).	1
max_participations	int | None	Max participations per example (None ⇒ n_steps).	None
noise_multiplier	float	Gaussian noise multiplier. The clipped-gradient max_norm is read from each ClippedPytree input.	required
key	RngKey	Explicit RNG key for deterministic randomness.	required
compute_dtype	dtype	Dtype used for the underlying torch.randn and the linear-combination arithmetic. Defaults to torch.float32 to match :func:opake.dpsgd.noise.gaussian_noise — sampling Gaussians in bf16/fp16 introduces discretization that distorts the noise distribution. Type stability on the public boundary is preserved: the input pytree's dtype is matched on output (input upcast to compute_dtype, noise added, downcast at return).	float32
second_moment_strategy	MfStrategy | None	Optional explicit strategy recipe for the squared-gradient stream.	None

Returns:

Type	Description
tuple[Callable[..., tuple[Any, MFNoiseState | SecondMomentMFNoiseState]], MFNoiseState | SecondMomentMFNoiseState]	A tuple (noise_fn, state) for the training loop.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/noise/_mf_gaussian_noise.py

def mf_gaussian_noise(
    grad_template: Any,
    strategy: MfStrategy,
    *,
    n_steps: int,
    min_sep: int = 1,
    max_participations: int | None = None,
    noise_multiplier: float,
    key: RngKey,
    compute_dtype: torch.dtype = torch.float32,
    second_moment_strategy: MfStrategy | None = None,
) -> tuple[
    Callable[..., tuple[Any, MFNoiseState | SecondMomentMFNoiseState]],
    MFNoiseState | SecondMomentMFNoiseState,
]:
    """Create a correlated noise mechanism for the given MF strategy.

    The polymorphic :meth:`MfStrategy.streaming_matrix` query selects the
    appropriate streaming primitive at the supplied amplification
    context; :class:`LambdaCgdStrategy` uses a PRNG-replay path (no
    streaming matrix is materialized).

    The returned ``noise_fn`` dispatches on its input type:

    - ``ClippedPytree`` → ``NoisedPytree`` (single-stream noise).
    - ``SecondMomentClippingOutput`` → ``SecondMomentNoiseOutput``
      (paired-stream noise; only when ``second_moment_strategy`` was
      supplied at construction).

    The paired-stream release uses the sensitivity-proportional joint
    Mahalanobis allocation
    (:func:`~opake.api.engine.noise_allocation.paired_noise_stddevs`)
    with the MF translation ``nm / ‖C₁‖`` as the joint effective
    multiplier so the joint PLD matches the single-stream MF Gaussian
    accountant at ``(noise_multiplier, ‖C₁‖)``.  ``PerGroup``
    ``max_norm`` is supported on both streams.

    Args:
        grad_template: Pytree with same structure/shapes as gradients.
        strategy: MF strategy recipe from one of the factory functions.
        n_steps: Total training steps (the amplifier's horizon).
        min_sep: Minimum separation between participations (default 1).
        max_participations: Max participations per example
            (``None`` ⇒ ``n_steps``).
        noise_multiplier: Gaussian noise multiplier.  The clipped-gradient
            ``max_norm`` is read from each ``ClippedPytree`` input.
        key: Explicit RNG key for deterministic randomness.
        compute_dtype: Dtype used for the underlying ``torch.randn`` and the
            linear-combination arithmetic.  Defaults to ``torch.float32`` to
            match :func:`opake.dpsgd.noise.gaussian_noise` — sampling
            Gaussians in bf16/fp16 introduces discretization that distorts
            the noise distribution. Type stability on the public boundary is
            preserved: the input pytree's dtype is matched on output (input
            upcast to ``compute_dtype``, noise added, downcast at return).
        second_moment_strategy: Optional explicit strategy recipe for
            the squared-gradient stream.

    Returns:
        A tuple ``(noise_fn, state)`` for the training loop.
    """
    resolved_noise_multiplier = _resolve_noise_multiplier(noise_multiplier)
    if not isinstance(key, RngKey):
        raise InputTypeError(*(f"key must be RngKey, got {type(key)}",))

    if second_moment_strategy is not None:
        return make_second_moment_mf_noise(
            grad_template,
            strategy,
            second_moment_strategy,
            n_steps=n_steps,
            min_sep=min_sep,
            max_participations=max_participations,
            noise_multiplier=resolved_noise_multiplier,
            key=key,
            compute_dtype=compute_dtype,
        )

    raw_noise_fn, raw_state, row_l2_at = _make_raw_mf_noise(
        grad_template,
        strategy,
        n_steps=n_steps,
        min_sep=min_sep,
        max_participations=max_participations,
        key=key,
        compute_dtype=compute_dtype,
    )

    def noise_fn(
        clipped_grads: Any,
        st: MFNoiseState,
    ) -> tuple[NoisedPytree, MFNoiseState]:
        with record_function("opake::mf_gaussian_noise"):
            return _noise_fn_impl(clipped_grads, st)

    def _noise_fn_impl(
        clipped_grads: Any,
        st: MFNoiseState,
    ) -> tuple[NoisedPytree, MFNoiseState]:
        if isinstance(clipped_grads, SecondMomentClippingOutput):
            raise InputTypeError(
                *(
                    "mf_gaussian_noise was constructed without `second_moment_strategy` "
                    "and cannot consume SecondMomentClippingOutput inputs.  Either "
                    "pass a single-stream ClippedPytree, or rebuild the noise "
                    "function with `second_moment_strategy=...`.",
                )
            )
        clipped_grads = _expect_clipped(clipped_grads, op="mf_gaussian_noise")
        max_norm = _validate_constant_max_norm(
            clipped_grads, st._first_max_norm, op="mf_gaussian_noise"
        )
        if isinstance(max_norm, PerGroup):
            base_stddev = per_group_noise_stddev(max_norm, resolved_noise_multiplier)
        else:
            # Non-private run: force base σ = 0 before the ``nm * max_norm``
            # product so a disabled clip (max_norm = +inf) gives ``0`` rather
            # than ``0 * inf = NaN``.
            base_stddev = (
                0.0
                if resolved_noise_multiplier == 0.0
                else resolved_noise_multiplier * max_norm
            )
        noisy_tree, new_state = raw_noise_fn(
            clipped_grads.pytree,
            st,
            stddev=base_stddev,
        )
        # Realized per-step σ = base σ · ‖row_t(C^-1)‖.  ``base_stddev``
        # is scalar or :class:`PerGroup`; both broadcast with the scalar
        # ``row_l2`` via ``PerGroup.__mul__`` / float multiplication.
        # Step we just applied noise to is ``st._step_counter`` (the
        # pre-increment value); ``raw_noise_fn`` advances counter inside.
        row_l2 = row_l2_at(st._step_counter)
        realized_stddev = base_stddev * row_l2
        sync_fp = mf_per_group_sync_fingerprint_for_latch(st, max_norm)
        return (
            NoisedPytree(
                pytree=noisy_tree,
                max_norm=clipped_grads.max_norm,
                noise_stddev=realized_stddev,
            ),
            replace(
                new_state,
                _first_max_norm=max_norm,
                _first_max_norm_sync_fingerprint=sync_fp,
            ),
        )

    return noise_fn, raw_state

Strategies

band_mf_strategy(*, bands, momentum=1.0, lr_schedule=None)

Create a BandMF strategy recipe.

The factory only validates static recipe knobs. The Toeplitz optimization happens lazily when a strategy method is first called with an amplifier-supplied n_steps.

Parameters:

Name	Type	Description	Default
bands	int	Number of bands in the Toeplitz matrix (>= 1).	required
momentum	float	Polyak momentum coefficient (default 1.0 = prefix-sum).	1.0
lr_schedule	Schedule | None	Optional :data:opake.scheduling.types.Schedule (Callable[[int], float]). Materialised at [0, n_steps) when the strategy first sees the amplifier's n_steps. Custom schedules must be deterministic, side-effect-free, and remain immutable for the strategy's lifetime.	None

Returns:

Name	Type	Description
A	BandMfStrategy	class:BandMfStrategy recipe.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/noise/_band_mf.py

def band_mf_strategy(
    *,
    bands: int,
    momentum: float = 1.0,
    lr_schedule: Schedule | None = None,
) -> BandMfStrategy:
    """Create a BandMF strategy recipe.

    The factory only validates static recipe knobs.  The Toeplitz
    optimization happens lazily when a strategy method is first called
    with an amplifier-supplied ``n_steps``.

    Args:
        bands: Number of bands in the Toeplitz matrix (>= 1).
        momentum: Polyak momentum coefficient (default 1.0 = prefix-sum).
        lr_schedule: Optional :data:`opake.scheduling.types.Schedule`
            (``Callable[[int], float]``).  Materialised at ``[0, n_steps)``
            when the strategy first sees the amplifier's ``n_steps``. Custom
            schedules must be deterministic, side-effect-free, and remain
            immutable for the strategy's lifetime.

    Returns:
        A :class:`BandMfStrategy` recipe.
    """
    # Recipe bounds live in ``BandMfStrategy.__post_init__``.
    return BandMfStrategy(
        bands=bands,
        momentum=momentum,
        lr_schedule=lr_schedule,
    )

blt_strategy(*, max_buffers=10, momentum=1.0, lr_schedule=None)

Create a BLT (Buffered Linear Toeplitz) strategy recipe.

The L-BFGS optimization runs lazily inside the strategy methods when an amplifier supplies the participation context (n_steps, min_sep, max_participations) — once per recipe per context, via a module-level cache shared with the noise side.

Parameters:

Name	Type	Description	Default
max_buffers	int	Maximum BLT buffer count for the optimizer (default 10).	10
momentum	float	Polyak momentum (default 1.0 = prefix-sum workload).	1.0
lr_schedule	Schedule | None	Optional :data:opake.scheduling.types.Schedule (Callable[[int], float]). Materialised at [0, n_steps) when the strategy first sees the amplifier's n_steps. Custom schedules must be deterministic, side-effect-free, and remain immutable for the strategy's lifetime.	None

Returns:

Name	Type	Description
A	BltStrategy	class:BltStrategy recipe.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/noise/_blt.py

def blt_strategy(
    *,
    max_buffers: int = 10,
    momentum: float = 1.0,
    lr_schedule: Schedule | None = None,
) -> BltStrategy:
    """Create a BLT (Buffered Linear Toeplitz) strategy recipe.

    The L-BFGS optimization runs lazily inside the strategy methods
    when an amplifier supplies the participation context
    ``(n_steps, min_sep, max_participations)`` — once per recipe per
    context, via a module-level cache shared with the noise side.

    Args:
        max_buffers: Maximum BLT buffer count for the optimizer (default 10).
        momentum: Polyak momentum (default 1.0 = prefix-sum workload).
        lr_schedule: Optional :data:`opake.scheduling.types.Schedule`
            (``Callable[[int], float]``).  Materialised at ``[0, n_steps)``
            when the strategy first sees the amplifier's ``n_steps``. Custom
            schedules must be deterministic, side-effect-free, and remain
            immutable for the strategy's lifetime.

    Returns:
        A :class:`BltStrategy` recipe.
    """
    # max_buffers/momentum bounds live in ``BltStrategy.__post_init__``.
    return BltStrategy(
        max_buffers=max_buffers,
        momentum=momentum,
        lr_schedule=lr_schedule,
    )

lambda_cgd_strategy(*, lambda_, normalized=True, lr_schedule=None)

Create a DP-lambda-CGD strategy recipe (bandwidth=2, PRNG-replay noise).

Parameters:

Name	Type	Description	Default
lambda_	float	Correlation coefficient in [0, 1).	required
normalized	bool	Use column-normalized matrix (default True).	True
lr_schedule	Schedule | None	Deprecated compatibility argument. Only None is accepted. Pass learning-rate schedules to the optimizer instead.	None

Returns:

Name	Type	Description
A	LambdaCgdStrategy	class:LambdaCgdStrategy recipe.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/noise/_lambda_cgd.py

def lambda_cgd_strategy(
    *,
    lambda_: float,
    normalized: bool = True,
    lr_schedule: Schedule | None = None,
) -> LambdaCgdStrategy:
    """Create a DP-lambda-CGD strategy recipe (bandwidth=2, PRNG-replay noise).

    Args:
        lambda_: Correlation coefficient in [0, 1).
        normalized: Use column-normalized matrix (default True).
        lr_schedule: Deprecated compatibility argument. Only ``None`` is
            accepted. Pass learning-rate schedules to the optimizer instead.

    Returns:
        A :class:`LambdaCgdStrategy` recipe.
    """
    return LambdaCgdStrategy(
        lambda_=lambda_,
        normalized=normalized,
        lr_schedule=lr_schedule,
    )

bisr_strategy(*, bandwidth, normalized=True, momentum=0.0, lr_schedule=None, inv_coefficients=None)

Create a BISR (Banded Inverse Square Root) strategy recipe.

Parameters:

Name	Type	Description	Default
bandwidth	int	BISR bandwidth p (>= 2).	required
normalized	bool	Use column-normalized matrix (default True).	True
momentum	float	The paper's beta in the alpha=1 specialization. Must be in [0, 1) (default 0), matching Lemma 1's 0 <= beta < alpha <= 1 hypothesis.	0.0
lr_schedule	Schedule | None	Deprecated compatibility argument. Only None is accepted. Pass learning-rate schedules to the optimizer instead.	None
inv_coefficients	Sequence[float] | None	Explicit :math:C^{-1} coefficients (default BISR optimal).	None

Returns:

Name	Type	Description
A	BisrStrategy	class:BisrStrategy recipe.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/noise/_bisr.py

def bisr_strategy(
    *,
    bandwidth: int,
    normalized: bool = True,
    momentum: float = 0.0,
    lr_schedule: Schedule | None = None,
    inv_coefficients: Sequence[float] | None = None,
) -> BisrStrategy:
    """Create a BISR (Banded Inverse Square Root) strategy recipe.

    Args:
        bandwidth: BISR bandwidth p (>= 2).
        normalized: Use column-normalized matrix (default True).
        momentum: The paper's beta in the alpha=1 specialization. Must be in
            ``[0, 1)`` (default 0), matching Lemma 1's
            ``0 <= beta < alpha <= 1`` hypothesis.
        lr_schedule: Deprecated compatibility argument. Only ``None`` is
            accepted. Pass learning-rate schedules to the optimizer instead.
        inv_coefficients: Explicit :math:`C^{-1}` coefficients (default BISR optimal).

    Returns:
        A :class:`BisrStrategy` recipe.
    """
    return BisrStrategy(
        bandwidth=bandwidth,
        normalized=normalized,
        momentum=momentum,
        lr_schedule=lr_schedule,
        inv_coefficients=(
            tuple(inv_coefficients) if inv_coefficients is not None else None
        ),
    )

identity_strategy()

Create an identity (DP-SGD) noise strategy.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/noise/_identity.py

def identity_strategy() -> IdentityStrategy:
    """Create an identity (DP-SGD) noise strategy."""
    return IdentityStrategy()

State Classes

NoiseState

Bases: ABC

Base class for noise state.

All noise functions (Gaussian and MF) return a state object that inherits from this class, providing a unified interface for step tracking and RNG key management.

Attributes:

Name	Type	Description
_step_counter	int	Number of noise_fn calls made.
_rng_key	RngKey	Immutable RNG key for deterministic per-step derivation.

Source code in packages/opake-engine/src/opake/api/engine/types.py

class NoiseState(ABC):
    """Base class for noise state.

    All noise functions (Gaussian and MF) return a state object that
    inherits from this class, providing a unified
    interface for step tracking and RNG key management.

    Attributes:
        _step_counter: Number of ``noise_fn`` calls made.
        _rng_key: Immutable RNG key for deterministic per-step derivation.
    """

    _step_counter: int
    """Number of ``noise_fn`` calls made."""

    _rng_key: RngKey
    """Immutable RNG key for deterministic per-step derivation."""

SecondMomentNoiseOutput

Bases: NamedTuple

Noise output with private first and second moment streams.

Attributes:

Name	Type	Description
noisy_grads	NoisedPytree	Noised clipped gradients for the optimizer's first moment / update direction.
noisy_squared_grads	NoisedPytree	Noised element-wise squared clipped gradients for optimizers with a second moment accumulator.

Source code in packages/opake-engine/src/opake/api/engine/types.py

class SecondMomentNoiseOutput(NamedTuple):
    """Noise output with private first and second moment streams.

    Attributes:
        noisy_grads: Noised clipped gradients for the optimizer's first
            moment / update direction.
        noisy_squared_grads: Noised element-wise squared clipped gradients
            for optimizers with a second moment accumulator.
    """

    noisy_grads: NoisedPytree
    noisy_squared_grads: NoisedPytree

GaussianNoiseState dataclass

Bases: NoiseState

Immutable state for Gaussian noise generation.

Holds an immutable RNG key for deterministic per-step derivation. Noise for step t is generated from fold_in(_rng_key, GAUSSIAN_STREAM_FOLD, t).

Attributes:

Name	Type	Description
_step_counter	int	Number of noise_fn calls made.
_rng_key	RngKey	Immutable RNG key for deterministic derivation.

Source code in packages/opake-dpsgd/src/opake/api/dpsgd/noise/_gaussian.py

@dataclasses.dataclass(frozen=True)
class GaussianNoiseState(NoiseState):
    """Immutable state for Gaussian noise generation.

    Holds an immutable RNG key for deterministic per-step derivation.
    Noise for step ``t`` is generated from
    ``fold_in(_rng_key, GAUSSIAN_STREAM_FOLD, t)``.

    Attributes:
        _step_counter: Number of noise_fn calls made.
        _rng_key: Immutable RNG key for deterministic derivation.
    """

    _step_counter: int
    _rng_key: RngKey

MFNoiseState dataclass

Bases: NoiseState

State for matrix factorization noise.

Restoration is template-driven. A saved state must be loaded into a template built with the same strategy recipe, horizon, compute dtype, and parameter-tree structure. The trainer checkpoint records and checks that execution provenance before restoring this state; standalone callers are responsible for preserving the same construction inputs.

Attributes:

Name	Type	Description
_inner_state	Any	Internal state (streaming matrix state or step counter).
_step_counter	int	Number of noise_fn calls made.
_rng_key	RngKey	Immutable RNG key for deterministic per-step derivation.
_first_max_norm	float | PerGroup | None	ClippedPytree.max_norm from the first call (scalar or :class:~opake.types.PerGroup), latched by the dispatcher to enforce constant per-step sensitivity. None until the first call. MF privacy analyses assume the per-step sensitivity is constant across the sequence. Both fixed clipping (:func:opake.dpftrl.clipping.clipped_grad) and AUTO-S clipping (:func:opake.dpftrl.clipping.auto_clipped_grad) satisfy this assumption — the per-record bound is fixed at construction and does not depend on data. Adaptive clipping (:func:opake.dpsgd.clipping.adaptive_clipped_grad) breaks it, so the dispatcher rejects subsequent calls whose max_norm differs.
_first_max_norm_sync_fingerprint	int | None	Deterministic 64-bit fingerprint (computed at latch time) for cheap cross-rank checks in :func:sync_mf_noise_state, for both scalar and PerGroup norms. None before the first call.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/noise/_engine.py

@dataclasses.dataclass(frozen=True)
class MFNoiseState(NoiseState):
    """State for matrix factorization noise.

    Restoration is template-driven. A saved state must be loaded into a
    template built with the same strategy recipe, horizon, compute dtype, and
    parameter-tree structure. The trainer checkpoint records and checks that
    execution provenance before restoring this state; standalone callers are
    responsible for preserving the same construction inputs.

    Attributes:
        _inner_state: Internal state (streaming matrix state or step counter).
        _step_counter: Number of noise_fn calls made.
        _rng_key: Immutable RNG key for deterministic per-step derivation.
        _first_max_norm: ``ClippedPytree.max_norm`` from the first call (scalar
            or :class:`~opake.types.PerGroup`), latched by the dispatcher to
            enforce constant per-step sensitivity.  ``None`` until the first
            call.  MF privacy analyses assume the per-step sensitivity is
            constant across the sequence.  Both fixed clipping
            (:func:`opake.dpftrl.clipping.clipped_grad`) and AUTO-S clipping
            (:func:`opake.dpftrl.clipping.auto_clipped_grad`) satisfy this
            assumption — the per-record bound is fixed at construction and
            does not depend on data.  Adaptive clipping
            (:func:`opake.dpsgd.clipping.adaptive_clipped_grad`) breaks it,
            so the dispatcher rejects subsequent calls whose ``max_norm``
            differs.
        _first_max_norm_sync_fingerprint: Deterministic 64-bit fingerprint
            (computed at latch time) for cheap cross-rank checks in
            :func:`sync_mf_noise_state`, for both scalar and ``PerGroup`` norms.
            ``None`` before the first call.
    """

    _inner_state: Any
    _step_counter: int
    _rng_key: RngKey
    _first_max_norm: float | PerGroup | None = None
    _first_max_norm_sync_fingerprint: int | None = None

SecondMomentMFNoiseState dataclass

Bases: NoiseState

Internal state for MF noise with private second moments.

Source code in packages/opake-dpftrl/src/opake/api/dpftrl/noise/_second_moment.py

@dataclasses.dataclass(frozen=True)
class SecondMomentMFNoiseState(NoiseState):
    """Internal state for MF noise with private second moments."""

    _first_state: MFNoiseState
    _second_state: MFNoiseState

    @property
    def _step_counter(self) -> int:  # type: ignore[override]
        return self._first_state._step_counter

    @property
    def _rng_key(self) -> RngKey:  # type: ignore[override]
        return self._first_state._rng_key

Distributed Synchronization

Use opake.distributed.sync(state) — it auto-dispatches on the state's type to the right sync function. GaussianNoiseState and MFNoiseState both register handlers at import time, so no named sync call is required.