I'm moving away from centralized training to explore decentralized model training for a privacy-focused mobile app. I specifically want to implement the Federated Averaging (FedAvg) algorithm to keep user data on-device. However, I am stuck on the aggregation logic—how do I ensure the global model weights are correctly updated when I have a high volume of clients with non-IID data? Is there a standard way to weight the averages based on local sample sizes to ensure the model converges efficiently?
3 answers
Implementing FedAvg effectively requires a strict 4-step cycle: broadcast, local update, upload, and aggregation. The "averaging" part isn't just a simple mean; for best results, you must use a weighted average where each client's contribution is proportional to the number of local data samples they processed. In my work with PyTorch, I initialize a global model on a coordinator, then in each round, I select a random subset of clients. Each client performs several epochs of SGD locally. When the server receives the weights, it calculates the new global weights using $w_{t+1} = \sum \frac{n_k}{n} w_{t+1}^k$. This ensures that a client with 1,000 images has a larger influence on the gradient than one with only 10, preventing the model from drifting toward outliers.
That weighted average formula makes sense for keeping the model balanced, but how do you handle clients that drop off mid-training? If a client with a large dataset fails to upload its weights due to a poor connection, doesn't that significantly skew the global model for that specific round?
In my experience with TensorFlow Federated, it is often better to use a small learning rate on the server-side optimizer to "smooth out" the updates from heterogeneous clients.
I agree with Matthew. A server-side learning rate acts like a stabilizer. I also suggest checking out "FedProx" if your client data is extremely non-IID, as it adds a proximal term to the local objective to keep updates closer to the global model.
That is a classic "straggler" problem in Federated Learning, Barbara. To solve this, I usually implement a timeout threshold at the server level. If a client doesn't report back within the window, the server just proceeds with the available updates and scales the weights accordingly. I've found that using the Flower framework (flwr) handles this quite gracefully. It has built-in strategies to define a minimum number of fit clients before the aggregation even starts, which prevents a single round from being ruined by a few disconnected nodes.