Train a Model with DP-FTRL¶

This notebook trains a logistic regression model on synthetic data using DP-FTRL — differential privacy with correlated Gaussian noise across training steps via matrix factorization. Compared to DP-SGD (tutorial) it improves the privacy-utility trade-off on cumulative updates, at the cost of fixing the training length up front and committing to a strategy at calibration time.

Prerequisites: DP-SGD Training.

Components exercised: band_mf_strategy, mf_gaussian_noise, opake.dpftrl.clipping.clipped_grad, dpftrl_acc.poisson + dpftrl_acc.mf_gaussian, acc.calibrate.

In [1]:

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

import opake.accounting as acc
import opake.dpftrl.accounting as dpftrl_acc
from opake.dpftrl.clipping import clipped_grad
from opake.dpftrl.noise import band_mf_strategy, mf_gaussian_noise
from opake.dpftrl.sampling import CyclicPoissonSampler
from opake.functional import make_functional
from opake.random import key

torch.manual_seed(42)
np.random.seed(42)

import matplotlib.pyplot as plt import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset import opake.accounting as acc import opake.dpftrl.accounting as dpftrl_acc from opake.dpftrl.clipping import clipped_grad from opake.dpftrl.noise import band_mf_strategy, mf_gaussian_noise from opake.dpftrl.sampling import CyclicPoissonSampler from opake.functional import make_functional from opake.random import key torch.manual_seed(42) np.random.seed(42)

Synthetic data¶

A binary classification task with 10 features and 2000 examples — same shape as the DP-SGD tutorial so the two pipelines can be compared directly.

In [2]:

n_samples, n_features = 2000, 10
X = torch.randn(n_samples, n_features)
y = (X[:, 0] + 0.5 * X[:, 1] > 0).float()

print(f"Dataset: {n_samples} samples, {n_features} features")
print(f"Positive class: {y.mean():.1%}")

n_samples, n_features = 2000, 10 X = torch.randn(n_samples, n_features) y = (X[:, 0] + 0.5 * X[:, 1] > 0).float() print(f"Dataset: {n_samples} samples, {n_features} features") print(f"Positive class: {y.mean():.1%}")

Dataset: 2000 samples, 10 features
Positive class: 51.0%

1. Strategy and training length¶

DP-FTRL accountants describe whole training runs. Pick the number of steps and the bandwidth (number of off-diagonal bands) up front; the strategy object holds the sensitivity and lower-triangular coefficients used by both the noise mechanism and the accountant.

We use band_mf_strategy here — a banded Toeplitz factorization that is a strong general-purpose default. See DP-FTRL mechanisms for alternatives (BLT, BSR, BISR, DP-λCGD, identity).

In [3]:

batch_size = 64
n_steps = 500
bands = 16
# Cyclic Poisson activates one of `bands` groups per step, so this
# eligible-example rate gives `batch_size` expected records per batch.
sample_rate = batch_size * bands / n_samples

strategy = band_mf_strategy(bands=bands)

print(f"Strategy: band-MF, n_steps={n_steps}, bands={bands}")
print(
    "Single-participation sensitivity: "
    f"{strategy.sensitivity(n_steps=n_steps, min_sep=n_steps, max_participations=1):.4f}"
)

batch_size = 64 n_steps = 500 bands = 16 # Cyclic Poisson activates one of `bands` groups per step, so this # eligible-example rate gives `batch_size` expected records per batch. sample_rate = batch_size * bands / n_samples strategy = band_mf_strategy(bands=bands) print(f"Strategy: band-MF, n_steps={n_steps}, bands={bands}") print( "Single-participation sensitivity: " f"{strategy.sensitivity(n_steps=n_steps, min_sep=n_steps, max_participations=1):.4f}" )

Strategy: band-MF, n_steps=500, bands=16

Single-participation sensitivity: 1.0000

2. Calibration¶

Wrap the strategy into a DpProcess describing the full run, then solve for the noise multiplier that achieves the target ε. The same strategy object goes into mf_gaussian_noise later — that's how DP correctness is preserved.

In [4]:

target_epsilon = 3.0
target_delta = 1e-5

