Painless Federated Learning: An Interplay of Line-Search and Extrapolation

TL;DR

This paper introduces FedSLS and FedExpSLS algorithms that adapt line search and extrapolation techniques to federated learning, improving convergence rates and performance in heterogeneous and noisy data environments.

Contribution

It proposes novel federated stochastic line search and extrapolation methods with theoretical convergence guarantees and empirical validation.

Findings

FedSLS achieves linear convergence for strongly convex problems.

FedExpSLS improves empirical performance with extrapolation.

Methods outperform popular federated algorithms across various tasks.

Abstract

The classical line search for learning rate (LR) tuning in the stochastic gradient descent (SGD) algorithm can tame the convergence slowdown due to data-sampling noise. In a federated setting, wherein the client heterogeneity introduces a slowdown to the global convergence, line search can be relevantly adapted. In this work, we show that a stochastic variant of line search tames the heterogeneity in federated optimization in addition to that due to client-local gradient noise. To this end, we introduce Federated Stochastic Line Search (FedSLS) algorithm and show that it achieves deterministic rates in expectation. Specifically, FedSLS offers linear convergence for strongly convex objectives even with partial client participation. Recently, the extrapolation of the server's LR has shown promise for improved empirical performance for federated learning. To benefit from extrapolation, we extend FedSLS to Federated Extrapolated Stochastic Line Search (FedExpSLS) and prove its convergence. Our extensive empirical results show that the proposed methods perform at par or better than the popular federated learning algorithms across many convex and non-convex problems.