Cohort Squeeze: Beyond a Single Communication Round per Cohort in Cross-Device Federated Learning

## Abstract

Abstract Virtually all federated learning (FL) methods, including FedAvg, operate in the following manner: i) an orchestrating server sends the current model parameters to a cohort of clients selected via certain rule, ii) these clients then independently perform a local training procedure (e.g., via SGD or Adam) using their own training data, and iii) the resulting models are shipped to the server for aggregation. This process is repeated until a model of suitable quality is found. A notable feature of these methods is that each cohort is involved in a single communication round with the server only. In this work we challenge this algorithmic design primitive and investigate whether it is possible to “squeeze more juice” out of each cohort than what is possible in a single communication round. Surprisingly, we find that this is indeed the case, and our approach leads to up to 74% reduction in the total communication cost needed to train a FL model in the cross-device setting. Our method is based on a novel variant of the stochastic proximal point method (SPPM-AS) which supports a large collection of client sampling procedures some of which lead to further gains when compared to classical client selection approaches.

## 1 Introduction

Federated Learning (FL) is increasingly recognized for its ability to enable collaborative training of a global model across heterogeneous clients, while preserving privacy (McMahan et al., , ; Kairouz et al., ; Li et al., ; Karimireddy et al., ; Mishchenko et al., ; Malinovsky et al., ; Yi et al., ). This approach is particularly noteworthy in cross-device FL, involving the coordination of millions of mobile devices by a central server for training purposes (Kairouz et al., ). This setting is characterized by intermittent connectivity and limited resources. Consequently, only a subset of client devices participates in each communication round. Typically, the server samples a batch of clients (referred to as a *cohort* in FL), and each selected client trains the model received from the server using its local data. Then, the server aggregates the results sent from the selected cohort.
Another notable limitation of this approach is the constraint that prevents workers from storing states (operating in a stateless regime), thereby eliminating the possibility of employing variance reduction techniques.
We will consider a reformulation of the cross-device objective that assumes a finite number of workers being selected with uniform probabilities. Given that, in practice, only a finite number of devices is considered, i.e. the following finite-sum objective is considered:

$$ $\displaystyle\min_{x\in\mathbb{R}^{d}}f(x)\coloneqq\frac{1}{n}\sum_{i=1}^{n}f_ {i}(x).$ (1) $$

This reformulation aligns more closely with empirical observations and enhances understanding for illustrative purposes.
The extension to the expectation form of the following theory can be found in Section F.4.

Current representative approaches in the cross-device setting include FedAvg and FedProx. In our work, we introduce a method by generalizing stochastic proximal point method with arbitray sampling and term as SPPM-AS. This new method is inspired by the stochastic proximal point method (SPPM), a technique notable for its ability to converge under arbitrarily large learning rates and its flexibility in incorporating various solvers to perform proximal steps. This adaptability makes SPPM highly suitable for cross-device FL (Li et al., ; Yuan & Li, , ; Khaled & Jin, ; Lin et al., ).
Additionally, we introduce support for an arbitrary cohort sampling strategy, accompanied by a theoretical analysis. We present novel strategies that include support for client clustering, which demonstrate both theoretical and practical improvements.

Another interesting parameter that allows for control is the number of local communications.
Two distinct types of communication, *global* and *local*, are considered.
A *global* iteration is defined as a single round of communication between the server and all participating clients.
On the other hand, *local* communication rounds are synchronizations that take place within a chosen cohort.
Additionally, we introduce the concept of total communication cost, which includes both local and global communication iterations, to measure the overall efficiency of the communication process.
The total communication cost naturally depends on several factors. These include the local algorithm used to calculate the prox, the global stepsize, and the sampling technique.

Figure: Figure 1: The total communication cost (defined as $TK$) with the number of local communication rounds $K$ needed to reach the target accuracy $\epsilon$ for the chosen cohort in each global iteration. The dashed red line depicts the communication cost of the FedAvg algorithm. Markers indicate the $TK$ value for different learning rates $\gamma$ of our algorithm SPPM-AS.
Refer to caption: x1.png

Previous results on cross-device settings consider only one local communication round for the selected cohort (Li et al., ; Reddi et al., ; Li et al., ; Wang et al., , ; Xu et al., ; Malinovsky et al., ; Jhunjhunwala et al., ; Sun et al., , ).
Our experimental findings reveal that *increasing the number of local communication rounds within a chosen cohort per global iteration can indeed lower the total communication cost needed to reach a desired global accuracy level*, which we denote as $\varepsilon$.
Figure 1 illustrates the relationship between total communication costs and the number of local communication rounds.
Assume that the cost of communication per round is $1$ unit. $K$ represents the number of local communication rounds per global iteration for the selected cohort, while $T$ signifies the *minimum* number of global iterations needed to achieve the accuracy threshold $\epsilon$. Then, the total cost incurred by our method can be expressed as $TK$.
For comparison, the dashed line in the figure shows the total cost for the FedAvg algorithm, which always sets $K$ to $1$, directly equating the number of global iterations to total costs. Our results across various datasets identify the optimal $K$ for each learning rate to achieve $\epsilon$-accuracy.
Figure 1 shows that adding more local communication rounds within each global iteration can lead to a significant reduction in the overall communication cost. For example, when the learning rate is set to $1000$, the optimal cost is reached with $10$ local communication rounds, making $K=10$ a more efficient choice compared to a smaller number. On the other hand, at a lower learning rate of $100$, the optimal cost of $12$ is reached with $K=3$. This pattern indicates that as we increase the number of local communication rounds, the total cost can be reduced, and the optimal number of local communication rounds tends to increase with higher learning rates.

Our key *contributions* are summarized as follows:

$\bullet$ We present and analyze SPPM-AS, a novel approach within the stochastic proximal point method framework tailored for cross-device federated learning, which supports arbitrary sampling strategies. Additionally, we provide an analysis of standard sampling techniques and introduce new techniques based on clustering approaches. These novel techniques are theoretically analyzed, offering a thorough comparison between different methods.

$\bullet$ Our numerical experiments, conducted on both convex logistic regression models and non-convex neural networks, demonstrate that the introduced framework enables fine-tuning of parameters to surpass existing state-of-the-art cross-device algorithms. Most notably, we found that increasing the number of local communication rounds within the selected cohort is an effective strategy for reducing the overall communication costs necessary to achieve a specified target accuracy threshold.

$\bullet$ We offer practical guidance on the proper selection of parameters for federated learning applications. Specifically, we examine the potential choices of solvers for proximal operations, considering both convex and non-convex optimization regimes. Our experiments compare first-order and second-order solvers to identify the most effective ones.

## 2 Method

In this section, we explore efficient stochastic proximal point methods with arbitrary sampling for cross-device FL to optimize the objective (1). Throughout the paper, we denote $[n]\coloneqq\left\{1,\dots,n\right\}$. Our approach builds on the following assumptions.

###### Assumption 1 .

Function $f_{i}:\mathbb{R}^{d}\to\mathbb{R}$ is differentiable for all samples $i\in[n]$.

This implies that the function $f$ is differentiable. The order of differentiation and summation can be interchanged due to the additive property of the gradient operator.

$$ $\displaystyle\nabla f(x)\stackrel{{\scriptstyle Eqn.~{}(\ref{eqn:obj_1})}}{{=} }\nabla\left[\frac{1}{n}\sum_{i=1}^{n}f_{i}(x)\right]=\frac{1}{n}\sum_{i=1}^{n }\nabla f_{i}(x).$ $$

###### Assumption 2 .

Function $f_{i}:\mathbb{R}^{d}\to\mathbb{R}$ is $\mu$-strongly convex for all samples $i\in\left[n\right]$, where $\mu>0$. That is,
$f_{i}(y)+\left\langle\nabla f_{i}(y),x-y\right\rangle+\frac{\mu}{2}{\left
\lVert x-y\right\rVert}^{2}\leq f_{i}(x),$ for all $x,y\in\mathbb{R}^{d}$.

This implies that $f$ is $\mu$-strongly convex and hence has a unique minimizer, which we denote by $x_{\star}$. We know that $\nabla f(x_{\star})=0$. Notably, we do *not* assume $f$ to be $L$-smooth.

###### Assumption 1 .

###### Assumption 2 .

### 2.1 Sampling Distribution

Let $\mathcal{S}$ be a probability distribution over the $2^{n}$ subsets of $[n]$. Given a random set $S\sim\mathcal{S}$, we define

$$ $\displaystyle p_{i}\coloneqq\operatorname{Prob}(i\in{S}),\quad i\in[n].$ $$

We restrict our attention to proper and nonvacuous random sets.

###### Assumption 3 .

$\mathcal{S}$ is proper (i.e., $p_{i}>0$ for all $i\in[n]$) and nonvacuous (i.e., $\operatorname{Prob}({S}=\emptyset)=0$).

Let $C$ be the selected cohort. Given $\emptyset\neq C\subseteq[n]$ and $i\in[n]$, we define

$$ $\displaystyle v_{i}(C):=\begin{cases}\frac{1}{p_{i}}&i\in C\\ 0&i\notin C\end{cases}\ \Rightarrow\ f_{C}(x):=\frac{1}{n}\sum_{i=1}^{n}v_{i}( C)f_{i}(x){=}\sum_{i\in C}\frac{1}{np_{i}}f_{i}(x).$ (2) $$

Note that $v_{i}(S)$ is a random variable and $f_{S}$ is a random function. By construction, $\mathrm{E}_{S\sim\mathcal{S}}\left[v_{i}(S)\right]=1$ for all $i\in[n]$, and hence

$$ $\displaystyle\mathrm{E}_{{S}\sim\mathcal{S}}\left[f_{{S}}(x)\right]=\mathrm{E} _{{S}\sim\mathcal{S}}\left[\frac{1}{n}\sum_{i=1}^{n}v_{i}(S)f_{i}(x)\right]= \frac{1}{n}\sum_{i=1}^{n}\mathrm{E}_{S\sim\mathcal{S}}\left[v_{i}(S)\right]f_{ i}(x)=\frac{1}{n}\sum_{i=1}^{n}f_{i}(x)=f(x).$ $$

Therefore, the optimization problem in Equation 1 is equivalent to the stochastic optimization problem

$$ $\displaystyle\min_{x\in\mathbb{R}^{d}}\left\{f(x):=\mathrm{E}_{S\sim\mathcal{S }}\left[f_{S}(x)\right]\right\}.$ (3) $$

Further, if for each $C\subset[n]$ we let $p_{C}:=\operatorname{Prob}(S=C)$, then $f$ can be written in the equivalent form

$$ $\displaystyle f(x)=\mathbb{E}_{S\sim\mathcal{S}}\left[f_{S}(x)\right]=\sum_{C \subseteq[n]}p_{C}f_{C}(x)=\sum_{C\subseteq[n],p_{C}>0}p_{C}f_{C}(x).$ (4) $$

###### Assumption 3 .

### 2.2 Core Algorithm

Applying SPPM (Khaled & Jin, ) to Equation 3, we arrive at stochastic proximal point method with arbitrary sampling (SPPM-AS, Algorithm 1):

$$ $x_{t+1}=\operatorname{prox}_{\gamma f_{S_{t}}}\left(x_{t}\right),$ $$

where $S_{t}\sim\mathcal{S}$.

Figure: Algorithm 1 Stochastic Proximal Point Method with Arbitrary Sampling (SPPM-AS)

###### Theorem 4 (Convergence of SPPM-AS ) .

Let 1 (differentiability) and 2 (strong convexity) hold. Let $\mathcal{S}$ be a sampling satisfying 3, and define

$$ $\displaystyle\mu_{\mathrm{AS}}:=\min_{C\subseteq[n],p_{C}>0}\sum_{i\in C}\frac {\mu_{i}}{np_{i}},\quad\sigma_{\star,\mathrm{AS}}^{2}:=\sum_{C\subseteq[n],p_{ C}>0}p_{C}\left\|\nabla f_{C}\left(x_{\star}\right)\right\|^{2}.$ (5) $$

Let $x_{0}\in\mathbb{R}^{d}$ be an arbitrary starting point. Then for any $t\geq 0$ and any $\gamma>0$, the iterates of SPPM-AS (Algorithm 1) satisfy

$$ $\mathrm{E}\left[\left\|x_{t}-x_{\star}\right\|^{2}\right]\leq\left(\frac{1}{1+ \gamma\mu_{\mathrm{AS}}}\right)^{2t}\left\|x_{0}-x_{\star}\right\|^{2}+\frac{ \gamma\sigma_{\star,\mathrm{AS}}^{2}}{\gamma\mu_{\mathrm{AS}}^{2}+2\mu_{ \mathrm{AS}}}.$ $$

###### Theorem 4 (Convergence of SPPM-AS ) .

#### Theorem interpretation.

In the theorem presented above, there are two main terms: $\left(\nicefrac{{1}}{{(1+\gamma\mu_{\mathrm{AS}})}}\right)^{2t}$ and $\nicefrac{{\gamma\sigma_{\star,\mathrm{AS}}^{2}}}{{(\gamma\mu_{\mathrm{AS}}^{2
}+2\mu_{\mathrm{AS}})}}$, which define the convergence speed and neighborhood, respectively. Additionally, there are three hyperparameters to control the behavior: $\gamma$ (the global learning rate), $\mathrm{AS}$ (the sampling type), and $T$ (the number of global iterations). In the following paragraphs, we will explore special cases to provide a clear intuition of how the SPPM-AS theory works.

#### Interpolation regime.

Consider the interpolation regime, characterized by $\sigma_{\star,\mathrm{AS}}^{2}=0$ . Since we can use arbitrarily large $\gamma>0$, we obtain an arbitrarily fast convergence rate.
Indeed, $\left(\nicefrac{{1}}{{(1+\gamma\mu_{\mathrm{AS}})}}\right)^{2t}$ can be made arbitrarily small for any fixed $t\geq 1$, even $t=1$, by choosing $\gamma$ large enough. However, this is not surprising, since now $f$ and all functions $f_{\xi}$ share a single minimizer, $x_{\star}$, and hence it is possible to find it by sampling a small batch of functions even a single function $f_{\xi}$, and minimizing it, which is what the prox does, as long as $\gamma$ is large enough.

#### A single step travels far.

Observe that for $\gamma=\nicefrac{{1}}{{\mu_{\mathrm{AS}}}}$, we have $\nicefrac{{\gamma\sigma^{2}_{\star,\mathrm{AS}}}}{{(\gamma\mu_{\mathrm{AS}}^{2
}+2\mu_{\mathrm{AS}})}}=\nicefrac{{\sigma^{2}_{\star,\mathrm{AS}}}}{{3\mu_{
\mathrm{AS}}^{2}}}$. In fact, the convergence neighborhood $\nicefrac{{\gamma\sigma^{2}_{\star,\mathrm{AS}}}}{{(\gamma\mu_{\mathrm{AS}}^{2
}+2\mu_{\mathrm{AS}})}}$ is bounded above by three times this quantity irrespective of the choice of the stepsize. Indeed,
$\frac{\gamma\sigma^{2}_{\star,\mathrm{AS}}}{\gamma\mu_{\mathrm{AS}}^{2}+2\mu_{
\mathrm{AS}}}\leq\min\left\{\frac{\sigma^{2}_{\star,\mathrm{AS}}}{\mu_{\mathrm
{AS}}^{2}},\frac{\gamma\sigma^{2}_{\star,\mathrm{AS}}}{\mu_{\mathrm{AS}}}
\right\}\leq\frac{\sigma^{2}_{\star,\mathrm{AS}}}{\mu_{\mathrm{AS}}^{2}}.$
That means that no matter how far the starting point $x_{0}$ is from the optimal solution $x_{\star}$, if we choose the stepsize $\gamma$ to be large enough, then we can get a decent-quality solution after a single iteration of SPPM-AS already! Indeed, if we choose $\gamma$ large enough so that
$\left(\nicefrac{{1}}{{1+\gamma\mu_{\mathrm{AS}}}}\right)^{2}{\left\lVert x_{0}
-x_{\star}\right\rVert}^{2}\leq\delta,$
where $\delta>0$ is chosen arbitrarily, then for $t=1$ we get
$\mathbb{E}\left[{\left\lVert x_{1}-x_{\star}\right\rVert}^{2}\right]\leq\delta
+\nicefrac{{\sigma^{2}_{\star,\mathrm{AS}}}}{{\mu_{\mathrm{AS}}^{2}}}.$

#### Iteration complexity.

We have seen above that an accuracy arbitrarily close to (but not reaching) $\nicefrac{{\sigma^{2}_{\star,\mathrm{AS}}}}{{\mu_{\mathrm{AS}}^{2}}}$ can be achieved via a single step of the method, provided that the stepsize $\gamma$ is large enough.
Assume now that we aim for $\epsilon$ accuracy, where $\epsilon\leq\nicefrac{{\sigma^{2}_{\star,\mathrm{AS}}}}{{\mu_{\mathrm{AS}}^{2}}}$. We can show that with the stepsize $\gamma=\nicefrac{{\varepsilon\mu_{\mathrm{AS}}}}{{\sigma^{2}_{\star,\mathrm{AS
}}}}$, we get
$\mathrm{E}\left[\left\|x_{t}-x_{\star}\right\|^{2}\right]\leq\varepsilon$
provided that
$t\geq\left(\frac{\sigma^{2}_{\star,\mathrm{AS}}}{2\varepsilon\mu_{\mathrm{AS}}
^{2}}+\frac{1}{2}\right)\log\left(\frac{2\left\|x_{0}-x_{\star}\right\|^{2}}{
\varepsilon}\right).$
We provide the proof in Section F.5.
To ensure thoroughness, we present in Section F.9 the lemma of the inexact formulation for SPPM-AS, which offers greater practicality for empirical experimentation. Further insights are provided in the subsequent experimental section.

#### General Framework.

With freedom to choose arbitrary algorithms for solving the proximal operator one can see that SPPM-AS is generalization for such renowned methods as FedProx (Li et al., ) and FedAvg (McMahan et al., ). A more particular overview of FedProx-SPPM-AS is presented in further Section B.4.

### 2.3 Arbitrary Sampling Examples

Details on simple Full Sampling (FS) and Nonuniform Sampling (NS) are provided in Section B.2. In this section, we focus more intently on the sampling strategies that are of particular interest to us.

#### Nice Sampling (NICE).

