Federated Learning for Privacy-Preserving Training

"I keep my weights close, but my gradients closer. Sharing is caring, as long as nobody sees the data."

A Farsighted Guard, Gradient-Hoarding AI Agent

50.3.1 Federated Learning Fundamentals

Centralized training collects all data on a central server. Federated learning inverts this: the model travels to the data, not the other way around.

The FedAvg Algorithm

McMahan et al. (2017) introduced Federated Averaging (FedAvg), the foundational FL algorithm. Each communication round runs four steps:

1. The server sends the current global model $w_t$ to a subset of $K$ clients
2. Each client $k$ trains the model on its local data for $E$ local epochs, producing updated weights $w_{t+1}^k$
3. Clients send their updates $\Delta w^k = w_{t+1}^k - w_t$ back to the server
4. The server aggregates updates using a weighted average: $w_{t+1} = w_t + \frac{1}{K} \sum_{k=1}^{K} \Delta w^k$

The weighted average is typically proportional to each client's dataset size: clients with more data contribute proportionally more to the global update.

Key Challenges

• Non-IID data: Client datasets are typically not independent and identically distributed. A hospital in cardiology sees different conditions than one in pediatrics. This data heterogeneity causes client updates to diverge, slowing convergence.
• Communication cost: Transmitting full model gradients for a 7B parameter model requires ~28 GB per round (FP32). With hundreds of clients and many rounds, bandwidth becomes the bottleneck.
• Stragglers: Clients with slower hardware or limited connectivity delay each round. Asynchronous FL protocols mitigate this but introduce staleness in gradients.
• Privacy leakage: Even gradient updates can leak information about training data through gradient inversion attacks (Section 50.3.3). Differential privacy or secure aggregation is needed.

50.3.2 Federated Fine-Tuning of LLMs

Full federated pretraining of LLMs is prohibitively expensive: the communication cost of exchanging billions of parameters each round is impractical. Instead, the dominant approach is federated fine-tuning, where a pretrained base model is adapted to domain-specific data held by multiple parties.

Federated LoRA (FFA-LoRA)

LoRA adapters are a natural fit for federated learning because they are small. Instead of exchanging 7B parameters each round, clients exchange only the low-rank adapter matrices (typically 1-10M parameters), reducing communication cost by 100-1000x.

import copy
import torch
def federated_lora_round(global_adapter, client_datasets, base_model, config):
    """One round of federated LoRA fine-tuning."""
    client_adapters = []
    client_sizes = []
    for dataset in client_datasets:
        # Each client starts from the global adapter
        local_adapter = copy.deepcopy(global_adapter)
        local_model = apply_lora(base_model, local_adapter)
        # Local training (E epochs on client data)
        optimizer = torch.optim.AdamW(local_adapter.parameters(), lr=config.lr)
        for epoch in range(config.local_epochs):
            for batch in dataset:
                loss = local_model(**batch).loss
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()
                client_adapters.append(local_adapter)
                client_sizes.append(len(dataset))
                # Server aggregation: weighted average of adapter weights
                total_size = sum(client_sizes)
                new_adapter = copy.deepcopy(global_adapter)
                with torch.no_grad():
                    for name, param in new_adapter.named_parameters():
                        weighted_sum = sum(
                            (size / total_size) * dict(adapter.named_parameters())[name]
                            for adapter, size in zip(client_adapters, client_sizes)
                            )
                        param.copy_(weighted_sum)
                        return new_adapter
                        # Training loop
                        global_adapter = initialize_lora(rank=16, target_modules=["q_proj", "v_proj"])
                        for round_num in range(config.num_rounds):
                            global_adapter = federated_lora_round(
                                global_adapter, client_datasets, base_model, config
                                )

Heterogeneous Federated Fine-Tuning

In practice, FL clients span a wide hardware range: some run A100 GPUs, others run consumer cards or CPUs. Heterogeneous FL lets each client pick adapter configuration (different LoRA ranks, for example) to match its compute budget, and the server runs rank-adaptive aggregation across the heterogeneous updates.

50.3.3 Privacy-Preserving Techniques in Federated LLM Training

Federated learning alone does not guarantee privacy. Gradient updates can be reverse-engineered to reconstruct training data, especially for text. Additional privacy mechanisms are essential. To see why these mechanisms are not optional, it helps to make the threat concrete: the very gradient a well-behaved client uploads can be inverted back into the private example that produced it.

Gradient Inversion: The Threat the Defenses Answer

