Federated learning improves by tailoring memory to each client

Beyond Aggregate Drift: Optimal Client-Specific Windows in Federated Learning

Machine Learning

Summary

When many devices or clients learn together without sharing all their data, the data they have can change at different rates and have different noise levels. This paper studies how much past data each client should remember to best keep up with changes. The authors propose giving each client its own memory size, based on how quickly its data changes and how noisy it is, instead of using the same memory for all clients. This approach reduces errors in learning compared to using a shared memory window.

What this means in practice

  • For mobile app developers: Improve federated learning updates on phones by customizing how much past data each device uses to adjust for local data changes individually.
  • For edge computing teams: Optimize distributed learning across edge devices with varying data flows by allocating client-specific memory windows to reduce learning errors.

Authors

Yuriy Dorn, Dmitrii Vishovan, Kirill Novoselov, Ilya Kuruzov, Darina Dvinskikh, Alexander Gasnikov

Abstract

We study federated learning from continuously arriving data: each client samples from its own distribution, these distributions drift over time, and the server must track the minimizer of the clients' weighted average loss. Clients are not alike as they collect data at different rates, with different noise, and change at different speeds. The main question is how much history should each one keep? A single shared window must serve every client at once, and so makes a fast-changing client carry the sampling noise of a slow, nearly stationary one. We give each client a memory of its own, sized to that client's noise, arrival rate, and drift, and refresh it by restarted accelerated method that starts from a common point. The tracking error is then bounded by \[ R\left(\sum_m w_m^{3/2}\,σ_m\sqrt{γ_m/λ_m}\right)^{2/3}, \] for clients with weight $w_m$, noise $σ_m$, arrival rate $λ_m$, and drift rate $γ_m$ on a domain of diameter $R$. This is never worse than what any common window can achieve, and strictly better whenever noise and drift are spread unevenly across clients. A lower bound that puts different groups of clients on different drift clocks shows this quantity is the right one. The two bounds agree up to a constant whenever the clients occupy boundedly many refresh scales, and in general differ only by the sixth root of the number of scales. Experiments show the gain comes from the memory rather than from the solver: it persists across accelerated, stochastic-gradient, and one-step methods, and against a separately tuned common-memory baseline.