Original: On Convergence of FedProx: Local Dissimilarity Invariant Bounds, Non-smoothness and Beyond
Revised: On Convergence of FedProx: Local Dissimilarity Invariant Bounds, Non-smoothness and Beyond

Original: Anonymous

Original: Abstract
Revised: Abstract

Original: The FedProx algorithm is a simple yet powerful distributed proximal point optimization method widely used for federated learning (FL) over heterogeneous data.
Revised: The FedProx algorithm is a simple yet powerful distributed proximal point optimization method widely used for federated learning (FL) over heterogeneous data.

Original: Despite its popularity and remarkable success witnessed in practice, the theoreticalunderstanding of FedProx is largely underinvestigated: the appealing convergencebehavior of FedProx is so far characterized under certain non-standard and unrealistic dissimilarity assumptions of local functions, and the results are limited tosmooth optimization problems.
Revised: 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.

Original: In order to remedy these deﬁciencies, we developa novel local dissimilarity invariant convergence theory for FedProx and its minibatch stochastic extension through the lens of algorithmic stability.
Revised: In order to remedy these deﬁciencies, we develop a novel local dissimilarity invariant convergence theory for FedProx and its minibatch stochastic extension through the lens of algorithmic stability.

Original: As a result, wecontribute to derive several new and deeper insights into FedProx for non-convexfederated optimization including: 1) convergence guarantees independent on local dissimilaritytype conditions; 2) convergence guarantees for non-smooth FLproblems; and 3) linear speedup with respect to size of minibatch and number ofsampled devices.
Revised: 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.

Original: Our theory for the ﬁrst time reveals that local dissimilarity andsmoothness are not must-have for FedProx to get favorable complexity bounds.
Revised: Our theory for the ﬁrst time reveals that local dissimilarity and smoothness are not must-have for FedProx to get favorable complexity bounds.

Original:  Introduction
Revised:  Introduction

Original: Federated Learning (FL) has recently emerged as a promising paradigm for communication-efﬁcientdistributed learning on remote devices, such as smartphones, internet of things or agents [Koneˇcn ` yet al., 2016, Yang et al., 2019].
Revised: Federated Learning (FL) has recently emerged as a promising paradigm for communication-efﬁcient distributed learning on remote devices, such as smartphones, internet of things, or agents (Kone ˇ cn ` y et al., 2016; Yang et al., 2019).

Original: The goal of FL is to collaboratively train a shared model that worksfavorably for all the local data but without requiring the learners to transmit raw data across thenetwork.
Revised: The goal of FL is to collaboratively train a shared model that works favorably for all the local data but without requiring the learners to transmit raw data across the network.

Original: The principle of optimizing a global model while keeping data localized can be beneﬁcialfor both computational efﬁciency and data privacy [Bhowmick et al., 2018].
Revised: The principle of optimizing a global model while keeping data localized can be beneﬁcial for both computational efﬁciency and data privacy (Bhowmick et al., 2018).

Original: While resemblingthe classic distributed learning regimes, there are two most distinct features associated with FL: 1)large statistical heterogeneity of local data mainly due to the non-iid manner of data generalizationand collection across the devices [Hard et al., 2020]; and 2) partial participants of devices in thenetwork mainly due to the massive number of devices.
Revised: While resembling the classic distributed learning regimes, there are two most distinct features associated with FL: 1) 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.

Original: These fundamental challenges make FL highlydemanding to tackle, both in terms of optimization algorithm design and in terms of theoreticalunderstanding of convergence behavior [Li et al., 2020a].
Revised: These fundamental challenges make FL highly demanding to tackle, both in terms of optimization algorithm design and in terms of theoretical understanding of convergence behavior (Li et al., 2020a).

Original: FL is most conventionally formulated as the following problem of global population risk minimizationaveraged over a set of M devices:
Revised: FL is most conventionally formulated as the following problem of global population risk minimization averaged over a set of M devices:

Original: [Equation]
Revised: [Equation]

Original:  where R ( m ) is the local population risk on device m , (cid:96) ( m ) : R p × Z ( m ) (cid:55)→ R + is a non-negativeloss function whose value (cid:96) ( w ; Z ( m ) ) measures the loss over a random data point Z
Revised: [Equation] ) where R ( m ) is the local population risk on device m , (cid:96) ( m ) : R p × Z ( m ) (cid:55)→ R + is a non-negative loss function whose value (cid:96) ( w ; Z ( m ) ) measures the loss over a random data point Z

Original: ( m ) ∈
Revised: ( m ) ∈

Original: Z ( m ) with  parameter w , D ( m ) represents an underlying random data distribution over Z ( m ) .
Revised: Z ( m ) with 36th Neural Systems (NeurIPS 2022). parameter w , D ( m ) represents an underlying random data distribution over Z ( m ) .

Original: Since the datadistribution is typically unknown, the following empirical risk minimization (ERM) version of (1) isoften considered alternatively:
Revised: Since the data distribution is typically unknown, the following empirical risk minimization (ERM) version of (1) is often considered alternatively:

Original: [Equation]
Revised: [Equation]

Original: where R ( m ) erm is the local empirical risk over the training sample D ( m )
Revised: where R ( m ) erm is the local empirical risk over the training sample D ( m )

Original: = { z ( m ) i } N m i
Revised: = { z ( m ) i } N m i

Original: =1 on device m .
Revised: =1 on device m .

Original: Thesample size N m may vary signiﬁcantly across devices, which can be regarded as another source ofdata heterogeneity.
Revised: The sample size N m may vary signiﬁcantly across devices, which can be regarded as another source of data heterogeneity.

Original: Federated optimization algorithms for solving (1) or (2) have attracted signiﬁcantresearch interest from both academia and industry, with a rich body of efﬁcient solutions developedthat can ﬂexibly adapt to the communication-computation tradeoffs and data/system heterogeneity.
Revised: Federated optimization algorithms for solving (1) or (2) have attracted signiﬁcant research interest from both academia and industry, with a rich body of efﬁcient solutions developed that can ﬂexibly adapt to the communication-computation tradeoffs and data/system heterogeneity.

Original: Several popularly used FL algorithms for this setting include FedAvg [McMahan et al., 2017],
Revised: Several popularly used FL algorithms for this setting include FedAvg (McMahan et al., 2017),

Original: FedProx [Li et al., 2020b], SCAFFOLD [Karimireddy et al., 2020], and FedPD [Zhang et al., 2020], toname a few.
Revised: FedProx (Li et al., 2020b), SCAFFOLD (Karimireddy et al., 2020), and FedPD (Zhang et al., 2020), to name a few.

Original: A consensus among these methods on communication-efﬁcient implementation is tryingto extensively update the local models (e.g., with plenty epochs of local optimization) over subsetsof devices so as to quickly ﬁnd an optimal global model using a minimal number of inter-devicecommunication rounds for model aggregation.
Revised: A consensus among these methods on communication-efﬁcient implementation is trying to extensively update the local models (e.g., with plenty epochs of local optimization) over subsets of devices so as to quickly ﬁnd an optimal global model using a minimal number of inter-device communication rounds for model aggregation.

Original: In this paper, we revisit the FedProx algorithm which is one of the most prominent frameworks forheterogeneous federated optimization.
Revised: In this paper, we revisit the FedProx algorithm which is one of the most prominent frameworks for heterogeneous federated optimization.

Original: Reasons for the interests of FedProx include implementationsimplicity, low communication cost, promise in dealing with data heterogeneity and tolerance topartial participants of devices [Li et al., 2020b].
Revised: 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).

Original: We analyze its convergence behavior, exposeproblems, and propose alternatives more suitable for scaling up and generalization.
Revised: We analyze its convergence behavior, expose problems, and propose alternatives more suitable for scaling up and generalization.

Original: We contributeto derive some new and deeper theoretical insights into the algorithm from a novel perspective ofalgorithmic stability  theory.
Revised: We contribute to derive several new and deeper theoretical insights into the algorithm from a novel perspective of algorithmic stability and generalization theory.

Original: 1.1 Review of FedProx
Revised: 1.1 Review of FedProx

Original: For solving FL problems in the presence of data heterogeneity, methods such as FedAvg based onlocal stochastic gradient descent (SGD) can fail to converge in practice when the selected devicesperform too many local updates [Li et al., 2020b].
Revised: For solving FL problems in the presence of data heterogeneity, methods such as 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).

Original: To mitigate this issue, FedProx [Li et al., 2020b]was recently proposed for solving the empirical FL problem (2) using the (inexact) proximal pointupdate for local optimization.
Revised: 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.

Original: The beneﬁts of FedProx include: 1) it provides more stable localupdates 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, evenunder partial participation and very dissimilar amounts of local updates [Li et al., 2020a].
Revised: The beneﬁts 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).

Original: Morespeciﬁcally, at each time instance t , FedProx uniformly randomly selects a subset I t ⊆
Revised: More speciﬁcally, at each time instance t , FedProx uniformly randomly selects a subset I t ⊆

