summaries: This paper proposes and analyzes a framework for Federated optimization called FedLin. It combines client-specific local steps with (possibly) biased compression at the server and/or clients and guarantees linear convergence to the exact minimum. Algorithmic contributions are supported with lower bounds for the case of infrequent communication and carefully designed synthetic experiments on least squares and logistic regression problems, which highlight theoretical claims.

summaries: This paper considers a standard distributed optimization setup in which clients periodically coordinate with a server to train a global model. The focus is on the deterministic setting, in which each worker computes full local gradients. It is first highlighted that standard algorithms for this problem do not converge to the true global minimizer without trading off speed (for example through diminishing step-sizes).  Thus, the FedLin algorithm, which converges linearly to the true minimizer, is introduced to get the best of both worlds. Several theorems covering various settings (strongly convex, convex, non-convex, stochastic) are given, and show that FedLin matches the communication complexity of parallel GD, regardless of the number of local steps used. Then, a matching lower bound is given for FedLin to show that the strongly convex convergence result is tight.    Then, several gradient sparsification schemes are investigated (at the server and at the workers), which show that server-side sparsification is far less impactful than client-side sparsification. The impact of gradient sparsification is evaluated experimentally in Section 7.

summaries: This paper proposes a new algorithm called *FedLin* which,   T.1) queries a full gradient oracle,   T.2) de-biases the local gradients using a correction term (with the last synchronized iterate in memory) to deal with objective heterogeneity,  T.3) error corrects to deal with compression at the client and the server, and  T.4) uses a client-specific learning rate to deal with the different number of local steps on each client.   Thus, FedLin is able to accommodate both objective and system heterogeneity along with sparse gradients and local steps. In doing so it is an early work that can deal with these multiple aspects of federated learning at once, albeit at the cost of an expensive gradient oracle.   The paper provides the following theoretical guarantees for FedLin,  R.1) matching upper and lower bounds in the strongly convex-smooth setting without compression,  R.2) upper bounds without compression in the convex-smooth setting and the non-convex setting,  R.3) an iterate sub-optimality recursion for the strongly convex-smooth setting while using a stochastic gradient oracle,  R.4) upper bounds for client and server compression in the strongly convex-smooth setting.  See the main review for comments on these results.

summaries: This paper proposes an optimization method for federated learning, FedLin, that is specifically designed to ensure fast convergence even in the presence of heterogeneous data. Notably, the authors show that FedLin converges to the global optimum (or just a critical point, in the case of non-convex functions) without decaying the client learning rate. This is in contrast to methods such as FedAvg, FedProx, and FedNova, which require decaying the client loss in order to circumvent "objective mismatch", where the methods are actually optimizing an altered surrogate loss.  This work also derives lower bounds matching the convergence of FedLin on quadratics, that explain the price and cost of performing multiple local update steps. They also analyze convergence using gradient sparsification at both the client and server. Finally, the paper shows the empirical benefits of FedLin over other methods on synthetic least squares and logistic regression problems.

abstractText: We consider a standard federated learning (FL) setup where a group of clients periodically coordinate with a central server to train a statistical model. We develop a general algorithmic framework called FedLin to tackle some of the key challenges intrinsic to FL, namely objective heterogeneity, systems heterogeneity, and infrequent and imprecise communication. Our framework is motivated by the observation that under these challenges, various existing FL algorithms suffer from a fundamental speed-accuracy conflict: they either guarantee linear convergence but to an incorrect point, or convergence to the global minimum but at a sub-linear rate, i.e., fast convergence comes at the expense of accuracy. In contrast, when the clients’ local loss functions are smooth and strongly convex, we show that FedLin guarantees linear convergence to the global minimum, despite arbitrary objective and systems heterogeneity. We then establish matching upper and lower bounds on the convergence rate of FedLin that highlight the effects of infrequent, periodic communication. Finally, we show that FedLin preserves linear convergence rates under aggressive gradient sparsification, and quantify the effect of the compression level on the convergence rate. Notably, our work is the first to provide tight linear convergence rate guarantees, and constitutes the first comprehensive analysis of gradient sparsification in FL.

authors:
Jakub Konečnỳ
H Brendan McMahan
Daniel Ramage
Peter Richtárik

title: Federated optimization: Distributed machine learning for on-device intelligence

venue: arXiv preprint arXiv:1610.02527,

authors:
Brendan McMahan
Eider Moore
Daniel Ramage
Seth Hampson
Blaise Aguera y Arcas

title: Communication-efficient learning of deep networks from decentralized data

venue: In Artificial Intelligence and Statistics,

authors:
Keith Bonawitz
Hubert Eichner
Wolfgang Grieskamp
Dzmitry Huba
Alex Ingerman
Vladimir Ivanov
Chloe Kiddon
Jakub Konečnỳ
Stefano Mazzocchi
H Brendan McMahan

title: Towards federated learning at scale: System design

venue: arXiv preprint arXiv:1902.01046,

authors:
Peter Kairouz
H Brendan McMahan
Brendan Avent
Aurélien Bellet
Mehdi Bennis
Arjun Nitin Bhagoji
Keith Bonawitz
Zachary Charles
Graham Cormode
Rachel Cummings

title: Advances and open problems in federated learning

venue: arXiv preprint arXiv:1912.04977,

authors:
Tian Li
Anit Kumar Sahu
Ameet Talwalkar
Virginia Smith

title: Federated learning: Challenges, methods, and future directions

venue: IEEE Signal Processing Magazine,

authors:
Xiang Li
Kaixuan Huang
Wenhao Yang
Shusen Wang
Zhihua Zhang

title: On the convergence of fedavg on non-iid data

authors:
Grigory Malinovskiy
Dmitry Kovalev
Elnur Gasanov
Laurent Condat
Peter Richtarik

title: From local sgd to local fixed-point methods for federated learning

venue: In International Conference on Machine Learning,

authors:
Zachary Charles
Jakub Konečnỳ

title: On the outsized importance of learning rates in local update methods

venue: arXiv preprint arXiv:2007.00878,

authors:
Zachary Charles
Jakub

title: Konečnỳ. Convergence and accuracy trade-offs in federated learning and meta-learning

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
Reese Pathak
Martin J Wainwright

title: Fedsplit: An algorithmic framework for fast federated optimization

venue: arXiv preprint arXiv:2005.05238,

authors:
Sai Praneeth Karimireddy
Satyen Kale
Mehryar Mohri
Sashank Reddi
Sebastian Stich
Ananda Theertha Suresh

title: Scaffold: Stochastic controlled averaging for federated learning

title: Local sgd converges fast and communicates little

venue: arXiv preprint arXiv:1805.09767,

authors:
Jianyu Wang
Gauri Joshi

title: Cooperative sgd: A unified framework for the design and analysis of communication-efficient sgd algorithms

venue: arXiv preprint arXiv:1808.07576,

authors:
Artin Spiridonoff
Alex Olshevsky
Ioannis Ch Paschalidis

title: Local sgd with a communication overhead depending only on the number of workers

venue: arXiv preprint arXiv:2006.02582,

authors:
Amirhossein Reisizadeh
Aryan Mokhtari
Hamed Hassani
Ali Jadbabaie
Ramtin Pedarsani

title: Fedpaq: A communication-efficient federated learning method with periodic averaging and quantization

authors:
Farzin Haddadpour
Mohammad Mahdi Kamani
Mehrdad Mahdavi
Viveck Cadambe

title: Local sgd with periodic averaging: Tighter analysis and adaptive synchronization

venue: In Advances in Neural Information Processing Systems,

authors:
Blake Woodworth
Kumar Kshitij Patel
Sebastian U Stich
Zhen Dai
Brian Bullins
H Brendan McMahan
Ohad Shamir
Nathan Srebro

title: Is local sgd better than minibatch sgd

authors:
Ahmed Khaled
Konstantin Mishchenko
Peter Richtárik

title: First analysis of local gd on heterogeneous data

venue: arXiv preprint arXiv:1909.04715,

title: Tighter theory for local sgd on identical and heterogeneous data

authors:
Farzin Haddadpour
Mehrdad Mahdavi

title: On the convergence of local descent methods in federated learning

venue: arXiv preprint arXiv:1910.14425,

authors:
Anastasia Koloskova
Nicolas Loizou
Sadra Boreiri
Martin Jaggi
Sebastian U Stich

title: A unified theory of decentralized sgd with changing topology and local updates

authors:
Anit Kumar Sahu
Tian Li
Maziar Sanjabi
Manzil Zaheer
Ameet Talwalkar
Virginia Smith

title: On the convergence of federated optimization in heterogeneous networks

venue: arXiv preprint arXiv:1812.06127,

authors:
Jianyu Wang
Qinghua Liu
Hao Liang
Gauri Joshi
H Vincent Poor

title: Tackling the objective inconsistency problem in heterogeneous federated optimization

venue: Advances in Neural Information Processing Systems,

authors:
Alham Fikri Aji
Kenneth Heafield

title: Sparse communication for distributed gradient descent

venue: arXiv preprint arXiv:1704.05021,

authors:
Yujun Lin
Song Han
Huizi Mao
Yu Wang
William J Dally

title: Deep gradient compression: Reducing the communication bandwidth for distributed training

venue: arXiv preprint arXiv:1712.01887,

authors:
Dan Alistarh
Demjan Grubic
Jerry Li
Ryota Tomioka
Milan Vojnovic

title: Qsgd: Communication-efficient sgd via gradient quantization and encoding

authors:
Wei Wen
Cong Xu
Feng Yan
Chunpeng Wu
Yandan Wang
Yiran Chen
Hai Li

title: Terngrad: Ternary gradients to reduce communication in distributed deep learning

venue: In Advances in neural information processing systems,

authors:
Sarit Khirirat
Hamid Reza Feyzmahdavian
Mikael Johansson

title: Distributed learning with compressed gradients

venue: arXiv preprint arXiv:1806.06573,

authors:
Konstantin Mishchenko
Eduard Gorbunov
Martin Takáč
Peter Richtárik

title: Distributed learning with compressed gradient differences

venue: arXiv preprint arXiv:1901.09269,

authors:
Zhize Li
Dmitry Kovalev
Xun Qian
Peter Richtárik

title: Acceleration for compressed gradient descent in distributed and federated optimization

venue: arXiv preprint arXiv:2002.11364,

authors:
Samuel Horváth
Dmitry Kovalev
Konstantin Mishchenko
Sebastian Stich
Peter Richtárik

title: Stochastic distributed learning with gradient quantization and variance reduction

authors:
Farzin Haddadpour
Mohammad Mahdi Kamani
Aryan Mokhtari
Mehrdad Mahdavi

title: Federated learning with compression: Unified analysis and sharp guarantees

venue: arXiv preprint arXiv:2007.01154,

authors:
Frank Seide
Hao Fu
Jasha Droppo
Gang Li
Dong Yu

title: 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech dnns

venue: In Fifteenth Annual Conference of the International Speech Communication Association,

authors:
Sebastian U Stich
Jean-Baptiste Cordonnier
Martin Jaggi

title: Sparsified sgd with memory

authors:
Dan Alistarh
Torsten Hoefler
Mikael Johansson
Nikola Konstantinov
Sarit Khirirat
Cédric Renggli

title: The convergence of sparsified gradient methods

authors:
Guannan Qu
Na Li

title: Harnessing smoothness to accelerate distributed optimization

venue: IEEE Transactions on Control of Network Systems,

authors:
Angelia Nedic
Alex Olshevsky
Wei Shi

title: Achieving geometric convergence for distributed optimization over time-varying graphs

venue: SIAM Journal on Optimization,

authors:
Chenguang Xi
Ran Xin
Usman A Khan

title: Add-opt: Accelerated distributed directed optimization

venue: IEEE Transactions on Automatic Control,

