Federated Training Loop

The server selects a subset of available edge devices to participate in the current training round. Selection strategies balance system efficiency, statistical representativeness, and resource constraints.

• Common Criteria: Device connectivity, available compute/memory, battery level, and data relevance.
• Partial Participation: A practical reality where only a fraction of the total client pool is active per round.
• Straggler Mitigation: Techniques like deadline-based selection prevent slow devices from bottlenecking the entire round.

The server transmits the current global model parameters to all selected clients. This phase is communication-bound and often employs model compression techniques.

• Payload: The full set of neural network weights or a compressed/delta version.
• Efficiency Focus: Techniques like pruning, quantization, and knowledge distillation reduce broadcast size.
• Security Consideration: The model architecture is public, but the broadcast must be integrity-checked to prevent model poisoning attacks.

Each client performs Stochastic Gradient Descent (SGD) on its private, on-device dataset using the received global model as the initialization. This is the core computational workload.

• Local Epochs (E): The number of complete passes over the local data before stopping. A key hyperparameter.
• Client Drift: The primary challenge, where local optimization causes models to diverge from the global objective due to non-IID data.
• Algorithms: Vanilla SGD, or modified objectives like FedProx (adds a proximal term) to constrain drift.

Clients compute and send their model update (or model delta) back to the server. This is often the most sensitive communication, as updates can leak information about the private training data.

• The Update: Typically the difference between the final local weights and the initial global weights (Δ = w_local - w_global).
• Privacy & Security: Updates are protected via secure aggregation, homomorphic encryption, or differential privacy noise addition.
• Compression: Update compression via sparsification or quantization is critical to reduce uplink bandwidth costs.

The server combines the received client updates to form a new global model. The standard method is weighted averaging, where each client's update is weighted by its local dataset size.

• FedAvg Formula: w_global_new = Σ (n_k / n) * w_local_k, where n_k is client k's data size and n is the total data size.
• Synchronous vs. Asynchronous: The server either waits for all selected clients (synchronous) or updates immediately upon receipt (asynchronous).
• Advanced Aggregators: FedOpt framework uses optimizers like Adam on the server; SCAFFOLD uses control variates to correct drift.

The aggregated update is applied to the global model, concluding one communication round. The server then evaluates the new model's performance to determine loop continuation.

• Convergence Check: The server monitors loss/accuracy on a held-out validation set (which must not violate privacy).
• Metrics: Evaluation uses federated evaluation metrics that account for statistical heterogeneity across the network.
• Termination: The loop repeats until a target accuracy is met, convergence plateaus, or a maximum round count is reached.

The core assumption of independent and identically distributed (IID) data is violated in federated learning. Each client's local dataset has a unique distribution, leading to client drift where local models diverge from the global objective. This heterogeneity is the primary cause of slow convergence and reduced final model accuracy. For example, a keyboard next-word prediction model trained across users will see vastly different vocabularies and writing styles.

• Mitigation: Algorithms like FedProx, SCAFFOLD, and FedBN are designed to correct for this drift.

Clients vary dramatically in hardware (CPU, GPU, NPU), memory, network connectivity (5G, Wi-Fi, cellular), and availability (on/offline). This creates stragglers—slow clients that bottleneck synchronous training rounds. Managing this requires robust client selection and aggregation protocols.

• Strategies: Asynchronous aggregation updates the model immediately upon receiving any client update. Deadline-based selection only includes clients that respond within a time window. Partial client participation is a default design choice, not an exception.

Transmitting full model updates (e.g., for a 100M parameter model) over constrained or metered edge networks is prohibitively expensive and slow. The communication cost often dominates total training time, not local computation.

• Optimization Techniques: Update compression via quantization (reducing numerical precision) and sparsification (sending only the largest gradient values). Gradient clipping bounds update size and aids stability. Efficient protocols reduce the frequency of communication rounds.

While raw data never leaves the device, transmitting model updates can still leak sensitive information through techniques like model inversion or membership inference attacks. Providing formal privacy guarantees is a critical engineering requirement.

• Core Techniques: Secure Aggregation uses multi-party computation to sum client updates so the server only sees the aggregate, not individual contributions. Differential Privacy adds calibrated noise to updates before they are sent, providing a mathematically rigorous privacy bound.

The server cannot vet data on remote devices, making the system vulnerable to adversarial clients performing data poisoning (sending malicious updates) or model poisoning. A single malicious actor can corrupt the global model.