Original: [ M ] ofdevices and introduces for each device ξ ∈
Revised: [ M ] of devices and introduces for each device ξ ∈

Original: I t the following proximal point ERM sub-problem forlocal update around the previous global model w t − 1 :
Revised: I t the following proximal point ERM sub-problem for local update around the previous global model w t − 1 :

Original: [Equation]
Revised: [Equation]

Original: where η t > 0 is the learning rate that controls the impact of the proximal term.
Revised: where η t > 0 is the learning rate that controls the impact of the proximal term.

Original: Then the global modelis updated by uniformly aggregating those local updates from I t as
Revised: Then the global model is updated by uniformly aggregating those local updates from I t as

Original: [Equation]
Revised: [Equation]

Original: In the extreme case of allowing η t → + ∞ in (3), FedProx reduces to the regime of FedAvg if using
Revised: In the extreme case of allowing η t → + ∞ in (3), FedProx reduces to the regime of FedAvg if using

Original: SGD for local optimization.
Revised: SGD for local optimization.

Original: Since its inception, FedProx and its variants have received signiﬁcantinterests in research [Pathak and Wainwright, 2020, Nguyen et al., 2020, Li et al., 2019a] and becomean algorithm of choice in application areas such as automatous driving [Donevski et al., 2021] andcomputer vision [He et al., 2021].
Revised: Since its inception, FedProx and its variants have received signiﬁcant interests in research (Pathak and Wainwright, 2020; Nguyen et al., 2020; Li et al., 2019a) and become an algorithm of choice in application areas such as automatous driving (Donevski et al., 2021) and computer vision (He et al., 2021).

Original: Theoretically, FedProx comes with convergence guaranteesunder the following bounded local gradient dissimilarity assumption that captures the statisticalheterogeneity of local objectives across the network:
Revised: 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:

Original: Deﬁnition 1 ( ( B, H ) -LGD) .
Revised: Deﬁnition 1 ( ( B, H ) -LGD) .

Original: We say the local functions R ( m ) have ( B, H ) -local gradient dissimilarity (LGD) if the following holds for all w ∈ R p :
Revised: We say the local functions R ( m ) have ( B, H ) -local gradient dissimilarity (LGD) if the following holds for all w ∈ R p :

Original: [Equation]
Revised: [Equation]

Original: The  deﬁnition naturally extends to the local empirical risks { R ( m ) erm } Mm =1 .
Revised: The above deﬁnition naturally extends to the local empirical risks { R ( m ) erm } Mm =1 .

Original: Specially in the homogenous setting where R ( m ) ≡ ¯ R , ∀ m ∈ [ M ] , we have B = 1 and H = 0 .
Revised: Specially in the homogenous setting where R ( m ) ≡ ¯ R , ∀ m ∈ [ M ] , we have B = 1 and H = 0 .

Original: Under( B, 0) -LGD and some regularization condition on the modulus B , it was shown that FedProx fornon-convex problems requires T = O (cid:0) 1 (cid:15) (cid:1) rounds of inter-device communication to reach an (cid:15) -stationary solution, i.e., 1 T (cid:80)
Revised: Under ( B, 0) -LGD and some regularization conditions on the modulus B , it was shown that FedProx for non-convex problems requires T = O (cid:0) 1 (cid:15) (cid:1) rounds of inter-device communication to reach an (cid:15) -stationary solution in the sense of 1 T (cid:80)

Original: Tt =

Revised: Tt =

Original: (cid:107)∇ ¯ R erm ( w t ) (cid:107) 2 ≤ (cid:15) [Li et al., 2020b].
Revised: (cid:107)∇ ¯ R erm ( w t ) (cid:107) 2 ≤ (cid:15) (Li et al., 2020b).

Original: Similar guarantees have alsobeen established for a variant of FedProx with non-uniform model aggregation [Nguyen et al., 2020].
Revised: Similar guarantees have also been established for a variant of FedProx with non-uniform model aggregation (Nguyen et al., 2020).

Original: Open issues and motivation.
Revised: Open issues and motivation.

Original: 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, restrictiveproblem regimes and expensive local oracle cost that remain open for exploration, as speciﬁed below.
Revised: 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 speciﬁed below.

Original: • Local dissimilarity  .
Revised: • Local dissimilarity condition .

Original: The appealing convergence behavior of FedProx is so far characterized undera key but non-standard ( B, H ) -LGD  (cf.
Revised: The appealing convergence behavior of FedProx is so far characterized under a key but non-standard ( B, 0) -LGD condition (cf.

Original: Deﬁnition 1)condition with B > 0and H = 0 .
Revised: Deﬁnition 1) with B > 0 .

Original: Such acondition is obviously unrealistic in practice: it essentially requires the local objectives share thesame stationary point as the global objective since (cid:107)∇ ¯ R erm ( w ) (cid:107) = 

Revised: Such a condition is obviously unrealistic in practice: it essentially requires the local objectives share the same stationary point as the global objective since (cid:107)∇ ¯ R erm ( w ) (cid:107) =

Original: implies (cid:107)∇ R ( m ) erm ( w ) (cid:107) = 0 forall m ∈
Revised: implies (cid:107)∇ R ( m ) erm ( w ) (cid:107) = 0 for all m ∈

Original: [ M ] .
Revised: [ M ] .

Original: However, if the optima of R ( m ) erm are exactly (or even approximately) the same, therewould be little point in distributing data across devices for federated learning.
Revised: However, if the optima of R ( m ) erm are exactly (or even approximately) the same, there would be little point in distributing data across devices for federated learning.

Original: It is thus desirable tounderstand the convergence behavior of FedProx for heterogeneous FL without imposing stringentlocal dissimilarity conditions like ( B, 0) -LGD with B > 0 .•
Revised: It is thus desirable to understand the convergence behavior of FedProx for heterogeneous FL without imposing stringent local dissimilarity conditions like ( B, 0) -LGD.

Original:  Non-smooth optimization .
Revised: • Non-smooth optimization .

Original: The existing convergence guarantees of FedProx are only availablefor FL with smooth losses.
Revised: The existing convergence guarantees of FedProx are only available for FL with smooth losses.

Original: More often than not, however, FL applications involve non-smoothobjectives due to the popularity of non-smooth losses (e g., hinge loss and absolute loss) in machinelearning, and training deep neural networks with non-smooth activation like ReLU.
Revised: 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.

Original: Therefore, it isdesirable to understand the convergence behavior of FedProx in non-smooth problem regimes.• Local oracle complexity .
Revised: Therefore, it is desirable to understand the convergence behavior of FedProx in non-smooth problem regimes. • Local oracle complexity .

Original: Unlike the (stochastic) ﬁrst-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 tobe expensive to solve even approximately per-iteration.
Revised: Unlike the (stochastic) ﬁrst-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.

Original: Plus, due to the potentially imbalanceddata distribution over devices, the computational overload of the proximal point oracle couldvary signiﬁcantly across the network.
Revised: Plus, due to the potentially imbalanced data distribution over devices, the computational overload of the proximal point oracle could vary signiﬁcantly across the network.

Original: Therefore, it is important to investigate whether using thestochastic approximation to the proximal point oracle (3) can provably improve the computationalefﬁciency of FedProx .
Revised: Therefore, it is important to investigate whether using minibatch stochastic approximation to the proximal point oracle (3) can provably improve the computational efﬁciency of FedProx .

Original: Last but not least, existing convergence analysis of FedProx mainly focuses on the empirical FLproblem (2).
Revised: Last but not least, existing convergence analysis of FedProx mainly focuses on the empirical FL problem (2).

Original: The optimality in terms of the population FL problem (1) is not yet clear for FedProx .
Revised: The optimality in terms of the population FL problem (1) is not yet clear for FedProx .

Original: The primary goal of this work is to remedy these theoretical issues simultaneously, so as to lay amore solid theoretical foundation for the popularly applied FedProx algorithm.
Revised: The primary goal of this work is to remedy these theoretical issues simultaneously, so as to lay a more solid theoretical foundation for the popularly applied FedProx algorithm.

Original: 1.2 Our Contributions
Revised: 1.2 Our Contributions

Original: In this paper, we make progress towards understanding the convergence behavior of FedProx fornon-convex heterogenous FL under weaker and more realistic conditions.
Revised: In this paper, we make progress towards understanding the convergence behavior of FedProx for non-convex heterogenous FL under weaker yet more realistic conditions.

Original: The main results are a setof local dissimilarity invariant bounds for smooth or non-smooth problems.
Revised: The main results are a set of local dissimilarity invariant bounds for smooth or non-smooth problems.

Original: Main results for the vanilla FedProx .
Revised: Main results for the vanilla FedProx .

Original: As a starting point to address the restrictiveness of localdissimilarity assumption, we provide a novel convergence analysis for the vanilla FedProx algorithmindependent of local dissimilarity type conditions.
Revised: As a starting point to address the restrictiveness of local dissimilarity assumption, we provide a novel convergence analysis for the vanilla FedProx algorithm invariant to the ( B, 0) -LGD condition.