authors:
Shi Pu
Angelia Nedić

title: Distributed stochastic gradient tracking methods

venue: Mathematical Programming,

authors:
Ran Xin
Anit Kumar Sahu
Usman A Khan
Soummya Kar

title: Distributed stochastic optimization with gradient tracking over strongly-connected networks

venue: In Proc. of the 58th IEEE Conference on Decision and Control (CDC),

authors:
Arkadi Nemirovski
Anatoli Juditsky
Guanghui Lan
Alexander Shapiro

title: Robust stochastic approximation approach to stochastic programming

venue: SIAM Journal on optimization,

authors:
Eduard Gorbunov
Filip Hanzely
Peter Richtárik

title: Local sgd: Unified theory and new efficient methods

authors:
Amirhossein Reisizadeh
Isidoros Tziotis
Hamed Hassani
Aryan Mokhtari
Ramtin Pedarsani

title: Straggler-resilient federated learning: Leveraging the interplay between statistical accuracy and system heterogeneity

venue: arXiv preprint arXiv:2012.14453,

authors:
Sanghamitra Dutta
Jianyu Wang
Gauri Joshi

title: Slow and stale gradients can win the race

venue: arXiv preprint arXiv:2003.10579,

title: Adaptive communication strategies to achieve the best errorruntime trade-off in local-update sgd

venue: arXiv preprint arXiv:1810.08313,

authors:
Rawad Bitar
Mary Wootters
Salim El Rouayheb

title: Stochastic gradient coding for straggler mitigation in distributed learning

venue: IEEE Journal on Selected Areas in Information Theory,

authors:
Nuwan Ferdinand
Stark C Draper

title: Anytime stochastic gradient descent: A time to hear from all the workers

venue: In 2018 56th Annual Allerton Conference on Communication, Control, and Computing (Allerton),

authors:
Amirhossein Reisizadeh
Hossein Taheri
Aryan Mokhtari
Hamed Hassani
Ramtin Pedarsani

title: Robust and communication-efficient collaborative learning

authors:
Aleksandr Beznosikov
Samuel Horváth
Peter Richtárik
Mher Safaryan

title: On biased compression for distributed learning

venue: arXiv preprint arXiv:2002.12410,

title: Scalable distributed dnn training using commodity gpu cloud computing

venue: In Sixteenth Annual Conference of the International Speech Communication Association,

authors:
Sai Praneeth Karimireddy
Quentin Rebjock
Sebastian U Stich
Martin Jaggi

title: Error feedback fixes signsgd and other gradient compression schemes

authors:
Sebastian U Stich
Sai Praneeth Karimireddy

title: The error-feedback framework: Better rates for sgd with delayed gradients and compressed communication

venue: arXiv preprint arXiv:1909.05350,

authors:
Eduard Gorbunov
Dmitry Kovalev
Dmitry Makarenko
Peter Richtárik

title: Linearly converging error compensated sgd

title: Introductory lectures on convex optimization: A basic course, volume 87

venue: Springer Science & Business Media,

title: Convex optimization: Algorithms and complexity

venue: arXiv preprint arXiv:1405.4980,

heading: 1 Introduction

text:
In a canonical federated learning (FL) architecture, a set S of clients periodically communicate with a central server to find a global statistical model that solves the following problem [1–5]:
min x∈Rd
f(x), where f(x) = 1
m m∑ i=1 fi(x). (1)
Here, m is the number of clients, fi : Rd → R is the local objective (loss) function of client i, and f(x) is the global objective function. Some of the core distinguishing tenets of the FL paradigm are as follows [1–5]. First, due to privacy considerations, clients cannot directly share their local training data with the server. Second, differences in the clients’ data-sets may cause the clients to have nonidentical loss functions with different minima - this is known as statistical or objective heterogeneity. Third, due to variability in hardware (CPU, memory) and power (battery level), i.e., due to systems or device heterogeneity, the client devices may have different computation speeds; in particular, this may lead to slow and straggling devices that affect convergence guarantees. Finally, communicationefficiency is a major concern, dictating the need to reduce the number of communication rounds, and also the size of the messages transmitted in each round. The above considerations pose unique technical challenges that we aim to address in this paper.
In a typical FL setting, to reduce the number of communication rounds, clients perform multiple local training steps in isolation before communicating with the server. Due to such local steps, the popular
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
FedAvg algorithm suffers from a “client-drift phenomenon" under objective heterogeneity [6–11]: the local iterates of each client drift-off towards the minimum of their own local loss function, leading to slow convergence rates. For analysis on FedAvg, we refer the reader to [6, 8, 12–21]. Recently, several new algorithms such as FedProx [22], SCAFFOLD [11], FedSplit [10], and FedNova [23] have been proposed as improvements to FedAvg. Despite these advances, there remain gaps in our understanding of the extent to which these algorithms match the guarantees of a centralized baseline.1
For instance, even for simple, deterministic settings, FedProx [22] and FedNova [23] exhibit a fundamental speed-accuracy conflict under objective heterogeneity; see [8, 9] and Section 2. Specifically, with constant step-sizes, these algorithms converge linearly, but potentially to an incorrect point. Thus, convergence to the minimum of the global loss function necessitates diminishing step-sizes, which, in turn, leads to sub-linear convergence. Thus, fast convergence comes at the expense of accuracy. Although SCAFFOLD [11] and FedSplit [10] employ variance-reduction and operatorsplitting techniques, respectively, to tackle objective heterogeneity, it is not known whether the rates in these papers are tight. More importantly, neither SCAFFOLD nor FedSplit account for the effects of systems heterogeneity or compression, both of which are key challenges in FL. Indeed, due to systems heterogeneity, the number of local steps may vary across clients, causing some clients to make much less progress than others in each round [23]. Moreover, while empirical studies [24, 25] have revealed significant benefits of biased sparsification, theoretical guarantees for such methods in a federated setting have remained elusive. In this context, our contributions are as follows.
• A New Algorithm: Motivated by the above concerns, we develop a general algorithmic framework called FedLin that simultaneously accounts for objective heterogeneity, systems heterogeneity, and gradient sparsification. The key components of FedLin include a gradient correction term in the local update rule that exploits memory; the use of client-specific learning rates; and error-feedback mechanisms at the clients and the server.
• Matching Centralized Rates: For smooth and strongly convex losses, we show that FedLin converges to the global minimum linearly in the deterministic setting, and with a O(1/T ) rate for a general stochastic oracle model, thereby matching centralized rates (up to constants). We then present matching rates for smooth, convex and non-convex settings as well. Importantly, our results hold under arbitrary objective and systems heterogeneity. In contrast, the only other work in FL (as far as we are aware) that investigates both objective and systems heterogeneity [23] provides results only for the non-convex setting, under a bounded dissimilarity assumption. Moreover, the FedNova algorithm in [23] suffers from the speed-accuracy conflict, while FedLin does not.
• Quantifying the Price of Multiple Local Steps: We establish a lower bound for FedLin that matches the upper-bound we obtain for smooth, strongly convex losses. In doing so, we provide the first (as far as we are aware) tight linear convergence rate analysis. Our lower bound highlights the price paid for performing multiple local steps, i.e., the effect of infrequent communication on the convergence rate. In particular, our analysis reveals, perhaps surprisingly, that there exist simple instances (involving quadratic losses) for which performing multiple local steps does not improve the rate of convergence, indicating that even mild statistical heterogeneity can hurt. Our analysis also provides valuable insights into the limitations of gradient-tracking/variance-reduction techniques.
• Analyzing the Impacts of Gradient Sparsification at Server and at Clients: While several works explore the effect of unbiased random quantization in distributed settings [26–31], there are only a handful of papers [15, 32] that also consider the effect of local steps in FL. Different from all these works, we explore the impacts of sparsifying gradients using a biased TOP-k operator, both at the server side and at the clients. Our results in this context (i) constitute the first formal study of gradient sparsification in a federated setting; (ii) reveal key differences between up-link and down-link compression; and (iii) quantify the effect of the compression level on the convergence rate. Notably, FedLin preserves linear convergence rates despite aggressive gradient sparsification.
Basic Notation and Terminology: Referring to (1), let x∗ ∈ argminx∈Rd f(x), and x∗i ∈ argminx∈Rd fi(x). Every FL algorithm mentioned in this paper operates in rounds t ∈ {1, . . . , T}. In each round t, every client performs a certain number of local steps in isolation, starting from a common global model x̄t. We will denote by x (t) i,` client i’s estimate of the model at the `-th local step of round t. In particular, x(t)i,0 = x̄t,∀i ∈ S. 1By a centralized baseline, we refer to a setup where each client can communicate with every other client at all time steps via the server.

heading: 2 Motivation: Speed-Accuracy Trade-Off

