SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kilian Pfeiffer
kilian.pfeiffer@kit.edu
Ramin Khalili
Huawei Research Center Munich
ramin.khalili@huawei.com
Jörg Henkel
henkel@kit.edu
Federated learning (FL) is usually performed on resource-constrained edge devices, e.g., with limited memory for the computation. If the required memory to train a model exceeds this limit, the device will be excluded from the training. This can lead to a lower accuracy as valuable data and computation resources are excluded from training, also causing bias and unfairness. The FL training process should be adjusted to such constraints. The state-of-the-art techniques propose training subsets of the FL model at constrained devices, reducing their resource requirements for training. However, these techniques largely limit the co-adaptation among parameters of the model and are highly inefficient, as we show: it is actually better to train a smaller (less accurate) model by the system where all the devices can train the model end-to-end than applying such techniques. We propose a new method that enables successive freezing and training of the parameters of the FL model at devices, reducing the training's resource requirements at the devices while still allowing enough co-adaptation between parameters. We show through extensive experimental evaluation that our technique greatly improves the accuracy of the trained model (by 52.4 p.p.) compared with the state of the art, efficiently aggregating the computation capacity available on distributed devices.
Federated learning (FL) has achieved impressive results in many domains and is proposed for several use cases, such as healthcare, transportation, and robotics [1, 2, 3, 4, 5, 6]. As data in FL is not processed centrally but usually on resource-constrained edge devices, training machine learning (ML) models impose a large computational burden on these devices [7]. Additionally, FL requires communication, specifically exchanging ML model parameters from the devices to a centralized entity for aggregation. Extensive research has been done to lower the communication overhead required for FL, e.g., on the use of quantization in the communication [8, 9] or sketched updates [10]. Similarly, techniques such as partial updates [11], asynchronous aggregation [12, 13], and tier-based aggregation [14, 15] have been proposed to lower and account for varying computational throughput. While constrained computation throughput and communication capabilities can slow down FL convergence, high memory requirements for training that are imposed on devices can exclude devices completely from the FL system. This is, for example, the case in Google GBoard [16], where devices that do not have 2GB of memory for training are removed. Excluding devices from training
lowers the reachable accuracy, as fewer devices participate in the training, also resulting in bias and unfairness [17].
Several techniques have been proposed to tackle these constraints, where the main idea is to train a lower complexity submodel on the devices and embed the trained submodel into the full high-capacity server model. A submodel is typically created by scaling the width of the neural network (NN), e.g., using a subset of convolutional filters per NN layer. There exist several variations of the technique [9, 18, 19, 20]. In particular, Caldas et al. [9] propose Federated Dropout (FD), which randomly, per round and per device, selects NN filters that are trained. Alam et al. [20] propose FedRolex, a sliding window approach, where all devices train the same submodel, and in each FL round, the used filter indices are shifted by one. While both these techniques allow training within given memory constraints, our results show (Fig. 1) that they perform worse than a straightforward baseline, i.e., using a smaller NN model that can be trained by all devices end-to-end. We evaluate CIFAR10, FEMNIST, and TinyImageNet in an FL setting using ResNet and scale the width of the NN down s.t. we achieve a  $2 - 8 \times$  reduction in training memory. We observe that training the same small model at all devices outperforms FedRolex and FD w.r.t to the final accuracy and convergence speed (we expect similar results for other subset-derived techniques), especially when enforcing a  $4 \times$  and  $8 \times$  memory reduction (as also evaluated in [20]).
Small Model
FedRolex [20]
—FD [9]
Figure 1: Accuracy of FedRolex [20] and FD [9] compared to a small model using different NN topologies and datasets after  $2500 / 1000$  FL rounds. In the case of  $1\times$ , both techniques are equivalent to vanilla Federated Averaging (FedAvg). Hyperparameters of the experiments are given in Section 3.
Our results indicate that applying these techniques is rather harmful. This is as a large part of filters/parameters has to be dropped during each round of training at each device, extremely limiting the co-adaptation between parameters. Hence, the gradients for the subset of parameters that are trained on devices are calculated without considering the error of the parameters that reside on the server (more details in Appendix A). Motivated by these observations, we propose a new technique that enables successive freezing and training of the parameters of the FL model at devices, reducing the training's resource requirements at the devices while allowing a higher co-adaptation between parameters. Instead of switching between subsets of the model on an FL-round basis, we train the same parameters for several rounds and successively switch to a larger model. To obey the same memory constraints as in [9, 20], we train early layers using the full width while utilizing a scaled-down NN head. We then freeze the early layers and expand the head layers' width. By freezing early layers, no activation has to be kept in memory, hence, we lower the memory footprint. But still, the error of these frozen parameters is included in the calculation of the gradient of the new subset of parameters. We apply this technique successively till all parameters of the model are trained.
In summary, we make the following novel contributions:
- We empirically show that employing current state-of-the-art techniques, FD [9] and FedRolex [20], for memory-constrained systems can actually hurt the performance.
- We propose a novel training scheme called Successive Layer Training  $(SLT)^{1}$ , which addresses the shortcomings of previous techniques by successively adding more parameters to the training, successively freezing layers, and reusing a scaled-down NN head.
- Our evaluation of common NN topologies, such as ResNet and DenseNet, shows that SLT reaches significantly higher accuracies in independent and identically distributed (iid) and non-iid CIFAR, FEMNIST, and TinyImageNet training compared to the state of the art. Also, SLT provides a much faster convergence, reducing the communication overhead to reach a certain level of accuracy by over  $10 \times$  compared with FD and FedRolex. The same
<sup>1</sup>The source code of SLT is available at https://github.com/k1l1/SLT.
Algorithm 1: Successive Layer Training:  $w$  and  $W$  label the set of all layers' parameters.
Requires: Number of rounds  $R$ , devices  $\mathcal{C}$ , number of devices per round  $|\mathcal{C}^{(r)}|$ , configurations  $S_{n}$ $n \in [1, \dots, N]$ , that satisfy constraint  $m$ , init. parameters  $W^{(1)}$
Server:
foreach round  $r = 1,2,\ldots ,R$  do  $\mathcal{C}^{(r)}\gets$  select  $|\mathcal{C}^{(r)}|$  random devices out of C  $w^{(r)},S^{(r)}\gets$  ConfigurationSelection(W(r),r)   
foreach device  $c\in \mathcal{C}^{(r)}$  in parallel do  $w^{(r)},S^{(r)}$  receive from server  $w^{(r,c)}\gets$  DeviceTraining  $(w^{(r)},S^{(r)})$  upload  $w_{j}^{(r,c)}$  to server  $\forall j\in \{j:K_F <   j\leq K\}$    
end   
 $w_{j}^{(r + 1)}\gets \frac{1}{|\mathcal{C}^{(r)}|}\sum_{c\in \mathcal{C}^{(r)}}w_{j}^{(r,c)}\quad //averaging$  of trained layers   
 $W^{(r + 1)}\gets w^{(r + 1)}\quad //$  layers get embedded into server model
