SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Michael Crawshaw*  
Department of Computer Science  
George Mason University  
Fairfax, VA 22030, USA  
m crawsha@gmu.edu
Yajie Bao*  
School of Mathematical Sciences  
Shanghai Jiao Tong University  
Shanghai, China  
baoyajie2019stat@sjtu.edu.cn
Mingrui Liu†  
Department of Computer Science  
George Mason University  
Fairfax, VA 22030, USA  
mingruil@gmu.edu
We study the problem of Federated Learning (FL) under client subsampling and data heterogeneity with an objective function that has potentially unbounded smoothness. This problem is motivated by empirical evidence that the class of relaxed smooth functions, where the Lipschitz constant of the gradient scales linearly with the gradient norm, closely resembles the loss functions of certain neural networks such as recurrent neural networks (RNNs) with possibly exploding gradient. We introduce EPISODE++, the first algorithm to solve this problem. It maintains historical statistics for each client to construct control variates and decide clipping behavior for sampled clients in the current round. We prove that EPISODE++ achieves linear speedup in the number of participating clients, reduced communication rounds, and resilience to data heterogeneity. Our upper bound proof relies on novel techniques of recursively bounding the client updates under unbounded smoothness and client subsampling, together with a refined high probability analysis. In addition, we prove a lower bound showing that the convergence rate of a special case of clipped minibatch SGD (without randomness in the stochastic gradient and with randomness in client subsampling) suffers from an explicit dependence on the maximum gradient norm of the objective in a sublevel set, which may be large. This effectively demonstrates that applying gradient clipping to minibatch SGD in our setting does not eliminate the problem of exploding gradients. Our lower bound is based on new constructions of hard instances tailored to client subsampling and a novel analysis of the trajectory of the algorithm in the presence of clipping. Lastly, we provide an experimental evaluation of EPISODE++ when training RNNs on federated text classification tasks, demonstrating that EPISODE++ outperforms strong baselines in FL. The code is available at https://github.com/MingruiLiu-ML-Lab/episode_plusplus.
*Equal Contribution. †Corresponding Author.
Table 1: Best complexity to find an  $\epsilon$ -stationary point for various methods and settings. The setting column describes the features of the setting in which the algorithm can solve the problem: (Re) denotes relaxed smoothness, (H) heterogeneous data, and (S) client subsampling.  $\sigma$ : stochastic gradient noise,  $\kappa$ : client heterogeneity,  $\Delta$ : objective gap at the initial solution,  $I$ : local steps,  $R$ : communication rounds,  $T = RI$ : iteration complexity,  $N$ : number of clients,  $S$ : number of subsampled clients.  $\dagger$  denotes a high probability guarantee.  $\tilde{O}(\cdot)$  and  $\tilde{\Omega}(\cdot)$  omit logarithmic terms.
Method	Communication Complexity (R)	Best Iteration Complexity	Largest I to guarantee linear speedup	Setting
Local SGD [48]	O\(\left(\frac{\Delta L\sigma^2}{NI\epsilon^4}+\frac{\Delta L\kappa^2NI}{\sigma^2\epsilon^2}+\frac{\Delta LN}{\epsilon^2}\right)\)	O\(\left(\frac{L\sigma^2}{N\epsilon^4}\right)\)	O\(\left(\frac{\sigma^2}{\kappa N\epsilon}\right)\)	(H)
SCAFFOLD [24]	O\(\left(\frac{\Delta L\sigma^2}{SIE^4}+\frac{\Delta L}{\epsilon^2}\right)\)	O\(\left(\frac{\Delta L\sigma^2}{S\epsilon^4}\right)\)	O\(\left(\frac{\sigma^2}{N\epsilon^2}\right)\)	(H), (S)
CELGC [32]	O\(\left(\frac{\Delta L_0\sigma^2}{NI\epsilon^4}\right)\)	O\(\left(\frac{\Delta L_0\sigma^2}{N\epsilon^4}\right)\)	O\(\left(\frac{\sigma}{N\epsilon}\right)\)	(Re)
EPISODE [9]	O\(\left(\frac{\Delta L_0\sigma^2}{NI\epsilon^4}+\frac{\Delta(L_0+L_1(\kappa+\sigma))}{\epsilon^2}\right)\)\(1+\frac{\sigma}{\epsilon}\))	O\(\left(\frac{\Delta L_0\sigma^2}{N\epsilon^4}\right)\)	O\(\left(\frac{L_0\sigma^2}{(L_0+L_1(\kappa+\sigma)(1+\frac{\sigma}{\epsilon}))N\epsilon^2}\right)\)	(Re), (H)
EPISODE++ (Theorem 1)†	\(\tilde{O}\left(\frac{\Delta L_0\sigma^2}{SIE^4}+\frac{\Delta(L_0+L_1(\kappa+\rho\sigma)}{I\epsilon^3}\frac{L_0}{L_1\rho}\right)\)	\(\tilde{O}\left(\frac{\Delta L_0\sigma^2}{S\epsilon^4}\right)\)	\(\tilde{O}\left(\frac{L_0\sigma^2}{(L_0+L_1(\kappa+\rho\sigma))( \sigma+\frac{L_0}{L_1\rho})S\epsilon}\right)\)	(Re), (H), (S)
Clipped Minibatch SGD (Theorem 2)	\(\tilde{\Omega}\left(\frac{\Delta L_1M}{\epsilon^2}\right)\)	\(\tilde{\Omega}\left(\frac{\Delta L_1MI}{\epsilon^2}\right)\)	-	(Re), (H), (S)
Federated Learning (FL) [33, 23] is a distributed learning paradigm in which many clients collaboratively train a machine learning model while communicating over a network, which preserves privacy and leverages parallelism across many clients. Minimizing communication cost, accounting for data heterogeneity across clients, and allowing for partial client participation are core principles of FL. Interest in FL has grown in recent years, especially with user-facing applications such as next word prediction on smartphones [19].
Optimization is a central part of FL algorithms, and most work on non-convex optimization assumes a smooth objective in both the single machine setting [15, 16, 1] and the FL setting [41, 48, 24, 26]. However, recent work [52, 8] has shown empirical evidence that certain neural networks (LSTMs [21], and Transformers [42]) do not satisfy this assumption, but do satisfy a weaker condition known as relaxed smoothness [52]. Under relaxed smoothness, techniques such as gradient clipping [37] are essential for avoiding exploding gradients. To avoid the negative effects of exploding gradients, GD without gradient clipping requires a step size inversely proportional to the maximum gradient norm of the objective in a sublevel set (denoted as  $M$ ), resulting in very slow convergence. The usefulness of gradient clipping under relaxed smoothness matches observations of training these neural networks in practice, for example on natural language tasks, which are common in FL [23].
However, little work yet exists for FL in the relaxed smoothness setting. Liu et al. [32] introduced a communication-efficient gradient clipping algorithm for FL under relaxed smoothness, with the additional assumption of homogeneous client data and a distributional assumption on the noise of stochastic gradients. The EPISODE algorithm [9] was subsequently introduced to handle heterogeneous client data in this setting, but requires full client participation, that is, that every client participates in every communication round. This significantly decreases the practical applicability of EPISODE, since full client participation is rarely achievable with large-scale FL in practice [23].
In this work, we introduce EPISODE++, the first algorithm for FL under relaxed smoothness, client heterogeneity, and client subsampling. EPISODE++ maintains statistics of the history of gradients for each client, and uses these statistics to (1) correct each local update step to approximate an update on the global loss and (2) determine at which steps the clipping operation should be performed. We prove that EPISODE++ achieves linear speedup in the number of participating clients, has reduced communication cost, and enjoys convergence rate independent of client heterogeneity.
A previous line of work in the non-convex smooth stochastic setting [45, 47, 27] compares FL algorithms against a classical baseline: minibatch SGD [40]. Despite the impressive results of newer
algorithms in practice, only a few works have proven that some FL algorithms have a theoretical advantage over minibatch SGD [24, 27]. It is therefore natural to ask what is the analogue of minibatch SGD in the relaxed smooth setting, and whether it is possible to improve upon this analogue. In this work we consider clipped minibatch SGD, which is obtained by applying gradient clipping to each update of minibatch SGD. We demonstrate a surprising negative result for clipped minibatch SGD: under client subsampling, relaxed smoothness, and client heterogeneity, the convergence rate of clipped minibatch SGD still depends on  $M$ . This implies that gradient clipping does not help minibatch SGD in this setting.
- We introduce EPISODE++, the first algorithm for FL under relaxed smoothness, client heterogeneity, and client subsampling. We prove that EPISODE++ achieves linear speedup in the number of participating clients, reduced communication cost, and has convergence rate independent of client heterogeneity. Table 1 shows a detailed comparison of the complexity of various algorithms. To achieve this result, we introduce novel techniques of recursively bounding the client updates in the presence of unbounded smoothness and client data heterogeneity, together with a refined high probability analysis.
- We demonstrate a lower bound for clipped minibatch SGD in which the convergence rate depends on  $M$  (the maximum gradient norm of the objective function in a sublevel set). This shows that, in our setting, clipped minibatch SGD is susceptible to exploding gradients, and avoiding them requires a very small learning rate, which slows down convergence. Our lower bound is based on new constructions of hard instances tailored to client subsampling and a novel analysis of the trajectory of the algorithm in the presence of clipping.
- We empirically evaluate EPISODE++ against strong baselines when training RNNs on federated text classification tasks. The results show that EPISODE++ consistently outperforms baselines across various client participation ratios and is resilient to heterogeneous client data, which is consistent with our theory.
Federated Learning FL was proposed by [33], where the authors designed the FedAvg algorithm, which is also referred to as local SGD in the literature [41, 31, 48]. The local SGD algorithm has been analyzed in various settings, including convex smooth setting [41, 11, 28, 24, 45, 47, 26, 49, 17, 25], convex composite setting [50, 2, 35], and nonconvex smooth setting [22, 43, 31, 18, 48, 28, 24, 38, 53, 26]. There is a line of work which specially considered client partial participation in FL [6, 13, 24, 29, 44, 7] under convex or nonconvex smooth settings. Recently, Liu et al. [32] and Crawshaw et al. [9] considered FL with nonconvex and relaxed smooth functions for homogeneous and heterogeneous data respectively. However, they assume full client participation and neither of them are applicable to the case of client subsampling.
Relaxed Smoothness Relaxed smoothness was proposed by [52] as a relaxation of the standard smoothness condition, which is used to model the exploding gradient problem in training deep neural networks such as recurrent neural networks [36, 37] and long-short term memory networks [52], language models [14, 34] and transformers [8]. Zhang et al. [52] proved that gradient clipping converges faster than any fixed step size gradient descent for relaxed smooth functions. The complexity bound in [52] was further improved by [51]. Recently, there is a line of work which considered different algorithms and various analyses under relaxed smoothness [8, 39, 12]. However, all of them focused on single machine setting and may not be applicable to FL setting.
Lower bounds in Federated Learning There are several lower bound results for FL algorithms. Woodworth et al. [45, 47] compared minibatch SGD and local SGD in the regime of federated stochastic convex optimization setting for homogeneous and heterogeneous data and established lower bounds for local SGD. Woodworth et al. [46] proved a min-max complexity of distributed stochastic convex optimization for any intermittent communication algorithm. Glasgow et al. [17] established improved lower bounds of local SGD in convex optimization setting for both homogeneous and heterogeneous data. However, all of these lower bounds are not applicable to our setting where the problem instance is relaxed smooth with heterogeneous data and client subsampling.
We consider federated learning with heterogeneous and stochastic objectives, where the goal is to minimize the average loss function across  $N$  clients. For  $i \in [N]$ , let  $f_{i}(\boldsymbol{x}) = \mathbb{E}_{\xi \sim \mathcal{D}_{i}}[F_{i}(\boldsymbol{x}; \xi)]$  be the objective of the  $i$ -th client, where  $\mathcal{D}_{i}$  is the underlying data distribution of the  $i$ -th client. Then the global objective is
$$
\min  _ {\boldsymbol {x} \in \mathbb {R} ^ {d}} \left\{f (\boldsymbol {x}) := \frac {1}{N} \sum_ {i = 1} ^ {N} f _ {i} (\boldsymbol {x}) \right\}. \tag {1}
$$
Since each  $f_{i}$  is not necessarily convex, we consider the problem of finding an  $\epsilon$ -stationary point, that is, a point  $\pmb{x} \in \mathbb{R}^{d}$  such that  $\| \nabla f(\pmb{x}) \| \leq \epsilon$ .
Most works on non-convex optimization [15, 16, 1] consider the case where each  $f_{i}$  is  $L$ -smooth, so that  $\| \nabla^2 f_i(\pmb{x})\| \leq L$  for every  $\pmb{x}\in \mathbb{R}^d$ . However, several works [52, 8] have shown empirical evidence that objective functions corresponding to some types of neural networks (such as LSTMs [21] and Transformers [42]) do not satisfy this condition, but do satisfy a strictly weaker condition known as  $(L_0,L_1)$ -smoothness, or relaxed smoothness. A second-order differentiable function  $g:\mathbb{R}^d\to \mathbb{R}$  is called  $(L_0,L_1)$ -smooth if  $\| \nabla^2 g(\pmb {x})\| \leq L_0 + L_1\| \nabla g(\pmb {x})\|$  for all  $\pmb {x}\in \mathbb{R}^{d}$ . Notice that any  $L$ -smooth function is  $(L,0)$ -smooth.
In this work we consider the problem described in (1) under the following assumptions.
Assumption 1. (i) Denoting by  $\mathbf{x}_0$  the initial iterate, there exists some  $\Delta > 0$  such that  $f(\mathbf{x}_0) - \min_{\mathbf{x} \in \mathbb{R}^d} f(\mathbf{x}) \leq \Delta$ . (ii) Each  $f_i$  and  $f$  is  $(L_0, L_1)$ -smooth. (iii) There exist  $\kappa \geq 0$ ,  $\rho \geq 1$  such that  $\| \nabla f_i(\mathbf{x}) \| \leq \kappa + \rho \| \nabla f(\mathbf{x}) \|$  for all  $\mathbf{x} \in \mathbb{R}^d$ . (iv) There exists  $\sigma \geq 0$  such that  $\mathbb{E}_{\xi \sim \mathcal{D}_i} [\nabla F_i(\mathbf{x}; \xi)] = \nabla f_i(x)$  and  $\| \nabla F_i(\mathbf{x}; \xi) - \nabla f_i(\mathbf{x}) \| \leq \sigma$  almost surely for  $\xi \sim \mathcal{D}_i$ .
Assumption  $(i)$  is standard in non-convex optimization [15, 16]. Assumption  $(ii)$  is typically used in the FL literature [32, 9]. Assumption  $(iii)$  is used in heterogeneous federated learning [24] and describes the heterogeneity between client objectives: if  $\kappa = 0$  and  $\rho = 1$ , then all client objectives  $f_{i}$  are equal. Assumption  $(iv)$  is common in the relaxed smoothness setting [52, 51, 32, 8, 9].
In addition, we consider the case of partial client participation in federated learning, also known as client subsampling. With partial participation, only  $S$  out of  $N$  clients will participate in each communication round, which exacerbates the issue of client heterogeneity.
4 Algorithm and Convergence Analysis
4.1 Main Challenges and Algorithm Design
We first illustrate why existing algorithms such as SCAFFOLD [24] and EPISODE [9] are not able to handle heterogeneous data, relaxed smoothness and client subsampling simultaneously. The analysis of SCAFFOLD crucially requires the function to be  $L$ -smooth to recursively bound (1) the lag error from client subsampling and (2) client drift from local updates, but this argument is not applicable for relaxed smooth functions whose gradient information changes quickly. EPISODE has convergence guarantees for relaxed smooth functions and heterogeneous data, but only with full client participation. A naive variant of EPISODE in the client subsampling case does not work: the indicator of gradient clipping is based only on information from clients participating in the current round and ignores information from unsampled clients, which introduces non-negligible bias from client heterogeneity.
To address these challenges, we design a new algorithm named EPISODE++, which is presented in Algorithm 1. Similar to EPISODE [9], our algorithm utilizes episodic gradient clipping, which determines whether a clipping operation will be performed depending on the size of the average control variate  $G_{r}$ . This means that during each round, either (1) all clients perform a normalized update for all steps, or (2) all clients will perform an unnormalized update for all steps. However, different from EPISODE, our algorithm corrects local updates with control variates  $G_{r}^{i}$  computed as the averaged stochastic gradient over the previous round in which client  $i$  participated, as in SCAFFOLD. As we will show in our proof, the episodic gradient clipping together with the update correction strategy allows the algorithm to progress in a stable manner: and it will sufficiently decrease the objective value and also avoid the negative effect of possibly exploding gradients.
Algorithm 1 EPISODE++
1: Initialize  $\bar{\pmb{x}}_0,\pmb{G}_0^i\gets \nabla F_i(\bar{\pmb{x}}_0,\tilde{\xi}_i),\pmb{G}_0\gets \frac{1}{N}\sum_{i = 1}^{N}\pmb{G}_0^i$    
2: for  $r = 0,1,\ldots ,R - 1$  do   
3: Sample  $\mathcal{S}_r\subset [N]$  uniformly at random such that  $|\mathcal{S}_r| = S$    
4: for  $i\in \mathcal{S}_r$  do   
5:  $\pmb{x}_{r,0}^{i}\gets \bar{\pmb{x}}_{r}$    
6: for  $k = 0,\dots ,I - 1$  do   
7: Sample  $\nabla F_i(\pmb{x}_{r,k}^i;\xi_{r,k}^i)$  , where  $\xi_{r,k}^{i}\sim \mathcal{D}_{i}$    
8:  $\pmb{g}_{r,k}^{i}\gets \nabla F_{i}(\pmb{x}_{r,k}^{i};\xi_{r,k}^{i}) - \pmb{G}_{r}^{i} + \pmb{G}_{r}$    
9:  $\pmb{x}_{r,k + 1}^{i}\gets \pmb{x}_{r,k}^{i} - \eta \pmb{g}_{r,k}^{i}\mathbb{1}_{\| \pmb{G}_r\| \leq \gamma /\eta} - \gamma \frac{\pmb{g}_{r,k}^{i}}{\| \pmb{g}_{r,k}^{i}\|}\mathbb{1}_{\| \pmb{G}_r\| \geq \gamma /\eta}$    
10: end for   
11:  $G_{r + 1}^{i}\gets \frac{1}{I}\sum_{k = 0}^{I - 1}\nabla F_{i}(x_{r,k}^{i};\xi_{r,k}^{i})$    
12:  $\Delta G_r^i\gets G_{r + 1}^i -G_r^i$    
13: end for   
14: Update  $\bar{\pmb{x}}_{r + 1}\gets \frac{1}{S}\sum_{i\in S_r}\pmb{x}_{r,I}^i$    
15: Update  $G_{r + 1}\gets G_r + \frac{1}{N}\sum_{i\in S_r}\Delta G_r^i$    
16: Denote  $G_{r + 1}^{i}\gets G_{r}^{i}$  for all  $i\notin S_r$  3   
17: end for
4.2 Convergence Result
The following result proves that EPISODE++ converges to an  $\epsilon$ -stationary point with high probability.
Theorem 1. Let  $\epsilon \leq \frac{AL_0}{16BL_1\rho}$  and  $\delta \in (0,1)$ . Denote  $K = \left\lceil \frac{\log(RN / \delta)}{\log(N / (N - S))} \right\rceil$ ,  $\Gamma_1 \coloneqq AL_0 + BL_1\kappa + 4BL_1\rho\left(2\sigma + \frac{\gamma}{\eta}\right)$  and  $\Gamma_2 \coloneqq 64\left(\kappa + 5\rho\left(2\sigma + \frac{\gamma}{\eta}\right)\right)^2\left(\frac{1}{S} - \frac{1}{N}\right)$ . If
$$
\eta \leq \min  \left\{\frac {1}{9 0 (K + 1) \Gamma_ {1} I}, \frac {\epsilon}{3 2 \Gamma_ {1} I \left(7 4 \sigma + \frac {A L _ {0}}{B L _ {1} \rho}\right)}, \frac {S \epsilon^ {2}}{2 1 6 A L _ {0} \sigma^ {2} \log \frac {1}{\delta}}, \frac {\Delta}{\log \frac {1}{\delta}} \min  \left\{\frac {1}{1 2 \sigma^ {2}}, \frac {1}{\Gamma_ {2} I} \right\} \right\}, \tag {2}
$$
and  $\gamma = \left(72\sigma +\frac{AL_0}{BL_1\rho}\right)\eta$  then Algorithm 1 satisfies  $\frac{1}{R}\sum_{r = 0}^{R - 1}\| \nabla f(\bar{\pmb{x}}_r)\| \leq 35\epsilon$  with probability at least  $1 - 15\delta$ , as long as  $R\geq \frac{8\Delta}{\epsilon^2\eta I}$ .
The above result holds for a wide range of  $\eta$  and  $I$ , and for any noise level  $\sigma$ . The corollary below summarizes the best possible iteration complexity and communication complexity implied by Theorem 1. The full proofs of Theorem 1 and Corollary 1 can be found in Appendix A.5 and A.6.
Corollary 1. Suppose  $\sigma >0$ . If, under the setting of Theorem 1, we additionally choose  $\epsilon \leq \min \left\{\frac{AL_0}{16BL_1\rho},\frac{16\left(74\sigma + \frac{AL_0}{BL_1\rho}\right)}{45(K + 1)},\sqrt{\frac{18\Delta AL_0}{S}},\frac{32\Delta\Gamma_1\left(74\sigma + \frac{AL_0}{BL_1\rho}\right)}{\Gamma_2\log\frac{1}{\delta}}\right\}$ ,  $\eta$  as large as possible under (2), and  $I \leq \frac{27AL_0\sigma^2\log\frac{1}{\delta}}{4\epsilon\Gamma_1S\left(74\sigma + \frac{AL_0}{BL_1\rho}\right)}$ , then Algorithm 1 has iteration complexity  $RI = O\left(\frac{\Delta L_0\sigma^2\log\frac{1}{\delta}}{S\epsilon^4}\right)$ . If additionally  $I = \Theta \left(\frac{L_0\sigma^2\log\frac{1}{\delta}}{\epsilon S(L_0 + L_1(\kappa + \rho\sigma))(\sigma + \frac{L_0}{L_1\rho})}\right)$ , then Algorithm 1 has communication complexity  $R = O\left(\frac{\Delta(L_0 + L_1(\kappa + \rho\sigma))(\sigma + \frac{L_0}{L_1\rho})}{\epsilon^3}\right)$ .
<sup>3</sup>Although it appears that this operation must be performed by unsampled clients, this is just an artifact of the notation: unsampled clients do not have to execute any operations.
4.3 Proof Sketch
In this section we provide a sketch for the proof of Theorem 1. We wish to establish the descent inequality for the global objective function in each round by applying Lemma A.3 in [51]:
$$
f \left(\bar {\boldsymbol {x}} _ {r + 1}\right) \leq f \left(\bar {\boldsymbol {x}} _ {r}\right) + \left\langle \nabla f \left(\bar {\boldsymbol {x}} _ {r}\right), \bar {\boldsymbol {x}} _ {r + 1} - \bar {\boldsymbol {x}} _ {r} \right\rangle + \frac {A L _ {0} + B L _ {1} \| \nabla f \left(\bar {\boldsymbol {x}} _ {r}\right) \|}{2} \| \bar {\boldsymbol {x}} _ {r + 1} - \bar {\boldsymbol {x}} _ {r} \| ^ {2}, \tag {3}
$$
where  $A = 1 + e^{C} - \frac{e^{C} - 1}{C}$ ,  $B = \frac{e^{C} - 1}{C}$  and  $C > 0$  is an absolute constant. However, using this descent inequality requires  $\| \bar{\pmb{x}}_{r + 1} - \bar{\pmb{x}}_r\| \leq C / L_1$ . Achieving a universal bound on the distance over two consecutive rounds is nontrivial in the presence of client subsampling, data heterogeneity and relaxed smoothness, which requires a new analysis.
Compared with SCAFFOLD [24] and EPISODE [9], the main difficulty of analyzing EPISODE++ lies in controlling the distance between local weights  $\boldsymbol{x}_{r,k}^{i}$  and the synchronization point  $\bar{\boldsymbol{x}}_r$  under relaxed smoothness and client subsampling. In EPISODE, the magnitude of  $g_{r,k}^{i}$  can be bounded in terms of  $\| \nabla f_i(\boldsymbol{x}_{r,k}^i) - \nabla f_i(\bar{\boldsymbol{x}}_r)\|$  and  $\| G_r\|$ , since  $G_r^i$  is evaluated at the synchronization point  $\bar{\boldsymbol{x}}_r$ . Then the distance can be bounded recursively when clipping does not happen (i.e.,  $\| G_r\| \leq \gamma /\eta$ ). However, EPISODE++ utilizes historical gradients to construct the indicator  $G_{r}$ , which means the increment also depends on the lag error  $\| G_r^i -\nabla f_i(\bar{\boldsymbol{x}}_r)\|$ . Due to relaxed smoothness, we cannot bound the lag error as in SCAFFOLD [24] because we do not know the distance between  $\bar{\boldsymbol{x}}_r$  and where  $G_r^i$  is evaluated. In the extreme case that client  $i$  has never been sampled before round  $r$ , we have  $G_r^i = \nabla F_i(\bar{\boldsymbol{x}}_0;\tilde{\xi}_i)$ , so the lag error can be unbounded when  $r$  is large.
To address the issue mentioned above, we analyze the convergence of EPISODE++ in a high-probability framework. Let  $\pmb{y}_{r,k}^{i}$  denote the local model of client  $i$  at step  $k$  during the most recent round in which client  $i$  participated (before round  $r$ ), and let  $q_{r}^{i}$  denote the index of this round. Since client subsampling independent across rounds,  $\mathbb{P}(q_r^i - r \geq K) = (1 - S / N)^K$ . Therefore, we introduce the event  $\mathcal{E} := \{\max_{0 \leq r \leq R-1, 1 \leq i \leq N} r - q_r^i \leq K\}$ , where  $K$  is a logarithmic factor such that  $\mathbb{P}(\mathcal{E}) \geq 1 - \delta$ . Under  $\mathcal{E}$ , we obtain the following lemma to bound the distance of local updates.
Lemma 1. Suppose that (9) holds. Then for any  $r \geq 0$  we have
$$
\mathbb {1} _ {\mathcal {E}} \cdot \max  _ {k \in [ I ]} \| \boldsymbol {x} _ {r, k} ^ {i} - \bar {\boldsymbol {x}} _ {r} \| \leq 2 I (2 \sigma \eta + \gamma), \quad f o r a n y i \in \mathcal {S} _ {r}, \tag {4}
$$