On Convergence of FedProx: Local Dissimilarity Invariant Bounds, Non-smoothness and Beyond
Authors: Xiaotong Yuan
Ping Li
Venue: NeurIPS 2022 Accept (2022)
Keywords: Federated learning
FedProx
Minibatch stochastic proximal point methods
Uniform stability
Non-convex optimization
Non-smooth optimization

Abstract: The \FedProx~algorithm is a simple yet powerful distributed proximal point optimization method widely used for federated learning (FL) over heterogeneous data. Despite its popularity and remarkable success witnessed in practice, the theoretical understanding of FedProx is largely underinvestigated: the appealing convergence behavior of \FedProx~is so far characterized under certain non-standard and unrealistic dissimilarity assumptions of local functions, and the results are limited to smooth optimization problems. In order to remedy these deficiencies, we develop a novel local dissimilarity invariant convergence theory for \FedProx~and its minibatch stochastic extension through the lens of algorithmic stability. As a result, we contribute to derive several new and deeper insights into \FedProx~for non-convex federated optimization including: 1) convergence guarantees invariant to certain stringent local dissimilarity conditions; 2) convergence guarantees for non-smooth FL problems; and 3) linear speedup with respect to size of minibatch and number of sampled devices. Our theory for the first time reveals that local dissimilarity and smoothness are not must-have for \FedProx~to get favorable complexity bounds.