Original: 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
Revised: 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

Original: [Equation]
Revised: [Equation]

Original: where I is the number devices randomly selected for local update at each iteration.
Revised: where I is the number devices randomly selected for local update at each iteration.

Original: If all the devicesparticipate in the local updates for every round, i.e. I t =
Revised: If all the devices participate in the local updates for every round, i.e. I t =

Original: [ M ] , the rate of convergence can be improvedto O ( 1 T 2 / 3 ) .
Revised: [ M ] , the rate of convergence can be improved to O ( 1 T 2 / 3 ) .

Original: For T < I 3 , the rate in (4) is dominated by O ( 1 T 2 / 3 ) which gives the communicationcomplexity 1 (cid:15) 3 / 2 to achieve an (cid:15) -stationary solution.
Revised: For T < I 3 , the rate in (4) is dominated by O ( 1 T 2 / 3 ) which gives the communication complexity 1 (cid:15) 3 / 2 to achieve an (cid:15) -stationary solution.

Original: On the other hand when T ≥ I 3 , the rate isdominated by O ( 1 √ TI ) which gives the communication complexity 1 I(cid:15) 2 .
Revised: On the other hand when T ≥ I 3 , the rate is dominated by O ( 1 √ TI ) which gives the communication complexity 1 I(cid:15) 2 .

Original: Compared to the alreadyknown O ( 1 (cid:15) ) complexity bound of FedProx under the unrealistic ( B, 0) -LGD condition [Li et al., 2020b], our rate in (4) is slower but it holds without needing to impose stringent regularity conditionson the dissimilarity of local functions, and it reveals the effect of device  sampling for acceleratingconvergence.
Revised: Compared to the already known O ( 1 (cid:15) ) complexity bound of FedProx under the unrealistic ( B, 0) -LGD condition (Li et al., 2020b), our rate in (4) is slower but it holds without needing to impose stringent regularity conditions on the dissimilarity of local functions, and it reveals the beneﬁt of device minibatch sampling for accelerating convergence.

Original: Further for non-smooth and non-convex problems, we establish in Theorem 2 thefollowing rate of convergence
Revised: Further for non-smooth and non-convex problems, we establish in Theorem 2 the following rate of convergence

Original: [Equation]
Revised: [Equation]

Original: which is invariant to the number of selected devices in each round.
Revised: which is not dependent on the number of selected devices in each round.

Original: In the case of I = O (1) ,the bounds in (4) and (5) are comparable, which indicates that smoothness is not must-have for
Revised: In the case of I = O (1) , the bounds in (4) and (5) are comparable, which indicates that smoothness is not must-have for

Original: FedProx to get sharper convergence bound especially with low participation ratio.
Revised: FedProx to get sharper convergence bound especially with low participation ratio.

Original: On the other endwhen I =
Revised: On the other end when I =

Original: O ( M ) , the bound (5) for non-smooth problems is slower than the bound (4) for smoothfunctions in large-scale networks.
Revised: O ( M ) , the bound (5) for non-smooth problems is slower than the bound (4) for smooth functions in large-scale networks.

Original: Main results for minibatch stochastic FedProx .
Revised: Main results for minibatch stochastic FedProx .

Original: Then as the chief contribution of the present work,we propose a minibatch stochastic extension of FedProx along with its population optimizationperformance analysis from a novel perspective of algorithmic stability theory.
Revised: 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.

Original: Inspired by the recentsuccess of minibatch stochastic proximal point methods (MSPP) [Li et al., 2014, Wang et al., 2017,
Revised: Inspired by the recent success of minibatch stochastic proximal point methods (MSPP) (Li et al., 2014; Wang et al., 2017;

Original: Asi et al., 2020, Deng and Gao, 2021], we propose to implement FedProx using MSPP as the localupdate oracle.
Revised: Asi et al., 2020; Deng and Gao, 2021), we propose to implement FedProx using MSPP as the local update oracle.

Original: The resulting method, which is referred to as FedMSPP , is expected to attain improvedtrade-off between computation, communication and memory efﬁciency for large-scale FL.
Revised: The resulting method, which is referred to as FedMSPP , is expected to attain improved trade-off between computation, communication and memory efﬁciency for large-scale FL.

Original: In the caseof imbalanced data distribution, minibatching is also beneﬁcial for making the local computationmore balanced across the devices.
Revised: In the case of imbalanced data distribution, minibatching is also beneﬁcial for making the local computation more balanced across the devices.

Original: Based on some extended uniform stability arguments for gradients,we show in Theorem 3 the followinglocal dissimilarity invariant rate of convergence for FedMSPP interms of population optimality:
Revised: 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) -LGD condition:

Original: [Equation]
Revised: [Equation]

Original: where b is the minibatch size of local update.
Revised: where b is the minibatch size of local update.

Original: For empirical FL, identical bound holds under samplingaccording to empirical distribution.
Revised: For empirical FL, identical bound holds under sampling according to empirical distribution.

Original: For T < ( bI ) 3 , the rate in (6) is dominated by O ( 1 T 2 / 3 ) whichgives the communication complexity 1 (cid:15) 3 / 2 , and it matches that of the vanilla FedProx .
Revised: For T < ( bI ) 3 , the rate in (6) is dominated by O ( 1 T 2 / 3 ) which gives the communication complexity 1 (cid:15) 3 / 2 , and it matches that of the vanilla FedProx .

Original: For sufﬁcientlylarge T ≥ ( bI ) 3 , the rate is dominated by O ( 1 √ TbI ) which gives the communication complexity 1 bI(cid:15) 2 .
Revised: For sufﬁciently large T ≥ ( bI ) 3 , the rate is dominated by O ( 1 √ TbI ) which gives the communication complexity 1 bI(cid:15) 2 .

Original: This shows that local minibatching and device sampling are both beneﬁcial for linearly speedingup communication.
Revised: This shows that local minibatching and device sampling are both beneﬁcial for linearly speeding up communication.

Original: Further, when applied to non-smooth problems, we cansimilarly show that FedMSPP converges at the rate of
Revised: Further, when applied to non-smooth problems, we can show that FedMSPP converges at the rate of

Original: [Equation]
Revised: [Equation]

Original: which is comparable to that of (6) when b =
Revised: which is comparable to that of (6) when b =

Original: O (1) and I =
Revised: O (1) and I =

Original: O (1) , but without showing the effect oflinear speedup with respect to b and I .
Revised: O (1) , but without witnessing the effect of linear speedup with respect to b and I .

Original: Comparison with prior results.
Revised: Comparison with prior results.

Original: In Table 1, we summarize our communication complexity boundsfor FedProx ( FedMSPP ) and compare them with several related heterogeneous FL algorithms interms of the dependency on local dissimilarity, applicability to non-smooth problems and toleranceto partial participation.
Revised: 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.

Original: A few observations are in order.
Revised: A few observations are in order.

Original: First , regarding the requirement of localdissimilarity, all of our O ( 1 (cid:15) 2 ) bounds are independent of local dissimilarity conditions, and they arecomparable to those of SCAFFOLD and FCO (for convex problems) which are also invariant to localdissimilarity.
Revised: First , regarding the local dissimilarity condition, all of our O ( 1 (cid:15) 2 ) bounds are not dependent on the ( B, H ) -LGD type conditions, and they are comparable to those of SCAFFOLD and FCO (for convex problems) which are also invariant to local dissimilarity conditions.

Original: Second , with regard to the applicability to non-smooth optimization, our convergenceguarantees in Theorem 2 and Theorem 4 are established for non-smooth and weakly convex functions.
Revised: 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.

Original: While FCO is the only one in the other considered algorithms that can be applied to non-smoothproblems, it is customized for federated convex composite optimization with potentially non-smoothregularizers [Yuan et al., 2021].
Revised: While FCO is the only one in the other considered algorithms that can be applied to non-smooth problems, it is customized for federated convex composite optimization with potentially non-smooth regularizers (Yuan et al., 2021).

Original: Third , in terms of tolerance to partial participation, all of our resultsare robust to device sampling, and the O ( 1 bI(cid:15) 2 ) bound in Theorem 3 for FedMSPP is comparable tothe best known results under partial participation as achieved by FedAvg and SCAFFOLD .
Revised: Third , in terms of tolerance to partial participation, all of our results are robust to device sampling, and the O ( 1 bI(cid:15) 2 ) bound in Theorem 3 for FedMSPP is comparable to the best known results under partial participation as achieved by FedAvg and SCAFFOLD .

Original: If assumingthat all the devices participate in local update for each communication round and under certain localdissimilarity conditions, substantially faster O ( 1 (cid:15) ) bounds are possible for STEM and FedPD , while the
Revised: If assuming that all the devices participate in local update for each communication round and using momentum acceleration techniques, substantially faster O ( 1 (cid:15) ) bounds are possible for STEM and FedPD , while the