result = acc.calibrate(
    acc.epsilon_budget(target_epsilon, delta=target_delta),
    lambda nm: dpftrl_acc.poisson(
        dpftrl_acc.mf_gaussian(nm, strategy),
        sample_rate=sample_rate,
        n_steps=n_steps,
    ),
    param_min=0.1,
    param_max=20.0,
)
noise_multiplier = result.param
print(f"Calibrated noise multiplier: {noise_multiplier:.4f}")

target_epsilon = 3.0 target_delta = 1e-5 result = acc.calibrate( acc.epsilon_budget(target_epsilon, delta=target_delta), lambda nm: dpftrl_acc.poisson( dpftrl_acc.mf_gaussian(nm, strategy), sample_rate=sample_rate, n_steps=n_steps, ), param_min=0.1, param_max=20.0, ) noise_multiplier = result.param print(f"Calibrated noise multiplier: {noise_multiplier:.4f}")

Calibrated noise multiplier: 4.1981

3. Functional model and per-example loss¶

In [5]:

model = nn.Linear(n_features, 1)
fmodel, params = make_functional(model)


def loss_fn(params, x, y):
    logit = fmodel(params, x.unsqueeze(0)).squeeze()
    return F.binary_cross_entropy_with_logits(logit, y)

model = nn.Linear(n_features, 1) fmodel, params = make_functional(model) def loss_fn(params, x, y): logit = fmodel(params, x.unsqueeze(0)).squeeze() return F.binary_cross_entropy_with_logits(logit, y)

4. Clipped gradient and MF noise¶

Clipping comes from opake.dpftrl.clipping (same engine primitive, stack-context-correct path). mf_gaussian_noise takes the parameter template so it can pre-allocate streaming-matrix state with the right shapes; the per-step contribution bound is fixed at construction time and must stay constant across the run.

In [6]:

clipping_norm = 1.0

grad_fn, clip_state = clipped_grad(
    loss_fn,
    argnums=0,
    batch_argnums=(1, 2),
    clipping_norm=clipping_norm,
    normalize_by=batch_size,
)

# `mf_gaussian_noise` reads structure from a parameter template. The per-step
# contribution bound is `clipping_norm × noise_multiplier`, latched on
# the first noised step and required to stay constant for the rest of
# the run — that's why adaptive clipping isn't compatible with DP-FTRL.
noise_fn, noise_state = mf_gaussian_noise(
    params,
    strategy,
    n_steps=n_steps,
    noise_multiplier=noise_multiplier,
    key=key(0),
)

print(f"Clip norm: {clipping_norm}")
print(f"Per-record sensitivity: {clipping_norm / batch_size:.4f}")
print(f"Noise multiplier: {noise_multiplier:.4f}")

clipping_norm = 1.0 grad_fn, clip_state = clipped_grad( loss_fn, argnums=0, batch_argnums=(1, 2), clipping_norm=clipping_norm, normalize_by=batch_size, ) # `mf_gaussian_noise` reads structure from a parameter template. The per-step # contribution bound is `clipping_norm × noise_multiplier`, latched on # the first noised step and required to stay constant for the rest of # the run — that's why adaptive clipping isn't compatible with DP-FTRL. noise_fn, noise_state = mf_gaussian_noise( params, strategy, n_steps=n_steps, noise_multiplier=noise_multiplier, key=key(0), ) print(f"Clip norm: {clipping_norm}") print(f"Per-record sensitivity: {clipping_norm / batch_size:.4f}") print(f"Noise multiplier: {noise_multiplier:.4f}")

Clip norm: 1.0
Per-record sensitivity: 0.0156
Noise multiplier: 4.1981

5. Training loop¶

The structure mirrors DP-SGD: clip then noise. The only difference is the noise call routes through the MF streaming matrix — successive calls produce correlated noise that partially cancels on cumulative updates.

In [7]:

lr = 0.1
losses = []
dataset = TensorDataset(X, y)
sampler = CyclicPoissonSampler(
    dataset,
    sample_rate=sample_rate,
    bands=bands,
    n_steps=n_steps,
    key=key(1),
)
loader = DataLoader(dataset, batch_sampler=sampler)