• Defensive Architectures: Robust aggregation rules like median-based or trimmed mean aggregation, which discard outlier updates. Reputation systems track client contribution history. Cryptographic verification of client identity and update integrity.

Coordinating thousands of ephemeral devices requires a robust orchestrator (e.g., TensorFlow Federated, Flower, NVIDIA FLARE) to manage the loop's state: model versioning, client eligibility, update logging, and fault tolerance. The server must handle clients dropping out mid-round without corrupting the aggregation process.

• Key Functions: Reliable checkpointing of the global model, idempotent operations for retries, and comprehensive telemetry for monitoring convergence and participant behavior across the entire federated network.

Federated Averaging (FedAvg) is the foundational iterative algorithm that orchestrates the federated training loop. It defines the standard sequence: server broadcasts a global model, selected clients train locally, and the server aggregates their updates via a weighted average.

• Core Mechanism: Clients perform multiple local epochs of SGD on their private data before sending a model delta.
• Aggregation: The server computes a new global model: w_global = Σ (n_k / N) * w_local^k, where weights are proportional to local dataset sizes.
• Purpose: Enables collaborative model training without centralizing raw data, forming the basis for all advanced variants.

A communication round is a single complete cycle of the federated training loop, representing the fundamental unit of federated learning progress.

• Phases: 1) Server-to-Client Broadcast: The global model is distributed. 2) Local Training: Clients compute updates over local epochs. 3) Client-to-Server Transmission: Model deltas are sent. 4) Server Aggregation: Updates are combined via weighted averaging.
• Synchronous vs. Asynchronous: In synchronous aggregation, the server waits for all selected clients. Asynchronous aggregation updates the model immediately upon receiving any client update, requiring careful handling of staleness.
• Bottleneck: Communication rounds are often the primary bottleneck, leading to techniques for update compression and straggler mitigation.

Client drift is the divergence of local client models from the global objective due to optimization on statistically heterogeneous (non-IID) data, which can destabilize convergence.

• Cause: When local data distributions differ significantly, local SGD steps pull models in conflicting directions.
• Impact: Manifests as reduced final accuracy, unstable loss curves, and slower convergence.
• Mitigation Algorithms:
  • FedProx: Adds a proximal term to the local loss, penalizing updates that stray too far from the global model.
  • SCAFFOLD: Uses control variates (correction terms) to reduce the variance between client updates.
  • FedNova: Normalizes client updates based on their local step counts to ensure consistent contribution magnitude.

Secure aggregation protocols are cryptographic methods that allow a server to compute the sum of client model updates without learning any individual client's contribution, providing an additional layer of privacy beyond the federated learning architecture.

• Principle: Uses techniques like secret sharing and homomorphic encryption. Clients can encrypt or mask their updates so that only the aggregated sum is decryptable by the server.
• Benefit: Protects against a curious or malicious server attempting to perform model inversion or membership inference attacks on individual updates.
• Challenge: Adds computational and communication overhead, which must be balanced against the privacy requirements of the use case (e.g., healthcare with HIPAA concerns).

Personalized federated learning refers to techniques that produce models tailored to individual clients' local data distributions, addressing the inherent statistical heterogeneity challenge.

• Motivation: A single global model may perform poorly on all clients if their data is highly non-IID.
• Key Approaches:
  • Personalization Layers (FedPer): Only the base model layers are learned collaboratively; the final classification/regression layers are trained locally and kept private.
  • Model Interpolation: Train a global model and local models, then combine them via fine-tuning or weighted averaging for inference.
  • Meta-Learning Frameworks: Use algorithms like Per-FedAvg to learn a model initialization that can be quickly adapted to any client with few local steps.

The FedOpt framework generalizes federated averaging by allowing the server to apply adaptive optimization algorithms (like Adam, Yogi, or SGD with momentum) to the aggregated client updates, rather than using simple averaging.

• Core Idea: Treats the aggregated client update as a pseudo-gradient. The server then applies an optimizer: w_t+1 = w_t - η * ServerOpt(Δ_aggregated).
• Algorithms:
  • FedAdam: Applies the Adam optimizer on the server, adapting the learning rate per parameter based on past update moments.
  • FedAvgM: Incorporates a server momentum term to smooth the update direction and accelerate convergence.
• Benefit: Can lead to faster convergence and better handling of heterogeneous client landscapes compared to vanilla FedAvg.