Original: O ( 1 (cid:15) 3 / 2 ) bounds can be achieved by FedAvg [Khanduri et al., 2021].
Revised: O ( 1 (cid:15) 3 / 2 ) bounds can be achieved by FedAvg (Khanduri et al., 2021).

Original: To summarize the comparison, 
Revised: To summarize the comparison, our ( B, H ) -LGD invariant convergence bounds for FedProx ( FedMSPP ) are comparable to the bestknown 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.

Original: Method
Revised: [Table]

Original: Work
Revised: Work

Original: Commun.
Revised: Commun.

Original: Complex.
Revised: Complex.

Original: LD Independ.
Revised: LD Condition

Original: NS
Revised: NS

Original: PP
Revised: PP

Original: [Table]
Revised: [Table]

Original: [Li

Original: [Table]

Original: al., 2020b]

Original: [Table]

Original: [Table]

Original: [Khanduri

Original: [Table]

Original: [Table]

Original: [Karimireddy

Original: [Table]

Original: [Table]

Original: [Zhang et al., 2020]

Original: [Table]

Original: [Table]

Original: [Table]

Original: our local dissimilarity invariant convergence bounds for FedProx ( FedMSPP ) are comparable to thebest-known rates in the identical setting, while covering the generic non-smooth and non-convexcases which to our knowledge so far has not been possible for other FL algorithms.

Original: Highlight of contributions.
Revised: Highlight of contributions.

Original: The theoretical contributions of this work are highlighted as follows:
Revised: The theoretical contributions of this work are highlighted as follows:

Original: • From the perspective of algorithmic stability theory, we provide a set of novel local dissimilarityinvariant convergence guarantees for the widely used FedProx algorithm for non-convex heteroge- neous FL, with smooth or non-smooth local functions.
Revised: • From the perspective of algorithmic stability theory, we provide a set of novel local dissimilarity invariant convergence guarantees for the widely used FedProx algorithm for non-convex heterogeneous FL, with smooth or non-smooth local functions.

Original: Our theory for the ﬁrst time reveals thatlocal dissimilarity and smoothness are not necessary to guarantee the convergence of FedProx withreasonable rates.
Revised: Our theory for the ﬁrst time reveals that local dissimilarity and smoothness are not necessary to guarantee the convergence of FedProx with reasonable rates.

Original: • We present FedMSPP as a minibatch stochastic extension of FedProx and analyze itsconvergence behavior in terms of population optimality, again without assumingany type of local dissimilarityconditions.
Revised: • 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.

Original: The main result provably shows that FedMSPP enjoys linear speedup in terms ofminibatching size and partial participation ratio.
Revised: Particularly for smooth problems, our result provably shows that FedMSPP enjoys linear speedup in terms of minibatching size and partial participation ratio

Original: Paper organization.
Revised: Paper organization.

Original: In Section 2 we present our local dissimilarity invariant convergence analysis forthe vanilla FedProx with smooth or non-smooth loss functions. 
Revised: In Section 2 we present our local dissimilarity invariant convergence analysis for the vanilla FedProx with smooth or non-smooth loss functions. In Section 3 we propose FedMSPP as

Original: In Section 3 we propose FedMSPP asa minibatch stochastic extension of FedProx and analyze its convergence behavior throughthe lens ofalgorithmic stability theory.
Revised: a minibatch stochastic extension of FedProx and analyze its convergence behavior through

Original:  The concluding remarks are made in Section 4.
Revised: 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.

Original: Finally, all the technicalproofs and some additional related work are relegated to the appendix sections.
Revised: Finally, all the technical proofs and some additional related work are relegated to the appendix sections.

Original:  Convergence of FedProx
Revised:  Convergence of FedProx

Original: We begin by providing an improved analysis for the vanilla FedProx independent of the localdissimilarity type conditions.
Revised: We begin by providing an improved analysis for the vanilla FedProx which is not relying on the ( B, H ) -LGD type conditions.

Original: We ﬁrst introduce notations that will be used in the analysis to follow.
Revised: We ﬁrst introduce notations that will be used in the analysis to follow.

Original: Notations.
Revised: Notations.

Original: Throughout the paper, we use [ n ] to denote the set { 1 , ..., n } , (cid:107) · (cid:107) to denote the Euclideannorm and (cid:104)· , ·(cid:105) to denote the Euclidean inner product.
Revised: Throughout the paper, we use [ n ] to denote the set { 1 , ..., n } , (cid:107) · (cid:107) to denote the Euclidean norm and (cid:104)· , ·(cid:105) to denote the Euclidean inner product.

Original: We say a function f is G -Lipschitz continuousif | f ( w ) − f ( w (cid:48) )
Revised: We say a function f is G -Lipschitz continuous if | f ( w ) − f ( w (cid:48) )

Original: | ≤ G (cid:107) w − w (cid:48) (cid:107) for all w, w (cid:48)
Revised: | ≤ G (cid:107) w − w (cid:48) (cid:107) for all w, w (cid:48)

Original: ∈ R p , and it is L -smooth if |∇ f ( w ) − ∇ f ( w (cid:48) )
Revised: ∈ R p , and it is L -smooth if |∇ f ( w ) − ∇ f ( w (cid:48) )

Original: | ≤
Revised: | ≤

Original: L (cid:107) w − w (cid:48) (cid:107) for all w, w (cid:48)
Revised: L (cid:107) w − w (cid:48) (cid:107) for all w, w (cid:48)

Original: ∈ R p .
Revised: ∈ R p .

Original: Moreover, we say f is ν -weakly convex
Revised: Moreover, we say f is ν -weakly convex

Original: if for any w, w (cid:48) ∈ R p ,
Revised: if for any w, w (cid:48) ∈ R p ,

Original: [Equation]
Revised: [Equation]

Original: where ∂f ( w (cid:48) ) represents a subgradient of f evaluated at w (cid:48) .
Revised: where ∂f ( w (cid:48) ) represents a subgradient of f evaluated at w (cid:48) .

Original: We denote by
Revised: We denote by

Original: [Equation]

Original: [Equation]

Original: min u

Original: [Equation]
Revised: [Equation]

Original: the η -Moreau-envelope of f , and by
Revised: the η -Moreau-envelope of f , and by

Original: [Equation]
Revised: [Equation]

Original: the proximal mapping associated with f .
Revised: the proximal mapping associated with f .

Original: We also need to access the following deﬁnition of inexactlocal update oracle for FedProx .
Revised: We also need to access the following deﬁnition of inexact local update oracle for FedProx .

Original: Deﬁnition 2 (Local inexact oracle of FedProx ) .
Revised: Deﬁnition 2 (Local inexact oracle of FedProx ) .

Original: Suppose that the local proximal point regularizedobjective Q ( m ) erm ( w ; w t − 1 ) (cf. (3) ) admits a global minimizer.
Revised: Suppose that the local proximal point regularized objective Q ( m ) erm ( w ; w t − 1 ) (cf. (3) ) admits a global minimizer.

Original: For each time instance t , we say thatthe local update oracle of FedProx is ε t -inexactly solved with sub-optimality ε t ≥ 

Revised: For each time instance t , we say that the local update oracle of FedProx is ε t -inexactly solved with sub-optimality ε t ≥

Original: if
Revised: if

Original: [Equation]
Revised: [Equation]

Original: We assume that the objective value gap ¯∆ erm : = ¯ R erm ( w
Revised: Throughout our analysis, we assume that the objective value gap ¯∆ erm :=

Original: [Equation]
Revised: [Equation]

Original: − min w ∈ R p ¯ R [Equation] ) is bounded.
Revised: min w ∈ R p ¯ R erm ( w ) is bounded.

Original: 2.1 Results for Smooth Problems
Revised: 2.1 Results for Smooth Problems

Original: The following theorem is our main result on the convergence rate of FedProx for smooth andnon-convex federated optimization problems.
Revised: The following theorem is our main result on the convergence rate of FedProx for smooth and non-convex federated optimization problems.

Revised: A proof of this result is deferred to Appendix B.1.

Original: Theorem 1.
Revised: Theorem 1.

Original: Assume that for each m ∈
Revised: Assume that for each m ∈

Original: [ M ] , the loss function (cid:96) ( m ) is G -Lipschitz and L -smoothwith respect to its ﬁrst argument.
Revised: [ M ] , the loss function (cid:96) ( m ) is G -Lipschitz and L -smooth with respect to its ﬁrst argument.

Original: Set | I t | ≡ I and η t ≡ 13 L min (cid:110) 1 T 1 / 3 , (cid:113) IT (cid:111) .
Revised: Set | I t | ≡ I and η t ≡ 13 L min (cid:110) 1 T 1 / 3 , (cid:113) IT (cid:111) .

Original: Suppose that the localupdate oracle of FedProx is ε t -inexactly solved with ε t ≤ min (cid:110)
Revised: Suppose that the local update oracle of FedProx is ε t -inexactly solved with ε t ≤ min (cid:110)

Revised: 2 L 2 G 2 η 3 t I 2 ( Lη t +1) ,

Original: G 2 L √ I , Gη t  I  (cid:111) .
Revised: G 2 η t 2 I Lη t +1) (cid:111) .

Original: Let t ∗ be an indexuniformly randomly chosen in { 0 , 1 , ..., T − 1 } .
Revised: Let t ∗ be an index uniformly randomly chosen in { 0 , 1 , ..., T − 1 } .

Original: Then,
Revised: Then it holds that

Original: [Equation]
Revised: [Equation]

Original: (cid:46) (cid:0) L ¯∆ erm + G 2 (cid:1) max (cid:26) 1 T 2 / 3 , 1 √
Revised: [Table] 1 √

Original: TI (cid:27) .
Revised: TI (cid:27) .

Original: Proof.

Original: A proof of this result is deferred to Appendix B.1.

Original: A few remarks are in order.
Revised: A few remarks are in order.

Original: Remark 1.
Revised: Remark 1.

Original: Compared to the O ( 1 T ) bound from Li et al.
Revised: Compared to the O ( 1 T ) bound from Li et al.

Original: [2020b], our rate established in Theorem 1is slower
Revised: (2020b), our rate established in Theoremis slower

Original: but it is valid without assuming the unrealistic ( B, 0) -LGD
Revised: but it is valid without assuming the unrealistic ( B, 0) -LGD

Original: conditions and imposing strongregularization conditions on I [see, e.g., Li et al., 2020b, Remark 5].
Revised: conditions and imposing strong regularization conditions on I (see, e.g., Li et al., 2020b, Remark 5).

Original: Moreover, the dominant term1 √ TI in our bound reveals the beneﬁt of device sampling for linear speedup which is not clear in theoriginal analysis of Li et al.
Revised: Moreover, the dominant term√ TI in our bound reveals the beneﬁt of device sampling for linear speedup which is not clear in the previous analysis by Li et al.

Original: [2020b].
Revised: (2020b).

Original: Remark 2.
Revised: Remark 2.

Original: In the extreme case of full device participation, i.e., I t ≡ [ M ] , the terms related to I in
Revised: In the extreme case of full device participation, i.e., I t ≡ [ M ] , the terms related to I in

Original: Theorem 1 can be removed and thus the convergence rate becomes 1 T 2 / 3 under η t
Revised: Theorem 1 can be removed and thus the convergence rate becomes 1 T 2 / 3 under η t

Original: = O ( 1 LT 1 / 3 ) .
Revised: = O ( 1 LT 1 / 3 ) .

Original: Inthis same setting, we comment that the rate can also be improved to O ( 1 T ) using our proof augmentsif ( B, 0) -LGD is additionally assumed.
Revised: In this same setting, we comment that the rate can also be improved to O ( 1 T ) using our proof augments if ( B, 0) -LGD is additionally assumed.

Original: Remark 3.
Revised: Remark 3.

Original: The G -Lipschitz-loss assumption in Theorem 1 can be alternatively replaced by thebounded gradient condition as commonly used in the analysis of FL algorithms [Li et al., 2020b,
Revised: The G -Lipschitz-loss assumption in Theorem 1 can be alternatively replaced by the bounded gradient condition as commonly used in the analysis of FL algorithms (Li et al., 2020b;

Original: Zhang et al., 2020].
Revised: Zhang et al., 2020).

Original: Despite that our analysis does not explicitly access to any local dissimilarityconditions, the assumed G -Lipschitz (or bounded gradient) condition actually implies that the localobjective gradients are not too dissimilar, which shares a close spirit to the typically assumed (0 , H )-
Revised: Despite that our analysis does not rely on the ( B, 0) -LGD condition, the assumed G -Lipschitz (or bounded gradient) condition actually implies that the local objective gradients are not too dissimilar, which shares a close spirit to the typically assumed (0 , H )

Original: LGD condition [Karimireddy et al., 2020] and inter-client-variance condition [Khanduri et al., 2021].
Revised: -LGD condition (Karimireddy et al., 2020) and inter-client-variance condition (Khanduri et al., 2021).

Original: It is noteworthy that these mentioned client heterogeneity conditions are substantially milder than the( B, 0)
Revised: It is noteworthy that these mentioned client heterogeneity conditions are substantially milder than the ( B, 0)

Original: -LGD condition as required in the original analysis of FedProx .
Revised: -LGD condition as required in the original analysis of FedProx .

Original: 2.2 Results for Non-smooth Problems
Revised: 2.2 Results for Non-smooth Problems

Original: Now we turn to study the convergence of FedProx for weakly convex but not necessarily smoothproblems.
Revised: Now we turn to study the convergence of FedProx for weakly convex but not necessarily smooth problems.

Original: For the sake of presentation clarity, we work on the exact FedProx in which the localupdate oracle is assumed to be exactly solved, i.e. ε t ≡ 0 .
Revised: For the sake of presentation clarity, we work on the exact FedProx in which the local update oracle is assumed to be exactly solved, i.e. ε t ≡ 0 .

Original: Extension to the inexact case is more orless straightforward, though with somewhat more involved perturbation treatment.
Revised: Extension to the inexact case is more or less straightforward, though with somewhat more involved perturbation treatments.

Original: We assume thatthe objective value gap ¯∆ erm ,ρ :=
Revised: In the analysis to follow, we assume that the objective value gap ¯∆ erm ,ρ :=

Original: ¯ R erm ,ρ ( w 0 ) − min w ¯ R erm ,ρ ( w ) associated with ρ -Moreau-envelopeof ¯ R erm is bounded.
Revised: ¯ R erm ,ρ ( w 0 ) − min w ¯ R erm ,ρ ( w ) associated with ρ -Moreau-envelope of ¯ R erm is bounded.

Original: The following is our main result on the convergence of FedProx for non-smoothand weakly convex problems.
Revised: The following is our main result on the convergence of FedProx for non-smooth and weakly convex problems.

Original: Theorem 2.
Revised: Theorem 2.

Original: Assume that for each m ∈
Revised: Assume that for each m ∈

Original: [ M ] , the loss function (cid:96) ( m ) is G -Lipschitz
Revised: [ M ] , the loss function (cid:96) ( m ) is G -Lipschitz

Original: and ν -weaklyconvex with respect to its ﬁrst argument.
Revised: and ν -weakly convex with respect to its ﬁrst argument.

Original: Set η t ≡ ρ √ T for arbitrary ρ < 12 ν .
Revised: Set η t ≡ ρ √ T for arbitrary ρ < 12 ν .

Original: Suppose that the localupdate oracle of FedProx is exactly solved with ε t ≡ 0 .
Revised: Suppose that the local update oracle of FedProx is exactly solved with ε t ≡ 0 .

Original: Let t ∗ be an index uniformly randomlychosen in { 0 , 1 , ..., T − 1 } .
Revised: Let t ∗ be an index uniformly randomly chosen in { 0 , 1 , ..., T − 1 } .

Original: Then it holds that
Revised: Then it holds that

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: The proof technique is inspired by the arguments from Davis and Drusvyatskiy [2019] developed for analyzing stochastic model-based algorithms, with several new elements along developedfor handling the challenges introduced by the model averaging and partial participation mechanismsassociated with FedProx .
Revised: The proof technique is inspired by the arguments from Davis and Drusvyatskiy (2019) developed 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 .

Original: A particular crux here is that due to the random subset model aggregationof w t = 1 | I t | (cid:80) ξ ∈
Revised: A particular crux here is that due to the random subset model aggregation of w t = 1 | I t | (cid:80) ξ ∈

Original: I t w ( ξ ) t , the local function values R ( ξ ) erm ( w t ) are no longer independent to each otherthough ξ is uniformly random.
Revised: I t w ( ξ ) t , the local function values R ( ξ ) erm ( w t ) are no longer independent of each other though ξ is uniformly random.

Original: As a consequence, 1 | I t | (cid:80) ξ ∈
Revised: As a consequence, 1 | I t | (cid:80) ξ ∈

Original: I t R ( ξ ) erm ( w t ) is not an unbiased estimationof ¯ R erm ( w t ) .
Revised: I t R ( ξ ) erm ( w t ) is not an unbiased estimation of ¯ R erm ( w t ) .

Original: To overcome this technical obstacle, we make use of a key observation that w ( m ) twill be almost surely close enough to w t − 1 if the learning rate η t is small enough (which is thecase in our choice of η t ), and thus we can replace the former with the latter whenever beneﬁcialbut without introducing too much approximation error.
Revised: To overcome this technical obstacle, we make use of a key observation that w ( m ) t will be almost surely close enough to w t − 1 if the learning rate η t is small enough (which is the case in our choice of η t ), and thus we can replace the former with the latter whenever beneﬁcial but without introducing too much approximation error.

Original: A full proof of this result can be found in
Revised: A full proof of this result can be found in

Original: Appendix B.2.
Revised: Appendix B.2.

Original: A few comments are in order.
Revised: A few comments are in order.

Original: Remark 4.
Revised: Remark 4.

Original: To our best knowledge, Theorem 2 is the ﬁrst convergence guarantee for FL algorithmsapplicable to generic non-smooth and weakly convex problems.
Revised: To our best knowledge, Theorem 2 is the ﬁrst convergence guarantee for FL algorithms applicable to generic non-smooth and weakly convex problems.

Original: This is in sharp contrast with
Revised: This is in sharp contrast with

Original: FCO [Yuan et al., 2021] which focuses on composite convex and non-smooth problems such as (cid:96) 1estimation, or Fed-HT
Revised: FCO (Yuan et al., 2021) which focuses on composite convex and non-smooth problems such as (cid:96) 1estimation, or Fed-HT

Original: [Tong et al., 2020] which is specially customized for cardinality-constrainedsparse learning problems where the non-convexity essentially arises from the (cid:96) 0 -constraint.
Revised: (Tong et al., 2020) which is specially customized for cardinality-constrained sparse learning problems where the non-convexity essentially arises from the cardinality constraint.

Original: Remark 5.
Revised: Remark 5.

Original: Let us consider ¯ w t ∗ := prox ρ ¯ R erm ( w t ∗ ) , the proximal mapping of w t ∗ associated with¯ R erm .
Revised: Let us consider ¯ w t ∗ := prox ρ ¯ R erm ( w t ∗ ) , the proximal mapping of w t ∗ associated with ¯ R erm .

Original: In view of a feature of Moreau envelope to characterize stationarity [Davis and Drusvyatskiy,2019], if w t ∗ has small gradient norm (cid:13)(cid:13) ∇ ¯ R erm ,ρ ( w t ∗ ) (cid:13)(cid:13) , then ¯ w t ∗ must be a near-stationary solutionand w t ∗ stays in the proximity of ¯ w t ∗ due to the identity (cid:107) w t ∗ − ¯ w t ∗ (cid:107) =
Revised: In view of a feature of Moreau envelope to characterize stationarity (Davis and Drusvyatskiy, 2019), if w t ∗ has small gradient norm (cid:13)(cid:13) ∇ ¯ R erm ,ρ ( w t ∗ ) (cid:13)(cid:13) , then ¯ w t ∗ must be a near-stationary solution and w t ∗ stays in the proximity of ¯ w t ∗ due to the identity (cid:107) w t ∗ − ¯ w t ∗ (cid:107) =

Original: ρ (cid:13)(cid:13) ∇ ¯ R erm ,ρ ( w t ∗ ) (cid:13)(cid:13) .
Revised: ρ (cid:13)(cid:13) ∇ ¯ R erm ,ρ ( w t ∗ ) (cid:13)(cid:13) .

Original: Therefore,the bound in Theorem 2 suggests that in expectation ¯ w t ∗ converges to a stationary solution and w t ∗converges to ¯ w t ∗ , both at the rate of O ( 1 √ T ) .
Revised: Therefore, the bound in Theorem 2 suggests that in expectation ¯ w t ∗ converges to a stationary solution and w t ∗ converges to ¯ w t ∗ , both at the rate of O ( 1 √ T ) .

Original: Remark 6.
Revised: Remark 6.

Original: We comment that the bound in Theorem 2 is not dependent on I , the number of selecteddevices.
Revised: We comment that the bound in Theorem 2 is not dependent on I , the number of selected devices.

Original: For I 
Revised: On one hand, for I =

Original: = O (1) and sufﬁciently large T >
Revised: O (1) and sufﬁciently large T >

Original: O
Revised: O

Original: ( I 3 ) , the bounds Theorem 1 and Theorem 2 arecomparable to each other, which demonstrates that the smoothness is not must-have for FedProx toget sharper convergence bound with small device sampling rate.
Revised: ( I 3 ) , the bounds Theorem 1 and Theorem 2 are comparable to each other, which demonstrates that the smoothness is not must-have for FedProx to get sharper convergence bound with small device sampling rate.

Original: However, in the near-full participation setting where I =
Revised: On the other hand, in the near-full participation setting where I =

Original: O
Revised: O

Original: ( M ) , the bound in Theorem 2 for non-smooth problems will be slower when M is large.
Revised: ( M ) , the bound in Theorem 2 for non-smooth problems will be slower when M is large.

Original: Extremely when I t =
Revised: Extremely when I t =

Original: [ M ] , the O ( 1 √ T ) bound is substantially inferior to the smooth casewhich has improved rate of O ( 1 T 2 / 3 ) as discussed in Remark 2.
Revised: [ M ] , the O ( 1 √ T ) bound is substantially inferior to the smooth case which has improved rate of O ( 1 T 2 / 3 ) as discussed in Remark 2.

Original:  Convergence of FedProx with Stochastic Minibatching
Revised:  Convergence of FedProx with Stochastic Minibatching

Original: When it comes to the implementation of FedProx , a notable challenge is that the local proximalpoint update oracle (3) is by itself a full-batch ERM problem which would be expensive to solveeven approximately in large-scale settings.
Revised: When it comes to 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.

Original: Moreover, in the settings where the data distribution overdevices is highly imbalanced, the computational overload of local update could vary signiﬁcantlyacross the network, which impairs communication efﬁciency.
Revised: Moreover, in the settings where the data distribution over devices is highly imbalanced, the computational overload of local update could vary signiﬁcantly across the network, which impairs communication efﬁciency.

Original: It is thus desirable to seek stochasticapproximation schemes for hopefully improving the local oracle update efﬁciency and overloadbalance of FedProx .
Revised: It is thus desirable to seek stochastic approximation schemes for hopefully improving the local oracle update efﬁciency and overload balance of FedProx .

Original: To this end, inspired by the recent success of minibatch stochastic proximal pointmethods (MSPP)
Revised: To this end, inspired by the recent success of minibatch stochastic proximal point methods (MSPP)

Original: [Asi et al., 2020, Deng and Gao, 2021], we propose to implement FedProx using
Revised: (Asi et al., 2020; Deng and Gao, 2021), we propose to implement FedProx using

Original: MSPP as the local stochastic optimization oracle.
Revised: MSPP as the local stochastic optimization oracle.

Original: More precisely, let B ( m ) t= { z ( m ) i,t } bi =1 i.i.d.
Revised: More precisely, let B ( m ) t

Original: ∼ ( D ( m ) ) bbe a minibatch of b i.i.d.
Revised: [Equation] be a minibatch of b i.i.d.

Original: samples drawn from the distribution D ( m ) at device m and time instancet ≥ 1 .
Revised: samples drawn from the distribution D ( m ) at device m and time instance t ≥ 1 .

Original: We denote
Revised: We denote

Original: [Equation]
Revised: [Equation]

Original: Input :Minibatch size b ; learning rates { γ t } t ∈ [ T ] .
Revised: Input :Minibatch size b ; learning rates { γ t } t ∈ [ T ] .

Original: Output : w T .
Revised: Output : w T .

Original: Initialization Set w 0 , e.g., typically as a zero vector.
Revised: Initialization Set w 0 , e.g., typically as a zero vector.

Original: for t = 1 , 2 , ..., T do /
Revised: for t = 1 , 2 , ..., T do /

Original: Device selection and model broadcast on the server */
Revised: Device selection and model broadcast on the server */

Original: Server uniformly randomly selects a subset I t ⊆
Revised: Server uniformly randomly selects a subset I t ⊆

Original: [ M ] of devices and sends w t − 1 to all the selected devices; /* Local model updates on the selected devices */ for ξ ∈
Revised: [ M ] of devices and sends w t − 1 to all the selected devices; /* Local model updates on the selected devices */ for ξ ∈

Original: I t in parallel do
Revised: I t in parallel do

Original: Device ξ samples a minibatch B ( ξ ) t
Revised: Device ξ samples a minibatch B ( ξ ) t

Original: [Equation]
Revised: [Equation]

Original: Device ξ inexactly updates the its local model as
Revised: Device ξ inexactly updates the its local model as

Original: [Equation]
Revised: [Equation]

Original: where R ( ξ ) B ( ξ ) t ( w ) is given by (7).
Revised: where R ( ξ ) B ( ξ ) t ( w ) is given by (7).

Original: Device ξ sends w ( ξ ) t back to server.
Revised: Device ξ sends w ( ξ ) t back to server.

Original: end /*
Revised: end /*

Original: Model aggregation on the server */
Revised: Model aggregation on the server */

Original: Sever aggregates the local models received from I t to update the global model as
Revised: Sever aggregates the local models received from I t to update the global model as

Original: [Equation]
Revised: [Equation]

Original: end as the local minibatch empirical risk function over B ( m ) t .
Revised: end as the local minibatch empirical risk function over B ( m ) t .

Original: The only modiﬁcation we propose tomake here is to replace the empirical risk R ( m ) erm ( w ) in the original update form (3) with its minibatchcounterpart R ( m ) B ( m ) t ( w ) .
Revised: Here, the only modiﬁcation we propose to make is to replace the empirical risk R ( m ) erm ( w ) in the original update form (3) with its minibatch counterpart R ( m ) B ( m ) t ( w ) .

Original: The resultant FL framework, which we refer to as FedMSPP (Federated
Revised: The resultant FL framework, which we refer to as FedMSPP (Federated

Original: MSPP), is outlined in Algorithm 1.
Revised: MSPP), is outlined in Algorithm 1.

Original: Clearly, the vanilla FedProx is a special case of FedMSPP whenapplied to the federated ERM form (2) with full data batch B ( m ) t ≡ D ( m ) .
Revised: Clearly, the vanilla FedProx is a special case of FedMSPP when applied to the federated ERM form (2) with full data batch B ( [Equation]

Original: 3.1 Results for Smooth Problems
Revised: 3.1 Results for Smooth Problems

Original: We ﬁrst analyze the convergence rate of FedMSPP for smooth and non-convex problems using thetools borrowed from algorithmic stability theory.
Revised: We ﬁrst analyze the convergence rate of FedMSPP for smooth and non-convex problems using the tools borrowed from algorithmic stability theory.

Original: Analogous to the Deﬁnition 2, we introduce thefollowing deﬁnition of inexact local update oracle for FedMSPP .
Revised: Analogous to the Deﬁnition 2, we introduce the following deﬁnition of inexact local update oracle for FedMSPP .

Original: Deﬁnition 3 (Local inexact oracle of FedMSPP ) .
Revised: Deﬁnition 3 (Local inexact oracle of FedMSPP ) .

Original: Suppose that the local proximal point regularizedobjective Q ( m ) B [Equation] t ( w ; w t − 1 ) (cf.
Revised: Suppose that the local proximal point regularized objective Q ( m ) B ( m ) t ( w ; w t − 1 ) (cf.

Original: (8) ) admits a global minimizer.
Revised: (8) ) admits a global minimizer.

Original: For each time instance t , we say thatthe local update oracle of FedMSPP is ε t -inexactly solved with sub-optimality ε t ≥ 

Revised: For each time instance t , we say that the local update oracle of FedMSPP is ε t -inexactly solved with sub-optimality ε t ≥

Original: if
Revised: if

Original: [Equation]
Revised: [Equation]

Original: We also assume that the population value gap ¯∆
Revised: We also assume that the population value gap ¯∆

Original: =
Revised: =

Original: ¯ R ( w (0) )
Revised: ¯ R ( w (0) )

Original: − min w ∈ R p ¯ R ( w ) is bounded.
Revised: − min w ∈ R p ¯ R ( w ) is bounded.

Original: Thefollowing theorem is our main result on FedMSPP for smooth and non-convex FL problems.
Revised: The following theorem is our main result on FedMSPP for smooth and non-convex FL problems.

Original: Theorem 3.
Revised: Theorem 3.

Original: Assume that for each m ∈
Revised: Assume that for each m ∈

Original: [ M ] , the loss function (cid:96) ( m ) is G -Lipschitz and L -smooth withrespect to its ﬁrst argument.
Revised: [ M ] , the loss function (cid:96) ( m ) is G -Lipschitz and L -smooth with respect to its ﬁrst argument.

Original: Set | I t | ≡ I and η t ≡ 18 L min (cid:26)
Revised: Set | I t | ≡ I and η t ≡ 18 L min [Equation]

Original: T 1 / 3 , (cid:113) bIT (cid:27) .
Revised: [Equation] (cid:113) bIT (cid:27) .

Original: Suppose that the localupdate oracle of FedMSPP is ε t -inexactly solved with ε t ≤ min (cid:110)
Revised: Suppose that the local update oracle of FedMSPP is ε t -inexactly solved with ε t ≤ min (cid:110)

Original: G 2 L , G 2 η t 8 b 2 , Gη t 2 bI  (cid:111) .
Revised: G 2 η t 2( Lη t +1) , G 2 η t 8 b 2 , L 2 G 2 η 3 t 2 bI ( Lη t +1) (cid:111) .

Revised: Let t ∗ be an index uniformly randomly chosen in { 0 , 1 , ..., T − 1 } .

Original: Then we have
Revised: Then it holds that

Original: T T − 1 (cid:88) t =
Revised: [Equation] [Table] [Equation]

Original: E (cid:2) (cid:107)∇ ¯ R ( w t ) (cid:107) 2 (cid:3) (cid:46) (cid:0) L ¯∆ + G 2 (cid:1) max (cid:26) 1 T 2 / 3 , 1 √ TbI
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Let us consider d ( m
Revised: Let us consider d ( m

Original: [Equation]
Revised: [Equation]

Original: t
Revised: t

Original: [Equation]
Revised: [Equation]

Original: ) which is roughly the local update direction on devicem , in the sense that w ( m ) t ≈ w t − [Equation] − 
Revised: ) which is roughly the local update direction on device m , in the sense that w ( m ) t ≈ w t − 1 − η

Original: [Equation]
Revised: [Equation]

Original: d (
Revised: d (

Original: [Equation]
Revised: [Equation]

Original:  given that the local update oracle is solved to sufﬁcientaccuracy.
Revised: t given that the local update oracle is solved to sufﬁcient accuracy.

Original: As a key ingredient of our proof, we show via some extended uniform stability argumentsin terms of gradients (see Lemma 3) that the averaged directions d t := 1 | I t | (cid:80) ξ ∈
Revised: As a key ingredient of our proof, we show via some extended uniform stability arguments in terms of gradients (see Lemma 3) that the averaged directions d t := 1 | I t | (cid:80) ξ ∈

Original: I t d ( ξ ) t aligns wellwith the global gradient ∇ ¯ R ( w t − 1 ) in expectation (see Lemma 11).
Revised: I t d ( ξ ) t aligns well with the global gradient ∇ ¯ R ( w t − 1 ) in expectation (see Lemma 11).

Original: Therefore, in average it roughlyholds that w t = 1 | I t | (cid:80) ξ ∈ 
Revised: Therefore, in average it roughly holds that w t = 1 | I t | (cid:80) ξ ∈ I t w ( ξ ) t ≈

Original: I
Revised: [Equation]

Original: t w (ξ ) t
Revised: R (

Original: [Equation]
Revised: [Equation]

Original: R
Revised: t

Original: [Equation]
Revised: [Equation]

Original: 1 ) , which suggests that w t isupdated roughly along the direction of global gradient descent and thus guarantees quick convergence.
Revised: 1 ) , which suggests that w t is updated roughly along the direction of global gradient descent and thus is expected to converge quickly.

Original: Based on this novel analysis, we are free of  imposingany kind of local dissimilarity  conditions onlocal objectives.
Revised: Based on this novel analysis, we are free of explicitly imposing local dissimilarity type conditions on local objectives.

Original: See Appendix C.1 for a full proof of this result.
Revised: See Appendix C.1 for a full proof of this result.

Original: Remark 7.
Revised: Remark 7.

Original: For T ≥ ( bI ) 3 , the bound in Theorem 3 is dominated by O ( 1 √ TbI ) which gives thecommunication complexity 1 bI(cid:15) 2 .
Revised: For T ≥ ( bI ) 3 , the bound in Theorem 3 is dominated by O ( 1 √ TbI ) which gives the communication complexity 1 bI(cid:15) 2 .

Original: This shows that FedMSPP enjoys linear speedup  both in the size oflocal minibatching andin the size of device sampling.
Revised: This shows that FedMSPP enjoys linear speedup with respect to both local minibatching and device sampling sizes.

Original: Remark 8.
Revised: Remark 8.

Original: While the bound in Theorem 3 is derived for the population form of FL in (1) , identicalbound naturally holds for the empirical form (2) under minibatch sampling according to local dataempirical distribution.
Revised: 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.

Original: 3.2 Results for Non-smooth Problems
Revised: 3.2 Results for Non-smooth Problems

Original: Analogues to FedProx , we can further show that FedMSPP converges reasonably well when appliedto weakly convex and non-smooth problems.
Revised: Analogues to FedProx , we can further show that FedMSPP converges reasonably well when applied to weakly convex and non-smooth problems.

Original: We assume that the objective value gap ¯∆ ρ :
Revised: In the analysis to follow, we assume that the objective value gap ¯∆ ρ :

Original: =
Revised: =

Original: ¯ R ρ ( w 0 ) − min w ¯ R ρ ( w ) associated with ρ -Moreau-envelope of ¯ R is bounded.
Revised: ¯ R ρ ( w 0 ) − min w ¯ R ρ ( w ) associated with ρ -Moreau-envelope of ¯ R is bounded.

Original: The following is ourmain result in this line.
Revised: The following is our main result in this line.

Original: Theorem 4.
Revised: Theorem 4.

Original: Assume that for each m ∈
Revised: Assume that for each m ∈

Original: [ M ] , the loss function (cid:96) ( m ) is G -Lipschitz
Revised: [ M ] , the loss function (cid:96) ( m ) is G -Lipschitz

Original: and ν -weaklyconvex with respect to its ﬁrst argument.
Revised: and ν -weakly convex with respect to its ﬁrst argument.

Original: Set η t ≡ ρ √ T for arbitrary ρ < 12 ν .
Revised: Set η t ≡ ρ √ T for arbitrary ρ < 12 ν .

Original: Suppose that the localupdate oracle of FedMSPP is exactly solved with ε t ≡ 0 .
Revised: Suppose that the local update oracle of FedMSPP is exactly solved with ε t ≡ 0 .

Original: Let t ∗ be an index uniformly randomlychosen in { 0 , 1 , ..., T − 1 } .
Revised: Let t ∗ be an index uniformly randomly chosen in { 0 , 1 , ..., T − 1 } .

Original: Then it holds that
Revised: Then it holds that

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: The proof argument is a slight adaptation of that of Theorem 2 to the population FL setup (1)with FedMSPP .
Revised: The proof argument is a slight adaptation of that of Theorem 2 to the population FL setup (1) with FedMSPP .

Original: For the sake of completeness, a full proof is reproduced in Appendix C.2.
Revised: For the sake of completeness, a full proof is reproduced in Appendix C.2.

Original: We comment in passing that the discussions made in Remarks 4-6 immediately extend to Theorem 4.  
Revised: We comment in passing that the discussions made in Remarks 4-6 about Theorem 2 extend to Theorem 4.  Additional Related Work

Revised: The present work is situated at the intersection of federated learning, stochastic proximal point optimization and algorithmic stability theory.

Original: Conclusions
Revised: We next brieﬂy review some additional work in these lines of research that are closely related to ours. Heterogenous federated learning. The presence of device heterogeneity features a key distinction between FL and classic distributed learning. The most commonly used FL method is FedAvg (McMahan et al., 2017), where the local update oracle is formed as multi-epoch SGD.

Revised: FedAvg was early analyzed for identical functions (Stich, 2018; Stich and Karimireddy, 2020) under the name of local

Revised: SGD.

Revised: 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., 2019b;

Revised: Woodworth et al., 2020;

Revised: Khaled et al., 2020; Reddi et al., 2020; Khanduri et al., 2021).

Revised: As another representative FL method,

Revised: FedProx (Li et al., 2020b) has recently been proposed to apply averaged proximal point updates to solve heterogeneous federated minimization problems.

Revised: 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 Deﬁnition 1) to measure data heterogeneity (Li et al., 2020b; Pathak and

Revised: Wainwright, 2020; Nguyen et al., 2020).

Revised: This assumption was relaxed by FedPD (Zhang et al., 2020) inside a meta-framework of primal-dual optimization.

Revised: The SCAFFOLD (Karimireddy et al., 2020) and

Revised: 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.

Revised: For composite non-smooth FL problems, the FCO proposed in Yuan et al.

Revised: (2021) employs a server dual averaging procedure to circumvent the curse of primal averaging suffered by FedAvg .

Revised: 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.

Revised: Minibatch stochastic proximal point methods.

Revised: The proposed FedMSPP algorithm is a variant of

Revised: 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.

Revised: 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;

Revised: Wang et al., 2017; Asi and Duchi, 2019; Deng and Gao, 2021).

