Adaptive Federated Learning with Auto-Tuned Clients
Authors: Junhyung Lyle Kim
Taha Toghani
Cesar A Uribe
Anastasios Kyrillidis
Venue: ICLR 2024 poster (2024)
Keywords: federated learning
distributed optimization
adaptive method
Primary area: optimization
TL;DR: We propose a simple step size rule for SGD that enables each client to use its own adaptive step size that shows great performance without additional tuning in different FL settings.

Abstract: Federated learning (FL) is a distributed machine learning framework where the global model of a central server is trained via multiple collaborative steps by participating clients without sharing their data. While being a flexible framework, where the distribution of local data, participation rate, and computing power of each client can greatly vary, such flexibility gives rise to many new challenges, especially in the hyperparameter tuning on the client side. We propose $\Delta$-SGD, a simple step size rule for SGD that enables each client to use its own step size by adapting to the local smoothness of the function each client is optimizing. We provide theoretical and empirical results where the benefit of the client adaptivity is shown in various FL scenarios.

Published as a conference paper at ICLR 2024
ADAPTIVE FEDERATED LEARNING WITH AUTO-TUNED CLIENTS
Junhyung Lyle Kim" Mohammad Taha Toghani César A. Uribe & Anastasios Kyrillidis" Department of Computer Science, 'Dpparment of Electrical and Computer Engineering Rice University, Houston, TX 77005, USA (jlylekim, mitoghani, cauribe, mastasioosf&rice.em
ABSTRACT
Federated learning (FL) is distributed machine learning framework where the global model of central server is trained via multiple collaborative steps by participating clients without sharing their data. While being flexible framework, where the distribution of local data, participation rate, and computing power each client can greatly vary, such flexibility gives rise many new challenges, especially in the hyperparameter tuning on the client side. We propose A-SGD, simple step size rule for SGD that enables each client use own step size by adapting to the local smoothness of the function each client is optimizing. We provide theoretical and empirical results where the benefit of the client adaptivity is shown in various FI. scenarios.
INTRODUCTION
Federated Learning (FL.) is machine learning framework that enables multiple clients to collabora- tively learn global model in a distributed manner. Each client trains the model on their local data, and then sends only the updated model parameters to a central server for aggregation. Mathematically, FL aims to solve the following optimization problem:
min f(x)
= 7 Σ f((z),
ZER
where f((x) == RE-DD, (F(x,, =]] is the loss function of the i-th client, and m the number of clients.
A key property of FL its flexibility in how various clients participate in the overall training procedure. The number of clients, their participation rates, and computing power available to each client can vary and change at any time during the training. Additionally, the local data of each client is not shared with others, resulting in better data privacy (McMahan et al., 2017; Agarwal et al., 2018).
While advantageous, such flexibility also introduces plethora of new challenges, notably: i) how the server aggregates the local information coming from each client, and ii) how to make sure each client meaningfully "learns" using their local data and computing device. The first challenge was partially addressed in Reddi et al. (2021), where adaptive optimization methods such as Adam (Kingma & Ba, 2014) was utilized in the aggregation step. Yet, the second challenge remains largely unaddressed.
Local data of each client is not shared, which intrinsically introduces heterogeneity in terms of the size and the distribution local datasets. That is. D; differs for each client as well as the number of samples= ~ DD Consequently. [x(): can vastly differ from client to client, making the problem in (1) hard to optimize. Moreover, the sheer amount of local updates is far larger than the number of aggregation steps, due the high communication cost-typically 3-4x orders of magnitude more expensive than local computation- distributed settings (Lan et al., 2020).
As result, extensive fine-tuning of the client optimizers is often required achieve good performance in FL scenarios. For instance, experimental results of the well-known FedAvg algorithm were obtained after performing grid-search of typically 11-13 step sizes of the clients' SGD (McMahan et al., 2017, Section 3), as SGD (and its variants) are highly sensitive to the step size (Toulis & Airoldi, 2017; Assran & Rabbat, 2020; Kim et al., 2022b). Similarly, in Reddi al. (2021), 6 different client

---

Published as a conference paper at ICLR 2024
CIFAR 10 trainec with ResNet 18
MNIST trained with CNN
CIFAR
100 trained with ResNet 50
1.0
0.8
0.6
Test Acturary
0.8
0.5
0.6
0,4
0.4
0.
0.4
022
0.2
0.2
0.1
C.0
0
500
10:00
1500
2000
0
100
200
30C
400
500
0
50.0
1000
1500
2000
Number of Rounds
Number Rounds
Numbbr Rounds
[A]
(8)
(C)
A-SGD (ours)
SGD
SGD (decay)
SGDM
SGDM (decay)
Acam
Adagrad
Figure 1: Illustration of the effect of not properly tuning the client step sizes. In (A), each client optimizer uses the best step size from grid-search. Then, the same step size from (A) intentionally used in settings (B) and (C). Only A-SGD works well across all settings without additional tuning.
step sizes were grid-searched for different tasks, and not surprisingly, each task requires a different client step size to obtain the best result, regardless of the server-side adaptivity (Reddi et al., 2021, Table 8). Importantly, even these "fine-tunings" are done under the setting that all clients use the same step size, which is sub-optimal, given that fi can be vastly different per client; we analyze this further in Section 2.
Initial examination as motivation. The implication of not properly tuning the client optimizer is highlighted in Figure We plot the progress of test accuracies for different client optimizers, where, for all the other test cases, we intentionally use the same step size rules that were fine-tuned for the task in Figure ((A). There, we train a ResNet-18 for CIFAR- 10 dataset classification within an FL. setting, where the best step sizes for each (client) optimizer was used after grid-search; defer the experimental details to Section 4. Hence, all methods perform reasonably well, although A-SGD, our proposed method, achieves noticeably better test accuracy when compared to non-adaptive SGD variants -e.g., a 5% difference in final classification accuracy from SGDM- and comparable final accuracy only with adaptive SGD variants, like Adam and Adagrad.
In Figure IB), we train a shallow CNN for MNIST classification, using the same step sizes from (A). MNIST classification is now considered an "easy" task, and therefore SGD with the same constant and decaying step sizes from Figure 1(A) works well. However, with momentum, SGDM exhibits highly oscillating behavior, which results in slow progress and poor final accuracy, especially without decaying the step size. Adaptive optimizers, like Adam and Adagrad, show similar behavior, falling short in achieving good final accuracy, compared their performance in the case of Figure 1(A).
Similarly, in Figure I(C), we plot the test accuracy progress for CIFAR-106 classification trained on a ResNet-50, again using the same step size rules as before. Contrary to Figure 1(B), SGD with momentum (SGDM) works better than SGD, both with the constant and the decaying step sizes. Adam becomes a "good optimizer" again, but its "sibling." Adagrad, performs worse than SGDM. On the contrary, our proposed method, A-SGD which we introduce in Section 3, achieves superior performance in all cases without any additional tuning.
The above empirical observations beg answers to important and non-trivial questions in training FL tasks using variants of SGD methods as the client optimizer: Should the momentum be used? Should the step size be decayed? so. when? Unfortunately, Figure 1 indicates that the answers these questions highly vary depending on the setting; once the dataset itself or how the dataset is distributed among different clients changes, once the model architecture changes, the client optimizers have to be properly re-tuned to ensure good performance. Perhaps surprisingly. the same holds for adaptive methods like Adagrad (Duchi al., 201 1) and Adam (Kingma & Ba, 2014).
Our hypothesis and contributions. Our paper takes stab in this direction: we propose DELTA- SGD (DistributEd LocaliTy Adaptive SGD), simple adaptive distributed SGD scheme, that can automatically tune its step size based on the available local data. We will refer to our algorithm as SGG in the rest of the text. Our contributions can be summarized as follows:
We propose A-SGD, which has two implications in FL settings: i) each client can its own step size, and ii) each client's step size adapts to the local smoothness of fü -hence LocaliTy Adaptive- and can even increase during local iterations. Moreover, due to the simplicity of the proposed step
2

---

Published as a conference paper at ICLR 2024
size, A-SGD is agnostic to the loss function and the server optimizer, and thus be combined with methods that use different loss functions such as FedProx (Li et al., 2020) or MOON (Li et al., 2021), adaptive server methods such as FedAdam (Reddi et al., 2021).
We provide convergence analysis of A-SGD in general nonconvex setting (Theorem 1). We also prove convergence convex setting; due to space constraints, we only state the result of the general nonconvex case in Theorem and defer other theorems and proofs to Appendix A.
We evaluate our approach on several benchmark datasets and demonstrate that A A-SGD achieves superior performance compared to other state-of-the-art FL. methods. Our experiments show that A-SGID can effectively adapt the client step size to the underlying local data distribution, and achieve convergence of the global model without any additional tuning. Our approach can help overcome the client step size tuning challenge in FL and enable more efficient and effective collaborative learning in distributed systems.
2 PRELIMINARIES AND RELATED WORK
bit of background on optimization theory. Arguably the most fundamental optimization algorithm for minimizing function f(x): : R + E is the gradient descent (GD), which iterates with the step size 1he as: 21+1 = 21 1)p V (zt). Under the assumption that f is globally L-smooth, that is:
| ffyiill < L- 1a-ll--- 1o---l yll
Vx, E Rd.
the step size v)- = 1 is the "optimal" step size for GD, which converges for convex f at the rate:
L == - -((+1)) & 2(21+ +
Limitations of popular step sizes. In practice, however, the constants like L are rarely known (Boyd al., 2004). As such, there has been a plethora of efforts in the optimization community develop universal step size rule or method that does not require a priori knowledge of problem constants so that an optimization algorithm can work as pplug-and-play (Nesterov, 2015; Orabona Pál, 2016; Levy et al., 2018). A major lines of work on this direction include: i) line-search (Armijo, 1966; Paquette & Scheinberg, 2020), úi) adaptive learning rate (e.g., Adagrad (Duchi et al., 201 and Adam (Kingma & Ba, 2014)), and iii) Polyak step size (Polyak, 1969), to name few.
Unfortunately, the pursuit of finding the "ideal" step size still active (Loizou et al., 2021; Defazio & Mishchenko, 2023; Bernstein et al., 2023), the aforementioned approaches have limitations. For instance, to use line search, solving a sub-routine is required, inevitably incurring additional evaluations of the function and/or the gradient. To use methods like Adagrad and Adam, the knowledge of D, the distance from the initial point to the solution set, is required to ensure good performance (Defazio & Mishchenko, 2023). Similarly, Polyak step size requires knowledge of / (x*) (Hazan & Kakade, 2019), which is not always known a priori, and is unclear what value to use for an arbitrary loss function and model; see also Section 4 where not properly estimating D or f(*) resulting in suboptimal performance of relevant client optimizers in FL scenarios.
Implications in distributed/FL scenarios. The global L-smoothness in (2) needs be satisfied for all r and y. implying even finite L can be arbitrarily large. This results in a small step size, leading to slow convergence, as seen in (3). Malitsky& & Mishchenko (2020) attempted to resolve this issue by proposing a step size for (centralized) GD that depends on the local smoothness of f. which by definition is smaller than the global smoothness constant L in (2). A-SG inspired from this work.
Naturally, the challenge is aggravated in the distributed case. To solve (1) under the assumption that each f; is T.-smooth, the step size of the form 1/ Tmmax where Temax == max; T42 often used (Yuan et al., 2016; Scaman et al., 2017: Qu & Li, 2019), to ensure convergence (Uribe et al., 2020). Yet. one can easily imagine a situation where L, >> L, for & j, in which case the convergence of f (x) can be arbitrarily slow by using step size of the form 1/Lmax-
Therefore, to ensure that each client learns useful information in FL settings as in (1), ideally: a) each agent should be able to use its own step size, instead of crude ones like 1/ Lm for all agents; ii) the individual step size should be "locally adaptive" to the function fú, (i.e., even 1/L, can be too crude, analogously to the centralized GD); and iii) the step size should not depend on problem constants
3

