LoopyLuci Differential Privacy for ML Training
Implementing differentially private machine learning using DP-SGD and related techniques to train models with formal privacy guarantees.
When to Use
- Training on sensitive data (medical records, financial data, personal information)
- Regulatory compliance (GDPR, HIPAA, CCPA)
- Publishing models trained on user data without leaking individual records
- Federated learning with privacy guarantees
- Building privacy-preserving AI systems
Core Concepts
Privacy Budget (ε, δ)
# ε (epsilon): privacy loss budget. Lower = more privacy.
# δ (delta): probability of privacy breach (typically < 1/N where N = dataset size)
# Typical values: ε=1 (strong privacy), ε=8 (weak privacy), δ=1e-5
DP-SGD Algorithm
import torch
import torch.nn as nn
from torch.optim import Optimizer
class DPSGD(Optimizer):
"""Differentially Private SGD with gradient clipping and noise.
Key modifications to standard SGD:
1. Clip per-sample gradients to L2 norm C
2. Add Gaussian noise scaled by C * σ
3. Account privacy budget via moments accountant
"""
def __init__(self, params, lr=1e-3, clip_norm=1.0, noise_multiplier=1.1,
batch_size=256, dataset_size=50000, delta=1e-5):
defaults = dict(lr=lr, clip_norm=clip_norm, noise_multiplier=noise_multiplier,
batch_size=batch_size, dataset_size=dataset_size)
super().__init__(params, defaults)
self.delta = delta
self.epsilon = None # Will be computed via moments accountant
self.steps_taken = 0
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
params = [p for p in group['params'] if p.grad is not None]
if not params:
continue
# Step 1: Per-sample gradient clipping
# We need per-sample gradients, which requires:
# - Setting `loss.backward(retain_graph=True)` for each sample in batch
# - OR using a per-sample gradient computation (Ghost Clipping)
# Here we assume micro-batch approach
# For each parameter, clip its gradient
total_norm = torch.norm(
torch.stack([torch.norm(p.grad, 2) for p in params])
)
clip_coef = min(1, group['clip_norm'] / (total_norm + 1e-8))
for p in params:
p.grad *= clip_coef # Clip to C
# Step 2: Add Gaussian noise
noise = torch.normal(
mean=0,
std=group['clip_norm'] * group['noise_multiplier'],
size=p.grad.shape,
device=p.grad.device
)
p.grad += noise / group['batch_size']
# Step 3: Update parameters
p.add_(p.grad, alpha=-group['lr'])
self.steps_taken += 1
return loss
Per-Sample Gradient Computation
class PerSampleGradientComputer:
"""Efficient per-sample gradient computation for DP-SGD.
Without this, you'd need to backward() per sample (very slow).
Ghost clipping computes per-sample norms efficiently.
"""
def compute_per_sample_norms(self, model, x, y, loss_fn):
"""Compute per-sample gradient norms without storing per-sample grads."""
# Forward pass with a single backward
logits = model(x)
loss = loss_fn(logits, y)
# Compute per-sample gradients (requires grad in loss function)
# Efficient implementation depends on layer type:
per_sample_norms = {}
for name, param in model.named_parameters():
if param.grad is None:
continue
# For linear layers: per-sample grad = outer product
if isinstance(param, nn.Linear):
# Ghost clipping: compute per-sample gradient norms efficiently
pass
# Clip and aggregate
total_clipped = self.clip_and_aggregate(per_sample_norms, clip_norm=1.0)
return total_clipped
Privacy Accounting
Moments Accountant
class MomentsAccountant:
"""Tracks privacy budget (ε) across training steps using Rényi DP."""
def __init__(self, noise_multiplier, delta=1e-5):
self.noise_multiplier = noise_multiplier
self.delta = delta
self.orders = [1 + x/10.0 for x in range(1, 100)] + list(range(12, 64))
self.renyi_divergences = {}
def compute_epsilon(self, steps, sampling_rate):
"""Compute total privacy spend after `steps` steps."""
# Rényi DP composition
eps = float('inf')
for order in self.orders:
# RDP for Gaussian mechanism
rdp = order * sampling_rate / (2 * self.noise_multiplier**2)
rdp_total = rdp * steps
# Convert from RDP to (ε, δ)-DP
eps_order = rdp_total - np.log(self.delta) / (order - 1)
eps = min(eps, eps_order)
return eps
def get_privacy_spent(self, steps_per_epoch, num_epochs, batch_size, dataset_size):
"""Get total (ε, δ) privacy guarantee."""
sampling_rate = batch_size / dataset_size
total_steps = steps_per_epoch * num_epochs
epsilon = self.compute_epsilon(total_steps, sampling_rate)
return epsilon, self.delta
Privacy Budget Scheduler
class PrivacyBudgetScheduler:
"""Schedule noise_multiplier across training to hit a target ε budget."""
def __init__(self, target_epsilon, total_steps, dataset_size, batch_size, delta=1e-5):
self.target_eps = target_epsilon
self.total_steps = total_steps
self.dataset_size = dataset_size
self.batch_size = batch_size
self.delta = delta
self.sampling_rate = batch_size / dataset_size
def find_noise_multiplier(self):
"""Binary search for noise_multiplier that meets budget."""
lo, hi = 0.1, 10.0
for _ in range(20): # Binary search iterations
mid = (lo + hi) / 2
accountant = MomentsAccountant(mid, self.delta)
eps = accountant.compute_epsilon(self.total_steps, self.sampling_rate)
if eps > self.target_eps:
lo = mid # More noise needed
else:
hi = mid # Less noise might work
return hi # Return smallest noise that meets budget
Training Loop
def dp_train(model, train_loader, epochs=50, lr=1e-3,
clip_norm=1.0, noise_multiplier=1.1, target_epsilon=8.0):
"""Full DP-SGD training loop with privacy accounting."""
dataset_size = len(train_loader.dataset)
batch_size = train_loader.batch_size
delta = 1 / (dataset_size * 10) # δ < 1/dataset_size
optimizer = DPSGD(model.parameters(), lr=lr, clip_norm=clip_norm,
noise_multiplier=noise_multiplier, batch_size=batch_size,
dataset_size=dataset_size, delta=delta)
accountant = MomentsAccountant(noise_multiplier, delta)
for epoch in range(epochs):
for x, y in train_loader:
# Standard forward + backward (per-sample gradients needed for DP)
logits = model(x)
loss = F.cross_entropy(logits, y)
loss.backward()
optimizer.step()
optimizer.zero_grad()
# Privacy accounting
eps = accountant.get_privacy_spent(
len(train_loader), epoch + 1, batch_size, dataset_size
)
print(f"Epoch {epoch}: ε={eps[0]:.2f}, δ={eps[1]:.2e}")
if eps[0] > target_epsilon:
print(f"Reached target privacy budget ({target_epsilon}). Stopping.")
break
Techniques for Better Privacy/Utility Trade-off
# 1. Virtual batch size: accumulate gradients across micro-batches
# before clipping + noise
class VirtualBatchDP:
def __init__(self, micro_batch_size=32, virtual_batch_size=256):
self.micro_batch_size = micro_batch_size
self.virtual_batch_size = virtual_batch_size
def accumulate_gradients(self, model, dataset):
"""Accumulate gradients over virtual batch."""
model.zero_grad()
for i in range(0, len(dataset), self.micro_batch_size):
micro_batch = dataset[i:i + self.micro_batch_size]
loss = model(micro_batch)
loss.backward()
# Don't step yet — accumulate across micro-batches
# Now we have accumulated gradients for the virtual batch
# Clip and add noise at virtual_batch_size granularity
self.clip_and_noise(model)
# 2. Weight standardization + DP
# Standardize weights after each step (improves stability with DP noise)
# 3. Public pre-training + private fine-tuning
# Pre-train on public data, fine-tune with DP on private data
Common Pitfalls
- Per-sample gradient cost — naive per-sample backward is N times slower; use ghost clipping or JAX vmap
- Privacy budget underestimation — moments accountant assumes Poisson sampling; fixed-size batches leak information
- Hyperparameter sensitivity — DP training is very sensitive to LR and clipping norm; tune carefully
- Model capacity — larger models need more privacy budget for same utility; use simpler models for strong privacy
- Batch size effects — small batches have worse privacy/utility tradeoff; use large virtual batches (1024+)
- Evaluation leakage — don't use validation set for early stopping without accounting for it in privacy budget
Verification Checklist
- Privacy budget (ε, δ) reported transparently
- Per-sample gradient computation verified (norms computed correctly)
- Implementation produces same results as reference (e.g., Opacus for PyTorch)
- Training converges under privacy guarantees
- Utility drop quantified (accuracy vs. non-private baseline)
- Dataset-specific δ computed as < 1/N
- Accounting uses correct sampling rate (Poisson vs. shuffled)
See Also
- adversarial-ml-robustness — robustness against adversarial inputs
- federated-learning — distributed training with privacy
- ai-fairness-mitigation — fairness alongside privacy
- custom-training-loops — customizing training loops