Revised: However, it is not yet known if FedProx or FedMSPP can achieve similar strong guarantees for non-smooth heterogenous FL problems.

Revised: Algorithmic stability.

Revised: Our analysis for FedMSPP builds largely upon the classic algorithmic stability theory.

Revised: Since the seminal work of Bousquet and Elisseeff (Bousquet and Elisseeff, 2002), algorithmic stability has been serving as a powerful proxy for establishing strong generalization bounds (Zhang, 2003; Mukherjee et al., 2006; Shalev-Shwartz et al., 2010).

Revised: Particularly, the state-of-the-art risk bounds of strongly convex ERM are offered by approaches based on the notion of uniform stability (Feldman and Vondrak, 2018, 2019; Bousquet et al., 2020; Klochkov and Zhivotovskiy, 2021).

Revised: It was shown by Hardt et al.

Revised: that the solution obtained via SGD is stable for smooth convex or non-convex loss functions.

Revised: For non-smooth convex losses, the stability induced generalization bounds have been established for SGD (Lei and Ying, 2020; Bassily et al., 2020).

Revised: Through the lens of uniform algorithmic stability, convergence rates of MSPP have been studied for non-smooth and convex (Wang et al., 2017), or weakly convex (Deng and Gao, 2021) losses.

Revised:  Conclusions