text:
To motivate our work, we first show how some recently proposed FL algorithms, namely FedProx [22] and FedNova [23], exhibit a fundamental speed-accuracy trade-off even in simple, deterministic settings. Specifically, we show that these schemes do not, in general, guarantee convergence to the global minimum with constant step-sizes. This, in turn, necessitates diminishing step-sizes, leading to sub-linear convergence rates. Our analysis here is inspired by that in [8] for FedAvg. We consider a deterministic quadratic model where the local loss function of client i is given by
fi(x) = 1/2‖A1/2i (x− ci)‖2, where Ai is a symmetric positive-definite matrix. We begin by assuming that all clients perform the same number of local steps H . The following is the FedProx update rule where a proximal term is added to mitigate client-drift:
x (t) i,`+1 = x (t) i,` − η ( ∇fi(x(t)i,` ) + β(x (t) i,` − x̄t) ) , ` = 0, . . . ,H − 1; x̄t+1 = 1
m ∑ i∈S x (t) i,H . (2)
Proposition 1. For any step-size η > 0, T rounds of FedProx amount to performing T rounds of parallel GD on the surrogate optimization problem given by
min x
1
m ∑ i∈S 1 2 ∥∥∥∥(H−1∑ `=0 [I − η(Ai + βI)]`Ai )1/2 (x− ci) ∥∥∥∥2. (3)
Proposition 1 shows that even when clients perform the same number of local updates, FedProx minimizes a surrogate objective function (3) whose minimum may not, in general, coincide with the minimum of the original problem. When β = 0, FedProx reduces to FedAvg, and our observations continue to hold. To capture systems heterogeneity as in [23], suppose now that client i performs τi local steps. Define τeff , 1/m ∑ i∈S τi and αi , τeff/τi, ∀i ∈ S. The update rule of FedNova relies on normalized aggregation of cumulative local gradients, and is given by
x (t) i,`+1 = x (t) i,` − η∇fi(x (t) i,` ); x̄t+1 = x̄t −
η
m ∑ i∈S αi τi−1∑ `=0 ∇fi(x(t)i,` ), (4)
where ` = 0, . . . , τi − 1, i ∈ S. Although FedNova can accommodate any local solver whose accumulated gradients are expressible as a linear combination of local gradients, we choose gradient descent, a simple solver, to isolate the impact of normalized aggregation - the essence of FedNova.
Algorithm 1 FedLin 1: Input: Client step-sizes ηi, i ∈ S, compression levels δc and δs, initial iterate x̄1 ∈ Rd, g1 = ∇f(x̄1), initial compression errors ρi,1 = 0,∀i ∈ S and e1 = 0
2: for t = 1, . . . , T do 3: for i = 1, . . . ,m do 4: for ` = 0, . . . , τi − 1 do 5: x(t)i,`+1 ← x (t) i,` − ηi(∇fi(x (t) i,` )−∇fi(x̄t) + gt); x (t) i,0 = x̄t 6: end for 7: Transmit x(t)i,τi to server 8: end for 9: Server transmits x̄t+1 = 1/m ∑ i∈S x (t) i,τi
10: for i = 1, . . . ,m do 11: Transmit hi,t+1 = Cδc(ρi,t +∇fi(x̄t+1)) to server 12: ρi,t+1 ← ρi,t +∇fi(x̄t+1)− hi,t+1 13: end for 14: Server transmits gt+1 = Cδs(et + 1/m ∑ i∈S hi,t+1)
15: et+1 ← et + 1/m ∑ i∈S hi,t+1 − gt+1 16: end for
Proposition 2. For any step-size η > 0, T rounds of FedNova amount to performing T rounds of parallel GD on the surrogate optimization problem given by
min x
1
m ∑ i∈S 1 2 ∥∥∥∥( τi−1∑ `=0 [I − ηAi]`αiAi )1/2 (x− ci) ∥∥∥∥2. (5)
For the proofs of Propositions 1 and 2, see Appendix B. Proposition 2 shows that in the presence of both objective and systems heterogeneity, FedNova minimizes a surrogate loss function whose minimum may not coincide with x∗.2 Observe from (3) and (5) that using a larger learning rate η introduces more distortion to the original problem. In Figure 1, we see how FedProx and FedNova both converge to incorrect minimizers, even for simple instances with two clients and deterministic, quadratic losses. In contrast, FedLin, our proposed approach that we develop in the next section, guarantees linear convergence to the global minimum.
Main Takeaway: The main message we want to convey here is that even for deterministic settings, there are non-trivial challenges posed by objective and systems heterogeneity that only get amplified when one additionally considers biased compression. For such scenarios, it is not at all apparent whether (and to what extent) one can match even the basic centralized benchmark of achieving linear convergence for smooth, strongly convex loss functions. To focus on the above unresolved issues, we will primarily consider a deterministic model in this paper. Nonetheless, the general approach we develop applies to the stochastic setting as well, as aptly demonstrated by Theorem 4 in Section 4.
3 Proposed Algorithm: FedLin
In this section, we develop our proposed algorithm FedLin, formally described in Algorithm 1. FedLin is initialized from a common global iterate x̄1 ∈ Rd. For simplicity, we assume that g1 = ∇f(x̄1), i.e., every client has access to the true gradient of f(·) initially; we can allow g1 to be arbitrary as well without affecting the convergence guarantees. FedLin proceeds in rounds: in each round t, starting from a common global model x̄t, each client i performs τi local training steps in parallel, as per line 5 of Algorithm 1. The key features of our local update rule are as follows: exploiting past gradients to account for objective heterogeneity, using client-specific step-sizes to tackle systems heterogeneity, and employing error-feedback to account for gradient sparsification. We now discuss each of these features in detail.
2In a follow-up work to [8], the authors in [9] generalize their framework to encompass proximal methods such as FedProx as well. As such, Propositions 1 and 2 in this section turn out to be special cases of the results in [9]. We were not aware of [9] at the time of submission of this paper.
To gain intuition regarding the local step in line 5, note that the ideal local update at client i is x
(t) i,`+1 = x (t) i,` − ηi∇f(x (t) i,` ). However, this requires client i to have access to the gradients of all other clients - which it does not, since clients do not communicate between rounds. To get around this, client i exploits memory, and uses the gradient of the global function ∇f(x̄t) from the beginning of round t (when the clients last communicated) as a guiding direction in its update rule. However, since ∇f(x̄t) is evaluated at a stale point x(t)i,0 = x̄t, client i subtracts off ∇fi(x̄t) from ∇f(x̄t), and adds in the most recently evaluated gradient∇fi(x(t)i,` ). This results in the update rule: x
(t) i,`+1 = x (t) i,` − ηi(∇fi(x (t) i,` )−∇fi(x̄t) +∇f(x̄t)). Our local update rule in line 5 is precisely of the above form, where gt is an inexact version of ∇f(x̄t) to account for gradient sparsification. When each client i performs τi local-steps, our analysis reveals that the bound on the drift-term ‖xi,` − x̄t‖ scales linearly in τi (see Lemma 9 in Appendix F). Accordingly, to compensate for such drift at client i, the step-size ηi needs to be chosen to vary inversely with the number of local steps τi. In fact, the requirement that ηi ∝ 1/τi also turns out to be necessary (see Theorem 5), providing further motivation for the choice of client-specific learning rates in FedLin.
To explain the gradient sparsification module, let us denote by Cδ : Rd → Rd the TOP-k operator, where δ = d/k, and k ∈ {1, . . . , d}. Given any x ∈ Rd, let Eδ(x) be a set containing the indices of the k largest-magnitude components of x. Then, the TOP-k operator we consider is given by (Cδ(x))j = (x)j if j ∈ Eδ(x), and (Cδ(x))j = 0 otherwise. Here, we use (x)j to denote the j-th component of a vector x. Clearly, a larger δ implies more aggressive compression. We employ a standard error-feedback mechanism [33–35] at both the server and the clients to account for gradient sparsification. At client i, ρi,t represents the accumulated error due to gradient sparsification. At the end of round t, instead of just compressing∇fi(x̄t+1), client i instead compresses∇fi(x̄t+1) + ρi,t, to account for gradient coordinates not transmitted in the past. It then updates the aggregate error via line 12. An analogous description applies to the error-feedback scheme at the server, where et is the aggregate error at the beginning of round t. The parameters of FedLin are the client step-sizes {ηi}i∈S , and the compression levels δc and δs at the clients and at the server, respectively. We now comment on some related algorithmic ideas.
Related Algorithmic Approaches: In the related but different setting of distributed optimization, we note that the idea of exploiting past gradients has been used to design gradient-tracking algorithms [36–40]. In the context of FL, this idea is also related to the variance-reduction technique employed in SCAFFOLD [11]. A major difference of FedLin with the above works is that none of them consider the effect of systems heterogeneity or biased compression. In particular, accounting for the inexact gradient term gt in our update rule introduces new technical challenges that we address in this paper.
There are some additional basic differences between FedLin and SCAFFOLD. To see this, consider the update rule of FedLin without sparsification: x(t)i,`+1 = x (t) i,` − ηi(∇fi(x (t) i,` )−∇fi(x̄t) +∇f(x̄t)). Now suppose the global model x̄t at the beginning of round t has already converged to x∗. Since x
(t) i,0 = x̄t,∀i ∈ S , and ∇f(x∗) = 0, it is easy to see that the iterates of the clients do not evolve any further, as one would ideally want. Thus, the global optimum x∗ can be viewed as a fixed-point of the FedLin update rule. Adapting to our notation, and considering the case when there is no noise in the gradients, the update rule of SCAFFOLD takes the form x(t)i,`+1 = x (t) i,` − η(∇fi(x (t) i,` )− ci + c), where ci is a ‘control-variate’ maintained by client i, and c is the average of the ci’s. Importantly, the control variates {ci}i∈S used in round t of SCAFFOLD contain stale terms from round t− 1. As a result, even if x̄t = x∗, it may very well be that (∇fi(x̄t) − ci + c) 6= 0, causing the iterates of the clients to move away from x∗, and requiring further rounds of communication to average out the imbalance. Thus, the fixed-point property we discussed for FedLin does not hold in general for SCAFFOLD. Our simulations in Section 7 reveal that FedLin converges much faster relative to SCAFFOLD on a simple linear regression model; we conjecture it is precisely due to the reason described above.
Keeping aside the differences due to systems heterogeneity and compression, the FedSVRG algorithm in [1] includes a similar gradient correction term as in FedLin, but makes use of certain additional diagonal scaling and pre-conditioning matrices. Although promising empirical results are reported for FedSVRG in [1], these results come with no supporting theoretical guarantees of convergence. In contrast, we will develop rigorous complexity guarantees for FedLin in the following sections. Specifically, we will show that FedLin guarantees linear convergence rates despite the challenges of objective heterogeneity, systems heterogeneity, and aggressive gradient sparsification.

heading: 4 Matching Centralized Rates under Objective and Systems Heterogeneity

text:
In this section, we will analyze the performance of FedLin in the face of both objective and systems heterogeneity. To focus solely on the effects of client heterogeneity, we will assume throughout this section that there is no gradient sparsification, i.e., δc = δs = 1. Accordingly, observe that ρi,t = 0, et = 0,∀i ∈ S,∀t ∈ {1, . . . , T}. Thus, the local update rule for FedLin simplifies to
x (t) i,`+1 = x (t) i,` − ηi(∇fi(x (t) i,` )−∇fi(x̄t) +∇f(x̄t)). (6)
Let us denote by κ = L/µ the condition number of an L-smooth and µ-strongly convex function. Also, let ηi = η̄/τi,∀i ∈ S, where η̄ ∈ (0, 1) is a flexible parameter that we will specify based on context. We are now ready to state the main results of this section. Theorem 1. (Strongly convex case) Suppose each fi(x) is L-smooth and µ-strongly convex. Moreover, suppose τi ≥ 1,∀i ∈ S, and δc = δs = 1. Then, with ηi = 16Lτi ,∀i ∈ S, FedLin guarantees:
f(x̄T+1)− f(x∗) ≤ (
1− 1 6κ
)T (f(x̄1)− f(x∗)).
Theorem 2. (Convex case) Suppose each fi(x) is L-smooth and convex. Moreover, suppose τi ≥ 1,∀i ∈ S, and δc = δs = 1. Then, with ηi = 110Lτi ,∀i ∈ S, FedLin guarantees:
f
( 1
T T∑ t=1 x̄t
) − f(x∗) ≤ 10L
T
( ‖x̄1 − x∗‖2 − ‖x̄T+1 − x∗‖2 ) .
Theorem 3. (Non-convex case) Suppose each fi(x) is L-smooth. Moreover, suppose τi ≥ 1,∀i ∈ S , and δc = δs = 1. Then, with ηi = 126Lτi ,∀i ∈ S, FedLin guarantees:
min t∈[T ]
‖∇f(x̄t)‖2 ≤ 52L
T (f(x̄1)− f(x̄T+1)). (7)
Noisy Case Analysis: We now analyze the performance of FedLin under a general stochastic oracle model. For each i ∈ S and x ∈ Rd, let qi(x) be an unbiased estimate of the gradient ∇fi(x) with variance bounded above by σ2. We consider the update rule: x(t)i,`+1 = x (t) i,` − ηi(qi(x (t) i,` )− qi(x̄t) +
q(x̄t)), where q(x) , 1/m ∑ i∈S qi(x),∀x ∈ Rd. We then have the following result. Theorem 4. (Strongly convex case with noise) Consider the above stochastic oracle model. Suppose each fi(x) is L-smooth and µ-strongly convex. Moreover, suppose τi ≥ 1,∀i ∈ S , and δc = δs = 1. For each i ∈ S , let ηi = η̄τi , where η̄ ∈ (0, 1) satisfies η̄ < 1 6L . Then, ∀t ∈ [T ], FedLin guarantees:
E[‖x̄t+1 − x∗‖2] ≤ (
1− η̄µ 2
) E[‖x̄t − x∗‖2] + 25η̄2σ2. (8)
The proofs of Theorems 1, 2, 3, and 4 are provided in Appendix F.
Main Takeaways: From Theorems 1, 2, and 3, we note that FedLin matches the convergence guarantees of centralized gradient descent (up to constants) for smooth, strongly convex, convex, and non-convex settings, respectively. As far as we are aware, this is the first work to provide such comprehensive guarantees under arbitrary objective and systems heterogeneity. In fact, all our results continue to hold even when the operating speeds of the client machines vary across rounds, i.e., τi is allowed to be a function of t. Each client i can simply adjust its learning rate ηi ∝ 1/τi(t) locally to account for such variations. The bound for the noisy case in Theorem 4 resembles that of centralized SGD [41]: with a time-varying parameter η̄t = O(1/t), we get the standard O(1/T ) rate after T rounds (using the exact same arguments as in [41]). The key thing to note here is that despite arbitrary heterogeneity, the assumptions we make on the stochastic gradients are the same as those made in the analysis of centralized SGD: unbiased gradients with bounded variance, nothing more.
Comparison with Related Work: In the recent paper [10], the authors propose FedSplit, and analyze it in a deterministic setting. For strongly-convex and smooth loss functions, FedSplit guarantees linear convergence, but only to a non-vanishing neighborhood of x∗. Thus, like FedAvg [2], FedProx [22], and FedNova [23], FedSplit fails to guarantee exact linear convergence to x∗. Empirically, we observe that FedSplit diverges on certain instances; see Appendix J. Compared to
these algorithms, we see from Theorem 1 that FedLin guarantees linear convergence to x∗. Notably, the linear convergence rate we obtain in Theorem 1 under both objective and systems heterogeneity is the best rate we know of in FL, and matches that of SCAFFOLD [11] where only objective heterogeneity is considered.3 The model of systems heterogeneity we study is taken from [23], where the authors provide guarantees only for the non-convex case under a bounded dissimilarity assumption. In contrast, our results cover all the three standard settings - strongly-convex, convex, and non-convex - without requiring any bounded dissimilarity assumption. For further related work on straggler-robust distributed learning algorithms (without objective heterogeneity or local steps), see [43–48].

heading: 4.1 The Price of Infrequent Communication

text:
In this section, we take a closer look at the effect of performing multiple local steps on the convergence rate. To do so, we assume that all clients perform the same number of local steps H , i.e., there is no communication for H consecutive time-steps between two communication rounds. Now consider a centralized baseline where each client can communicate with every other client at all times (i.e., even between rounds). In this case, since each client can always access∇f(x), gradient descent yields
f(x̄T+1)− f(x∗) ≤ exp(− 1
κ TH)(f(x̄1)− f(x∗)) (9)
after T rounds, with H synchronized local iterations within each round. Based on Theorem 1, observe that we lose out by a factor of H in the exponent relative to the centralized baseline. Notably, both in the centralized case, and in FedLin, each client queries the gradient of its local objective H times in each round, thereby making TH gradient queries over T rounds. Thus, relative to a centralized baseline, FedLin incurs the same computational cost in terms of gradient queries, and reduces communication by a factor of H , at the expense of a convergence rate that is slower by a factor of H . We emphasize here that just as with FedLin, H does not show up in the convergence rate (exponent) of algorithms like FedSplit [10] and SCAFFOLD [11] either.
The primary reason for the slower convergence rate (relative to a centralized baseline) stems from the need to set η ∝ 1/H to mitigate client-drift under objective heterogeneity. At this stage, one may conjecture that the above requirement is simply an artifact of a conservative analysis of Algorithm 1, and that a more refined analysis will reveal the utility of performing more local steps even in the heterogeneous setting. Our next result suggests otherwise; for a proof, see Appendix E. Theorem 5. (Lower bound for FedLin) Suppose δc = δs = 1, and τi = H, ηi = η,∀i ∈ S. Then, given anyL ≥ 14 andH ≥ 2, there exists an instance involving 2 clients where each fi(x), i ∈ {1, 2}, is 1-strongly convex and L-smooth, and an initial condition x̄1, such that FedLin initialized from x̄1 generates a sequence of iterates {x̄t} satisfying the following for any T ≥ 1: ‖x̄T+1 − x∗‖2 ≥ exp (−4T )‖x̄1 − x∗‖2; f(x̄T+1)− f(x∗) ≥ exp(−4T )(f(x̄1)− f(x∗)). (10)
Main Takeaways: There are several key implications of Theorem 5. First, it complements Theorem 1 by providing a matching lower bound. We believe ours is the first work to provide a tight linear convergence rate analysis: [11] and [10] only provide upper-bounds for SCAFFOLD and FedSplit, respectively. Second, our analysis of Theorem 5 in Appendix E indicates that there are problem instances where setting η ∝ 1/H is in fact necessary to guarantee convergence to x∗. As a result, for such problem instances, no matter how many local steps H each client performs, the error at the end of T rounds remains bounded below by an H-independent quantity, as is apparent from (10). Perhaps surprisingly, we show in Appendix E that the lower bound in Theorem 5 even applies to simple instances with non-identical quadratic losses (across clients) where every fi(x) has the same minimum! This is particularly insightful since it highlights the limitations of exploiting stale gradient terms in the local update rule (as is done in both FedLin and SCAFFOLD), and suggests the need for more informed updating schemes that explicitly take into account the level of statistical heterogeneity.
Proof Idea for Theorem 5: To establish Theorem 5, we set up an instance involving two clients with quadratic loss functions. Our main idea is to relate the convergence of FedLin to the Schur stability of an appropriate discrete-time linear time-invariant (LTI) system. Based on this connection, we show that guaranteeing stability necessitates setting η ∝ 1/H , which immediately leads to the lower bound. We believe that the same technique can be used to establish a similar lower bound for SCAFFOLD.
3In a concurrent work [42], the authors develop a linearly converging algorithm called S-Local-SVRG for the finite-sum setting, but neither consider systems heterogeneity nor compression. Moreover, unlike the lower bound we develop for FedLin in Theorem 5, no lower bounds are provided for S-Local-SVRG in [42].

heading: 5 Gradient Sparsification at Server

text:
In this section, our focus will be on addressing the following question: For strongly convex and smooth deterministic functions, and in the presence of both objective and systems heterogeneity, can we still hope for linear convergence to x∗ when gradients are sparsified at the server? Interestingly, we will show that not only is it possible to converge linearly to x∗, it is possible to do so without any error-feedback. Moreover, this claim holds regardless of how aggressive the server is in its sparsification scheme: it may even transmit just a single component of the aggregated gradient vector.
To isolate the impact of server-level sparsification, we will assume throughout this section that gradients are not sparsified at the clients, i.e., δc = 1. Consequently, hi,t+1 = ∇fi(x̄t+1),∀i ∈ S,∀t ∈ {1, . . . , T}. We begin by considering a simpler variant of FedLin with no error-feedback at the server side, i.e., line 15 is skipped, and gt+1 in line 14 of Algo. 1 is instead updated as follows
gt+1 = Cδs
( 1
m ∑ i∈S ∇fi(x̄t+1) ) = Cδs (∇f(x̄t+1)) . (11)
Theorem 6. (Sparsification at server with no error-feedback) Suppose each fi(x) is L-smooth and µ-strongly convex. Moreover, suppose τi ≥ 1,∀i ∈ S, and δc = 1. Consider a variant of FedLin, where line 14 is replaced by equation (11), and line 15 is skipped, i.e., there is no error-feedback. Then, with ηi = 12(2+ √ δs)Lτi ,∀i ∈ S, this variant of FedLin guarantees
f(x̄T+1)− f(x∗) ≤ ( 1− 1
2δs ( 2 + √ δs ) κ
)T (f(x̄1)− f(x∗)).
Main Takeaways: From Theorem 6, we see that even without error-feedback, it is possible to linearly converge to x∗; the rate of convergence, however, is inversely proportional to δ 3 2 s . Thus, Theorem 6 quantifies the trade-off between the level of sparsification at the server, and the rate of convergence. When there is no gradient compression, i.e., when δs = 1, we exactly recover Theorem 1.
One may ask: Is there any potential benefit to employing error-feedback when gradients are sparsified at the server? Our next result answers this question in the affirmative. Theorem 7. (Sparsification at server with error-feedback) Suppose each fi(x) is L-smooth and µ-strongly convex. Moreover, suppose τi ≥ 1,∀i ∈ S, and δc = 1. Let the step-size for client i be chosen as ηi = 172Lδsτi . Then, FedLin guarantees:
f(x̄T+1)− f(x∗) ≤ 2κ (
1− 1 96δsκ
)T (f(x̄1)− f(x∗)) .
For proofs of Theorems 6 and 7, see Appendix G and I.
Main Takeaways: Comparing the guarantee of Theorem 6 with that of Theorem 7, we note that the convergence rate is inversely proportional to δ 3 2 s in the former, and inversely proportional to δs in the latter. Thus, the main message here is that employing error-feedback leads to a faster convergence rate by improving the dependence of the rate on δs.

heading: 6 Gradient Sparsification at Clients

text:
In this section, we will turn our attention to the case when gradients are sparsified at the clients prior to being transmitted to the server. Throughout this section, we will assume that gradients are not compressed any further at the server side, i.e., δs = 1. To proceed, we will need to make the following bounded gradient dissimilarity assumption. Assumption 1. There exist constants C ≥ 1 and D ≥ 0 such that the following holds ∀x ∈ Rd:
1
m m∑ i=1 ‖∇fi(x)‖2 ≤ C‖∇f(x)‖2 +D. (12)
The following is the main result of this section; for a proof, see Appendix H.
Theorem 8. (Sparsification at clients with error-feedback) Suppose each fi(x) is L-smooth and µ-strongly convex, and suppose Assumption 1 holds. Moreover, suppose τi ≥ 1,∀i ∈ S , and δs = 1. Let the step-size for client i be chosen as ηi = η̄τi , where η̄ ∈ (0, 1) satisfies η̄ ≤ 1 72LδcC
. Then, FedLin guarantees:
‖x̄T+1 − x∗‖2 ≤ 2 (
1− 3 4 η̄µ
)T ‖x̄1 − x∗‖2 + 16
3 η̄
( 6
δcC + δc
) D
µ . (13)
Main Takeaways: Intuitively, one would expect that sparsifying gradients at each client prior to aggregation at the server would inject more errors than when gradients are first accurately aggregated at the server, and then the aggregated gradient vector is sparsified: Theorems 6 and 8 support this intuition. For the former, we neither required error-feedback nor Assumption 1 to guarantee linear convergence to the global minimum x∗; for the latter, even with error-feedback and the bounded gradient dissimilarity assumption, we can establish linear convergence to only a neighborhood of x∗, in general. From (13), we note that the size of this neighborhood scales linearly with D - a measure of objective heterogeneity. In particular, when D = 0, the iterates x̄t converge exactly to x∗. Remark 1. To the best of our knowledge, our results in Sections 5 and 6 constitute the first formal analysis of biased gradient sparsification in FL. In particular, we significantly generalize the recent results in [49] for a single worker to a multi-client FL setting with both objective and systems heterogeneity. To arrive at these results, we develop a new potential-function based proof technique in Appendix H. For more related work on compression in distributed learning, see Appendix A.
Extensions: We studied the effect of compressing information at the server and at the clients separately, with the goal of identifying the key differences between each of these mechanisms. The analysis techniques we developed in the process pave the way for studying various natural extensions: (i) combined sparsification at both the clients and the server; (ii) gradient sparsification in tandem with model parameter compression; and (iii) stochastic counterparts of Theorems 6, 7, and 8.

heading: 7 Experimental Results

text:
In this section, we provide numerical results for FedLin on a least squares problem to validate our theory. In Appendix K, we also provide additional numerical results on a logistic regression problem.
For now, we consider the following least squares regression problem:
min x∈Rd f(x) = min x∈Rd
1
m m∑ i=1 1 2 ‖Aix− bi‖2, (14)
where Ai ∈ R500×100 is a design matrix and bi ∈ R500 is a response vector. The client objective functions, fi(x) are strongly convex. Assuming that all design matrices are full column rank, problem (14) admits a unique minimizer. To generate synthetic data, for each client i ∈ S = {1, . . . , 20}, we generate Ai and bi according to the model bi = Aixi + εi, where xi is a weight vector and εi ∈ R500 is a disturbance. In particular, we generate [Ai]jk i.i.d.∼ N (0, 1), and εi ∼ N (0, 0.5I500), ∀i ∈ S. To capture statistical heterogeneity, the entries of the local true parameter of client i are modeled as [xi]k ∼ N (ui, 1), k ∈ {1, . . . , 100}, where ui ∼ N (0, α) and α ≥ 0. Hence, α controls the level of statistical heterogeneity. To model the effect of systems heterogeneity, for each client i ∈ S, the number of local steps is drawn uniformly and independently from [2, 100]. We will primarily focus on a deterministic setting here for our experiments; in Appendix L, we evaluate FedLin on a standard stochastic oracle model. Our experiments in Appendix L reveal that under a noisy oracle, FedLin guarantees linear convergence to a ball around the true minimum, exactly as suggested by Thm. 4.
Gradient Sparsification at Server. We first consider a variant of FedLin where gradient sparsification is implemented only at the server side and without any error-feedback. In particular, we consider the cases where δs ∈ {2, 4}, which correspond to the implementation of a TOP-50 and a TOP-25 operator, respectively. For comparison, we also plot the resulting performance when no gradient sparsification is implemented at the server. To examine the effect of statistical heterogeneity on the performance of FedLin, we generate two synthetic datasets corresponding to two different levels of heterogeneity in the clients’ local objectives, namely α = 10 and α = 50. As illustrated in Fig. 2, irrespective of the level of gradient sparsification on the server side, FedLin achieves linear convergence to the true minimum in the presence of both objective and systems heterogeneity, confirming Theorem 6. Also, both the convergence speed and accuracy of FedLin remain unaffected as the level of heterogeneity in the clients’ objective functions increases.
0 10 20 30 40
10 -20
10 -10
10 0
0 10 20 30 40
10 -20
10 -10
10 0
0 50 100 10
-2
10 -1
10 0
10 1
0 50 100 10
-1
10 0
10 1
Figure 3: Client-side sparsification results for FedLin. The constant η̄ is fixed at 5× 10−4. Left: α = 1. Right: α = 10.
0 100 200 300 400
10 -20
10 -10
10 0
0 100 200 300 400
10 -5
10 0
Figure 4: Comparison of FedLin with SCAFFOLD. (Left) Deterministic setting. (Right) General stochastic oracle model: unbiased gradients with variance σ = 10−1.
Comparison with SCAFFOLD. We now compare FedLin with SCAFFOLD on the least squares regression setup described above. To make a fair comparison, we assume that there is no systems heterogeneity or gradient compression. For implementing SCAFFOLD, we use Option II in their paper [11] for updating the control variates. We set the number of local steps H = 20, the statistical heterogeneity parameter α = 10, and use a step-size of 10−3 for both algorithms (the step-size was tuned to get best results). For the deterministic setting, we note from Fig. 4 that FedLin con-
verges much faster compared to SCAFFOLD. This trend persists when we perturb the gradients with zero-mean Gaussian noise with variance σ = 10−1. We conjecture that the faster convergence of FedLin stems from the fact that it uses less stale gradient correction terms relative to the control variates of SCAFFOLD; see the discussion about the fixed point property of FedLin in Sec. 3.

heading: 8 Conclusion

text:
We developed a novel algorithmic framework called FedLin to tackle some of the key challenges in FL, namely objective heterogeneity, systems heterogeneity, and imprecise communication. We showed that FedLin enjoys strong theoretical guarantees: (i) FedLin matches centralized rates, and, in particular, guarantees linear convergence to the global minimum under arbitrary objective and systems heterogeneity; and (ii) preserves linear convergence rates despite aggressive gradient sparsification. We also established a tight lower-bound for FedLin, highlighting that even mild statistical heterogeneity can end up hurting convergence rates - this is the first such result in FL. Our current approach requires two passes of communication between the clients and the server in each round. Moreover, our analysis does not account for partial client participation. As future work, we plan to address these limitations. We also plan to investigate other federated learning formulations (beyond supervised learning) where statistical heterogeneity can potentially help.
Acknowledgement: This work was supported by NSF Award 1837253, NSF CAREER award CIF 1943064, and the Air Force Office of Scientific Research Young Investigator Program (AFOSR-YIP) under award FA9550-20-1-0111.

title: Linear Convergence in Federated Learning: Tackling Client Heterogeneity and Sparse Gradients


summaries: This paper introduces a (possibly) cheaper approximation of Sliced Wasserstein (SW) distance. The proposed technique is based on calculating 1D Wasserstein distance between Gaussian projection directions. This is motivated by the fact that Wasserstein distance between Gaussian distributions admits a closed-form solution, which leads to an easy compute of the proposed approximate SW and discard a large number of uniformly sampling random projections for the Monte Carlo approximation of the original SW distance.

summaries: This paper develops a novel deterministic approximation for the Sliced-Wasserstein Distance (SWD), based on the fact that a random linear projection of a high-dimensional vector has an approximately Gaussian distribution. In many cases, the approximation involves no tuning parameters and can be computed much more quickly than the usual Monte Carlo estimate of SWD, especially for very high-dimensional data. If the $d$ dimensions of the data are not too strongly correlated, then the approximation converges no slower than $d^{-1/8}$ as $d \to \infty$. Experiments with synthetic data validate this theoretical result, as well as the computational benefits of the proposed approximation. Finally, the paper demonstrates an application to image generation.

summaries: The paper presents a method to estimate the Sliced-Wasserstein distance. The method is based on the Gaussian approximation of the projected data, thus the SW distance could be approximated through the SW distance between Gaussians which has a closed-form. Hence, the method is better than the Monte Carlo SW distance in term of computational speed. The authors also demonstrate their method on toy data set and real data sets, i.e. MNIST and Celeb A.

summaries: The paper proposes a new approach to approximate the sliced Wasserstein distance. In particular, the paper extends the central limit theorems for random projections to derive a deterministic approximation of sliced Wasserstein (based on the closed form of Wasserstein distance between two Gaussians) that is efficient (low error) in high dimension under centering assumption (weak dependence assumption). On the experimental side, the paper illustrates the derived theory on the synthetic data and compares the proposed approach to the conventional Monte Carlo approximation of the sliced Wasserstein distance.

abstractText: The Sliced-Wasserstein distance (SW) is being increasingly used in machine learning applications as an alternative to the Wasserstein distance and offers significant computational and statistical benefits. Since it is defined as an expectation over random projections, SW is commonly approximated by Monte Carlo. We adopt a new perspective to approximate SW by making use of the concentration of measure phenomenon: under mild assumptions, one-dimensional projections of a highdimensional random vector are approximately Gaussian. Based on this observation, we develop a simple deterministic approximation for SW. Our method does not require sampling a number of random projections, and is therefore both accurate and easy to use compared to the usual Monte Carlo approximation. We derive nonasymptotical guarantees for our approach, and show that the approximation error goes to zero as the dimension increases, under a weak dependence condition on the data distribution. We validate our theoretical findings on synthetic datasets, and illustrate the proposed approximation on a generative modeling problem.

authors:
Julien Rabin
Gabriel Peyré
Julie Delon
Marc Bernot

title: Wasserstein Barycenter and Its Application to Texture Mixing

venue: Scale Space and Variational Methods in Computer Vision,

authors:
Nicolas Bonneel
Julien Rabin
Gabriel Peyré
Hanspeter Pfister

title: Sliced and Radon Wasserstein barycenters of measures

venue: Journal of Mathematical Imaging and Vision,

authors:
Soheil Kolouri
Yang Zou
Gustavo K Rohde

title: Sliced-Wasserstein Kernels for Probability distributions

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Mathieu Carrière
Marco Cuturi
Steve Oudot

title: Sliced Wasserstein Kernel for Persistence Diagrams

venue: Proceedings of the 34th International Conference on Machine Learning,

authors:
Kimia Nadjahi
Valentin Bortoli
Alain Durmus
Roland Badeau
Umut Simsekli

title: Approximate Bayesian Computation with the Sliced-Wasserstein

venue: distance. pages 5470–5474,

authors:
Soheil Kolouri
Gustavo K Rohde
Heiko Hoffmann

title: Sliced Wasserstein distance for Learning Gaussian Mixture Models

authors:
Soheil Kolouri
Phillip E. Pope
Charles E. Martin
Gustavo K. Rohde

title: Sliced Wasserstein Auto- Encoders

venue: In International Conference on Learning Representations,

authors:
Ishan Deshpande
Ziyu Zhang
Alexander Schwing

title: Generative Modeling using the Sliced Wasserstein Distance

authors:
Jiqing Wu
Zhiwu Huang
Dinesh Acharya
Wen Li
Janine Thoma
Danda Pani Paudel
Luc Van Gool

title: Sliced Wasserstein Generative Models

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Antoine Liutkus
Umut Simsekli
Szymon Majewski
Alain Durmus
Fabian-Robert Stöter

title: Sliced- Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions

venue: Proceedings of the 36th International Conference on Machine Learning,

title: Unidimensional and Evolution Methods for Optimal Transportation

venue: PhD thesis, Paris

authors:
Kimia Nadjahi
Alain Durmus
Umut Simsekli
Roland Badeau

title: Asymptotic Guarantees for Learning Generative Models with the Sliced-Wasserstein distance

venue: Advances in Neural Information Processing Systems

authors:
Erhan Bayraktar
Gaoyue Guo

title: Strong equivalence between metrics of Wasserstein type, 2019

title: The speed of mean Glivenko-Cantelli convergence

venue: Ann. Math. Statist., 40(1):40–50,

authors:
Nicolas Fournier
Arnaud Guillin

title: On the rate of convergence in Wasserstein distance of the empirical measure

venue: Probability Theory and Related Fields,

authors:
Jonathan Weed
Francis Bach

title: Sharp asymptotic and finite-sample rates of convergence of empirical measures in Wasserstein distance

authors:
Kimia Nadjahi
Alain Durmus
Lénaïc Chizat
Soheil Kolouri
Shahin Shahrampour
Umut Simsekli

title: Statistical and Topological Properties of Sliced Probability Divergences

venue: Advances in Neural Information Processing Systems,

authors:
Cheng Meng
Yuan Ke
Jingyi Zhang
Mengrui Zhang
Wenxuan Zhong
Ping Ma

title: Large-scale optimal transport map estimation using projection pursuit

authors:
Ishan Deshpande
Yuan-Ting Hu
Ruoyu Sun
Ayis Pyrros
Nasir Siddiqui
Sanmi Koyejo
Zhizhen Zhao
David Forsyth
Alexander G Schwing

title: Max-Sliced Wasserstein distance and its use for GANs

authors:
Soheil Kolouri
Kimia Nadjahi
Umut Simsekli
Roland Badeau
Gustavo Rohde

title: Generalized Sliced Wasserstein Distances

authors:
Khai Nguyen
Nhat Ho
Tung Pham
Hung Bui

title: Distributional Sliced-Wasserstein and Applications to Generative Modeling

title: Typical distributions of linear functionals in finite dimensional spaces of high dimension

venue: Soviet Math. Dokl.,

authors:
Persi Diaconis
David Freedman

title: Asymptotics of Graphical Projection Pursuit

venue: The Annals of Statistics,

title: Conditional central limit theorems for Gaussian projections

venue: IEEE International Symposium on Information Theory (ISIT),

authors:
Sebastian Goldt
Bruno Loureiro
Galen Reeves
Florent Krzakala
Marc Mézard
Lenka Zdeborová

title: The Gaussian equivalence of generative models for learning with shallow neural networks

venue: Proceedings of Machine Learning Research,

authors:
D.C Dowson
B.V Landau

title: The Fréchet distance between multivariate normal distributions

venue: Journal of Multivariate Analysis,

authors:
Svetlozar T Rachev
Ludger Rüschendorf

title: Mass Transportation Problems: Volume I: Theory, volume 1

authors:
Gabriel Peyré
Marco Cuturi

title: Computational Optimal Transport: With Applications to Data Science

venue: Foundations and Trends R © in Machine Learning,

authors:
Peter Hall
Ker-Chau Li

title: On almost Linearity of Low Dimensional Projections from High Dimensional Data

title: Sudakov’s typical marginals, random linear functionals and a conditional central limit theorem

authors:
Milla Anttila
Keith Ball
Irini Perissinaki

title: The Central Limit Problem for Convex Bodies

venue: Transactions of the American Mathematical Society,

title: On concentration of distributions of random weighted sums

venue: The Annals of Probability,

title: Approximation of Projections of Random Vectors

venue: Journal of Theoretical Probability,

authors:
Lutz Dümbgen
Perla Del Conte-Zerial

title: On low-dimensional projections of high-dimensional distributions

venue: Institute of Mathematical Statistics Collections,

title: On the conditional distributions of low-dimensional projections from high-dimensional data

authors:
Paul Doukhan
Michael H. Neumann

title: Probability and moment inequalities for sums of weakly dependent random variables, with applications

venue: Stochastic Processes and their Applications,

authors:
Paul Doukhan
Sana Louhichi

title: A new weak dependence condition and applications to moment inequalities

title: The notion of Ψ-weak dependence and its applications to bootstrapping time series

venue: Probability Surveys,

authors:
Ziwei Liu
Ping Luo
Xiaogang Wang
Xiaoou Tang

title: Deep Learning Face Attributes in the Wild

venue: In Proceedings of International Conference on Computer Vision (ICCV),

heading: 1 Introduction

text:
Recent years have witnessed the emergence of numerical methods inspired by optimal transport (OT) to solve machine learning problems. In particular, Wasserstein distances are a core ingredient of OT and define metrics between probability measures. Despite their nice theoretical properties, they are in general computationally expensive in large-scale settings. Several workarounds that scale better to large problems have been developed, and include the Sliced-Wasserstein distance (SW, [1, 2]).
The SW metric is a computationally cheaper alternative to Wasserstein as it exploits the analytical form of the Wasserstein distance between univariate distributions. More precisely, consider two random variables X and Y in Rd with respective distributions µ and ν, and denote by θ?]µ, θ?] ν the univariate distributions of the projections of X,Y along θ ∈ Rd. SW then compares µ and ν by computing E[W(θ?]µ, θ?] ν)], where the expectation E is taken with respect to θ uniformly distributed on the unit sphere, and W is the Wasserstein distance.
In practice, the expectation is typically estimated by Monte Carlo: one uniformly draws L projection directions {θl}Ll=1, and approximates SW with L−1 ∑L i=1 W ( (θl) ? ]µ, (θl) ? ]ν ) . Since the Wasserstein
∗Corresponding author: kimia.nadjahi@telecom-paris.fr
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
distance between univariate distributions can easily be computed in closed-form, this scheme leads to significant computational benefits as compared to the Wasserstein distance, provided that L is not chosen too large. SW has then been successfully applied in several practical tasks, such as classification [3, 4], Bayesian inference [5], the computation of barycenters of measures [1, 6], and implicit generative modeling [7–12]. Besides, SW has been shown to offer nice theoretical properties as well. Indeed, it satisfies the metric axioms [13], the estimators obtained by minimizing SW are asymptotically consistent [14], the convergence in SW is equivalent to the convergence in Wasserstein [14, 15], and even though the sample complexity of Wasserstein grows exponentially with the data dimension [16–18], the sample complexity of SW does not depend on the dimension [19]. However, the latter study also demonstrated with a theoretical error bound, that the quality of the Monte Carlo estimate of SW depends on the number of projections and the variance of the one-dimensional Wasserstein distances [19, Theorem 6]. In other words, to ensure that the induced approximation error is reasonably small, one might need to choose a large value for L, which inevitably increases the computational complexity of SW. Alternative approaches have been proposed to overcome this issue, and mainly consist in picking more “informative” projection directions: e.g., SW based on orthogonal projections [10, 20], maximum SW [21], generalized SW distances [22] and distributional SW distances [23].
In this paper, we adopt a different perspective and leverage concentration results on random projections to approximate SW: previous work showed that, under relatively mild conditions, the typical distribution of low-dimensional projections of high-dimensional random variables is close to some Gaussian law [24, 25]. Recently, this phenomenon has been illustrated with a bound in terms of the Wasserstein distance [26]: let {Xi}di=1 be a sequence of real random variables with distribution µd, such that X1, . . . , Xd are independent with finite fourth-order moments; then, E[W(θ?]µd,Nµd)2] goes to zero as d increases, where Nµd is a univariate Gaussian distribution whose variance depends on µd and the expectation is taken with respect to a Gaussian variable θ. This result has very recently been used to bound the “maximum-sliced distance” between any probability measure and its Gaussian approximation [27]. In our work, we use it to design a novel technique that estimates SW with a simple deterministic formula. As opposed to Monte Carlo, our method does not depend on a finite set of random projections, therefore it eliminates the need of tuning the hyperparameter L and can lead to a significant computational time reduction. Besides, our proposal is quite different from the aforementioned variants of SW which consist in selecting “informative” projection directions: these alternatives are defined as optimization problems whose resolution is challenging (e.g., [23, Section 3.2]) and are then computed by finding an approximate solution. This incurs an additional computational cost and estimation error, while our method directly approximates SW (thus, does not define an alternative distance) via simple deterministic operations, does not rely on any hyperparameters, and comes with theoretical guarantees on its induced error.
The important steps to formulate our approximate SW are summarized as follows. We first define an alternative SW whose projection directions are drawn from the same Gaussian distribution as in [26], instead of uniformly on the unit sphere, and establish its relation with the original SW. By combining this property with [26, Theorem 1], we bound the absolute difference between SW applied to any two probability measures µd, νd on Rd, and the Wasserstein distance between the univariate Gaussians Nµd , Nνd . Then, we explain why the mean parameters of µd and νd should be zero for the approximation error to decrease as d grows. Nevertheless, we show that it is not a limiting factor, by exploiting the following decomposition of SW: SW between µd, νd can be equivalently written as the sum of the difference between their means and the SW between the centered versions of µd, νd.
Our approach then consists in estimating SW between the centered versions with the Wasserstein term between Gaussian approximations to meet the zero-means condition, and recover SW between the original measures via the aforementioned property. Since the Wasserstein distance between Gaussian distributions admits a closed-form solution, our approximate SW is very easy to compute, and faster than the Monte Carlo estimate obtained with a large number of projections. We derive nonasymptotical guarantees on the error induced by our approach. Specifically, we define a weak dependence condition under which the error is shown to go to zero with increasing d. Our theoretical results are then validated with experiments conducted on synthetic data. Finally, we leverage our theoretical insights to design a novel adversarial framework for a typical generative modeling problem in machine learning, and illustrate its advantages in terms of accuracy and computational time, over
generative models based on the Monte Carlo estimate of SW. Our empirical results can be reproduced with our open source code2.

heading: 2 Background

text: We first give some background on optimal transport distances and concentration of measure for random projections. All random variables are defined on a probability space (Ω,F ,P) with associated expectation operator E. We denote N∗ = N \ {0}, and for d ∈ N∗, P(Rd) is the set of probability measures on Rd.

heading: 2.1 Optimal transport distances

text:
Let p ∈ [1,+∞) and Pp(Rd) = { µ ∈ P(Rd) : ∫ Rd ‖x‖ p dµ(x) < +∞ } be the set of probability measures on Rd with finite moment of order p. The Wasserstein distance of order p between any µ, ν ∈ Pp(Rd) is defined as
Wpp(µ, ν) = inf π∈Π(µ,ν) ∫ Rd×Rd ‖x− y‖p dπ(x, y) , (1)
where ‖ · ‖ denotes the Euclidean norm, and Π(µ, ν) the set of probability measures on Rd × Rd whose marginals with respect to the first and second variables are given by µ and ν respectively. In some particular settings, Wp is relatively easy to compute since the optimization problem in (1) admits a closed-form solution: we give two examples that will be useful in the rest of the paper.
Gaussian distributions. Denote by N(m,Σ) the Gaussian distribution on Rd with mean m ∈ Rd and covariance matrix Σ ∈ Rd×d symmetric positive-definite. The Wasserstein distance between two Gaussian distributions, also known as the Wasserstein-Bures metric, is given by [28]
W22{N(m1,Σ1),N(m2,Σ2)} = ‖m1 −m2‖2 + Tr [ Σ1 + Σ2 − 2 ( Σ 1/2 1 Σ2Σ 1/2 1 )1/2] , (2)
where Tr is the trace operator.
Univariate distributions. Consider µ, ν ∈ Pp(R), and denote by F−1µ and F−1ν the quantile functions of µ and ν respectively. By [29, Theorem 3.1.2.(a)],
Wpp(µ, ν) = ∫ 1 0 ∣∣F−1µ (t)− F−1ν (t)∣∣p dt . (3) If µ = n−1 ∑n i=1 δxi and ν = n −1∑n i=1 δyi , with {xi}ni=1, {yi}ni=1 ⊂ Rn and δz the Dirac distribution with mass on z, (3) can simply be calculated by sorting {xi}ni=1 and {yi}ni=1 as x(1) ≤ . . . ≤ x(n) and y(1) ≤ . . . ≤ y(n). Indeed, in this case, Wpp(µ, ν) = n−1 ∑n i=1 |x(i) − y(i)|p. However, when the empirical distributions are multivariate, Wp(µ, ν) is not analytically available in general, so its computation is expensive: the standard methods used to solve the linear program in (1) have a worst-case computational complexity in O(n3 log n), and tend to have a super-cubic cost in practice [30, Chapter 3].
The Sliced-Wasserstein distance [1, 2] defines a practical alternative metric by leveraging the computational efficiency of Wp for univariate distributions. Let Sd−1 be the d-dimensional unit sphere and σ the uniform distribution on Sd−1. For θ ∈ Sd−1, θ? : Rd → R denotes the linear form x 7→ 〈θ, x〉 with 〈·, ·〉 the Euclidean inner-product. Then, SW of order p ∈ [1,∞) between µ, ν ∈ Pp(Rd) is
SWpp(µ, ν) = ∫ Sd−1 Wpp(θ ? ]µ, θ ? ] ν)dσ(θ) , (4)
where for any measurable function f : Rd → R and ξ ∈ P(Rd), f]ξ is the push-forward measure of ξ by f : for any Borel set A in R, f]ζ(A) = ζ(f−1(A)), with f−1(A) = {x ∈ Rd : f(x) ∈ A}. Since θ?]µ, θ ? ] ν are univariate distributions, the Wasserstein distances in (4) are conveniently computed using (3). Besides, in practical applications, the expected value in (4) is typically approximated with a standard Monte Carlo method:
2See https://github.com/kimiandj/fast_sw
SWpp,L(µ, ν) = (1/L) L∑ l=1 Wpp ( (θl) ? ]µ, (θl) ? ]ν ) , with {θl}Ll=1 i.i.d. from σ . (5)
Computing SWp,L between two empirical distributions then amounts to projecting sets of n observations in Rd along L directions, and sorting the projected data. The resulting computational complexity is O(Ldn + Ln log n), which is more efficient than Wp in general. This complexity means that the Monte Carlo estimate is more expensive when d, n and L increase, and it is often unclear how L should be chosen in order to control the approximation error; see [19, Theorem 6].

heading: 2.2 Central limit theorems for random projections

text:
There is a rich literature on the typical behavior of one-dimensional random projections of highdimensional vectors. To be more specific, let (θi)i∈N∗ be i.i.d. standard one-dimensional Gaussian random variables and (Xi)i∈N∗ be a sequence of one-dimensional random variables. Denote for any d ∈ N∗, θ1:d = {θi}di=1 and X1:d = {Xi}di=1. Several central limits theorems ensure that, under relatively mild conditions, the sequence of distributions of d−1/2 〈θ1:d, X1:d〉 ∈ R given θ1:d ∈ Rd converges in distribution to a Gaussian random variable in probability. This line of work goes back to [24, 25], whose contributions have then been sharpened and generalized in [31–38]. In particular, a recent study [26] gives a quantitative version of this phenomenon. More precisely, denote for any d ∈ N∗ by µXd , the distribution of X1:d (i.e., the joint distribution of X1, X2, . . . , Xd) and γd the zero-mean Gaussian distribution with covariance matrix (1/d)Id. Assume that for any d ∈ N∗, µXd ∈ P2(Rd). Then, [26, Theorem 1] shows that there exists a universal constant C ≥ 0 such that∫
Rd W22
( θ?]µ X d ,N ( 0, d−1m2(µ X d ) )) dγd(θ) ≤ CΞd(µXd ) , with (6)
Ξd(µ X d ) = d −1{α(µXd ) + ( m2(µ X d )β1(µ X d ) )1/2 + m2(µ X d ) 1/5β2(µ X d ) 4/5} , (7)
m2(µ X d ) = E [ ‖X1:d‖2 ] , α(µXd ) = E [∣∣∣‖X1:d‖2 − m2(µXd )∣∣∣] , βq(µXd ) = E 1q [|〈X1:d, X ′1:d〉|q] , (8)
where q ∈ {1, 2} and (X ′i)i∈N∗ is an independent copy of (Xi)i∈N∗ . A formal statement of this result is also given for completeness in the supplement.

heading: 3 Approximate Sliced-Wasserstein distance based on concentration of random projections

text: We develop a novel method to approximate the Sliced-Wasserstein distance of order 2, by extending the bound in (6) and deriving novel properties for SW. We then derive nonasymptotical guarantees of the corresponding approximation error, which ensure that our estimate is accurate for highdimensional data under a weak dependence condition.

heading: 3.1 Sliced-Wasserstein distance with Gaussian projections

text:
First, to enable the use of (6) for the analysis of SW, we introduce a variant of SWp (4) whose projections are drawn from the Gaussian distribution considered in (6), instead of uniformly on the sphere. The Sliced-Wasserstein distance of order p ∈ [1,+∞) based on Gaussian projections is defined for any µ, ν ∈ Pp(Rd) as
S̃W p
p(µ, ν) = ∫ Rd Wpp(θ ? ]µ, θ ? ] ν)dγd(θ) . (9)
In the next proposition, we establish a simple mathematical relation between traditional SW and the newly introduced one: we prove that S̃Wp is equal to SWp up to a proportionality constant that only depends on the data dimension d and the order p.
Proposition 1. Let p ∈ [1,+∞). Then, S̃Wp (9) is related to SWp (4) as follows: for any µ, ν ∈ Pp(Rd), S̃Wp(µ, ν) = (2/d)1/2 { Γ(d/2 + p/2) / Γ(d/2) }1/p SWp(µ, ν), where Γ is the Gamma function.
Since (6) only applies to the Wasserstein distance of order 2, we will focus on SW of that same order in the rest of the paper. In this case, SW with Gaussian projections is equal to the original SW. Indeed, we can show that the constant (2/d)1/2 {Γ(d/2 + p/2) / Γ(d/2)}1/p defined in Proposition 1 is equal to 1 when p = 2, by using the property Γ(d/2 + 1) = (d/2)Γ(d/2).

heading: 3.2 Approximate Sliced-Wasserstein distance

text:
Our next result is an easy consequence of (6) and Proposition 1, and shows that the absolute difference between SW2(µd, νd) and W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} for any µd, νd ∈ P2(Rd), is bounded from above by Ξd(µd) + Ξd(νd) (7).
Theorem 1. There exists a universal constant C > 0 such that for any µd, νd ∈ P2(Rd),∣∣SW2(µd, νd)−W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))}∣∣ ≤ C(Ξd(µd) + Ξd(νd))1/2 , (10) where, for ξd ∈ {µd, νd}, Ξd(ξd) and m2(ξd) are defined in (7) and (8) respectively.
Since W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} has a closed-form solution by (2), it provides a computationally efficient approximation of SW2(µd, νd) whose accuracy is quantified by Theorem 1. Next, we identify settings where this approximation is accurate, by analyzing the error Ξd(µd) + Ξd(νd).
Our first observation is that µd and νd should have zero means for the error to go to zero as d→ +∞, and we develop a novel approximation of SW that takes into account this constraint. Going back to the definition of Ξd(µXd ) in (7), setting X̄i = Xi − E[Xi] and X̄ ′i = X ′i − E[X ′i], we get
m2(µ X d ) = E[‖X̄1:d‖2] + ‖E[X1:d]‖ 2 (11)
β22(µ X d ) = E [〈 X̄1:d, X̄ ′ 1:d 〉2] + 4E [〈 E[X1:d], X̄1:d 〉2] + ‖E[X1:d]‖4 . (12)
By Equations (11) and (12), since in practice the norm of the mean E[X1:d] is expected to increase linearly with d1/2 at least, so are m2(µXd ) and β2(µ X d ) as functions of d. As a consequence, Ξd(µ X d ) cannot be shown to converge to 0 as d → ∞ in this setting, but only to be bounded. However, if the data are centered, the norm of the mean is zero, thus Ξd(µXd ) might be decreasing. Therefore, we derive a convenient formula to compute SW2(µd, νd) from SW2(µ̄d, ν̄d) where for any ξd ∈ P2(Rd), ξ̄d is the centered version of ξ, i.e. the pushforward measure of ξd by x 7→ x−mξd with mξd = ∫ Rd y dξd(y). This result is the last ingredient to formulate our approximation of SW.
Proposition 2. Let µd, νd ∈ P2(Rd) with respective means mµd ,mνd . Then, the Sliced-Wasserstein distance of order 2 can be decomposed as
SW22(µd, νd) = SW 2 2(µ̄d, ν̄d) + (1/d)‖mµd −mνd‖2 . (13)
Based on Proposition 2, instead of estimating SW2(µd, νd) with W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} directly, we propose approximating SW2(µ̄d, ν̄d) with W2{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))} and then using (13). This strategy yields our final approximation of SW, which is defined for any µd, νd ∈ P2(Rd) as
ŜW 2 2(µd, νd) = W 2 2{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))}+ (1/d)‖mµd −mνd‖2 , (14)
where for ξd ∈ {µ̄d, ν̄d}, m2(ξd) is defined in (8). Note that (14) can be simplified since by (2), W22{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))} = d−1(m2(µ̄d)1/2 − m2(ν̄d)1/2)2. Besides, if µd and νd are both supported on a finite set of points, ŜW2(µd, νd) has a closed-form expression: given ξd = n−1 ∑n j=1 δx(j) ∈ P2(Rd) with x(j) ∈ Rd for j ∈ {1, . . . , n}, we then have mξd = n−1 ∑n j=1 x
(j), and m2(ξd) = n−1 ∑n j=1 ‖x(j)‖2. The associated computational complexity is therefore in O(dn).
Hence, we introduced an alternative technique to estimate SW which does not rely on a finite set of random projections, as opposed to the commonly used Monte Carlo technique (5). Our approach thus eliminates the need for practitioners to tune the number of projections L, but also to sort the projected data. As a consequence, it is more efficient to compute ŜW2(µd, νd) than SW2,L(µd, νd) for large L. We illustrate this latter point with empirical results in Section 4.

heading: 3.3 Error analysis under weak dependence

text:
We have discussed why centering the data is important to ensure that the approximation error goes to zero with increasing d. Next, we introduce a weak dependence condition under which the error is guaranteed to decrease as d increases.
We first consider a setting mentioned in [26] where µd = µ(1)⊗· · ·⊗µ(d) and νd = ν(1)⊗· · ·⊗ν(d), ⊗ denoting the tensor product of measures, and µ(j), ν(j) ∈ P4(R) for j ∈ {1, . . . , d}. We prove in this case that W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} converges to SW2(µd, νd) at a rate of d−1/8. This result is reported in the supplementary document, and can be interpreted as an extension of [26, Corollary 3] for SW.
We emphasize that the assumptions of this first setting severely restrict the scope of application of our approximation method: in several statistical and machine learning tasks, the random variables of interest {Xi}di=1 are not independent from each other (e.g. for image data, each Xi typically represents the value of a pixel at a certain position, thus depends on the neighboring pixels). Therefore, we relax this independence condition by considering a concept of ‘weak dependence’ inspired by [39] and properly defined in Definition 1. Definition 1. Let (Xj)j∈N∗ be a stationary sequence of one-dimensional random variables with mean zero, i.e. Xi and Xj have the same distribution for any i, j ∈ N∗ and E[X1] = 0. We say that (Xj)j∈N∗ is fourth-order weakly dependent if there exist some constant K ≥ 0 and a nonincreasing sequence of real coefficients {ρ(n)}n∈N such that, for any i, j ∈ N∗, i ≤ j,
|Cov(X2i , X2j )| ≤ Kρ(j − i) , |Cov(Xi, Xj)| ≤ Kρ(j − i) . (15) In addition, the sequence {ρ(n)}n∈N satisfies ∑+∞ n=0 ρ(n) ≤ ρ∞ < +∞.
Intuitively, in practical applications, the weak dependence condition would essentially require the components of the observations not to exhibit strong correlations; yet, they are allowed to depend on each other. Furthermore, since our weak dependence condition is weaker than the one introduced in [39, Theorem 1], it is satisfied by the various examples of models described in [39, Section 5]. We present some of them below, to illustrate Definition 1 more clearly. 1) Gaussian processes and associated processes [40, Section 3.1], provided that they are stationary. 2) Bernoulli shifts: Xt = H(εt, . . . , εt−r) for t ∈ N∗, where H : Rr+1 → R is a measurable
function and (εi)i∈N∗ is a sequence of i.i.d. real random variables. A simple example of such process is given by moving-average models.
3) Autoregressive models, defined as Xt = f(Xt−1, . . . , Xt−r) + εt for t ∈ N∗, where (εi)i∈N∗ a sequence of i.i.d. real random variables with E |ε1| <∞, and |f(u1, . . . , ur)− f(v1, . . . , vr)| ≤∑r i=1 ai |ui − vi| for some a1, . . . , ar ≥ 0 such that (∑r i=1 ai )1/r < 1.
We then consider a sequence of fourth-order weakly dependent random variables (Xj)j∈N∗ , and prove that Ξd(µXd ) goes to zero as d → ∞, with a rate of convergence depending on {ρ(n)}n∈N. This result is given in the supplementary document, and helps us refine Theorem 1 under this weak dependence condition: the next corollary establishes that the error approaches 0 at a rate of d−1/8. Corollary 1. Let (Xj)j∈N∗ and (Yj)j∈N∗ be sequences of random variables which are fourth-order weakly dependent. Set for any d ∈ N∗, X1:d = {Xj}dj=1 and Y1:d = {Yj}dj=1, and denote by µd, νd the distributions of X1:d, Y1:d respectively. Then, there exists a universal constant C > 0 such that∣∣SW2(µd, νd)−W2(N(0, d−1m2(µd)),N(0, d−1m2(νd)))∣∣ ≤ Cd−1/8. Hence, by replacing the independence condition of the first setting with weak dependence, we broaden the scope of application whilst guaranteeing that the approximation error goes to zero as d increases. We finally note that in these two settings, the data are required to have zero mean, which is automatically verified with our approximation method since we estimate SW between the centered distributions (see eq. (14)).

heading: 4 Experiments

text:
Synthetic experiments. The goal of these experiments is to illustrate our theoretical results derived in Section 3. In each setting, we generate two sets of d-dimensional samples, denoted by {x(j)}nj=1
and {y(j)}nj=1 with n = 104 and x(j), y(j) ∈ Rd for j ∈ {1, . . . , n}. We then approximate SW between their empirical distributions in P2(Rd), given by µd = n−1 ∑n j=1 δx(j) and νd =
n−1 ∑n j=1 δy(j) .
First, we analyze the consequences of centering data. Here, {x(j)}nj=1, {y(j)}nj=1 are n independent samples from Gaussian or Gamma distributions: see the supplementary document for more details. We compute |W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} − SW2(µd, νd)| on the one hand, and |W2{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))} − SW2(µ̄d, ν̄d)| on the other hand. In the Gaussian case, the exact value of SW2 is known (we report it in the supplementary document), while for the Gamma distributions, it is approximated with Monte Carlo based on 2 × 104 random projections. Figures 1(a) and 1(b) show that the error goes to zero as d increases if the data are centered. This confirms our analysis provided in Section 3.2 about the influence of the mean, and in Section 3.3 on sequences of independent random variables.
Next, we consider autoregressive processes of order one (AR(1)). An AR(1) process is defined as X1 = ε1 and, for t ∈ N∗, Xt = αXt−1 + εt, where α ∈ [0, 1] and (εi)i∈N∗ is an i.i.d. sequence of real random variables with E[ε1] = 0 and finite second-order moment. If α < 1, the process has a stationary distribution and (Xj)j∈N∗ satisfies the weak dependence condition in its stationary regime [41]. In practice, we generate a sample by using this recursion formula for 104 + d steps, and keeping the last d samples. The discarded samples correspond to a “burn-in” phase which helps reaching the stationary solution of the process. We generate {x(j)}nj=1 and {y(j)}nj=1 using the same distribution for the noise (either a Gaussian or Student’s t-distribution, as described in the supplementary document). This means that both datasets come from the same distribution, thus the exact value of SW2 is zero. We plot on Figures 1(c) and 1(d) the approximation error according to d ∈ [10, 103] for different values of α. The error converges to zero with increasing d, which is consistent with Corollary 1.
Note that Figure 1 exhibits rate of convergence that are better than the one in d−1/8 derived in Section 3.3: in Figure 1(b), the slope is approximately −0.45 (Gaussian) and −0.7 (Gamma), and in Figures 1(c) and 1(d), it is on average −0.35. This suggests that our theoretical bounds might be improved, and we further investigate this aspect for the Gaussian case: we consider the case where {x(j)}nj=1, {y(j)}nj=1 are n independent samples from Gaussian distributions with diagonal
covariance matrices, and we prove that E|W2{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))}−SW2(µ̄d, ν̄d)| goes to 0 as dn→ +∞ with a convergence rate in d−1/2n−1/2. We provide the complete statement and formal proof in Section S3.1 (Proposition S3). This result is consistent with Figure 1(b), and is a first encouraging step towards the following research direction: we will study if our proofs and the ones in [26] can be refined when assuming additional structure on the distributions (e.g., sub-Gaussian and sub-exponential), in order to identify the settings under which our current bounds are tight or can be improved.
Finally, we compare our approximation scheme against the standard Monte Carlo estimation, in terms of accuracy and computation time. We use the same setting as in Figure 1(b), where the n samples are independently drawn from Gamma distributions. We compute ŜW2(µd, νd) (14) and SW2,L(µd, νd) (5) with L ∈ {100, 1000, 5000}, and we compare each approximation with SW2,2×104(µd, νd), which we consider as the exact value of SW. Figure 2 reports the approximation error and computation time of each scheme for d ∈ [10, 104], and shows that our method is more accurate and faster than Monte Carlo. In particular, when d = 103, the average computation time of our technique is 0.02s, while the second best approximation (Monte Carlo with L = 5000) takes more than 8s. Besides, we observe that Monte Carlo is very sensitive to the hyperparameters, since it loses accuracy when L decreases and gets slower as L and d increase. This observation is consistent with the computational complexity of SW2,L recalled in Section 2.1. On the other hand, our approximation scheme is extremely efficient even for large d and n, since it is based on a simple deterministic formula which does not require projecting and sorting data along random directions.
Image generation. Finally, we leverage our theoretical insights to design a novel method for a typical generative modeling application. The problem consists in tuning a neural network that takes as input k-dimensional samples from a reference distribution (e.g., uniform or Gaussian), to generate images of dimension d > k. During the training phase, the parameters of the network are updated by iteratively minimizing a dissimilarity measure between the dataset to fit and the generated images.
In [9], the dissimilarity measure is Monte Carlo SW approximated with 104 random projections, and the resulting generative model is called the Sliced-Wasserstein generator (SWG). This model performs well on moderately high-dimensional image datasets (e.g., 28 × 28 for MNIST images [42]). However, for very large dimensions (e.g., 64 × 64 × 3 for the CelebA dataset [43]), Monte Carlo SW requires more than 104 random projections to capture relevant information, which leads to very expensive training iterations and potential memory issues. To offer better scalability, SWG can be augmented with a discriminator network [9, Section 3.2] that aims at finding a lower-dimensional space in which the two projected datasets are clearly distinguishable. The intuition behind this heuristic is that the more distinct the two datasets are from each other, the fewer projection directions Monte Carlo SW requires to provide useful information. The training then consists in optimizing the generator’s and discriminator’s objective functions in an alternating fashion.
Our novel approach builds on SWG and modifies the saddle-point problem in [9, Section 3.2]: motivated by the gain in accuracy and time illustrated in Figure 2 on high-dimensional datasets, we propose to replace Monte Carlo SW with our approximate SW (14) in the generator’s objective; then, to make sure that our approximation is accurate, we regularize the discriminator’s objective:
max ψ
L(ψ) + λ1 ∥∥Cov[d′ψ(X)]∥∥2F + λ1∥∥Cov[d′ψ(gφ(Z))]∥∥2F (16)
+ λ2 E [ ‖d′ψ(X)‖−2 ] + λ2 E [ ‖d′ψ(gφ(Z))‖−2 ] (17)
where L is the discriminator’s loss used in SWG, gφ and d′ψ are the generator’s last layer and the discriminator’s penultimate layer respectively (parameterized by φ, ψ), X and Z are the random variables corresponding to the images to fit and the generator’s input, Cov denotes the covariance matrix, ‖ · ‖F the Frobenius norm, and λ1, λ2 ≥ 0. The regularization in (16) enforces the weak dependence condition (Corollary 1), while (17) prevents the network to converge to d′ψ = 0. We call this generative adversarial network regularized deterministic SWG (reg-det-SWG).
To investigate the consequences of (i) regularizing the discriminator, and (ii) replacing the Monte Carlo SW with our approximation, we design another model, called regularized SWG (reg-SWG): similarly to SWG, the generator minimizes SW2,104 , but the discriminator’s objective is regularized as in (16), (17). We then compare reg-det-SWG against SWG and reg-SWG, by training the models on MNIST and CelebA and measuring their respective training time and Fréchet Inception Distances (FID, [44]): see Table 1. We used the same network architectures for all methods, and tuned (λ1, λ2) via cross-validation: more details on the experimental setup are given in Section S4. First, we observe that the regularized models produce images of higher quality, since reg-SWG and reg-det-SWG return lower FID values than SWG. The FID of reg-SWG and reg-det-SWG are close for both datasets, thus the two models seem to yield similar performances. Hence, we report in Figure 3 the images generated by SWG and reg-det-SWG only.
The training process is more expensive when regularizing the discriminator: the average running time per epoch is higher for the regularized models. We also observe that reg-det-SWG is faster than reg-SWG, which is consistent with the fact that our approximation method is faster than Monte Carlo on high-dimensional settings. To further illustrate this point, we reported the average time spent in computing the generative loss per epoch, i.e. SW2,104 for SWG and reg-SWG, and ŜW2 for reg-det-SWG: see column TSW in Table 1. On GPU, reg-det-SWG is at least 15 times faster than SWG and reg-SWG on MNIST, and 6 times faster on CelebA. Note that the models were trained using PyTorch, thus Monte Carlo SW benefits from a GPU-accelerated implementation of the sorting operation (with the function torch.sort). We also reported the computation times when models are trained on CPU. In this case, computing ŜW2 takes at most less than 3s per epoch, whereas the Monte Carlo estimation executes in several minutes (e.g., approximately 45min on CelebA). As a result, the total training time is almost the same for reg-det-SWG and SWG on CelebA, and the lowest for reg-det-SWG on MNIST.

heading: 5 Conclusion

text:
In this work, we presented a novel method to approximate the Sliced-Wasserstein distance of order 2, which relies on the concentration-of-measure phenomenon for random projections. The resulting method computes SW with simple deterministic operations, which are computationally efficient even on high-dimensional settings and do not require any hyperparameters. We proved nonasymptotical guarantees showing that, under a weak dependence condition, the approximation error goes to zero as the dimension increases. Our theoretical findings are then illustrated with experiments on synthetic datasets. Motivated by the computational efficiency and accuracy of our approximate SW, we finally designed a novel approach for image generation that leverages our theoretical insights. As compared to generative models based on SW estimated with Monte Carlo, our framework produces images of higher quality with further computational benefits. This encourages the use of our approximate SW on other algorithms that rely on Monte Carlo SW, e.g. autoencoders [8] or normalizing flows [12].
The weak dependence condition can be inappropriate to describe the underlying geometry of real data in ML applications, and in that case, approximating SW with our method seems inadequate. To overcome this problem, we encourage practitioners to resort to models where real data are represented by features that can be made weakly dependent. This strategy has proven successful in our image generation experiment: the reg-det-SWG model uses our approximation to compare two sets of features (instead of the raw images) whose covariance matrices are regularized to enforce weak dependence. Since many ML techniques make use of features and regularizers, we believe that our methodology is not restrictive and can then be applied to other standard problems than image generation. Besides, our weak dependence condition in Definition 1 is weaker than the one in [39], which is a notion commonly used in statistics.
Our empirical results on synthetic data show that the approximation error goes to zero with a faster convergence rate than the one we proved. Then, the main current limitation of our framework is that our theoretical convergence rate in d−1/8 might be slower than necessary. We proved that the overall approximation error is upper-bounded by a term in d−1/2 when comparing Gaussians with diagonal covariance matrices, and the improvement of our error bounds for other specific distributions is left for future work. On the other hand, the extension of our methodology to variants of SW is another challenging future research direction. To the best of our knowledge, the literature on the concentration of measure phenomenon focuses on linear random projections, therefore the derivation of deterministic approximations for SW based on nonlinear projections seems highly nontrivial. A more promising direction would be to generalize our approach to SW based on k-dimensional linear projection by leveraging the bound in [26, Theorem 1] for k > 1.
Since this paper is focused on developing a theoretically-grounded novel method to estimate a distance between probability distributions, we believe it will not pose any negative societal or ethical consequence. On the other hand, as demonstrated in Section 4, our contribution provides tools to speed up existing machine learning algorithms on CPU, which is useful when powerful hardware resources are not available, or when their use is deliberately avoided for environmental purposes.
Acknowledgments and Disclosure of Funding
This work is partly supported by the industrial chair “Data Science & Artificial Intelligence for Digitalized Industry & Services” from Télécom Paris. Umut Şimşekli’s research is supported by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute). Pierre E. Jacob gratefully acknowledges support by the National Science Foundation through grant DMS1844695. Alain Durmus acknowledges support of the Lagrange Mathematical and Computing Research Center. Kimia Nadjahi is grateful to Pierre Colombo for his helpful advice on how to train the generative models in Section 4 on Télécom Paris’s GPUs.

title: Fast Approximation of the Sliced-Wasserstein Distance Using Concentration of Random Projections