Why Your Phone Isn't Average
In a traditional datacenter training setup, data is shuffled so that every mini-batch is statistically representative of the full dataset. This property is called IID: Independent and Identically Distributed. Each batch looks roughly like a random sample from the global distribution.
In Federated Learning, data is inherently Non-IID — each device holds data that reflects only that user's behavior, demographics, and context:
- User A is a wildlife photographer: their phone has 10,000 photos of birds.
- User B is a chef: their phone has photos of food and recipe texts.
- User C types exclusively in French.
- User D is a teenager who uses slang that User C has never encountered.
When User A's phone trains the global model, it pushes the model toward predicting birds. When User B's phone trains next, it pushes toward food. These updates can actively conflict, slowing convergence or causing the model to oscillate rather than improve.
This phenomenon — where local models drift away from the global optimum due to non-representative local data — is called Client Drift. It is the primary reason why FedAvg, despite its elegance, can fail to converge on highly heterogeneous data.
Mitigations for Non-IID Data
FedProx (Federated Proximal) is the most widely adopted fix. It modifies the local training objective by adding a proximal term — a penalty that discourages the local model from straying too far from the global model during local training:
Minimize: F_k(w) + (μ/2) * ||w - w_t||²
The second term (μ/2) * ||w - w_t||² acts like a rubber band anchoring the local model to the global starting point. The hyperparameter μ controls the strength of this anchor. Higher μ = less drift but slower local learning. Lower μ = FedProx degrades to FedAvg.
Other approaches include:
- SCAFFOLD: Uses control variates to correct for client drift at the gradient level, achieving faster convergence than FedProx on highly heterogeneous data.
- FedNova: Normalizes local updates by the number of local steps taken, preventing devices that train longer from dominating the aggregation.
- Clustered FL: Groups devices with similar data distributions and trains separate global models per cluster, then merges. Useful when the population has genuinely distinct subgroups (e.g., language clusters).
The choice of algorithm depends on the degree of data heterogeneity (measurable via the Earth Mover's Distance between local and global distributions) and the communication budget. In practice, FedProx with a tuned μ is the default starting point for cross-silo deployments.