Original: In this paper, we have exposed three shortcomings of the prior analysis for FedProx in unrealisticassumptions about local dissimilarity, inapplicability to non-smooth problems and expensive (andpotentially imbalanced) computational cost of local update.
Revised: 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.

Original: In order to tackle these issues, wedeveloped a novel convergence theory for the vanilla FedProx and its minibatch stochastic variant,
Revised: In order to tackle these issues, we developed a novel convergence theory for the vanilla FedProx and its minibatch stochastic variant,

Original: FedMSPP , through the lens of algorithmic stability theory.
Revised: FedMSPP , through the lens of algorithmic stability theory.

Original: At nutshell, our results reveal that withminimal modiﬁcations, FedProx is able to kill three birds with one stone: it enjoys favorable ratesof convergence which are simultaneously invariant to  local dissimilarity, applicable to smooth ornon-smooth problems, and scaling linearly with respect to local minibatch size and device samplingratio for smooth problems.
Revised: In a nutshell, our results reveal that with minimal modiﬁcations, 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-smooth problems, and scaling linearly with respect to local minibatch size and device sampling ratio for smooth problems.

Original: To the best of our knowledge, the present work is the ﬁrst theoreticalcontribution that achieves all these appealing properties in a single FL framework.
Revised: To the best of our knowledge, the present work is the ﬁrst theoretical contribution that achieves all these appealing properties in a single FL framework.

