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: 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 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.
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 develop a 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, 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
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

Original: FL problems; and 3) linear speedup with respect to size of minibatch and number of sampled devices.
Revised: 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 and smoothness 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ﬁcient distributed learning on remote devices, such as smartphones, internet of things, or agents (Kone ˇ cn ` y et al., 2016;
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;

Original: Yang et al., 2019).
Revised: Yang et al., 2019).

Original: 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.
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ﬁcial for 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 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.
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 highly demanding to tackle, both in terms of optimization algorithm design and in terms of theoretical understanding 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 minimization averaged 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] [Equation]
Revised: [Equation] [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-negative loss function whose value (cid:96) ( w ; Z ( m ) ) measures the loss over a random data point Z
Revised: ) 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 36th
Revised: Z ( m ) with 36th

Original: Neural
Revised: Neural

Original: Systems (NeurIPS 2022).
Revised: Systems (NeurIPS 2022).

Original: parameter w , D ( m ) represents an underlying random data distribution over Z ( m ) .
Revised: parameter w , D ( m ) represents an underlying random data distribution over Z ( m ) .

Original: Since the data distribution is typically unknown, the following empirical risk minimization (ERM) version of (1) is often 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: The sample size N m may vary signiﬁcantly across devices, which can be regarded as another source of data 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ﬁ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.
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), to name 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 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.
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 for heterogeneous 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 implementation simplicity, low communication cost, promise in dealing with data heterogeneity and tolerance to partial participation 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, expose problems, 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 contribute to derive several new and deeper theoretical insights into the algorithm from a novel perspective of algorithmic stability and generalization 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 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).
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 point update 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 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).
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: More speciﬁ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 ] of devices and introduces for each device ξ ∈
Revised: [ M ] of devices and introduces for each device ξ ∈

Original: I t the following proximal point ERM sub-problem for local 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 model is 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ﬁ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).
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 guarantees under the following bounded local gradient dissimilarity assumption that captures the statistical heterogeneity 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 above 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 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)
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 also been 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, restrictive problem 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 condition .
Revised: • Local dissimilarity condition .

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

Original: -LGD condition (cf.
Revised: -LGD condition (cf.

Original: Deﬁnition 1) with B > 0 .
Revised: Deﬁnition 1) with B > 0 .

Original: 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) = 

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 for all 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, there would 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 to understand the convergence behavior of FedProx for heterogeneous FL without imposing stringent local dissimilarity conditions like ( B, 0) -LGD.
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 available for 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-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.
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 is desirable to understand the convergence behavior of FedProx in non-smooth problem regimes.
Revised: Therefore, it is desirable to understand the convergence behavior of FedProx in non-smooth problem regimes.

Original: • Local oracle complexity .
Revised: • 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 to be 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 imbalanced data distribution over devices, the computational overload of the proximal point oracle could vary 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 minibatch stochastic approximation to the proximal point oracle (3) can provably improve the computational efﬁ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 FL problem (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 a more 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 for non-convex heterogenous FL under weaker yet 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 set of 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 local dissimilarity assumption, we provide a novel convergence analysis for the vanilla FedProx algorithm invariant to the ( B, 0) -LGD condition.
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 devices participate 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 improved to 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 communication complexity 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 is dominated 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 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.
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
Revised: Further for non-smooth and non-convex problems, we establish in

Original: Theorem 2 the following rate of convergence
Revised: Theorem 2 the following rate of convergence

Original: [Equation]
Revised: [Equation]

Original: which is not dependent on 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 end when 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 smooth functions 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 optimization performance 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 recent success 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 local update 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 improved trade-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 case of imbalanced data distribution, minibatching is also beneﬁcial for making the local computation more 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 following rate of convergence for FedMSPP in terms of population optimality, which is also invariant to the ( B, 0)
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)

Original: -LGD condition:
Revised: -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 sampling according 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 ) which gives 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ﬁciently large 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 speeding up 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 can 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 witnessing the effect of linear 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 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.
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 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.
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 convergence guarantees 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-smooth problems, it is customized for federated convex composite optimization with potentially non-smooth regularizers (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 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 .
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 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
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, 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.
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: [Table]
Revised: [Table]

Original: Work
Revised: Work

Original: Commun.
Revised: Commun.

Original: Complex.
Revised: Complex.

Original: LD Condition
Revised: LD Condition

Original: NS
Revised: NS

Original: PP
Revised: PP

Original: [Table]
Revised: [Table]

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 dissimilarity invariant convergence guarantees for the widely used FedProx algorithm for non-convex heterogeneous 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 that local dissimilarity and smoothness are not necessary to guarantee the convergence of FedProx with reasonable 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 its population optimization performance for both smooth and non-smooth FL problems, again without assuming local dissimilarity conditions.
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: Particularly for smooth problems, our result provably shows that FedMSPP enjoys linear speedup in terms of minibatching 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 for the 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.

Original: In Section 3 we propose FedMSPP as
Revised: In Section 3 we propose FedMSPP as

Original: a minibatch
Revised: a minibatch

Original: stochastic
Revised: stochastic

Original: extension of
Revised: extension of

Original: FedProx and analyze its convergence behavior through
Revised: FedProx and analyze its convergence behavior through

Original: the lens
Revised: the lens

Original: of algorithmic stability theory.
Revised: of algorithmic stability theory.

Original: In Section 4, we discuss some additional related work on the topics covered by this paper.
Revised: In Section 4, we discuss some additional related work on the topics covered by this paper.

Original: The concluding remarks are made in Section 5.
Revised: The concluding remarks are made in Section 5.

Original: Finally, all the technical proofs 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 which is not relying on the ( B, H ) -LGD 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 Euclidean norm 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 continuous if | 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]
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 inexact local 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 regularized objective Q ( m ) erm ( w ; w t − 1 ) (cf.
Revised: Suppose that the local proximal point regularized objective Q ( m ) erm ( w ; w t − 1 ) (cf.

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

Original: For each time instance t , we say that the 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: Throughout our analysis, we assume that the objective value gap ¯∆ erm :=
Revised: Throughout our analysis, we assume that the objective value gap ¯∆ erm :=

Original: [Equation]
Revised: [Equation]

Original: min w ∈ R p ¯ R erm ( w ) 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 and non-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.

Original: A proof of this result is deferred to Appendix B.1.
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 -smooth with 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 local update 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)

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

Original: G 2 η t 2 I
Revised: G 2 η t 2 I

Original: Lη t +1) (cid:111) .
Revised: Lη t +1) (cid:111) .

Original: Let t ∗ be an index uniformly randomly chosen 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: [Table]
Revised: [Table]

Original: 1 √
Revised: 1 √

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

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 Theoremis 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 strong regularization 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 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.
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: 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.
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 the bounded 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 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 )
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 smooth problems.
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 local update 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 or less straightforward, though with somewhat more involved perturbation treatments.
Revised: Extension to the inexact case is more or less straightforward, though with somewhat more involved perturbation treatments.

Original: In the analysis to follow, we assume that the 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-envelope of ¯ 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
Revised: The following is our main result on the convergence of

Original: FedProx for non-smooth and weakly convex problems.
Revised: 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 ν -weakly convex 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 local update 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 randomly chosen 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 developed for handling the challenges introduced by the model averaging and partial participation mechanisms associated 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 aggregation of 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 of each other though ξ 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 estimation of ¯ 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 ) 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.
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 algorithms applicable 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-constrained sparse learning problems where the non-convexity essentially arises from the cardinality 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 ¯
Revised: Let us consider ¯ w t ∗ := prox ρ ¯ R erm ( w t ∗ ) , the proximal mapping of w t ∗ associated with ¯

Original: R erm .
Revised: 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 solution and 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 selected devices.
Revised: We comment that the bound in Theorem 2 is not dependent on I , the number of selected devices.

Original: On one hand, 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
Revised: ( I 3 ) , the bounds Theorem 1 and

Original: Theorem 2 are comparable to each other, which demonstrates that the smoothness is not must-have for
Revised: Theorem 2 are comparable to each other, which demonstrates that the smoothness is not must-have for

Original: FedProx to get sharper convergence bound with small device sampling rate.
Revised: FedProx to get sharper convergence bound with small device sampling rate.

Original: On the other hand, 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 case which 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 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.
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 over devices is highly imbalanced, the computational overload of local update could vary signiﬁcantly across 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 stochastic approximation schemes for hopefully improving the local oracle update efﬁciency and overload balance 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 point methods (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
Revised: More precisely, let B ( m ) t

Original: [Equation]
Revised: [Equation]

Original: be a minibatch of b i.i.d.
Revised: be a minibatch of b i.i.d.

Original: samples drawn from the distribution D ( m ) at device m and time instance t ≥ 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: 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 ) .
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 when applied to the federated ERM form (2) with full data batch B (
Revised: Clearly, the vanilla FedProx is a special case of FedMSPP when applied to the federated ERM form (2) with full data batch B (

Original: [Equation]
Revised: [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 the tools 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 the following 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 regularized objective Q ( m ) B ( m ) 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 that the 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: The following 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 with 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 ≡ 18 L min
Revised: Set | I t | ≡ I and η t ≡ 18 L min

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: (cid:113) bIT (cid:27) .
Revised: (cid:113) bIT (cid:27) .

Original: Suppose that the local update 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 η 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: 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) .

Original: Let t ∗ be an index uniformly randomly chosen 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: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
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 device m , in the sense that w ( m ) t ≈ w t − 1 − η
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: t given that the local update oracle is solved to sufﬁcient accuracy.
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 arguments in 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 well with 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 roughly holds that w t = 1 | I t | (cid:80) ξ ∈ I t w ( ξ ) t ≈
Revised: Therefore, in average it roughly holds that w t = 1 | I t | (cid:80) ξ ∈ I t w ( ξ ) t ≈

Original: [Equation]
Revised: [Equation]

Original: R (
Revised: R (

Original: [Equation]
Revised: [Equation]

Original: t
Revised: t

Original: [Equation]
Revised: [Equation]

Original: 1 ) , which suggests that w t is updated roughly along the direction of global gradient descent and thus is expected to converge quickly.
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 explicitly imposing local dissimilarity type conditions on local 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 the communication 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 with respect to both local minibatching and device sampling sizes.
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) , an identical bound naturally holds for the empirical form (2) under minibatch sampling according to local data empirical 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 applied to 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: In the analysis to follow, 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 our main 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 ν -weakly convex 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 local update 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 randomly chosen 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 about Theorem 2 extend to
Revised: We comment in passing that the discussions made in Remarks 4-6 about Theorem 2 extend to

Original: Theorem 4.
Revised: Theorem 4.

Original:  Additional Related Work
Revised:  Additional Related Work

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

Original: We next brieﬂy review some additional work in these lines of research that are closely related to ours.
Revised: We next brieﬂy review some additional work in these lines of research that are closely related to ours.

Original: Heterogenous federated learning.
Revised: Heterogenous federated learning.

Original: The presence of device heterogeneity features a key distinction between FL and classic distributed learning.
Revised: The presence of device heterogeneity features a key distinction between FL and classic distributed learning.

Original: The most commonly used FL method is FedAvg (McMahan et al., 2017), where the local update oracle is formed as multi-epoch SGD.
Revised: The most commonly used FL method is FedAvg (McMahan et al., 2017), where the local update oracle is formed as multi-epoch SGD.

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

Original: SGD.
Revised: SGD.

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

Original: Woodworth et al., 2020;
Revised: Woodworth et al., 2020;

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

Original: As another representative FL method,
Revised: As another representative FL method,

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

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

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

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

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

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

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

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

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

Original: Minibatch stochastic proximal point methods.
Revised: Minibatch stochastic proximal point methods.

Original: The proposed FedMSPP algorithm is a variant of
Revised: The proposed FedMSPP algorithm is a variant of

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

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

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

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

Original: Algorithmic stability.
Revised: Algorithmic stability.

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

Original: 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: 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).

Original: 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: 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).

Original: It was shown by Hardt et al.
Revised: It was shown by Hardt et al.

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

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

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

Original:  Conclusions
Revised:  Conclusions

Original: 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.
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, we developed 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: 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.
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 theoretical contribution 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 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.
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 for smooth problems in Theorem 1 and Theorem 3, and thus we view it as a limitation of the techniques used by 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: In the smooth-loss case, the comparison in Table 1 suggests that our results in
Revised: In the smooth-loss case, the comparison in Table 1 suggests that our results in

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

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: The authors sincerely thank the anonymous reviewers for their insightful comments on this paper.
Revised: The authors sincerely thank the anonymous reviewers for their insightful comments which are truly helpful for improving this paper.

Original: U21B
Revised: Xiao-Tong Yuan is funded in part by the National Key Research and Development Program of China under Grant No. 2018AAA0100400, and in part by the Natural Science Foundation of China (NSFC) under Grant No. U21B2049, No.61876090 and No.61936005.

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] 2.
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 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.
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]
Revised: [N/A]

Original: We did not use any of such assets in this work.
Revised: 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.