Training on Sensitive Data: DP-SGD and Federated Learning on TensorFlow

Most of the interesting datasets our clients own are also the ones their legal team is least willing to move: claims histories, clinical notes, transaction ledgers, keyboard and app telemetry. The usual answer — "copy it into a bucket and train" — is getting harder to defend as data-protection regimes and AI governance rules tighten, and "we deleted the raw data afterwards" is not a privacy guarantee.

TensorFlow has two mature answers to this, and they solve different problems:

• Differential privacy (DP) bounds what any single training record can contribute to the final weights. It protects against memorisation and membership inference after the data is centralised. Tooling: TensorFlow Privacy.
• Federated learning (FL) keeps the data on the device or in the tenant's own environment and moves only model updates. It changes where training happens. Tooling: TensorFlow Federated.

They compose: production federated systems normally apply DP to the aggregated updates. This tutorial builds the smallest honest version of each, then talks about the part clients actually care about — the accuracy you pay for the guarantee.

Setup

pip install -U tensorflow tensorflow-privacy
# federated section (separate, pinned environment - see note below)
pip install -U tensorflow-federated

A practical warning before you spend an afternoon on it: tensorflow-privacy and tensorflow-federated are pinned to specific TensorFlow versions and to each other. Put them in their own virtualenv or container, not in the environment where you keep the rest of your stack current. If you are on the TensorFlow release train described in our 2.21 upgrade guide, expect to hold the privacy environment a version or two behind.

Part 1: DP-SGD in three changes

DP-SGD modifies gradient descent in two places. Per-example gradients are clipped to a fixed L2 norm, so no single record can dominate an update, and calibrated Gaussian noise is added to the summed gradients before the step. That is the whole idea; the rest is bookkeeping.

Start from an ordinary Keras model:

import tensorflow as tf
import keras

num_features = 64

def build_model():
    return keras.Sequential([
        keras.layers.Input(shape=(num_features,)),
        keras.layers.Dense(128, activation="relu"),
        keras.layers.Dense(64, activation="relu"),
        keras.layers.Dense(1, activation="sigmoid"),
    ])

Now swap the optimiser and the loss reduction:

from tensorflow_privacy import DPKerasAdamOptimizer

BATCH = 256
MICROBATCHES = 256      # 1 microbatch per example = tightest clipping, slowest
L2_CLIP = 1.0
NOISE_MULTIPLIER = 1.1

model = build_model()
model.compile(
    optimizer=DPKerasAdamOptimizer(
        l2_norm_clip=L2_CLIP,
        noise_multiplier=NOISE_MULTIPLIER,
        num_microbatches=MICROBATCHES,
        learning_rate=1e-3,
    ),
    # per-example losses are required: no reduction
    loss=keras.losses.BinaryCrossentropy(
        reduction=tf.losses.Reduction.NONE
    ),
    metrics=[keras.metrics.AUC(name="auc")],
)

Three things that bite people here:

1. The loss must be unreduced. DP-SGD needs a per-example loss vector to compute per-example gradients. A scalar loss silently breaks the guarantee, which is why the optimiser rejects it.
2. num_microbatches must divide the batch size, and the batch size must be fixed — call dataset.batch(BATCH, drop_remainder=True). Ragged final batches invalidate the sampling assumption behind the accounting.
3. Microbatching is the cost dial. num_microbatches=batch_size gives true per-example clipping and is slow; larger microbatches (say 32 examples each) run much faster but clip groups, which weakens the effective guarantee for a given noise level.

Train as usual:

train_ds = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(10_000)
    .batch(BATCH, drop_remainder=True)
    .prefetch(tf.data.AUTOTUNE)
)

EPOCHS = 15
model.fit(train_ds, validation_data=val_ds, epochs=EPOCHS)

Accounting: the number you actually report

A DP model without a computed epsilon is marketing, not privacy. TensorFlow Privacy ships an accountant:

from tensorflow_privacy.privacy.analysis import compute_dp_sgd_privacy_lib as accounting

n_train = x_train.shape[0]

stats = accounting.compute_dp_sgd_privacy_statement(
    number_of_examples=n_train,
    batch_size=BATCH,
    num_epochs=EPOCHS,
    noise_multiplier=NOISE_MULTIPLIER,
    delta=1.0 / n_train,
    used_microbatching=True,
)
print(stats)

That prints an (epsilon, delta) pair for the run. Rules of thumb we use when advising clients: epsilon below ~1 is a strong, defensible guarantee; 1-10 is the range most published applied work lives in; above ~20 you should describe it as "noise added" rather than as a meaningful privacy guarantee, because the bound no longer excludes much.