Original: Limitations.
Revised: Limitations.

Original: While our results in Theorem 2 and Theorem 4 for the ﬁrst time guarantee that FedProx and FedMSPPconverge asymptotically for non-smooth and weak-convex problems, thecorresponding rates of convergence so far cannot demonstrate any linear speedup effort with respectto device sampling ratio and local minibatch size.
Revised: While our results in Theorem 2 and Theorem 4 for the ﬁrst time guarantee the nonasymptotic convergence of FedProx and FedMSPP for non-smooth and weakly-convex problems, the corresponding rates of convergence so far cannot demonstrate any linear speedup effort with respect to device sampling ratio and local minibatching size.

Original: This is as opposed to what have been shown forsmooth problems in Theorem 1 and Theorem 3, and thus we view it as a limitation of  our analysis.
Revised: This is as opposed to what have been shown for smooth problems in Theorem 1 and Theorem 3, and thus we view it as a limitation of the techniques used by our analysis.

Original: Also for non-smooth FL problems, the current convergence results are only limited to the ideal casethat the  local update oracles are exactly solved.
Revised: In the smooth-loss case, the comparison in Table 1 suggests that our results in Theorem 1 and Theorem 3 show no 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. Acknowledgments and Disclosure of Funding

Original: We leave a more detailed understanding of how theerrors in the local oracle update would affect the ﬁnal result for future investigation.
Revised: The authors sincerely thank the anonymous reviewers for their insightful comments on this paper. U21B

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?