The warning above ("gradients can leak training data") names an attack without showing it. The attack is gradient inversion, also called gradient leakage, formalized as Deep Leakage from Gradients (DLG) by Zhu et al. (2019). Consider an honest-but-curious server (or any party that observes one client's update). The client computed its gradient on a private example $(x, y)$ by backpropagating a loss $\mathcal{L}(f_\theta(x), y)$ through the shared model $f_\theta$, and uploaded the resulting per-example gradient $g^{*} = \nabla_\theta \mathcal{L}(f_\theta(x), y)$. The attacker knows $\theta$ (it is the global model everyone shares) and observes $g^{*}$, but not $x$ or $y$. The question is whether $(x, y)$ can be recovered from $g^{*}$ alone.

DLG answers yes by treating reconstruction as an optimization problem. The attacker initializes a dummy input $x'$ and dummy label $y'$ from random noise, runs the same forward and backward pass to obtain a dummy gradient $\nabla_\theta \mathcal{L}(f_\theta(x'), y')$, and then descends on $(x', y')$ to make that dummy gradient match the observed one. The objective is the gradient-matching loss:

Crucially the model weights $\theta$ are frozen; the optimization variables are the synthetic data $(x', y')$ themselves. Each step of gradient descent on this objective nudges $x'$ toward an input whose gradient signature reproduces $g^{*}$, and because the gradient is a near-unique fingerprint of the example for a small batch (the regime below), the recovered $x'^{\,*}$ converges to a pixel-level or token-level copy of the client's private $x$. The refinement iDLG (Zhao et al., 2020) removes the need to optimize the label at all: for the standard cross-entropy loss, the sign of the gradient of the final (classification) layer reveals the ground-truth class directly, because only the true-label logit receives a negative gradient. Reading $y$ off the last-layer gradient sign makes the remaining reconstruction of $x$ both faster and more reliable.

Why it works, and when it fails. A gradient is a high-dimensional, smooth function of the input: for a single example the map $x \mapsto \nabla_\theta \mathcal{L}(f_\theta(x), y)$ has so many output coordinates (one per parameter, often millions) constraining so few input coordinates (the pixels or token embeddings of one example) that it is locally near-invertible. The system is heavily over-determined, so matching the gradient pins down the input almost uniquely. This immediately explains the two structural defenses. First, large batches: when a client's update is the average gradient over $B$ examples, the observed vector $\frac{1}{B}\sum_{i} \nabla_\theta \mathcal{L}(f_\theta(x_i), y_i)$ is a sum, and the attacker must disentangle $B$ unknown inputs from one averaged signature, an under-determined problem that degrades rapidly as $B$ grows. Second, aggregation: secure aggregation forces the server to observe only the sum over many clients, pushing the effective batch to thousands and destroying the per-example signal the attack depends on. The threat is therefore sharpest exactly where FL is most naive: per-example or tiny-batch updates sent in the clear.

Code Fragment 50.3.3 below implements the gradient-matching loop of the objective above so the leakage is not merely asserted. It is deliberately a skeleton: it assumes a tiny model and a single example, prints the matching loss decreasing as $x'$ converges, and omits the cosine-similarity refinements and total-variation image priors that production attacks (and the Geiping et al. 2020 "Inverting Gradients" follow-up) add for high-resolution recovery.

# Gradient inversion (DLG) skeleton: recover a private input x
# from an observed gradient g_star, with model weights frozen.
# Optimization variable is the DUMMY input x_dummy, not the model.
import torch

def gradient_inversion(model, g_star, input_shape, true_label,
                       steps=300, lr=0.1):
    # Start the dummy input from random noise; label is read off
    # the last-layer gradient sign (iDLG), so we treat it as known.
    x_dummy = torch.randn(input_shape, requires_grad=True)
    optimizer = torch.optim.Adam([x_dummy], lr=lr)
    loss_fn = torch.nn.CrossEntropyLoss()

    for step in range(steps):
        optimizer.zero_grad()
        # Forward + backward on the dummy example to get its gradient.
        pred = model(x_dummy)
        dummy_loss = loss_fn(pred, true_label)
        dummy_grad = torch.autograd.grad(
            dummy_loss, model.parameters(), create_graph=True
        )
        # Gradient-matching objective: || dummy_grad - g_star ||^2
        match = sum(((dg - gs) ** 2).sum()
                    for dg, gs in zip(dummy_grad, g_star))
        match.backward()      # backprop THROUGH the gradient, into x_dummy
        optimizer.step()
        if step % 50 == 0:
            print(f"step {step:3d}  match_loss = {match.item():.4f}")
    return x_dummy.detach()

# Representative output (loss falls as x_dummy approaches the private x):
# step   0  match_loss = 14.8271
# step  50  match_loss =  2.1043
# step 100  match_loss =  0.3389
# step 150  match_loss =  0.0412
# step 200  match_loss =  0.0036

With the threat demonstrated, the two defenses already named in this section read as direct countermeasures rather than abstract good practice. Secure aggregation (detailed below) attacks the observability of $g^{*}$: the server never sees an individual gradient, only the masked sum, so there is no per-client signature to invert. DP-FL attacks the information content of $g^{*}$: clipping bounds how much any one example can move the gradient and the added Gaussian noise blurs the fingerprint, so the gradient-matching minimizer no longer lands on the true $x$. The residual risk is real and must be stated plainly: secure aggregation still leaks the aggregate (a sufficiently small cohort or a colluding majority can re-expose individuals), and DP only bounds reconstruction in expectation under its $(\epsilon, \delta)$ budget, so a loose budget still permits partial leakage. This is the privacy-utility tradeoff in concrete form: tighter clipping and larger noise shrink the attacker's reconstruction fidelity but also degrade the gradient signal the global model learns from, which is why, as noted below, the tradeoff bites harder for high-dimensional LLMs than for small classifiers.

Differential Privacy with Federated Learning (DP-FL)

Combining differential privacy with FL adds formal privacy guarantees. Each client clips and adds noise to their gradients before sending them to the server:

1. Gradient clipping: Bound the L2 norm of each client's update to a maximum value $C$, limiting any single example's influence
2. Noise addition: Add calibrated Gaussian noise $\mathcal{N}(0, \sigma^2 C^2 I)$ to the clipped gradients
3. Privacy accounting: Track the cumulative privacy budget $(\epsilon, \delta)$ across rounds using the moments accountant or Renyi DP

The privacy-utility tradeoff is more severe for LLMs than for smaller models: the high dimensionality means more noise is needed to achieve the same privacy guarantee, and LLMs are more sensitive to noisy gradients during fine-tuning.

Secure Aggregation

Secure aggregation protocols ensure the server can compute the aggregate update without seeing any individual client's contribution. This is achieved through cryptographic techniques (secret sharing, homomorphic encryption) that allow addition over encrypted values. Google's implementation in Gboard uses secure aggregation for federated next-word prediction on mobile devices.

50.3.4 Applications and Use Cases

50.3.5 Frameworks and Tools

• Flower (flwr): The most popular open-source FL framework. Framework-agnostic (supports PyTorch, TensorFlow, JAX), with built-in strategies for FedAvg, FedProx, and FedOpt. Provides simulation mode for development and deployment mode for production.
• NVIDIA FLARE: Enterprise FL framework with built-in privacy (DP, homomorphic encryption), designed for healthcare and regulated industries. Integrates with NVIDIA's NeMo for LLM fine-tuning.
• PySyft (OpenMined): Privacy-preserving ML framework supporting federated learning, differential privacy, and secure multi-party computation. Python-native with remote execution capabilities.
• FedML: Platform supporting federated training on heterogeneous devices (cloud, edge, mobile) with MLOps integration for monitoring and experiment tracking.

import torch
import flwr as fl
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import get_peft_model, LoraConfig
# Client: each participant runs this on their local data
class LLMClient(fl.client.NumPyClient):
    def __init__(self, model_name, local_dataset):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        base = AutoModelForCausalLM.from_pretrained(model_name)
        lora_config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"])
        self.model = get_peft_model(base, lora_config)
        self.dataset = local_dataset
    def get_parameters(self, config):
        """Return only LoRA adapter parameters."""
        return [
            val.cpu().numpy()
            for name, val in self.model.named_parameters()
            if "lora" in name
            ]
    def fit(self, parameters, config):
        """Train locally for E epochs, return updated adapter weights."""
        self.set_parameters(parameters)
        train_local(self.model, self.dataset, epochs=config.get("local_epochs", 3))
        return self.get_parameters(config), len(self.dataset), {}
    def set_parameters(self, parameters):
        lora_params = [
            (name, param) for name, param in self.model.named_parameters()
            if "lora" in name
            ]
        for (name, param), new_val in zip(lora_params, parameters):
            param.data = torch.tensor(new_val).to(param.device)
            # Server: aggregation strategy
            strategy = fl.server.strategy.FedAvg(
                min_fit_clients=3,
                min_available_clients=5,
                fraction_fit=0.5, # Sample 50% of available clients per round
                )
            # Launch federated training
            fl.server.start_server(
                server_address="0.0.0.0:8080",
                config=fl.server.ServerConfig(num_rounds=20),
                strategy=strategy,
                )

50.3.6 Challenges and Limitations

• Data heterogeneity (non-IID): The most persistent challenge. When client data distributions differ significantly, FedAvg can converge to a suboptimal model or diverge entirely. Algorithms like FedProx (adds a proximal term penalizing divergence from the global model) and SCAFFOLD (uses control variates to correct client drift) partially address this.
• Model poisoning: A malicious client can submit adversarial updates that degrade or backdoor the global model. Byzantine-robust aggregation methods (e.g., coordinate-wise median, trimmed mean) provide some defense.
• Evaluation without data access: The server cannot directly evaluate the global model on client data. Techniques include federated evaluation (clients evaluate locally and report metrics) and holdout validation sets.
• Convergence speed: FL typically requires many more communication rounds than centralized training to reach equivalent quality, especially with non-IID data and privacy constraints.
• Intellectual property concerns: Even without raw data sharing, the resulting model encodes knowledge from all participants. Ownership and usage rights of the federated model must be defined contractually.