Epsilon grows with epochs and shrinks with noise and with dataset size. The lever nobody expects is more data: at fixed epsilon, a 10x larger dataset buys back most of the accuracy DP costs you.

Verifying, not assuming

The complement to accounting is an empirical attack. Train a shadow-model membership-inference test against the DP model and its non-private twin and compare attacker AUC:

from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack import (
    membership_inference_attack as mia,
)

The non-private baseline routinely leaks — attacker AUC comfortably above 0.5, especially on small, high-dimensional, over-trained models. Showing the same attack flattened to near chance on the DP model is the single most persuasive artifact you can hand a data-protection officer.

What DP costs

Budget for it up front, so the trade-off is a decision and not a surprise:

• Accuracy. A few points of AUC on tabular problems at moderate epsilon; more on small datasets, less on large ones.
• Throughput. Per-example clipping is typically 2-5x slower per step. Mixed precision and XLA still help — see our training-speed tutorial — but vectorised per-example gradients are the dominant cost.
• Fairness. DP noise hurts rare subgroups first, because they have the least signal to survive clipping and noise. Always slice your evaluation by subgroup before shipping a DP model; this is where an unexamined privacy win turns into a bias problem.
• Tuning. Hyperparameter search itself consumes privacy budget if you tune on the sensitive data. Tune on public or synthetic data where you can.

One practical shortcut: fine-tune, don't pretrain. Pretraining a backbone on public data and applying DP only to the last layers or to a LoRA adapter on the sensitive corpus gives a far better privacy/utility curve than training from scratch under DP. That pairs naturally with the KerasHub workflow in our Gemma fine-tuning tutorial.

Part 2: a federated averaging baseline

When the data cannot be centralised at all — separate hospitals, separate tenants, on-device keyboards — you move the computation instead. TensorFlow Federated expresses this as: each client trains locally for a few steps, the server averages the resulting weight deltas, repeat.

TFF wants a model_fn that constructs a fresh, uncompiled model plus its input spec:

import tensorflow_federated as tff

def model_fn():
    return tff.learning.models.from_keras_model(
        build_model(),
        input_spec=example_dataset.element_spec,
        loss=keras.losses.BinaryCrossentropy(),
        metrics=[keras.metrics.AUC(name="auc")],
    )

process = tff.learning.algorithms.build_weighted_fed_avg(
    model_fn,
    client_optimizer_fn=lambda: keras.optimizers.SGD(0.02),
    server_optimizer_fn=lambda: keras.optimizers.SGD(1.0),
)

state = process.initialize()
for round_num in range(1, 51):
    sampled = sample_clients(client_datasets, k=10)   # your cohort policy
    result = process.next(state, sampled)
    state = result.state
    print(round_num, result.metrics["client_work"]["train"])

Notes from doing this on real engagements:

• model_fn must build the model inside the function. TFF traces it into a computation; a captured, already-built Keras model will fail or, worse, share state across clients.
• Non-IID data is the whole difficulty. Clients whose label distributions differ badly make plain FedAvg oscillate. Fewer local steps per round, server-side momentum, and more clients per round all help.
• Simulation is not deployment. The Python runtime is for research and for proving the algorithm converges. Cross-silo production deployments need per-participant infrastructure and a secure aggregation story, which is a systems project, not a modelling one.
• Combine with DP. tff.learning.ddp_aggregator / DP aggregators add noise to the aggregate, giving user-level DP — a stronger and more meaningful claim than record-level DP for device data.

Choosing between them

Constraint	Reach for
Data can be centralised, but memorisation or membership inference is the risk	DP-SGD
Data cannot legally or physically leave its silo	Federated learning
Regulator wants a written, quantified guarantee	DP-SGD with a published epsilon
Consumer devices, thousands of participants	FL with user-level DP
Small dataset, tight accuracy target	Neither — try public pretraining plus DP fine-tuning of a small head

Most engagements land on the last row: a public backbone, a narrow DP fine-tune, and a documented epsilon. That is far cheaper than a federated deployment and, for a single-tenant dataset, usually just as defensible.

Whatever you build, treat the privacy claim like any other model artifact: version the epsilon, the clipping norm, the noise multiplier, and the batch/epoch counts alongside the checkpoint, and re-run the membership-inference test in CI whenever the training data changes. A guarantee you cannot reproduce is not a guarantee.

---

Building ML on data your compliance team is nervous about? Get in touch — our TensorFlow consultants can pressure-test the privacy design before you commit to an architecture.