Choose $\tau\in[n]$ and let $S$ be a random subset of $[n]$ of size $\tau$ chosen uniformly at random. Then $p_{i}=\nicefrac{{\tau}}{{n}}$ for all $i\in[n]$. Moreover, let $\binom{n}{\tau}$ represents the number of combinations of $n$ taken $\tau$ at a time, $p_{C}=\frac{1}{\binom{n}{\tau}}$ whenever $|C|=\tau$ and $p_{C}=0$ otherwise. So,

$$ $\mu_{\mathrm{AS}}=\mu_{\mathrm{NICE}}(\tau):=\min_{C\subseteq[n],p_{C}>0}\sum_ {i\in C}\frac{\mu_{i}}{np_{i}}=\min_{C\subseteq[n],|C|=\tau}\frac{1}{\tau}\sum _{i\in C}\mu_{i},$ $$

$$ $\displaystyle\sigma_{\star,\mathrm{AS}}^{2}$ $\displaystyle=\sigma_{\star,\mathrm{NICE}}^{2}(\tau):=\sum_{C\subseteq[n],p_{C }>0}p_{C}\left\|\nabla f_{C}\left(x_{\star}\right)\right\|^{2}\stackrel{{ \scriptstyle Eqn.~{}(\ref{eqn:main_8001})}}{{=}}\sum_{C\subseteq[n],|C|=\tau} \frac{1}{\binom{n}{\tau}}\left\|\frac{1}{\tau}\sum_{i\in C}\nabla f_{i}\left(x _{\star}\right)\right\|^{2}.$ $$

It can be shown that $\mu_{\mathrm{NICE}}(\tau)$ is a *nondecreasing* function of $\tau$ (Section F.6). So, as the minibatch size $\tau$ increases, the strong convexity constant $\mu_{\mathrm{NICE}}(\tau)$ can only improve. Since $\mu_{\mathrm{NICE}}(1)=\min_{i}\mu_{i}$ and $\mu_{\mathrm{NICE}}(n)=$ $\frac{1}{n}\sum_{i=1}^{n}\mu_{i}$, the value of $\mu_{\mathrm{NICE}}(\tau)$ interpolates these two extreme cases as $\tau$ varies between 1 and $n$. Conversely, $\sigma_{\star,\mathrm{NICE}}^{2}(\tau)=\frac{\nicefrac{{n}}{{\tau}}-1}{n-1}
\sigma_{\star,\mathrm{NICE}}^{2}(1)$ is a nonincreasing function, reaching a value of $\sigma_{\star,\mathrm{NICE}}^{2}(n)=0$, as explained in Section F.6.

#### Block Sampling (BS).

Let $C_{1},\ldots,C_{b}$ be a partition of $[n]$ into $b$ nonempty blocks. For each $i\in[n]$, let $B(i)$ indicate which block $i$ belongs to. In other words, $i\in C_{j}$ if $B(i)=j$. Let $S=C_{j}$ with probability $q_{j}>0$, where $\sum_{j}q_{j}=1$. Then $p_{i}=q_{B(i)}$, and hence Equation 5 takes on the form

$$ $\displaystyle\mu_{\mathrm{AS}}=\mu_{\mathrm{BS}}:=\min_{j\in[b]}\frac{1}{nq_{j }}\sum_{i\in C_{j}}\mu_{i},\quad\sigma_{\star,\mathrm{AS}}^{2}=\sigma_{\star, \mathrm{BS}}^{2}:=\sum_{j\in[b]}q_{j}\left\|\sum_{i\in C_{j}}\frac{1}{np_{i}} \nabla f_{i}\left(x_{\star}\right)\right\|^{2}.$ $$

*Considering two extreme cases:* If $b=1$, then SPPM-BS = SPPM-FS = PPM. So, indeed, we recover the same rate as SPPM-FS. If $b=n$, then SPPM-BS = SPPM-NS. So, indeed, we recover the same rate as SPPM-NS. We provide the detailed analysis in Section B.3.

#### Stratified Sampling (SS).

Let $C_{1},\ldots,C_{b}$ be a partition of $[n]$ into $b$ nonempty blocks, as before. For each $i\in[n]$, let $B(i)$ indicate which block does $i$ belong to. In other words, $i\in C_{j}$ iff $B(i)=j$. Now, for each $j\in[b]$ pick $\xi_{j}\in C_{j}$ uniformly at random, and define $S=\cup_{j\in[b]}\left\{\xi_{j}\right\}$. Clearly, $p_{i}=\frac{1}{\left|C_{B(i)}\right|}$. Let’s denote $\mathbf{i}_{b}\coloneqq(i_{1},\cdots,i_{b}),\mathbf{C}_{b}\coloneqq C_{1}
\times\cdots\times C_{b}$. Then, Equation 5 take on the form

$$ $\displaystyle\mu_{\mathrm{AS}}=\mu_{\mathrm{SS}}:=\min_{\mathbf{i}_{b}\in \mathbf{C}_{b}}\sum_{j=1}^{b}\frac{\mu_{i_{j}}\left|C_{j}\right|}{n},\quad \sigma_{\star,\mathrm{AS}}^{2}=\sigma_{\star,\mathrm{SS}}^{2}:=\sum_{\mathbf{i }_{b}\in\mathbf{C}_{b}}\left(\prod_{j=1}^{b}\frac{1}{\left|C_{j}\right|}\right )\left\|\sum_{j=1}^{b}\frac{\left|C_{j}\right|}{n}\nabla f_{i_{j}}\left(x_{ \star}\right)\right\|^{2}.$ $$

*Considering two extreme cases:* If $b=1$, then SPPM-SS = SPPM-US. So, indeed, we recover the same rate as SPPM-US. If $b=n$, then SPPM-SS = SPPM-FS. So, indeed, we recover the same rate as SPPM-FS. We provide the detailed analysis in Section B.3.

###### Lemma 5 (Variance Reduction Due to Stratified Sampling) .

Consider the stratified sampling. For each $j\in[b]$, define

$$ $\sigma_{j}^{2}:=\max_{i\in C_{j}}\left\|\nabla f_{i}\left(x_{\star}\right)- \frac{1}{\left|C_{j}\right|}\sum_{l\in C_{j}}\nabla f_{l}\left(x_{\star}\right )\right\|^{2}.$ $$

In words, $\sigma_{j}^{2}$ is the maximal squared distance of a gradient (at the optimum) from the mean of the gradients (at optimum) within cluster $C_{j}$. Then

$$ $\displaystyle\sigma_{\star,\mathrm{SS}}^{2}\leq\frac{b}{n^{2}}\sum_{j=1}^{b} \left|C_{j}\right|^{2}\sigma_{j}^{2}\leq b\max\left\{\sigma_{1}^{2},\ldots, \sigma_{b}^{2}\right\}.$ $$

In Lemma 5, we demonstrate that stratified sampling outperforms block sampling due to reduced variance. Note that in the scenario of complete inter-cluster homogeneity, where $\sigma_{j}^{2}=0$ for all $j$, both bounds imply that $0=\sigma^{2}_{\star,\mathrm{SS}}\leq\sigma^{2}_{\star,\mathrm{BS}}$.

###### Lemma 5 (Variance Reduction Due to Stratified Sampling) .

#### Stratified Sampling Outperforms Block Sampling and Nice Sampling in Convergence Neighborhood.

We theoretically compare stratified sampling with block sampling and nice sampling, advocating for stratified sampling as the superior method for future clustering experiments due to its optimal variance properties. We begin with the assumption of $b$ clusters of uniform size $b$ (14), which simplifies the analysis by enabling comparisons of various sampling methods, all with the same sampling size, $b$: $b$-nice sampling, stratified sampling with $b$ clusters, and block sampling where all clusters are of uniform size $b$. Furthermore, we introduce the concept of optimal clustering for stratified sampling (noted as $\mathcal{C}_{b,\mathrm{SS}}$, Definition 15) in response to a counterexample where block sampling and nice sampling achieve lower variance than stratified sampling (Example 1). Finally, we compare neighborhoods using the stated assumption.

###### Lemma 6 .

Given 14, the following holds: $\sigma_{\star,\mathrm{SS}}^{2}\left(\mathcal{C}_{b,\mathrm{SS}}\right)\leq
\sigma_{\star,\mathrm{NICE}}^{2}$ for arbitrary $b$.

Lemma 6 demonstrates that, under specific conditions, the stratified sampling neighborhood is preferable to that of nice sampling. One might assume that, under the same assumptions, a similar assertion could be made for showing that block sampling is inferior to stratified sampling . However, this has only been verified for the simplified case where both the block size and the number of blocks are $b=2$, as detailed in Section F.8.

###### Lemma 6 .

## 3 Experiments

#### Practical Decision-Making with SPPM-AS .

In our analysis of SPPM-AS, guided by theoretical foundations of Theorem 4 and empirical evidence summarized in Table 1, we explore practical decision-making for varying scenarios. This includes adjustments in hyperparameters within the framework $KT(\epsilon,\mathcal{S},\gamma,\mathcal{A}\left(K\right))$. Here, $\epsilon$ represents accuracy goal, $\mathcal{S}$ represents the sampling distribution, $\gamma$ is representing global learning rate (proximal operator parameter), $\mathcal{A}$ denotes the proximal optimization algorithm, while $K$ denotes the number of local communication rounds. In table 1,

Table: Table 1: $KT(\epsilon,\mathcal{S},\gamma,\mathcal{A}\left(K\right))$

| HP | Control | $KT(\cdots)$ | Exp. |
| --- | --- | --- | --- |
| $\gamma$ | $\gamma\uparrow$ | $KT\downarrow$, $\epsilon\uparrow{\textsuperscript{\color[rgb]{0,0,1}\definecolor[named]{<br>pgfstrokecolor}{rgb}{0,0,1}(1)}}$ | D.2 |
| optimal $\left(\gamma,K\right)\uparrow$ | $\downarrow$ | 3.3 |  |
| $\mathcal{A}$ | $\mu$-convex + BFGS/CG | $\downarrow$ compared to LocalGD | 3.3 |
| NonCVX and Hierarchical FL + Adam with tuned lr | $\downarrow$ compared to LocalGD | 3.7 |  |

- (1)
$\epsilon$ is a convergence neighborhood or accuracy.

we summarize how changes on following hyperparameters will influence target metric. With increasing learning rate $\gamma$ one achieves faster convergence with smaller accuracy, also noted as accuracy-rate tradeoff. Our primary observation that with an increase in both the learning rate, $\gamma$, and the number of local steps, $K$, leads to an improvement in the convergence rate. Employing various local solvers for proximal operators also shows an improvement in the convergence rate compared to FedAvg in both convex and non-convex cases.

#### Practical Decision-Making with SPPM-AS .

In our analysis of SPPM-AS, guided by theoretical foundations of Theorem 4 and empirical evidence summarized in Table 1, we explore practical decision-making for varying scenarios. This includes adjustments in hyperparameters within the framework $KT(\epsilon,\mathcal{S},\gamma,\mathcal{A}\left(K\right))$. Here, $\epsilon$ represents accuracy goal, $\mathcal{S}$ represents the sampling distribution, $\gamma$ is representing global learning rate (proximal operator parameter), $\mathcal{A}$ denotes the proximal optimization algorithm, while $K$ denotes the number of local communication rounds. In table 1,

### 3.1 Objective and Datasets

Our analysis begins with logistic regression with a convex $l_{2}$ regularizer, which can be represented as:

$$ $f_{i}(x)\coloneqq\frac{1}{n_{i}}\sum_{j=1}^{n_{i}}\log\left(1+\exp(-b_{i,j}x^{ T}a_{i,j})\right)+\frac{\mu}{2}\|x\|^{2},$ $$

where $\mu$ is the regularization parameter, $n_{i}$ denotes the total number of data points at client $i$, $a_{i,j}$ are the feature vectors, and $b_{i,j}\in\{-1,1\}$ are the corresponding labels. Each function $f_{i}$ exhibits $\mu$-strong convexity and $L_{i}$-smoothness, with $L_{i}$ computed as $\frac{1}{4n_{i}}\sum_{j=1}^{n_{i}}\|a_{i,j}\|^{2}+\mu$. For our experiments, we set $\mu$ to 0.1.

Our study utilized datasets from the LibSVM repository (Chang & Lin, ), including mushrooms, a6a, ijcnn1.bz2, and a9a. We divided these into feature-wise heterogeneous non-iid splits for FL, detailed in Section C.1, with a default cohort size of $10$.
We primarily examined logistic regression, finding results consistent with our theoretical framework, as discussed extensively in Section 3.3 through Section D.2. Additional neural network experiments are detailed in Section 3.7 and Appendix E.

Figure: (a) base
Refer to caption: x3.png

### 3.2 On Choosing Sampling Strategy

As shown in Section 2.3, multiple sampling techniques exist. We propose using clustering approach in conjuction with SPPM-SS as the default sampling strategy for all our experiments. The stratified sampling optimal clustering is impractical due to the difficulty in finding $x_{\star}$; therefore, we employ a clustering heuristic that aligns with the concept of creating homogeneous worker groups. One such method is K-means, which we use by default. More details on our clustering approach can be found in the Appendix C.1.
We compare various sampling techniques in the left panel of Figure 3. Extensive ablations verified the efficiency of stratified sampling over other strategies, due to variance reduction (Lemma 5).

Figure: Figure 3: The first column compares sampling methods, while the right two columns analyze convergence relative to popular baselines. $\gamma=1.0$.
Refer to caption: x7.png

### 3.3 Communication Cost Reduction through Increased Local Communication Rounds

In this study, we investigate whether increasing the number of local communication rounds, denoted as $K$, in our proposed algorithm SPPM-SS, can lead to a decrease in the total communication cost required to converge to a predetermined global accuracy $\epsilon>0$. In Figure 1, we analyzed various datasets, including a6a and mushrooms, confirming that higher local communication rounds reduce communication costs, especially with larger learning rates. Our study includes both self-ablation of SPPM-SS across different learning rate scales and comparisons with the widely-used cross-device FL method LocalGD (or FedAvg) on the selected cohort. Ablation studies were conducted with a large empirical learning rate of $0.1$, a smaller rate of $0.01$, and an optimal rate as discussed by Khaled & Richtárik (), alongside minibatch sampling described by Gower et al. ().

In Figure 2, we present more extensive ablations. Specifically, we set the base method (Figure 2(a)) using the dataset a6a, a proximal solver BFGS, and $\epsilon=5\cdot 10^{-3}$. In Figure 2(b), we explore the use of an alternative solver, CG (Conjugate Gradient), noting some differences in outcomes. For instance, with a learning rate $\gamma=1000$, the optimal $K$ with CG becomes 7, lower than 10 in the base setting using BFGS. In Figure 2(c), we investigate the impact of varying $\epsilon=10^{-2}$. Our findings consistently show SPPM-SS’s significant performance superiority over LocalGD.

### 3.4 Evaluating the Performance of Various Solvers 𝒜 𝒜 \mathcal{A} caligraphic_A

We further explore the impact of various solvers on optimizing the proximal operators, showcasing representative methods in Table 2 in the Section A.3. A detailed overview and comparison of local optimizers listed in the table are provided in Section A.3, given the extensive range of candidate options available. To emphasize key factors, we compare the performance of first-order methods, such as the Conjugate Gradient (CG) method (Hestenes et al., ), against second-order methods, like the Broyden-Fletcher-Goldfarb-Shanno (BFGS) algorithm (Broyden, ; Shanno, ), in the context of strongly convex settings. For non-convex settings, where first-order methods are prevalent in deep learning experiments, we examine an ablation among popular first-order local solvers, specifically choosing MimeLite (Karimireddy et al., ) and FedOpt (Reddi et al., ). The comparisons of different solvers for strongly convex settings are presented in Figure 2(b), with the non-convex comparison included in the appendix. Upon comparing first-order and second-order solvers in strongly convex settings, we observed that CG outperforms BFGS for our specific problem. In neural network experiments, MimeLite-Adam was found to be more effective than FedOpt variations. However, it is important to note that all these solvers are viable options that have led to impressive performance outcomes.

### 3.5 Comparative Analysis with Baseline Algorithms

In this section, we conduct an extensive comparison with several established cross-device FL baseline algorithms. Specifically, we examine MB-GD (MiniBatch Gradient Descent with partial client participation), and MB-LocalGD, which is the local gradient descent variant of MB-GD. We default the number of local iterations to 5 and adopt the optimal learning rate as suggested by Gower et al. (). To ensure a fair comparison, the cohort size $\left|C\right|$ is fixed at 10 for all minibatch methods, including our proposed SPPM-SS. The results of this comparative analysis are depicted in Figure 3.
Our findings reveal that SPPM-SS consistently achieves convergence within a significantly smaller neighborhood when compared to the existing baselines. Notably, in contrast to MB-GD and MB-LocalGD, SPPM-SS is capable of utilizing arbitrarily large learning rates. This attribute allows for faster convergence, although it does result in a larger neighborhood size.

Figure: Figure 4: The left column shows the Server-hub-client hierarchical FL architecture. For the right two columns: on the left, communication cost for achieving 70% accuracy in hierarchical FL ($c_{1}=0.05$, $c_{2}=1$); on the right, convergence with optimal hyperparameters ($c_{1}=0.05$, $c_{2}=1$).
Refer to caption: x10.png

### 3.6 Hierarchical Federated Learning

We extend our analysis to a hub-based hierarchical FL structure, as conceptualized in the left part of Figure 4. This structure envisions a cluster directly connected to $m$ hubs, with each hub $m_{i}$ serving $n_{i}$ clients. The clients, grouped based on criteria such as region, communicate exclusively with their respective regional hub, which in turn communicates with the central server. Given the inherent nature of this hierarchical model, the communication cost $c_{1}$ from each client to its hub is consistently lower than the cost $c_{2}$ from each hub to the server. We define communication from clients to hubs as local communication and from hubs to the server as global communication. Under SPPM-SS, the total cost is expressed as $(c_{1}K+c_{2})T_{\operatorname{SPPM-SS}}$, while for LocalGD, it is $(c_{1}+c_{2})T_{\operatorname{LocalGD}}$. As established in Section 3.3, $T_{\operatorname{SPPM-SS}}$ demonstrates significant improvement in total communication costs compared to LocalGD within a hierarchical setting. Our objective is to illustrate this by contrasting the standard FL setting, depicted in Figure 2(a) with parameters $c_{1}=1$ and $c_{2}=0$, against the hierarchical FL structure, which assumes $c_{1}=0.1$ and $c_{2}=1$, as shown in Figure 2(d). Given the variation in $c_{1}$ and $c_{2}$ values between these settings, a direct comparison of absolute communication costs is impractical. Therefore, our analysis focuses on the ratio of communication cost reduction in comparison to LocalGD. For the base setting, LocalGD’s optimal total communication cost is 39 with 12 local iterations, whereas for SPPM-SS ($\gamma=1000$), it is reduced to 10 with 10 local and 1 global communication rounds, amounting to a 74.36% reduction. With the hierarchical FL structure in Figure 2(d), SPPM-SS achieves an even more remarkable communication cost reduction of 94.87%. Further ablation studies on varying local communication cost $c_{1}$ in the Section D.3 corroborate these findings.

