Optimal client training time improves federated learning accuracy and fairness

Local Epochs as Optimal Stopping: Theory and Evidence in Federated Learning

Machine LearningDistributed, Parallel, and Cluster Computing

Summary

Federated learning involves many devices training a shared model by doing some local work before sending updates. The challenge is deciding how long each device should train locally before communicating. The authors model this decision as a smart stopping problem that balances improvement against drift (errors from local changes). They create a practical stopping rule called FedStop that adapts training length and shows better worst-case accuracy without losing average performance.

What this means in practice

  • For machine learning engineers: Use FedStop to dynamically adjust client training time and improve model robustness in federated learning setups.
  • For mobile device developers: Apply the stopping rule to reduce communication and energy cost by letting devices stop training locally at an optimal point.

Authors

Ke Xiao, Qiyuan Wang, Christos Anagnostopoulos

Abstract

How long should each client train before communicating? Federated learning usually answers with a fixed, hand-tuned number of local epochs, yet an extra epoch helps only while the loss it removes outweighs the client drift it adds. We cast local training as finite-horizon optimal stopping: a client trains while the next epoch's expected loss reduction exceeds its priced squared drift. We prove that this one-stage look-ahead threshold is optimal whenever stopping decisions are nested, derive closed forms for geometric and power-law returns, and show that an estimated rule loses utility only through wrong-signed decisions. Epoch-level evidence supports the model's structure (geometric diminishing returns and additive drift) and reveals a simple regularity: the stage of training, not the client, sets the budget. The resulting rule, FedStop, is never more than one point below tuned FedAvg in mean accuracy. Where its budgets vary, it raises worst-decile accuracy by 4-8 points over tuned FedAvg and an adaptive-interval baseline, and every one of these gains remains significant after a Holm correction. Against the client-side stopping rules ALT and FLASH, it has the highest CIFAR-10 worst decile on every seed and is level elsewhere. It stays stable where per-client early stopping collapses, and each stopping decision takes constant time.