Locally Adaptive Federated Learning
Authors: Sohom Mukherjee
Nicolas Loizou
Sebastian U Stich
Venue: Accepted by TMLR ()

Abstract: Federated learning is a paradigm of distributed machine learning in which multiple clients coordinate with a central server to learn a model, without sharing their own training data. Standard federated optimization methods such as Federated Averaging (FedAvg) ensure balance among the clients by using the same stepsize for local updates on all clients. However, this means that all clients need to respect the global geometry of the function which could yield slow convergence. In this work, we propose locally adaptive federated learning algorithms, that leverage the local geometric information for each client function. We show that such locally adaptive methods with uncoordinated stepsizes across all clients can be particularly efficient in interpolated (overparameterized) settings, and analyze their convergence in the presence of heterogeneous data for convex and strongly convex settings. We validate our theoretical claims by performing illustrative experiments for both i.i.d. non-i.i.d. cases. Our proposed algorithms match the optimization performance of tuned FedAvg in the convex setting, outperform FedAvg as well as state-of-the-art adaptive federated algorithms like FedAMS for non-convex experiments, and come with superior generalization performance.

Published in Transactions on Machine Learning Research (10/2024)
Locally Adaptive Federated Learning
Sohom Mukherjec'
sohom. schon.mukkiericchhni-uuurrrrrr
Siiiss-Mzzmmlianns Universität Würzburg
Nicolas Loizou Johns Hopkins University
nloizou@jhu. edu
Scbastian U. Stich
stich de
CISPA Helmholts Center for Information Security
Reviewed on OpenReview: https: www.r/eviww..ee///e/eeee net/ forum? id= auLdS1 iuKW
Abstract
Federated learning is a paradigm of distributed machine learning in which multiple clients coordinate with central server to learn a model, without sharing their own training data. Standard federated optimization methods such as Federated Averaging (FedAvg) ensure bal- ance among the clients by using the same stepsize for local updates all clients. However, this means that all clients need to respect the global geometry of the function which could yield slow convergence. In this work, we propose locally adaptive federated learning algo- rithms, that leverage the local geometric information for each client function. We show that such locally adaptive methods with uncoordinated stepsizes across all clients can be particu- larly efficient in interpolated (overparameterized) settings, and analyze their convergence in the presence of heterogencous data for convex and strongly convex settings. We validate our theoretical claims by performing illust.rative experiments for both i.i.d. non-i.i.d cases. Our proposed algorithms match the optimization performance of tuned FedAvg in the convex setting, outperform FedAvg as well as state-of-the-art adaptive federated algorithms like FedAMS for non-convex experiments, and come with superior generalization performance.
1 Introduction
Federated Learning (FL) Kairouz et al. 2021 has become popular as collaborative learning paradigm where multiple clients jointly train machine learning model without sharing their local data. Despite the recent success FL, of federated optimization methods like FedAvg (McMahan et al. 2017) still face various challenges in practical scenarios such as not being able to adapt according to the training Synamics-FodAAre using vanilla SGD updates with constant stepsizes maybe unsuitable for heavy- fail stochastic gradient noise distributions, arising frequently in training large-scale models such ViT (Dosovitskiy et al. )021). Such settings benefit from adaptive stepsizes, which use some optimization statistics (o.g., loss history, gradient norm).
In the centralized setting, adaptive methods such as Adam (Kingma & Ba 2014) and AdaGrad Duchi let al.) 2011) have succeeded obtaining superior empirical performance over SGD for various machine learning tasks. However, extending adaptive methods to the federated setting remains a challenging task, and majority of the recently proposed adaptive federated methods such as FedAdam (Reddi et al. 2021 and FedAMS (Wang al., 2022a) consider only server-side adaptivity, ie.,, essentially adaptivity only in the aggregation step. Some methods like Local-AMSGrad (Chen et al. 2020) and Adcal-AdaAlte (Xie et al. 2019h do consider local (client-side) adaptivity, but they perform some form of stepsize aggregation in the communication round, thereby using the same stepsize on all clients.
"oork done the CISPA Helmholtz Center for Information Security

---

Published in Transactions on Machine Learning Research (10/2024)
Using the same stepsize for all clients, that needs to respect the geometry of the global function, can yield sub-optimal convergence. To harness the full power of adaptivity for federated optimization, we argue that it makes sense to use fully locally adaptive stepsizes on each client to capture the local geometric information of each objective function (Wang et al. 2021). thereby leading to faster convergence. However, such a change is non trivial, as federated optiriization with uncoordinated stepsizes on different clients might. not convergence. The analysis of locally adaptive methods necessitates developing new proof techniques extending the existing error-feedback framework (Stich, 2018) for federated optimization algorithuns (that originally works for equal stepsizes) to work for fully un-coordinated local (client) stepsizes. In this work, we provide affirmative answers to the following open questions:
(a) Can local adaplivity for federated optimization be useful (faster convergence) (b) Can design such locally adaptive federated optimization algorithm that provably converges?
To answer (a), we shall see a concrete case Example | along with an illustration in Figure showing locally adaptive stepsizes can substantially speed up convergence. For designing fully locally adaptive method for federated optimization, we need an adaptive stepsize that would be optimal for each client function. Inspired by the Polyak stepsize (Polyak) 1987) which is designed for gradient descent convex functions, Loizou ct al. (2021) recently proposed the stochastic Polyak step-size (SPS) for SGD. SPS comes with strong convergence guarantees, and needs less luning compared to other adaptive methods like Adam and AdaGrad. We propose the FedSPS algorithm by incorporating the SPS stepsize in the local client updates. We obtain exact convergence of our locally adaptive FodSPS when interpolation condition is satisfied (overparameterized case common in deep learning problems), and convergence to a neighbourhood for the general case. Reminiscing the fact that Li et al. (2020b) showed FedAvg needs decaying stepsizes to converge under heterogeneity, extend our method to a decreasing stepsize version FedDecSPS (following ideas from DecSPS (Orvicto et al. 2022 ): that provides exact convergence in practice, for the general non-interpolating setting without the aforementioned small stepsize assumption. Finally, we experimentally observe that the optimization performance of FedSPS is always on par or better than that of tuned FedAvg and FedAMS, and FedDeoSPS is particularly efficient in non-interpolating settings.
Contributions. We summarize our contributions as follows:
We show that local adaptivity can lead to substantially faster convergence. We design the first fully locally adaptive method for federated learning called FedSPS, and prove sublinear and linear convergence the optimum, for convex and strongly convex cases, respectively, under interpolation (Theorem 3 This is in contrast to existing adaptive federated methods such as FedAdam and FedAMS, both of which employ adaptivity only for server aggregation.
For real-world FL scenarios (such as when the interpolation condition is not satisfied due to client helerogeneity); we propose a practically motivated algorithm FedDecSPS that enjoys local adaplivily and exact convergence also in the non-interpolating regime due to decreasing stepsizes.
We empirically verify our theoretical claims by performing relevant illustrative experiments to show that our method requires less tuning compared state-of-the algorithms which need extensive grid scarch. We also obtain competitive performance (both optimization as well as genoralization) of the proposed FedSPS and FedDecSPS compared to tuned FedAvg and FedAMS for the convex and non-convex cases in i.i.d. as well as non-i.i.d. settings.
1.1 Additional related Work
Adaptive gradient methods and SPS. Recently, adaptive stepsize methods that use some optimization statistics have become popular for deep learning applications. Such methods, including Adam (Kingma & Bal 2014) and AdaGrad (Duchi et al.; 2011), work well in practice and also theoretically convergent under different levels of smoothness (Rodomanov et al. 2024). A further adaptive method with sound theoretical guarantces is the Polyak stepsize (Polyak) 1987). which has boen recently extended to the stochastic setting by Loizou et al. (2021) and termed stochastic Polyak stepsize (SPS). Extensions of SPS have been proposed for solving structured non-convex problems (Gower et al. 2021a and in the update rule of stochastic mirror
2

---

Published in Transactions on Machine Learning Research (10/2024)
descent D'Orazie et al. 2021 Further follow-up works have come up with various ways to overcome the
limitations of vanilla SPS, such when optimal stochastic loss values are not. known (Orvieto et. al. 2022:
Gower et al. 2022), or when the interpolation condition does not hold (Orvieto et al. 2022 Gower et al.
2021b Jiang & Stich 2024), as well as a proximal variant for tackling regularization terms (Schaipp et al. 2023
Adaptive federated optimization. Reddi al. (021] provide general framework for adaptive federated optimization (FedOpt), including particular instances such FedAdam and FedYogi, by using the corre- sponding centralized adaptive methods as the server optimizer. Several works followed on the idea of server side adaptivity, some recent ones being CD-Adam (Wang et al. 2022b) and FedAMS (Wang et al. 2022a) Fully locally adaptive stepsizes on the client side have not been explored before, except in one concurrent work Kim et al. 2023). Their proposed method is based on estimator for the inverse local Lipschitz constant from (Malitsky Mishchenko 2019), and analyses only the non-convex setting a strong bounded gradient assumption. Our work is also related to tackling the heterogencity issue in foderated learning, wherein clients differ from each other with respect. to local datasets, or compute capabilities, or both. FedNova (Wang al., 2020) proposes solution to this problem by using locally adjusted stepsizes but suffers from speed-accuracy trade-off. FedLin (Mitra et al. 2021 | and CLIMB (Shen et al. 2022) further improve this line of work. However, none of these existing methods leverage the geometry of the local loss landscape to tackle client heterogeneity. Another line of work studies federated optimization with local solvers (Li et al. 2020a; Jiang et al. 2024a b).
2 Problem setup
In this work, we consider the following sum-structured (cross-silo) federated optimization problem
min xER
f(x)
f((x)
(1)
where the components fe: Rd + R are distributed among local clients and are given in stochastic form f((x) == (e-DD/ {)), where D. denotes the distribution of & over parameter space Ω on client i E [n]. Standard empirical risk minimization is an important special case of this problem, when each D; presents a finite number m, of elements (%,,...,E,,. Then ft can be rewritten as fi(x) m. do not make any restrictive assumptions on the data distributions D;, so our analysis covers the case of heterogeneous (non-i.i.d.) data where D; + Dj, Vi + and the local minima x == arg min fi(x), can be different from the global minimizer of
L Ej= Fx(, gj)- We
3 Locally Adaptive Federated Optimization
In the following, we provide background federated optimization, and the (stochastic) Polyak stepsize. This is followed by Fxample to motivate how local adaptivity with (stochastic) Polyak stepsizes can help to improve convergence especially in the interpolation regime, Finally, we outline our proposed method FedSPS to solve
3.1 Background and Motivation
Federated averaging. A common approach to solving 1 in the distributed setting is FedAvg McMahan et al. 2017) also known as Local SGD (Stich 2018). This involves the clients performing a local step of SGD in each iteration, and the clients communicate with a central server after every iterations their iterates are averaged on the server, and sent. back to all clients. FedAvg corresponds to the special case of Algorithm with constant stepsizes 7/ = %o (Line 4).
PS and SPS. Considering the centralized setting (n = 1) of finite-sum optimization on a single worker minxeRd [i(x) ==
Fi [x we introduce the PS as well as the SPS as below:
3

---

Published in Transactions on Machine Learning Research (10/2024)
Deterministic Polyak stepsize. The convergence analysis of Gradient Descent (GD) for con- vex function fi(x) involves the inequality -x I |x - x=1 27/1 (fi(x:) fi fi())) +
oi [V)[(()))]* the right-hand side of which is minimized by the PS 7 = [+1()))
[(x))-1i
Stochastic Polyak stepsize. We use the notion of SPSmax from the original paper Loizou
ct al.
2021). The SPSmax stepsize for SGD (with single stochastic sample) is given by %
where Fi x Ex, 76 is upper bound the -x (l.)
min
stepsize that controls the size of neighbourhood (7h trades-off adaplivity for accuracy), and c > 0 is a constant scaling factor. Instead of using the optimal function values of each stochastic function as in the original SPS paper (Loizou et al., 2021), we use the lower bound on the function values i S Fi, which is easier to obtain for many practical tasks as shown in (Orvieto et al. 2022.. It should also be noted that using the lower bound on function values lead to the same convergence result as in original SPS using optimal function values, as shown in Theorem of Orvieto et al. (2022 )
10
P.
30
23
3311
10"
30
Figure Illustration for Example] 1 showing local adaptivity can improve convergence. We run SGD with constant, global SPS, and locally adaptive SPS stepsizes (with c = 0.5, 1.0), for functions fi(z) = x², fa(x) = fx,, where stochastic noise was simulated by adding Gaussian noise with mean 0, and standard deviation 10 to the gradients.
Example 1 (Local adaptivity using Polyak stepsizes can improve convergence). For a parameter consider the finite sum optimization problem mingeR [f(x) = Ei Ei with f1(x) =ga, f2(x)=111))) 1 in the interpolation regime. If we solve this problem using mini-batch GD, F(+1 = It - - we are required to choose stepsize: 7 SS2/ 2/T,, where I= - to enable convergence, and therefore Sa(a log !) steps are needed. However, if we solve the same problem using locally adaptive distributed GD of the form a+t1 = zt 2 ((nVii(z)) + 72 (2())) then the complerity can be near-constant. Concretely, for any slepsizes E 1 ; ]. with 71 = J. 15 =1 (which also includes the Polyak slepsizes corresponding to functions fi and f2), the iteration complexity is Oloo }), which can be arbitrarily better than Ma log 2) when a + x0. Note that the observation made in this example can also be extended to the stochastic case of SGD with SPS, as illustrated in Figure 7
3.2 Proposed Method
Motivated by the previous example on the benefit of local adaptivity, we now turn to design such a locally adaptive federated optimization algorithm with provable convergence guarantees. As stated before, we need an adaptive stepsize that is optimal for each client function, and we choose the SPS stepsize for this purpose. In the following, we describe our proposed method FedSPS.
FedSPS. We propose a fully locally (i.c., client-side) adaptive federated optimization algorithm FodSPS (Algorithm [ with asynchronous stepsizes, i.e., the stepsizes are different. across the clients, and also across the local steps for particular client. The FedSPS stepsize for client i and local iteration t will be given by
Ex((, -6
71 min
7h
(2)
V. (x(,
4

---

Published in Transactions on Machine Learning Research (10/2024)
Algorithm 1 FedSPS: Federated averaging with fully locally adaptive stepsizes.
Input: x Xo, Vi [n]]
1: for 0,1, 0,., T-11 do for each client i= 1, n in parallel do
2:
sample E: compute
(x(.
FedSPS: 7/ = min -) 76
4:
b local stochastic Polyak stepsize
<lle)
if
+++11 is multiple of then
6:
xi+1 = Ei= (xi -
> communication round
else
8:
x:+1 = x: vigg
b local step
9: end
10: end for
11: for
where 76 > O are constants as explained before, č is the sample at time t on worker i, Fx {i) is the stochastic loss, gì == (i) is the stochastic gradient, and E S Fr = infgrED, Rx is a lower bound on the minima all functions on worker i. Since the loss functions are non-negative for most practical machine learning tasks, we can use e as discussed before, for running our algorithms. We analyse FedSPS in the strongly convex and convex settings and prove convergence guarantees (Theorem 3
We would like to stress that 7 and are free hyperparameters (in the sense that they do not depend theoretically on any problem-dependent parameters) and we demonstrate that they require minimal tuning through relevant sensitivity analysis in Section 0 This is an advantage over the lcarning rate paramcter of Avd and FedAMS which theoretically depend on L,, and require extensive grid search in practice (Section F.4. The notations used throughout can be extended to the bini-batch setting as described in Appendix
Remark 2 (Alternative design choices). Note that there can be various alternative design choices for incor- poraling SPS in FedAug. We tried some varionts such PedSPS-Normakes using cliend and server side correction to occount for solution bias (Wang et al., 2021) due to asynchronous stepsizes (Appendiz D.1). We also introduce FedSPS- Globel in the Appendia D. that uses aggregation of stepsizes in communication rounds, similar to [Chen el al., 2020; Xie el al.) 2019). However, none of these other design choices provided any practical advantages over our proposed FedSPS.
4 Convergence analysis of FedSPS
4.1 Assumptions on the objective function and noise
Assumption (L-smoothness). Each function F((x, Rd x Ω4 R, i [n] is differentiable for each supp(D:) and there erists constant. L 2 O such that for each x,y Rdd & supp(D,):
(V Exl S L x
Note that Assumption implies L-smoothness of each fx(x) and of fxx). The assumption of each F((x, E being smooth is often used federated and decentralized optimization literature (for e.g., (Koloskova et al. 2020 Assumption la), or (Koloskova et al. 2022 Assumption 3)).
Assumption 2 (μ-convexityy. There exists constant μ such that for each for each i [n], E supp(D,) and for all x,y R':
F.(y F(( (VF( E).y x) + / lly x||2
(4)
For some of our results, we assume μ-strong convexity for parameter μ > 0, or convexity (when μ = 0). Furthermore we assume (as mentioned in the introduction) access stochastic functions F,(x {) on each client with Fe (x,(x,)) Vx((x), FEgDD E( fx(x).
5

---

Published in Transactions on Machine Learning Research (10/2024)
Finite optimal objective difference. For each [n] we denote fi inf f((x). Recall that we defined F == infgeD.xeR* Fx E), and need knowledge of lower bounds, E s F2 for our algorithm. We define the quantity
n of = n
R Σ 71 4
(5)
that will appear in our complexity estimates, and thus we implicitly assume that of < co (finite optimal objective difference). Moreover, of also acts our measure of heterogeneity between clients. This is in line with previous works on federated optimization in non-i.i.d setting, such as Li et al. (020b)) that used P f"= as the hetcrogencity mcasure. We can relate of to the more standard mcasures of function het- erogencity =² 1 [V
(x(x)] and gradient variance o² E Eg [V, (*) [f(()]
in the federated literature Koloskova et al. 2022 Wang et al.) 2022a as shown in the following proposition (proof in Appendix C.2p. For the case of convex functions, it suffices (Kolskooa et al. 2020) to compare with 22 - fi of = - Eg.
VF((x' fx((x)) calculated at the global optimum
x* arg meg f(x). We can observe that of is actually stronger assumption than bounded noise optimum (S., o.), but weaker than uniformly bounded noisc ().
Proposition 1 (Comparison of hetcrogeneity measures). Using the definitions of oî Se, and of as defined above, we have: (a) 3 S 2Lo and (b) of 2L.7.
4.2 Convergence of fully locally adaptive FedSPS
In this section we provide the convergence guarantees of FedSPS on sums of convex (or strongly convex) functions. We do not make any restriction on N6 and thus denote this as the fully locally adaptive setting that is of most. interest to us. The primary theoretical challenge was to extend the error-feedback framework (that originally works for equal stepsizes) (Mania et al. 2017 Stich & Karimireddy 2020) to work for fully un-coordinated local stepsizes, and we do this for the first time in our work. All proofs are provided in Appendix B
Theorem (Convergence of FedSPS). Assume that Assumptions 1 and 2 hold and 2 27², then after T itcrations (T/T communication rounds) of FedSPS (Algorithm it holds (a) Convex case:
a
2 Ta
4,003
[x)
| x*1
(6)
Tn
a
where a= min { ). If > 0, and c 2 4 we have (b) Strongly convex case:
fo
E x x+1 < A(1 - ATT ||x0 -x1
(7)
where A 1 and X, p
Σ 11=1 x/.
The convergence criterion of the first result 6 is non-standard, as involves the average of iterates x; on the left hand side, and not the average X, more commonly used. However, note that every T-th iteration these quantitics are the same, and thus our result implics convergence of 1 (T-1)/ Tn
07).
Moreover, in the interpolation case all & f*
Remark 4 (Minimal need for hyperparanter tuning). The parameler 7 is a user selected inpul parameler to determine the number of local steps, and 7b trades-off adoptivity (potentially faster convergence for large (h) and accuracy (higher for small 76). Moreover, as only depends on the input parameter T and not on properties of the function (e.g. L or μ),, is also free parameter. The algorithm provably converges (up to the indicated accuracy) for any choice of these parameters. The lower bounds E can be set to zero for many

---

Published in Transactions on Machine Learning Research (10/2024)
machine learning problems discussed before. Therefore, we effectively reduce the need for hyperparame- ter luning compored to previous methods like MS whose convergence depended on problem--ependeet parameters.
Comparison with SPS (Loizou et al,, 2021). We will now compare our results to Loizou et al. (2021)
that studied SPS for a single worker (n = 1). First, we note that in the strongly convex case, we almost recover Theorem 3.1 in Loizou et al. (2021). The only differences are that they have A =1 and allow weaker
bounds on the paramcter again recover (Loizou al. 2021 Theorem 3.4) up to constants and the stronger condition on (vs. c
1/2, vs. our 4), but we match other constants. In the convex case, we
in their case).
(Special case I) Exact convergence of FedSPS in interpolation regime: We highlight the linear convergence of FedSPS in the interpolation case (ay 0) in the following corollary.
Corollary 5 (Linear Convergence of FedSPS under Interpolation). Assume interpolation, of = 0 and let the assumptions of Theorem be satisfied with 0. Then
E E - x 11 < A(1 - pa)" 1R0 - -x11
(8)
(Special case II) Exact convergence of FedSPS in the small stepsize regime: Theorem 3 shows convergence of FedSPS to a neighborhood of the solution. Decreasing ?b (smaller than Ze can improve the accuracy, but the error is at least (((f)) even when % + 0. This issue also persistent in the original work on SPSmax (Loizou et al. 2021, Corr. 3.3). However, we remark when the stepsize upper bound Th is chosen extremely small not allowing for adaptivity -then FedSPS becomes identical constant stepsize FedAvg. This is not reflected in Theorem X that cannot recover the exact convergence known for FedAvg. We address this in the next theorem, proving exact convergence of small stepsize FedSPS (equivalent to analysis of FedAvg with of assumption).
Theorem 6 (Convergence of small stepsize FedSPS). Assume that Assumptions and hold and % min 2cL } ), then after T iterations of FedSPS (Algorithm) i it holds (a) Convex case:
( 1/o0 x=1 + Th Lo 1 1811907)
- - - Thh
(9)
[=0
and when 0, (b) Strongly conven case:
x²11 -x11 -1000-111 (1
(1580 pr
Loz + 7o
L 75
E E x=11==00
11-1)
(10)
This theorem shows that by choosing an appropriately small 76 any arbitrary target accuracy 0 can be obtained. We are only aware of Li al. (200bb) that studies FedAvg under similar assumptions := f" - -Eff measuring heterogeneity). However, their analysis additionally required bounded stochastic gradients and their convergence rales are weaker (e.g., not recovering linear convergence under interpolation when 0² 0).
5 Decreasing FedSPS for exact convergence
In the previous section, we have proved that FedSPS converges in the interpolation setting irrespective of the value of the slepsize parameter "T- However, many practical federated learning scenarios such as those involving heterogeneous clients constitute the non-interpolating setting (a, > 0). Here, we need to choose small value of % ensure convergence, trading-off adaptivity for achicving exact convergence. We might recall that Li al. (2020b) proved choosing decaying stepsize is necessary for convergence of FedAvg under heterogeneity. In this section, draw inspiration from the decreasing SPS stepsize DecSPS (Orvieto)
7

---

Published in Transactions on Machine Learning Research (10/2024)
1.5
Tran Len (og)
y-
1.25
Tran Lowe (0g)
C-10.2
0.3
0:73
200
420
-00
300
502
-20
Roand
Rumber Loca Stessir
(a) Effect varying 7o-
(b) Rffect of varying T6-
(c) Effeett of varying
(d) Optimal versus 7.
Figure 2: Sensitivity analysis of FedSPS to hyperparameters for convex logistic regression on the MNIST dataset (i..d.) without client sampling. (a) Comparing the effect of varying 7 on FedSPS and varying Y on FedAvg convergence FedAvg is more sensitive to changes in 7 while FedSPS is insensitive changes in to 78. (b) Effect of varying ?6 on FedSPS stepsize adaptivity adaptivity lost if 7o is chosen too small. (c) Small works well practice (r 5). (d) Optimal versus T, showing that there is no dependence.
et al. 2022), to develop FedDecSPS, that achieves exact convergence in practical non-interpolating scenarios without compromising adaptivity.
FedDeeSPS. In order to obtain exact convergence to arbitrary accuracy (without the small stepsize as- sumption) in the heterogeneous setting with of > O, we propose heuristic decreasing stepsize ver- sion of FedSPS, called FedDeeSPS. The FedDecSPS stepsize for client i and local iteration given by -)
[VE[ G- 1 where (c)) is any non-decreasing positive sequence of real num- bers. We also set c_1 co >-1 = 7h- Experiments involving heterogeneous clients in Section demonstrate the practical convergence benefits of FedDecSPS in non-interpolating scttings.
min of c
TediVS bf
Feckvg FrovMss PepROwN ME
(wot
Faskig yt
Falkig (ners s)
Trah tbgg
Tahh Lame tog)
Feetss
12
FecDactde tran
04
20
-00
200
330
600
5t0
150
202
320
400
530
200
2003
400
300
103
2003
400
507
(a) Convex MNIST (L..d.).
(b) Convex w8a (L..d.).
(c) Convex (non-i.i.d.).
(d) Stepsizee for (c).
(b) w8a dataset (iidd. with client sampling). (c) MNIST dataset (non-i.idd with client sampling). (d) Average stepsize across all clients for FedSPS and FedDeeSPSS corresponding to (c). Performance of FedSPS matches that of FedAvg with best taned local learning rate for the i.i.d. cases, and outperforms in the non-i.i.d. case.
Figure 3: Comparison for convex logistic regression. (a) MNIST dataset (Lidd. without client sampling).
6 Experiments
Experimental setup. For all federated training experiments we have 500 communication rounds (the no. of communication rounds being T/T as per our notation), 5 local steps on each client (r = 5, unless otherwise specified for some ablation experiments), and a batch size of 20 (IB) = 20). We perform experiments in the i.i.d. well as non-i.i.d. settings. Results are reported for both settings without client sampling (10 clients) and with client sampling (10 clients sampled uniformly at random from 100 clients with participation fraction 0.1, and data split among all 100 clients) i.c., = 10 active clients throughout. The i.i.d. experiments involve randomly shuflling the data and equally splitting the data between clients.
For non-i.i.d experiments with the MNIST dataset (LeCun et al. 2010), we assign every client samples from exactly two classes of the datasct, the splits being non-overlapping and balanced with each client having same number of samples (Li et al.] 2020b). For non-i.i.d. experiments with the CETAR-10/CIFARE we the Dirichlet distribution over classes following the proposal in Hsu et al. (009)) The degree of
datasets,
8

---

Published in Transactions on Machine Learning Research (10/2024)
heterogeneity is controlled by a, the concentration paramcter for Dirichlet distribution and report the results for a = 1, = 0.1 and = 0.01. Our code is based on publicly available repositories for SPS and FedAM and will be made available upon acceptance.
FedSPS. The implementation is done according to Algorithm I Since all our experimental settings involve non-negative loss functions, we can use the lower bound & = (Orvieto et al. 2022). ; throughout. In the following, we perform empirical sensitivity analysis for the free hyperparameters 76 and c, concluding that our method is indeed insensitive to changes in these parameters.
We start with benchmarking our method by running some initial convex experiments performing classification of the MNIST (i..d.)) dataset with & logistic regression model, without client sampling, In Figure 2 a), we compare the effect of varying "h E {1,5, 103 on FedSPS, and varying 7 E (0.1, 0.013 on FedAvg. We find that FedAvg is not robust to changing stepsize converging well for stepsize 0.1, but very slow convergence for stepsize 0.01. On the contrary, all instances of FedSPS converge to a neighbourhood of the optimum the size of the neighbourhood being proportional perform an ablation study to understand the effect of varying SPS scaling parameter on the convergence in Figure 2(c). For number of local steps T = 5, we vary from 0.01 to 40 (i.e., of the order of square of r). Unlike what is predicted by our theory, empirically we observe that small e works better and larger leads to slower convergence. Moreover, all values of e E (0.01,0.1 0.1, 1.0] have similarly good convergence, thereby implying our method is robust to this hyperparameter and needs no tuning. We provide additional plots for r E {10, 20, 50, 100} local steps in Appendix
To as suggested by the theory. We now fix 78 1, and
F.2 to confirm that this observation is valid across all values of T, and plot the optimal value of e versus T for each case in Figure dd). Gaining insights from above experiments we fix c 0.5 for all further experiments.
FedDecSPS. We evaluate the performance of FedDecSPS with e, - co <ov/+I Similar to the sensitivity analysis of FedSPS towards e (Figure 2), we performed ablations studies for a fixed value of 76 and varying Oo as well as r. The observation is same as the previons case: the optimal value of oo does not scale according to T as suggested by theory and we fix oo = 0.5 for all experiments. Similarly we fix % = 1, following similar observations as before. We compare the convergence of FedSPS and FedDecSPS for the case of heterogeneous data on clients (i.e., of > 0) in Figure 3 (c) and (d), as well as Figure We observe that our practically motivated FedDeeSPS performs better in such non-interpolating settings, as expected.
AedAvg and FedAMS. We compare the performance of our methods- FedSPS and FedDecSPS against the FedAvg baseline, and the state-of-the-art adaptive federated algorithm FedAMS Wang et al. (2022a). FedAvg and FedAMS need extensive tuning using grid scarch for client learning rate ni server learning rate 1/ as well as max stabilization factor e, and 31, B3- We refer readers to Appendix for details on the grid search performed and the optimal set of hyperparameters.
Convex comparison. For the convex setting of logistic regression on MNIST dataset (ii.d. setting), without client sampling, now compare FedSPS with FedAvg and FedAMS in Figure B(). We see that the convergence of FedSPS matches that of the best tuned FedAug. Note that while the best tuned FedAMS slightly outperforms our method, it requires considerable tuming depicted by the large margin between best and worst learning rate performances. For additional convex experiments in the more practical setting with client sampling, we take the problem of binary classification of LIBSVM Chang & Lin (011I) datasets (w8a, mushrooms, ijenn, phishing, a9a) with logistic regression model in the i.i.d. setting. We report the performance on w8a in Figure 3 Bb), where FedSPS again converges similarly
tuned FedAvg, and better than FedAMS. We defer rest of the LIBSVM dataset plots to Appendix F] In the non-i.i.d. case compare our proposed FedSPS and FedDeeSPS to the FedAvg bascline, adaptive federated methods FedAMS and FedADAM, as well as another state-of-the-an federated method MIME (Karimireddy et al. 2021 In this setting DedDecSPS does better than FedSPS, and our methods outperform the best tuned FedAvg, FedADAM and MIME.
Non-convex comparison. For non-convex experiments, we look at multi-class classification of MNIST dataset using LeNet architecture (Leunn et al.) 1998), CIFAR-10 dataset using ResNet18 architecture (He et al. 2016), and CIFAR-100 dataset using Res Net50 architecture, in the i.i.d. as well as non-i.i.d. setting
1PSS http:///gghhubbbco///ssaaaaaaatt FedAMS http:://ggthaaccm//jnggggghhhhaa

---

Published in Transactions on Machine Learning Research (10/2024)
Table 1: Non-convex experiments with CIFAR-10/CIFAR- dataset and ResNet18, Reeesee0 architecture for differ- ent levels of heterogeneity. FedSPS outperforms other non-adaptive, locally adaptive and globally adaptive methods in the homogeneous setting, and FedDecSPSS performs best in the non-lid settings. We report the mean and standard deviation for runs with three random seeds.
Dataset /Model
Non-iidness (Dirichlet a)
Optimizer
CIFAR-10
CIFAR-100
ResNet18
ResNet50
FedAvg
86.71 + 0.41
56.34 + 0.52
FedAMS
88.66 + 0.87
64.32 + 0.91
FedADAM
87.79 + 0.76
63.91 + 0.76
MIME
87.53 + 0.41
63.78 + 0.45
a =
Local-AdaAlter
86.98 = 0.23
60.34 + 0.27
Local-AMISGra
86.89 + 0.92
61.02 + 0.89
FedSPS
89.92 = 0.13
65.24 = 0.19
FedDecSPS
89.91 + 0.15
65.11 + 0.23
FedAvg
82.68 + 0.51
47.91 + 0.56
FedAMS
84.98 +: 0.89
62.24 + 0.72
FedADAM
83.62 = 0.68
61.67 + 0.67
MTME
83.41 + 0.45
61.39 + 0.34
0.1
Aocal-Ada
82.49 + 0.27
57.72 + 0.23
Local-AMMGral
82.39 + 0.91
58.81 + 0.84
FedSPS
84.74 + 0.16
62.91 + 0.16
FedDecSPS
85.78 + 0.14
63.05 + 0.18
FedAvg
27.71 + 0.39
23.22 + 0.45
FedAMS
31.94 + 0.88
30.61 + 0.85
FedADAM
30.15 + 0.71
28.96 + 0.73
MIME
30.35 = 0.51
27.98 + 0.50
a 0.01
Local-AdaAlter
28.87 + 0.32
26.23 + 0.30
Local-AMSGrad
29.18 + 0.84
26.21 + 0.76
FedSPS
31.15 + 0.19
30.66 + 0.23
FedDecSPS
32.57 + 0.16
31.29 = 0.14
10

---

Published in Transactions on Machine Learning Research (10/2024)
FedAMS FedAug
Test Accuracy
1.00
Train
0.75
FedAMS
0.75
FedAug
FedSPS
100
200
300
400
500
100
200
300
400
500
Round
Round
(a) MNIST (i.id.).
1.0
FedAMS
0.30
FedAvg
FedSPS
0.25
FedDecSPS
lest Accuracy
Train Lows
0.20
0.6
0.15
FedAMS
0.4
FedAAg
FedSPS
0.05
0.2
FedDecSPS
400
100
Round
Round
(b) MNIST (non-i.i.d.).
Figure Non-convex MNIST experiments with client sampling. (a) Non-convex case of LeNet on MNIST dataset (ii.d.). (b) Non-convex of LeNet MNIST dataset (nou-ii.d.). First column represents training loss, second column test accuracy. Convergence of FedSPS is very close to that of FedAvg with the best possible tuned local learning rate. Moreover, FedSPS converges better than FedAMS for the non-convex MNIST case (both i.i.d. and non-i.i.d.), and also offers superior generalization performance than FedAMS.
FedsPS
109
FedAMS
0.8
Fed
Train Loss (log)
lest Accuracy
0.6
10-1
FedSPS
FedAMS
10-2
0.2
FedAvg
200
300
100
300
Round
Round
Figure Non-convex CIFAR-10 (1.id.) experiments with client sampling and ResNet18 architecture. First column represents training loss, second column is test accuracy. FedSPS converges better than FedAvg and FedAMS, and also offers superior generalization performance.
(with client sampling), in Figures and Table 1 For the upper bound on stepsizes, the smoothing
technique for rest of the experiments, as suggested by Loizou et al. (2021) for avoiding sudden fluctuations in the slepsize. For a client E [x]] and iteration 4, the adaplive ileration-depender upper bound is given by %h. = 2/81/mingh..- where [B] is the batch-size, m, is the number data examples on that client and we fix 76.0 = In Figure (MNIST), we find that the convergence of FedSPS very close that of FedAvg with the best possible tuned local learning rate, while outperforming FedAMS. In Figure (CIFAR-10), FedSPS outperforms tuned FedAvg and FedAMS in terms of both training loss and test accuracy. We report the
11

---

Published in Transactions on Machine Learning Research (10/2024)
results in the non-ii.d. setting of CLEAR-10/CIFAR/AA for different levels heterogencity in Table ] and find that FedDecSPS performs the best in heterogeneous settings closely followed by FedSPS.
7 Conclusion
In this paper, we show that locally adaptive federated optimization can lead to faster convergence by harness- ing the geometric information of local objective functions. This is especially beneficial in the interpolating setting, which arises commonly for overparameterized deep learning problems. We propose locally adaptive federated optimization algorithm FedSPS, by incorporating the stochastic Polyak stepsize in local steps, and prove sublinear and linear convergence to neighbourhood for convex and strongly convex cases, respectively, We further extend our method to the decreasing stepsize version FedDecSPS, that enables exact convergence even in practical non-interpolating FL settings without compromising adaptivity. We perform relevant lustrative experiments to show that our proposed method is relatively insensitive to the hyperparameters involved, thereby requiring less tuning compared other state-of- the-art federated algorithms. Moreover, our methods perform good better than tuned FedAvg and FedAMS for convex as well non-convex experiments in i.idd and non-i.idd settings.
Acknowledgments
This work was supported by the Helmholtz Association's Initiative and Networking Fund on the HAICORE@FZJ partition. Sebastian Stich thanks the partial financial support from a Google Rescarch Scholar Award 2023. Nicolas Loizou acknowledges support from CISCO Research.
References
Chih-Chung Chang and Chih-Jen Lin. Libsvm: a library for support vector machines. ACM transactions on intelligent systems and technology (TIST), 2(3):1-27, 2011.
Xiangyi Chen, Xiaoyun Li, and Ping Li. Toward communication efficient adaptive gradient method. Proceedings of the 2020 ACM-IMS on Foundations of Data Science Conference, pp. 119-128, 2020.
Ryan D'Orazio, Nicolas Loizou, Issam Laradji, and Ioannis Mitliagkas. Stochastic mirror descent: Con- vergence analysis and adaptive variants via the mirror stochastic polyak stepsize. arXiv preprint 0rXiv:2110.15415 2021.
Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL http::///ppaaaeeeeeeeeee///rrrrr
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Robert Gower, Othmane Sebbouh, and Nicolas Loizou. Sgd for structured nonconvex functions: Learning rates, minibatching and interpolation. In International Conference Artificial Intelligence and Statistics, pp. 1313-1323. PMLR, 2021a.
Robert M Gower, Aaron Defazio, and Michael Rabbat. Stochastic polyak stepsize with a moving target. arXiv preprint XrXi::2100. 1 1851, 2021b.
Robert M Gower, Mathieu Blondel, Nidham Gazagnadou, and Fabian Pedregosa. Cutting some slack for sgd with adaptive polyak stepsizes. arXie preprint Xrxi::200.. 12328, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEFE conference on computer vision and pattern recognition, pp. 770-77 2016.
12

---

Published in Transactions on Machine Learning Research (10/2024)
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXin: 19990093,,,3 2019.
Xiaowen Jiang and Sebastian U Stich. Adaptive sgd with polyak stepsize and line-search: Robust convergence and variance reduction. Advances in Neural Information Processing Systems, 36, 2024.
Xiaowen Jiang, Anton Rodomanov, and Sebastian Stich. Federated optimization with doubly regularized drift correction. arXiv preprint arXiv:2404.088. 08447, 2024a.
Xiaowen Jiang, Anton Rodomanov, and Sebastian U Stich. Stabilized proximal-point methods for federated optimization. arviv preprint (xXx::247..0088, 2024b.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mchdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trends0 in Machine Learning. 14(1-2):1-210 2021.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Breaking the centralized barrier for cross-device federated learning. In M. Ranzalo, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 28663 28676. Cur- ran Associates, Inc., 2021. URL attpps://proceeddigg..eeerrs fbbbbb/ee7ccct//6s5500997777777-
cc/ wwwppppp.ccpppp.p..pppppp files//ppper2222/11ee
pdf
Junhyung Lyle Kim, Mohammad Taha Toghani, César Uribe, and Anastasios Kyrillidis. Adaptive feder- ated learning with auto-tuned clients. orXiv preprint arXiv:2906. 111223000110,,, 11201, 2023.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arxiv preprint arXiv:1412.680 2014.
Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian U. Stich. A unified theory of decentralized SGD with changing topology and local updales. In 37th International Conference Machine Learning (ICML). PMLR, 2020.
Anastasiia Koloskova, Sebastian U Stich, and Martin Jaggi. Sharper convergence guarantees for asynchronous sgd for distributed and federated learning, Advances in Newral Information Processing Systems, 35:17202 17215, 2022.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to docu- ment recognition. Proceedings of the IEEE, 86(11):2278 2324, 1998.
Yann LeCun, Corinna Cortes, and Chris Burges. Mnist handwritten digit database, 2010.
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, 2020a.
Xiang Li, Wenhao Yang, Shusen Wang, and Zhihua Zhang. Communicaiion-elicieen local decent ralized raceetr sgd
methods. arxiv preprint arXiv:1910.00 09126, 2019.
Xiang Li, Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang. On the convergence of fc- davg on non-iid data. In International Conference on Learning Representations, 2020b. URL https:
Xiaoyu Li and Francesco Orabona. On the convergence of stochastic gradient descent with adaptive stepsizes. In The 22nd international conference on artificial intelligence and statistics, pp. 083-992. PMLR, 2019.
Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacostc-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.
13

---

Published in Transactions on Machine Learning Research (10/2024)
Yura Malitsky and Konstantin Mishchenko. Adaptive gradient descent without descent. arXiv preprint arXiv:1910009522 2019.
Horia Mania, Xinghao Pan, Dimitris Papailiopoulos, Benjamin Recht, Kannan Ramchandran, and Michael Jordan. Perturbed iterate analysis for asynchronous stochastic optimization. STAM Journal on Optimiza- tion, 27(4):2202 2229, 2017.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effiiiell learning of deep networks from decentralized data. In Artificial intelligence and statistics, pp. 1273-1282. PMLR, 2017.
Aritra Mitra, Rayana Jaafar, George J Pappas, and Hamed Hassani. Linear convergence in federated learning: Tackling client heterogeneity and sparse gradients. Advances in Newral Information Processing Systems, 84:14606-14619, 2021.
Yurii Nesterov. Introductory lectures on conver optimization: basic course, volume 87. Springer Science & Business Media, 2003.
Antonio Orvieto, Simon Lacoste-Julien, and Nicolas Loizou. Dynamics of SGD with stochastic polyak stepsizes: Truly adaptive variants and convergence to exact solution. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
URL http::/pppparrriiiee.eeeerrrraae
Boris T Polyak. Introduction to optimization. optimization software. Inc., Publications Division, New York, 1:32, 1987.
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 on Learning Representations, 2021. URL htttee///ppaaoooossssssse/oooord
Anton Rodomanov, Xiaowen Jiang, and Sebastian Stich. Universality of adagrad stepsizes for stochastic optimization: Inexact oracle, acceleration and variance reduction. arXiv preprint arXiv:2406. 06398, 2024.
Fabian Schaipp, Robert M Gower, and Michael Ulbrich. A stochastic proximal polyak step size. arXic preprint trXfi:2301.0490 2023.
Zebang Shen, Juan Cervino, Hamed Hassani, and Alejandro Ribeiro. An agnostic approach to federated learning with class imbalance. In International Conference on Learning Representations, 2022.
Scbastian Stich. Local sgd converges fast and communicates little. arXiv preprint 807.09767. 2018.
Sebastian U. Stich. Unified optimal analysis of the (stochastic) gradient method. 04X004997, 2019. URL https://aaxxvvorr//bs/////
arXiv preprint
Sebastian Stich and Sai P. Karimireddy. The error-feedback framework: Better rates for SGD with delayed gradients and compressed communication. Journal of Machine Learning Research (JMLR), 2020.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistenncy problem in heterogeneous federated optimization. Advances in neural information processing systems, 33: 7611-7623, 2020.
Jianyu Wang, Zheng Xu, Zachary Garrett, Zachary Charles, Luyang Liu, and Gauri Joshi. Local adaptivity in federated learning: Convergence and consistency. arXiv preprint ar Xrxt:22006. 023000 2021.
Yujia Wang, Lu Lin, and Jinghui Chen. Communicationn-lfiiceeen adaptive federated learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learn- ing Research, pp. 11802-2888.. PMLR, 17-23 Jul 2022a. URL https: ittps://proceainng..n www.mmr......11221 wang22o.htm]
14

---

Published in Transactions on Machine Learning Research (10/2024)
Yujia Wang, Lu Lin, and Jinghui Chen. Communiction-coppesseee adaptive gradient method for distributed nonconvex optimization. In International Conference on Arlificial Intelligence and Stalislics, pp. 6292 6320. PMLR, 2022b.
Cong Xie, Oluwasanmi Koyejo, Indranil Gupta, and Haibin Lin. Local adaalter: Communication-eficeet stochastic gradicnt descent with adaptive learning rates. orXiv preprint SE::::911.000900 2019.
15

---

Published in Transactions on Machine Learning Research (10/2024)
Appendix
The Appendix is organized as follows. We bogin by introducing some general definitions and incqualities used throughout the proofs, in Section A Proofs for convergence analysis of FedSPS are provided in Section including convex and strongly convex cases. Section C provides some additional theoretical details for FedSPS. Alternative design choices for FedSPS, namely FodSPS-Normalized and FedSPS-Global are described in Section D We provide some theoretical results for FedSPS in the non-convex setting in Section Finally, additional experiments for FedSPS and RedSPS-Global in convex settings are provided Section
B
Contents of Appendix
A Technical preliminaries
17
General definitions
A.2 General inequalities
17
B Convergence analysis of FedSPS
18
B.1 FedSPS inequalities
18
B.2
Proof for convex case of
B.2.1
Difference Lemmas
18
B.2.2
Proof of Theorem
(general valid for all stepsizes)
................................
20
B.2.3
Proof of Theorem 6 (special case with small step sizes)
................................
22
B.3
Proof lor strongly convex case of PedSPS
23
................................
B.3.1
Proof of Theorem (general case valid for all stepsizes)
23
B.3.2
Proof of Theorem
(special case with small step sizes)
24
C
Additional details for FedSPS
24
Extending FedSPS to the mini-batch setting
................................
24
C.2
Comparison of heterogeneity measures
24
C.2.1 Prool of Proposition
24
D
Alternative design choices for FedSPS
25
D.1
WedSPS-Normalized
25
D.I.I
Algorithm
................................
26
D.1.2
PedSPS-Normalizee experimental details
................................
26
0.2
FedSPS-Cloha
26
D.2.1
Algorithm
26
D.2.2
Choice of aggregation formula for FedSPS-Global
27
D.2.3
FedSPS-Global experimental details
28
E
Extension of FedSPS to non-convex setting
29
16

---

Published in Transactions on Machine Learning Research (10/2024)
E.1
Convergence analysis for non-convex functions
29
Proof of Theorem 18
29
F
Additional experiments
31
Visualization of FedSPS stepsize statistics
31
2 Effect of varying с on convergence
32
Additional convex experiments
Hyperparameters tuning for FedAvg and FedAMS
32
A
Technical preliminaries
Let us present some basic definitions and inequalities used in the proofs throughout the appendix.
A.1 General definitions
Definition 1 (Convexity). The function f IR. R, is conver, if for all x,y Rd
f(x) 2 f(y) + (Vf(y)xx-))
(11)
Definition 2 (I-smooth). The function f R IR, is L-smooth, if there erists constant L such that for all y E Rd
[Vx f(y)) s -1|x-yl1
(12)
or equivalently (Neslerov
2003
Lemana 1.2.3)
f(x) f(y) + (VI(y), y) L [| y12
(13)
Lemma 7 Li & Orabona (0019
Lemma 4). For L-smooth function I2d - R, hoving an optima at
x" E Rd we have for any Rd
[V (T
(14)
Lemma 8 OOvieto et al. the following bound holds
2022
Lemma B.2). For L-smooth and μr- stronggyy convez function R: - R,
1 f(x) - S f(x)- -
(15)
2eL.
cx(][(x)||=
[[(((()) 2cp
where l infx [(x), and e is a lower bound 8 f+.
A.2 General inequalities
Lemma 9. For orbitrary set of vectors {a:) a; E Rdd
1
1
(16)
n
i=1
Lemma 10. For arbilrary set of n vectors {a: R a; E Rdd
a
< n
a
(17)
17

---

Published in Transactions on Machine Learning Research (10/2024)
Lemma 11. For given two vectors a. b R
2 2 b) I v |||| + + [|]l
A4 444 0.
(18)
Lemma 12. For given two vectors a, b E Rd
a a bll 2 (1+ ||ll + (1 (1+x
YA > 0.
(19)
Lemma 13. For any random variable X
E E| =
(20)
B Convergence analysis of FedSPS
B.1 FedSPS inequalities
Lemma 14 (Upper and lower bounds on FodSPS stepsizes). Under the assumption that each F3 is L-ssooth (Assumption 7), the stepsizes of FedSPS follow for all rounds i [T-1] 1]
S 70 Tb.
(21)
where a min { 2 7
Proof. This lemma easily obtained by plugging in the definition FedSPS stepsizes into Lemma FedSPS stepsizes ot. follow the following fundamental inequality
15
Ihgg 12 S Le (F((), - E
(22)
Proof. We observe that it holds from the definition of FedSPS stepsizes
F(x+, { E ccll cc|VF((x]. (i)1?
|ngii|| = 7 min
7b ) [|
[F] <t) fz).
B.2 Proof for convex case of FedSPS
Proof outline and notation. Before we start with the proof, let us recall the notation: below we will frequently use gi denote the random component sampled by worker i itcration t, so gi FF {i) and
the stepsize 7/ min F((x).))-( (F(( 76 } We define the (possibly virtual) average iterate X, Σ x/-
We follow the proof template for FedAvg (Local SGD) Koloskova et al. 2020 Stich 2019 deriving a difference lemma to bound Rg = + - - x x|1 and plugging it into the decrease ||x x* 2 Our
poof introduces differences various points like the difference lemma (55 and the docrease (29..
incorporate local adaptivity using properties of FedSPS stepsizes (Lemma 14 and Lemma 15) and the finite
optimal objective difference assumption (a) oo).
B.2.1 Difference Lemmas
Lemma 16 (Bound difference RG for convex case with any stepsizes). For 2 2r², the variance of the is bounded as
iterates between the clients Rg == 1 -
1x
xi] lll
Re < 2nt
Σ
( 5) ft),
(23)
-=(1-1)-k()
where - 1 k(t) denotes the inder of the last communication round (k 1).
18

---

Published in Transactions on Machine Learning Research (10/2024)
Proof. We use the property that X, xj for every t that is multiple of Therefore, there exists k()) TT-1 1 such that R((-1)-b(1) 0. So we have,
n
Σ
R₁ =
Σ
2
j=(t-1) ()-)(()
The inequality follows by the fact the v is the mean of the terms in the sum. With the incquality Ei=1 ||all| S r Ei= ||all for vectors a, E R and the property of the Polyak Stepsize we deduce:
where v == E 1)= 7/8/-
R
l1s8811
(24)
(1-1))-k(0 k(t)
(x) {X) - -61
(25)
i=l
11-1))
3
x1 E7)
(=(1-1)-k()
where we used the assumption c 27 to obtain the last inequality.
Lemma 17 (Bound on difference R₂ for convex case with small stepsizes). For 76 S 201T the variance of the iterates between the clients R, = E | xill², is bounded as
Σ
ER < 3LT
B[f(x)) - T
Σ of
(26)
-=(1-1)-R(1)
j=(t-1)-A(()
where - 1 k(1) denotes the index of the last commumication round (h S 1).
Proof. Again we derive a bound on the variance on the iterated between the clients. We use the property above. We will now prove that it holds
that 8, = xf for every t that is multiple of r. Define k(t)
R2
<
32
2
Σ
[
(27)
1)- k(t)
12
641
Σ
[F,(2, - (7).
1=
k(0)
and consequently (after taking expectation) and using 76 S Z0LF
ER S 3
Σ
-M
Σ
o7.
j=(t-1)--(1)
(t
Note that if t is multiple of 7, then R2 = 0 and there is nothing to prove. Otherwise note that
R
Σ
X x + VF =
1=1
19

---

Published in Transactions on Machine Learning Research (10/2024)
where E vi F Et) denotes the average of the client updates. With the inequality l|a + b1²² 5
(1 +) ||a|| + 2T |bll for T 2 1, a, b. Rd, we continue:
Rs + 2T
R441
1+
11 1 Fx vll
1
2T
1 1
R4
+
F. x 1I
277
Re
||V x 1
n
1
4
{ (8. - (+)1))
1
Rs
12
1
4 2
(I.² R xXIl +|V +1)1
1+
+
1 + 2
R +
4776 n
Σ
(x +0
where we used 76 S 1. SSiST.
1027. Equation
27 now follows by unrolling, and noting that (1+2)* 5 for all
B.2.2 Proof of Theorem 3 (a) (general case valid for all stepsizes)
Distance to optimality. We now proceed by using the definition of the virtual average: Rc-1 1 + (x) -gg)). à
2
8,41 - x31 < ||x, x11
X
x"
+
I
1=1
i=1
2
2 Σ
= |8, x+1
x
nig/
+
(xt x Ag8
n
iml
x11
2
+
m
++++ n
-+- Σ
n
[x, xi|1 + +1>/gi 1*)
m
n
2
2
x
n///
+
+ R
(28)
71
1=1
20

---

Published in Transactions on Machine Learning Research (10/2024)
Upper bound (valid for arbitrary stepsizes). We now proceed in similar fashion outlined in the
proof of [Loizou et al.
2021 Theorem 3.4). Using the property (22) of the FeSPS stepsize we get:
22
2
2
X++1
x
X
x'
ng!i)
+
M (Fx((
(t) fi] + R.
(29)
11
nc
i=1
[x
x'
1|
2
7
(x + F((x .
MFFF 02)- - -6, Re
71
nC
2
X
x"1 11
[x 4 - - - -xx
+
4 - Gl + Rr
| - x ||
1 ト 7 (x
-2
2 " Σ M .1) [67 + Re: (30)
1=1
We use the assumption >> 2 and simplify further
|x+1- - x [x=-
(1 41 200 + Re,
where we defined of = 1 E [F(('' . A:}- After rearranging:
1 (x) 4 f S X x=1 ||x++1 - -x112= + 200 + Rr.
We now plug in the bound on Rs calculated above in equation
23
(xf. (t) - 07 I & x=112 - -1|x,11 - -x1112 + 2%
Σ
[FFF(() - -77
2n
iml jm k(t)
X - -1|/-11++1111 - xx111 + 20 +
Σ
31F - -11
2nt
iml
We now sum this equation up from t= to T
- 1, and divide by T:
T-1
T
[ < - -
||*;
- x11 - -1x+111
T 1=
T 1=
T-1
27b
Σ of +
[[ -71,
T
2n
by noting that each component in the last term appears at most 7 times in the sum. We can now move the last term to the left:
1 r [x(x + -57 T
+
||x0 x* 12 + 2%b T Σ of
T
2n
1=0
It remains to note 7/ 1/20 a == min
J }b therefore:
n ((((() 21 Σ olF(((( 47)
The attentive reader will observe that any e > 1/2 would be sufficient to show convergence of the function suboptimality (note that Loizon et al. 2021 used e 2 3. but only showed convergence of the iterate distance to optimality).
21

---

Published in Transactions on Machine Learning Research (10/2024)
To summarize:
T-1
1 4) S To
1
x*1
+
2Td
1
Σ ot
(31)
T
(=0
t=0
We now take full expectation to get the claimed result.
B.2.3 Proof of Theorem 6 (a) (special case with small step sizes)
We now give an additional bound that tightens the result when the parameter % is chosen to be a small value, smaller than z As a consequence of this assumption (see Equation (15)), it holds of = m for all and e [n]. The proof now follows a similar template as above, but we can now directly utilize the imposed upper bound % (v.s. the bound on used previously).
Upper bound (valid for small stepsizes). 1 E (x) - -(gg))
Using the definition of the virtual average: Re+1
2 n
xt+1 - S ||x;
x²1
Ag.
+
Ihg
(32)
12
We now observe
(xi= x", Mgg)
-7 (xi- x gl)
- 7/ (x} x", gi)
13
1i (1) [x(x (2 - L
| xill"
-78
a2 (xf - -x*,8)
(33)
L
Fi (t 4))
|| xil1 Fx(x) Fx(x*, Et)
%
xi
[F((2e, Et) Fx(x*,
|| xil??
(34)
2
Therefore
||x+11 x"
X x"
1
2 Σ
M [(E +
Fx(x'x [t)]
L
7e Xt
xi1
+
Σ
n
lhig |
(35)
We now use the observation that %6 = no (by the assumption that % is small we could have made this simplification also earlier), and continue:
2nb
[ xx11 12 2
1
F((x. ())
LLR +
g
n
We now use
19
2 V {)) - (2)( +2)
21.2 X - x (2
and the assumption % <
1oz. to obtain
[x x11 < |x x1
276
[i(( {t) F((xxxxb) + 2y0LR, +
2 n
(X +
n
n
i=1
2%b
4T72 2
< x xx112
[F, { Fx(x*, + 2T +
[x $ -67
n
22

---

Published in Transactions on Machine Learning Research (10/2024)
To simplify, let's take expectation over the randomness in iteration t:
E x - x²1 < | -x f* [(R 4L - - - oo]l
Because f(x)) f* 2 and 76 S 107,
1 we can further simplify:
E X x 2 s x x*12 rolf(xe) - 271+2 4L
(36)
After re-arranging and taking expectation:
7t [xx( l - E x;
- X++1 - x=11² + 2 4L7807
We now plug in the bound on R₂ from Equation
26
7h E[f(x)) M < E x
x=1
2%
Σ
- [+ 8
3r
j=((-I)-A(()
and after summing over t =
to - 1, and dividing by T:
Σ T6 Bi){)) - M
1 T
Σ
| E
x 1 E E||1++11111 -
1-11 -x*1)) +
27
T-1
Σ
- [1 4 + 1.0
37
and consequently:
T Σ
3 T'y
-
E ||x0 x=12 + 12(m + +32r218)
(37)
B.3
Proof for strongly convex case of FedSPS
Here outline the proof in the strongly convex As the proof is closely following the arguments from the convex case (just with the difference of applying p-strong convexity whenever convexity was used), we just highlight the main differences for the readers.
B.3.1 Proof of Theorem 3 (b) (general case valid for all stepsizes)
In the respective proof in the convex case, convexity used in Equation 29 instead, we obtain
If strong convexity
2
2
M <t) - Fx(x", at) + Thr |xx x" x11
T
2
4 [4]++ Rr
nc
< |2,-
x" x1
(x +) - -F(("", .t: + 2 |x x' $11
2 Σ
2:14
X - xX1
T=1
+ nc
M (x), 4 - - R
n
< ||x - x2
[F + F((x*, . Xx ||3
n
2 Σ
(X 4 - +
23

---

Published in Transactions on Machine Learning Research (10/2024)
x+1 and the last inequality used Nip S 1/2 S 1.2 By applying the lower bound the stepsize, that convergence is governed by linearly decreasing term:
where the second inequality used | -x1] S
+ +2 -
2 M(A()) - Fx 5D1 M MIP.( - E7 + 2R.
| x s (1 pa) |xx x²1
while the remaining terms are the same (up to a small change in the constant in front of Re) as in the convex case. The increased constant can be handled by imposing slightly stronger condition on (here used 2 4r², compared to c 2 2r in the convex case). The rest of the proof follows by the same steps, with one additional departure: when summing
the inequalities over the iterations = 0 to T 1 we need to use an appropriate weighing to benefit from the linear decrease. This technique is standard in the analysis (illustrated in Stich (2019) and carried out in the setting a residual Rg in (Stich & Karimireddy 2020 Proof of Thcorem 16) and with a residual R2 distributed sctting in Koloskova ct al. 2020 Proof of Theorem
2, in particular Lemma 15).
B.3.2 Proof of Theorem (b) (special case with small step sizes)
Convexity used in Equation 33 If we j-strong convexity instead, obtain:
-(x+ - 1xxx,,n80 Vegi n'g
- Fi Ex) Fx (1)
2 ||x xill 2 Ft( - Fx(x', (t)
11
2 || x x1(²²
S - (x( - F,(x*, [t] + t/ |x -xxxit xi11
Vél ||x 1=1
(38)
2
The last term is essential to get a linear decrease in the first term. For illustration, Equation (36) now changes
where the last inequality used | - 1 s ||x xi 2 + 2 x x
?
E E x² (1 76/4 || xx1 - 476 4L
(the constant in front of E[Re] increased by a factor of two because of the estimate in 33 and can be controlled by the slightly stronger condition on 76). From there, the proof very standard and follows exactly the template outlined in e.g. (Stich & Karimireddy 2020 Proof of Theorem 16) (Koloskova al. Proof of Theorem in particular Lemma 15).
2020
Additional details for FedSPS
C.1
Extending FedSPS to the mini-batch setting
Note that for the sake of simplicity in notation of the proposed method and algorithms, we used a sin- gle stochastic sample e. However, this can be trivially extended to the mini-batch setting. For a batch Bi the stochastic gradient would become [8 EgeB; (F {), and the stochastic loss would become
|& Egeny (x) We use this mini-batch setting for our practical experiments.
C.2
Comparison of heterogeneity measures
C.2.1 Proof of Proposition 1
Proof. We look at the two cases separately as follows:
Function heterogeneity. We recall the definitions C = 1 EE |V(((*) vx(*))) - and of 1 Ei (fi(*)) -f). The global optimal point is denoted by x", and the local optima by x Note that n [x(x)) 0, Vfi(xt) 0, and make use of these facts in our proof.
24

---

Published in Transactions on Machine Learning Research (10/2024)
Assuming L smoothness of each fi(x,, apply Lemma 7 to each of them with xx to obtain
[v((x*) - / < 2L (Ji(x)) - (T) Vi E [n].
(39)
Using this result we can bound Ç from above as follows (noting that vf(x*) = O, vfi(x;)) 0):
2 - - V5(())1
[Vi((** V)((5))11
39
2L (f((x)) f2)
7
21 (J((x) (()
2L
Gradient variance. We recall the definitions o² = 1 - E. | { - )| and of = 1 E's (f(x') -.). The global optimal point is denoted by x*, local optimum client i by x and the local optima of the functions on worker i by x Note that Vf((x)) = 0, (F = 0 and we make use of these facts in our proof.
n
Assuming L smoothness of each F((x, Et), we can apply Lemma to each function on any client i, with x = x* to obtain
VF((xt)) V 5 ||| s 2L (F((x' - 5(x)..5)) vi E [m]. D..
(40)
Using this result we can bound of from above follows
Eg VV:((" ,5- V)((x))1
20
Egi
VF:(x" s') ,')
12
1 P
Egi | 54) (F et) I1")
72
2L
Eg: Fx(x' (²) - F. (x.L.
7
2L
(Eei [F((x" - -2:1)
2L
Σ (fi(x")- - -7)
n
20
Alternative design choices for FedSPS
D.1 FedSPS-Normalized
Drawing ideas from Wang et al. (021) which employs client and server side normalization to account for any solution bias due to asynchronous local stepsizes (especially for heterogeneous clients), we design FedSPS-
25

---

Published in Transactions on Machine Learning Research (10/2024)
Algorithm 2 FodSPS-Normalized: Fully locally adaptive FedSPS with normalization to account for heterogeneity as suggested in Wang ct al. 2021
Input: x Xb,
[n]
1: for t 0, 1, , - do
2:
for each client i = =, ;n in parallel do
3:
sample L compute gi == VFi(x).
RedSPS-Normalized: 7i = min (,(())
4:
local stochastic Polyak stepsize
clls;
5:
if
1++1 is multiple of then
6:
x241 2
Σ == of (xi vigi)
aggregation incorporating normalizations
='=0
8:
else
xi+1 xi -igi
local step
10:
y' = v¹ +
b client normalization factor
11:
end
12:
end for
13:
Normalized. In the following we provided a detailed description of the algorithm, as well as comparison lo our main proposed method FedSPS.
D.1.1 Algorithm
we calculate the client correction factor of for each client i E [n]] the sum of local stepsizes (since the last aggregation round). During the server aggregation step, the pscudo-gradients then normalized using the client and server normalization factors, to remove any bias due to local heterogeneity.
RedaPS--Nommalized. FedSPPS-Normalized (Algorid (Algorithm 2).
D.1.2 FedSPS-Normalized experimental details
The implementation is done according to Algorithm 2 All remarks we made about the choice of scaling parameter upper bound on stepsize "B. and E in the previous section on FedSPS also remain same here. Figure 0 shows a comparison of FedSPS and FedSPS-Normalized for the convex case of logistic regression on MNIST dataset. For the homogeneous case (a) normalization has no effect this is as expected since the client and server side correction factors are equal and balance out. We expect to observe some changes for the heterogencous sctting (b), but here the FedSPS-Normalized actually performs very slightly worse than FedSPS. We conclude that normalization does not lead to any significant performance gain, and intuitively attribute this to the fact the FedSPS stepsizes have low inter-client variance (Figure 7
D.2 FedSPS-Global
For the sake of comparison with existing locally adaptive FL algorithms such as Local-AMSGrad Chen et al. (2020) and Local-Ada Xie al. (2019) (both of which perform some form of stepsize aggregation on the server), we introduce synchronous stepsize version of our algorithm called
D.2.1 Algorithm
FedSPS-Global. In FedSPS-Global (Algorithm the stepsize is constant across all clients and also for local steps for particular client, for particular round There can be several choices for the aggregation formula for the "global" stepsize- we use simple average of the local SPS stepsizes as given below
Fx(x f - V Ei]
7
min
(41)
26

---

Published in Transactions on Machine Learning Research (10/2024)
FedSPS FedSPS-Nermalize
FedSPS
4x
10-1
FedSPS Normaiized
[lgg]
Train (log).
x
10-2
700
400
100
200
300
400
500
MNIST Lid.
(b) MNIST non-i.idd
Figure 6: Comparison of FedSPS and Normai-Normllized for the convex case of logistic regression on MNIST dataset (Li.d. and non-ii.d.). the homogeneous case (a) normalization has no effect, while for the heterogeneous case (b) normalization slightly deteriorates performance.
Algorithm FedSPS-Global: FedSPS with global stepsize aggregation
Input: 7= = 7o, x) Xu, Vi [n]
1: for t= 0, 1....
do
2:
for each client i == 1, n in parallel do
3:
sample ćl, and compute stochastic gradient gi = VFi(xi,
if a multiple of then
5:
x = Σ (xi - -g:)
> communication round
6:
FedSPS-Global: 7= Ei min
global stepsize
<|
else
8:
x xi -gg
b local step
9: end
end for
11: end for
We provide a theoretical justification of this choice of aggregation formula, and compare it with other possible choices in Appendix D.2.2 As it is apparcnt from 411). need to use stale quantities from the last round, for caleulating the stepsize for the current round. This is justified by the empirical observation that stepsizes do not vary too much between consecutive rounds, and is reasonable assumption that has also been made in previous work on adaptive distributed optimization Xie et al. (2019). Due to the staleness in update of the stepsize in this method, we need to provide a starting stepsize To to the algorithm, and this stepsize is used all clients till the first communication round happens. FedSPS-Global can be thought of more as a heuristic method using cheap stepsize computations, and offering similar empirical performance as FedSPS.
D.2.2 Choice of aggregation formula for FedSPS-Global
We can have various choices for aggregating the local stepsizes to calculate the global stepsize. The first obvions choice would be to perform a simple average of the local stepsizes across all clients. This is given by the aggregation formula 7(0) 1 n
Σn= min [(xx)=+; <|li:|? 76 which is used in our proposed method (Algorithm
1 [((()))-)) |x:
30. Two other plausible choices of aggregation formula could be 70h min
7b and
7 (c) min { 1 Ei (x(0))-)))))
c| Σ", x: 7b/ Among these choices, y00) represents the "correct" " SPS stepsize if we
follow the original definition and replace batches with clients in the distributed setup. In the following Proposition we show theoretically that (b) min {7(a), (()). Experimentally we find the simple averaging of local stepsizes ie., (a) to work best, followed closely by g(a), while the "correct." SPS stepsize 7(()) explodes
27

---

Published in Transactions on Machine Learning Research (10/2024)
in practice. Therefore, we choose y(0) our aggregation formula and this has some added advantages for the associated proofs. We feel that reason behind the good performance of FedSPS-Global is the low inter-client and intra-client variance of the stepsizes explained in Section 6 this is the reason behind why simple averaging of the local stepsizes work.
Proposition 2 (Global SPS stepsize aggregation formula). Using the definitions of the three aggregation formulae for synchronous SPS stepsizes 7(%) 700) and y as defined above in Section D.2.2 have the following inequalities
g(0) z 1b)
7(a) 2 1b)
combining which we get 710) < min {,(4)
Proof. We look at the two cases separately as follows:
1. From Lemma 10 it is casy to observe that
12
Σ g
72
g
which can be rearranged as
lf, E7
(f(
R
E1
1 E gi
Σ
g
The required statement follows trivially from the above inequality.
From Chebyshev's 's inequality we have
4
En
lf(x))
Eg]
1
72
13
From AM-HM inequality we obtain
Di=l
Σ i=1
|| gg
Plugging in this into the above we get
1 Σ
[f( 0j
g
1
Σ
g!
n
The required statement follows trivially from the above inequality.
D.2.3 FedSPS-Global experimental details
The implementation is done according to Algorithm 37 All remarks we made about the choice of scaling parameter e, upper bound on stepsize Tb, and E in the previous section on FedSPS also remain same here. As mentioned before, this method needs an extra hyperparameter Yo, that is the stepsize across all clients until the first communication round. Empiriaally client 0 at iteration 0, works quite well in practice. This is explained by the fact that the SPS stepsizes have low inter-client and intra-client variance (Figure 7 Experimentally we find that FedSPS-GIa Global converges almost identically to the locally adaptive FedSPS (Figure and Figure
found that setting To 70: i.e. the local SPS stepsize for
Note that for any experiment in the Appendix that involves FedSPS-Global refer FedSPS (from the main paper) as FedSPS-Toor the legends of plots, for the sake clarity between the approaches.
28

---

Published in Transactions on Machine Learning Research (10/2024)
E Extension of FedSPS to non-convex setting
We outline here some thcoretical results of extending FedSPS to the non-convex setting. Our non-convex results have the limitation of requiring small stepsizes, but we do not need the additional strong assumption of bounded stochastic gradients used in previous work on adaptive federated optimization.
E.1 Convergence analysis for non-convex functions
We now discuss the convergence of FedSPS on non non-convex functions. Unfortunately, it is required to impose additional assumptions this section. This is mainly due that the Polyak stepsize was designed for convex objectives and additional assumptions are needed in the non-convex setting.
Assumption 3 (Bounded variance). Each function fs: IRd + R. i E [n] has stochastic gradient with bounded local voriance, that is, for all x RR,
E (VF((x,)) V
(42)
Moreover, global variance of the loss function on each client is bounded, that is, for all x E Rd,
VVf(())VV((lle)))
(43)
This assumption is frequently used in the analysis of FedAvg Koloskova al. (2020), as well as adaptive federated methods Reddi et al. (2021); Wang et al. (2022a). The local variance denotes randomness in stochastic gradients of clients, while the global variance represents heterogeneity between clients. Note that
= 0 corresponds to the i.i.d. setting. We further note that Equation 43) corresponds to the variance
Assumption in Loizou et al. (2021) (their Equation (8) with p = 1, 5 = (²) that they also required for the analysis in the non-convex setting.
The following theorem applies to FedSPS and FedAvg with small stepsizes:
Theorem 18 (Convergence of FedSPS for non-convex case). Under Assumptions 1 and if 7b
min 2 2 then after T steps, the iterates of FedSPS and FedAug satisfy
Fu T
La²
+ 78 L + ++{?)
Φ = O
+
(44)
n
where F7 == min E ||V and Fo = f(x0) fr
The proof of this theorem again relies on the assumption that the stepsize is very small, but otherwise follows the template of earlier work Li et al. (2019); Koloskova el al. (2020) and precisely recovers their result. For the sake of completeness we still add proof in the Appendix, but do not claim novelty here. The theorem states that when using constant stepsize, the algorithms reach neighborhood of the solution, where the neighbourhood is dependent on the local variance and global variance ñ, and for the i.i.d. case of Ç 0 the neighbourhood is smaller and less dependent on number of local steps By choosing an appropriately small To. any arbitrary target accuracy can be reached.
A limitation of our result is that it only applies to the small stepsize regime, when the adaptivity governed by the stepsize bound on 76- However, when comparing to other theoretical work on adaptive federated learning algorithms we observe that related work has similar (or stronger) limitations, as e.g. both Wang et al. (2022a) (FedAMS) and Reddi et al. (2021) (FedAdam) require uniform bound the stochastic gradients (we do not need) and also require effective stepsizes smaller than T
1 similar to our case.
E.1.1 Proof of Theorem 18
Small Stepsize. We start by the observation that it must hold oo = 7o for all iterations and clients i. This is due to our strong assumption the small stepsizes.
29

---

Published in Transactions on Machine Learning Research (10/2024)
Bound on R. Similar as in the convex case, we need bound on E[R.]. We note that when deriving the bound in Equation (7 we did not use any convexity assumption, so this bound also holds for arbitrary
smooth functions:
32 R S
Σ i=
Σ
|I(((,,
()))))))))
and consequently (after taking expectation) combined with Assumption 3
E R. S 3278T
Σ
[X +
+
(((-)))-(0))
Σ
+2,
(45)
161.2
j=(t-1))k(()
where the last estimate followed by 76 S 25TI.'
Decrease. We study the (virtual) average x By smoothness and definition of X,+1 we have:
( L
f(x+11) ((x))
(46)
10
The scalar product can be estimated as follows:
ff(())) = (V)())) VF. ()) + + ( -FF
(xt)) (x, 4) |
VF - (F)( 1|
2n
(F))))). (F)
|V/((1))1 2 1.²
n
and after taking expectation:
- 12
gi)
n
We now also use smoothness estimate the last term in 46
E
F { - (xxt
+
-
fo (x:)
72
0² +2
1 n
11 +2
Σ
v(((t)) vxt(
n
n
0²
+2
[xx)) R..
n
By combining these estimates, and using 74 < 1
T , we arrive al:
7h [Vx(x;)] 2| + 2 + 3L
0²
E f(x+1) S fxxx))
21[[f(x))]1
2
+ +2L Rs
2
2
2 L2 + Lo²
< s(x.) - Th [V]((x)))
4
2n.
30

---

Published in Transactions on Machine Learning Research (10/2024)
9
FesAMS (best
0.45
Fediing (best y/
Train Look (leg
103
FodAvg (wanst y) FidSP5
Test Accuracy
Stepsize
0.40
0.90
0.89
0.35
(best
0.88
FedAvg (best y)
Intra-client
FedSPS
Inter-client
400
100
300
400
500
400
500
Training loss and stepsizes for FedSPS (L..d.).
0.7
10-1
Foddiig RodAMS
0.0
Inter-client
0.6
Train Loos (log)
Test Accurncy
0.6
FedbiPs
0.5
Stepsize
FedAvg
FedAMS
0.4
0.4
FedADAM
FedSPS
0.2
FedDecSPS
100
400
500
100
500
(b) Training loss, test accuracy stepsizes FedSPS (non-i.id.)).
Figure 7: Visualization of FedSPS stepsize statisties for convex case of logistic regression on MNIST dalaset (ii.d. and uon-i..d.). Left column represents traiming loss, middle column lest accuracy, right column stepsize plots, "intra-client" refers to selecting any client and plotting the mean and standard deviation of stepsizes for that client across 7 local steps, for a particular communication round. "inter-client" refers plotting the mean and standard deviation of stepsizes across all clients, for particular communication round. We extend (a) beyond the best and worst stepsizes and report the performance for all stepszies in Table
We now re-arrange and plug in the estimate Rg from
45
Th
(x:)]l fxx) f(x++1) +
2he Lo
76
Σ
E B ( +64
4
2n
8T
--(1-1)-R(4)
and after summing from / = to T - and dividing by T:
28. LO2
7h
[x 11
(E f(x)) f(xx1))
2
2n
T-1
+
% 8T
Σ
f|[(f(x))]
implying that
T-1 E
TbL02
|
fxx)) <
(f(Xo) - -f*)
+ (6022222200/774
8T t=0
Tho
2n
F Additional experiments
F.1 Visualization of FedSPS stepsize statistics
Intuitively may seem at first glance that using fully locally adaptive stepsizes can lead poor convergence due to different stepsizes each client. However, have already in theory and well verified
31

---

Published in Transactions on Machine Learning Research (10/2024)
c-po1
C-0.::
0-201
6:6a1
6-8.1
0-633
0-61
10
0-1.3
c>320
ce200
0-300
05200
9050
20
500.
10000
300
320
0.0
530
130
2000
420
552
150
203
200
500
100
200
302
Round
Round
(a) 7 10.
(b) 7 20.
(c) 7 50.
(d) T 100.
Figure Additional experiment the effect of varying convergence of FedSPS for different values of In practice, there no square law relation between and т, and any small value & 0.5 works well.
Table 2: Test accuracy for different values of for convex case of logistic regression on MNIST dataset (1..d.)) extending the results of Figure (a) for multiple values of local stepsizes. The best and worst stepsize values are highlighted in bold. We see that there is substantial gap between the best other stepsizes.
FedAvg
FedAMS
0.0001
76.38
80.75
0.001
78.27
81,52
0.01
79.97
83.05
0.1
91.81
91.23
1
79.32
82.85
in experiments, that this is not the case our fully locally adaptive FedSPS indeed converges. In this remark, try to shed further light into the convergence of FedSPS by looking the stepsize statistics across clients. Figure plots the "intra-client" and "inter-client" stepsize plots for i.i.d. and non-i.i.d experiments. "intra- client" " visualizes the variance in stepsizes for a particular client across different local steps. "inter-client" visualizes the the variance in stepsizes across all clients. We can notice that both these variances are small, explaining the good convergence behaviour of FedSPS.
F.2 Effect of varying e on convergence
We provide additional experiments on the effect of varying the FedSPS scale parameter e on convergence, for different number of local steps T {10, 20, 50, 100} (convex case logistic regression on i.idd. MNIST dataset without client sampling) in Figure Similarly as before, we observe that smaller e results in better convergence, and any {0.01, 0.1, 0.5) works well. Moreover, the effect of varying the convergence is same for all 7 clarifying that in practice there is square law relation between and т, unlike as suggested by the theory.
F.3 Additional convex experiments
Additional convex experiments for the rest of the LIBSVM datasets (ii.d.) have been shown in Figure g The first column represents the training loss and the second column the test accuracy. The third column represents the FedSPS stepsize statistics as described before in Section F.1 We can make similar observations as before the main paper-the the proposed FedSPS methods (without tuning) converge as well as or slightly better than Avd and FedAdam with the best tuned hyperparameters. The stepsize plots again highlight the low inter-client and intra-client stepsize variance.
F.4 Hyperparameters tuning for FedAvg and FedAMS
Here we provide more details on hyperparameters for each datasct and model needed by Fed and FedAMS. We perform grid search for the client learning rate m (0.0001, 0.001, 0.01, 0.1, 1.0), and server learning rate E (0.001,0.01, 0.1, 1.0). Moreover, for FedAMS also choose 31 0.9, 32 0.99, and the max
32

---

Published in Transactions on Machine Learning Research (10/2024)
Table Hyperparameters for FedAvg and FedAMS.
Dataset
Architecture
?]
LIBSVM
Logistic Regression
1.0
1 0.01
Logistic Regression
MNIST
0.1 1 0.001
LeNet.
stabilization factor E {10- 10-4, 1,10-,, 1,104,10-11
1,00,,101 ,100). The grid scarch leads to the following set of
optimal hyperparameters presented in Table 3
33

---

Published in Transactions on Machine Learning Research (10/2024)
FedAgg
1.0
RudSPS-Local
Accurary
0.8
Train L235
0.90
0.2
0.94
0.2
Peddig
0.92
Ped5PS-Local
Intra-client
0.0
0.50
FedSP5-Global
0.2
Intercient
1.00
500
500
100
500
Round
mushrooms train.
mushrooms
(c) mushrooms stepsize.
0.325
FesAMS PeaAlig
1.040
Intra cllett
Inter-dient
0.300
Pec5P5-iocal
1.035
Tain L295
0.275
FooSPs Global
Accurncy
0.91
Stupsze
1.030
0.250
1.025
0.50
FedAMS
1.020
0.000
0.89
FedSP5 Lotal FedSPS-Slobal
1.015
(d) ijenn train.
(e) ijenn test.
(f) ijenn stepsize.
FedAMS
0.95
FedAwg
1.014
client
0.50
dient
Tain Lo50
FudSPS-- Clobal
Accurery
0.85
1.012
0.4
1.010
0.80
0.75
FodAMS
1.000
0.2
0.76
Fedlug
Fed5P5
1.006
Fod5P5 Stobal
100
200
300
400
500
100
200
300
400
500
130
200
300
500
Round
Round
Round
{g) phishing train.
(h) phishing test.
(1) phishing stepsize.
0.50
FecAMS
Feoling
1.0
LaccPS-Lacal PiesPS-cisbal
0.45
Train Lo55
Aenume
Stepsire
0.750
6.725
FodAvg
0.4
PedSPS Local
6.700
FeoSPS-Glebel
0.2
Intercilent
0.30
200
300
500
130
500
Round
a9a train.
(k) a9a test.
(1) a9a stepsize.
Figure 9: Additional convex experiments for the LIBSVM datasets (ii.d.). As mentioned in Section D.2.3 FedSPS
(from the main paper) is referred to FedSPS-Local here in the legends, to distinguish it clearly from FedSPS- Global..
34


Reviews:
[{"review_id": "l0qia8T4me", "reviewer": "Reviewer_tjPx", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "[S1] I believe this paper sets to tackle a very relevant problem such as the client mismatch in Federated Learning. \n\n[S2] Paper is well-written and ideas are presented in an easy-to-follow way. The contributions are clearly stated and the goal of the paper is explicitly stated. \n\n[S3] To the best of my knowledge, the proofs are correct and well explained. \n\n[W1] It seems to me that this work is a minor extension of :\n\nNicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste-Julien. Stochastic polyak stepsize for sgd: An adaptive learning rate for fast convergence. In International Conference on Artificial\nIntelligence and Statistics, pp. 1306–1314. PMLR, 2021.\n\nThis paper presents results that closely resemble the ones presented in Loizou et at. I fail to understand the need for the two papers to exit, the explanation that the authors provide is: \n\n> Comparison with SPS (Loizou et al., 2021). We will now compare our results to Loizou et al. (2021)\nthat studied SPS for a single worker (n = 1). First, we note that in the strongly convex case, we almost\nrecover Theorem 3.1 in Loizou et al. (2021). The only differences are that they have A = 1 and allow weaker\nbounds on the parameter c (c ≥ 1/2, vs. our c > 4), but we match other constants. In the convex case, we\nagain recover (Loizou et al., 2021, Theorem 3.4) up to constants and the stronger condition on c (vs. c > 1\nin their case).\n> \n\nThe authors should further clarify the differences between these two papers. \n\n[W2] The assumption of convexity is a strong assumption. There is no real-world problem with neural networks as the variable to optimize over is convex. Can the authors further explain why these results are relevant in a context where nothing is convex? \n\n[W3] The algorithm is based on the Polyak stepsize, which requires the use of the optimal functional value, which is often not known in practice. In the experiment section, the authors chose a lower bound for this value based on the loss function. But, in more challenging situations, in which the lower bound is unknown, what is the effect of choosing a “bad” lower bound? How can this be circumvented? \n\n[W4] The experiments in the non-i.i.d. case use a dataset in which each client has the same number of samples. Can the authors also run an experiment with different number of samples per client? Same number of samples is a very unrealistic situation. \n\n[W5] There are ‘tricks’ in the experiments that prevent the paper to showcase the advantages of the proposed methodology. For example “we use the smoothing technique for rest of the experiments”. This technique is not explained, what is the impact of using it? \n\n[W6] The authors do not show any Algorithm for FedDecSPS. Similar to [W1], it seems to me that there is little novelty here, at this is just an adaptation of the work:\n\nAntonio Orvieto, Simon Lacoste-Julien, and Nicolas Loizou. Dynamics of SGD with stochastic polyak stepsizes: Truly adaptive variants and convergence to exact solution. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=lUyAaz-iA4u.\n\n[W7] The authors do not include DecSPS as a baseline, nor some variant of SPS that can be locally run. Why is this the case?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes an algorithm to tackle the heterogeneity problem in federated learning. Typically, papers assume agents to be homogeneous, and therefore, in the presence of differences, suboptimal results are rendered. This paper proposes to use agent’s information for locally adapting. In particular, this paper proposes to use adaptive methods with uncoordinated step sizes and show convergence both in convex as well as strongly convex settings. The paper presents results showing the benefits of the proposed method. \nIn all, the paper is well-written and the ideas are clearly presented. My main concerns are regarding the novelty of the paper, the relevance of the theoretical results (convexity), and the experiment section.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "[C1] The figures in the experiments are difficult to interpret (Figure 4A Right, Figure 5A Right, Figure 5B Right to name a few). They are small, and they leave very little space for the regime in which the GNN converged. Can the authors provide figures in which the last iterates are more clearly shown? \n\n[C2] I suggest having one line per question in the top of page 2.\n\n[C3] I believe these works are related and not cited:\n - Mitra, Aritra, et al. \"Linear convergence in federated learning: Tackling client heterogeneity and sparse gradients.\" Advances in Neural Information Processing Systems 34 (2021): 14606-14619.\n - Z Shen, et al. \"An agnostic approach to federated learning with class imbalance\" ICLR 2022.", "broader_impact_concerns": "See Strengths And Weaknesses.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "m64KcHKK5r", "reviewer": "Reviewer_PSuz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The utilization of local geometry with Polyak stepsize is principled, theoretically sound, and leads to clear improvements in empirical scenarios.\n2. The robustness to choice of hyperparameter is a significant advantage over most algorithms in this space and can be very useful in practical applications.\n\nWeaknesses:\n\n1. While the authors claim that $\\sigma_f$ as defined in Section 4.1 can act as a measure of heterogeneity between clients it is unclear to me how it can do so if (as is later stated in Section 6), they use $l_i = 0$ since that will give the same value of $\\sigma_f = f*$ as far as I can tell. I could also not find any mention of an alternate choice of $l_i$ for heterogeneous settings.\n\n2. Along the same lines, it is also not clear to me why FedSPS is better than FedAvg in the interpolating $\\sigma_f = 0$. Going by the explanation in Section 4.1 if $\\sigma_f$ is a measure of heterogeneity then $\\sigma_f = 0$ refers to the homogeneous or i.i.d case in which case the local geometry of the function should not matter, yet Fig 5a suggests that FedSPS converges much faster than FedAvg even in this case. Please clarify why that is the case.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces FedSPS, a locally adaptive federated learning algorithm based on stochastic Polyak stepsize that leverages local geometric information for each client function. The authors also provide a modified version of the algorithm, FedDecSPS which uses a decaying upper bound on the stepsize to provide improved performance in non-interpolating settings. The convergence of FedSPS is theoretically analyzed and empirical results show that FedSPS and FedDecSPS lead to faster convergence than existing adaptive and non-adaptive baselines in both i.i.d and non-i.i.d settings.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. There appears to be a dependence on 'j' in the definition of $\\gamma_t$ in the review of Stochastic Polyak Stepsize on Page 3 (Section 3.1). Should it be $\\gamma_{t}^{j}$?\n\n2. Can you also clarify why using a lower bound on the optimal function value is adequate for convergence in Stochastic Polyak stepsize?\n\n3. The last line of the caption in Fig4 mentions FedSPS-Local and FedSPS-Global but I don't see them mentioned in the legends. Please fix this.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "5ZGCiet3Rb", "reviewer": "Reviewer_7kgh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- Clear proofs for the main claims\n- Relatively well written paper\n- Experiments show improvements upon the baseline\n- Method is robust to hyper parameters (which is important for FL where hyperparameter tuning can be tricky in practice)\n\n### Weaknesses\n- Methodologically, this work is an application of the existing SPS and DecSPS algorithms to the federated setting\n- There are no error-bars reported, so in some cases it is hard to gauge whether the improvement is statistically significant\n- Experiments in the non-convex setting are relatively toy and some important baselines are missing", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes FedSPS and FedDecSPS, two algorithms for federated optimization. FedSPS and FedDecSPS are both based on prior works, i.e., SPS and DecSPS, which allow for automatic tuning of the learning rate during training. The main novelty compared to other methods for federated optimization, such as FedAdam, is that the authors perform learning rate adaptivity (via SPS) at the client level, instead of the server. This allows for a learning rate that can adapt to the local geometry of the loss and can lead to better convergence. The authors prove that FedSPS converges at (the vicinity of) the optimum for (non-)i.i.d. data in the convex case, whereas for the non-convex case they prove convergence to a stationary point (up to a user-defined accuracy threshold). Experimentally, the authors show that both FedSPS and FedDecSPS are relatively insensitive to hyper parameters and they can achieve comparable or better performance against tuned baseline methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I am on the positive side for this work; while the idea feels a bit incremental the results are interesting and the proofs are presented clearly. I only have a few questions / requests:\n\n- It would be beneficial if the authors show the results FedAvg and FedAMS obtain with each hyperparameter, or at least the average performance. While it does seem that there can be a big gap between the worst and best hyperparameter, there might still be a reasonable range of hyperparameters with similar performance. \n- The authors discuss FedSPS-Local and FedSPS-Global in the experiment section / caption of figure 4, but these are not introduced and are not present in the figure. \n- The authors mention that there are some methods that do client side adaptivity (e.g., Local-AMSGrad and Local-AdaAlter) but they do not compare FedSPS and FedDecSPS with any of them. Adding comparisons with them would help in better understanding the importance of FedSPS and FedDecSPS relative to prior art. \n- Some notion of statistical significance / error bars would be appreciated for the results.\n- More challenging datasets (e.g., CIFAR 100, TinyImagenet, StackOverflow) and architectures (e.g., transformers) would be appreciated and further strengthen this work.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper considers the setting when the servers in Federated learning are hetererogenous and their local functions have different geometry. Thus they propose running Fed Avg but with different learning rates on the servers. The paper identifies the benefits of this approach over the standard homogenous setting. The reviewers found the paper clear and of interest to the community. One reviewer was concerned with the novelty aspect of this paper but that's not a necessity for the paper for TMLR. I recommend accept.
Decision: Accept (poster)