### 3.7 Neural Network Evaluations

Our empirical analysis includes experiments on Convolutional Neural Networks (CNNs) using the FEMNIST dataset, as described by Caldas et al. (). We designed the experiments to include a total of 100 clients, with each client representing data from a unique user, thereby introducing natural heterogeneity into our study. We employed the Nice sampling strategy with a cohort size of 10. In contrast to logistic regression models, here we utilize training accuracy as a surrogate for the target accuracy $\epsilon$. For the optimization of the proximal operator, we selected the Adam optimizer, with the learning rate meticulously fine-tuned over a linear grid. Detailed descriptions of the training procedures and the CNN architecture are provided in the Appendix E.

Our analysis primarily focuses on the hierarchical FL structure. Initially, we draw a comparison between our proposed method, SPPM-AS, and LocalGD. The crux of our investigation is the total communication cost required to achieve a predetermined level of accuracy, with findings detailed in the right part of Figure 4. Significantly, SPPM-AS demonstrates enhanced performance with the integration of multiple local communication rounds. Notably, the optimal number of these rounds tends to increase alongside the parameter $\gamma$. For each configuration, the convergence patterns corresponding to the sets of optimally tuned hyperparameters are depicted in Figure 4.

## 4 Conclusion

Our research challenges the conventional single-round communication model in federated learning by presenting a novel approach where cohorts participate in multiple communication rounds. This adjustment leads to a significant 74% reduction in communication costs, underscoring the efficacy of extending cohort engagement beyond traditional limits. Our method, SPPM-AS, equipped with diverse client sampling procedures, contributes substantially to this efficiency. This foundational work showcases a pivotal shift in federated learning strategies. Future work could focus on improving algorithmic robustness and ensuring privacy compliance.

## Appendix A Related Work

### A.1 Cross-Device Federated Learning

This paper delves into the realm of Federated Learning (FL), focusing on the cross-device variant, which presents unique and significant challenges. In FL, two predominant settings are recognized: cross-silo and cross-device scenarios, as detailed in Table 1 of Kairouz et al., . The primary distinction lies in the nature of the clients: cross-silo FL typically involves various organizations holding substantial data, whereas cross-device FL engages a vast array of mobile or IoT devices.
In cross-device FL, the complexity is heightened by the inability to maintain a persistent hidden state for each client, unlike in cross-silo environments. This factor renders certain approaches impractical, particularly those reliant on stateful clients participating consistently across all rounds. Given the sheer volume of clients in cross-device FL, formulating and analyzing outcomes in an expectation form is more appropriate, but more complex than in finite-sum scenarios.

The pioneering and perhaps most renowned algorithm in cross-device FL is FedAvg (McMahan et al., ) and implemented in applications like Google’s mobile keyboard (Hard et al., ; Yang et al., ; Ramaswamy et al., ). However, it is noteworthy that popular accelerated training algorithms such as Scaffold (Karimireddy et al., ) and ProxSkip (Mishchenko et al., ) are not aligned with our focus due to their reliance on memorizing the hidden state for each client, which is applicable for cross-device FL.
Our research pivots on a novel variant within the cross-device framework. Once the cohort are selected for each global communication round, these cohorts engage in what we term as ‘local communications’ multiple times. The crux of our study is to investigate whether increasing the number of local communication rounds can effectively reduce the total communication cost to converge to a targeted accuracy.

### A.2 Stochastic Proximal Point Method

Our exploration in this paper centers on the Stochastic Proximal Point Method (SPPM), a method extensively studied for its convergence properties. Initially termed as the incremental proximal point method by Bertsekas (), it was shown to converge nonasymptotically under the assumption of Lipschitz continuity for each $f_{i}$. Following this, Ryu & Boyd () examined the convergence rates of SPPM, noting its resilience to inaccuracies in learning rate settings, contrasting with the behavior of Stochastic Gradient Descent (SGD).
Further developments in SPPM’s application were seen in the works of Patrascu & Necoara (), who analyzed its effectiveness in constrained optimization, incorporating random projections. Asi & Duchi () expanded the scope of SPPM by studying a generalized method, AProx, providing insights into its stability and convergence rates under convex conditions. The research by Asi et al. () and Chadha et al. () further extended these findings, focusing on minibatching and convergence under interpolation in the AProx framework.

In the realm of federated learning, particularly concerning non-convex optimization, SPPM is also known as FedProx, as discussed in works like those of Li et al. () and Yuan & Li (). However, it is noted that in non-convex scenarios, the performance of FedProx/SPPM in terms of convergence rates does not surpass that of SGD. Beyond federated learning, the versatility of SPPM is evident in its application to matrix and tensor completion such as in the work of Bumin & Huang (). Moreover, SPPM has been adapted for efficient implementation in a variety of optimization problems, as shown by Shtoff ().
While non-convex SPPM analysis presents significant challenges, with a full understanding of its convex counterpart still unfolding, recent studies such as the one by Khaled & Jin () have reported enhanced convergence by leveraging second-order similarity. Diverging from this approach, our contribution is the development of an efficient minibatch SPPM method SPPM-AS that shows improved results without depending on such assumptions. Significantly, we also provide the first empirical evidence that increasing local communication rounds in finding the proximal point can lead to a reduction in total communication costs.

### A.3 Local Solvers

**Table 2: Local optimizers for solving the proximal subproblem.**
| Setting | 1st order | 2nd order |
| --- | --- | --- |
| Strongly-Convex | Conjugate Gradients (CG)<br>Accelerated GD<br>Local GD<br>Scaffnew | BFGS<br>AICN<br>LocalNewton |
| Nonconvex | Mime-Adam<br>FedAdam-AdaGrad<br>FedSpeed | Apollo<br>OASIS |

In the exploration of local solvers for the SPPM-AS algorithm, the focus is on evaluating the performance impact of various inexact proximal solvers within federated learning settings, spanning both strongly convex and non-convex objectives. Here’s a simple summary of the algorithms discussed:

FedAdagrad-AdaGrad (Wang et al., ): Adapts AdaGrad for both client and server sides within federated learning, introducing local and global corrections to address optimizer state handling and solution bias.

BFGS (Broyden, ; Fletcher, ; Goldfarb, ; Shanno, ): A quasi-Newton method that approximates the inverse Hessian matrix to improve optimization efficiency, particularly effective in strongly convex settings but with limitations in distributed implementations.

AICN (Hanzely et al., ): Offers a global $O(1/k^{2})$ convergence rate under a semi-strong self-concordance assumption, streamlining Newton’s method without the need for line searches.

LocalNewton (Bischoff et al., ): Enhances local optimization steps with second-order information and global line search, showing efficacy in heterogeneous data scenarios despite a lack of extensive theoretical grounding.

Fed-LAMB (Karimi et al., ): Extends the LAMB optimizer to federated settings, incorporating layer-wise and dimension-wise adaptivity to accelerate deep neural network training.

FedSpeed (Sun et al., ): Aims to overcome non-vanishing biases and client-drift in federated learning through prox-correction and gradient perturbation steps, demonstrating effectiveness in image classification tasks.

Mime-Adam (Karimireddy et al., ): Mitigates client drift in federated learning by integrating global optimizer states and an SVRG-style correction term, enhancing the adaptability of Adam to distributed settings.

OASIS (Jahani et al., ): Utilizes local curvature information for gradient scaling, providing an adaptive, hyperparameter-light approach that excels in handling ill-conditioned problems.

Apollo (Ma, ): A quasi-Newton method that dynamically incorporates curvature information, showing improved efficiency and performance over first-order methods in deep learning applications.

Each algorithm contributes uniquely to the landscape of local solvers in federated learning, ranging from enhanced adaptivity and efficiency to addressing specific challenges such as bias, drift, and computational overhead.

## Appendix B Theoretical Overview and Recommendations

### B.1 Parameter Control

We have explored the effects of changing the hyperparameters of SPPM-AS on its theoretical properties, as summarized in Table 3. This summary shows that as the learning rate increases, the number of iterations required to achieve a target accuracy decreases, though this comes with an increase in neighborhood size. Focusing on sampling strategies, for SPPM-NICE employing NICE sampling, an increase in the sampling size $\tau_{\mathcal{S}}$ results in fewer iterations ($T$) and a smaller neighborhood. Furthermore, given that stratified sampling outperforms both block sampling and NICE sampling, we recommend adopting stratified sampling, as advised by Lemma 5.

**Table 3: Theoretical summary**
| Hyperparameter | Control | Rate (T) | Neighborhood |
| --- | --- | --- | --- |
| $\gamma$ | $\uparrow$ | $\downarrow$ | $\uparrow$ |
| $\mathcal{S}$ | $\tau_{\mathcal{S}}\uparrow$^(1) | $\downarrow$ | $\downarrow$ |
| Stratified sampling optimal clustering instead of $\mathrm{BS}$ or $\mathrm{NICE}$ sampling | $\downarrow$ | Lemma 5 |  |

### B.2 Comparison of Sampling Strategies

#### Full Sampling (FS).

Let $S=[n]$ with probability 1. Then SPPM-AS applied to Equation 9 becomes PPM (Moreau, ; Martinet, ) for minimizing $f$.
Moreover, in this case, we have $p_{i}=1$ for all $i\in[n]$ and Equation 5 takes on the form

$$ $\mu_{\mathrm{AS}}=\mu_{\mathrm{FS}}:=\frac{1}{n}\sum_{i=1}^{n}\mu_{i},\quad \sigma_{\star,\mathrm{AS}}^{2}=\sigma_{\star,\mathrm{FS}}^{2}:=0.$ $$

Note that $\mu_{\text{FS }}$ is the strong convexity constant of $f$, and that the neighborhood size is zero, as we would expect.

#### Nonuniform Sampling (NS).

Let $S=\{i\}$ with probability $p_{i}>0$, where $\sum_{i}p_{i}=1$. Then Equation 5 takes on the form

$$ $\displaystyle\mu_{\mathrm{AS}}=\mu_{\mathrm{NS}}:=\min_{i}\frac{\mu_{i}}{np_{i }},\quad\sigma_{\star,\mathrm{AS}}^{2}=\sigma_{\star,\mathrm{NS}}^{2}:=\frac{1 }{n}\sum_{i=1}^{n}\frac{1}{np_{i}}\left\|\nabla f_{i}\left(x_{\star}\right) \right\|^{2}.$ $$

If we take $p_{i}=\frac{\mu_{i}}{\sum_{j=1}^{n}\mu_{j}}$ for all $i\in[n]$, we shall refer to Algorithm 1 as SPPM with importance sampling (SPPM-IS). In this case,

$$ $\mu_{\mathrm{NS}}=\mu_{\mathrm{IS}}:=\frac{1}{n}\sum_{i=1}^{n}\mu_{i},\quad \sigma_{\star,\mathrm{NS}}^{2}=\sigma_{\star,\mathrm{IS}}^{2}:=\frac{\sum_{i=1 }^{n}\mu_{i}}{n}\sum_{i=1}^{n}\frac{\left\|\nabla f_{i}\left(x_{\star}\right) \right\|^{2}}{n\mu_{i}}.$ $$

This choice maximizes the value of $\mu_{\mathrm{NS}}$ (and hence minimizes the first part of the convergence rate) over the choice of the probabilities.

Table 4 summarizes the parameters associated with various sampling strategies, serving as a concise overview of the methodologies discussed in the main text. This summary facilitates a quick comparison and reference.

**Table 4: Arbitrary samplings comparison.**
| Setting/Requirement | $\mu_{\mathrm{AS}}$ | $\sigma_{\star,\mathrm{AS}}$ |
| --- | --- | --- |
| Full | $\frac{1}{n}\sum_{i=1}^{n}\mu_{i}$ | 0 |
| Non-Uniform | $\min_{i}\frac{\mu_{i}}{np_{i}}$ | $\frac{1}{n}\sum_{i=1}^{n}\frac{1}{np_{i}}\left\|\nabla f_{i}\left(x_{\star}<br>\right)\right\|^{2}$ |
| Nice | $\min_{C\subseteq[n],|C|=\tau}\frac{1}{\tau}\sum_{i\in C}\mu_{i}$ | $\sum_{C\subseteq[n],|C|=\tau}\frac{1}{\binom{n}{\tau}}\left\|\frac{1}{\tau}<br>\sum_{i\in C}\nabla f_{i}\left(x_{\star}\right)\right\|^{2}$ |
| Block | $\min_{j\in[b]}\frac{1}{nq_{j}}\sum_{i\in C_{j}}\mu_{i}$ | $\sum_{j\in[b]}q_{j}\left\|\sum_{i\in C_{j}}\frac{1}{np_{i}}\nabla f_{i}\left(x<br>_{\star}\right)\right\|^{2}$ |
| Stratified | $\min_{\mathbf{i}_{b}\in\mathbf{C}_{b}}\sum_{j=1}^{b}\frac{\mu_{i_{j}}\left|C_{<br>j}\right|}{n}$ | $\sum_{\mathbf{i}_{b}\in\mathbf{C}_{b}}\left(\prod_{j=1}^{b}\frac{1}{\left|C_{j<br>}\right|}\right)\left\|\sum_{j=1}^{b}\frac{\left|C_{j}\right|}{n}\nabla f_{i_{<br>j}}\left(x_{\star}\right)\right\|^{2}$ |
|  | Upper bound: $\frac{b}{n^{2}}\sum_{j=1}^{b}\left|C_{j}\right|^{2}\sigma_{j}^{2}$ |  |

### B.3 Extreme Cases of Block Sampling and Stratified Sampling

#### Extreme cases of block sampling.

We now consider two extreme cases:

- •
If $b=1$, then SPPM-BS = SPPM-FS = PPM. Let’s see, as a sanity check, whether we recover the right rate as well. We have $q_{1}=1,C_{1}=[n],p_{i}=1$ for all $i\in[n]$, and the expressions for $\mu_{\mathrm{AS}}$ and $\sigma_{\star,\text{ BS }}^{2}$ simplify to
$\displaystyle\mu_{\mathrm{BS}}=\mu_{\mathrm{FS}}:=\frac{1}{n}\sum_{i=1}^{n}\mu
_{i},\sigma_{\star,\mathrm{BS}}^{2}=\sigma_{\star,\mathrm{FS}}^{2}:=0.$
So, indeed, we recover the same rate as SPPM-FS.
- •
If $b=n$, then SPPM-BS = SPPM-NS. Let’s see, as a sanity check, whether we recover the right rate as well. We have $C_{i}=\{i\}$ and $q_{i}=p_{i}$ for all $i\in[n]$, and the expressions for $\mu_{\mathrm{AS}}$ and $\sigma_{\star,\mathrm{BS}}^{2}$ simplify to
$\displaystyle\mu_{\mathrm{BS}}=\mu_{\mathrm{NS}}:=\min_{i\in[n]}\frac{\mu_{i}}
{np_{i}},\quad\sigma_{\star,\mathrm{BS}}^{2}=\sigma_{\star,\mathrm{NS}}^{2}:=
\frac{1}{n}\sum_{i=1}^{n}\frac{1}{np_{i}}\left\|\nabla f_{i}\left(x_{\star}
\right)\right\|^{2}.$
So, indeed, we recover the same rate as SPPM-NS.

#### Extreme cases of stratified sampling.

We now consider two extreme cases:

- •
If $b=1$, then SPPM-SS = SPPM-US. Let’s see, as a sanity check, whether we recover the right rate as well. We have $C_{1}=[n],\left|C_{1}\right|=n,\left(\prod_{j=1}^{b}\frac{1}{\left|C_{j}\right
|}\right)=\frac{1}{n}$ and hence
$\displaystyle\mu_{\mathrm{SS}}=\mu_{\mathrm{US}}:=\min_{i}\mu_{i},\quad\sigma_
{\star,\mathrm{SS}}^{2}=\sigma_{\star,\mathrm{US}}^{2}:=\frac{1}{n}\sum_{i=1}^
{n}\left\|\nabla f_{i}\left(x_{\star}\right)\right\|^{2}.$
So, indeed, we recover the same rate as SPPM-US.
- •
If $b=n$, then SPPM-SS = SPPM-FS. Let’s see, as a sanity check, whether we recover the right rate as well. We have $C_{i}=\{i\}$ for all $i\in[n],\left(\prod_{j=1}^{b}\frac{1}{\left|C_{j}\right|}\right)=1$, and hence
$\displaystyle\mu_{\mathrm{SS}}=\mu_{\mathrm{FS}}:=\frac{1}{n}\sum_{i=1}^{n}\mu
_{i},\quad\sigma_{\star,\mathrm{SS}}^{2}=\sigma_{\star,\mathrm{FS}}^{2}:=0.$
So, indeed, we recover the same rate as SPPM-FS.

### B.4 Federated Averaging SPPM Baselines

In this section we propose two new algorithms based on Federated Averaging principle. Since to the best of our knowledge there are no federated averaging analyses within the same assumptions, we provide analysis of modified versions of SPPM-AS.

#### Averaging on prox γ ⁢ f i subscript prox 𝛾 subscript 𝑓 𝑖 \mathop{\mathrm{prox}}\nolimits_{\gamma f_{i}} roman_prox start_POSTSUBSCRIPT italic_γ italic_f start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_POSTSUBSCRIPT .

We introduce FedProx-SPPM-AS (see Algorithm 2), which is inspired by the principles of FedProx (Li et al., ). Unlike the traditional approach where a proximal operator is computed for the chosen cohort as a whole, in FedProx-SPPM-AS, we compute and then average the proximal operators calculated for each member within the cohort. However, this algorithm is not a simple case of SPPM-AS because it does not directly estimate the proximal operator at each step.

