Federated Averaging (FedAvg)

FedAvg operates via a repeated communication round structure, forming its fundamental training loop:

1. Server Broadcast: The central server sends the current global model to a selected subset of clients.
2. Local Training: Each client downloads the model and performs local epochs of Stochastic Gradient Descent (SGD) on its private data, creating a local model.
3. Update Transmission: Clients compute a model delta (the difference between local and global parameters) and send it to the server.
4. Secure Aggregation: The server performs update aggregation, typically a weighted average based on client dataset sizes, to produce a new global model. This loop continues until convergence, ensuring raw training data never leaves the client device.

A defining challenge FedAvg must address is statistical heterogeneity, where client data is not independent and identically distributed (non-IID). This is the norm in real-world federated settings (e.g., different user typing habits on phones).

• Client Drift: Local optimization on divergent data causes client models to diverge from the global objective, hindering convergence.
• FedAvg's Implicit Robustness: By averaging many local updates, FedAvg empirically demonstrates a degree of inherent robustness to heterogeneity.
• Explicit Mitigations: Advanced variants like FedProx (adds a proximal term to the local loss) and SCAFFOLD (uses control variates) were developed to explicitly correct for client drift, improving upon basic FedAvg under high heterogeneity.

FedAvg is designed for a realistic partial client participation scenario, where only a fraction of the total client pool is available each round due to factors like:

• Connectivity: Devices are intermittently online (mobile phones, IoT sensors).
• Resource Constraints: Devices may be charging, idle, or have limited compute budget.
• System Heterogeneity: Devices have varying hardware capabilities (stragglers). The algorithm does not require all clients to participate every round. Server-side weighted averaging naturally accounts for which clients contributed. Straggler mitigation techniques, such as dropping slow clients after a deadline, are often layered on top to maintain system efficiency.

FedAvg's primary method for reducing communication costs—the bottleneck in cross-device FL—is to perform substantial local computation (multiple local epochs) between communication rounds.

• FedAvg vs. FedSGD: Compared to the baseline FedSGD (where clients perform one gradient step per round), FedAvg can converge in far fewer rounds by amortizing the communication cost over more local computation.
• Trade-off: Excessive local computation (too many epochs) can exacerbate client drift under non-IID data. Tuning the number of local epochs is critical.
• Complementary Techniques: FedAvg is often combined with update compression methods (quantization, sparsification) to further reduce the size of transmitted model deltas.

The canonical FedAvg algorithm uses a weighted averaging scheme for update aggregation on the server. The new global model parameters (\theta_{t+1}) are computed as: (\theta_{t+1} = \sum_{k \in S_t} \frac{n_k}{n} \theta_{t+1}^k) Where (S_t) is the set of participating clients, (n_k) is the number of samples on client (k), and (n) is the total samples across participating clients. This weighting aligns the global objective with the overall data distribution.

• Synchronous Aggregation: The standard approach waits for all selected clients in a round.
• Generalization (FedOpt): The FedOpt framework generalizes this step, allowing the server to apply adaptive optimizers like FedAdam or FedAvgM (with momentum) to the aggregated updates for potentially faster convergence.

FedAvg provides a foundational privacy-by-architecture benefit: raw training data remains on-device. However, transmitted model updates can potentially leak information about the underlying data through inference attacks.

• Privacy Baseline: FedAvg alone does not offer formal privacy guarantees. It reduces the attack surface compared to data centralization.
• Enhancing Privacy: FedAvg is the essential substrate upon which formal privacy techniques are integrated:
  • Secure Aggregation Protocols: Cryptographically combine updates so the server only sees the aggregate, not individual contributions.
  • Differential Privacy: Adding calibrated noise to client updates before aggregation provides mathematically bounded privacy guarantees. Thus, FedAvg enables privacy-preserving learning but must be combined with additional mechanisms for robust, verifiable privacy.

The core assumption of independent and identically distributed (IID) data is routinely violated in federated settings. Each client's local dataset is drawn from a unique distribution (e.g., different user writing styles on phones, regional health data patterns). This statistical heterogeneity causes client drift, where local models diverge from the global objective, severely hindering convergence and final model accuracy. Mitigation strategies include FedProx, SCAFFOLD, and personalized approaches like FedPer.

Clients vary drastically in their hardware (CPU, memory), connectivity (bandwidth, latency), and availability (battery, network). This system heterogeneity creates stragglers—slow or intermittently available devices that bottleneck the synchronous Federated Training Loop. Solutions involve:

• Asynchronous Aggregation protocols.
• Deadline-based client selection.
• Partial client participation per round.
• Update compression to reduce transmission size for bandwidth-constrained devices.

Transmitting full model updates (millions of parameters) over wireless networks in each communication round is prohibitively expensive and slow. The communication cost often dominates total training time. Key techniques to improve efficiency include:

• Local Epochs > 1 to compute more locally per transmission.
• Update Compression via quantization (reducing numerical precision) and sparsification (sending only the largest gradient values).
• Federated Dropout to train smaller sub-networks.

Federated learning protects raw data but exposes model updates to inference attacks. Model inversion or membership inference attacks can potentially reconstruct training data characteristics from shared gradients. Defenses require a layered approach:

• Secure Aggregation Protocols using cryptographic multi-party computation to combine updates without revealing individual contributions.
• Differential Privacy by adding calibrated noise to client updates before transmission.
• Gradient Clipping to bound update magnitude, aiding both stability and privacy.

The open participation model means some clients may be unreliable (sending corrupted updates due to faulty sensors) or actively Byzantine (sending malicious updates to poison the global model). This requires robust aggregation and validation mechanisms:

• Robust Aggregation Rules (e.g., median, trimmed mean) instead of simple weighted averaging.
• Update validation via reputation systems or anomaly detection on received model deltas.
• Federated learning attack mitigation frameworks to filter outlying updates.

Monitoring convergence and evaluating the final global model is challenging without access to a centralized test set. Performance may vary wildly across clients due to statistical heterogeneity. Engineers must implement:

• Federated evaluation metrics computed securely across held-out client data.
• Personalized Federated Learning (e.g., FedRep, FedBN) to tailor models per client when a single global model performs poorly.
• Careful tracking of loss and accuracy per client cohort to detect fairness issues.

FedSGD (Federated Stochastic Gradient Descent) is the most basic federated optimization baseline. In each communication round, each client computes the gradient on its local data and sends this gradient (not a full model update) to the server. The server then aggregates these gradients, typically via a weighted average, and applies a single global update. It is mathematically equivalent to FedAvg configured with one local epoch and a batch size equal to the client's entire dataset. Its high communication frequency makes it impractical for most real-world edge deployments but serves as an important theoretical benchmark.

Client drift is a critical convergence challenge in federated learning caused by statistical heterogeneity (non-IID data). As clients perform multiple local epochs of training on their unique data distributions, their local models optimize for their local objectives, causing them to diverge from the ideal global objective. This divergence introduces variance into the updates sent to the server, slowing convergence and potentially leading the global model to a sub-optimal solution. Algorithms like FedProx and SCAFFOLD are specifically designed to mitigate client drift.

FedProx is a federated optimization algorithm that directly addresses client drift. It modifies the local training objective on each client by adding a proximal term. This term penalizes the local model parameters for straying too far from the initial global model parameters received at the start of the round. The strength of this penalty is controlled by a hyperparameter (μ). By constraining local updates, FedProx ensures clients' solutions remain closer together, leading to more stable convergence and improved performance on heterogeneous data, especially when clients perform many local epochs.

SCAFFOLD (Stochastic Controlled Averaging for Federated Learning) is an advanced algorithm that uses control variates to correct for client drift. Each client and the server maintain a control variable that estimates the update direction. The key innovation is that clients use these control variates to adjust their local gradient steps, effectively removing the client-specific "drift" component from their updates. This correction reduces the variance between client updates, enabling significantly faster and more robust convergence under non-IID data conditions compared to standard FedAvg.

The FedOpt framework generalizes the server-side aggregation step of federated learning. Instead of using a simple (weighted) average of client updates, the server applies adaptive optimization algorithms—like those used in centralized deep learning—to the aggregated updates. For example:

• FedAdam: Applies the Adam optimizer to aggregated client gradients.
• FedYogi: A variant of FedAdam with a different adaptive learning rate adjustment.
• FedAvgM: Incorporates server momentum into the aggregation step. This framework allows the global model update to account for past update history, often leading to faster convergence and better handling of heterogeneous client participation.

These are two fundamental coordination protocols for the federated training loop:

• Synchronous Aggregation (used in standard FedAvg): The server waits for updates from all clients selected in a communication round before aggregating them and updating the global model. This is simpler and provides strong convergence guarantees but is vulnerable to stragglers (slow clients) that bottleneck the entire round.

• Asynchronous Aggregation: The server updates the global model immediately upon receiving an update from any single client. This eliminates waiting time and improves system efficiency but introduces complexity, as updates are based on stale global models, which can harm convergence stability and requires careful algorithmic design.