---

Published as a conference paper at ICLR 2024
like L, μ4, which are often unknown. In Section 3, we introduce DistributEd LocaliTy Adaptive SGD (A-SGD), which satisfies all of the aforementioned desiderata.
Related work on FL. The FL literature is vast, here we focus the results that closely relate to work. The FedAvg algorithm (McMahan et al., 2017) is one of the simplest FL algorithms, which averages client parameters after some local iterations. Reddi al. (2021) showed that FedAvg is a special case meta-algorithm, where both the clients and the server use SGD optimizer, with the server learning rate being To handle data heterogeneity and model drift, they proposed using adaptive optimizers at the server. This results in algorithms such as FedAdam and Yogiogi, which respectively use the Adam (Kingma & Ba, 2014) and Yogi (Zaheer et al., 2018) as server optimizers. Our approach is orthogonal and complimentary to this, as we propose a client adaptive optimizer that can be easily combined with these server-side aggregation methods (c.f., Appendix B.4).
Other approaches have handled the heterogeneity by changing the loss function. FedProx (Li et al., 2020) adds an 22-norm proximal term to to handle data heterogeneity. Similarly, MOON (Li et al., 2021) uses the model-contrastive loss between the current and previous models. Again, our proposed method can be seamlessly combined with these approaches (c.f., Table 2b, Appendix B.6 and B.5),
The closest related work to ours is concurrent work by Mukherjee et al. (2023). There, authors utilize the Stochastic Polyak Stepsize (Loizou et al., 2021) in FL settings. We do include SPS results in Section 4. There some previous works that considered client adaptivity, including AdaAlter (Xie et al., 2019a) and Wang et al. (2021). Both works utilize an adaptive client optimizer similar to Adagrad. However, AdaAlter incurs twice the memory compared FedAvg, AdaAlter requires communicating both the model parameters as well as the "client accumulators" similarly, Wang et al. (2021) requires complicated local and global correction steps achieve good performance. In short, previous works that utilize client adaptivity require modification in server-side aggregation; without such heuristies, Adagrad (Duchi et al., 2011) exhibits suboptimal performance, as seen in Table
Lastly, another line of works attempts to handle heterogeneity by utilizing control-variates (Pra- neeth Karimireddy et al., 2019; Liang al., 2019; Karimireddy et al., 2020), a similar idea to variance reduction from single-objective optimization (Johnson & Zhang, 2013; Gower et al., 2020). While theoretically appealing. these methods either require periodic full gradient computation or are not usable under partial client participation, as all the clients must maintain state throughout all rounds.
DELTA(A)-SGD: DISTRIBUTED LOCALITY ADAPTIVE SGD
We now introduce A-SGD. In its simplest form, client i at communication round t uses the step size:
v = min { [ VAGE V 0z 1i- 0:-1 = vi-1/1i22
The first part of minf..) approximates the (inverse of) local smoothness of ft.' and the second part controls how fast oi can increase. Indeed, A-SGD with (4) enjoys the following decrease in the Lyapunov function:
| 1=1
x
2
2
The O4g (1i(at)- (((a)))((a+))))
2m
<
10, - x-
112
(5)
2m
when fi's are assumed to be convex (c.f., Theorem 5 in Appendix A). For the Fl. settings, we extend (4) by including stochasticity and local iterations, as summarized in Algorithm and visualized in Figure 7. For brevity, we use V (x()) = T Exes V (F, =) to denote the stochastic gradients with batch size [8] = b.
NNttcee that [V/(+)) V.((at
< 2n1 ||aq wi || x Leallai z/-1ll-
4

---

Published as a conference paper at ICLR 2024
Algorithm DELTA(A)-SGD) DistributEd LocaliTy Adaptive SGD
1: input: z0 R.d 10. 00.7 0, and p E (0, 1). do
2: for each round = 0, 1......7
7-1 7-
3: sample a subset S, of clients with size |Sy] = m
4: for each machine in parallel for i E S, do
5: set o²,0 = Ztt
6:
set Mio0 = and 07.0 = 8
7:
for local step A E [K] do
8:
f 2 k- ph,k- fi
> update local parameter with A-SGD
9:
M = min
+| -
+ - Mi.k
2 ( Vi x Dll
10:
8i n /
(line 9 & 10) update locality adaptive step size
11: end for
12: end for
13: I++1
1 Eies, 4/,K S.
server-side aggregation
end for
15:
return ZT
We make a few remarks of Algorithm 1. First, the input 0o can be quite arbitrary, as it can be corrected, per client level, in the first local iteration (line 10); similarly for vo > 0, although no should be sufficiently small to prevent divergence in the first local step. Second, we include the "amplifier" y to the first condition of step size (line 9), but this is only needed for Theorem 1.2 Last, fz shows up twice: in updating zi, (line 8) and ni,k (line 9). Thus, one can use the same or different batches; we use the same batches in experiments to prevent additional gradient evaluations.
3.1 CONVERGENCE ANALYSIS
Technical challenges. The main difference of analyzing Algorithm compared other decentralized optimization algorithms is that the step size ni, not only depends on the round and local steps k. but also on i, the client. To deal with the client-dependent step size, we require a slightly ntn-standard assumption the dissimilarity between f and f.. as detailed below.
Assumption There exist nonnegative constants a, p. and G such that for all E [M] and x E TRd,
(,, ()i(2)) <o?
(bounded variance)
(1a)
[V/[(>)]| G [V)(=) [= {.)) (strong growth of dissimilarity)
(bounded gradient)
(Ib)
(1c)
Assumption Ia has been standard in stochastic optimization literature (Ghadimi Lan, 2013; Stich, 2018; Khaled et al., 2020); recently, this assumption has been relaxed to weaker ones such expected smoothness (Gower et al., 2021), but this is out of scope of this work. Assumption 1 fairly standard in nonconvex optimization (Zaheer et al., 2018: 8; Ward al., 2020), and often used in FL setting (Xie et al., 201 9a;b; Reddi et al., 2021). Assumption 1c reminiscent of the strong growth assumption in stochastic optimization (Schmidt & Roux, 2013), which still used in recent works (Cevher & Vü, 2019; Vaswani et al., 201 9a;b). To the best of our knowledge, this is the first theoretical analysis of the FL setting where clients can use their own step size. We now present the main theorem.
Theorem Let Assumption hold, with p O(1). Further, suppose that y O k k and = O(r). Then, the following property holds for Algorithm 1, for T sufficiently large:
T-1
V1 /T
i i2 T
Σ
E ( <
+
+
T
VT3
For all our experiments, use the default value 2 from the original implementation in https: //gi thub. com/ ymal tsky/ adatt iveptive_t//bb http:///bb/mmttrrr/ttoo///////// zer.py.
5

---

Published as a conference paper at ICLR 2024
where 11=max { f(o)) - f(x²)} and V2 22= 22= (2= are global constants, with = [B] being the batch size; L, is a constant at most the maximum of local smoothness, i.e,, maxi, Tri, where Lig the local smoothness of fa round t.
The convergence result in Theorem I implies a sublinear convergence to an E-first order stationary point with at least T (e 2) communication rounds. We remind again that the conditions on y and no are only required for our theory.' Importantly, we did not assume fi is globally L-smooth, a standard assumption in nonconvex optimization and FL literature (Ward et al., 2020; Koloskova et al., 2020; Li al., 2020; Reddi al., 2021); instead,
obtain smaller quantity, i., through our analysis; for space limitation, we defer the details and the proof to Appendix A.
EXPERIMENTAL SETUP AND RESULTS
We now introduce the experimental setup and discuss the results. Our implementation can be found in https: //t///giiie github. com/ plylekimm aattttuuuuuttttt FL.
Datasets and models. We consider image classification and text classification tasks. We use four datasets for image classification: MNIST, FMNIST, CIFAR-10, and CIFAR-100 (Krizhevsky et al., 2009). For MNIST and FMNIST, we train shallow CNN with two convolutional and two fully- connected layers, followed by dropout and ReLU activations. For CIFAR-10 10, we train a ResNet-18 (He al., 2016). For CIFAR-I00, we train both ResNet-18 and ResNet--0 to study the effect of changing the model architecture. For text classification, we use two datasets: DBpedia and AGnews datasets (Zhang et al., 2015), and train DistilIBERT (Sanh et al., 2019) for classification.
For each dataset, we create a federated version by randomly partitioning the training data among 100 clients (50 for text classification), with each client getting 500 examples. To control the level of non-iidness, we apply latent Dirichlet allocation (LDA) over the labels following Hsu et al. (2019), where the degree class heterogeneity can be parameterized by the Dirichlet concentration parameter a. The class distribution of different a's is visualized in Figure 8 in Appendix B.9.
FL setup and optimizers. We fix the number of clients to be 100 for the image classification, and 50 for the text classification; randomly sample 10% participating clients. We perform E' local epochs of training over each client's s dataset. We utilize mini-batch gradients of size b = 64 for the local gradient
image classification, and 16 for the text classification, leading K z | | b steps; we use E = 1 for all settings. For client optimizers, we compare stochastic gradient descent (SGD), SGD with momentum (SGDM), adaptive methods including Adam (Kingma & Ba, 2014), Adagrad (Duchi et al., 2011) 1), SGD with stochastic Polyak step size (SPS) (Loizou et al., 2021), and our proposed method: A-SGD in Algorithm 1, As our focus is client adaptivity, we mostly present the results using FedAvg (McMahan et al., 2017) as the server optimizer; additional results using FedAdam can be found Appendix B.4.
Hyperparameters. For each optimizer, we perform a grid search of learning rates on single task: CIFAR-10 classification trained with ResNet-18, with Dirichlet concentration parameter a 0.1: for the rest of the settings, we use the same learning rates. For SGD, we perform a grid search with 7 E (0.01, 0.05, 0.1, 0.5). For SGDM, we use the same grid for r) and use momentum parameter 3 0.9. To properly account for the SGD(M) fine-tuning typically done in practice, we also test dividing the step size by 10 after 50%, and again by 10 after 75% the total training rounds (LR decay). For Adam and Adagrad, we grid search with 77 E {0.001, 0.01, 0.1). For SPS, use the default setting of the official implementation." 4 For SGG,,, we append in front of the second
condition: / 1 11+58 50 Th.K- following Malitsky Mishchenko (2020), and use 0.1 for all experiments." 5 Finally, for the number of rounds T, we use 500 for MNIST, 1000 for FMNIST, 2000 for CIFAR-10 10 and CIFAR- 100, and 100 for the text classification tasks.
In all experiments, we use the default settings 7 = 2, m10 = 0.2, and 6g 1 without additional tuning. 'Le., we use f = 0, and e 0.5, for the SPS step size: f((x)= found in hst : / / thub. com/ 1 ssaamaarad ji / /sps.
The official implementation can be
We also demonstrate in Appendix B.1 that 5 has very little impact on the final accuracy of A-SGD.

---

Published as a conference paper at ICLR 2024
4.1 RESULTS
We clarify again that the best step size for each client optimizer tuned via grid-search for single task: CIFAR- 10 classification trained with a ResNet-18, with Dirichlet concentration parameter 0.1. We then intentionally use the same step size in all other tasks to highlight two points: A-SGD works well without any tuning across different datasets, model architectures, and degrees of heterogeneity; ii) other optimizers perform suboptimally without additional tuning.
Non-iidness
Optimizer
Dataset / Model
MNIST
FMNIST CNN
CIFAR-10 ResNct-18
CIFAR-100
CIFAR-100
Dir(a
CNN
ResNct-18
ResNct-50
SGD
_8.3.
86.3
87.7 u(.7
57.7 (44.)
53.0 53.0_(12.8)
SGD (L)
97.8u(0.7)
86.31(1-0
(7.8
61.94(0.0)
60.9
SGDM
98.5u(0.0;
85.2
88.7
58.8 (3.1))
60.5 60.5
SGDM (4)
(8.4
87.21(0.2)
89.3
(1.44(u.5)
63.3 (2.6)
Q ==1
Adam
94.7u(3.8)
71.8
89.4
55.61(0.))
61.44(4.4)
Adagrad
(4.31(3422)
45.5 (445)
86.61
53.5 (88.4)
51.9. (13.9)
SPS
10.1 ((8.4)
85.91(1.4)
82.7
1.0 (60.9)
50.0 (13.8)
A-SGD
98.4
87.31(0.0)
89.8
61.5
65.81(0.0)
SGD
98.1.(0.0)
(3.6)
72.1 (2.1.(12.9)
((.4)(
(4.2.
SGD (4)
98.0.(0.1)
84.74(1.7)
((.4)
59.31(1.8)
48.7 48.7_(15.4)
SGDM
97.6_(0.3)
83.6
(9.6
58.8 (2.3)
52.3 (218_
SGDM (4)
98.0_(0.1)
86.11(0.3)
77.9 17.91(7.1)
60.44(0.7)
52.8 (11.3)
a = 0.1
Adam
96.4_(1.7)
80.4
85.0
(5.4
(8.2
Adagrad
89.9.(8.2)
46.3 ((4.1
84.1
49.6,(11.5)
48.0, 48.0.(16.1)
SPS
96.0
85.0
70.3 70.3
42.2
42.2 (21.9)
A-SGD
98.1u
86.4
84.5
61.1
64.1410.0)
SGD
(6.8
79.0
22.6
30.5
24.31
SGD (1)
97.2L(0.3)
19.3
33.9
30.3
24.61(6.6)
SGDM
77.9
15.7
28.4
24.8
22.0
SGDM (4)
94.0 (3.5_
79.5
29.0
20.9 ((10.9)
14.7 (16.7)
= 0.01
Adam
80.8 ((6.7)
60.6
22. (211))11.8
18.2 ((3.6))
22.6 (2.6)(8.8)
Adagrad
72.4 (22..)
45.9 ((34.3)
12.5.(21.4)
25.8 ((6.))
22.2
SPS
69.7 (2277)
44.0 (36.2)
21.5
22.0
17.4 17.4_(14.0) (14.0)
A-SGD
97.5
80.2
31.6
31.84(000
31.44(0.0)
Table 1: Experimental results based the settings detailed in Section 4. Best accuracy (+ 0.5%) for each task are shown in bold. Subscripts is the performance difference from the best result, and highlighted in pink when it is bigger than 2%. The down-arrow symbol (L) indicates step-wise learning rate decay, where the step sizes are divided by 10 after 50%, and another by 10 after 75% of the total rounds.
CIFAR trainec with ResNet 18
FMNIST trained CNN
CIFAR-1 100 trained with ResNet-50
90
Test Accuracy
80
60
70
40
30
30
10
0.01
0.1
0.1 Dirich et Parameter
0.01
Dirichiet Parameser
Dirichlet Parameter
(B)
[C)
(ours)
SGD
SdDide
SCCM
(deaayy
Adam
Adagrad
SPS
Figure 2: The effect of stronger heterogeneity on different client optimizers, induced by the Dirich- let concentration parameter a E (0.01,0.1,11, A-SGD remains robust performance in cases, whereas other methods show significant performance degradation when changing the level of hetero- geneity a, or when changing the setting (modelJarchitecturr)..
7

---

Published as a conference paper at ICLR 2024
Changing the level of non-iidness. We first investigate how the performance of different client optimizers degrade in increasing degrees of heterogeneity by varying the concentration parameter {1, 11,, multiplied to the prior of Dirichlet distribution, following Hsu et al. (2019).
Three illustrative cases are visualized in Figure 2. We remind that the step sizes are tuned for the task in Figure 2(A). For this task, with o 1 (i.e., closer to iid), all methods perform better than a = 0,1, as expected. With a = 0.01, which is highly non-iid (c.f., Figure 8 in Appendix B.9), we see a significant drop in performance for all methods. SGD with LR decay and A-SGD perform the best, while adaptive methods like Adam (85% + 22.1% and Adagrad (84.1% + 12.5%) degrade noticeably more than other methods.
Figure 2(B) shows the results of training FMNIST with a CNN. This problem is generally considered casier than CIFAR-10 10 classification. The experiments show that the performance degradation with varying a's is much milder in this case. Interestingly, adaptive methods such as Adam, Adagrad, and SPS perform much worse than other methods
Lastly, in Figure 2(C), results for C1FAR-100 classification trained with ResNet-50 are illustrated. A-SGD exhibits superior performance all cases of a. Unlike MNIST and FMNIST, Adam enjoys the second (a = 0.1) or the third (ar = 1) best performance in this task, complicating how one should "tune" Adam for the task at hand. Other methods, including SGD with and without momentum/l.t decay, Adagrad, and SPS, perform much worse than A-SGD.
Changing the model architecture/dataset. For this remark, let us focus on CIFAR-IO0 100 trained on ResNet-18 versus on ResNet-50, with a 0.1, illustrated in Figure 3(A). SGD and SGDM (both with and without LR decay), Adagrad, and SPS perform worse using ResNet-50 than ResNet- 18. This is a counter-intuitive behavior, one would expect to get better accuracy by using a more powerful model, although the step sizes are tuned on ResNet- 18, A-SGD is an exception: without any additional tuning, A-SGD can improve its performance. Adam also improves similarly, but the achieved accuracy is significantly worse than that of A-SGD.
We now focus on cases where the dataset changes, but the model architecture remains the same. We mainly consider two cases, illustrated in Figure 3(B) and (C): When CNN is trained for classifying MNIST versus FMNIST, and when ResNet- is trained for CIFAR-10 versus CIFAR-100 100. From (B), one can infer the difficulty in tuning SGDM: without the LR decay, SGDM only achieves around 78% accuracy for MNIST classification, while SGD without LR decay still achieves over 96% accuracy. For FMNIST on the other hand, SGDM performs relatively well, but adaptive methods like Adam, Adagrad, and SPS degrade significantly.
Aggravating the complexity of fine-tuning SGD(M) for MNIST, one can observe from (C) a similar trend CIFAR-10 trained with ResNet-18. In that case, A-SG does not achieve the best test accuracy (although it is the second best with pretty big margin with the rest), while SGD with LR decay does. However, without LR decay, the accuracy achieved by SGD drops more than 11% Again, adaptive methods like Adam, Adagrad, and SPS perform suboptimally in this case.
Same Datasel Different Model
Same Model & Different Dalasel
Satte Model & Dif Difeerent Datasel
100
25
Test Accuracy
60
30
5D
25
40
20-
30
40
15
20
10
20
0
CII CID
CITAR-100 100
MNIST CNN
MNNIST CNN
CI ResNee-1
CIFAR-10 Resket 18
ResMeer-1 88
ResNet-50
(N)
(B)
[ Adagrad
A-53D
SGD
SGD (decayy
SGCM
(dDM(decay
Adam
SPS
Figure 3: The effect of changing the dataset and the model architecture different client optimizers. A-SGD remains superior performance without additional tuning when model or dataset changes, whereas other methods often degrades in performance. (A): CIFAR-100 trained with ResNet-18 versus Re (a = 0.1), (B): MNIST versus FMNIST trained with CNN (a = 0.01), (C): CIFAR-10 10 versus CIFAR- 100 trained with R8 (a 0.01).
8

---

Published as a conference paper at ICLR 2024
Changing the domain: text classification. In Table 2a, make bigger leap, and test the performance of different client optimizers for text classification tasks. We use the DBpedia and AGnews datasets (Zhang et al., 2015), and train DistilIBERT (Sanh et al., 2019) for classification. Again, the best step size for each client optimizer was tuned for the CIFAR-10 image classification trained with ResNet-18, 18, for 0.1. Thus, the type of the dataset changed completely from image to text, as well as the model architecture: from vision to language.
Not surprisingly, four methods (SGDM, SGDM(L), Adam, and Adagrad) ended up learning nothing (i.e., the achieved accuracy is 1/(# of labels)), indicating the fine-tuned step sizes for these methods using CIFAR- 10 classification task was too large for the considered text classification tasks. A-SGD, on the other hand, still achieves competitive accuracies without additional tuning, thanks to the locality adaptive step size. Interestingly, SGD,, SGD(J), and SPS worked well for the text classification tasks, which is contrary to their suboptimal performances for image classification tasks in Table 1.
Changing the loss function. In Table 2b, we illustrate how A-SGD performs when combined with a different loss function, such as FedProx (Li et al., 2020). FedProx aims to address heterogeneity; we thus used the lowest concentration parameter, a = 0.01. We only present a subset of results: (additional results in Appendix B.6). The results
CIFAR- 10 100) classification with ResNet- show that A SGD again outperforms all other methods with a sigmificant margin: A-SGD is not only the best, but also the performance difference is at least 2.5%, and large as 26.9%. Yet, compared to the original setting without the proximal term in Table 1, it is unclear whether or not the additional proximal term helps. For A-SGD, the accuracy increased for CIFAR-10 with ResNet--8, but slightly got worse for CIFAR- 100 with ResNet-50; a similar story applies to other methods.
Text classification
Dataset Model
FedProx
Dataset / Model
Agnews
Dbpedia
0 = 0.01 Optimizer
CIFAR-10
CIFAR-100
Optimizer
DistilBBER
RcsNet-18
ResNet-50
SGD
91.1,(0.55
96.0u(2.9)
SGD
20.0
(559)
SGD (L)
91.61(0.0)
48.7.(
SGD (L)
31.3 ((2.)
20.2
SGDM
25.0 25.0_(0a.6)
7.11
SGDM
29.3
23.8 (378))7.2)
SGDM (L)
25.0 (66.6)
7.1 1.11
SGDM (L)
25.3 ((8.))
15.0 (16.0()).0)
Adam
25.0 (c00.)
7.1 (091.00
Adam
28.1 ((0.7)
22.6.(8.4)
Adagrad
25.0 (6.00666)) 6)
7.1 (.1.8)
Adagrad
19.3 ((14.5
4.1 (26.9)
SPS
91.5
40.9
SPS
27.61(8.2)
16.5 16.55(14.5)
A-SGD
90.7
98.64(0.3)
A-SGD
33.8
31.0
(a) Experimental results for text classification
(b) Experimental results using FedProx loss.
Table 2: Additional experiments for a different domain (a), and a different loss function (b).
Other findings. We present all the omitted theorems and proofs in Appendix A, and additional empirical findings in Appendix B. Specifically, we demonstrate: the robustness of the performance of A-SGD (Appendix I.I), the effect of the different numbers of local iterations (Appendix B B.2), additional experiments when the size of local dataset different per client (Appendix B.3), additional experimental results using FedAdam (Reddi et al., 2021) (Appendix B.4), additional experiential results using MOON (Appendix B.5) and FedProx (Appendix B.6) loss functions, and three indepen- dent trials for a subset of tasks (Appendix B.7). We also visualize the step size conditions for A-SGD (Appendix B.8) as well as the degree of heterogeneity (Appendix B.9).
5 CONCLUSION
In this work, we proposed A-SGD, a distributed SGD scheme equipped with an adaptive step size that enables each client to use its own step size and adapts to the local smoothness of the function each client is optimizing. We proved the convergence of ASG in the general nonconvex setting and presented extensive empirical results, where the superiority of A-SGD is shown in various scenarios without any tuning. For future works, extending A-GG to a coordinate-wi -wiee step size in the spirit of Duchi al. (2011); Kingma& & Ba (2014), applying the step size to more complex algorithms like (local) factored gradient descent (Kim et al., 2022a), and enabling asynchronous updates (Assran al., 2020; Toghani et al., 2022; Nguyen et al., 2022) could be interesting directions.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work is supported by NSF FET: Small No. 1907936, NSF MLWINS CNS No. 2003137 (in collaboration with Intel). NSF CMMI No. 2037545, NSF CAREER award No. 2145629, NSF CIF No. 2008555, Rice InterDisciplinary Excellence Award (IDEA), NSF CCF No. 2211815, NSF DEB No. 2213568, and Lodieska Stockbridge Vaughn Fellowship.
REFERENCES
Naman Agarwal, Ananda Theertha Suresh, Felix Xinnan X Yu, Sanjiv Kumar, and Brendan McMahan. cpSGD: conmunication-effciinn and peterenttillyyppivvatt distributed SGD. Advances Neural Information Processing Systems, 31, 2018.
Larry Armijo. Minimization of functions having lipschitz continuous first partial derivatives. Pacific Journal of mathematics, 16(1):1-33 1966.
Mahmoud Assran and Michael Rabbat. On the convergence of nesterov's accelerated gradient method in stochastic settings. arXiv preprint arXiv: /2414, 2020.
Mahmoud Assran, Arda Aytekin, Hamid Reza Feyzmahdavian, Mikael Johansson, and Michael G Rabbat, Advances in asynchronous parallel and distributed optimization. Proceedings of the IEEE, 108(11):2013-201, 2020.
Jeremy Bernstein, Chris Mingard, Kevin Huang, Navid Azizan, and Yisong Yue. Automatic gradient descent: Deep learning without hyperparameters. arXiv preprint arXiv:2304.05187 2023.
Stephen Boyd, Stephen P Boyd, and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Volkan Cevher and Bang Công Vü. On the linear convergence of the stochastic gradient method with constant step-size. Optimization Letters, 1375:11177-1187, 2019.
Aaron Defazio and Konstantin Mishchenko. Learning-rate-fre learning by d-adaptation. In Interna- tional Conference on Machine Learning, pp. 7449-7479. PMLR, 2023.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, ([(4):2341-2368 2013.
Robert M Gower, Mark Schmidt, Francis Bach, and Peter Richtárik. seriance-reduced methods for machine learning. Proceedings of the IEEE, 111(1))1168-1983 2020.
Robert Gower, Peter Richtárik, and Francis Bach. Stochastic quasi-gradient methods: Variance reduction via jacobian sketching. Mathematical Programming, 188:135-192, 2021.
Elad Hazan and Sham Kakade. Revisiting the polyak step size. arXiv preprint arXiv: 1rXiv:1905.000 1905.00313, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.00 06335, 2019.
Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. Advances in neural information processing systems, 26, 2013.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv preprint 20l0:2000000000 2020.
10

---

Published as a conference paper at ICLR 2024
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local sgd on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics, pp. 4519-4529 PMLR, 2020.
Junhyung Lyle Kim, Mohammad Taha Toghani, César A Uribe, and Anastasios Kyrillidis. Local stochastic factored gradient descent for distributed quantum state tomography. IEEE Control Systems Letters, 7:199-204, 2022a.
Junhyung Lyle Kim, Panos Toulis, and Anastasios Kyrillidis. Convergence and stability of the stochastic proximal point algorithm with momentum. In Learning for Dynamics and Control Conference, pp. 1084-104.. PMLR, 2022b.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14X1::1412.69800 2014.
Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian Stich. A unified theory of decentralized sgd with changing topology and local updates. In International Conference on Machine Learning, pp. 5381-5393, PMLR, 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Guanghui Lan, Soomin Lee, and Yi Zhou. Communication-efficcenn algorithms for decentralized and stochastic optimization. Mathematical Programming. 180(1-2):377-24, 2020.
Kfir Levy, Alp Yurtsever, and Volkan Cevher. Online adaptive methods, universality and accelera- tion. Advances in neural information processing systems, 31, 2018.
Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 10713-10722, 2021.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429-450 2020.
Xianfeng Liang. Shuheng Shen, Jingchang Liu, Zhen Pan, Enhong Chen, and Yifei Cheng. Variance reduced local sgd with lower communication complexity. arXiv preprint arXiv:1912.12844, 2019.
Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste-Julien. Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence. In International Conference on Artificial Intelligence and Statistics, pp. 1306-1314. PMLR, 2021.
Yura Malitsky and Konstantin Mishchenko. Adaptive gradient descent without descent. In Proceed- ings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1702-6712. PMLR, 2020.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication-effcceenn learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
Sohom Mukherjee, Nicolas Loizou, and Sebastian Stich. Locally adaptive federated learning via stochastic polyak stepsizes. arXiv preprint arXiv:2307.066306, 2023.
Yu Nesterov. Universal gradient methods for convex optimization problems. Mathematical Program- ming, (52))32) 404, 2015.
John Nguyen, Kshitiz Malik, Hongyuan Zhan, Ashkan Yousefpour, Mike Rabbat, Mani Malek, and Dzmitry Huba. Federated learning with buffered asynchronous aggregation. In International Conference on Artificial Intelligence and Statistics, pp. 30113536600 PMLR, 2022.
Francesco Orabona and Dávid Pál. Coin betting and parameter-free online learning. Advances in Neural Information Processing Systems, 29, 2016.
Courtney Paquette and Katya Scheinberg. A stochastic line search method with expected complexity analysis. SIAM Journal Optimization, 30(1):349-376, 2020.
11

---

Published as a conference paper at ICLR 2024
Boris Teodorovich Polyak. Minimization of nonsmooth functionals. Zhurnal Vychislitel noi Matem- atiki Matematicheskoi Fiziki, 9(3):509 521, 1969.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning, arXiv e-prints, pp. arXiv-1910, 2019.
Guannan Qu and Na Li. Accelerated distributed nesterov gradient descent. IEEE Transactions Automatic Control, 65(6):2566-258) 2019.
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference Learning Representations, 2021.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert. a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108, 2019.
Kevin Scaman, Francis Bach, Sébastien Bubeck, Yin Tat Lee, and Laurent Massoulié. Optimal algorithms for smooth and strongly convex distributed optimization in networks. international conference on machine learning, pp. 3027-3036. PMLR, 2017.
Mark Schmidt and Nicolas Le Roux. Fast convergence of stochastic gradient descent under a strong growth condition. arXiv preprint arXiv: 00.330866770 2013.
Sebastian U Stich. Local sgd converges fast and communicates little. arXiv preprint arXiv: 2018.
Mohammad Taha Toghami, Soomin Lee, and César A Uribe. PersA-FL: Personalized Asynchronous Federated Learning. arXiv preprint arXiv:2210.01176, 2022.
Panos Toulis and Edoardo M Airoldi. Asymptotic and finite-sample properties of estimators based on stochastic gradients. 2017.
César Uribe, Soomin Lee, Alexander Gasnikov, and Angelia Nedić. A dual approach for optimal algorithms in distributed optimization over networks. In 2020 Information Theory and Applications Workshop (ITA), pp- 1-37. IEEE, 2020.
Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of sgd for over- parameterized models and an accelerated perceptron. In The 22nd international conference on artificial intelligence and statistics, pp. 1115-11200 1204. PMLR, 2019a.
Sharan Vaswani, Aaron Mishkin, Issam Laradji, Mark Schmidt, Gauthier Gidel, and Simon Lacoste- Julien. Painless stochastic gradient: Interpolation, line-search, and convergence rates. Advances in neural information processing systems, 32, 2019b.
Jianyu Wang, Zheng Xu, Zachary Garrett, Zachary Charles, Luyang Liu, and Gauri Joshi. Local adaptivity in federated learning: Convergence and consistency. arXiv preprint 2rX.v:2106.02305, 2021.
Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. The Journal of Machine Learning Research, 21(1):904799076 2020.
Cong Xie, Oluwasanmi Koyejo, Indranil Gupta, and Haibin Lin. Local adaalter: Communication- efficient stochastic gradient descent with adaptive learning rates. arXiv preprint 1rXiv: 2019a.
Cong Xie, Sanmi Koyejo, and Indranil Gupta. Asynchronous federated optimization. arXiv preprint 1003v03903..3934, 2019b.
Kun Yuan, Qing Ling, and Wotao Yin. On the convergence of decentralized gradient descent. SIAM Journal on Optimization, 26(3):1835-1855 2016.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. Advances in neural information processing systems, 31, 2018.
Xiang Zhang, Junbo Zhao, and Yann LeCun Character-level convolutional networks for text classification. Advances in neural information processing systems, 28, 2015.
12

---

Published as a conference paper at ICLR 2024
SUPPLEMENTARY MATERIALS FOR "ADAPTIVE FEDERATED LEARNING WITH AUTO-TUNED CLIENTS"
In this appendix, we provide all missing proofs that were not present in the main text, as well as additional plots and experiments. The appendix is organized as follows.
In Appendix A, we provide all the missing proofs. In particular:
In Section A. we provide the proof of Algorithm 1, under Assumption 1.
In Section A.2, we provide the proof of Algorithm, 1, with some modified assumptions, including that is convex for all E [m].
In Appendix B, we provide additional experiments, well miscellaneous plots that were missing in the main text due to the space constraints. Specifically:
In Section B.I, we provide some perspectives on the case of tuning Å-SGD in Algorithm 1. In Section B.2, we provide some perspectives on how A-SGD performs using different number of local iterations.
In Section B.3, we provide additional experimental results where each client has different number of samples as local dataset.
In Section B.4, we provide additional experimental results using FedAdam (Reddi et al., 2021) server-side adaptive method.
In Section B.5, we provide additional experimental results using MOON (Li et al., 2021) loss function.
In Section B.6, we provide additional experiemtns using FedProx (Li et al., 2020) function.
In Section B.7, we perform three independent trials for a subset of tasks for client-optimizees. We plot the average and the standard deviations in Figure 6, and report those values in Table 7. In Section B.8, visualize the step size conditions for A-SGD.
In Section B.9, we provide illustration the different level of heterogeneity induced by the Dirichlet concentration parameter a.
A MISSING PROOFS
We first introduce the following two inequalities, which will be used in the subsequent proofs:
1
(2,,
||,; 1|
2 ||,,119,
and
2
Σ
m
(8)
1=1
for any set of m vectors {2.)m with x;
Rd
We also provide the following lemma, which establishes L-smoothness based the local smoothness.
Lemma 2. For locally smooth function fü, the following inequality holds:
(u4) & f((x) + (Vf((), x) + L lly x11
Vx, E St,
(9)
2
where S.,, conv((=)),) (++)).
Proof. First, notice that by Assumptions (1a) and (Ib), stochastic gradients are also bounded for all E gd. Therefore, the local iterates, i.e., {x] kmo remain bounded. Hence, the set S := conv( (
1
k= ( U (Ki, x141) }), where conv(-) denotes the convex hull, is bounded. Due the fact
13

---

Published as a conference paper at ICLR 2024
that is locally smooth, ff is Lipschitz continuous bounded sets. This implies there exists positive constant Li, such that,
[V)( Vi((i)ll & L, - æl1 Vx,y St
(10)
Thus, we have
(fly - Vf((z). & Vv)(t) vf(2')] l1y Ill
Lielly - xl13, Vx, y St,
Let g(T) = fi(2 + Tyy - z)). Then, we have
(T) + ())-)))), x)
-'(r) (Vf(x+7(y--)) Vf(=)), y --x
[ x)) Vf(2)) 11yl-lly 1 =l
< < -
Using the mean value theorem, we have
f(()) 9(1) g(0) S. g'(r)dr
< (0) + / [g((0) + +L,, 11 =11
<0(0) 5((0) + Lid |1y - =x112
2
< f()) + (Vf()), y -x) =))+ + T.,t 11 x012
2
Now, define L, := max, L.,-- Given this definition, relate the local-smootee smoothnes constant İ of the average function with (individual) smoohhssss constants L; in the below lemma.
Lemma 3. The average of local smooth functions fr are at least ( Em L.)-local smooth.
Proof. Starting from (10), we have:
1
f
Σ
(Vii(x)
1=1
Vf((x)
m
m
L yl1
(11)
m
i. 1 vll-
where -a=max Lt,t-
(12)
Technically, (11) can provide slightly tighter bound for Theorem 1, but (12) simplifies the final expression. We also have the following form:
L f(++11 & f(x)) + (VI(x)) :1-1 - xx) | - 12
which can be obtained from Lemma 2.
14

---

Published as a conference paper at ICLR 2024
A.1
PROOF OF ALGORITHM UNDER ASSUMPTION 1
Proof. Based on Algorithm 1, we use the following notations throughout the proof:
F 1,0 It:
2t, -i,k-
V
vk [K],
Tt-1
[
k
SES,
where we denoted with Je+1 as the server parameter at round t + 1, which is the average of the local parameters oi,x for all client V i [m]] after running local iterations for V k [K].
We also recall that we use (z(()) | EEEB (z) to denote the stochastic gradients with batch size [B] b. Note that whenever we use this shorthand notation in the theory, imply that we draw a new independent batch of samples to compute the stochastic gradients.
We start from the smoothness of f()), which note again is obtained via Lemma 3, i.e., byproduct of our analysis, and we do not assume that fi's are globally L-smooth. Also note that Lemma 3 holds for {zo, T1,... J, i.e., the average of the local parameters visited by Algorithm
Proceeding, we have
(((+1) (((4))(())) ((x)) (VI(z))
Σ
K-1 Σ
+
+
Σ M (
S
Se
iES.
Taking expectation from the expression above, where we overload the notation such that the expec- tation is both with respect to the client sampling randomness (indexed with i) as well as the data sampling randomness (indexed with 2), that is IE(.) Ex, E, [E Fe, F,l,,,, where Fg is the natural filtration, we have:
Bf(x++1)
K-1
K-1
f(())) E Vf(1)),
Σ
T f((((.))
1
Σ Tie
(13)
2
Se
iES.
ieS,
Am0
K-1
<((x)) E
Σ
E
The
(14)
2m
K=U
k=0
m
K-1
Mhe D,
-(()) D. D
D, (f(
m
K K-1
The (f
(15)
2m
k=
where in the equality, we added and subtracted L
K-1 nhe V in the inner product term, and also replaced D, == 1. > K-1 n Then, using (7) and (13)-(15), we have
Σ - Eko
Dj |Vf((z))1
Ef(x++1) < f(x)) - D. [T
(16)
2
+
D
D
K
Th.
Vd(( -ffo)))
+
E
K
The
V
((( I|
2
D.
2m
1=1
k=0
1=1
k=0
Dj
-s(x))
[V
Vi.
2
2D,
K-1
+ 2m
Σ
E
Σ
Tie
(17)
1=1
k=0
15

---

Published as a conference paper at ICLR 2024
Dj
f((x))
[V)(2))
2
K-1 1
- - + V)((a))) -V(aa)-V(aa) |
+
E
The
2Dr
m
k=0
1
K- 1 Σ
R
k
- V - V(((())
2m
k=0
[)(()) - +
(18)
We now expand the terms in (18) and bound each term separately, as follows.
Dt
Bf(x++1 < f(x))=
[x)(2)))
2
K
V(GGER) ||
(i((a)))
Dr
K=0
A1
-V))))
D3
A2
27
K 1
E
1=1
k=0
Aa
K 1
(fil vi(aa))
E
Ni
1=1
k=0
A1
2
m
K
E
(f((zz)) vf(xo))
m
Ag
27. m
K-1
E
Σ Th
Vxxz))
(19)
m
Now, we provide an upper bound for each term in (19). Starting with A1, using (7), (8), and (12), we have
K
m
K
( 2
A1
E
2 E
(20)
Dem
Kİ²
K-1
- x ?
(mi,)) 2 E
Dom
KĪ²
(mi. 2 E Σ n.. 4 (((()t))))) ||
Dpm
K m K-1
Σ
(2Ex))
Σ
ni, (V.((.t))))) V)(ti V
Dpm 0=
16

---

Published as a conference paper at ICLR 2024
2KL² m
K-1
R-1
Thes ( - || + 10| _ mi viic(t)))
E
Σ
Dmm
k=0
8=
2=
2K12
K-1
( R V((())))) - v || + E VAA
)
Dam
2KL²
m K-1
Σ 1(Th.)))
) ( b
(T+ +-
Dam
(=1 k==
2
d" + G2
(21)
Dmm
k=
For A2, using (7) and (8), we have:
K
A2
(nh ) E Vf((o)) {
Dim
Kp
m K
f(()))
(nh,))2
(22)
Dmm
For Az, using (7), (8), and (12), we have:
2K
m
K-1
(.(()t))) - v)(aha)) ||
Ag E
(n
2 E
m
1=1
k=0
2K Lo²
K-1 (
(23)
mb
k=
For As, using (7), (8), and (12), we have:
2
K-1
V(((z)) - ViGCEta)1 ||?
A4
(ni,))
2 B
(24)
m
K=
2K L
m
K-1
xe x/.,0 1p
(ni,,) 2 E
m
k=
2K 33
K
(nftw) 22 E
ni. f (2ct,)) |
2KI3
m
K-1
Σ
( B
ni. (V V/(a V vA(C |
m
Σ K
Whe (V)((tt V /| + E
Σ mia V fa [9
Σ
(z 2
E
m
2=
4KΣ3
k(ri.))
(nnt))² E
v(((tte)) - V.ic /| +E v
4 m K-1
2 (nhE.) )2
Σ
Σ
+22
m
b
1=1 K=
4K2
x² + G2
(nite) )
(25)
m
i=
k=
17

---

Published as a conference paper at ICLR 2024
For As, using (7), (8), and (12), we have:
2
K
( E (f(( V T( ||
As
<
m
2K L
2 m K-1
4 ((t))
Σ (
(26)
m
For A6, using (7), (8), and (12), we have:
2KL
m
Aa
(
(27)
m
Putting all bounds together. Now, by putting (20) back into (19) and rearranging the terms, we obtain the following inequality:
Dr 2
Kp
m
K
m 2KL(1 +p)
(ni. 12
|
(1x))1
Dmm
m
i=1 K=
i=1 k=0
2K2
o²
5
+ G22
K
(Th )2) 2
< f(xt)) - f(x+11)
Dmm
k=
2K
ΣΣ
m K-1
(
4K2i3
G2
m
K
(
(28)
mb
m
k=
Now, it remains to notice that, by definition, ni.,- = Or)), for all i [m], k E (0, I,.., K-1 13, and
t > O. Therefore. D. = O()K). Thus, by dividing the above inequality by 2 we have:
p m K (1
4KL(1+P) + p)
m K-1
(vE.A)
(1i..) [[]
Dim
Dum
1=1 k=0
k=0
0(1)
Ol7K)
2(f(x)) - -f(eet1))
4Kio²
m K-1
(ni.))
Dr
Demb
0=
o +xx)
((TK)
2 8- + G2
m
K-1
(ni )2)
2
Dim
i=1
A=
O()2K2)
6- 8K2 1,3
1/+5 + G2)
m
K
(2)
(29)
Dim
k=
(()2K3)
for p= (1). For the simplicity of exposition, assuming p < 1 and T > (5/p)²- = 440, we obtain the below by averaging (29) for t = 0, 1, T-:
-1) [f((,)) < 0 ((00)- ) + t0(11) (2Ko²) + ) +0((2((2 02
b + + xx22)
T
KT
b
+ +0(13( b
13 02
+ (+))*
(30)
18

---

Published as a conference paper at ICLR 2024
Finally, by choosing = (x+T)) and defining v =
max { - f(0*)) and U2
(o + G² with b = [B] being the batch size, we arrive at the statement in Theorem 1, i.e.
Σ f(m)))
1
+ 0
L2
L
VT
A.2 PROOF OF ALGORITHM FOR CONVEX CASE
Here, we provide extension of the proof of Malitsky & Mishchenko (2020, Theorem 1) for the distributed version. Note that the proof we provide here is not exactly the same version of A-SGD presented in Algorithm I; in particular, we assume no local iterations, i.e., k for all i E [m], and every client participate, i., p = 1, without stochastic gradients.
We start by constructing Lyapunov function for the distributed objective in (1).
Lemma 4. Let fi Rd + R be convex with locally Lipschitz gradients. Then, the sequence {x+} generated by Algorithm I, assuming k and p = with full batch, satisfy the following:
[ 11
1
2
2 m
nt + (f((tt))) -
2m
< 11 - xx"1
+
|
2
11?
- (31)
2m
1=1
The above lemma, which we prove below, constructs a contracting Lyapunov function, which can be seen from the second condition on the step size n/!
Th & V 1+ Oini
= 1 & (1 +1+00 m.).
Proof.
We denote T22
1 -i=1 z.
1
| - =2 -1120
T/ f(()) - x11
(32)
m
||
x
11
n
2
x²,
n V (x())
(33)
We will first bound the second term in (33). Using
Σ
aill
m
Σm aa 12 we have
v (
Σ (at
m
x1+1)
2
m
| o
xtt1l? 11.
(34)
m
To bound || - 2 11,, observe that
||= 21+1 =21lot 2444 11 ||l| 21+1 1|2
on -2/=1> ||= x 11?
-2n/(V - V/i ( x xx)) - ||:; 12-1l?
+
(35)
For the first term in (35), using Sccchy-Scwwarz and Young's inequalities as well as ||= | Ta (IVha)-AALL.. V,l we have:
2nd (V((tt) - V)(()-)), - & 2n (V(()))
V.A(s
||044 ==fll
19

---

Published as a conference paper at ICLR 2024
< || ال=1ll || zill
< | - x x 112 + 2 31 - - -x/1?
2
For the third term in (35), by convexity of [1()), we have
(v[(((( s/+1 - x4) = 2 vh
.-- (1iz) - x/, wi+1 - x²t)
20 - Z V v(((())
< 20in)((((t2_1) - f1(1)))
Together, we have
1
|
11?
blla - 11? 2 - =i11 + -
(36)
Now to bound the last term in (33), we have
-2 22 x², 1 m
m
Σ m f((z()
2/. vf(((())
x
m
m
E=
2 m
-
(37)
m
where in the equality we used the fact that oa = Fe, for 1 Putting (36) and (37) back to (33), we have
m
m
|,844 x=1 112
||,r.
x=1 112
AV
2 2x x² Σ n (V(((z))
|
1=1
-0/-11/111 1/-111 12-111--11 112-1111 111 - + |z
i=1
22(01 82
Th
+2 m
- ((((x"))((xx))))
(38)
Averaging (38) for i = 1..... m. notice that the first term in (38) disappears, as Eg Em
1
21 + (f((tt) - f((x*)) + 1 || - =[]1]]
| - xx 2² +
m
m
||
x²1
- + 1 loo 11
m
Theorem S. Let f. Rd - R be convex with locally Lipschitz gradients. Then, the sequence {x1} generated by Algorithm 1, assuming k = 1 and p 1 with full batch, satisfy the following:
m
2D12
(I((3t) - -(((")))) < 2ît +1'
(39)
m
where
5. 11 + 0))a + + af
S
of (1 0%)
mh- or h++
S: (1 + 03) +
Σ a
Σ n
+ n/
km
and L is a constant. Further, x is any minimum of f((x) for all i E [m]) (39) implies convergence for problem (1).
20

---

Published as a conference paper at ICLR 2024
Proof. We start by telescoping (31). Then, we arrive at
1
2
| x"
x111²
- f((x'))
2m
m
1-1
(nf
02)
T (1=))))))))
(40)
2
< 1|11- x2
+ m Σ ni (f((z)) f((z))) = D. (41)
2m
Observe that (40) is nonnegative by the definition of n/ implying the iterates {æi} are bounded. Now, define the set S == conv((=h, ..,...>>, }), which is bounded as the convex hull of bounded points. Therefore, fi is Lipschitz continuous on S. That is, there exist L, such that
[Vf((z) ff(ill 1 yll,
Vx,y E S.
Also note that (40) is of the form afif( - f((=))] + Σ aflf((a)))))) - fa(x))]
Ei= a (x(() - f((x))] The sum of the coefficients equals:
1-1
Σ at = 11 + B,) + Σinl (in + +1) - n (1+10k+11 = =ni0 + Σ
7k = St (42)
k=
Therefore, by Jensen's inequality, we have
1
D > |z |
| zbll + 2S (f((t)) fc(x*)) > (.((E) - (43)
2m
1=1
which implies f((ātt) fi(x')
<
D' 28
where
a = n/(1 0f )of + Ete at
S
af v (1 + 0%) - nh+ 0/+11
S (1 (1 + 6;) + Σ
nk 1
(44)
k=
k=1
Now, notice that by definition of m/: we have
| x/-1و
m f f((=_)))
Vi [m]
2L,
Also, notice that v/ (44), we have
bn))
= (2 assuming for simplicity that nó 1. Thus, going back to
%
S = Σ mh 0/
20.
4Î
k=1
t
t
1
21.
412
2L
4,.²'
where L = max; L, Finally, we have
1
2 m Σ
(f((()))) fi(x'))
1=1
1 m
Σ
1
1
(f((Et) ff((x)))
>
+
m
2L
412
which implies
2DÎ
1 ((())){((z'))
2 +1
1=1
21

---

Published as a conference paper at ICLR 2024
B ADDITIONAL EXPERIMENTS AND PLOTS
B.1 EASE OF TUNING
In this subsection, we show the effect of using different parameter á in the step size of A-SGD. As mentioned in Section 4, for A-SGD, we append in front of the second condition, and hence the step size becomes
- (. - V
Mi. min
80 17
The experimental results presented in Table 1 were obtained using value of å = 0.1. For y we remind the readers that we did not change from the default value 7 = the original implementa- tion in https:// thub. com/ yma tssky/ ceeepttiveeette www.looaaee/ppttteeeeoee optimizer.py.
In Figure 4, we display the final accuracy achieved when using different values of å from the set {0.01, 0.1, Interestingly, we observe that has very little impact on the final accuracy. This suggests that A-SGD is remarkably robust and does not require much tuning. while consistently achieving higher test accuracy compared to other methods in the majority of cases as shown in Table 1.
CIFAR- 10 trained with ResNet-18
CIFAR-100 crained with ResNet-19
CIFAR-100 trained with ResNet-50
60
65
1est Accuracy
80
55
60
55
50-
60
50
50
45
45
40
40
40
35
35
30
30
0.01
0.1
0.01
0.1
0.0
0.1
paremeeter (A)
SGD parameter
SGD parameter
(B)
(C)
Dirichlet
Dirichlet 0.1
airichle:a=1
Figure 4: Effect of using different in the second condition of the step size of A-SGD.
B.2 EFFECT OF DIFFERENT NUMBER OF LOCAL ITERATIONS
In this subsection, we show the effect of the different number of local epochs, while fixing the experimental setup to be the classification CIFAR-10M dataset, trained with ResNet-50. We show for all three cases of the Dirichlet concentration parameter a E (0.01, 0. The result is shown in Figure 5.
As mentioned in the main text, instead of performing local iterations, we instead perform local epochs E, similarly to (Reddi et al., 2021; Li et al., 2020). All the results in Table use 1, and in Figure 5, we show the results for E' {1, 2, 3}. Note that in terms of E, the considered numbers might be small difference, but in terms actual local gradient steps, they differ significantly.
As can be seen, using higher E' leads to slightly faster convergence in all cases of However, the upeed-up seems marginal compared to the amount of extra computation time requires (i.e., using E = 2 takes roughly twice more total wall-clock time than using E = 1). Put differently, S-SGD performs well with only using a single local epoch per client (E 1). and still can achieve great performance.
22

---

Published as a conference paper at ICLR 2024
Dirichlet Parameter a 0.01
Dirich et Parameter a 0.1
Diricniet Parameter a
0.25
0:30
06
0.5
0.25
0.5
0.5
0.20
0.4
04
0.15
0.
03
0.10
0.
0.
005
0.1
0.1
D00
0.0
0.0
500
1000
1500
2000
500
1000
1500
2000
0
300
1000
1500
2000
Number of Rounds (A)
Number of Rounds (3)
Number of Rounds 6C)
Local epoch 1=
Local epoch
Local epoch =
Figure 5: Effect of the different number of local epochs.
B.3 ADDITIONAL EXPERIMENTS USING DIFFERENT NUMBER OF LOCAL DATA PER CLIENT
Non-iidness
Oplimizer
Dataset / Model
CIFAR-10
CIFAR-100 ResNet-50
Dir(a p)
ResNet-18
SGD
80.7
53.5
SGD (L)
78.8
53.6 (3.9)(3.9)
SGDM
75.0 (5.7)
53.9 (336))
SGDM (4)
66.61(14.1)
53.1
0.1
Adam
79.9
51.1 (66.4)
Adagrad
19.3
44.5 (13.0)
SPS
64.4 64.4
37.2 (20.3)
A-SGD
80.4
47.5
Table 3: Experimental results using different number of local data per client. Performance difference within 0.5% of the best result for each task are shown in bold. Subscripts a) is the performance difference from the best result and is highlighted in pink when the difference is bigger than 2%. The symbol (1) appended to SGD and SGDM indicates step-wise learning rate decay, where the step sizes are divided by 10 after 50%, and another by 10 after 75% of the total rounds.
In this subsection, provide additional experimental results, where the size of the local dataset each client has differs. In particular, instead of distributing 500 samples to each client as done in Section 4 of the main text, we distribute n; E [100, 500] samples for each client i, where n; is the number of samples that client has, and is randomly sampled from the range [100, 500). We keep all other settings the same.
Then, in the server aggregation step (line 13 of Algorithm 1), we compute the weighted average of the local parameters, instead of the plain average, as suggested in (McMahan al., 2017), and used for instance in (Reddi et al., 2021).
We experiment in two settings: CIFAR-I0 dataset classification trained with a ResNer-18 18 (He et al., 2016), with Dirichlet concentration parameter 0.1, which was the setting where the algorithms were sened using grid search; please see details in Section 4. We also provide results for CIFAR- 100 dataset classification trained with ResNet-50, with the same
The result can be found in Table 3. Similarly to the case with same number of data per client, i.e., 500 for all i E [m), we see that A-SGD achieves good performance without any additional tuning. In particular, for CIFAR-I0 10 classification trained with a ResNet-18, A-SGD achieves the second best performance after SGD (without LR decay), with small margin. Interestingly, SGDM (both with and without LR decay) perform much worse than SGD. This contrasts with the same setting in Table where SGDM performed better than SGD, which again complicates how one should tune client optimizers in the FL setting. For CIFAR-100 classification trained with ResNet-50, A-SGD outperforms all other methods with large margin.
23

---

Published as a conference paper at ICLR 2024
B.4 ADDITIONAL EXPERIMENTS USING FEDADAM
As mentioned in Section of the main text, since our work focuses on client adaptivity, our method can seamlessly be combined with server-side adaptivity like FedAdam (Reddi et al., 2021). Here, we provide additional experiments using the FedAdam server optimizer.
A big downside of the adaptive server-side approach like FedAdam is that it adds at least one more hyperparameter to tune: the step size for the global optimizer, set aside the two "momentum" parameters (3, and (2) in the case of Adam. This necessity to tune additional parameters is quite orthogonal to the approach we take, where we try to remove the need for step size tuning.
D.2], 9 different server-side global optimizer step size is grid-scarched, and Reddi et al. (2021) 8] shows the best performing server learning rate differs for each task.
Indeed, as can be seen in Reddi et al. (2021)
That being said, we simply used the default step size and momentum values of Adam (from Pytorch), to showcase that our method can seamlessly combined with FedAdam. Due to time constraints, we tested up to CIFAR-100 trained with a ResNet-18, so that all datasets are covered. The result can be found below.
For three out of four cases, A-SGD equipped with Adam server-side optimizer achieves the best accuracy with large margin. For the last case, CIFAR-100 trained with ResNet- SGD with decaying step size achieves noticeably better accuracy than A-SGD; however, we refer the readers to Table I where for the same setting, A-SGD equipped with simple averaging achieves 61.1% accuracy, without any tuning.
Additional experiments using FedAdam (Reddi et al., 2021)
Non-iddesss
Optimizer
Dataset / Model
MNIST CNN
FMNIST
CIFAR-10 ResNet-18
CIFAR-100
Dir(a p)
CNN
ResNet-18
SGD
97.3
83.7.(2.1)
52.0
46.7u(2.0)
SGD (L)
96.44(1.4)
80.9
49.1
49.2
SGDM
97.5u(0.4
84.61(1.2)
53.71(10.1)
13.3
SGDM (1)
96.41(1.5)
81.8 (4.8)
53.3 (10.5
16.8 ((3.4(
a = 0.1
Adam
96.44(1.5)
81.55(4.3)
27.8
38.31(10.9)
Adagrad
95.7)(2.2)
82.1
[0.44
1.0
SPS
96.6g.1.3)
85.0g(0.8)
21.6
1.6
A-SD
97.9
85.8
63.8
41.9
Table 4: Additional experiments using FedAdam (Reddi et al., 2021).
B.5 ADDITIONAL EXPERIMENTS USING MOON
Here, we provide additional experimental results using MOON (Li et al., 2021), which utilizes model-contrastive learning to handle heterogeneous data. In this sense, the goal of MOON is very similar to that of FedProx (Li et al., 2020); yet, due to the modellcontrastive learning nature, it can incur bigger memory overhead compared to FedProx, as one needs to keep track of the previous model(s) to compute the representation of each local batch from the local model of last round, as well as the global model.
As such, we only ran using MOON for CIFAR-10 classification using a ResNet-18; for the CIFAR- 100 classification using ResNet-50, computing resource out of memory run the same configuration we ran in the main text.
The results summarized in Table 5. As can be seen, even ignoring the fact that MOON utilizes more memory, MOON does not help for most of the optimizers- it actually often results in worse final accuracies than the corresponding results in Table from the main text.
24

---

Published as a conference paper at ICLR 2024
Non-iddness
Optimizer
Dataset Model
CIFAR-10
Dir(a p)
ResNet-18
SGD
78.2
SGD (4)
((.29
SGDM
76.41(6.7)
SGDM (4)
75.5
= 0.1
Adam
82.44(0.6)
Adagrad
81.3g(1.8)
SPS
A-SGD
(3.1.(0.0)
Table 5: Experimental results using MOON loss function. Performance differences within 0.5% of the best result for each task are shown in bold. Subscripts is the performance difference from the best result and is highlighted in pink when the difference is bigger than 2%. The symbol (L) appended SGD and SGDM indicates step-wise learning rate decay, where the step sizes are divided by 10 after 50%, and another by 10 after 75% of the total rounds.
B.6 ADDITIONAL EXPERIMENTS USING FEDPROX
Here, we provide the complete result of Table 2b, conducting the experiment using FedProx loss function for the most heterogeneous case The result can be found below, where A-SGD achieves the best accuracy in three out of five cases (FMNIST trained with a CNN, CIFAR- trained with a ResNet-18, and CIFAR- with ResNet-50); in the remaining two cases (MNIST trained with CNN and CIFAR-100 trained with ResNet-18), A- SGD achieves second-bess best accuracy with close margin to the best.
0.01) for all datasets and architecture considered.
Additional experiments using FedProx loss function (Li et al., 2020)
Non-iidness
Optimizer
Dataset Model
MNIST CNN
FMNIST CNN
CIFAR-10 ResNet-18
CIFAR-100
CIFAR-100
Dir(a p)
ResNet-18
ResNet-50
SGD
95.7.(1.4)
(9.0,(1.4)
(13.8)
30.3
25.2.(0.9)
SGD (L)
97.2
79.3
31.3
29.5
20.2 (00..
SGDM
73.8 L(28.))
72.8
29.3
22.9.(7.6)
23.8 (T.4)
SGDM (L)
81.2
78.0_
25.3
19.8_(10.5)
15.0
o 0.01
Adam
82.3 L(14.8)
65.6
28.1 28.1
24.8 24.8((5.0)
22.6
Adagrad
76.3 ((2.9))
51.2 4)
19.3
[94.5 (14.5)
12.9 (279.
4.1 (28.9)
SPS
85.5
(28.4)))
27.6
21.3
(6.5.
SGG
96.99(0.26)
80.6
33.8
29.74(0.6)
31.0,(0.0)
Table 6: Additional experiments using FedProx loss function.
B.7 ADDITIONAL PLOT WITH STANDARD DEVIATION (3 RANDOM SEEDS)
In this section, we repeat for three independent trials (using three random seeds) for two tasks: CIFAR-10 classification trained with ResNet- 18, and FMNIST classification trained with shallow CNN, both with Dirichlet a 0.1,. The aim is to study the robustness of A-SGD in comparison with other client optimizers.
We remind the readers that CIFAR-10 classification trained with a ResNet- 18 with Dirichlet a = 0.1 where performed grid-search for each client optimizer, and thus each are using the best step size that we tried. Then, for FMNIST classification, the same step sizes from the previous task are used. The result can be found below.
25

---

Published as a conference paper at ICLR 2024
We that A SGG not only achieves the best average test accuracy, but also achieves very small standard deviation. Critically, none of the other client optimizers achieve good performance in both settings, other than the exception of A-SGD.
CIFAR 10 trained with Net. (3 random seeds)
0.8
Accuracy
0.5
0.4
0.2
0.0
500
1250 1750
Number of Rounds
3-500 (ours)
SGD Ielecas)
(desayy
Adagra
SGD
SGOM
Adtim
SPS
FMNIST tra ned with CNN random seecs)
0.8
Accuracy
0.5
0.2
0.3
250
Number of Rounds
3:500 (ours)
SGD Ieleces)
(decess)
Adegrad
SGD
SG/M
Adamn
sPS
FAR-100 trained with ResNet-50 random seeds)
Accuracy
0.4
0.2
0.0
500
1250 Number Rounds
750
1000
1500 1750 200G
3,5G2 lours)
SGD Idecsp)
SCDMideces)
Adegrad
SGD
S6ow
Adam
Figure 6: Additional plots for CIRAR-10 classification trained with 18, FMNIST classifica- tion trained with a CNN, and CIFAR- 100 classification trained with a ResNet-50, repeated for three independent trials. The average is plotted, with shaded indicating the standard deviation.
B.8 STEP SIZE PLOOT FOR A-SGD
Here, we visualize the step size conditions for A-SGD to see how the proposed step size looks like in practice. We only plot the first 300 epochs, as otherwise the plot gets quite messy.
From Figure 7, we can see that both conditions for ni are indeed necessary. The first condition, plotted green, approximates the local smoothness of f., but can get quite oscillatory. The second condition, plotted in blue, effectively restricts the first condition from taking too large values.
26

---

Published as a conference paper at ICLR 2024
Non-tidness
Optimizer
Dataset / Model
Dir(a p)
FMNIST/CNN
CIFAR-100/ReSNCC
Average
Std. Dev.
Average
Std. Dev.
Average
Std. Dev.
SGD
73. ((0.5))
0.0263
83.43
0.0178
49.44u(14.58)
0.0376
SGD (4)
67.79u:16.10
0.1550
77.83
0.0815
40.26
0.1687
SGDM
77.59
0.0142
83.27
0.0186
53.80
0.0149
SGDM (4)
72.11 ((1.77)
0.0611
81.361(3.83)
0.0606
46.77
(17.22))
0.0762
0.1
Adam
83 23.24,10 (3.224((0.55)
0.0126
79.98
0.0228
58.89 18,891. L(5.11)
0.0098
Adagrad
83.86
0.0048
53.06_(92.18)
0.0487
49.93
0.0177
SPS
68.38
0.0192
84.59
0.0154
43.16
0.0072
A-SGD
(0000)
0.0052
85.21 ((0.00)
0.0152
64.02
0.0051
Table CIRAR-1O classification trained with a ResNet-18 and FMNIST classification trained with CNN, and CITAR-100 classification trained with ResNet-50, repeated for three independent trials. The average among three trials and the standard deviation for each client optimizer are reported.
Step size conditions for A SGD
0.5
0.4
0.3
0.2
0.0
50
100
150
200
250
300
Number of Rounds
16,
*adu -
Figure 7: Step size conditions A-SGD In green, we plot the first condition, and in blue, we plot the second condition for n which is plotted in red color taking the minimum of the two conditions.
B.9 ILLUSTRATION OF THE DEGREE OF IETEROGENEITY
As mentioned in the main text, the level of "non-iddees"" is controlled using latent Dirichlet allocation (LDA) applied to the labels. This approach based on (Hsu al., 2019) and involves assigning each client multinomial distribution over the labels. This distribution determines the labels from which the client's training examples are drawn. Specifically, the local training examples for each client are sampled from categorical distribution over N classes, parameterized by a vector q. The vector is drawn from a Dirichlet distribution with a concentration parameter a times prior distribution vector over N classes.
To vary the level of "non-iidness." different values of a are considered: 0.01, 0.1, and 1. A larger corresponds to settings that are closer to i.i.d. scenarios, while smaller values introduce more diversity and non-iidness among the clients, as visualized in Figure 8. As can be seen, for a = 1, each client (each row y-axis) have fairly diverse classes, similarly to the i.i.d case; as decrease a, the number of classes that each client has access to decreases. With & 0.01, there are man clients who only have access single couple classes.
27

---

Published as a conference paper at ICLR 2024
Dirichlet a= 1
Dirichlet = 0.1
Dirichlet a= 0.01
Clients
Clients
Clients
Class distribution
Class distribution
Class distribution
Figure 8: Illustration of the degree of heterogeneity induced by using different concentration pa- rameter a for Dirichlet distribution, for 10 dataset (10 colors) and 100 clients (100 rows y-axis).
28


Reviews:
[{"review_id": "Heb4ZUeHKx", "reviewer": "Reviewer_WDyG", "summary": "Federated learning tasks are sensitive to the learning rate selection in client level. The paper argues it may be beneficial if we could enable different learning rate or learning rate scheduler per client due to the highly heterogeneous nature of FL clients. The paper proposes an auto-tuned learning rate scheduler, that could enable learning rate adapting to each client. The paper theoretically show the convergence of the proposed approach and experimental results demonstrate the effectiveness of $\\Delta$-SGD.", "questions": "My main questions lie in the previous weaknesses section.\n\nAlso an optional question: though it may be true there is no convergence guarantee given to the varying step size across clients, there are various papers that give convergence guarantee to the scenario where clients can have different number of local iterations, which seems to be a bit related. Is there any connection or difference in proving these two scenarios?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Disclaimer: the reviewer is not very familiar with the hyperparameter auto-tuning in centralized computing or in distributed federated computing. Thus, I may not fairly assess the novelty of the technique proposed by this paper. \n\nThe paper is well written and easy to follow. I find the paper enjoyable to read. \n\nThe hyperparameter tuning in client level is tedious and existing literature typically by default set universal learning rate for each client. Therefore, the paper tackles an under-explored problem and proposes a simple and effective approach.", "weaknesses": "- Novelty and technical challenge: Client-level optimization is orthogonal to server-level optimization. Therefore, using auto-tuning, which has been studied extensively in the context of centralized computing, in client level is not a very challenging transferral, as we could directly use any auto-tuner directly in client level and check the performance of FL tasks. \n\nI am wondering whether the auto-tuner, used or partially inspired by any practice in centralized computing. And is there any unique challenge if we simply combine any centralized auto-tuner to FL clients?\n\n- Insufficient comparison with auto-tuner baseline: As I mentioned in the last point, it should not be very challenging to directly deploy auto-tuners to FL clients. However, there is limited comparison to this important baseline, i.e., some representative autotuners developed in centralized computing directly used in FL clients.\n\n- Non-standard Assumptions used in theory: bounded gradient is strong assumption, and a bit contradictory to the heterogeneous setting the paper is motivated from. But it may be understandable to use bounded gradient as many existing literature especially who study adaptive optimizers also use it. Strong growth of dissimilarity is less standard and few papers in FL use it. Not sure whether it has any realistic and intuitive correspondence in FL tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9mVZaRwWae", "reviewer": "Reviewer_WWBM", "summary": "The authors provide a new learning rate schedule in federated learning scenarios. Each client can adjust their own learning rates by their own local gradients. The authors show the convergence analysis of the proposed algorithm in the non-convex case. The experiments show the benefit of the proposed algorithm.", "questions": "Can you provide some detailed verification of the proof in the appendix?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The authors introduce a new learning rate schedule for each local client in the federated learning. Meanwhile, they show the convergence under mild assumptions in the nonconvex case.\n\n2. The experiments show that the proposed algorithm is more robust than other algorithms w.r.t. non-i.i.dness in federated learning.", "weaknesses": "The proof does not seem to be correct to me.\n\n***For the proof of Lemma 2:***\n\n(i)  $\\eta_t^{i}$ is used for updating from $x_t$ to $x_{t+1}$, but in the proof it seems like the $\\eta_t^{i}$ is used for updating from $x_{t-1}$ to $x_t$. If the latter case is necessary for the proof, how can we get $x_t$ and $\\nabla f(x_t)$ without knowing $\\eta_t^{i}$.\n       \n(ii) In the algorithm $\\eta_x^{i}$ is based on the gradient of minibatch estimation, how to change them into the true gradient.\n\n***For the proof in A.1:***\n  \nInequality from (12) to (13) does not seem to be correct. Because $\\eta_{t,k}^i$ depends on $\\tilde{\\nabla} f_i (x_{t,k}^t)$, we need to show that the expected value of $\\eta_{t,k}^i \\tilde{\\nabla} f_i (x_{t,k}^t)$. Since they are not independent, it is trivial to see that $E [\\eta_{t,k}^i \\tilde{\\nabla} f_i (x_{t,k}^t]) = E [\\eta_{t,k}^i E[\\tilde{\\nabla} f_i (x_{t,k}^t)]]$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9zgDCc1ahC", "reviewer": "Reviewer_RQMs", "summary": "The authors proposed using a locality adaptive step size rule for client-side training in Federated Learning frameworks. Their method requires almost no hyperparameter tuning. They showed the superiority of their algorithm by comparing it with other adaptive client-side step-size methods experimentally. They inspired this technique from a similar technique proposed for centralized training. They proved its convergence guarantees for both convex and nonconvex functions.", "questions": "1.\tIn (4), client $i$ seems to use exact gradients to choose the step size. Shouldn’t it be $\\tilde{\\nabla}{f_i(x_t^i)}- \\tilde{\\nabla}{f_i(x_{t-1}^i)}$, or how does client $i$ know true gradient? I guess, it is just a notation thing because there doesn’t seem any problem in Algorithm 1.\n2.\tI couldn’t get the intuition behind setting the initial step $\\eta_{t,0}^i$ to an arbitrary $\\eta_0$ at the beginning of each local training after the first one (line 6 in Algorithm 1). Wouldn’t it be good to use the final step size of each client's latest local training?\n3.\tWhat parameter(s) did you tune for $\\Delta-SGD$ except $\\gamma$?\n4.\t“This is a counter-intuitive behavior, as one would expect to get better accuracy by using a more powerful model.” $\\rightarrow$ Here, mentioning Adagrad, Adam, and SPS is finetuned on ResNet18 can be useful. I guess, it should be one of the reasons. Also, it would highlight your method doesn’t need any additional tuning.\n5.\tIn (B) and (C) of Figure 3, have the parameters been tuned on MNIST+CNN and CIFAR-10+ResNet-18, respectively?\n6.\tYou may enrich the study by reporting how the learning rate changes across time and across clients in $\\Delta-SGD$.  Also, you can compare this with the best learning rate of the other methods. \n7.\tThe two parts of local learning rate calculation (line 9 in Algorithm 1) are clear in how they are used in the theoretical analysis. It may be good to include an experimental report to see in how many iterations the learning rate is updated based on the first part $\\left(\\frac{\\gamma \\lVert x^i_{t, k} - x^i_{t, k-1} \\rVert}{2 \\lVert \\tilde{\\nabla} f_i(x^i_{t, k}) - \\tilde{\\nabla} f_i(x^i_{t, k-1}) \\rVert}\\right)$ or the second part $\\left(\\sqrt{1 + \\theta^i_{t,k-1}} \\eta^i_{t,k-1}\\right)$. It may show that the minimum of those two is useful in practice as well.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1.\tFederated learning frameworks need a careful design to handle the heterogeneity across clients. A local learning rate best for a client may be suboptimal for another. The authors successfully apply a similar technique proposed for centralized ML to the federated problem. Also, they prove the convergence.\n2.\tTheir method allows them to use a milder “L-smooth” assumption than the one commonly used in nonconvex optimization analyses. They use the maximum of “average local smoothness over local steps” across clients and rounds.", "weaknesses": "1.\tWhen you presented the results based on the learning rate tuned on one model/dataset and used for all, the baseline methods were not as successful as yours. One natural question is how good those methods are if we tune for each model/dataset separately. Are all methods comparable? Or, how much the not optimum learning rates of baselines are different from the ones optimal for each dataset? \n2.\tThe bounded gradient assumption can be seen as a strong assumption, and it doesn’t hold for many common functions.  Can those terms be bounded using the other two assumptions, bounded variance and dissimilarity?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZOY4gEuIvE", "reviewer": "Reviewer_Yzoy", "summary": "This paper demonstrates an interesting phenomenon where not properly tuning the local step size of the GD algorithm might negatively impact performance of an FL system. To address this, this paper proposes an adaptive scheme for setting local step sizes called $\\Delta$-SGD, and provides an analysis of its convergence under mild assumptions. Finally, the paper provides a thorough empirical demonstration of the proposed method, comparing it with other optimizers when used by FedAvg.", "questions": "1.The result presentation is a little bit confusing. The small number on the right of Table 1 is the performance difference from the best result, but what is this \"best result\" referring to? Best among \\alpha = 1.0/ 0.1/ 0.01 ?\n\n2. For all methods, the authors performed grid search on CIFAR-10 and apply it on other test scenarios. Why not performing grid search on each individual setting? Is it because of the expensive cost of hyperparameter tuning?\n\n3. If hyperparameter tuning is the bottleneck (and the motivation), how will this method compare to the method FedEx proposed by the paper Federated Hyperparameter Tuning: Challenges, Baselines, and Connections to Weight-Sharing (Khodak et al., 2021)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is clearly written and well motivated.\n\nTheorem 1 and its proof in Appendix A seems to be the main contribution. I have checked and believed that the proofs are sound.\n\nThe empirical results presented are generally promising.", "weaknesses": "1. Novelty:\n- I believe $\\Delta$-SGD is algorithmically similar to the Adaptive SGD method in the paper \"Adaptive Gradient Descent without Descent\" (Malitsky and Mishchenko, 2020). While there are some attempts to cite and discuss this paper, I think the authors should be more upfront about it, as it changes the perceived contribution entirely.\n- To the best of my knowledge, Malitsky and Mishchenko did not deal with the FL setting, so I believe theorem 1 and its proof are novel. However, I think the authors should provide a proof sketch to help readers understand the nuance of the proving strategy. What is hard/non-trivial about applying Adaptive SGD in the FL setting?\n\n2. Significance:\n- Despite the claim that $\\Delta$-SGD is complementary to other FL methods that perform server adaptation this is not demonstrated in the experiments.\n- I think Table 2b should be expanded to become one of the main result (conduct over all datasets, rather than as an ablation study on 2 datasets).\n-  The gain in performance is not significant enough in many cases. The authors should provide some error bar or standard deviation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)