Locally Adaptive Federated Learning via Stochastic Polyak Stepsizes
Sohom Mukherjee Nicolas Loizou Sebastian U. Stich
State-of-the-art federated learning algorithms such as FedAvg require
carefully tuned stepsizes to achieve their best performance. The improvements
proposed by existing adaptive federated methods involve tuning of additional
hyperparameters such as momentum parameters, and consider adaptivity only in
the server aggregation round, but not locally. These methods can be inefficient
in many practical scenarios because they require excessive tuning of
hyperparameters and do not capture local geometric information. In this work,
we extend the recently proposed stochastic Polyak stepsize (SPS) to the
federated learning setting, and propose new locally adaptive and nearly
parameter-free distributed SPS variants (FedSPS and FedDecSPS). We prove that
FedSPS converges linearly in strongly convex and sublinearly in convex settings
when the interpolation condition (overparametrization) is satisfied, and
converges to a neighborhood of the solution in the general case. We extend our
proposed method to a decreasing stepsize version FedDecSPS, that converges also
when the interpolation condition does not hold. We validate our theoretical
claims by performing illustrative convex experiments. Our proposed algorithms
match the optimization performance of FedAvg with the best tuned
hyperparameters in the i.i.d. case, and outperform FedAvg in the non-i.i.d.
case.