On Convergence of FedProx: Local Dissimilarity Invariant Bounds, Non-smoothness and Beyond
Xiao-Tong Yuan
Ping Li Linkedln Ads 700 Bellevue Way NE, Bellevue, WA, USA pinlielinkedin.co
Nanjing Univ. of Information Sci. & Tech.
219 Ningliu Rd, Nanjing, China
Abstract
The FedProx algorithm is simple yet powerful distributed proximal point opti- mization method widely used for federated learning (FL) over heterogeneous data, Despite its popularity and remarkable success witnessed in practice, the theoretical understanding of FedProx is largely underinvestigated: the appealing convergence behavior of FedProx is far characterized under certain non-standard and unre- alistic dissimilarity assumptions of local functions, and the results are limited to develop
smooth optimization problems. In order remedy these deficiencies, a novel local dissimilarity invariant convergence theory for FedProx and its mini- batch stochastic extension through the lens of algorithmic stability. As a result, we contribute to derive several new and deeper insights into FedProx for non-convex federated optimization including: 1) convergence guarantees invariant to certain stringent local dissimilarity conditions; 2) convergence guarantees for non-smooth FL problems; and 3) linear speedup with respect size of minibatch and number of sampled devices. Our theory for the first time reveals that local dissimilarity and smoothness are not must-have for FedProx to get favorable complexity bounds.
1 Introduction
Federated Learning (FL) has recently emerged as promising paradigm for communication-effiienn distributed learning on remote devices, such as smartphones, internet of things, or agents (Konečný et al., 2016; Yang et al., 2019). The goal of FL is to collaboratively train shared model that works favorably for all the local data but without requiring the learners to transmit raw data across the network. The principle of optimizing global model while keeping data localized can be beneficial for both computational efficiency and data privacy (Bhowmick et al., 2018). While resembling the classic distributed learning regimes, there large statistical heterogeneity of local data mainly due to the non-iid manner of data generalization and collection across the devices (Hard et al., 2020); and 2) partial participation of devices in the network mainly due to the massive number of devices. These fundamental challenges make FL highly demanding tackle, both in terms of optimization algorithm design and in terms of theoretical understanding of convergence behavior (Li et al., 2020a).
most distinct features associated with FL: I)
FL is most conventionally formulated the following problem of global population risk minimization averaged over a set of M devices:
M Σ (r((m)() 7PLm
E Exomr (emm: 3(m)))
min R(w)=
(1)
M
where (m() is the local population risk on device m, ((m) RP x mm) + R+ is a non-negative
loss function whose value (w; ((")) measures the loss over a random data point Z (m)) E (m))
with parameter w. (mm) represents an underlying random data distribution over ( Since the data distribution is typically unknown, the following empirical risk minimization (ERM) version of (1) is
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

often considered alternatively:
M
Nm
min Rerm (w) MERP
1 M
Rim) crm
(w)
=
e
(())
(w: 24 (())
(2)
Nm
m=1
1=1
where Rarn is the local empirical risk over the training sample ((m) {2] (m) Nm device The
sample size Nn may vary significantly across devices, which can be regarded as another source of data heterogeneity. Federated optimization algorithms for solving (1) or (2) have attracted significant research interest from both academia and industry, with a rich body of efficient solutions developed that flexibly adapt the communication-compucee Several popularly used FL algorithms for this setting include FedAvg (McMahan et al., 2017), FedProx (Li et al., 2020b), SCAFFOLD (Karimireddy et al., 2020), and FedPD (Zhang al., 2020), to name a few. A consensus among these methods on communication--ffficcii implementation is trying extensively update the local models (e.g., with plenty epochs of local optimization) over subsets of devices so as to quickly find an optimal global model using a minimal number of inter-device communication rounds for model aggregation.
tradeoffss and data/system heterogeneity.
In this paper, we revisit the FedProx algorithm which is one of the most prominent frameworks for heterogeneous federated optimization. Reasons for the interests of FedProx include implementation simplicity, low communication cost, promise in dealing with data heterogeneity and tolerance to partial participation of devices (Li et al., 2020b). We analyze its convergence behavior, expose problems, and propose alternatives more suitable for scaling up and generalization. We contribute to derive several new and deeper theoretical insights into the algorithm from a novel perspective of algorithmic stability and generalization theory.
1.1 Review of FedProx
For solving FL problems in the presence of data heterogencity, methods such FedAvg based on local stochastic gradient descent (SGD) can fail to converge in practice when the selected devices perform too many local updates (Li et al., 2020b). To mitigate this issue, FedProx (Li et al., 2020b) was recently proposed for solving the empirical FL problem (2) using the (inexact) proximal point update for local optimization. The benefits of FedProx include: 1) it provides more stable local updates by explicitly enforcing the local optimization in the vicinity of the global model to date; 2) the method comes with convergence guarantees for both convex and non-convex functions, even under partial participation and very dissimilar amounts of local updates (Li et al., 2020a). More specifically, each time instance I, FedProx uniformly randomly selects subset T, S [M] of devices and introduces for each device 4 the following proximal point ERM sub- for local update around the previous global model w;-1i
arg min { Q (w; w,-1) == (E)) (w) + erm
Wi
|| w
(3)
WEBP
21//
where nh 0 is the learning rate that controls the impact of the proximal term. Then the global model is updated by uniformly aggregating those local updates from 4 as
1 T.
(9) SEA
Wt
In the extreme case of allowing V/ -> +00 in (3), FedProx reduces to the regime of FedAvg if using SGD for local optimization. Since its inception, FedProx and its variants have received significant interests in research (Pathak and Wainwright, 2020; Nguyen et al., 2020; Li et al., 2019) and become an algorithm of choice in application areas such as automatous driving (Donevski al., 2021) and computer vision (He et al., 2021). Theoretically, FedProx comes with convergence guarantees under the following bounded local gradient dissimilarity assumption that captures the statistical heterogeneity of local objectives across the network:
Definition 1 ((B, H)-LGD). We say the local functions R(m) have (B, H -local gradient dissimi- larity (LGD) if the following holds for all w E IRP:
M
Σ MA
R mm)
(w))
B B
M

---

The above definition naturally extends to the local empirical risks { {Rarn) m==1- Specially in the homogenous setting where R! = R. Vm E [M], we have B = 1 and H = O. Under (B, 0)-LG and some regularization conditions on the modulus B, it was shown that FedProx for non-convex problems requires T O (+) rounds of inter-device communication to reach an e-stationary solution in the sense of + Ei= (ezr(( < c (Li et al., 2020b). Similar guarantees have also been established for variant of FedProx with non-uniform model aggregation (Nguyen et al., 2020).
Open issues and motivation. In spite of the remarkable success achieved by FedProx and its variants, there are still a number of important theoretical issues regarding the unrealistic assumptions, restrictive problem regimes and expensive local oracle cost that remain open for exploration, as specified below.
Local dissimilarity condition. The appealing convergence behavior of FedProx is so far charac- terized under a key but non-standard (B, 00 condition (cf. Definition 1) with B > 0. Such a condition is obviously unrealistic in practice: it essentially requires the local objectives share the same stationary point as the global objective since Rerm (w) = 0 implies ||V R!m) (w)ll = 0 for
all m [M]. However, if the optima of Rarn are exactly (or even approximately) the same, there would be little point in distributing data across devices for federated learning. It is thus desirable to understand the convergence behavior of FedProw for heterogeneous FL without imposing stringent local dissimilarity conditions like (B, D)-LG.
Non-smooth optimization. The existing convergence guarantees of FedProx are only available for FL with smooth losses. More often than not, however, FL applications involve non-smooth objectives due to the popularity of non-smooth losses (e.g.. hinge loss and absolute loss) in machine learning, and training deep neural networks with non-smooth activation like ReLU. Therefore, it is desirable to understand the convergence behavior of FedProw in non-smooth problem regimes.
Local oracle complexity. Unlike the (stochastic) first-order oracles such as SGD used by FedAvg. the proximal point oracle (3) for local update is by itself a full-batch ERM problem which tends to be expensive to solve even approximately per-iteration. Plus, due to the potentially imbalanced data distribution over devices, the computational overload of the proximal point oracle could vary significantly across the network. Therefore, it is important to investigate whether using minibatch stochastic approximation to the proximal point oracle (3) can provably improve the computational efficiency of FedProa
Last but not least, existing convergence analysis of FedProx mainly focuses on the empirical FL problem (2). The optimality in terms of the population FL problem (1) is not yet clear for FedProx. to lay
The primary goal of this work is remedy these theoretical issues simultaneously, more solid theoretical foundation for the popularly applied FedProx algorithm.
1.2 Our Contributions
In this paper, make progress towards understanding the convergence behavior of FedProx for non-convex heterogenous FL under weaker yet more realistic conditions. The main results are a set of local dissimilarity invariant bounds for smooth non-smooth problems.
Main results for the vanilla FedProx. As a starting point to address the restrictiveness of local dissimilarity assumption, we provide novel convergence analysis for the vanilla FedProx algorithm invariant to the (B, O)- condition. For smooth and non-convex optimization problems, our result in Theorem 1 shows that the rate of convergence to a stationary point is upper bounded by
T E
V (w=))l²] S max
1 T2/3
1 VTI
(4)
where I is the number devices randomly selected for local update at each iteration. If all the devices participate in the local updates for every round, i.c. I = [M]], the rate of convergence can be improved to O( (2))). ). For T < 13, the rate in (4) is dominated by O ) which gives the communication complexity 2// to achieve an c-stationary solution. On the other hand when T 2 13, the rate is T Compared to the already known O complexity bound of FedProx under the unrealistic (B, O)-L.GD condition (Li et al., 2020b), our rate in (4) is slower but it holds without needing impose stringent regularity conditions
dominated by O (TT) which gives the communication complexity 1
3

---

on the dissimilarity of local functions, and it reveals the benefit of device minibatch sampling for accelerating convergence. Further for non-smooth and weakly convex problems, we establish in Theorem 2 the following rate of convergence regarding proper />Morrau--envelopes of Remn
1 - E ( (wx) (w))) S VT t=1
(5)
T
The bound is not dependent on the number of selected devices in each round. In the case of I (1), the bounds in (4) and (5) are comparable, which indicates that smoothness is not must-have for FedProx to get sharper convergence bound especially with low participation ratio. On the other end when / = ()), the bound (5) for non-smooth problems is slower than the bound (4) for smooth functions in large-scale networks.
Main results for minibatch stochastic FedProx. Then as the chief contribution of the present work, we propose a minibatch stochastic extension of FedProx along with its population optimization performance analysis from a novel perspective of algorithmic stability theory. Inspired by the recent success of minibatch stochastic proximal point methods (MSPP) (Li et al., 2014; Wang et al., 2017; Asi et al., 2020: Deng and Gao, 2021), we propose to implement FedProx using MSPP as the local update oracle. The resulting method, which is referred to as FedMSPP, is expected to attain improved trade-off off between computation, communication and memory efficiency for large-scale FL. In the case of imbalanced data distribution, minibatching is also beneficial for making the local computation more balanced across the devices. Based on some extended uniform stability arguments for gradients, we show in Theorem 3 the following rate of convergence for FedMSPP in terms of population optimality, which is also invariant to the (B, 0)-LG condition:
T Σ E [[VRR((
S max
(6)
T2/3 VTbl
where is the minibatch size of local update. For empirical FL, FL. identical bound holds under sampling according to empirical distribution. For T < (bZ)², the rate in (6) is dominated by O (T773)) which gives the communication complexity ,3/2 and it matches that of the vanilla FedProx. For sufficiently large 2 (bT)², the rate is dominated by O ((7T)) which gives the communication complexity nте. This shows that local minibatching and device sampling are both beneficial for linearly speeding up communication. Further, when applied to non-smooth problems, we show in Theorem that FedMSPP converges at the following rate with respect to proper p-Moreau-envelopes of R:
T-1
[ (wx) (1) x VT
T 1=0
which is comparable to that of (6) when = (() and I 0(1), but without witnessing the effect of linear speedup with respect to b and T.
Comparison with prior results. In Table 1, we summarize our communication complexity bounds for FedProx (FedMSPP) and compare them with several related heterogeneous FL algorithms in terms of the conditions on local dissimilarity, applicability to non-smooth problems and tolerance to partial participation. A few observations our ( bounds are not dependent on the (B, H LLLG type conditions, and they are comparable to those of SCAFFOLD and FCO (for convex problems) which are also invariant to local dissimilarity conditions. Second, with regard to the applicability to non-smooth optimization, our convergence guarantees in Theorem 2 and Theorem 4 are established for non-smooth and weakly convex functions. While FCD is the only one in the other considered algorithms that can be applied to non-smooth problems, it customized for federated convex composite optimization with potentially non-smooth regularizers (Yuan et al., 2021). Third, in terms of tolerance to partial participation, all of our results are robust device sampling, and the O( ()==z) bound in Theorem 3 for FedMSPP is comparable to the best known results under partial participation as achieved by FedAvg and SCAFFOLD. If assuming that all the devices participate in local update for each communication round and using momentum acceleration techniques, substantially faster O bounds are possible for STEM and FedPD, while the
in order. First, regarding the local dissimilarity condition, all of
O GATZ bounds can be achieved by FedAvg (Khanduri et al., 2021). To summarize the comparison, our (B, H)-LGD invariant convergence bounds for FedProx (FedMSPP) are comparable to the best- known rates in the identical setting, while covering the generic non-smooth and non-convex cases which to our knowledge has not been provably possible for other FL algorithms.

---

Method
Work
Commun. Complex.
LD Condition
NS
PP
FedProx
(Li al., 2020b)
(B, (B,0)-LG
/
Theorem (ours)
7
Te
3/2
Theorem 2 (ours)
7
FedMSPP
Theorem 3 (ours)
X
Theorem 4 (ours)
FedAvg
(Karimireddy et al., 2020)
hiea +
3/2 +
(B, H)-GG (0, H)-LGD
X
(Yu et al., 2019)
bi
Mb
X
XXX
(Khanduri et al., 2021)
biea + (M/T)a/8
SCAFFOLD
(Karimireddy et al., 2020)
O
X
FedPD
(Zhang et al., 2020)
X
X
STEM
(Khanduri et al., 2021)
X
X
FCO
(Yuan et al., 2021)
x
(convex composite)
Table I: Comparison of heterogeneous FL algorithms in terms of communication complexity bounds for reaching an c-stationary solution, local dissimilarity (LD) condition, applicability to non-smooth (NS) functions and tolerance to partial participation (PP). Except for FCD, all the results listed are for non-convex functions. The involved quantities are M: total number of devices; /: number of chosen devices for partial participation; minibatch size for local stochastic optimization.
Highlight of contributions. The theoretical contributions of this work are highlighted follows:
From the perspective of algorithmic stability theory, we provide set of novel local dissimilarity invariant convergence guarantees for the widely used FedProx algorithm for non-convex heteroge- neous FL, with smooth non-smooth local functions. Our theory for the first time reveals that local dissimilarity and smoothness are not necessary to guarantee the convergence of FedProx with reasonable rates.
We present FedMSPP as a minibatch stochastic extension of FedProx and analyze its population optimization performance for both smooth and non-smooth FL problems, again without assuming local dissimilarity conditions. Particularly for smooth problems, our result provably shows that FedMSPP enjoys linear speedup in terms of minibatching size and partial participation ratio.
Paper organization. In Section we present our local dissimilarity invariant convergence analysis for the vanilla FedProx with smooth non-smooth loss functions. In Section we propose FedMSPP as a minibatch stochastic extension of FedProx and analyze its convergence behavior through the lens of algorithmic stability theory. In Section 4, we discuss some additional related work on the topics covered by this paper. The concluding remarks are made in Section 5. Finally, all the technical proofs and some additional related work are relegated the appendix sections.
2 Convergence of FedProx
We begin by providing an improved analysis for the vanilla FedProx which is not relying on the
(B, H) H-LG type conditions. We first introduce notations that will be used in the analysis to follow.
Notations. Throughout the paper, we use [n] to denote the set {1....,n}
n). || to denote the Euclidean norm and (--)) to denote the Euclidean inner product. We say function f is G-Lipschitz continuous if (ww - f(w')l < Gllu
ww - w/l for all w, w' e RP, and is L-smooth if V f((ww vf(w)) < L w w'll for all w E Rº. Moreover, say ff is p-weakly convex if for any w, w R.P
f(w) > f(w'((()))))))))))) + (Of(w'), w - w') || - w 11
where 0f(w') represents a subgradient of f evaluated at w'. We denote by
fo(w) = min f(u) + ||u - w|? 2m
5

---

the ny-Moreuu-envelope of f. and by
1 1|u w|? 2nj
prox, prox,f (16) = arg min [(u) +
the proximal mapping associated with f. We also need access the following definition of inexact local update oracle for FedProx.
Definition (Local inexact oracle of FedProx), Suppose that the local proximal point regularized objective Q (m) (w:; w(-1) (cf. (3)) admits a global minimizer. For each time instance t, we say that the local update oracle of FedProo is Ex-inexactly solved with sub-optimality &, if
Q (w) fmm) w;-1) I min Qumi (w; wx-1) + Et.
Throughout our analysis, we focus on the case where the devices are sampled with replacement, while all the results extend well to the regime of sampling without replacement.
2.1 Results for Smooth Problems
The following theorem is main result the convergence rate of FedProx for smooth and non- convex federated optimization problems. A proof of this result is deferred to Appendix B.1. We
assume that the initial sub-optimality A = Remm(wo) min Rem (w) is bounded.
Theorem 1. Assume that for each m E [M], the loss function g(m) is G-Lipschitz and L-smooth
with respect its first argument. Set [t] I and nn ar. min (T/8: TJ. Suppose that the local
GL update oracle of FedProw is Ex-inexactly solved with &1 min [ +1j 21(L }. Let t be an index uniformly randomly chosen in {0, I..,T T 13. Then it holds that
G²
E R
)1 x LA + G²) max T2/3
VTI
few remarks are in order.
Remark 1. Compared to the (+) bound from Li et al. (2020b), our rate established in Theorem I is slower but it is valid without assuming the unrealistic (B, 0)-LGD conditions and imposing strong regularization conditions on (see, e.g., Li et al., 2020b, Remark 5). Moreover, the dominant term VTI in our bound reveals the benefit of device sampling for linear speedup which is not clear in the previous analysis by Li et al. (2020b).
Remark In the extreme case of full device participation, i.e., It [M], the terms related in Theorem can be removed and thus the convergence rate becomes TH73 under тH O (TTTTT)) ). In this same setting, we comment that the rate can also be improved to (+)) using our proof augments if (B, O)-L additionally assumed.
Remark 3. The G-Lipschitz-loss assumption in Theorem can be alternatively replaced by the bounded gradient condition as commonty used the analysis of FL algorithms (Li al., 2020b: Zhang et al., 2020). Despite that our analysis does not rely on the (B, 0L-DD condition, the assumed G G-Lipschitz (or bounded gradient) condition actually implies that the local objective gradients are not too dissimilar, which shares close spirit to the typically assumed (0. HIGD condition (Karimireddy et al., 2020) and cnter-elient-varione condition (Khanduri et al., 202 ()21)). is noteworthy that these mentioned client heterogeneity conditions are substantially milder than the (B, O)-L condition as required in the original analysis of FedProa.
2.2 Results for Non-smooth Problems
Now we turn to study the convergence of FedProx for weakly convex but not necessarily smooth problems. For the sake of presentation clarity, work the exact FedProx in which the local update oracle is assumed to be exactly solved, i.e. & = O. Extension to the inexact case is more or less straightforward, though with somewhat more involved perturbation treatments. In the analysis to
follow, we assume that the initial subooptimality (10) == Rezm,p (wo) - min, Remm, associated with M-Moreau -eneeoope of Rem is bounded. The following our main result on the convergence of FedProx for non non-smooth and weakly convex problems.

---

Theorem 2. Assume that for each m E [M] the loss function g(m) is G-Lipschitz and v-weakly convex with respect to its first argument. Set Th - for arbitrary p < z.. Suppose that the local update oracle of FedProx is exactly solved with E = 0. Let t" be an index uniformly randomly chosen in {0, .,, 13. Then it holds that
A
enep + pC2
(0)
E V
PVT
Proof. The proof technique is inspired by the arguments from Davis and Drusvyatskiy (2019) devel- oped for analyzing stochastic model-based algorithms, with several new elements along developed for handling the challenges introduced by the model averaging and partial participation mechanisms associated with FedProx. A particular crux here is that due to the random subset model aggregation of w, |म. Eyer, w/(), the local function values Reeu (w;) are longer independent of each other
1. though ę uniformly random. As a consequence, [1,] EeEr, R (w:) not an unbiased estimation of Rerm( To overcome this technical obstacle, we make use of a key observation that whth) will be almost surely close enough to Wt-1 if the learning rate nt is small enough (which is the case in our choice of nt), and thus we can replace the former with the latter whenever beneficial but without introducing too much approximation error. See Appendix B.2 for a full proof of this result.
A few remarks are in order.
To our best knowledge, Theorem 2 the first convergence guarantee for FL algorithms applicable to generic non-smooth and weakly convex problems. This is in sharp contrast with FCO (Yuan et al., 2021) which focuses on composite convex and non-smooth problems such as t estimation, or Fed-HT (Tong et al., 2020) which is specially customized for sparse learning problems where the non-convexity essentially arises from the cardinality constraint.
Remark 4.
Remark 5. Let consider Wt- prox (we)), the proximal mapping of Wt- associated with Rarw- In view of a feature of Moreau envelope to characterize stationarity (Davis and Drusvyatskiy, 2019), if we- has small gradient norm || R Rera,p(we-) }I|, then Wy- must be a near-stationary solution and wi- stays in the proximity of Tr- due to the identity -m - w || = Rev.,r.e Therefore, the bound in Theorem 2 suggests that in expectation w.- converges to a stationary solution and wy- converges to Wr. both at the rate of O (T)).
Remark 6. We comment that the bound in Theorem 2 dependent on I, the number of selected devices. On one hand, for I = O(1) and sufficiently large T> O the bounds Theorem I and Theorem 2 are comparable to each other, which demonstrates that the smoothness not must-hav for FedProz to get sharper convergence bound with small device sampling rate. On the other hand, in the near-full participation setting where I = O(M), ) the bound in Theorem 2 for non-smooth problems will be slower when M is large. Extremely when It [M], the O ( bound substantially inferior the smooth case which has improved rate of O T discussed in Remark 2.
3 Convergence of FedProx with Stochastic Minibatching
When it comes the implementation of FedProx, a notable challenge is that the local proximal point update oracle (3) is by itself a full-batch ERM problem which would be expensive to solve even approximately in large-scale settings. Moreover, in the settings where the data distribution over devices is highly imbalanced, the computational overload of local update could vary significantly across the network, which impairs communication efficiency. It is thus desirable to seek stochastic approximation schemes for hopefully improving the local oracle update efficiency and overload balance of FedProx. To this end, inspired by the recent success of minibatch stochastic proximal point methods (MSPP) (Asi et al., 2020; Deng and Gao, 2021), we propose to implement FedProx using iid. MSPP as the local stochastic optimization oracle. More precisely, let B(()) = {=(_-)) } ~ (D(m)) be a minibatch of b i.i.d. samples drawn from the distribution ((m) device m and time instance
(m))
t 2 1. We denote
b glm) (w; 2/
[1)))) (w) = p(())
7

---

Algorithm 1: FedMSPP: Federated Minibatch Stochastic Proximal Point
Input :Minibatch size b; learning rates (ee))
Output:
Initialization Set wo, e.g., typically as zero vector.
for f == 1...... do
/* Device selection and model broadcast on the server
Server uniformly randomly selects subset 1, S [M] of devices and sends w to all the selected devices;
/* Local model updates the selected devices
for & E 1 in parallel do
Device & samples minibatch B(()
(2 Ji=1
iid (DIS))o.
Device & inexactly updates the its local model
{o(e) n3(0)
w
arg min
(w;; wi-1) == R(E)
(w) + 1/0 wi-111
(8)
wER"
(w) is given by (7). where R(()) g(e)
Device & sends with back server. end
/* Model aggregation on the server
Sever aggregates the local models received from 4, to update the global model as we प Eçen, w(()
end
the local minibatch empirical risk function over B(m) Here, the only modification propose to make is to replace the empirical risk R/ml (w) in the original update form (3) with its minibatch counterpart R. R!m)) (w). The resultant FL framework, which refer to FedMSPP (Federated MSPP), is outlined in Algorithm 1. Clearly, the vanilla FedProx is special case of FedMSPP when applied to the federated ERM form (2) with full data batch B( (mm).
3.1 Results for Smooth Problems
We first analyze the convergence rate FedMSPP for smooth and non-convex problems using the tools borrowed from algorithmic stability theory. Analogous to the Definition introduce the following definition of inexact local update oracle for FedMSPP.
Definition (Local inexact oracle of FedMSPP). Suppose that the local proximal point regularized objective Q (w: w;-1) (ef. (8)) admits a global minimizer. For each time instance I, we say that the local update oracle of FedMSPP is Et Es-inexactly solved with sub-optimality E 2 if
Q (m)) pil (w (m)) w I min Q (mm) (w: wi-1) Ex-
We also assume that the initial population sub-optimality (0) ((( - min R(w) is bounded. The following theorem is our main result on FeaMSPP for smooth and non-convex FL problems.
Theorem 3. Assume that for each m E [M], the loss function g(m)) is G-Lipschitz and L-smooth
with respect to its first argument. Set I = and T)s = BE min
Suppose that the
T17B
T
local update oracle of FedMSPP is En-inexactly solved with E S min
GP 2(Lmc+1)'
Ga
1.. (bI(
86
Let be an index uniformly randomly chosen {0, 1,.,, 1}. Then it holds that
| Rl || $ I (0) + G2) max T2/33 .TbI
E

---

Proof. Let consider d(m) V Vr B/m)
(mn( which is roughly the local update direction on device
m, in the sense that w fm) W4-1 mdl given that the local update oracle is solved to sufficient accuracy. As key ingredient our proof, we show via some extended uniform stability arguments in terms of gradients (see Lemma 3) that the averaged directions d, == Tr. EeEr, of aligns well with the global gradient V R 1) in expectation (see Lemma 11). Therefore, in average it roughly holds that w; पेत EeEr, w/{)) - n,d, = W 7) V R(w,,1), which suggests that wy updated roughly along the direction of global gradient descent and thus is expected converge quickly. Based on this novel analysis, we are free of explicitly imposing local dissimilarity type conditions on local objectives. See Appendix C.I for a full proof of this result.
Remark 7. For 2 ) the bound in Theorem 3 is dominated by (((A)) which gives the communication complexity nīez. This shows that FedMSPP enjoys linear speedup with respect both local minibatching and device sampling sizes.
While the bound in Theorem 3 is derived for the population form of FL in (1), an identical bound naturally holds for the empirical form (2) under minibatch sampling according to local data empirical distribution.
Remark
3.2 Results for Non-smooth Problems
Analogues to FedProx, , can further show that FedMSPP converges reasonably well when applicd to weakly convex and non-smooth problems. In the analysis to follow, we assume that the initial sub-optimality { Rw(w))= min (w)w) associated with e-Moorau-envelope of R is bounded. The following our main result in this line.
Theorem 4. Assume that for each m E [M] the loss function gim) is G-Lipschitz and v-weakly convex with respect to its first argument. Set ole - for arbitrary p < zu. Suppose that the local update oracle of FedMSPP is exactly solved with E = O. Let t be an index uniformly randomly chosen in {0, ' 13. Then it holds that
( + pG²
V [R )1²}
E
PVT
Proof. The proof argument is slight adaptation of that of Theorem 2 to the population FL setup (1) with FedMSPP. For the sake of completeness, a full proof is reproduced in Appendix C.2.
We comment in passing that the discussions made in Remarks 4-6 extend directly to Theorem 4.
Additional Related Work
The present work is situated at the intersection of federated learning, stochastic proximal point optimization and algorithmic stability theory. We next briefly review some additional work in these lines of research that are closely related to ours.
Heterogenous federated learning. The presence of device heterogeneity features key distinction between FL and classic distributed learning. The most commonly used FL method is FedAvg (M- cMahan al., 2017), where the local update oracle is formed multi-epoch SGD. FedAvg was early analyzed for identical functions (Stich, 2019; Stich and Karimireddy, 2020) under the name of local SGD. In heterogeneous setting, numerous recent studies have focused on the analysis of FedAvg and other variants under various notions of local dissimilarity (Li et al., 2020c; Woodworth al., 2020: Chen al., 2020; Khaled et al., 2020; Reddi al., 2021; Khanduri et al., 2021; Li et al., 2022; Chen et al., 2022; Zhao et al., 2022). As another representative FL method, FedProx (Li al., 2020b) has recently been proposed apply averaged proximal point updates solve het- erogeneous federated minimization problems. The theoretical guarantees of FedProx have been established for both convex and non-convex problems, but under a fairly stringent assumption of gradient similarity (see Definition 1) to measure data heterogeneity (Li et al., 2020b; Pathak and Wainwright, 2020; Nguyen et al., 2021). This assumption was relaxed by FedPD (Zhang et al., 2020) inside meta-frannt fett--rrameork of primal-dual optimization. The SCAFFOLD (Karimireddy et al., 2020) and

---

VRL-SGD (Liang et al., 2019) are two algorithms that utilize variance reduction techniques to correct the local update directions, achieving convergence guarantees independent of the data heterogeneity, For composite non-smooth FI. problems, the FCO proposed in Yuan et al. (2021) employs a server dual averaging procedure to circumvent the curse of primal averaging suffered by FedAvg. In sharp contrast to these prior works which either require certain stringent local dissimilarity conditions, or require full device participation, or only applicable to smooth problems, we show through a novel analysis based on algorithmic stability theory that the well-known FedProx can elegantly overcome all these shortcomings in a simple algorithmic framework. Another research direction in FL is to adopt compression methods for efficient communication (e.g., Haddadpour et al. (2020); Li and Li (2022)), which could also be applied to FedProx, as a topic for future investigation.
Minibatch stochastic proximal point methods. The proposed FedMSPP algorithm is variant of FedProx that simply replaces the local proximal point oracle with MSPP, which in each iteration updates the local model via (approximately) solving a proximal point estimator over a stochastic minibatch. The MSPP-type methods have been shown to attain a substantially improved iteration stability and adaptivity for large-scale machine learning, especially in non-smooth optimization settings (Li et al., 2014; Wang et al., 2017; Asi and Duchi, 2019; Deng and Gao, 2021). However, it is not yet known if FedProx or FedMSPP can achieve similar strong guarantees for non-smooth heterogenous FL problems.
Algorithmic stability. Our analysis for FedMSPP builds largely upon the classic algorithmic stability theory. Since the seminal work of Bousquet and Elissceff (2002), algorithmic stability has been serving as a powerful proxy for establishing strong generalization bounds (Zhang, 2003; Mukherjee et al., 2006; -Salre--- et al., 2010). Particularly, the state-of-the-ar risk bounds of strongly convex ERM are offered by approaches based on the notion of uniform stability (Feldman and Vondrák, 2018, 2019; Bousquet et al., 2020; Klochkov and Zhivotovskiy, 2021). was shown by Hardt et al. (2016) that the solution obtained via SGD is stable for smooth convex non-convex loss functions. For non-smooth convex losses, the stability induced generalization bounds have been established for SGD (Lei and Ying, 2020; Bassily et al., 2020: Yuan and Li, 2022). Through the lens of uniform algorithmic stability, convergence rates of MSPP have been studied for non- and convex (Wang& et al., 2017), or weakly convex (Deng and Gao, 2021) losses.
5 Conclusions
In this paper, we have exposed three shortcomings of the prior analysis for FedProx in unrealistic assumptions about local dissimilarity, inapplicability to non-smooth problems and expensive (and potentially imbalanced) computational cost of local update. In order to tackle these issues, we developed a novel convergence theory for the vanilla FedProx and its minibatch stochastic variant, FedMSPP, through the lens of algorithmic stability theory. In nutshell, our results reveal that with minimal modifications, FedProx is able to kill three birds with one stone: it enjoys favorable rates of convergence which are simultaneously invariant to certain stringent local dissimilarity conditions, applicable to smooth or non moo-smooth problems, and scaling linearly with respect to local minibatch size and device sampling ratio for smooth problems. To the best of our knowledge, the present work is the first theoretical contribution that achieves all these appealing properties in a single FL framework.
Limitations. While our results in Theorems 2 and 4 for the first time guarantee the mon-astmtootic convergence of FedProx and FedMSPP for non-smooth and weakly-convex problems, the correspond- ing rates of convergence so far cannot demonstrate any linear speedup effort w.r.t. device sampling ratio and local minibatching size. This as opposed to what have been shown for smooth problems in Theorems and 3; thus we view it as a potential limitation of the techniques used by our analysis. In the smooth- case, the comparison in Table suggests that our results in Theorem and Theorem show faster convergence rates than those of the existing FL methods based on local SGD update, despite that FedProx/FedMSPP requires a considerably more expensive oracle for local update.
A Acknowledgments and Disclosure of F Funding
The authors sincerely thank the anonymous reviewers for their insightful comments which are truly helpful for improving this paper. The research was conducted while both authors worked for the Baidu Cognitive Computing Lab. Xiao-Tomi Xiao-Tong Yuan is also funded in part by the National Key Research and Development Program of China under Grant No. 201 1018AAA0104000 and in part by the Natural Science Foundation of China (NSFC) under Grant No.U22 W.o.U21B204 BB249, No.6187600 11766090 and 60.61930005.
10

---

References
Hilal Asi and John C. Duchi. Stochastic (approximate) proximal point methods: Convergence, optimality, and adaptivity. SIAM J. Optim., 29(3):257-2290, 2019.
Hilal Asi, Karan N. Chadha, Gary Cheng, and John C. Duchi. Minibatch stochastic approximate proximal point methods. In Advances in Neural Information Processing Systems (NeurIPS), virtual, 2020.
Raef Bassily, Vitaly Feldman, Cristóbal Guzmán, and Kunal Talwar. Stability of stochastic gradient descent on nonsmooth convex losses. In Advances in Neural Information Processing Systems (NeurIPS), virtual, 2020.
Abhishek Bhowmick, John Duchi, Julien Freudiger, Gaurav Kapoor, and Ryan Rogers. Protec- tion against reconstruction and its applications in private federated learning. arXiv preprint arXiv:1812.00994, 2018.
Olivier Bousquet and André Elisseeff. Stability and generalization. J. Mach. Learn. Res., 2:499-526, 2002.
Olivier Bousquet, Yegor Klochkov, and Nikita Zhivotovskiy. Sharper bounds for uniformly stable algorithms. In Proceedings of the Conference on Learning Theory (COLT), pages 510-626, Virtual Event [Graz, Austria], 2020.
Xiangyi Chen, Xiaoyun Li, and Ping Li. Toward communication efficient adaptive gradient method. In Proceedings of the ACM-IMS Foundations of Data Science Conference (FODS), pages 119-128, Virtual Event, 2020.
Xiangyi Chen, Belhal Karimi, Weijie Zhao, and Ping Li. On the convergence of decentralized adaptive gradient methods. In Proceedings of The 14th Asian Conference on Machine Learning (ACML), Hyderabad, India, 2022.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. EMNIST: Extending MNIST to handwritten letters. In International Joint Conference Neural Networks (IJCNN), pages 2921-2926 Anchorage,AK, 2017.
Damek Davis and Dmitriy Drusvyatskiy, Stochastic model-based minimization of weakly convex functions. SIAM J. Optim., 29(1):007-239, 2019.
Qi Deng and Wenzhi Gao. Minibatch and momentum model-based methods for stochastic weakly convex optimization. In Advances in Neural Information Processing Systems (NeurIPS), pages 13115-23127, virtual, 2021.
Igor Donevski, Jimmy Jessen Nielsen, and Petar Popovski. On addressing heterogeneity in federated learning for autonomous vehicles connected to drone orchestrator. Frontiers in Communications and Networks, 2:1-14, 2021.
André Elisseeff, Theodoros Evgeniou, and Massimiliano Pontil. Stability of randomized learning algorithms. J. Mach. Learn. Res., 6:55-79, 2005.
Vitaly Feldman and Jan Vondrák. Generalization bounds for uniformly stable algorithms. In Advances in Neural Information Processing Systems (NeurIPS), pages 9770-9780, Montréal, Canada, 2018.
Jan
Vitaly Feldman and Jan Vondrák. High probability generalization bounds for uniformly stable algorithms with nearly optimal rate. In Proceedings of the Conference on Learning Theory (COLT), pages 1270-1279, Phoenix, AZ, 2019.
Alec Go, Richa Bhayani, and Lei Huang. Twitter sentiment classification using distant supervision. CS224N project report, Stanford, 2009.
Farzin Haddadpour, Belhal Karimi, Ping Li, and Xiaoyun Li. Fedsketch: Communication--fficcii and private federated learning via sketching. arXiv preprint arXiv:2008. 000::00000,00 2020.
11

---

Andrew Hard, Kurt Partridge, Cameron Nguyen, Niranjan Subrahmanya, Aishance Shah, Pai Zhu, Ignacio Lopez-Moreno. and Rajiv Mathews. Training keyword spotting models on non-iid data with federated learning. In Proceedings of the s Annual Conference of the International Speech Communication Association (Interspeech), pages 4343-4347 Virtual Event, Shanghai, China, 2020.
Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability stochastic gradient descent. In Proceedings of the 33nd International Conference on Machine Learning (ICML), pages 1225-1234, New York City, NY, 2016.
Chaoyang He, Alay Dilipbhai Shah, Zhenheng Tang, Di Fan Adarshan Naiynar Sivashunmugam, Keerti Bhogaraju, Mita Shimpi, Li Shen, Xiaowen Chu, Mahdi Soltanolkotabi, and Salman Avestimehr. FedCV: A federated learning framework for diverse computer vision tasks. arXiv preprint arXiv:2111..1001 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. SCAFFOLD: stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning (ICML), pages 5132-5143, Virtual Event, 2020.
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local SGD on identical and heterogeneous data. In Proceedings of the 23rd International Conference on Artificial Intelligence and Statisties (AISTATS), pages
Online [Palermo, Sicily, Italy], 2020.
Prashant Khanduri, Pranay Sharma, Haibo Yang. Mingyi Hong, Jia Liu, Ketan Rajawat, and Pramod K. Varshney. STEM: A stochastic two-sided momentum algorithm achieving near- rear-optimal sample and communication complexities for federated learning. In Advances in Neural Information Processing Systems (NeurIPS), pages
virtual, 2021.
Yegor Klochkov and Nikita Zhivotovskiy. Stability and deviation optimal risk bounds with con- vergence rate o11/n). In Advances in Neural Information Processing Systems (NeurIPS), pages 5065-5076, virtual, 2021.
Jakub Konečný, H Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv:1610.054p9 2016.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 86(11):221 (1:1)2)22-82244 1998.
Yunwen Lei and Yiming Ying. Fine fine-grained analysis of stability and generalization for stochastic gradient descent. In Proceedings of the 37th International Conference on Machine Learning (ICML), pages 5809-5819, Virtual Event, 2020.
Mu Li, Tong Zhang, Yuqiang Chen, and Alexander J. Smola. Efficient mini-batch training for stochastic optimization. In Proceedings of the 20th ACM SIGKDD International Conference Knowledge Discovery and Data Mining (KDD), pages 661-670 New York, NY, 2014.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. FedDANE: A federated newton-tyye method. In Proceedings of the 53rd Asilomar Conference on Signals, Systems, and Computers (Asilomar), pages 1227-1231, Pacific Grove, CA, 2019.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Process. Mag., 37(3):50-60. 2020a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems (MLSys), Austin, TX, 2020b.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang. and Zhihua Zhang. On the convergence of FedAvg on non-llI data. In Proceedings of the 8th International Conference on Learning Representations (ICLR), Addis Ababa, Ethiopia, 2020c.
12

---

Xiaoyun Li and Ping Li. Analysis of error feedback in federated non-convex optimization with biased compression. arXiv preprint arXiv:2211.14299 2022.
Xiaoyun Li, Belhal Karimi, and Ping Li. On distributed adaptive optimization with gradient compres- sion. In Proceedings of the Tenth International Conference on Learning Representations (ICLR), Virtual Event, 2022.
Xianfeng Liang, Shuheng Shen, Jingchang Liu, Zhen Pan, Enhong Chen, and Yifei Cheng. Variance reduced local SGD with lower communication complexity. arXiv preprint arXiv:191.. 12844, 2019.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. Communication-effie efficieen learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS), pages 1273 1282, Fort Lauderdale, FL. 2017.
Sayan Mukherjee, Partha Niyogi, Tomaso Poggio, and Ryan Rifkin. Learning theory: stability is sufficient for generalization and necessary and sufficient for consistency of empirical risk minimization. Advances in Computational Mathematics, 25(1-3):161-193, 2006.
Hung Nguyen, Vikash Sehwag, Seyyedali Hosseinalipour, Christopher Brinton, Mung Chiang, and H Vincent Poor. Fast-convergent federated learning. /EEE Journal Selected Areas in Communications, 39(1):201-218, 2020.
Hung T. Nguyen, Vikash Sehwag, Seyyedali Hosseinalipour, Christopher G. Brinton, Mung Chiang, and H. Vincent Poor. Fast-convergent federated learning. IEEE J. Sel. Areas Commun., 39(1): 201-218, 2021.
Reese Pathak and Martin J. Wainwright. Fedsplit: algorithmic framework for fast federated optimization. In Advances in Neural Information Processing Systems (NeurIPS), virtual, 2020.
Jeffrey Pennington, Richard Socher, and Christopher D. Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 1532-1543, Doha, Qatar, 2014.
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In Proceedings of the 9th International Conference on Learning Representations (ICLR), Virtual Event, Austria, 2021.
Omar Rivasplata, Csaba Szepesvári, John -aylor, Emilio Sun. PAC-Bayes bounds for stable algorithms with instance deeenne-de priors. In Advances in Neural Information Processing Systems (NeurIPS), pages 9234 9244, Montréal, Canada, 2018.
and Shiliang
Shai Shalev-Shwartz, Ohad Shamir, Nathan Srebro, and Karthik Sridharan. Learnability, stability and uniform convergence. J. Mach. Learn, Res., 11:2635-2670 2010.
Sebastian U. Stich. Local SGD converges fast and communicates little. In Proceedings of the 7th International Conference on Learning Representations (ICLR), New Orleans, LA, 2019.
Sebastian U Stich and Sai Praneeth Karimireddy. The error-feedback framework: Better rates for SGD with delayed gradients and compressed updates. J. Mach. Learn. Res., 21:1-36, 2020.
Qiangian Tong, Guannan Liang, Tan Zhu, and Jinbo Bi. Federated nonconvex sparse learning, arXiv preprint arXiv:2101.00052 2020.
Jialei Wang, Weiran Wang, and Nathan Srebro. Memory and communication efficient distributed stochastic optimization with minibatch prox. In Proceedings of the 30th Conference Learning Theory (COLT), pages 1882-1919, Amsterdam, The Netherlands, 2017.
Blake F. Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs local SGD for heteroge- neous distributed learning. In Advances in Neural Information Processing Systems (NeurIPS), virtual, 2020.
13

---

Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Trans. Intell. Syst. Technol., 101)::1::112:19, 2019.
Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted SGD with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the Thirty-Third AAAI Conference on Artificial Intelligence (AAAI), pages 5693-5700, Honolulu, HI, 2019.
Honglin Yuan, Manzil Zaheer, and Sashank J. Reddi. Federated composite optimization. In Pro- ceedings of the 38th International Conference on Machine Learning (ICML), pages 12253-12266, Virtual Event, 2021.
Xiao-Tong Yuan and Ping Li. Boosting the confidence of generalization for (2-Stable randomized learning algorithms. arXiv preprint arXiv:2206.03834. 2022.
Tong Zhang. Leave-one-out bounds for kernel methods. Neural Comput., 15(6):1377-1437 2003.
Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. FedPD: A federated learning framework with optimal rates and adaptivity to non-iid data. arXiv preprint arXiv:2005.11418,
2020.
Weijie Zhao, Xuewu Jiao, Mingqing Hu, Xiaoyun Li. Xiangyu Zhang, and Ping Li. Communication- efficient terabyte-scale model training framework for online advertising. In Proceedings of the IEEE International Conference on Big Data (IEEE BigData). Osaka, Japan, 2022.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes) In Section 1.2, we highlighted the core contributions made in this paper which are respectively expanded with details in Section 2 and 3.
(b) Did you describe the limitations of your work? [Yes]] We have provided remark on the limitations of our theory at the end of the main paper.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Our contribution is theoretical in nature. As far as we are aware of, there are no foreseeable societal ethical consequences for the present research.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] The key assumptions the siructure of loss functions (such as Lipschitzness, smoothness, weak-convexity), learning rates and local oracle update sub-optimality have all been explicitly stated in details in Theorems 1, 2, 4.
(b) Did you include complete proofs of all theoretical results? [Yes] Due space limit, all the full proofs of results are relegated to the appendix sections which can be found in the supplementary document. Particularly for Theorem 2 and Theorem 3, we further provide sketched proofs in the main submission to highlight the proof road maps.
3. If you ran experiments... [NIA] Our work is focused on providing deeper theoretical understandings of FedProz and its stochastic variants under milder conditions.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[N/A] did not use any of such assets this work.
5. If you used crowdsourcing or conducted research with human subjects... [N/A] The present
research was carried out with no crowdsourcing or human subjects involved in.
14


Reviews:
[{"review_id": "X8AlHndjNMX", "reviewer": "Reviewer_MGVY", "summary": "In this paper, the authors study FedProx with an improved convergence analysis without the unrealistic assumption of local gradient dissimilarity. Instead, they introduce a novel local dissimilarity invariant convergence theory for FedProx. FedProx for nonsmooth optimization and its minibatch stochastic extension are also studied with convergence guarantees and complexity bounds. \n", "questions": "- Typos: \n\t- Line 244: independent to → independent of", "limitations": "The authors have adequately addressed the limitations of their work.", "rating": 9, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n- This work gets rid of the usual yet unrealistic assumption of local gradient dissimilarity in the original analysis of FedProx\n- Both *smooth and nonconvex* and *nonsmooth and weakly convex* cases are studied with convergence rates \n- Minibatch stochastic extension of FedProx is also studied with convergence guarantees, which also enjoys linear speedup in terms of minibatch size and partial participation ratio \n\nWeaknesses: \n- Experimental results could be preliminary ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "SgG0p5W4ps", "reviewer": "Reviewer_Cx6F", "summary": "This paper studies convergence of a previously proposed algorithm called FedProx. FedProx has been fairly popular in the FL community but its analysis is pointed out to depend on unrealistic assumptions, namely, on bounded local gradient dissimilarity assumption. In this work, the authors propose an analysis that doesn't need this assumption. Moreover, they also study convergence of FedProx when the objective is non-smooth. Finally, they propose a variant of FedProx called FedMSPP, which uses proximal-point update in its inner loop. They prove convergence of FedMSPP and provide some numerical results for this method in the Appendix.\n\n# After rebuttal\nI decided to increase the score from 3 to 6 since the technical issue has been fixed (I also increased \"Soundness\" rating from 1 to 3). I still think that a more interesting story could be presented, with at least some result saying why FedProx would be a good method.", "questions": "### Main questions\n1. The proof of Lemma 5 seems to use $L$-smoothness of $Q_{\\mathrm{erm}}^{(m)}$ (when the gradient norm is upper bounded by functional gap). However, $Q_{\\mathrm{erm}}^{(m)}(w; w_{t-1})=R_{\\mathrm{erm}}^{(m)}(w) + \\frac{1}{2\\eta_t}\\Vert w-w_{t-1}\\Vert^2$ seems to be ($L+\\frac{1}{\\eta_t}$)-smooth. For small $\\eta_t$, this seems to make a big difference. As far as I can see, this may change the result of Lemma 5 to have a constant error, in which case the theory might break down.\n2. Could the authors comment on why FedProx would be better than simply computing a single gradient $\\nabla R_{\\mathrm{erm}}^{(\\xi)}(w)$ on each client $\\xi\\in I_t$? At least in the smooth case, it seems to me that there is no benefit from computing the prox. I think it makes sense to at least mention this in the \"Limitations\" section.\n3. The proof of Lemma 4 is problematic. First of all, you take expectation with respect to sampling $I_t$, but then you write that\n$\\mathbb{E}\\left[\\frac{1}{|I_t|}\\sum_{\\xi\\in I_t}d_t^{(\\xi)} \\right] = \\frac{1}{|I_t|}\\sum_{\\xi\\in I_t}\\mathbb{E}\\left[d_t^{(\\xi)} \\right].$\nThe problem here is that $I_t$ is stochastic so you cannot put it outisde expectation. The claim itself is correct, so this is just a minor comment.  \nYou also say that \"By uniform without-replacement sampling strategy\", but then you rely \"the independence among the indices in $I_t$\". The issue here is that sampling without replacement does not give independent indices. This can be easily fixed, for instance, see Lemma 1 in \"Random Reshuffling: Simple Analysis with Vast Improvements\" for the derivation of variance. Alternatively, you can just assume uniform sampling (which would probably make less sense since we talk about client sampling).\n\n### Typos:  \nPage 2: partial participants (participation) of devices  \nPage 5, Table 1: functions. The involved of quantities (should be no \"of\")  \nPage 9: At nutshell (in a nutshell)  \nPage 9: weak convex (weakly convex)  \nPage 16, line 561: (Proof) of Theorem 1. \"Proof\" is also missing on page 18, line 592; page 23, line 657; page 25, line 682.", "limitations": "The limitations were discussed well. I would be also interested in seeing a discussion on the utility of computing prox instead of computing gradients.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n1. The idea to study FedProx without bounded dissimilarity makes perfect sense to me. The method is quite popular in FL, where bounded dissimilarity is not always a reasonable assumption since clients data can be very heterogeneous.\n2. The obtained complexities look good compared to that of other methods (see Table 1). The theory covers partial participation and shows improvement in terms of client batch size in case of smooth functions. \n3. The work is well-written and the main contributions and ideas are stated clearly.\n\n## Weaknesses\n1. I do not understand one step in the proof of Lemma 5 (see \"Questions\" section of my review). This seems to be of fundamental importance, so if my question is not addressed, I would vote for rejection of the paper. On the other hand, if I miss something, I will be willing to drastically change the assigned score.\n2. As far as I can see, there is no improvement from performing more than a single local step. This is a common issue with FL methods, but it is still disappointing.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "NZJYc9nPE3z", "reviewer": "Reviewer_TmSS", "summary": "This is theoretical paper and it aims to study the convergence analysis of FedProx. The difference in this paper compared to existing works is that it does not assume the constraints on local dissimilarity, smoothness and the size of minibach. The novelty is to use algorithm stability theory to perform the analysis. In addition, the authors also provide a FedMSPP algorithm that is a minibatch stochastic extension of FedProx and perform its convergence analysis. ", "questions": "1. In Theorem 1, is the result obtained by assuming 1 data sample point used per iteration? \n\n2. What is the real novelty in terms of analysis, please indicate the key step that is different from the approach used in the literature?\n\n3. In Theorems 1 and 2,, what is the reason to use t* instead of averaging all the terms? \n\n", "limitations": "The convergence rate can be loose. First, it can be worse than what was given in the literature. The authors used a few times that the obtained convergence rate is comparable with those in the literature. It is unclear this is because of the mathematical tools used in this paper or this is the result of relaxing the assumption on the local dissimilarity and/or smoothness. Second, as mentioned before, Theorems 2 and 4 can be very loose since it is not a function of I and/or d (the authors also mentioned this in the end of the paper). If the results are too loose, the contribution of this paper may be weakened. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n\n1. The authors claim that their theory is the first time local dissimilarity and smoothness are not necessary to guarantee the convergence of FedProx with reasonbale rates. \n\n2. The authors propose FedMSPP algorithm and perform its convergence analysis. \n\nWeaknesses:\n\n1. The main novelty of this paper is a new analysis method without assuming local dissimilarity and smoothness. However, the authors failed to clearly explain the details of this novelty. \n\n2. The obtained convergence rates can be loose. For example, in Theorem 2, the convergence rate is not a function of I, which seems not make sense. In Theorem 4, the convergence rate is not a function of I and d, which also seems not make sense. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "DVbqFz1215w", "reviewer": "Reviewer_TUCL", "summary": "This paper studies the FedProx algorithm for nonconvex optimization. FedProx is a distributed version of the stochastic proximal point method that has found many applications in practice, and which generally seems to perform well in practice. FedProx was introduced and analyzed in [1], but the convergence rate given in [1] depends on a highly non-standard assumption, the strong growth condition, which this paper terms the (B, 0)-LGD assumption. The aim of this paper is to study the convergence of FedProx without this assumption; instead using the more-common assumption that each objective is Lipschitz. The authors give a convergence rate for smooth objectives (Theorem 1) that does not depend on the (B, 0)-LGD assumption, and a similar rate for non-smooth but weakly convex problems (Theorem 2)-- in addition, they develop minibatch variants of the algorithm and give convergence rates for it in Theorem 3 (smooth) and Theorem 4 (non-smooth).\n\n[1] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, Virginia Smith. Federated Optimization in Heterogeneous Networks, mlsys.org (2020).\n", "questions": "- Can your analysis be extended to use bounded variance, rather than the assumption that each function is Lipschitz?\n", "limitations": "See the previous sections for limitations. Potential negative societal impacts are not applicable here.\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Main points:\n- The convergence rate given by Theorem 1 is a welcome improvement upon the prior convergence result given in [1], as it does not require the (B, 0) condition and thus can be used in more general circumstances. On the other hand, this is quite limited by the alternative assumption used: the requirement that each function is Lipschitz itself is a control on the gradient dissimilarity, and does not allow the rate to show improvement with lower function dissimilarity.\n- The rate for non-smooth objectives given in Theorem 2 is also novel and interesting in its own right, as it adds to the arsenal of algorithms that are analyzed in the setting of nonconvex but weakly convex objectives. While the authors state that it is the first convergence rate for FL algorithms, the rate given by [2] for SGD can be easily applied to federated learning, by seeing the sampling of stochastic gradients as sampling clients.\n- The minibatch variants of the method are a nice addition, even though Theorem 4 does not show that minibatching results in any improvement for the algorithm under non-smoothness.\n- None of the rates is asymptotically better than ordinary distributed SGD. This is quite puzzling, as FedProx requires a much stronger computational oracle and seems to perform well in practice. However, this could be inherent to the algorithm, rather than a product of a bad analysis. Without a lower bound for FedProx, we can not tell.\n- The novelty of the proof for non-smooth problems is quite limited, especially since the proof shows no benefit from minibatching-- it follows the same outline as in [2].\n\nI believe that this paper improves over the existing analysis and is a good addition to the literature, despite the rate itself not being very good compared to the baseline SGD. While it relaxes the (B, 0) condition required in [1], it also adds the Lipschitz condition that is more restrictive compared to prior work. So while this paper has some good, it is also more restrictive in other parts.\n\n[1] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, Virginia Smith. Federated Optimization in Heterogeneous Networks, mlsys.org (2020).\n[2] Damek Davis, Dmitriy Drusvyatskiy Stochastic Model-Based Minimization of Weakly Convex Functions, arXiv:1803.06523 (2018).\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept