Achieving Dimension-Free Communication in Federated Learning via Zeroth-Order Optimization

TL;DR

This paper introduces DeComFL, a dimension-free federated learning algorithm using zeroth-order optimization, drastically reducing communication costs regardless of model size, with theoretical guarantees and practical validation on large models.

Contribution

DeComFL is the first dimension-free communication algorithm for federated learning leveraging zeroth-order methods, achieving constant communication per round and theoretical convergence guarantees.

Findings

Reduces communication from O(d) to O(1) per round

Achieves linear speedup with number of clients and local steps

Demonstrates significant practical reductions in communication overhead

Abstract

Federated Learning (FL) offers a promising framework for collaborative and privacy-preserving machine learning across distributed data sources. However, the substantial communication costs associated with FL significantly challenge its efficiency. Specifically, in each communication round, the communication costs scale linearly with the model's dimension, which presents a formidable obstacle, especially in large model scenarios. Despite various communication-efficient strategies, the intrinsic dimension-dependent communication cost remains a major bottleneck for current FL implementations. This paper proposes a novel dimension-free communication algorithm - DeComFL, which leverages the zeroth-order optimization techniques and reduces the communication cost from $\mathscr{O}(d)$ to $\mathscr{O}(1)$ by transmitting only a constant number of scalar values between clients and the server in each round, regardless of the dimension $d$ of the model parameters. Theoretically, in non-convex functions, we prove that our algorithm achieves state-of-the-art rates, which show a linear speedup of the number of clients and local steps under standard assumptions. With additional low effective rank assumption, we can further show the convergence rate is independent of the model dimension $d$ as well. Empirical evaluations, encompassing both classic deep learning training and large language model fine-tuning, demonstrate significant reductions in communication overhead. Notably, DeComFL achieves this by transmitting only around 1MB of data in total between the server and a client to fine-tune a model with billions of parameters. Our code is available at https://github.com/ZidongLiu/DeComFL.