Privacy-Preserving Analytics and Experimentation Questions

import torch
from opacus import PrivacyEngine
# model, dataloader, optimizer_d defined
privacy_engine = PrivacyEngine(
    module=disc,
    sample_rate=batch_size / len(dataset),
    alphas=[1 + x/10. for x in range(1, 100)],  # RDP orders
    noise_multiplier=noise_multiplier,
    max_grad_norm=clip_norm,
)
privacy_engine.attach(optimizer_d)

for epoch in range(epochs):
    for real_batch in loader:
        # standard GAN discriminator step
        optimizer_d.zero_grad()
        loss_d = discriminator_loss(real_batch, gen)
        loss_d.backward()
        optimizer_d.step()
    eps, best_alpha = privacy_engine.get_privacy_spent(delta=target_delta)
    print(f"Epoch {epoch} ε = {eps:.3f}, δ = {target_delta}")

from functorch import make_functional_with_buffers, vmap, grad

# make functional for generator
f_gen, params_gen, buffers = make_functional_with_buffers(gen)
def loss_fn(params, buffers, z, real_labels):
    fake = f_gen(params, buffers, z)
    return generator_loss_from_disc(fake, disc)

# compute per-sample grads
per_sample_grads = vmap(grad(loss_fn), in_dims=(None, None, 0, 0))(params_gen, buffers, z_batch, real_labels_batch)
# clip per-sample grads (L2 across parameters)
def clip_and_aggregate(per_sample_grads, max_norm):
    norms = torch.stack([torch.sqrt(sum((g**2).sum() for g in sample))) for sample in zip(*per_sample_grads)])
    scales = (max_norm / (norms + 1e-6)).clamp(max=1.0)
    # scale each sample's grads then sum and add noise
    aggregated = [ (g * scales.view(-1, *([1]* (g.dim()-1)))).sum(0) for g in per_sample_grads ]
    return aggregated

agg_grads = clip_and_aggregate(per_sample_grads, clip_norm)
# add Gaussian noise
for g in agg_grads:
    g += torch.normal(0, noise_multiplier*clip_norm, size=g.shape, device=g.device) / batch_size
# load into optimizer & step (manual param update or set gradients then optimizer.step)

import tensorflow as tf
from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras import DPKerasAdamOptimizer
from tensorflow_privacy.privacy.analysis.rdp_accountant import compute_rdp, get_privacy_spent

optimizer_d = DPKerasAdamOptimizer(
    l2_norm_clip=clip_norm,
    noise_multiplier=noise_multiplier,
    num_microbatches=batch_size,  # or set microbatches <= batch_size
    learning_rate=lr
)
# compile discriminator with optimizer_d and train via gradient tape
# compute RDP epsilon each epoch:
orders = [1 + x/10. for x in range(1, 100)]
sampling_prob = batch_size / dataset_size
rdp = compute_rdp(q=sampling_prob, noise_multiplier=noise_multiplier, steps=steps, orders=orders)
eps, opt_order = get_privacy_spent(orders, rdp, target_delta)
print(eps)