SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Irene Wang $^{1,3}$ , Prashant J. Nair $^{1}$ , Divya Mahajan $^{2,3}$
University of British Columbia<sup>1</sup>, Microsoft<sup>2</sup>, Georgia Institute of Technology<sup>3</sup>
irene.wang@gatech.edu prashantnair@ece.ubc.edu divya.mahajan@gatech.edu
Federated Learning (FL) allows machine learning models to train locally on individual mobile devices, synchronizing model updates via a shared server. This approach safeguards user privacy; however, it generates a heterogeneous training environment due to the varying performance capabilities of devices. As a result, "straggler" devices with lower performance often dictate the overall training time. In this work, we aim to alleviate this performance bottleneck due to stragglers by dynamically load balancing the training across the system. We introduce Invariant Dropout, a method that extracts a sub-model based on the weight update threshold, thereby minimizing potential impacts on accuracy. Building on this dropout technique, we develop an adaptive training framework, Federated Learning using Invariant Dropout (FLuID). FLuID offers a lightweight framework for sub-model extraction to regulate the computational intensity, thereby reducing the load on straggler devices without affecting model quality. Our method leverages neuron updates from non-straggler devices to construct a tailored sub-model for each straggler based on client performance profiling. Unlike prior work, FLuID can dynamically adapt to changes in stragglers as runtime conditions shift. We evaluate FLuID using five real-world mobile clients. The evaluations show that Invariant Dropout maintains baseline model efficiency while alleviating the performance bottleneck of stragglers through a dynamic and lightweight runtime approach.
Federated Learning (FL) enables machine learning models to train on edge and mobile devices, synchronizing the model updates via a common server [SB09, SVHN10, LSTS20, ZBW+22]. This method ensures user privacy as local data remains on the device, minimizing the risk of data breaches, leaks, or unauthorized access to sensitive information in the cloud  $\mathrm{[MMR^{+}17]}$ . However, FL introduces heterogeneity due to varying performance capabilities of the participating devices. Consequently, straggler devices with lower computational and network performance often dictate the overall training latency and throughput, as shown in Figure 1.
Previous research mitigates the impact of stragglers through asynchronous aggregation, where clients can communicate and update the server model independently and asynchronously [CCA $^{+}$ 21, CNSR20, XKG19]. While this approach alleviates some of the detrimental effects of stragglers, it can introduce staleness into the global model. This occurs because the gradients used to update the server model may rely on outdated parameters, resulting in inaccuracies that have the potential to slow down convergence and reduce overall accuracy [BHS20, CCA $^{+}$ 21, WUH $^{+}$ 21, HBGS19].
Other research proposes eliminating updates from slower devices entirely.  $\mathrm{[KMA^{+}19]}$ . However, this approach can introduce training bias, as it effectively excludes certain clients and their respective data.
Source code is available at https://github.com/iwang05/FLuID
Figure 1: Straggler's impact on FL performance. In synchronous FL, all clients, including stragglers, participate in global model aggregation.
To ensure the contribution of all devices, recent works in this area employ a dropout technique where the stragglers only train a subset of the global model  $\left[\mathrm{HLA}^{+}21, \mathrm{MGZP}22\right]$ . The overall accuracy of the model is determined by the subset of neurons that are dropped from the global model. Thus, previous work alleviates the training load on stragglers by either incurring training bias, creating performance-centric sub-models, or entirely reconstructing the sub-model.
Our paper proposes a novel dropout technique called Invariant Dropout, which identifies "invariant" neurons—those that train quickly and show little variation during training. We observe that the training (computational load) and transfer (communication load) of these invariant neurons to and from the straggler devices contribute minimally to the efficiency of the global model. Thus, they can be dropped. We observe that after only  $30\%$  of the training iterations,  $15\%-30\%$  of the neurons become invariant across CIFAR10 [Kri09], and LEAF datasets FEMNIST, and Shakespeare  $\left[\mathrm{CDW}^{+}18\right]$ . Building on this insight, we develop a dynamic framework called Federated Learning using Invariant Dropout (FLuID), which adjusts the sub-model size based on both the magnitude of neuron updates and the computational capabilities of the client devices. In addition to introducing a new dropout technique, FLuID, unlike prior work, periodically calibrates the sub-model size during runtime to account for changes in stragglers due to factors like low battery or network issues.
Overall, we face two challenges in building FLuID — identifying invariant neurons and establishing a dynamic lightweight method for straggler identification and sub-model size determination at runtime. FLuID addresses the first challenge by leveraging the non-straggler clients, which typically outnumber stragglers and train on the entire model, to identify invariant neurons. The server is not used for this purpose as it receives updates from stragglers running a sub-model. For the second challenge, FLuID uses a drop-threshold, below which neurons are dropped, allowing dynamic sub-model determination for each straggler. By profiling client training times and identifying stragglers during the initial epochs of each calibration step, FLuID can incrementally adjust the drop-threshold until the number of selected neurons matches the target sub-model size for stragglers.
Our evaluation of FLuID on various models, datasets, and real-world mobile devices shows an up to  $18\%$  speedup in performance. Additionally, it improves training accuracy by a maximum of 1.4 percentage points over the state-of-the-art Ordered Dropout, all while mitigating the computational performance overheads caused by straggler devices.
In the domain of heterogeneous client optimization, several prior work have tried to mitigate the effects of stragglers.
Dropout techniques. Federated Dropout [CKMT18, RKPH22] randomly drops portions of the global model, sending a sub-model to the slower devices. However, this can potentially impact accuracy. Subsequent work, Ordered Dropout  $\left[\mathrm{HLA}^{+}21\right.$  DDT21], mitigates this accuracy loss by regulating which neurons are dropped, either from the left or right portions of the global model. Other methods, such as those outlined in [MGZP22], use a low-rank approximation to identify over-parameterized neurons and generate tailored sub-models. Despite these advancements, none of these works consider the individual contribution of each neuron while creating a sub-model. In contrast, our work not only introduces a novel dropout technique that takes into account the contribution of each neuron but also
provides a framework capable of dynamically identifying slower devices and adjusting the dropout rate accordingly.
Server offloading strategies using split learning. The approach proposed by  $\left[\mathrm{WUH}^{+}21\right]$  employs split learning to offload part of the model to the server, while [HAA20] transfers their knowledge to a larger server-side CNN.  $\left[\mathrm{UWH}^{+}21\right]$  focuses on device mobility during Federated Learning, offloading training to edge servers. Finally,  $\left[\mathrm{WQR}^{+}22\right]$  introduces a compression method for split learning. However, unlike these methods, Invariant Dropout does not require data transfer out of the device and into the server, instead conducts all training on client devices. Nonetheless, if data movement was a possibility, Invariant Dropout can be stacked on top of these techniques to determine which part of the model needs to be offloaded to the server.
Communication Optimizations. To reduce learning time, [HWL20] proposes an online learning approach that reduces the communication overheads of Federated Learning by adaptively sparsifying gradients.  $\left[\mathrm{CSP}^{+}21\right]$  introduces a Federated Learning framework utilizing a probabilistic device selection method to improve FL convergence time.  $\left[\mathrm{XFD}^{+}21\right]$  proposes a framework incorporating overhead reduction techniques for efficient training on resource-limited edge devices, including pruning and quantization. Unlike these works, Invariant Dropout avoids both communication and computation overhead by dropping the least contributing "invariant" neurons, presenting a new insight compared to these prior works.
Model Pruning. PruneFL  $\left[\mathrm{JWK}^{+}22\right]$  introduces an approach to dynamically select model sizes during FL and reduce communication and computation overhead, thereby minimizing training time. Work in  $\left[\mathrm{MSA}^{+}21\right]$  prunes the global model for slower clients by excluding neurons with no or small activations. These approaches either generate a single sub-model for all clients, including non-stragglers, to train on or generate a static sub-model for the entire training process which can result in the stragglers permanently losing the opportunity to contribute to certain parts of the model. In contrast, Invariant Dropout enables the generation of multiple sub-models with various resource budgets and dynamically chooses a sub-model for each round based on current neuron contribution.
Coded Federated Learning. CodedPaddedFL and CodedSecAgg [SKRA21] utilize coding strategies to mitigate the impact of slower devices. By sharing an encoded version of their data with other clients, both schemes introduce redundancy in the client's local data. Therefore, during training, the computation of a subset of clients is sufficient to train the global model and the computations of straggling clients can be discarded without loss of information. While CodedPaddedFL combines one-time padding with gradient codes, codedSecAgg is based on Shamir's secret sharing. Unlike these methods, Invariant Dropout allows each client to keep their data local.
Training a Global Family of Models. SuperFed [KALT23] proposes co-training of multiple models to reduce training costs. It achieves this by sending subnetworks of different sizes to all clients. In contrast, FLuID focuses on optimizing the performance of stragglers by mostly training clients on the full global model and a smaller percentage (stragglers) on sub-models tailored to their capabilities.
3 Background and Motivation
3.1 Challenges in Federated Learning
In Federated learning's synchronous aggregation protocols, the server hosts a global model and regularly aggregates updates from clients, which are then redistributed to clients  $\left[\mathrm{CCA}^{+}21\right]$ . Thus, system heterogeneity remains a significant challenge in FL  $\left[\mathrm{DTS}^{+}22, \mathrm{HLA}^{+}21, \mathrm{LSTS}20\right]$ . As shown in Figure 2a, we observe significant differences in training times across five Android-based mobile phones, from 2018 to 2020, engaged in Federated Learning without any dropout mechanism in place. Depending on the dataset and machine learning model, training time can vary significantly, with differences of up to twofold across datasets such as CIFAR10, FEMNIST, and Shakespeare. The standard deviation between the training times of each client is 0.5, 22, and 21 seconds for FEMNIST, CIFAR10, and Shakespeare, respectively. This highlights the real-world challenge of system heterogeneity, where different devices, even if they are of a similar class (e.g., Android-based mobile phones) with just a few years' difference, can offer dramatically different computational performance. Furthermore, we find the performance of mobile devices can fluctuate over time due to varying network bandwidth and resource availability. In response to these observations, we develop
(a) The per-epoch training time of client devices in log scale.
(b) Compared accuracy of Ordered Dropout with baseline implementation of Federated Learning.
Figure 2: (a) Performance variation across mobile devices (in log scale), (b) accuracy implications of prior dropout (static) techniques.
FLuID, a system designed to dynamically identify stragglers and adjust the training load to balance performance with non-straggler devices.
3.2 Dropout Techniques
Model dropout is a technique that involves sending a subset of the global model, known as a sub-model, to stragglers for load balancing. There are two state-of-the-art proposals in this space: Federated Dropout [CKMT18] and Ordered Dropout from FjORD  $\left[\mathrm{HLA}^{+}21\right]$ . Federated Dropout randomly drops neurons, simplifying the selection of sub-models, but reducing the global model's accuracy. In contrast, Ordered Dropout systematically drops neurons by maintaining order within the sub-model. These works demonstrate that the sub-model's neuron selection is crucial for the global model's accuracy. In the context of this work, "neurons" refer to filters in convolutional (CONV) layers, activations in fully-connected (FC) layers, and hidden units in Long Short-Term Memory (LSTM) [HS97] layers.
Accuracy Implications with dropout: Figure 2b compares the testing accuracy of a non-dropout vanilla Federated Learning (FL) implementation with Ordered Dropout. This comparison is performed using five mobile devices, one of which is a straggler, across three different datasets: CIFAR10 [Kri09], FEMNIST, and Shakespeare  $\left[\mathrm{CDW}^{+}18\right]$ . As the sub-model size decreases, we observe that across all three datasets and machine learning models, Ordered Dropout experiences up to a 2.5 percentage point drop in accuracy. We vary the sub-model size from 0.5 (representing  $50\%$  of the global model) to 1 (the entire model). The results with 50-100 clients and  $20\%$  of them being stragglers are provided in Section 6.1.
4 Invariant Dropout
Invariant Dropout enables stragglers to only train on a sub-model consisting of neurons that 'vary' over time and contribute to the global model. Invariant dropout achieves this by selecting a subset  $(i)$  of sub-models  $(s_i)$  from a total sub-model distribution  $(S)$ , where each sub-model represents a different number of neurons and thus varies in compute and memory requirements. Appendix A.1 illustrates the temporal variation in the percentage of invariant neurons in the evaluated models. Note that all clients including the straggler still perform inference on the full model.
4.1 Proposed dropout mechanism
Invariant Dropout selects a sub-model from a set of sub-models,  $S$ , composed of  $m$  total sub-models, denoted as  $S = s_1, s_2, \ldots, s_m$ . Each sub-model  $s_i$  contains neuron layers  $a_1, a_2, \ldots, a_k$ . The size of the sub-model corresponds to a dropout rate  $r$  that is based on the drop-threshold  $th$  for the dropout mechanism. Given a change in neuron value,  $g = \Delta a$ , we can select the sub-model  $s_i = \{a_1, a_2, \ldots, a_k\}$  where the updates in neurons within the sub-model satisfy  $g \geq th \forall a_j \in s_i$  and  $1 \leq j \leq k$ . FLuID, which we describe below, selects  $s_i$  such that the compute utilization and memory footprint can effectively mitigate the inefficiency of the straggler.
Next, we assume that the system includes  $C$  clients with  $T$  stragglers and  $N$  non-stragglers.  $T$  and  $N$  are exclusive and independent subsets of  $C$ , where  $T \cup N = C$  and  $T \cap N = \phi$ . Invariant Dropout maintains a dropout rate  $r \in (0,1]$  per layer across the  $T$  stragglers. However, the complexity of selecting  $a_{j}(t + 1)$  will vary with each sub-model  $s_i$ . To reduce the computational overhead of selecting the appropriate sub-model, ID leverages non-straggler clients to provide directions on the set of  $a_{j}(t + 1)$ . The server computes over a subset of potential sub-models  $s_i$  to select the one with the maximum updates to the neurons.
4.2 Variance in Gradients with Invariant Dropout
Consider  $C$  clients participating in Federated Learning. The training data is represented as  $x_{c}c = 1^{C}$ , where  $c$  denotes one client, and each client has a corresponding loss function  $f c_{c=1}^{C}$ . Learning minimizes the loss function using the following optimization:  $f(w) := \frac{1}{C}\sum c = 1^{C}f_{c}(w)$ , where  $w_{t+1} = w_{t} - \eta_{t}(g(w_{t}))$ .
Invariant Dropout can be viewed as a sparse stochastic gradient vector, where each gradient has a certain probability of being dropped. The gradient vector is denoted as  $G = [g_{1}, g_{2}, \dots, g_{k}]$  where  $g \in R$  and each gradient has a probability of being retained and transmitted across the network, represented as  $[p_{1}, p_{2}, \dots, p_{k}]$ . The sparse vector for the stragglers is represented as  $G_{s}$ . The variance of the ID-based gradient vector can be represented as  $E(G_{s}^{2}) = \sum_{i=1}^{k} (g_{i}^{2} p_{i})$  [WWLZ18, XYXC21]. The variance of the dropout vector is a small factor deviation from the non-dropout gradient vector represented as follows:
$$
\min  \sum_ {i = 1} ^ {k} p _ {i}: \sum_ {i = 1} ^ {k} \left(g _ {i} ^ {2} p _ {i}\right) = (1 + \epsilon) \sum_ {i = 1} ^ {k} g _ {i} ^ {2} \tag {1}
$$
Invariant Dropout drops weights based on the  $g = \Delta a$  gradient across epochs. Thus, the probability of a gradient not getting dropped  $(p_i)$  is inversely proportional to the dropout rate  $r$ . This implies if  $|g_i| > |g_j|$  then  $p_i > p_j$ . Let's assume that the top-k magnitude of gradients are not dropped. Hence if  $G = [g_1, g_2, \dots, g_m]$  is sorted, then  $G = [g_1, g_2, \dots, g_k]$  has a  $p = 1$ , whereas  $G = [g_{k+1}, g_{k+2}, \dots, g_d]$  has a probability of  $p_i = \frac{|g_i|}{r}$ . This modifies the optimization problem in Equation 1 to:
$$
\sum_ {i = 1} ^ {k} g _ {i} ^ {2} + \sum_ {i = k + 1} ^ {m} \frac {\left| g _ {i} \right|}{r} - (1 + \epsilon) \sum_ {i = 1} ^ {m} g _ {i} ^ {2} = 0, \frac {\left| g _ {i} \right|}{r} \leq 1 \tag {2}
$$
Which implies that
$$
r = \frac {\sum_ {i = k + 1} ^ {m} \left| g _ {i} \right|}{(1 + \epsilon) \sum_ {i = 1} ^ {m} g _ {i} ^ {2} - \sum_ {i = 1} ^ {k} g _ {i} ^ {2}} \tag {3}
$$
As per the constraint  $\frac{|g_i|}{r} \leq 1$ :
$$
\left| g _ {i} \right| \left(\sum_ {i = k + 1} ^ {m} \left| g _ {i} \right|\right) \leq (1 + \epsilon) \sum_ {i = 1} ^ {m} g _ {i} ^ {2} - \sum_ {i = 1} ^ {k} g _ {i} ^ {2} \tag {4}
$$
Invariant Dropout retains the gradients with the greatest magnitude. As a result, the boundedness of the expected value from Equation 1 can be represented as follows:
$$
\sum_ {i = 1} ^ {m} p _ {i} = \sum_ {i = 1} ^ {k} p _ {i} + \sum_ {i = k + 1} ^ {m} p _ {i} \tag {5}
$$
$$
\sum_ {i = 1} ^ {m} p _ {i} = k + \left| g _ {i} \right| \sum_ {i = k + 1} ^ {m} \left(\frac {\sum_ {i = k + 1} ^ {m} \left| g _ {i} \right|}{(1 + \epsilon) \sum_ {i = 1} ^ {m} g _ {i} ^ {2} - \sum_ {i = 1} ^ {k} g _ {i} ^ {2}}\right) \tag {6}
$$
$$
\sum_ {i = 1} ^ {m} p _ {i} \leq k (1 + \epsilon) \tag {7}
$$
As such, the variance of the gradient in ID is bounded by Equation 7.
Figure 3: The workflow of FLuID. The non-stragglers are used to determine the neurons that are not updated within a set threshold. Thereafter, sub-models are dynamically created by dropping invariant neurons. These sub-models are sent to the straggler devices.
5 FLuID Framework
Figure 3 shows the workflow of Federated Learning using Invariant Dropout (FLuID). In FLuID, each calibration step includes straggler determination  $(T)$ , discovering invariant neurons  $(IN)$  and drop threshold  $(th)$ , and sub-model extraction  $(s_i)$ . Currently, the calibration occurs per epoch, i.e., one training run over the complete client dataset. However, the frequency of calibration can be reduced if the invariant neurons and stragglers do not significantly change over steps.
Algorithm 1 outlines the FLuID framework. At the onset of training, FLuID identifies the stragglers. To achieve this, FLuID runs the global model on all clients, including stragglers, and measures the performance delay between the end-to-end training time of the slowest client ( $T_{straggler}$ ) and the target time ( $T_{target}$ ). End-to-end training includes upload/download latency and communication time.
$T_{target}$  is the desired training time FLuID aims to achieve for stragglers. FLuID assigns  $T_{target}$  as the next slowest client's training time. This choice optimizes non-straggler idle time reduction. Note, FLuID can support any  $T_{target}$  value. Setting  $T_{target}$  lower than the next-slowest client's time offers no gain as non-stragglers cannot accelerate. Conversely, setting  $T_{target}$  above the next-slowest client's training time leads to longer idleness and suboptimal performance. The required speedup for stragglers is calculated as  $Speedup = \frac{T_{straggler}}{T_{target}}$ . The initial calibration is carried out in lines 6-9 and 18-22 of the algorithm. Note, it takes a few epochs to calibrate the initial threshold, stragglers, and submodel. In subsequent calibration steps, the global server continues to measure the training time of clients and thereby identifies if there is any change in the straggler cohort.
Tuning the performance of stragglers. The sub-model size is determined by calculating the dropout rate  $r$  as detailed in lines 18-21 of Algorithm 1. The value of  $r$  is selected to ensure that the updated straggler training time, denoted as  $T_{stragglernew}$ , is close to the target training time ( $T_{target}$ ). Figure 7 in Appendix A.3 demonstrates that across all evaluated datasets, the training time of all five mobile clients decreases linearly as the sub-model size decreases, and falls within  $10\%$  of the sub-model size. Using this insight, FLuID chooses an  $r$  that is closest to the inverse of the speedup.
Determining the drop threshold at runtime. Once FLuID has identified the stragglers and the dropout rate, it iteratively adjusts the threshold to drop as many invariant neurons as possible. The design of FLuID is inspired by the preliminary results regarding the characteristics of invariant neurons and their impact on the model accuracy. In Appendix A.2, we present results to quantify the impact of the threshold value on the number of invariant neurons during training. We observe that, in order to obtain the desired accuracy, it is critical to select a threshold that yields a number of invariant neurons as close as possible to the number of neurons to be dropped for the sub-model.
Lets assume  $w_{ijc}(t)$  represents the set of weights of the  $i$ th neuron in the  $j$ th layer for client  $c$  after the training epoch  $t$ . The percent difference  $g$  of the neuron, for each client  $c$ , is the minimum value of  $g$  as denoted by:  $g \geq \frac{w_{ijc}(t) - w_{ij}(t-1)}{w_{ij}(t-1)}$ . Neurons that are potential candidates for dropping are those whose weight updates are within the threshold  $(th)$  compared to the previous calibration point. Unfortunately, determining the appropriate  $th$  poses a few challenges.