for step, (xb, yb) in enumerate(loader):

    grads, clip_state = grad_fn(params, xb, yb, state=clip_state)
    noisy_grads, noise_state = noise_fn(grads, noise_state)

    params = tuple(p - lr * g for p, g in zip(params, noisy_grads.pytree))

    if (step + 1) % 50 == 0 or step == 0:
        with torch.no_grad():
            avg_loss = loss_fn(params, xb, yb).item()
        losses.append((step + 1, avg_loss))
        print(f"Step {step + 1:4d}/{n_steps}  loss={avg_loss:.4f}")

# Whole-process accounting: ε is computed once at the end, not per step.
process = dpftrl_acc.poisson(
    dpftrl_acc.mf_gaussian(noise_multiplier, strategy),
    sample_rate=sample_rate,
    n_steps=n_steps,
)
final_eps = process.epsilon_at(target_delta)
print(f"\nFinal privacy: (epsilon={final_eps:.2f}, delta={target_delta})")

lr = 0.1 losses = [] dataset = TensorDataset(X, y) sampler = CyclicPoissonSampler( dataset, sample_rate=sample_rate, bands=bands, n_steps=n_steps, key=key(1), ) loader = DataLoader(dataset, batch_sampler=sampler) for step, (xb, yb) in enumerate(loader): grads, clip_state = grad_fn(params, xb, yb, state=clip_state) noisy_grads, noise_state = noise_fn(grads, noise_state) params = tuple(p - lr * g for p, g in zip(params, noisy_grads.pytree)) if (step + 1) % 50 == 0 or step == 0: with torch.no_grad(): avg_loss = loss_fn(params, xb, yb).item() losses.append((step + 1, avg_loss)) print(f"Step {step + 1:4d}/{n_steps} loss={avg_loss:.4f}") # Whole-process accounting: ε is computed once at the end, not per step. process = dpftrl_acc.poisson( dpftrl_acc.mf_gaussian(noise_multiplier, strategy), sample_rate=sample_rate, n_steps=n_steps, ) final_eps = process.epsilon_at(target_delta) print(f"\nFinal privacy: (epsilon={final_eps:.2f}, delta={target_delta})")

Step    1/500  loss=0.7660
Step   50/500  loss=0.4617
Step  100/500  loss=0.3182
Step  150/500  loss=0.2299
Step  200/500  loss=0.2611
Step  250/500  loss=0.2317

Step  300/500  loss=0.1713
Step  350/500  loss=0.1830
Step  400/500  loss=0.1647
Step  450/500  loss=0.1814

Step  500/500  loss=0.1823

Final privacy: (epsilon=3.00, delta=1e-05)

6. Evaluation¶

In [8]:

X_test = torch.randn(500, n_features)
y_test = (X_test[:, 0] + 0.5 * X_test[:, 1] > 0).float()

with torch.no_grad():
    logits = fmodel(params, X_test).squeeze()
    preds = (logits > 0).float()
    accuracy = (preds == y_test).float().mean()

print(f"Test accuracy: {accuracy:.1%}")

X_test = torch.randn(500, n_features) y_test = (X_test[:, 0] + 0.5 * X_test[:, 1] > 0).float() with torch.no_grad(): logits = fmodel(params, X_test).squeeze() preds = (logits > 0).float() accuracy = (preds == y_test).float().mean() print(f"Test accuracy: {accuracy:.1%}")

Test accuracy: 99.2%

Summary¶

DP-FTRL has the same loop shape as DP-SGD with three differences:

1. Strategy first. band_mf_strategy(bands=...) (or one of the alternatives) configures the noise matrix. The training horizon is provided via n_steps to mf_gaussian_noise().
2. Single-shot accounting. The amplification factory takes n_steps — the privacy guarantee is computed once for the whole run, not composed step by step.
3. Correlated noise. mf_gaussian_noise produces noise correlated through the streaming matrix; cumulative updates see partial cancellation.

Adaptive clipping is not available under DP-FTRL — its threshold drifts across steps, violating the constant per-step sensitivity the MF privacy proof requires. AUTO-S clipping (opake.dpftrl.clipping.auto_clipped_grad) is compatible.

For other strategies (BLT for long runs, DP-λCGD for zero extra memory, BSR for the closed-form workload), see DP-FTRL mechanisms. For private second-moment estimation under MF noise, see Optimizers.