SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hanhan Zhou  
The George Washington University  
hanhan@gwu.edu
Guru Venkataramani  
The George Washington University  
guruv@gwu.edu
Tian Lan  
The George Washington University  
tlan@gwu.edu
Wenbo Ding Tsinghua-Berkeley Shenzhen Institute ding.wenbo@sz.tsinghua.edu.cn
Cross-device Federated Learning (FL) faces significant challenges where low-end clients that could potentially make unique contributions are excluded from training large models due to their resource bottlenecks. Recent research efforts have focused on model-heterogeneous FL, by extracting reduced-size models from the global model and applying them to local clients accordingly. Despite the empirical success, general theoretical guarantees of convergence on this method remain an open question. This paper presents a unifying framework for heterogeneous FL algorithms with online model extraction and provides a general convergence analysis for the first time. In particular, we prove that under certain sufficient conditions and for both IID and non-IID data, these algorithms converge to a stationary point of standard FL for general smooth cost functions. Moreover, we introduce the concept of minimum coverage index, together with model reduction noise, which will determine the convergence of heterogeneous federated learning, and therefore we advocate for a holistic approach that considers both factors to enhance the efficiency of heterogeneous federated learning.
Federated Learning (FL) is a machine learning paradigm that enables a massive number of distributed clients to collaborate and train a centralized global model without exposing their local data [1]. Heterogeneous FL is confronted with two fundamental challenges: (1) mobile and edge devices that are equipped with drastically different computation and communication capabilities are becoming the dominant source for FL [2], also known as device heterogeneity; (2) state-of-the-art machine learning model sizes have grown significantly over the years, limiting the participation of certain devices in training. This has prompted significant recent attention to a family of FL algorithms relying on training reduced-size heterogeneous local models (often obtained through extracting a subnet or pruning a shared global model) for global aggregation. It includes algorithms such as HeteroFL [3] that employ fixed heterogeneous local models, as well as algorithms like PruneFL [4] and FedRolex [5] that adaptively select and train pruned or partial models dynamically during training. However, the success of these algorithms has only been demonstrated empirically (e.g., [2, 4, 3]). Unlike standard FL that has received rigorous analysis [6, 7, 8, 9], the convergence of heterogeneous FL algorithms is still an open question.
This paper aims to answer the following questions: Given a heterogeneous FL algorithm that trains a shared global model through a sequence of time-varying and client-dependent local models, what conditions can guarantee its convergence? And intrinsically how do the resulting models compare to that of standard FL? There have been many existing efforts in establishing convergence guarantees for FL algorithms, such as the popular FedAvg [1], on both IID (independent and identically distributed data) and non-IID[9] data distributions, but all rely on the assumption that local models share the same uniform structure as the global model  ${}^{1}$ . Training heterogeneous local models, which could change both over time and across clients in FL is desirable due to its ability to adapt to resource constraints and training outcomes[10].
For general smooth cost functions and under standard FL assumptions, we prove that heterogeneous FL algorithms satisfying certain sufficient conditions can indeed converge to a neighborhood of a stationary point of standard FL (with a small optimality gap that is characterized in our analysis), at a rate of  $O\left(\frac{1}{\sqrt{Q}}\right)$  in  $Q$  communication rounds. Moreover, we show not only that FL algorithms involving local clients training different subnets (pruned or extracted from the global model) will converge, but also that the more they cover the parameters space in the global model, the faster the training will converge. Thus, local clients should be encouraged to train with reduced-size models that are of different subnets of the global model rather than pruning greedily. The work extends previous analysis on single-model adaptive pruning and subnetwork training[11, 12] to the FL context, where a fundamental challenge arises from FL's local update steps that cause heterogeneous local models (obtained by pruning the same global model or extracting a submodel) to diverge before the next aggregation. We prove
a new upperbound and show that the optimality gap (between heterogeneous and standard FL) is affected by both model-reduction noise and a new notion of minimum coverage index in FL (i.e., any parameters in the global model are included in at least  $\Gamma_{\mathrm{min}}$  local models).
The key contribution of this paper is to establish convergence conditions for federated learning algorithms that employ heterogeneous arbitrarily-pruned, time-varying, and client-dependent local models to converge to a stationary point of standard FL. Numerical evaluations validate the sufficient conditions established in our analysis. The results demonstrate the benefit of designing new model reduction strategies with respect to both model reduction noise and minimum coverage index.
Standard Federated Learning A standard FL problem considers a distributed optimization for N clients:
$$
\min  _ {\theta} \left\{F (\theta) \triangleq \sum_ {i = 1} ^ {N} p _ {i} F _ {i} (\theta) \right\}, \text {w i t h} F \left(\theta_ {i}\right) = \mathbb {E} _ {\xi \sim D _ {i}} l \left(\xi_ {i}, \theta_ {i}\right), \tag {1}
$$
where  $\theta$  is as set of trainable weights/parameters,  $F_{n}(\theta)$  is a cost function defined on data set  $D_{i}$  with respect to a user specified loss function  $l(x,\theta)$ , and  $p_i$  is the weight for the  $i$ -th client such that  $p_i \geq 0$  and  $\sum_{i=1}^{N} p_i = 1$ .
Figure 1: In this paper we show that instead of pruning small parameters greedily, local clients when applied with different local models not only will converge under certain conditions, it might even converge faster.
Throughout this paper, "non-IID data" means that the data among local clients are not independent and identically distributed. "Heterogeneous" means each client model obtained by model reduction from a global model can be different from the global model and other clients. "Dynamic" means time-varying, i.e. the model for one local client could change between each round.
The FL procedure, e.g., FedAvg [1], typically consists of a sequence of stochastic gradient descent steps performed distributedly on each local objective, followed by a central step collecting the workers' updated local parameters and computing an aggregated global parameter. For the  $q$ -th round of training, first, the central server broadcasts the latest global model parameters  $\theta_q$  to clients  $n = 1,\dots ,N$ , who perform local updates as follows:
$$
\theta_ {q, n, t} = \theta_ {q, n, t - 1} - \gamma \nabla F _ {n} \left(\theta_ {q, n, t - 1}; \xi_ {n, t - 1}\right) \text {w i t h} \theta_ {q, n, 0} = \theta_ {q}
$$
where  $\gamma$  is the local learning rate. After all available clients have concluded their local updates (in  $T$  epochs), the server will aggregate parameters from them and generate the new global model for the next round, i.e.,  $\theta_{q + 1} = \sum_{n = 1}^{N}p_{i}\theta_{q,n,T}$  The formulation captures FL with both IID and non-IID data distributions.
Federated Averaging and Related Convergence Analysis. FedAvg [1] is consi dered the first and the most commonly used federated learning algorithm. Several works have shown the convergence of FedAvg under several different settings with both homogeneous (IID) data [6, 13] and heterogeneous (non-IID) data [9, 7, 8] even with partial clients participation [14]. Specifically, [8] demonstrated LocalSGD achieves  $O\left(\frac{1}{\sqrt{NQ}}\right)$  convergence for non-convex optimization and [9] established a convergence rate of  $O\left(\frac{1}{Q}\right)$  for strongly convex problems on FedAvg, where Q is the number of SGDs and N is the number of participated clients.
Efficient and Heterogeneous FL through Neural Network Pruning and Sparsification. Several works [15, 16, 17, 18, 19, 20] are proposed to further reduce communication costs in FL. One direction is to use data compression such as quantization [21, 7, 22, 23], sketching [24, 25], split learning [26], learning with gradient sparsity [27] and sending the parameters selectively[28]. This type of work does not consider computation efficiency. There are also works that address the reduction of both computation and communication costs, including one way to utilize lossy compression and dropout techniques[29, 30]. Although early works mainly assume that all local models share the same architecture as the global model [31], recent works have empirically demonstrated that federated learning with heterogeneous client models to save both computation and communication is feasible. PruneFL[4] proposed an approach with adaptive parameter pruning during FL. [32] proposed FL with a personalized and structured sparse mask. FjORD[33] and HetroFL[3] proposed to generate heterogeneous local models as a subnet of the global network by extracting a static sub-models, Hermes[34] finds the small sub-network by applying the structured pruning. There are also researches on extracting a subnetwork dynamically, e.g. Federated Dropout[29] extracts submodels randomly and FedRolex[5] applies a rolling sub-model extraction.
Despite their empirical success, they either lack theoretical convergence analysis or are specific to their own work. PruneFL only shows a convergence of the proposed algorithm and does not ensure convergence to a solution of standard FL. Meanwhile, static subnet extraction like Hermes does not allow the pruned local networks to change over time nor develop general convergence conditions. Following Theorem 1 works like Hermes can now employ time-varying subnet extractions, rather than static subnets, while still guaranteeing the convergence to standard FL. The convergence of HeteroFL and FedRolex – which was not available – now follows directly from Theorem 1. Once PruneFL satisfies the conditions established in our Theorem 1, convergence to a solution of standard FL can be achieved, rather than simply converging to some point. In summary, our general convergence conditions in Theorem 1 can provide support to existing FL algorithms that employ heterogeneous local models, ensuring convergence to standard FL. It also enables the design of optimized pruning masks/models to improve the minimum coverage index and thus the resulting gap to standard FL.
4.1 Problem Formulation for FL with Heterogeneous Local models
Given an FL algorithm that trains heterogeneous local models for global aggregation, our goal is to analyze its convergence with respect to a stationary point of standard FL. We consider a general formulation where the heterogeneous local models can be obtained using any model reduction
strategies that are both (i) time-varying to enable online adjustment of reduced local models during the entire training process and (ii) different across FL clients with respect to their individual heterogeneous computing resource and network conditions. More formally, we denote the sequence of local models used by a heterogeneous FL algorithm by masks  $m_{q,n} \in \{0,1\}^{|\theta|}$ , which can vary at any round  $q$  and for any client  $n$ . Let  $\theta_q$  denote the global model at the beginning of round  $q$  and  $\odot$  be the element-wise product. Thus,  $\theta_q \odot m_{q,n}$  defines the trainable parameters of the reduced local model for client  $n$  in round  $q$ . Our goal is to find sufficient conditions on such masks  $m_{q,n} \forall q, n$  for the convergence of heterogeneous FL.
Here, we describe one around (say the  $q$ th) of the heterogeneous FL algorithm. First, the central server employs a given model reduction strategy  $\mathbb{P}(\cdot)$  to reduce the latest global model  $\theta_q$  and broadcast the resulting local models to clients:
$$
\theta_ {q, n, 0} = \theta_ {q} \cdot m _ {q, n}, \text {w i t h} m _ {q, n} = \mathbb {P} \left(\theta_ {q}, n, q\right), \forall n. \tag {2}
$$
We note that the model reduction strategy  $\mathbb{P}(\theta_q, n, q)$  can vary over time  $q$  and across clients  $n$  in heterogeneous FL. Each client  $n$  then trains the reduced local model by performing  $T$  local updates (in  $T$  epochs):
$$
\theta_ {q, n, t} = \theta_ {q, n, t - 1} - \gamma \nabla F _ {n} \left(\theta_ {q, n, t - 1}, \xi_ {n, t - 1}\right) \odot m _ {q, n}, f o r t = 1 \dots T,
$$
where  $\gamma$  is the learning rate and  $\xi_{n,t - 1}$  are independent samples uniformly drawn from local data  $D_{n}$  at client  $n$ . We note that  $\nabla F_{n}(\theta_{q,n,t - 1},\xi_{n,t - 1})\odot m_{q,n}$  is a local stochastic gradient evaluated using only local parameters in  $\theta_{q,n,t - 1}$  (available to the heterogeneous local model) and that only locally trainable parameters are updated by the stochastic gradient (via an element-wise product with  $m_{q,n}$ ).
Finally, the central server aggregates the local models  $\theta_{n,q,T} \forall n$  and produces an updated global model  $\theta_{q+1}$ . Due to the use of heterogeneous local models, each global parameter is included in a (potentially) different subset of the local models. Let  $\mathcal{N}_q^{(i)}$  be the set of clients, whose local models contain the  $i$ th modeling parameter in round  $q$ . That is  $n \in \mathcal{N}_q^{(i)}$  if  $m_{q,n}^{(i)} = 1$  and  $n \notin \mathcal{N}_q^{(i)}$  if  $m_{q,n}^{(i)} = 0$ . Global update of the  $i$ th parameter is performed by aggregating local models with the parameter available, i.e.,
$$
\theta_ {q + 1} ^ {(i)} = \frac {1}{| \mathcal {N} _ {q} ^ {(i)} |} \sum_ {n \in \mathcal {N} _ {q} ^ {(i)}} \theta_ {q, n, T} ^ {(i)}, \forall i, \tag {3}
$$
where  $|\mathcal{N}_q^{(i)}|$  is the number of local models containing the  $i$ th parameter. We summarize the algorithm details in Algorithm 1 and the explanation of the notations in the Appendix.
The fundamental challenge of convergence analysis mainly stems from the local updates in Eq.(3). While the heterogeneous models  $\theta_{q,n,0}$  provided to local clients at the beginning of round  $q$  are obtained from the same global model  $\theta_q$ , performing  $T$  local updates causes these heterogeneous local models to diverge before the next aggregation. In addition, each parameter is (potentially) aggregated over a different subset of local models in Eq.(3). These make existing convergence analysis intended for single-model adaptive pruning [11, 12] non-applicable to heterogeneous FL. The impact of local model divergence and the global aggregation of heterogeneous models must be characterized in order to establish convergence.
The formulation proposed above captures heterogeneous FL with any model pruning or sub-model extraction strategies since the resulting masks  $m_{q,n} \forall q, n$  can change over time  $q$  and across clients  $n$ . It incorporates many model reduction strategies (such as pruning, sparsification, and sub-model extraction) into heterogeneous FL, allowing the convergence results to be broadly applicable. It recovers many important FL algorithms recently proposed, including HeteroFL [3] that uses fixed masks  $m_{n}$ , PruneFL [4] that periodically trains a full-size local model with  $m_{n,q} = 1$ , Prune-andGrow [12] that can be viewed as a single-client version, as well as FedAvg [1] that employs full-size local models with  $m_{n,q} = 1$  at all clients. Our analysis establishes general conditions for any heterogeneous FL with arbitrary online model reduction that may vary over communication rounds to converge to standard FL.
2While a reduced local model has a smaller number of parameters than the global model. We adopt the notations in [4, 35, 12] and use  $\theta_q\odot m_{q,n}$  with an element-wise product to denote the pruned local model or the extracted submodel - only parameter corresponding to a 1-value in the mask is accessible and trainable in the local model.
Our paper establishes convergence conditions in the general form, which apply to both static and dynamically changing masks. With dynamically changing masks, we denote the reduced models/networks as  $\theta_{q,n,t}$ , which means that the model structure (with its corresponding mask) can change between each round of communications and during each local training epoch, and can be different from other local clients. We show that as long as the dynamic heterogenous FL framework can be framed as the setting above, our convergence analysis in this paper applies.
4.2 Notations and Assumptions
We make the following assumptions that are routinely employed in FL convergence analysis. In particular, Assumption 1 is a standard and common setting assuming Lipschitz continuous gradients. Assumption 2 follows from [12] (which is for a single-worker case) and implies the noise introduced by model reduction is bounded and quantified. This assumption is required for heterogeneous FL to converge to a stationary point of standard FL. Assumptions 3 and 4 are standard for FL convergence analysis following from [36, 37, 8, 9] and assume the stochastic gradients to be bounded and unbiased.
Assumption 1. (Smoothness). Cost functions  $F_{1}, \ldots, F_{N}$  are all  $L$ -smooth:  $\forall \theta, \phi \in \mathcal{R}^{d}$  and any  $n$ , we assume that there exists  $L > 0$ :
$$
\left\| \nabla F _ {n} (\theta) - \nabla F _ {n} (\phi) \right\| \leq L \| \theta - \phi \|. \tag {4}
$$
Assumption 2. (Model Reduction Noise). We assume that for some  $\delta^2 \in [0,1)$  and any  $q,n$ , the model reduction error is bounded by
$$
\left\| \theta_ {q} - \theta_ {q} \odot m _ {q, n} \right\| ^ {2} \leq \delta^ {2} \left\| \theta_ {q} \right\| ^ {2}. \tag {5}
$$
Assumption 3. (Bounded Gradient). The expected squared norm of stochastic gradients is bounded uniformly, i.e., for constant  $G > 0$  and any  $n, q, t$ :
$$
\mathbb {E} _ {\xi_ {q, n, t}} \left\| \nabla F _ {n} \left(\theta_ {q, n, t}, \xi_ {q, n, t}\right) \right\| ^ {2} \leq G. \tag {6}
$$
Assumption 4. (Gradient Noise for IID data). Under IID data distribution,  $\forall q, n, t$ , we assume a gradient estimate with bounded variance:
$$
\mathbb {E} _ {\xi_ {n, t}} \| \nabla F _ {n} \left(\theta_ {q, n, t}, \xi_ {n, t}\right) - \nabla F \left(\theta_ {q, n, t}\right) \| ^ {2} \leq \sigma^ {2} \tag {7}
$$
4.3 Convergence Analysis
We now analyze the convergence of heterogeneous FL for general smooth cost functions. We begin with introducing a new notion of minimum covering index, defined in this paper by
$$
\Gamma_ {\min } = \min  _ {q, i} | \mathcal {N} _ {q} ^ {(i)} |, \tag {8}
$$
where  $\Gamma_{\min}^3$  measures the minimum occurrence of the parameter in the local models in all rounds, considering  $|\mathcal{N}_q^{(i)}|$  is the number of heterogeneous local models containing the  $i$ th parameter. Intuitively, if a parameter is never included in any local models, it is impossible to update it. Thus conditions based on the covering index would be necessary for the convergence toward standard FL (with the same global model). All proofs for theorems and lemmas are collected in the Appendix with a brief proof outline provided here.
Theorem 1. Under Assumptions 1-4 and for arbitrary masks satisfying  $\Gamma_{\mathrm{min}} \geq 1$ , when choosing  $\gamma \leq 1/(6LT) \wedge \gamma \leq 1/(T\sqrt{Q})$ , heterogeneous FL converges to a small neighborhood of a stationary point of standard FL as follows:
$$
\frac {1}{Q} \sum_ {q = 1} ^ {Q} \mathbb {E} | | \nabla F (\theta_ {q}) | | ^ {2} \leq \frac {G _ {0}}{\sqrt {Q}} + \frac {V _ {0}}{T \sqrt {Q}} + \frac {H _ {0}}{Q} + \frac {I _ {0}}{\Gamma^ {*}} \cdot \frac {1}{Q} \sum_ {q = 1} ^ {Q} \mathbb {E} | | \theta_ {q} | | ^ {2}
$$
where  $G_0 = 4\mathbb{E}[F(\theta_0)]$ ,  $V_0 = 6LN\sigma^2 / (\Gamma^*)^2$ ,  $H_0 = 2L^2 NG / \Gamma^*$ , and  $I_0 = 3L^2\delta^2 N$  are constants depending on the initial model parameters and the gradient noise.
3We refer to  $\Gamma^{*}$  for all equations and derivations.
An obvious case here is that when  $\Gamma_{min} = 0$ , where there exists at least one parameter that is not covered by any of the local clients and all the client models can not cover the entire global model, we can consider the union of all local model parameters, the "largest common model" among them, as a new equivalent global model  $\hat{\theta}$  (which have a smaller size than  $\theta$ ). Then, each parameter in  $\hat{\theta}$  is covered in at least one local model. Thus Theorem 1 holds for  $\hat{\theta}$  instead and the convergence is proven - to a stationary point of  $\hat{\theta}$  rather than  $\theta$ .<sup>4</sup>
Assumption 5. (Gradient Noise for non-III data). Let  $\hat{g}_{q,t}^{(i)} = \frac{1}{|\mathcal{N}_q^{(i)}|}\sum_{n\in \mathcal{N}_q^{(i)}}\nabla F_n^{(i)}(\theta_{q,n,t},\xi_{n,t})$ . Under non-III data distribution, we assume  $\forall i,q,t$  a gradient estimate with bounded variance:
$$
\mathbb {E} _ {\boldsymbol {\xi}} \left\| \hat {g} _ {q, n, t} ^ {(i)} - \nabla F ^ {(i)} (\theta_ {q, n, t}) \right\| ^ {2} \leq \sigma^ {2}.
$$
Theorem 2. Under Assumptions 1-3 and 5, heterogeneous FL satisfying  $\Gamma_{\mathrm{min}} \geq 1$ , when choosing  $\gamma \leq 1 / \sqrt{TQ}$  and  $\gamma \leq 1 / (6LT)$ , heterogeneous FL converges to a small neighborhood of a stationary point of standard FL as follows:
$$
\frac {1}{Q} \sum_ {q = 1} ^ {Q} \mathbb {E} \| \nabla F \left(\theta_ {q}\right) \| ^ {2} \leq \frac {G _ {1}}{\sqrt {T Q}} + \frac {V _ {0}}{\sqrt {Q}} + \frac {I _ {0}}{\Gamma^ {*}} \cdot \frac {1}{Q} \sum_ {q = 1} ^ {Q} \mathbb {E} \| \theta_ {q} \| ^ {2} \tag {9}
$$
where  $G_{1} = 4\mathbb{E}[F(\theta_{0})] + 6LK\sigma^{2}$
Proof outline. There are a number of challenges in delivering main theorems. We begin the proof by analyzing the change of loss function in one round as the model goes from  $\theta_{q}$  to  $\theta_{q+1}$ , i.e.,  $F(\theta_{q+1}) - F(\theta_{1})$ . It includes three major steps: reducing the global model to obtain heterogeneous local models  $\theta_{q,n,0} = \theta_{q} \odot m_{q,n}$ , training local models in a distributed fashion to update  $\theta_{q,n,t}$ , and parameter aggregation to update the global model  $\theta_{q+1}$ .
Due to the use of heterogeneous local models whose masks  $m_{q,n}$  both vary over rounds and change for different workers, we first characterize the difference between local model  $\theta_{q,n,t}$  at any epoch  $t$  and global model  $\theta_q$  at the beginning of the current round. It is easy to see that this can be factorized into two parts: model reduction error  $\| \theta_{q,n,0} - \theta_q\| ^2$  and local training  $\| \theta_{q,n,t} - \theta_{q,n,0}\| ^2$ , which will be analyzed in Lemma 1.
Lemma 1. Under Assumption 2 and Assumption 3, for any  $q$ , we have:
$$
\sum_ {t = 1} ^ {T} \sum_ {n = 1} ^ {N} \mathbb {E} \| \theta_ {q, n, t - 1} - \theta_ {q} \| ^ {2} \leq \gamma^ {2} T ^ {2} N G + \delta^ {2} N T \cdot \mathbb {E} \| \theta_ {q} \| ^ {2} \tag {10}
$$
We characterize the impact of heterogeneous local models on global parameter updates. Specifically, we use an ideal local gradient  $\nabla F_{n}(\theta_{q})$  as a reference point and quantify the difference between aggregated local gradients and the ideal gradient. This will be presented in Lemma 2.
Lemma 2. Under Assumptions 1-3, for any  $q$ , we have:
$$
\sum_ {i = 1} ^ {K} \mathbb {E} \| \frac {1}{\Gamma_ {q} ^ {(i)} T} \sum_ {t = 1} ^ {T} \sum_ {n \in \mathcal {N} _ {q} ^ {(i)}} [ \nabla F _ {n} ^ {(i)} (\theta_ {q, n, t - 1}) - \nabla F _ {n} ^ {(i)} (\theta_ {q}) ] \| ^ {2} \leq \frac {L ^ {2} \gamma^ {2} T N G}{\Gamma^ {*}} + \frac {L ^ {2} \delta^ {2} N}{\Gamma^ {*}} \mathbb {E} \| \theta_ {q} \| ^ {2},
$$