Figure: Algorithm 2 Proximal Averaging SPPM-AS (FedProx-SPPM-AS)

Figure: Algorithm 3 Federated Averaging SPPM-AS (FedAvg-SPPM-AS)

Here, we employ a proof technique similar to that of Theorem 4 and obtain the following convergence.

###### Theorem 7 (FedProx-SPPM-AS convergence) .

labelthm:FedProx-SPPM-AS
Let the number of local iterations $K=1$, and assume that 1 (differentiability) and 2 (strong convexity) hold. Let $x_{0}\in\mathbb{R}^{d}$ be an arbitrary starting point. Then, for any $t\geq 0$ and any $\gamma>0$, the iterates of FedProx-SPPM (as described in Algorithm 2) satisfy:

$$ $\mathbb{E}\left[{\left\lVert x_{t}-x_{\star}\right\rVert}^{2}\right]\leq A_{ \mathcal{S}}^{t}\left\|x_{0}-x_{\star}\right\|^{2}+\frac{B_{\mathcal{S}}}{1-A_ {\mathcal{S}}},$ $$

where $A_{\mathcal{S}}\coloneqq\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\frac{1}{\left|
S_{t}\right|}\sum_{i\in S_{t}}\frac{1}{1+\gamma\mu_{i}}\right]$ and $B_{\mathcal{S}}\coloneqq\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\frac{1}{\left|
S_{t}\right|}\sum_{i\in S_{t}}\frac{\gamma}{(1+\gamma\mu_{i})\mu_{i}}{\left
\lVert\nabla f_{i}(x_{\star}))\right\rVert}^{2}\right]$.

###### Theorem 7 (FedProx-SPPM-AS convergence) .

#### Federated averaging for prox prox \mathop{\mathrm{prox}}\nolimits roman_prox approximation.

An alternative method involves estimating the proximal operator by averaging the proximal operators calculated for each worker’s function. We call it *Federated Averaging Stochastic Proximal Point Method* (FedAvg-SPPM-AS, see Algorithm 3).
(FedAvg-SPPM-AS, see Algorithm 3).

After selecting and fixing a sample of workers $S_{k}$, the main objective is to calculate the proximal operator. This can be accomplished by approximating the proximal calculation with the goal of minimizing $\tilde{f}_{S}(x)=f_{S}(x)+\frac{2}{\gamma}\left\|x-x_{t}\right\|^{2}$. It can be observed that this approach is equivalent to FedProx-SPPM-AS, as at each local step we calculate

$$ $\mathop{\mathrm{prox}}\nolimits_{\alpha\tilde{f_{i}}}(x_{k,t})\coloneqq \operatorname*{arg\,min}_{z\in\mathbb{R}^{d}}\left[\tilde{f_{i}}(z)+\frac{2}{ \alpha}{\left\lVert z-x_{k,t}\right\rVert}^{2}\right]=\operatorname*{arg\,min} _{z\in\mathbb{R}^{d}}\left[f_{i}(z)+\left(\frac{2}{\gamma}+\frac{2}{\alpha} \right){\left\lVert z-x_{k,t}\right\rVert}^{2}\right].$ $$

## Appendix C Training Details

### C.1 Non-IID Data Generation

In our study, we validate performance and compare the benefits of SPPM-AS over SPPM using well-known datasets such as mushrooms, a6a, w6a, and ijcnn1.bz2 from LibSVM (Chang & Lin, ). To ensure relevance to our research focus, we adopt a feature-wise non-IID setting, characterized by variation in feature distribution across clients. This variation is introduced by clustering the features using the K-means algorithm, with the number of clusters set to $10$ and the number of clients per cluster fixed at $10$ for simplicity. We visualize the clustered data using t-SNE in Figure 5, where we observe that the data are divided into 10 distinct clusters with significantly spaced cluster centers.

Figure: (a) mushrooms, 10 clusters
Refer to caption: x13.png

### C.2 Sampling

To simulate random sampling among clients within these 10 clusters, where each cluster comprises 10 clients, we consider two contrasting scenarios:

- •
Case I - SPPM-BS: Assuming clients within the same cluster share similar features and data distributions, sampling all clients from one cluster (i.e., $C=10$ clients) results in a homogeneous sample.
- •
Case II - SPPM-SS: Conversely, by traversing all 10 clusters and randomly sampling one client from each, we obtain a group of 10 clients representing maximum heterogeneity.

We hypothesize that any random sampling from the 100 clients will yield performance metrics lying between these two scenarios. In Figure 6, we examine the impact of sampling clients with varying degrees of heterogeneity using a fixed learning rate of $0.1$. Our findings indicate that heterogeneous sampling results in a significantly smaller convergence neighborhood $\sigma^{2}_{\star}$. This outcome is attributed to the broader global information captured through heterogeneous sampling, in contrast to homogeneous sampling, which increases the data volume without contributing additional global insights. As these two sampling strategies represent the extremes of arbitrary sampling, any random selection will fall between them in terms of performance. Given their equal cost and the superior performance of the SPPM-SS strategy in heterogeneous FL environments, we designate SPPM-SS as our default sampling approach.

Figure: Figure 6: Comparison with SPPM-SS and SPPM-BS samplings.
Refer to caption: x15.png

Figure: Algorithm 4 SPPM-AS Adaptation for Federated Learning

### C.3 SPPM-AS Algorithm Adaptation for Federated Learning

In the main text, Algorithm 1 outlines the general form of SPPM-AS. For the convenience of implementation in FL contexts and to facilitate a better understanding, we introduce a tailored version of the SPPM-AS algorithm specific to FL, designated as Algorithm 4. Notably, as block sampling is adopted as our default method, this adaptation of the algorithm specifically addresses the nuances of the block sampling approach. We also conducted arbitrary sampling on synthetic datasets and neural networks to demonstrate the algorithm’s versatility.

## Appendix D Additional Experiments on Logistic Regression

### D.1 Communication Cost on Various Datasets to a Target Accuracy

In Figure 1, we presented the total communication cost relative to the number of rounds required to achieve the target accuracy for the selected cohort. In this section, we provide more details on how is this figure was obtained and present additional results for various datasets.

Figure: Figure 7: Total communication cost with respect to the local communication round. For LocalGD, $K$ represents the local communication round $K$ for finding the prox of the current model. For LocalGD, we slightly abuse the x-axis, which represents the total number of local iterations, no local communication is required. We calculate the total communication cost to reach a fixed global accuracy $\epsilon$ such that ${\left\lVert x_{t}-x_{\star}\right\rVert}^{2}<\epsilon$. LocalGD, optim represents using the theoretical optimal stepsize of LocalGD with minibatch sampling.
Refer to caption: x18.png

Figure: Figure 8: $K=4$.
Refer to caption: x21.png

### D.2 Convergence Speed and σ ⋆ , SS 2 subscript superscript 𝜎 2 ⋆ SS \sigma^{2}_{\star,\mathrm{SS}} italic_σ start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT start_POSTSUBSCRIPT ⋆ , roman_SS end_POSTSUBSCRIPT Trade-Off

Unlike SGD-type methods such as MB-GD and MB-LocalGD, in which the largest allowed learning rate is $1/A$, where $A$ is a constant proportion to the smoothness of the function we want to optimize (Gower et al., ). For larger learning rate, SGD-type method may not converge and exploding. However, for stochastic proximal point methods, they have a very descent benefit of allowing arbitrary learning rate. In this section, we verify whether our proposed method can allow arbitrary learning rate and whether we can find something interesting. We considered different learning rate scale from 1e-5 to 1e+5. We randomly selected three learning rates [0.1, 1, 100] for visual representation with the results presented in Figure 9 and Figure 9. We found that a larger learning rate leads to a faster convergence rate but results in a much larger neighborhood, $\sigma^{2}_{\star,\mathrm{SS}}/\mu^{2}_{\mathrm{SS}}$. This can be considered a trade-off between convergence speed and neighborhood size, $\sigma^{2}_{\star,\mathrm{SS}}$. By default, we consider setting the learning rate to $1.0$ which has a good balance between the convergence speed and the neighborhood size.

In this section, we extend our analysis by providing additional results across a broader range of datasets and varying learning rates. Specifically, Figure 9 illustrates the outcomes using 4 local communication rounds ($K=4$), while Figure 9 details the results for 16 local communication rounds ($K=16$). Previously, in Figure 1, we explored the advantages of larger $K$ values. Here, our focus shifts to determining if similar trends are observable across different $K$ values. Through comprehensive evaluations on various datasets and multiple $K$ settings, we have confirmed that lower learning rates in SPPM-AS result in slower convergence speeds; however, they also lead to a smaller final convergence neighborhood.

### D.3 Additional Experiments on Hierarchical Federated Learning

In Figure 2(d) of the main text, we detail the total communication cost for hierarchical Federated Learning (FL) utilizing parameters $c_{1}=0.1$ and $c_{2}=1$ on the a6a dataset. Our findings reveal that SPPM-AS achieves a significant reduction in communication costs, amounting to $94.87\%$, compared with the conventional FL setting where $c_{1}=1$ and $c_{2}=1$, which shows a 74.36% reduction. In this section, we extend our analysis with comprehensive evaluations on additional datasets, namely ijcnn1.bz2, a9a, and mushrooms. Beyond considering $c_{1}=0.1$, we further explore the impact of reducing the local communication cost from each client to the corresponding hub to $c_{1}=0.05$. The results, presented in Figure 10 and the continued Figure 11, reinforce our observation: hierarchical FL consistently leads to further reductions in communication costs. A lower $c_{1}$ parameter correlates with even greater savings in communication overhead. These results not only align with our expectations but also underscore the efficacy of our proposed SPPM-AS in cross-device FL settings.

Figure: Figure 10: The total communication cost is analyzed with respect to the number of local communication rounds. For LocalGD, $K$ represents the local communication round used for finding the prox of the current model. In the case of LocalGD, we slightly abuse the x-axis to represent the total number of local iterations, as no local communication is required. We calculate the total communication cost needed to reach a fixed global accuracy $\epsilon$, such that ${\left\lVert x_{t}-x_{\star}\right\rVert}^{2}<\epsilon$. LocalGD, optim denotes the use of the theoretically optimal stepsize for LocalGD with minibatch sampling. Comparisons are made between different prox solvers (CG and BFGS).
Refer to caption: x23.png

Figure: (a) standard FL, $c_{1}=1,c_{2}=0$
Refer to caption: x26.png

## Appendix E Additional Neural Network Experiments

### E.1 Experiment Details

For our neural network experiments, we used the FEMNIST dataset (Caldas et al., ). Each client was created by uniformly selecting from user from original dataset, inherently introducing heterogeneity among clients. We tracked and reported key evaluation metrics—training and testing loss and accuracy—after every 5 global communication rounds. The test dataset was prepared by dividing each user’s data into a 9:1 ratio, following the partitioning approach of the FedLab framework (Zeng et al., ). For the SPPM-AS algorithm, we selected Adam as the optimizer for the proximal operator. The learning rate was determined through a grid search across the following range: $[0.0001,0.0005,0.001,0.005,0.01,0.05,0.1,0.5]$. The model architecture comprises a convolutional neural network (CNN) with the following layers: Conv2d(1, 32, 5), ReLU, Conv2d(32, 64, 5), MaxPool2d(2, 2), a fully connected (FC) layer with 128 units, ReLU, and another FC layer with 128 units, as specified in Table 5. Dropout, learning rate scheduling, gradient clipping, etc., were not used to improve the interpretability of results.

We explore various values of targeted training accuracy, as illustrated in Figure 12. This analysis helps us understand the impact of different accuracy thresholds on the model’s performance. For instance, we observe that as the target accuracy changes, SPPM-NICE consistently outperforms LocalGD in terms of total communication cost. As the target accuracy increases, the performance gap between these two algorithms also widens. Additionally, we perform ablation studies on different values of $c_{1}$, as shown in Figure 13, to assess their effects on the learning process. Here, we note that with $c_{2}=0.2$, SPPM-NICE performs similarly to LocalGD, suggesting that an increase in $c_{2}$ value could narrow the performance gap between SPPM-NICE and LocalGD.

**Table 5: Architecture of the CNN model for FEMNIST symbol recognition.**
| Layer | Output Shape | # of Trainable Parameters | Activation | Hyperparameters |
| --- | --- | --- | --- | --- |
| Input | (28, 28, 1) | 0 |  |  |
| Conv2d | (24, 24, 32) | 832 | ReLU | kernel size = 5; strides = (1, 1) |
| Conv2d | (10, 10, 64) | 51,264 | ReLU | kernel size = 5; strides = (1, 1) |
| MaxPool2d | (5, 5, 64) | 0 |  | pool size = (2, 2) |
| Flatten | 6400 | 0 |  |  |
| Dense | 128 | 819,328 | ReLU |  |
| Dense | 62 | 7,998 | softmax |  |

Figure: Figure 12: Varying targeted training accuracy level for SPPM-AS.
Refer to caption: x32.png

Figure: Figure 13: Varying $c_{1}$ cost.
Refer to caption: x35.png

### E.2 Convergence Analysis Compared with Baselines

Further, we compare SPPM-AS, SPPM, and LocalGD in Figure 15, placing a particular emphasis on evaluating the total computational complexity. This measure gains importance in scenarios where communication rounds are of secondary concern, thereby shifting the focus to the assessment of computational resource expenditure.

Figure: Figure 14: Different local solvers for prox baselines for training a CNN model over 100 workers using data from the FEMNIST dataset. The number of local communication rounds is fixed at 3 and the number of worker optimizer steps is fixed at 3. Nice sampling with a minibatch size of 10 is used. $\gamma$ is fixed at 1.0.
Refer to caption: x38.png

### E.3 Prox Solvers Baselines

We compare baselines from A.3 for training a CNN model over 100 workers using data from the FEMNIST dataset, as shown in Figure 14. The number of local communication rounds and worker optimizer steps is consistent among various solvers for the purpose of fair comparison. All local solvers optimize the local objective, which is prox on the selected cohort. The solvers compared are: LocalGD referred as FedSGD (McMahan et al., ) - the Federated Averaging algorithm with SGD as the worker optimizer, FedAdam - the Federated Averaging algorithm with Adam as the worker optimizer, FedAdam-Adam based on the FedOpt framework (Reddi et al., ), and finally MimeLite-Adam, which is based on the Mime (Karimireddy et al., ) framework and the Adam optimizer. The hyperparameter search included a double-level sweep of the optimizer learning rates: $[0.00001,0.0001,0.001,0.01,0.1]$, followed by $[0.25,0.5,1.0,2.5,5]*lr_{\text{best}}$. One can see that all methods perform similarly, with MimeLite-Adam and FedSGD converging better on the test data.

Figure: Figure 15: Accuracy compared with baselines.
Refer to caption: x39.png

## Appendix F Missing Proof and Additional Theoretical Analysis

### F.1 Facts Used in the Proof

###### Fact 1 (Differentiation of integral with a parameter (theorem 2.27 from Folland ( 1984 ) )) .

Suppose that $f:X\times[a,b]\rightarrow\mathbb{C}(-\infty<a<b<\infty)$ and that $f(\cdot,t):X\rightarrow\mathbb{C}$ is integrable for each $t\in[a,b]$. Let $F(t)=\int_{X}f(x,t)d\mu(x)$.

- a.
Suppose that there exists $g\in L^{1}(\mu)$ such that $|f(x,t)|\leq g(x)$ for all $x,t$. If $\lim_{t\rightarrow t_{0}}f(x,t)=f\left(x,t_{0}\right)$ for every $x$, then $\lim_{t\rightarrow t_{0}}F(t)=F\left(t_{0}\right)$; in particular, if $f(x,\cdot)$ is continuous for each $x$, then $F$ is continuous.
- b.
Suppose that $\partial f/\partial t$ exists and there is a $g\in L^{1}(\mu)$ such that $|(\partial f/\partial t)(x,t)|\leq$ $g(x)$ for all $x,t$. Then $F$ is differentiable and $F^{\prime}(x)=\int(\partial f/\partial t)(x,t)d\mu(x)$.

###### Fact 2 (Tower Property) .

For any random variables $X$ and $Y$, we have

$$ $\displaystyle\mathbb{E}\left[\mathbb{E}\left[X|Y\right]\right]=\mathbb{E}\left [X\right].$ $$

###### Fact 3 (Every point is a fixed point (Khaled & Jin, 2023 ) ) .

Let $\varphi:\mathbb{R}^{d}\to\mathbb{R}$ be a convex differentiable function. Then

$$ $\displaystyle\mathop{\mathrm{prox}}\nolimits_{\gamma\varphi}(x+\gamma\nabla \varphi(x))=x,\qquad\forall\gamma>0,\quad\forall x\in\mathbb{R}^{d}.$ $$

In particular, if $x_{\star}$ is a minimizer of $\varphi$, then $\mathop{\mathrm{prox}}\nolimits_{\gamma\varphi}(x_{\star})=x_{\star}$.

###### Proof.

Evaluating the proximity operator is equivalent to

$$ $\displaystyle\mathop{\mathrm{prox}}\nolimits_{\gamma\varphi}(y)=\operatorname* {arg\,min}_{x\in\mathbb{R}^{d}}\left(\varphi(x)+\frac{1}{2\gamma}{\left\lVert x -y\right\rVert}^{2}\right).$ $$

This is a strongly convex minimization problem for any $\gamma>0$, hence the (necessarily unique) minimizer $x=\mathop{\mathrm{prox}}\nolimits_{\gamma\varphi}(y)$ of this problem satisfies the first-order optimality condition

$$ $\displaystyle\nabla\varphi(x)+\frac{1}{\gamma}(x-y)=0.$ $$

Solving for $y$, we observe that this holds for $y=x+\gamma\nabla\phi(x)$. Therefore, $x=\mathop{\mathrm{prox}}\nolimits_{\gamma\varphi}(x+\gamma\nabla\varphi(x))$.
∎

###### Fact 4 (Contractivity of the prox (Mishchenko et al., 2022a ) ) .

If $\varphi$ is differentiable and $\mu$-strongly convex, then for all $\gamma>0$ and for any $x,y\in\mathbb{R}^{d}$ we have

$$ $\displaystyle\left\|\mathop{\mathrm{prox}}\nolimits_{\gamma\varphi}(x)-\mathop {\mathrm{prox}}\nolimits_{\gamma\varphi}(y)\right\|^{2}\leq\frac{1}{(1+\gamma \mu)^{2}}\left\|x-y\right\|^{2}.$ $$

###### Fact 5 (Recurrence ( Khaled & Jin, 2023 , Lemma 1)) .

Assume that a sequence $\{s_{t}\}_{t\geq 0}$ of positive real numbers for all $t\geq 0$ satisfies

$$ $\displaystyle s_{t+1}\leq as_{t}+b,$ $$

where $0<a<1$ and $b\geq 0$. Then the sequence for all $t\geq 0$ satisfies

$$ $\displaystyle s_{t}\leq a^{t}s_{0}+b\min\left\{t,\frac{1}{1-a}\right\}.$ $$

###### Proof.

Unrolling the recurrence, we get

$$ $\displaystyle s_{t}\leq as_{t-1}+b\leq a(as_{t-2}+b)+b\leq\cdots\leq a^{t}s_{0 }+b\sum_{i=0}^{t-1}a^{i}.$ $$

We can now bound the sum $\sum_{i=0}^{t-1}a^{i}$ in two different ways. First, since $a<1$, we get the estimate

$$ $\displaystyle\sum_{i=0}^{t-1}a^{i}\leq\sum^{t-1}_{i=0}1=t.$ $$

Second, we sum a geometic series

$$ $\displaystyle\sum_{i=0}^{t-1}a^{i}\leq\sum_{i=0}^{\inf}a^{i}=\frac{1}{1-a}.$ $$

Note that either of these bounds can be better. So, we apply the best of these bounds. Substituing the above two bounds gived the target inequality.
∎

###### Fact 1 (Differentiation of integral with a parameter (theorem 2.27 from Folland ( 1984 ) )) .

###### Fact 2 (Tower Property) .

###### Fact 3 (Every point is a fixed point (Khaled & Jin, 2023 ) ) .

###### Proof.

###### Fact 4 (Contractivity of the prox (Mishchenko et al., 2022a ) ) .

###### Fact 5 (Recurrence ( Khaled & Jin, 2023 , Lemma 1)) .

###### Proof.

### F.2 Simplified Proof of SPPM

We provide a simplified proof of SPPM (Khaled & Jin, ) in this section.
Using the fact that $x_{\star}=\mathop{\mathrm{prox}}\nolimits_{\gamma f_{\xi_{t}}}(x_{\star}+
\gamma\nabla f_{\xi_{t}}(x_{\star}))$ (see 3) and then applying contraction of the prox (4), we get

$$ $\displaystyle{\left\lVert x_{t+1}-x_{\star}\right\rVert}^{2}$ $\displaystyle={\left\lVert\mathop{\mathrm{prox}}\nolimits_{\gamma f_{\xi_{t}}} -x_{\star}\right\rVert}^{2}$ $\displaystyle\stackrel{{\scriptstyle(\lx@cref{creftype~refnum}{fact:fact1})}}{ {=}}{\left\lVert\mathop{\mathrm{prox}}\nolimits_{\gamma f_{\xi_{t}}}(x_{t})- \mathop{\mathrm{prox}}\nolimits_{\gamma f_{\xi_{t}}}(x_{\star}+\gamma\nabla f_ {\xi_{t}}(x_{\star}))\right\rVert}^{2}$ $\displaystyle\stackrel{{\scriptstyle(\lx@cref{creftype~refnum}{fact:fact2})}}{ {\leq}}\frac{1}{(1+\gamma\mu)^{2}}{\left\lVert x_{t}-(x_{\star}+\gamma\nabla f _{\xi_{t}}(x_{\star}))\right\rVert}^{2}$ $\displaystyle=\frac{1}{(1+\gamma\mu)^{2}}\left({\left\lVert x_{t}-x_{\star} \right\rVert}^{2}-2\gamma\left\langle\nabla f_{\xi_{t}}(x_{\star}),x_{t}-x_{ \star}\right\rangle+\gamma^{2}{\left\lVert\nabla f_{\xi_{t}}(x_{\star})\right \rVert}^{2}\right).$ $$

Taking expectation on both sides, conditioned on $x_{t}$, we get

$$ $\displaystyle\mathbb{E}\left[{\left\lVert x_{t+1}-x_{\star}\right\rVert}^{2}|x _{t}\right]$ $\displaystyle\leq\frac{1}{(1+\gamma\mu)^{2}}\left({\left\lVert x_{t}-x_{\star} \right\rVert}^{2}-2\gamma\left\langle\mathbb{E}\left[\nabla f_{\xi_{t}}(x_{ \star})\right],x_{t}-x_{\star}\right\rangle+\gamma^{2}\mathbb{E}\left[{\left \lVert\nabla f_{\xi_{t}}(x_{\star})\right\rVert}^{2}\right]\right)$ $\displaystyle=\frac{1}{(1+\gamma\mu)^{2}}\left({\left\lVert x_{t}-x_{\star} \right\rVert}^{2}+\gamma^{2}\sigma^{2}_{\star}\right),$ $$

where we used the fact that $\mathbb{E}\left[\nabla f_{\xi_{t}}(x_{\star})\right]=\nabla f(x_{\star})=0$ and $\sigma^{2}_{\star}\coloneqq\mathbb{E}\left[{\left\lVert\nabla f_{\xi_{t}}(x_{
\star})\right\rVert}^{2}\right]$. Taking expectation again and applying the tower property (2), we get

$$ $\displaystyle\mathbb{E}\left[{\left\lVert x_{t+1}-x_{\star}\right\rVert}^{2} \right]\leq\frac{1}{(1+\gamma\mu)^{2}}\left({\left\lVert x_{t}-x_{\star}\right \rVert}^{2}+\gamma^{2}\sigma^{2}_{\star}\right).$ $$

It only remains to solve the above recursion. Luckily, that is exactly what 5 does. In particular, we use it with $s_{t}=\mathbb{E}\left[{\left\lVert x_{t}-x_{\star}\right\rVert}^{2}\right],a=
\frac{1}{(1+\gamma\mu)^{2}}$ and $b=\frac{\gamma^{2}\sigma_{\star}^{2}}{(1+\gamma\mu)^{2}}$ to get

$$ $\displaystyle\mathbb{E}\left[{\left\lVert x_{t}-x_{\star}\right\rVert}^{2}\right]$ $\displaystyle\stackrel{{\scriptstyle(\lx@cref{creftype~refnum}{fact:fact3})}}{ {\leq}}\left(\frac{1}{1+\gamma\mu}\right)^{2t}{\left\lVert x_{0}-x_{\star} \right\rVert}^{2}+\frac{\gamma^{2}\sigma^{2}_{\star}}{(1+\gamma\mu)^{2}}\min \left\{t,\frac{(1+\gamma\mu)^{2}}{(1+\gamma\mu)^{2}-1}\right\}$ $\displaystyle\leq\left(\frac{1}{1+\gamma\mu}\right)^{2t}{\left\lVert x_{0}-x_{ \star}\right\rVert}^{2}+\frac{\gamma^{2}\sigma^{2}_{\star}}{(1+\gamma\mu)^{2}-1}$ $\displaystyle\leq\left(\frac{1}{1+\gamma\mu}\right)^{2t}{\left\lVert x_{0}-x_{ \star}\right\rVert}^{2}+\frac{\gamma\sigma^{2}_{\star}}{\gamma\mu^{2}+2\mu}.$ $$

### F.3 Missing Proof of Theorem 4

We first prove the following useful lemma.

###### Lemma 8 .

Let $\phi_{\xi}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be differentiable functions for almost all $\xi\sim\mathcal{D}$, with $\phi_{\xi}$ being $\mu_{\xi}$-strongly convex for almost all $\xi\sim\mathcal{D}$. Further, let $w_{\xi}$ be positive scalars. Then the function $\phi:=\mathbb{E}_{\xi\sim\mathcal{D}}\left[w_{\xi}\phi_{\xi}\right]$ is $\mu$-strongly convex with $\mu=\mathbb{E}_{\xi\sim\mathcal{D}}\left[w_{\xi}\mu_{\xi}\right]$.

###### Proof.

By assumption,

$$ $\phi_{\xi}(y)+\left\langle\nabla\phi_{\xi}(y),x-y\right\rangle+\frac{\mu_{\xi} }{2}\|x-y\|^{2}\leq\phi_{\xi}(x),\quad\text{for almost all }\xi\in\mathcal{D}, \forall x,y\in\mathbb{R}^{d}.$ $$

This means that

$$ $\mathbb{E}_{\xi\sim\mathcal{D}}\left[w_{\xi}\left(\phi_{\xi}(y)+\left\langle \nabla\phi_{\xi}(y),x-y\right\rangle+\frac{\mu_{\xi}}{2}\|x-y\|^{2}\right) \right]\leq\mathbb{E}_{\xi\sim\mathcal{D}}\left[w_{\xi}\phi_{\xi}(x)\right], \quad\forall x,y\in\mathbb{R}^{d},$ $$

which is equivalent to

$$ $\phi(y)+\langle\nabla\phi(y),x-y\rangle+\frac{\mathbb{E}_{\xi\sim\mathcal{D}} \left[w_{\xi}\mu_{\xi}\right]}{2}\|x-y\|^{2}\leq\phi(x),\quad\forall x,y\in \mathbb{R}^{d},$ $$

So, $\phi$ is $\mu$-strongly convex.
∎

Now, we are ready to prove our main Theorem 4.

###### Proof.

Let $C$ be any (necessarily nonempty) subset of $[n]$ such that $p_{C}>0$. Recall that in view of Equation 8 we have

$$ $f_{C}(x)=\mathbb{E}_{\xi\sim\mathcal{D}}\left[\frac{I\left(\xi\in C\right)}{p_ {\xi}}f_{\xi}(x)\right]$ $$

i.e., $f_{C}$ is a conic combination of the functions $\left\{f_{\xi}:\xi\in C\right\}$ with weights $w_{\xi}=\frac{I\left(\xi\in C\right)}{p_{\xi}}$. Since each $f_{\xi}$ is $\mu_{\xi}$-strongly convex, Lemma 8 says that $f_{C}$ is $\mu_{C}$-strongly convex with

$$ $\mu_{C}:=\mathbb{E}_{\xi\sim\mathcal{D}}\left[\frac{I\left(\xi\in C\right)\mu_ {\xi}}{p_{\xi}}\right].$ $$

So, every such $f_{C}$ is $\mu$-strongly convex with

$$ $\mu=\mu_{\mathrm{AS}}:=\min_{C\subseteq[n],p_{C}>0}\mathbb{E}_{\xi\sim\mathcal {D}}\left[\frac{I\left(\xi\in C\right)\mu_{\xi}}{p_{\xi}}\right].$ $$

Further, the quantity $\sigma_{\star}^{2}$ from (2.3) is equal to

$$ $\sigma_{\star}^{2}:=\mathrm{E}_{\xi\sim\mathcal{D}}\left[\left\|\nabla f_{\xi} \left(x_{\star}\right)\right\|^{2}\right]\stackrel{{\scriptstyle Eqn.~{}(\ref{ eqn:8010})}}{{=}}\sum_{C\subseteq[n],p_{C}>0}p_{C}\left\|\nabla f_{C}\left(x_{ \star}\right)\right\|^{2}:=\sigma_{\star,\mathrm{AS}}^{2}.$ $$

Incorporating Section F.2 into the above equation, we prove the theorem.
∎

###### Lemma 8 .

###### Proof.

###### Proof.

### F.4 Theory for Expectation Formulation

We will formally define our optimization objective, focusing
on minimization in expectation form. We consider

$$ $\displaystyle\min_{x\in\mathbb{R}^{d}}{f(x)\coloneqq\mathbb{E}_{\xi\sim \mathcal{D}}\left[f_{\xi}(x)\right]},$ (6) $$

where $f_{\xi}:\mathbb{R}^{d}\to\mathbb{R}$, $\xi\sim\mathcal{D}$ is a random variable following distribution $\mathcal{D}$.

###### Assumption 9 .

Function $f_{\xi}:\mathbb{R}^{d}\to\mathbb{R}$ is differentiable for almost all samples $\xi\sim\mathcal{D}$.

This implies that $f$ is differentiable. We will implicitly assume that the order of differentiation and expectation can be swapped (^1^11This assumption satisfies the conditions required for the theorem about differentiating an integral with a parameter (1).), which means that

$$ $\displaystyle\nabla f(x)\stackrel{{\scriptstyle Eqn.~{}(\ref{eqn:obj_1})}}{{=} }\nabla\mathbb{E}_{\xi\sim\mathcal{D}}\left[f_{\xi}(x)\right]=\mathbb{E}_{\xi \sim\mathcal{D}}\left[\nabla f_{\xi}(x)\right].$ $$

###### Assumption 10 .

Function $f_{\xi}:\mathbb{R}^{d}\to\mathbb{R}$ is $\mu$-strongly convex for almost all samples $\xi\sim\mathcal{D}$, where $\mu>0$. That is

$$ $\displaystyle f_{\xi}(y)+\left\langle\nabla f_{\xi},x-y\right\rangle+\frac{\mu }{2}{\left\lVert x-y\right\rVert}^{2}\leq f_{\xi}(x),$ $$

for all $x,y\in\mathbb{R}^{d}$.

This implies that $f$ is $\mu$-strongly convex, and hence $f$ has a unique minimizer, which we denote by $x_{\star}$. We know that $\nabla f(x_{\star})=0$. Notably, we do *not* assume $f$ to be $L$-smooth.

Let $\mathcal{S}$ be a probability distribution over all *finite* subsets of $\mathbb{N}$. Given a random set $S\sim\mathcal{S}$, we define

$$ $\displaystyle p_{i}\coloneqq\operatorname{Prob}(i\in{S}),\quad i\in\mathbb{N}.$ $$

We will restrict our attention to proper and nonvacuous random sets.

###### Assumption 11 .

$S$ is proper (i.e., $p_{i}>0$ for all $i\in\mathbb{N}$) and nonvacuous (i.e., $\operatorname{Prob}({S}=\emptyset)=0$).

Let $C$ be the selected cohort. Given $\emptyset\neq C\subset\mathbb{N}$ and $i\in\mathbb{N}$, we define

$$ $\displaystyle v_{i}(C):=\begin{cases}\frac{1}{p_{i}}&i\in C\\ 0&i\notin C,\end{cases}$ (7) $$

and

$$ $\displaystyle f_{C}(x):=\mathbb{E}_{\xi\sim\mathcal{D}}\left[v_{\xi}(C)f_{\xi} (x)\right]\stackrel{{\scriptstyle Eqn.~{}(\ref{eqn:8001})}}{{=}}\mathbb{E}_{ \xi\sim\mathcal{D}}\left[\frac{I\left(\xi\in C\right)}{p_{\xi}}f_{\xi}(x) \right].$ (8) $$

Note that $v_{i}(S)$ is a random variable and $f_{S}$ is a random function. By construction, $\mathrm{E}_{S\sim\mathcal{S}}\left[v_{i}(S)\right]=1$ for all $i\in\mathbb{N}$, and hence

$$ $\displaystyle\mathbb{E}_{{S}\sim\mathcal{S}}\left[f_{{S}}(x)\right]=\mathbb{E} _{{S}\sim\mathcal{S}}\left[\mathbb{E}_{\xi\sim\mathcal{D}}\left[v_{\xi}(C) \nabla f_{\xi}(x)\right]\right]$ $\displaystyle\qquad=\mathbb{E}_{\xi\sim\mathcal{D}}\left[\mathbb{E}_{{S}\sim \mathcal{S}}\left[v_{\xi}(S)\right]\nabla f_{\xi}(x)\right]=\mathbb{E}_{\xi \sim\mathcal{D}}\left[f_{\xi}(x)\right]=f(x).$ $$

Therefore, the optimization problem in Equation 1 is equivalent to the stochastic optimization problem

$$ $\displaystyle\min_{x\in\mathbb{R}^{d}}\left\{f(x):=\mathrm{E}_{S\sim\mathcal{S }}\left[f_{S}(x)\right]\right\}.$ (9) $$

Further, if for each $C\subset\mathbb{N}$ we let $p_{C}:=\operatorname{Prob}(S=C),f$ can be written in the equivalent form

$$ $\displaystyle f(x)=\mathbb{E}_{S\sim\mathcal{S}}\left[f_{S}(x)\right]=\sum_{C \subset\mathbb{N}}p_{C}f_{C}(x)=\sum_{C\subset\mathbb{N},p_{C}>0}p_{C}f_{C}(x).$ (10) $$

###### Theorem 12 (Main Theorem) .

Let 1 (diferentiability) and 2 (strong convexity) hold. Let $S$ be a random set satisfying 3, and define

$$ $\displaystyle\mu_{\mathrm{AS}}$ $\displaystyle:=\min_{C\subset\mathbb{N},p_{C}>0}\mathbb{E}_{\xi\sim\mathcal{D} }\left[\frac{I\left(\xi\in C\right)\mu_{\xi}}{p_{\xi}}\right],$ $\displaystyle\sigma_{\star,\mathrm{AS}}^{2}$ $\displaystyle:=\sum_{C\subset\mathbb{N},p_{C}>0}p_{C}\left\|\nabla f_{C}\left( x_{\star}\right)\right\|^{2}.$ (11) $$

Let $x_{0}\in\mathbb{R}^{d}$ be an arbitrary starting point. Then for any $t\geq 0$ and any $\gamma>0$, the iterates of SPPM-AS (Algorithm 1) satisfy

$$ $\mathrm{E}\left[\left\|x_{t}-x_{\star}\right\|^{2}\right]\leq\left(\frac{1}{1+ \gamma\mu_{\mathrm{AS}}}\right)^{2t}\left\|x_{0}-x_{\star}\right\|^{2}+\frac{ \gamma\sigma_{\star,\mathrm{AS}}^{2}}{\gamma\mu_{\mathrm{AS}}^{2}+2\mu_{ \mathrm{AS}}}.$ $$

###### Assumption 9 .

###### Assumption 10 .

###### Assumption 11 .

###### Theorem 12 (Main Theorem) .

### F.5 Missing Proof of Iteration Complexity of SPPM-AS

We have seen above that accuracy arbitrarily close to (but not reaching) $\nicefrac{{\sigma^{2}_{\star,\mathrm{AS}}}}{{\mu_{\mathrm{AS}}^{2}}}$ can be achieved via a single step of the method, provided the stepsize $\gamma$ is large enough. Assume now that we aim for $\epsilon$ accuracy where $\epsilon\leq\nicefrac{{\sigma^{2}_{\star,\mathrm{AS}}}}{{\mu_{\mathrm{AS}}^{2}}}$. Using the inequality $1-k\leq\exp(-k)$ which holds for all $k>0$, we get

$$ $\displaystyle\left(\frac{1}{1+\gamma\mu_{\mathrm{AS}}}\right)^{2t}=\left(1- \frac{\gamma\mu}{1+\gamma\mu_{\mathrm{AS}}}\right)^{2t}\leq\exp\left(-\frac{2 \gamma\mu_{\mathrm{AS}}t}{1+\gamma\mu_{\mathrm{AS}}}\right)$ $$

Therefore, provided that

$$ $t\geq\frac{1+\gamma\mu_{\mathrm{AS}}}{2\gamma\mu_{\mathrm{AS}}}\log\left(\frac {2\left\|x_{0}-x_{\star}\right\|^{2}}{\varepsilon}\right),$ $$

we get $\left(\frac{1}{1+\gamma\mu_{\mathrm{AS}}}\right)^{2t}\left\|x_{0}-x_{\star}
\right\|^{2}\leq\frac{\varepsilon}{2}$. Furthermore, as long as $\gamma\leq\frac{2\varepsilon\mu_{\mathrm{AS}}}{2\sigma_{\star,\mathrm{AS}}^{2}
-\varepsilon\mu_{\mathrm{AS}}^{2}}$ (this is true provided that the more restrictive but also more elegant-looking condition $\gamma\leq\nicefrac{{\varepsilon\mu_{\mathrm{AS}}}}{{\sigma_{\star,\mathrm{AS}
}^{2}}}$ holds), we get
$\frac{\gamma\sigma_{\star,\mathrm{AS}}^{2}}{\gamma\mu_{\mathrm{AS}}^{2}+2\mu_{
\mathrm{AS}}}\leq\frac{\varepsilon}{2}.$
Putting these observations together, we conclude that with the stepsize $\gamma=\nicefrac{{\varepsilon\mu_{\mathrm{AS}}}}{{\sigma_{\star,\mathrm{AS}}^{
2}}}$, we get
$\mathrm{E}\left[\left\|x_{t}-x_{\star}\right\|^{2}\right]\leq\varepsilon$
provided that

$$ $\displaystyle t\geq\frac{1+\gamma\mu_{\mathrm{AS}}}{2\gamma\mu_{\mathrm{AS}}} \log\frac{2\left\|x_{0}-x_{\star}\right\|^{2}}{\varepsilon}=\left(\frac{\sigma _{\star,\mathrm{AS}}^{2}}{2\varepsilon\mu_{\mathrm{AS}}^{2}}+\frac{1}{2}\right )\log\left(\frac{2\left\|x_{0}-x_{\star}\right\|^{2}}{\varepsilon}\right).$ $$

### F.6 σ ⋆ , NICE 2 ⁢ ( τ ) superscript subscript 𝜎 ⋆ NICE 2 𝜏 \sigma_{\star,\mathrm{NICE}}^{2}(\tau) italic_σ start_POSTSUBSCRIPT ⋆ , roman_NICE end_POSTSUBSCRIPT start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ( italic_τ ) and μ NICE ⁢ ( τ ) subscript 𝜇 NICE 𝜏 \mu_{\mathrm{NICE}}(\tau) italic_μ start_POSTSUBSCRIPT roman_NICE end_POSTSUBSCRIPT ( italic_τ ) are Monotonous Functions of τ 𝜏 \tau italic_τ

###### Lemma 13 .

For all $0\leq\tau\leq n-1$:

- 1.
$\mu_{\mathrm{NICE}}(\tau+1)\geq\mu_{\mathrm{NICE}}(\tau)$,
- 2.
$\sigma_{\star,\mathrm{NICE}}^{2}(\tau)=\frac{\frac{n}{\tau}-1}{n-1}\sigma_{
\star,\mathrm{NICE}}^{2}(1)\leq\frac{1}{\tau}\sigma_{\star,\mathrm{NICE}}^{2}(1)$.

###### Proof.

- 1.
Pick any $1\leq\tau<n$, and consider a set $C$ for which the minimum is attained in
$\mu_{\mathrm{NICE}}(\tau+1)=\min_{C\subseteq[n],|C|=\tau+1}\frac{1}{\tau+1}
\sum_{i\in C}\mu_{i}.$
Let $j=\arg\max_{i\in C}\mu_{i}$. That is, $\mu_{j}\geq\mu_{i}$ for all $i\in C$. Let $C_{j}$ be the set obtained from $C$ by removing the element $j$. Then $\left|C_{j}\right|=\tau$ and
$\mu_{j}=\max_{i\in C}\mu_{i}\geq\max_{i\in C_{j}}\mu_{i}\geq\frac{1}{\tau}\sum
_{i\in C_{j}}\mu_{i}.$
By adding $\sum_{i\in C_{j}}\mu_{i}$ to the above inequality, we obtain
$\mu_{j}+\sum_{i\in C_{j}}\mu_{i}\geq\frac{1}{\tau}\sum_{i\in C_{j}}\mu_{i}+
\sum_{i\in C_{j}}\mu_{i}.$
Observe that the left-hand side is equal to $\sum_{i\in C}\mu_{i}$, and the right-hand side is equal to $\frac{\tau+1}{\tau}\sum_{i\in C_{j}}\mu_{i}$. If we divide both sides by $\tau+1$, we obtain
$\frac{1}{\tau+1}\sum_{i\in C}\mu_{i}\geq\frac{1}{\tau}\sum_{i\in C_{j}}\mu_{i}.$
Since the left-hand side is equal to $\mu_{\mathrm{NICE}}(\tau+1)$, and the right hand side is an upper bound on $\mu_{\mathrm{NICE}}(\tau)$, we conclude that $\mu_{\mathrm{NICE}}(\tau+1)\geq\mu_{\mathrm{NICE}}(\tau)$.
- 2.
In view of (8) we have
$\displaystyle f_{C}(x)=\sum_{i\in C}\frac{1}{np_{i}}f_{i}(x).$
(12)
$\displaystyle\sigma_{\star,\mathrm{AS}}^{2}$
$\displaystyle=$
$\displaystyle\mathrm{E}_{S\sim\mathcal{S}}\left[\left\|\sum_{i\in S}\frac{1}{
np_{i}}\nabla f_{i}(x_{\star})\right\|^{2}\right]=\mathrm{E}_{S\sim\mathcal{S}
}\left[\left\|\sum_{i\in S}\frac{1}{\tau}\nabla f_{i}(x_{\star})\right\|^{2}\right]$
Let $\chi_{i}$ be the random variable defined by
$\displaystyle\chi_{j}=\left\{\begin{array}[]{ll}1&j\in S\\
0&j\notin S.\end{array}\right.$
(16)
It is easy to show that
$\displaystyle\mathbb{E}\!\left[\chi_{j}\right]=\operatorname{Prob}(j\in S)=
\frac{\tau}{n}.$
(17)
Let fix the cohort S. Let $\chi_{ij}$ be the random variable defined by
$\displaystyle\chi_{ij}=\left\{\begin{array}[]{ll}1&i\in S\text{ and }j\in S\\
0&\text{ otherwise}.\end{array}\right.$
(20)
Note that
$\displaystyle\chi_{ij}=\chi_{i}\chi_{j}.$
(21)
Further, it is easy to show that
$\displaystyle\mathbb{E}\!\left[\chi_{ij}\right]=\operatorname{Prob}(i\in S,j
\in S)=\frac{\tau(\tau-1)}{n(n-1)}.$
(22)
Denote $a_{i}:=\nabla f_{i}(x_{\star}).$
$\displaystyle\mathbb{E}\!\left[\left\|\frac{1}{\tau}\sum_{i\in S}a_{i}\right\|
^{2}\right]$
$\displaystyle=$
$\displaystyle\frac{1}{\tau^{2}}\mathbb{E}\!\left[\left\|\sum_{i\in S}a_{i}
\right\|^{2}\right]$
$\displaystyle=$
$\displaystyle\frac{1}{\tau^{2}}\mathbb{E}\!\left[\left\|\sum_{i=1}^{n}\chi_{i}
a_{i}\right\|^{2}\right]$
$\displaystyle=$
$\displaystyle\frac{1}{\tau^{2}}\mathbb{E}\!\left[\sum_{i=1}^{n}\left\|\chi_{i}
a_{i}\right\|^{2}+\sum_{i\neq j}\left\langle\chi_{i}a_{i},\chi_{j}a_{j}\right
\rangle\right]$
$\displaystyle=$
$\displaystyle\frac{1}{\tau^{2}}\mathbb{E}\!\left[\sum_{i=1}^{n}\left\|\chi_{i}
a_{i}\right\|^{2}+\sum_{i\neq j}\chi_{ij}\left\langle a_{i},a_{j}\right\rangle\right]$
$\displaystyle=$
$\displaystyle\frac{1}{\tau^{2}}\sum_{i=1}^{n}\mathbb{E}\!\left[\chi_{i}\right]
\left\|a_{i}\right\|^{2}+\sum_{i\neq j}\mathbb{E}\!\left[\chi_{ij}\right]\left
\langle a_{i},a_{j}\right\rangle$
$\displaystyle=$
$\displaystyle\frac{1}{\tau^{2}}\left(\frac{\tau}{n}\sum_{i=1}^{n}\left\|a_{i}
\right\|^{2}+\frac{\tau(\tau-1)}{n(n-1)}\sum_{i\neq j}\left\langle a_{i},a_{j}
\right\rangle\right)$
$\displaystyle=$
$\displaystyle\frac{1}{\tau n}\sum_{i=1}^{n}\left\|a_{i}\right\|^{2}+\frac{\tau
-1}{\tau n(n-1)}\sum_{i\neq j}\left\langle a_{i},a_{j}\right\rangle$
$\displaystyle=$
$\displaystyle\frac{1}{\tau n}\sum_{i=1}^{n}\left\|a_{i}\right\|^{2}+\frac{\tau
-1}{\tau n(n-1)}\left(\left\|\sum_{i=1}^{n}a_{j}\right\|^{2}-\sum_{i=1}^{n}
\left\|a_{i}\right\|^{2}\right)$
$\displaystyle=$
$\displaystyle\frac{n-\tau}{\tau(n-1)}\frac{1}{n}\sum_{i=1}^{n}\left\|a_{i}
\right\|^{2}+\frac{n(\tau-1)}{\tau(n-1)}\left\|\frac{1}{n}\sum_{i=1}^{n}a_{i}
\right\|^{2}$
$\displaystyle=$
$\displaystyle\frac{n-\tau}{\tau(n-1)}\frac{1}{n}\sum_{i=1}^{n}\left\|\nabla f_
{i}(x_{\star})\right\|^{2}+\frac{n(\tau-1)}{\tau(n-1)}\left\|\frac{1}{n}\sum_{
i=1}^{n}\nabla f_{i}(x_{\star})\right\|^{2}$
$\displaystyle=$
$\displaystyle\frac{n-\tau}{\tau(n-1)}\frac{1}{n}\sum_{i=1}^{n}\left\|\nabla f_
{i}(x_{\star})\right\|^{2}$
$\displaystyle\leq$
$\displaystyle\frac{1}{\tau}\frac{1}{n}\sum_{i=1}^{n}\left\|\nabla f_{i}(x_{
\star})\right\|^{2}$

∎

###### Lemma 13 .

###### Proof.

### F.7 Missing Proof of Lemma 5

For ease of notation, let $a_{i}=\nabla f_{i}\left(x_{\star}\right)$ and $\hat{z}_{j}=\left|C_{j}\right|a_{\xi_{j}}$, and recall that

$$ $\displaystyle\sigma_{\star,\mathrm{SS}}^{2}=\mathrm{E}_{\xi_{1},\ldots,\xi_{b} }\left[\left\|\frac{1}{n}\sum_{j=1}^{b}\hat{z}_{j}\right\|^{2}\right].$ (23) $$

where $\xi_{j}\in C_{j}$ is chosen uniformly at random. Further, for each $j\in[b]$, let $z_{j}=\sum_{i\in C_{j}}a_{i}$. Observe that $\sum_{j=1}^{b}z_{j}=\sum_{j=1}^{b}\sum_{i\in C_{j}}a_{i}=\sum_{i=1}^{n}a_{i}=
\nabla f\left(x_{\star}\right)=0$. Therefore,

$$ $\displaystyle\left\|\frac{1}{n}\sum_{j=1}^{b}\hat{z}_{j}\right\|^{2}$ $\displaystyle=\frac{1}{n^{2}}\left\|\sum_{j=1}^{b}\hat{z}_{j}-\sum_{j=1}^{b}z_ {j}\right\|^{2}$ $\displaystyle=\frac{b^{2}}{n^{2}}\left\|\frac{1}{b}\sum_{j=1}^{b}\left(\hat{z} _{j}-z_{j}\right)\right\|^{2}$ $\displaystyle\leq\frac{b^{2}}{n^{2}}\frac{1}{b}\sum_{j=1}^{b}\left\|\hat{z}_{j }-z_{j}\right\|^{2}$ $\displaystyle=\frac{b}{n^{2}}\sum_{j=1}^{b}\left\|\hat{z}_{j}-z_{j}\right\|^{2},$ (24) $$

where the inequality follows from convexity of the function $u\mapsto\|u\|^{2}$. Next,

$$ $\displaystyle\left\|\hat{z}_{j}-z_{j}\right\|^{2}=\left\|\left|C_{j}\right|a_{ \xi_{j}}-\sum_{i\in C_{j}}a_{i}\right\|^{2}=\left|C_{j}\right|^{2}\left\|a_{ \xi_{j}}-\frac{1}{\left|C_{j}\right|}\sum_{i\in C_{j}}a_{i}\right\|^{2}\leq \left|C_{j}\right|^{2}\sigma_{j}^{2}.$ (25) $$

By combining Equation 23, Section F.7 and Equation 25, we get

$$ $\displaystyle\sigma_{\star,\mathrm{SS}}^{2}\stackrel{{\scriptstyle Eqn.~{}( \ref{eqn:412})}}{{=}}\mathrm{E}_{\xi_{1},\ldots,\xi_{b}}\left[\left\|\frac{1}{ n}\sum_{j=1}^{b}\hat{z}_{j}\right\|^{2}\right]$ $\displaystyle\stackrel{{\scriptstyle Eqn.~{}(\ref{eqn:413})}}{{\leq}}\quad \mathrm{E}_{\xi_{1},\ldots,\xi_{b}}\left[\frac{b}{n^{2}}\sum_{j=1}^{b}\left\| \hat{z}_{j}-z_{j}\right\|^{2}\right]$ $\displaystyle\stackrel{{\scriptstyle Eqn.~{}(\ref{eqn:414})}}{{\leq}}\quad \mathrm{E}_{\xi_{1},\ldots,\xi_{b}}\left[\frac{b}{n^{2}}\sum_{j=1}^{b}\left|C_ {j}\right|^{2}\sigma_{j}^{2}\right]$ $\displaystyle=\frac{b}{n^{2}}\sum_{j=1}^{b}\left|C_{j}\right|^{2}\sigma_{j}^{2}.$ $$

The last expression can be further bounded as follows:

$$ $\frac{b}{n^{2}}\sum_{j=1}^{b}\left|C_{j}\right|^{2}\sigma_{j}^{2}\leq\frac{b}{ n^{2}}\left(\sum_{j=1}^{b}\left|C_{j}\right|^{2}\right)\max_{j}\sigma_{j}^{2} \leq\frac{b}{n^{2}}\left(\sum_{j=1}^{b}\left|C_{j}\right|\right)^{2}\max_{j} \sigma_{j}^{2}=b\max_{j}\sigma_{j}^{2},$ $$

where the second inequality follows from the relation $\|u\|_{2}\leq\|u\|_{1}$ between the $L_{2}$ and $L_{1}$ norms, and the last identity follows from the fact that $\sum_{j=1}^{b}\left|C_{j}\right|=n$.

### F.8 Stratified Sampling Against Block Sampling and Nice Sampling

In this section, we present a theoretical comparison of block sampling and its counterparts, providing a theoretical justification for selecting block sampling as the default clustering method in future experiments. Additionally, we compare various sampling methods, all with the same sampling size, $b$: $b$-nice sampling, block sampling with $b$ clusters, and block sampling, where all clusters are of uniform size $b$.

###### Assumption 14 .

For simplicity of comparison, we assume $b$ clusters, each of the same size, $b$:

$$ $\left|C_{1}\right|=\left|C_{2}\right|=\ldots=\left|C_{b}\right|=b.$ $$

It is crucial to acknowledge that, without specific assumptions, the comparison of different sampling methods may not provide meaningful insights. For instance, the scenario described in Lemma 5, characterized by complete inter-cluster homogeneity, demonstrates that block sampling achieves a variance term, denoted as $\sigma_{\star,\mathrm{SS}}^{2}$, which is lower than the variance terms associated with both block sampling and nice sampling. However, a subsequent example illustrates examples in which the variance term for block sampling surpasses those of block sampling and nice sampling.

###### Example 1 .

Without imposing any additional clustering assumptions, there exist examples for any arbitrary $n$, such that $\sigma^{2}_{\star,\mathrm{SS}}\geq\sigma^{2}_{\star,\mathrm{BS}}$ and $\sigma^{2}_{\star,\mathrm{SS}}\geq\sigma^{2}_{\star,\mathrm{NICE}}$.

###### Proof.

Counterexample when SS is worse in neighborhood than BS
Assume we have such clustering and $\nabla f_{i}(x_{\star})$ such that the centroids of each cluster are equal to zero: $\forall i\in[b]$, $\frac{1}{|C_{i}|}\sum_{j\in C_{i}}\nabla f_{j}(x_{\star})=0$. For instance, this can be achieved in the following case: The dimension is $d=2$, all clusters are of equal size $m$, then assign $\forall i\in[b]$, $\forall j\in C_{i}$, $\nabla f_{j}(x_{\star})=\left(Re\left(\omega^{mj+i}\right),Im\left(\omega^{mj+
i}\right)\right)$ where $\omega=\sqrt[n]{1}\in\mathbb{C}$. Let us calculate $\sigma_{\star,\mathrm{BS}}^{2}$:

$$ $\displaystyle\sigma_{\star,\mathrm{BS}}^{2}:=\sum_{j=1}^{b}q_{j}\left\|\sum_{i \in C_{j}}\frac{1}{np_{i}}\nabla f_{i}(x_{\star})\right\|^{2}=$ $\displaystyle=\frac{1}{n^{2}}\sum_{j=1}^{b}\frac{|C_{j}|^{2}}{q_{j}}\left\| \frac{1}{|C_{j}|}\sum_{i\in C_{j}}\nabla f_{i}(x_{\star})\right\|^{2}=0.$ $$

As a result:

$$ $\sigma_{\star,\mathrm{BS}}^{2}=0\leq\sigma_{\star,\mathrm{SS}}^{2}.$ $$

Counterexample when SS is worse in neighborhood than NICE
Here, we employ a similar proof technique as in the proof of Lemma 6.
Let us choose such clustering $\mathcal{C}_{b,\mathrm{SS},\max}=\operatorname*{arg\,max}_{\mathcal{C}_{b}}
\sigma^{2}_{\star,\mathrm{SS}}(\mathcal{C}_{b})$. Denote $\mathbf{i}_{b}\coloneqq(i_{1},\cdots,i_{b})$, $\mathbf{C}_{b}\coloneqq C_{1}\times\cdots\times C_{b}$, and $S_{\mathbf{i}_{b}}\coloneqq\left\|\frac{1}{\tau}\sum_{i\in\mathbf{i}_{b}}
\nabla f_{i}(x_{\star})\right\|$.

$$ $\displaystyle\sigma_{\star,\mathrm{NICE}}^{2}=\frac{1}{C(n,\tau)}\sum_{C \subseteq[n],|C|=\tau}\left\|\frac{1}{\tau}\sum_{i\in C}\nabla f_{i}(x_{\star} )\right\|^{2}$ $\displaystyle=\frac{1}{C(n,b)}\sum_{\mathbf{i}_{b}\subseteq[n]}S_{\mathbf{i}_{ b}}$ $\displaystyle\stackrel{{\scriptstyle 1}}{{=}}\frac{1}{\#_{\text{ clusterizations}}}\sum_{\mathcal{C}_{b}}\frac{1}{b^{b}}\sum_{\mathbf{i}_{b}\in \mathbf{C}_{b}}S_{\mathbf{i}_{b}}$ $\displaystyle=\frac{1}{\#_{\text{clusterizations}}}\sum_{\mathcal{C}_{b}} \sigma^{2}_{\star,\mathrm{SS}}(\mathcal{C}_{b})$ $\displaystyle\stackrel{{\scriptstyle 2}}{{\leq}}\sigma^{2}_{\star,\mathrm{SS}} (\mathcal{C}_{b,\mathrm{SS},\max}).$ $$

Equation 1 holds because, in every clusterization $\mathcal{C}_{b}$, there are $\frac{1}{b^{b}}$ possible sample combinations $\mathbf{i}_{b}$. Due to symmetry, one can conclude that each combination $S_{\mathbf{i}_{b}}$ is counted the same number of times. Equation 2 follows from the definition of $\mathcal{C}_{b,\mathrm{SS},\max}$.
For illustrative purposes, we can demonstrate this effect with a specific example. Let $n=4$ and define $\forall i~{}a_{i}=\nabla f_{i}(x^{*})\in\mathbb{R}^{2}$. Let $a_{1}=(0,1)^{T}$, $a_{2}=(1,0)^{T}$, $a_{3}=(0,-1)^{T}$, and $a_{4}=(-1,0)^{T}$. Then fix clustering $\mathcal{C}_{b}=\left\{C_{1}=\{a_{1},a_{3}\},C_{2}=\{a_{2},a_{4}\}\right\}$. Then:

$$ $\displaystyle\sigma_{\star,\mathrm{SS}}^{2}=\frac{1}{4}\sum_{\mathbf{i}_{b}\in \mathcal{C}_{b}}\left\|\frac{a_{i_{1}}+a_{i_{2}}}{2}\right\|^{2}$ $\displaystyle=\frac{1}{4}\sum_{\mathbf{i}_{b}\in\mathcal{C}_{b}}\left\|(\pm \frac{1}{2},\pm\frac{1}{2})\right\|^{2}$ $\displaystyle=\frac{1}{2}.$ $$

$$ $\displaystyle\sigma_{\star,\mathrm{NICE}}^{2}=\frac{1}{C(4,2)}\sum_{i<j}\left \|\frac{a_{i}+a_{j}}{2}\right\|^{2}$ $\displaystyle=\frac{1}{6}\sum_{i<j}\left\|\frac{a_{i}+a_{j}}{2}\right\|^{2}$ $\displaystyle=\frac{1}{6}\left(\left[\left\|\frac{a_{1}+a_{3}}{2}\right\|^{2}+ \left\|\frac{a_{2}+a_{4}}{2}\right\|^{2}\right]+2\times\left\|\frac{a_{i_{1}}+ a_{i_{2}}}{2}\right\|^{2}\right)$ $\displaystyle=\frac{1}{6}\left(0+2\times 2\times\frac{1}{2}\right)$ $\displaystyle=\frac{1}{3}$ $\displaystyle=\frac{2}{3}\times\sigma_{\star,\mathrm{SS}}^{2}$ $\displaystyle\leq\sigma_{\star,\mathrm{SS}}^{2}$ $$

∎

To select the optimal clustering, we will choose the clustering that minimizes $\sigma^{2}_{\star,\mathrm{SS}}$.

###### Definition 15 (Stratified sampling optimal clustering) .

Denote the clustering of workers into blocks as $\mathcal{C}_{b}:=\{C_{1},C_{2},\ldots,C_{b}\}$, such that the disjoint union of all clusters $C_{1}\cup C_{2}\cup\ldots\cup C_{b}=[n]$. Define *block sampling Optimal Clustering* as the clustering configuration that minimizes $\sigma_{\star,\mathrm{SS}}^{2}$, formally given by:

$$ $\mathcal{C}_{b,\mathrm{SS}}:=\operatorname*{arg\,min}_{\mathcal{C}_{b}}\sigma^ {2}_{\star,\mathrm{SS}}(\mathcal{C}_{b}).$ $$

###### Lemma 16 .

Given 14, the following holds: $\sigma_{\star,\mathrm{SS}}^{2}\left(\mathcal{C}_{b,\mathrm{SS}}\right)\leq
\sigma_{\star,\mathrm{NICE}}^{2}$  for arbitrary $b$.

###### Proof.

Denote $\mathbf{i}_{b}\coloneqq(i_{1},\cdots,i_{b})$, $\mathbf{C}_{b}\coloneqq C_{1}\times\cdots\times C_{b}$, and $S_{\mathbf{i}_{b}}\coloneqq\left\|\frac{1}{\tau}\sum_{i\in\mathbf{i}_{b}}
\nabla f_{i}(x_{\star})\right\|$.

$$ $\displaystyle\sigma_{\star,\mathrm{NICE}}^{2}=\frac{1}{C(n,\tau)}\sum_{C \subseteq[n],|C|=\tau}\left\|\frac{1}{\tau}\sum_{i\in C}\nabla f_{i}(x_{\star} )\right\|^{2}$ $\displaystyle=\frac{1}{C(n,b)}\sum_{\mathbf{i}_{b}\subseteq[n]}S_{\mathbf{i}_{ b}}$ $\displaystyle\stackrel{{\scriptstyle 1}}{{=}}\frac{1}{\#_{\text{ clusterizations}}}\sum_{\mathcal{C}_{b}}\frac{1}{b^{b}}\sum_{\mathbf{i}_{b}\in \mathbf{C}_{b}}S_{\mathbf{i}_{b}}$ $\displaystyle=\frac{1}{\#_{\text{clusterizations}}}\sum_{\mathcal{C}_{b}} \sigma^{2}_{\star,\mathrm{SS}}(\mathcal{C}_{b})$ $\displaystyle\stackrel{{\scriptstyle 2}}{{\geq}}\sigma^{2}_{\star,\mathrm{SS}} (\mathcal{C}_{b,\mathrm{SS},\min})$ $$

Equation 1 holds because, in every clusterization $\mathcal{C}_{b}$, there are $\frac{1}{b^{b}}$ possible sample combinations $\mathbf{i}_{b}$. Due to symmetry, one can conclude that each combination $S_{\mathbf{i}_{b}}$ is counted the same number of times. Equation 2 follows from the definition of $\mathcal{C}_{b,\mathrm{SS},\min}$ as the clustering that minimizes $\sigma^{2}_{\star,\mathrm{SS}}$, according to Definition 15.
∎

###### Example 2 .

Consider the number of clusters and the size of each cluster, with $b=2$, under 14. Then, $\sigma_{\star,\mathrm{SS}}^{2}\left(\mathcal{C}_{b,\mathrm{SS}}\right)\leq
\sigma^{2}_{\star,\mathrm{BS}}$.

###### Proof.

Let $n=4$, $b=2$. Denote $\forall i\enspace a_{i}=\nabla f_{i}(x_{*})$. Define $S^{2}:=\sum_{i<j}\left\|\frac{a_{i}+a_{j}}{2}\right\|^{2}$.

$$ $\displaystyle\sigma^{2}_{\star,\text{SS}}=\frac{1}{4}\left(S^{2}-\left\|\frac{ a_{C_{1}^{1}}+a_{C_{1}^{2}}}{2}\right\|^{2}-\left\|\frac{a_{C_{2}^{1}}+a_{C_{2 }^{2}}}{2}\right\|^{2}\right)$ $\displaystyle=\frac{1}{4}\left(S^{2}-2\sigma^{2}_{\star,\text{BS}}\right)$ $$

$\mathcal{C}_{b,\mathrm{SS}}$ clustering minimizes $\sigma^{2}_{\star,\text{SS}}$, thereby maximizing $\sigma^{2}_{\star,\text{BS}}$. Thus,

$$ $\displaystyle\sigma^{2}_{\star,\text{SS}}=\frac{1}{4}\left(\left[\left\|\frac{ a_{C_{1}^{1}}+a_{C_{2}^{1}}}{2}\right\|^{2}+\left\|\frac{a_{C_{1}^{2}}+a_{C_{2 }^{2}}}{2}\right\|^{2}\right]+\left[\left\|\frac{a_{C_{1}^{1}}+a_{C_{2}^{2}}}{ 2}\right\|^{2}+\left\|\frac{a_{C_{1}^{2}}+a_{C_{2}^{1}}}{2}\right\|^{2}\right]\right)$ $\displaystyle=\frac{1}{4}\left(2\sigma^{2}_{\star,\text{BS}}\left((C_{1}^{1},C _{2}^{1}),(C_{1}^{2},C_{2}^{2})\right)+2\sigma^{2}_{\star,\text{BS}}\left((C_{ 1}^{1},C_{2}^{2}),(C_{1}^{2},C_{2}^{1})\right)\right)$ $\displaystyle=\frac{1}{2}\left(\sigma^{2}_{\star,\text{BS}}\left((C_{1}^{1},C_ {2}^{1}),(C_{1}^{2},C_{2}^{2})\right)+\sigma^{2}_{\star,\text{BS}}\left((C_{1} ^{1},C_{2}^{2}),(C_{1}^{2},C_{2}^{1})\right)\right)$ $\displaystyle\leq\sigma^{2}_{\star,\text{BS}}.$ $$

∎

However, it is possible that this relationship might hold more generally. Empirical experiments for different configurations, such as $b=3$, support this possibility. For example, with $n=9$, $b=3$, and $d=10$, Python simulations where gradients $\nabla f_{i}$ are sampled from $\mathcal{N}(0,1)$ and $\mathcal{N}(e,1)$ across $1000$ independent trials, show that $\sigma^{2}_{\star,\text{SS}}\leq\sigma^{2}_{\star,\text{BS}}$.
Question of finding theoretical proof for arbitraty $n$ remains open and has yet to be addressed in the existing literature.

###### Assumption 14 .

###### Example 1 .

###### Proof.

###### Definition 15 (Stratified sampling optimal clustering) .

###### Lemma 16 .

###### Proof.

###### Example 2 .

###### Proof.

### F.9 Different Approaches of Federated Averaging

Proof of Theorem LABEL:thm:FedProx-SPPM-AS:

###### Proof.

$$ $\displaystyle{\left\lVert x_{t}-x_{\star}\right\rVert}^{2}$ $\displaystyle={\left\lVert\sum_{i\in S_{t}}\frac{1}{\left|S_{t}\right|}\mathop {\mathrm{prox}}\nolimits_{\gamma f_{i}}(x_{t-1})-\frac{1}{\left|S_{t}\right|} \sum_{i\in S_{t}}x_{\star}\right\rVert}^{2}$ $\displaystyle\stackrel{{\scriptstyle(\lx@cref{creftype~refnum}{fact:fact1})}}{ {=}}{\left\lVert\sum_{i\in S_{t}}\frac{1}{\left|S_{t}\right|}\left[\mathop{ \mathrm{prox}}\nolimits_{\gamma f_{i}}(x_{t-1})-\mathop{\mathrm{prox}} \nolimits_{\gamma f_{i}}(x_{\star}+\gamma\nabla f_{i}(x_{\star}))\right]\right \rVert}^{2}$ $\displaystyle\stackrel{{\scriptstyle Jensen}}{{\leq}}\sum_{i\in S_{t}}\frac{1} {\left|S_{t}\right|}{\left\lVert\left[\mathop{\mathrm{prox}}\nolimits_{\gamma f _{i}}(x_{t-1})-\mathop{\mathrm{prox}}\nolimits_{\gamma f_{i}}(x_{\star}+\gamma \nabla f_{i}(x_{\star}))\right]\right\rVert}^{2}$ $\displaystyle\stackrel{{\scriptstyle(\lx@cref{creftype~refnum}{fact:fact2})}}{ {\leq}}\sum_{i\in S_{t}}\frac{1}{\left|S_{t}\right|}\frac{1}{(1+\gamma\mu_{i}) ^{2}}{\left\lVert x_{t-1}-(x_{\star}+\gamma\nabla f_{i}(x_{\star}))\right \rVert}^{2}$ $$

$$ $\displaystyle\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[{\left\lVert x_{t}-x_{ \star}\right\rVert}^{2}|x_{t-1}\right]$ $\displaystyle\leq\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\sum_{i\in S_{t}}\frac {1}{\left|S_{t}\right|}\frac{1}{(1+\gamma\mu_{i})^{2}}{\left\lVert\left(x_{t-1 }-x_{\star}\right)-\gamma\nabla f_{i}(x_{\star}))\right\rVert}^{2}|x_{t-1}\right]$ $\displaystyle\stackrel{{\scriptstyle\text{Young, }\alpha_{i}>0}}{{\leq}} \mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\sum_{i\in S_{t}}\frac{1}{\left|S_{t} \right|}\frac{1}{(1+\gamma\mu_{i})^{2}}\left(\left(1+\alpha_{i}\right){\left \lVert x_{t-1}-x_{\star}\right\rVert}^{2}+\left(1+\alpha_{i}^{-1}\right){\left \lVert\gamma\nabla f_{i}(x_{\star}))\right\rVert}^{2}\right)|x_{t-1}\right]$ $\displaystyle\stackrel{{\scriptstyle\alpha_{i}=\gamma\mu_{i}}}{{=}}\mathbb{E}_ {S_{t}\sim\mathcal{S}}\left[\sum_{i\in S_{t}}\frac{1}{\left|S_{t}\right|}\frac {1}{(1+\gamma\mu_{i})^{2}}\left(\left(1+\gamma\mu_{i}\right){\left\lVert x_{t- 1}-x_{\star}\right\rVert}^{2}+\left(1+\frac{1}{\gamma\mu_{i}}\right){\left \lVert\gamma\nabla f_{i}(x_{\star}))\right\rVert}^{2}\right)|x_{t-1}\right]$ $\displaystyle=\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\sum_{i\in S_{t}}\frac{1} {\left|S_{t}\right|}\left(\frac{1}{1+\gamma\mu_{i}}{\left\lVert x_{t-1}-x_{ \star}\right\rVert}^{2}+\frac{\gamma}{(1+\gamma\mu_{i})\mu_{i}}{\left\lVert \nabla f_{i}(x_{\star}))\right\rVert}^{2}\right)|x_{t-1}\right]$ $\displaystyle=\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\frac{1}{\left|S_{t} \right|}\sum_{i\in S_{t}}\frac{1}{1+\gamma\mu_{i}}|x_{t-1}\right]{\left\lVert x _{t-1}-x_{\star}\right\rVert}^{2}+\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\frac {1}{\left|S_{t}\right|}\sum_{i\in S_{t}}\frac{\gamma}{(1+\gamma\mu_{i})\mu_{i} }{\left\lVert\nabla f_{i}(x_{\star}))\right\rVert}^{2}|x_{t-1}\right]$ $$

By applying tower property one can get the following:

$$ $\displaystyle\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[{\left\lVert x_{t}-x_{ \star}\right\rVert}^{2}\right]$ $\displaystyle=\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\frac{1}{\left|S_{t} \right|}\sum_{i\in S_{t}}\frac{1}{1+\gamma\mu_{i}}\right]{\left\lVert x_{t-1}- x_{\star}\right\rVert}^{2}+\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\frac{1}{ \left|S_{t}\right|}\sum_{i\in S_{t}}\frac{\gamma}{(1+\gamma\mu_{i})\mu_{i}}{ \left\lVert\nabla f_{i}(x_{\star}))\right\rVert}^{2}\right]$ $\displaystyle=A_{\mathcal{S}}{\left\lVert x_{t-1}-x_{\star}\right\rVert}^{2}+B _{\mathcal{S}}.$ $$

where $A_{\mathcal{S}}\coloneqq\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\frac{1}{\left|
S_{t}\right|}\sum_{i\in S_{t}}\frac{1}{1+\gamma\mu_{i}}\right]$ and $B_{\mathcal{S}}\coloneqq\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[\frac{1}{\left|
S_{t}\right|}\sum_{i\in S_{t}}\frac{\gamma}{(1+\gamma\mu_{i})\mu_{i}}{\left
\lVert\nabla f_{i}(x_{\star}))\right\rVert}^{2}\right]$.
By directly applying 5:

$$ $\mathbb{E}_{S_{t}\sim\mathcal{S}}\left[{\left\lVert x_{t}-x_{\star}\right \rVert}^{2}\right]\leq A_{\mathcal{S}}^{t}{\left\lVert x_{0}-x_{\star}\right \rVert}^{2}+\frac{B_{\mathcal{S}}}{1-A_{\mathcal{S}}}.$ $$

∎

###### Lemma 17 (Inexact formulation of SPPM-AS ) .

Let $b>0\in\mathbb{R}$ and define $\widetilde{\mathop{\mathrm{prox}}\nolimits}_{\gamma f}(x)$ such that $\forall x{\left\lVert\widetilde{\mathop{\mathrm{prox}}\nolimits}_{\gamma f}(x)
-\mathop{\mathrm{prox}}\nolimits_{\gamma f}(x)\right\rVert}^{2}\leq b$. Let 1 and 2 hold. Let $x_{0}\in\mathbb{R}^{d}$ be an arbitrary starting point. Then for any $t\geq 0$ and any $\gamma>0$, $s>0$, the iterates of SPPM-AS satisfy

$$ $\mathbb{E}\left[{\left\lVert x_{t}-x_{\star}\right\rVert}^{2}\right]\leq\left( \frac{1+s}{(1+\gamma\mu)^{2}}\right)^{t}{\left\lVert x_{0}-x_{\star}\right \rVert}^{2}+\frac{\left(1+s\right)\left(\gamma^{2}\sigma^{2}_{\star}+s^{-1}b(1 +\gamma\mu)^{2}\right)}{\gamma^{2}\mu^{2}+2\gamma\mu-s}.$ $$

###### Proof of Lemma 17 .

We provide more general version of SPPM proof

$$ $\displaystyle{\left\lVert x_{t+1}-x_{\star}\right\rVert}^{2}$ $\displaystyle={\left\lVert\widetilde{\mathop{\mathrm{prox}}\nolimits}_{\gamma f _{\xi_{t}}(x_{t})}-\mathop{\mathrm{prox}}\nolimits_{\gamma f_{\xi_{t}}}(x_{t}) +\mathop{\mathrm{prox}}\nolimits_{\gamma f_{\xi_{t}}}(x_{t})-x_{\star}\right \rVert}^{2}$ $\displaystyle\stackrel{{\scriptstyle Young,s>0}}{{\leq}}(1+s^{-1}){\left\lVert \widetilde{\mathop{\mathrm{prox}}\nolimits}_{\gamma f_{\xi_{t}}}(x_{t})- \mathop{\mathrm{prox}}\nolimits_{\gamma f_{\xi_{t}}}\right\rVert}^{2}(x_{t})+( 1+s){\left\lVert\mathop{\mathrm{prox}}\nolimits_{\gamma f_{\xi_{t}}}(x_{t})-x_ {\star}\right\rVert}^{2}$ $\displaystyle\leq(1+s^{-1})b+(1+s){\left\lVert\mathop{\mathrm{prox}}\nolimits_ {\gamma f_{\xi_{t}}}(x_{t})-x_{\star}\right\rVert}^{2}.$ $$

Then proof follows same path as proof Theorem 4 and we get

$$ $\displaystyle\mathbb{E}\left[{\left\lVert x_{t+1}-x_{\star}\right\rVert}^{2} \right]\leq(1+s^{-1})b+(1+s)\frac{1}{(1+\gamma\mu)^{2}}\left({\left\lVert x_{t }-x_{\star}\right\rVert}^{2}+\gamma^{2}\sigma^{2}_{\star}\right)$ $\displaystyle=\frac{1+s}{(1+\gamma\mu)^{2}}\left({\left\lVert x_{t}-x_{\star} \right\rVert}^{2}+\left[\gamma^{2}\sigma^{2}_{\star}+s^{-1}b(1+\gamma\mu)^{2} \right]\right).$ $$

azc It only remains to solve the above recursion. Luckily, that is exactly what 5 does. In particular, we use it with $s_{t}=\mathbb{E}\left[{\left\lVert x_{t}-x_{\star}\right\rVert}^{2}\right],A=
\frac{1+s}{(1+\gamma\mu)^{2}}$ and $B=\frac{\left(1+s\right)\left(\gamma^{2}\sigma^{2}_{\star}+s^{-1}b(1+\gamma\mu
)^{2}\right)}{(1+\gamma\mu)^{2}}$ to get

$$ $\displaystyle\mathbb{E}\left[{\left\lVert x_{t}-x_{\star}\right\rVert}^{2}\right]$ $\displaystyle\leq A^{t}{\left\lVert x_{0}-x_{\star}\right\rVert}^{2}+B\frac{1} {1-A}$ $\displaystyle\leq A^{t}{\left\lVert x_{0}-x_{\star}\right\rVert}^{2}+B\frac{(1 +\gamma\mu)^{2}}{(1+\gamma\mu)^{2}-1-s}$ $\displaystyle\leq A^{t}{\left\lVert x_{0}-x_{\star}\right\rVert}^{2}+\frac{ \left(1+s\right)\left(\gamma^{2}\sigma^{2}_{\star}+s^{-1}b(1+\gamma\mu)^{2} \right)}{(1+\gamma\mu)^{2}-1-s}$ $\displaystyle=\left(\frac{1+s}{(1+\gamma\mu)^{2}}\right)^{t}{\left\lVert x_{0} -x_{\star}\right\rVert}^{2}+\frac{\left(1+s\right)\left(\gamma^{2}\sigma^{2}_{ \star}+s^{-1}b(1+\gamma\mu)^{2}\right)}{\gamma^{2}\mu^{2}+2\gamma\mu-s}.$ $$

∎

###### Proof.

###### Lemma 17 (Inexact formulation of SPPM-AS ) .

###### Proof of Lemma 17 .

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Could theoretical convergence performance3 analysis of the proposed method be generalized to non-convex settings?

2. If yes to the above question, what are the major challenges in the theoretical analysis and how to overcome them?
soundness: 2
strengths: 1. The authors proposed a new variant of federated learning algorithms to further improve the communication efficiency in federated learning.

2. The experiments in this paper is comprehensive.
summary: This paper studied the problem of whether one can change the conventional operation in FL, where a cohort of client devices can be involved in multiple rounds of communication with the server. The authors proposed a variant of the stochastic proximal point method (SPPM-AS), which supports a large collection of client sampling procedures to lead to further gains compared to classical client selection approaches. The authors further conducted experiments to verify the performance of the proposed SPPM-AS algorithm.
weaknesses: 1. The federated learning setting considered in this paper are limited. Specifically, this paper assume that the objective function is strongly convex (Assumption 2.2). This not only significantly simplifies the theoretical analysis to achieve stronger convergence results, but also is not of practical interests since most ML models are non-convex.

2. The algorithm design is only a minor variation of the existing algorithmic framework in federated learning, which has been well-explored. Algorithmic ideas such as proximal point optimization, various types of client sampling, and multiple local updates for each cohort are not new and have been considered in the literature. It's unclear what are the major novelty in this paper.

3. As a consequence of the above factors, the theoretical contributions of this paper are marginal, since most of the proof details in Appendix F are quite standard.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Do any of the proposed algorithm variations achieve any theoretical speedup compared to Local GD with i.i.d. client sampling, beyond an improvement in constant factors?
2. Is there any way to execute optimal stratified sampling in practice?
3. How does your algorithm compare experimentally against baselines that use client selection strategies besides NICE sampling, e.g. Power-of-Choice [6]?
4. In the neural network experiments (Figure 4), how does SPPM-NICE compare against LocalGD when SPPM-NICE uses GD as a local prox solver instead of Adam?

[6] Jee Cho, Y., Wang, J. &amp; Joshi, G.. (2022).  Towards Understanding Biased Client Selection in Federated Learning . Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research.
soundness: 2
strengths: 1. The sentence-by-sentence writing is clear.
2. All of the proofs seem to be correct.
summary: This paper presents SPPM-AS, a variant of the stochastic proximal point method that supports various protocols for sampling data. For federated learning, this translates to a federated optimization algorithm that supports various protocols for sampling clients. The method is proven to converge to an $\epsilon$-approximate solution for strongly convex problems, and experiments show improvements compared to classical baselines.
weaknesses: 1. I don't see any significant theoretical improvement of the proposed algorithms.

    1a. The iteration complexity is $1/\epsilon$ (Line 212), which does not seem to improve upon FedAvg expect possibly in terms of constant factors. This is not too surprising considering that FedProx does not improve upon FedAvg, but there are already works in FL showing that the order of client sampling can improve convergence rate in terms of the $\epsilon$ dependence [1]. Therefore the iteration complexity shown in this paper is not a significant improvement.

    1b. The algorithm which enjoys theoretical guarantees cannot actually be implemented, because the choice of hyperparameters requires knowledge of the global minimum. The iteration complexity (Line 212) requires to choose the stepsize based on $\sigma_{\*,\text{AS}}^2$, which can only be computed with knowledge of $x_*$.

    1c. The only possibility for theoretical improvement is an improvement of constant factors from optimal stratified sampling (Lemma 2), but *optimal stratified sampling cannot be computed without knowledge of the global minimum*.

2. I don't see any significant experimental improvement due to issues with the experimental methodology.

    2a. The experimental evaluation only compares against naive baselines of Local GD and Minibatch GD. There are a huge number of works in FL that try to improve optimization with different client selection strategies, and these works are essentially ignored by this experimental evaluation (see [2] and its references). The authors' claim of $74$% improvement compares against the naive baseline, not against state-of-the-art (or any of the relevant existing work).

    2b. SPPM-SS cannot be implemented with real data. As I pointed out in Weakness #1c, optimal stratified sampling cannot be computed without knowledge of the global minimum. To run experiments, the authors instead use a clustering heuristic that stratifies clients according to features, clustered using K-means. However, it is unclear whether such a clustering procedure can be executed in a real federated learning scenario when client data must remain on-device. Without this, a significant portion of the experimental results (Figures 1, 2, part of 3) only describe an algorithm which cannot be implemented in practice.

    2c. The neural network experiments (Figure 4) may not be a fair comparison between LocalGD and SPPM-NICE. SPPM-NICE uses Adam as a local prox solver, which may not be a reasonable comparison against LocalGD, since LocalGD does not include any update preconditioning (known to be important for neural network training). It would be more appropriate to compare SPPM-NICE against LocalGD when SPPM-NICE uses GD as a local prox solver. An alternative is to compare SPPM-NICE w/ Adam against a local version of Adam, for example FedAdam. Appendix E.4 contains NN experiments with different local solvers (Figure 16), but I don't see exactly how these results related to those in Figure 4. It looks like the choice of local solver can create a gap of about 6\% in train accuracy, and this is described as "all methods perform similarly" (Line 1526), whereas a similar gap between LocalGD vs. SPPM-NICE in Figure 4 is described as "enhanced performance" (Line 513).

3. The paper exaggerates its own contribution and ignores relevant previous work. There are a huge number of works that improve federated optimization with different client selection strategies, which are ignored by this paper in terms of theory, experiments, and general framing (see [2] and its references). Some examples of exaggerated language that I find inappropriate:
- Abstract: "Virtually all FL methods operate in the following manner..." This claim is not accurate; there are many works in FL that use peer-to-peer communication [3], asynchronous communication [4], etc. Further, I fail to see how the proposed algorithms of this paper do not also fall into the category described in the abstract.
- Line 524: "This foundational work showcases a pivotal shift in federated learning strategies". I don't believe that this work departs very far at all from previous work in FL (e.g. [5] and related works). In my opinion, this kind of self-aggrandizing is not appropriate for a scientific publication.

4. The message of the paper is not totally coherent. The abstract talks about "cohort squeeze" and novel communication principles, but most of the paper actually deals with client selection strategies within the standard intermittent communication structure. The experiments discuss local vs. global communication (Section 3.6), which seems to be the connection to the "cohort squeeze" of the title and abstract, but this section makes up a very small part of the paper's technical content. Perhaps I have missed a connection between the content of the abstract and the content of the main text.

[1] Cho, Yae Jee, et al. "On the convergence of federated averaging with cyclic client participation." International Conference on Machine Learning. PMLR, 2023.

[2] Fu, Lei, et al. "Client selection in federated learning: Principles, challenges, and opportunities." IEEE Internet of Things Journal (2023).

[3] Beltrán, Enrique Tomás Martínez, et al. "Decentralized federated learning: Fundamentals, state of the art, frameworks, trends, and challenges." IEEE Communications Surveys & Tutorials (2023).

[4] Xu, Chenhao, et al. "Asynchronous federated learning on heterogeneous devices: A survey." Computer Science Review 50 (2023): 100595.

[5] Grudzień, Michał, Grigory Malinovsky, and Peter Richtárik. "Improving accelerated federated learning with compression and importance sampling." arXiv preprint arXiv:2306.03240 (2023).

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see my comments in the weakness part.
soundness: 2
strengths: - The paper is easy to read in general.
summary: This paper applies the stochastic proximal point method (SPPM) to federated learning. Convergence analysis of SPPM with strongly convex objectives are given, experiments showing that SPPM can reduce the total communication cost compared with FedAvg.
weaknesses: - In section 2.2, the author(s) discussed some properties of the SPPM-AS, but I cannot find the communication cost analysis of SSPM-AS, which is the most important factor of FL algorithms. Theoretically, how does the total communication cost of SPPM-AS compared with existing FL algorithms such as FedAvg, FedProx, SCAFFOLD, etc.
- Similar to the question above, how is $prox_{\gamma f_{S_t}} ( x_t )$ being solved? There must be some communication between $S_t$ during the optimization, how expensive is the communication?
- Table 1 is not very easy to read. I did not fully get the meaning between 313-323 when I read it for the first time.
- In line 340, how is $\tilde{f}_i$ defined?
- In experiments, how to solve the proximal point problem is kind of vague, what is the local communication cost and how is the local communication cost being controlled in each experiment?
- Federated leaning has been studied many years. The baseline methods in the experiments is limited (FedAvg), the author(s) should include some more recent FL algorithms.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: 1. The authors claim that increasing the number of local communication rounds can reduce the total cost. Does this claim hold for all numbers of local communication rounds, or is there a tradeoff between local communication rounds and total cost?
2. The authors state that the stratified sampling optimal clustering is impractical, so they employ a clustering heuristic which is K-means. What are the differences between these two methods? 
3. The authors indicate that stratified sampling outperforms nice sampling. Why do they provide the experiment results of CNN under nice sampling rather than stratified sampling?
soundness: 2
strengths: This paper introduces a new cross-device federated learning framework called SPPM-AS that supports arbitrary sampling strategies. The effectiveness of SPPM-AS is validated through both theoretical analysis and numerical experiments.
summary: Based on SPPM, this paper proposes SPPM-AS, a cross-device federated learning framework that supports arbitrary sampling strategies. The performance of SPPM-AS is evaluated both theoretically and numerically.
weaknesses: 1. The presentation of the paper needs to be improved. For example, it is not easy to follow the paper since a lot of important discussions and results are presented in the appendix. 
2. A detailed explanation of Algorithm 1 or SPPM should be provided to improve better reader understanding. 
3. The theoretical analysis is based on the strongly convex assumption. Extending the analysis to a more general non-convex setting would strengthen the paper.
4. The comparisons between different sampling methods are based on simplified settings, e.g., $b$ clusters of uniform size $b$, with blocking size and the number of blocks set as 2.
5. The authors only provide experiments on logistic regression using datasets from the LibSVM repository and on CNN with FEMNIST dataset, which are relatively simple. To better demonstrate the performance of SPPM-AS, experiments on more complex datasets (e.g., CIFAR-100, Shakespeare) and tasks (e.g., NLP) are recommended.


Minor:
1. Notations should be explained when they first appear in the paper, e.g., $n$.
2. In line 93, "dashed line" should be corrected to "dashed red line".
3. Abbreviations should be defined upon their first appearance in the paper, e.g., $HP$.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: See weaknesses.
soundness: 2
strengths: 1. The research provides a thorough theoretical underpinning to the SPPM-AS method, complete with convergence proofs. This not only bolsters the credibility of the approach but also offers a deeper understanding of its operational dynamics. The paper goes beyond mere theoretical exposition by delivering comprehensive interpretations of the theoretical outcomes, making the material more accessible and applicable for readers.

2. A significant aspect of the paper is its in-depth exploration of diverse sampling strategies, each accompanied by a detailed explanation and analysis. The authors present the sampling variance for each strategy and offer a comparative analysis, highlighting the nuances and implications of choosing one strategy over another. This meticulous examination of sampling strategies enriches the paper's contribution to the field of federated learning.

3. The empirical validation of the theoretical findings is a testament to the practical viability of the SPPM-AS method. Through a series of extensive experiments on both convex and non-convex models, the paper demonstrates the method's robustness and effectiveness in real-world scenarios. These experiments solidify the theoretical claims and showcase the method's potential to be integrated into existing federated learning frameworks, thereby bridging the gap between theory and practice.
summary: The paper presents an innovative method in the domain of federated learning, breaking away from the conventional approach where client cohorts interact with a central server solely once per training cycle. The authors have developed SPPM-AS (Stochastic Proximal Point Method with Arbitrary Sampling), a technique that facilitates additional communication rounds within each cohort, potentially slashing the overall communication expenditure needed for cross-device model training.

Theoretical underpinnings of SPPM-AS are thoroughly examined, with a focus on its convergence characteristics, and are juxtaposed with those of traditional methods. The study delves into the effects of various hyperparameters—including the learning rate and frequency of local communications—on algorithmic performance. Empirical evaluations conducted across both convex (logistic regression) and non-convex (neural network) models substantiate the method's proficiency in lowering communication expenses without compromising accuracy, and in some cases, even enhancing it over current methodologies.
weaknesses: The core novel contributions of this paper are still unclear to me. I appreciate the detailed explanations of the theoretical results and the examples with various concrete sampling strategies. However, the novel algorithm SPPM-AS seems to heavily rely on SPPM, which can already be applied directly to the federated learning setting (Equation 1). I want to understand the technical differences and contributions of SPPM-AS compared to the SPPM algorithm. Please provide a more explicit comparison between SPPM-AS and SPPM, highlighting the key technical differences and innovations. 

Moreover, It appears that this paper eliminates the need for the second-order similarity condition in SPPM. How eliminating the second-order similarity condition can be achieved in your proof is of great interest to me.

Last but not least, explain in more detail how the multiple communication rounds within cohorts contribute to the novelty of the approach.