ConfigurationSelection  $(W, r)$ :
\[
\begin{array}{l}
S \gets \text{LookupTable}(r, W) \\
w \gets W \quad // \text{scaling down NN head based on configuration} \\
\text{Return } w, S
\end{array}
\]
DeviceTraining  $(w,S)$  ..
freeze  $w_{j}$  for  $j\in \{0 <   j\leq K_F\}$  according to S   
foreach local mini-batch b do   
|  $w\gets w - \eta \nabla l(w,b)$    
end   
Return  $w$
behavior can be observed w.r.t. Floating Point Operations (FLOPs), where SLT requires  $10 \times$  fewer operations to reach a certain level of accuracy compared with FD and FedRolex.
- We study the performance of SLT in heterogeneous settings. We show that devices with different memory constraints can make a meaningful contribution to the global model, significantly outperforming the state-of-the-art techniques.
2 Methodology
2.1 Problem Statement and Setup
We consider a synchronous cross-device FL setting, where we have one server and a set of devices  $c \in \mathcal{C}$  as participants. There is a given ML model topology  $F$  on the FL server that is trained in a distributed manner for  $R$  rounds. Our goal is to maximize the accuracy of the model. Similar to FD and FedRox, we assume that a fixed number of devices  $|\mathcal{C}^{(r)}|$  out of  $\mathcal{C}$  participate in a round  $r \leq R$ . All devices are constrained in memory, and thus their training must not exceed this given memory constraint  $m_{\mathrm{constraint}}$ . In other terms, we assume that no participating device can train the server NN model end-to-end.
2.2 Successive Layer Training
The following describes our methodology of Successive Layer Training for convolutional neural networks (CNNs). Firstly, we rewrite  $F$  such that it is the consecutive operations of  $K$  layers, where each layer is defined as  $f_{k}$ ,  $k \in [1, \dots, K]$ . Each layer  $f_{k}$  has associated server parameters  $W_{k}$ . We label a convolution, followed by batch normalization and an activation function, a layer. Similar to [9, 20], we define a subset  $w_{k}$  of the layer parameters (server)  $W_{k}$  that is scaled down using  $s$  as
$$
w _ {k} = W _ {k} ^ {s, s} \quad w _ {k} \in \mathbb {R} ^ {\lfloor s P _ {k} \rfloor \times \lfloor s M _ {k} \rfloor} \quad W _ {k} \in \mathbb {R} ^ {P _ {k} \times M _ {k}}, \tag {1}
$$
Figure 2: Visualization of SLT. With each step,  $K_{F}$  and  $K_{T}$  are shifted by 1.  $W_{1},\ldots ,W_{K_{F}}$  denote the parameters that remain frozen during training,  $W_{K_T}$  denotes parameters of layer  $K_{T}$  that are fully trained, while  $W_{K_T + 1}^{1,s},\dots ,W_K^{s,s}$  denote the parameters of the scaled-down head using  $s$ .
where  $P_{k}$  labels the layer's input dimension,  $M_{k}$  labels the output dimension of the fully-sized server parameters, and  $s \in (0,1]$  is a scaling factor (we omit the filter kernel dimensions for brevity). To obey the memory constraint on the participating devices, we split the NN into three consecutive parts. The first part of the NN contains layers that are already trained and remain frozen on the devices. The second part contains layers that are being fully trained on the devices. The last part represents the NN's head. To train the remaining layers within the memory budget, the head's parameters are scaled down. Throughout the FL training, we successively switch the training configuration, s.t., the part of the NN that is being fully trained ( $s = 1$ ) moves from the first layer to the last layer. Thereby, successively, the remaining parameters from the scaled-down head are added. At the same time, we successively freeze more layers, starting with the first layer, to stay within the memory budget. We visualize the switching from one training configuration to the next in Fig. 2. The parts of the NN that are frozen, trained, and represent the head are labeled  $F_{F}$ ,  $F_{T}$ , and  $F_{H}$ . The resulting model that is trained on the devices can be described as  $F = F_{F} \circ F_{T} \circ F_{H}$ :
- The first part  $F_{F}$  labels the part of the NN that is frozen and where the server parameters  $W_{1},\ldots ,W_{K_{F}}$  do not get updated by the devices. Freezing the first part of the NN reduces the memory overhead during training, as in the forward pass, activations do not have to be stored for frozen layers. The frozen part of the NN is defined as  $F_{F}\coloneqq \bigcirc_{k\in \{k:0 < k < K_{F}\}}f_{k}$ , where layers  $1,\dots ,K_{F}$  remain frozen.
- The second part  $F_{T}$  labels the part of the NN that is being fully trained by the devices. The parameters  $W_{K_{F} + 1}, \ldots, W_{K_{T}}$  get updated during training. This part is defined as  $F_{T} := \bigcirc_{k \in \{k: K_{F} < k < K_{T}\}} f_{k}$ , s.t. layers  $K_{F} + 1, \ldots, K_{T}$  are fully trained.
- The last part  $F_{H}$  describes the NN's head that is scaled down using  $s$ , s.t.  $F_{H} := \bigcirc_{k \in \{k: K_{T} < k \leq K\}} f_{k}$ , where the scaled-down layers  $K_{T} + 1, \ldots, K$  are trained. The first layer of  $F_{H}$  scales down the parameters to the width of the head s.t.  $w_{K_{T} + 1} = W_{K_{T} + 1}^{1,s}$ , where  $W_{K_{T} + 1}^{1,s} \in \mathbb{R}^{P_{K_{T} + 1} \times \lfloor sM_{K_{T} + 1} \rfloor}$ . All consecutive layers are scaled down using  $s$ , s.t.  $w_{K_{T} + j} = W_{K_{T} + j}^{s,s} \forall j \in [2, \ldots, K - K_{T}]$ .
We define a set  $S$  as a training configuration that obeys the given memory constraint, s.t.  $S = \{K_F, K_T, s\}$  fully describes the mapping of layers  $f_k, k \in [1, \dots, K]$  into  $F_F, F_T, F_H$ , and the head's scale factor  $s$ . Each  $S$  has a respective memory footprint during training.  $m = \text{memory}(S)$  denotes the maximum memory that is utilized during training for a given configuration  $S$ . The maximum memory of a configuration can be determined by measurements or by calculating the size of weights, gradients, and activations that have to be kept in memory (see Section 3.2).
2.3 Configuration Selection
For each selected  $S$ , we aim to fully utilize the available memory. We define  $n$  as a configuration step in  $n \in [1, \dots, N]$ , where we add parameters to the training (i.e., fill up the remaining parameters of a head's layer). These steps are distributed over the total training rounds  $R$  (Fig. 3). We set  $K_{T} = K_{F} + 1$ , s.t. in each configuration step exactly one layer gets fully trained (filled up). We start with  $K_{F} = K_{T} = 0$  (consequently  $F = F_{H}$ ) to pre-train the head for a certain number of rounds. After pre-training, we increase  $K_{F}$  by one and apply  $K_{T} = K_{F} + 1$  (the first configuration has no
Figure 3: Visualization of the SLT training scheme with an exemplary 5-layer NN. The model is first pre-trained for  $r_p$  rounds. Following that, the model is trained for  $r_1 - r_p$ ,  $r_2 - r_1$ , and  $R - r_2$  rounds in configuration 1, 2, and  $N(= 3)$ , respectively.
frozen layers, i.e.,  $F = F_{T} \circ F_{H}$ ) and continue training $^{2}$ . We switch to the successive configuration by increasing  $K_{F}$  by one. Hence, for the training configuration at step  $n$ , we have  $K_{F} = n - 1$  and  $K_{T} = n$ , with  $s_{n}$  selected as follows:
$$
\max  s _ {n}, \text {s . t .} \quad \operatorname {m e m o r y} \left(S _ {n}\right) \leq m _ {\text {c o n s t r a i n t}} \wedge s _ {n} \leq s _ {j} \quad \forall j \in [ n + 1, \dots , N ], \tag {2}
$$
Equation (2) ensures that each configuration obeys the constraint  $m_{\mathrm{constraint}}$ . The second constraint in Eq. (2) enforces that  $s_n$  can only grow with increasing  $n$  to ensure that parameters of the head are only added throughout the training but not removed. We provide a justification for maximizing  $s$  instead of  $F_T$  by performing an ablation study in Appendix B. The configuration selection is performed offline. Lastly, we define the last step  $N$  where a given memory constraint  $(\mathrm{memory}(S_N))$  allows for  $s = 1$ . If this step is reached, we train with  $S_N$  for all remaining rounds since the memory budget allows to fully train all remaining parameters at once. We provide a visualization of the training process in Fig. 3 and outline SLT in Algorithm 1.
3 Experimental Evaluation
3.1 Experimental Setting and Hyperparameters
We evaluate SLT in an FL setting using PyTorch [21], where we distribute a share from the datasets CIFAR10, CIFAR100 [22], FEMNIST from the Leaf [23] benchmark, and TinyImageNet [24] to each device  $c \in \mathcal{C}$ , s.t. each device  $c$  has a local dataset  $\mathcal{D}_c$  of the same size. In each round  $r$ , a subset of devices  $\mathcal{C}^{(r)}$  is selected. We train with the optimizer stochastic gradient descent (SGD) with momentum of 0.9, an initial learning rate of  $\eta = 0.1$ , and apply cosine annealing to  $\eta = 0.01$  and a weight decay of  $1.0 \times 10^{-5}$ . We evaluate the vision models ResNet20, ResNet44 [25], and DenseNet40 [26]. For each experiment, we report the average accuracy and standard deviation of 3 independent seeds after  $R$  rounds of FL training. For CIFAR10, CIFAR100, and TinyImageNet, we evaluate a scenario with  $|\mathcal{C}| = 100$  devices, where each round  $|\mathcal{C}^{(r)}| = 10$  devices are actively participating. For FEMNIST, we evaluate with  $|\mathcal{C}| = 3550$  and  $|\mathcal{C}^{(r)}| = 35$ . We train for  $R = 2500$  rounds for CIFAR10, CIFAR100, and TinyImageNet, and  $R = 1000$  for FEMNIST. In each round  $r$ , each participating device iterates once over its local dataset. We apply standard image augmentation techniques like random cropping and horizontal and vertical flips to all datasets (horizontal and vertical flips are omitted for FEMNIST). An input resolution of  $3 \times 32 \times 32$  is used for CIFAR and FEMNIST (up-scaled from  $28 \times 28$ ) and  $3 \times 64 \times 64$  for TinyImageNet. We use batch size 32 and perform 363 experiments in total, with an average run-time of  $6\mathrm{h}$  on an NVIDIA Tesla V100.
Comparison with the state of the art: We compare SLT against several baselines. We introduce  $I_k^{(r)}$  as the set of indices of the output dimension of a layer  $k \in [1, \dots, K]$  where  $r \in [1, \dots, R]$  denotes the rounds. Consequently, for full-sized NNs  $|I_k^{(r)}| = M_k$ . The subset for training is scaled down by building a dense matrix using the indices from  $I_k^{(r)}$ , s.t. the scaled-down parameters are  $w_k \in \mathbb{R}^{\lfloor sP_k\rfloor \times \lfloor sM_k\rfloor}$ . The consecutive layer's input dimension indices are equal to the last layer's output indices. The first layer's input dimension is not scaled to feed all color channels into the NN.
Small model: Devices train a submodel where all filters per layer are scaled down by  $s$ , s.t. all devices can train the submodel. The remaining filters are trained end-to-end throughout the rounds.
2We discuss in the evaluation section how to decide the number of rounds a certain configuration should be trained before switching to the next configuration.
The same submodel is used for evaluation. The indices of the output dimension of a layer  $k$  are selected using  $I_k^{(r)} = I_k = \{i : 0 \leq i < \lfloor sM_k \rfloor\}$ .
FedRolex [20]: FedRolex creates a submodel by scaling the numbers of filters using  $s$ . Each device trains the same continuous block of indices. The filter indices trained on the devices are shifted in a rolling window fashion every round. The server averages the trained block and evaluates using the full server model  $(s = 1)$ . The indices  $I_k^{(r)}$  are selected with  $\hat{r} = r \mod M_k$  using
$$
I _ {k} ^ {(r)} = \left\{ \begin{array}{l l} \left\{\hat {r}, \hat {r} + 1, \dots , r + \lfloor s M _ {k} \rfloor - 1 \right\} & \text {i f} \hat {r} + \lfloor s M _ {k} \rfloor \leq M _ {k} \\ \left\{\hat {r}, \hat {r} + 1, \dots , M _ {k} - 1 \right\} \cup \left\{0, \dots , \hat {r} + \lfloor s M _ {k} \rfloor - 1 - M _ {k} \right\} & \text {o t h e r w i s e} \end{array} . \right. \tag {3}
$$
FD [9]: FD creates a submodel by scaling down the number of filters using  $s$ . The indices of the filters are randomly sampled per device per round on the server. Hence, the indices  $I_{k}^{(r,c)}$  of a device  $c$  of round  $r$  is a round-based per-device random selection of  $\lfloor sM_k\rfloor$  indices out of all  $M_{k}$  indices. The server aggregates the device-specific submodels after training and evaluates the full model  $(s = 1)$ .
3.2 Memory Footprint during Training
The high memory requirements during training can be split into three groups: Firstly, the weights of the NN have to be stored in memory. This is required both for the forward pass and the backward pass. Secondly, for the calculated gradients in the backward pass, the activation maps of the respective layers have to be kept in memory. Lastly, the calculated gradients have to be stored in memory. In state-of-the-art CNNs, the size of the activation map makes up for most of the memory requirements, while the size of the weights only plays a minor role. For ResNet44 and DenseNet40, we measure that activations make up for  $\sim 99\%$  of the required memory for training, while gradients and parameters account for the remaining  $1\%$ . Consequently, the required memory linearly reduces with  $s$  for FD and FedRolex, as the number of layer's output channels  $\lfloor sM_k\rfloor$  determines the activation map's size. Similarly, for SLT, we measure the maximum amount of memory that is required during training by counting the size of the activation maps, as well as the loaded weights and gradients in training. For frozen layers, it is only required to load the parameters in memory, while no activation maps and gradients have to be stored. For the fully trained layer  $K_{T}$ , it is required to store the layer's full parameters  $w_{K_T}$ , as well as the full-size activation map and gradients in memory. For all other layers (NN head), memory scales linearly with  $s$ . We provide implementation details in Appendix E.
We evaluate memory constraints that are given by scaling down  $s$  in FedRolex and FD by  $s_{\mathrm{FD/FedRolex}} \in [0.125, 0.25, 0.5, 1.0]$  for experiments with ResNet and [0.33, 0.66, 1.0] for DenseNet<sup>3</sup>. In SLT, for a given  $s_{\mathrm{FD/FedRolex}}$ , we adjust  $s_n$  for each step  $n$  in the following way
$$
\max  s _ {n}, \text {s . t .} \operatorname {m e m o r y} \left(S _ {n} ^ {\prime}\right) \leq \operatorname {m e m o r y} \left(s _ {\mathrm {F D / F e d R o l e x}}\right) \wedge s _ {n} \leq s _ {j} \quad \forall j \in [ n + 1, \dots , N ], \tag {4}
$$
to ensure that our technique obeys the same constraint as the baselines. If  $s_{\mathrm{FD/FedRolex}} = 1.0$ , all algorithms coincide with vanilla Federated Averaging (FedAvg) using the full server model.
We distribute the required steps  $N$  over the total rounds  $R$ , s.t. all parameters receive sufficient training. Specifically, we distribute the rounds based on the share of parameters that are added to the training within a configuration  $n$ . We calculate the number of all trained parameters  $Q$  by using  $K_{F}, K_{T}$ , and  $s$  s.t.
$$
Q \left(K _ {F}, K _ {T}, s\right) = \left(\sum_ {k \in \{k: K _ {F} <   k \leq K _ {T} \}} P _ {k} M _ {k}\right) + P _ {K _ {T} + 1} \lfloor s M _ {K _ {T} + 1} \rfloor + \sum_ {k \in \{k: K _ {T} + 1 <   k \leq K \}} \lfloor s P _ {k} \rfloor \lfloor s M _ {k} \rfloor , \tag {5}
$$
and use  $Q$  to calculate the share of rounds  $R_{n}$  for a step  $n$ . The share of rounds for pretraining is calculated using  $R_{\text{pretraining}} = R_{\frac{Q(0,0,s)}{Q(0,0,1)}}$ . For step 1,  $R_{1} = R_{\frac{Q(0,1,s)}{Q(0,0,1)}} - R_{\text{pretraining}}$ . For all steps  $n > 1$ , we calculate the rounds using
$$
R _ {n} = R \frac {Q (n - 1 , n , s _ {n}) - Q (n - 2 , n - 1 , s _ {n - 1})}{Q (0 , 0 , 1)}. \tag {6}
$$
Lastly, the switching point for pretraining is  $r_{\mathrm{pretraining}} = R_{\mathrm{pretraining}}$  and  $r_n = R\frac{Q(n - 1,n,s_n)}{Q(0,0,1)}$  for all steps  $n$ .
3For DenseNet, SLT only enables a reduction of  $3\times$ , as in DenseNet, specific layers have a significantly larger sized feature map than others, which limits our technique's effectiveness w.r.t memory reduction.