Original: [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.
Revised: [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.

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] We have provided a remark on the limitations of our theory at the end of the main paper.
Revised: [Yes] We have provided a remark on the limitations of our theory at the end of the main paper.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A]
Revised: [N/A]

Original: Our contribution is theoretical in nature.
Revised: Our contribution is theoretical in nature.

Original: As far as we are aware of, there are no foreseeable societal or ethical consequences for the present research.
Revised: As far as we are aware of, there are no foreseeable societal or ethical consequences for the present research.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] The keyassumptions on the structure 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, 3, 4.
Revised: [Yes] The key assumptions on the structure 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, 3, 4.

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] Due to space limit, all the full proofs of results are relegated to the appendix sections which can be found in the supplementary document.
Revised: [Yes] Due to space limit, all the full proofs of results are relegated to the appendix sections which can be found in the supplementary document.

Original: Particularly for Theorem 2 and Theorem 3, we further provide sketched proofs in the main submission to highlight the proof road maps.
Revised: Particularly for Theorem 2 and Theorem 3, we further provide sketched proofs in the main submission to highlight the proof road maps.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: [N/A]
Revised: [N/A]

Original: Our work is focused on providing deeper theoretical understandings of FedProx and its stochastic variants under milder conditions.
Revised: Our work is focused on providing deeper theoretical understandings of FedProx and its stochastic variants under milder conditions.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: [N/A] We did not use any of such assets in this work.
Revised: [N/A] We did not use any of such assets in this work.

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: [N/A]
Revised: [N/A]

Original: The present research was carried out with no crowdsourcing or human subjects involved in.
Revised: The present research was carried out with no crowdsourcing or human subjects involved in.