FedOBD: Opportunistic Block Dropout for Efficiently Training Large-scale Neural Networks through Federated Learning

1 Introduction

Over the years, machine learning techniques have been applied to solve a wide range of problems. The size of the learning model matters. Both manually designed neural architectures and those generated through neural architecture search [Elsken et al., 2019] demonstrate that scaling up model sizes can improve performance. However, it is a significant challenge to train a large-scale model, especially in a distributed manner (e.g., involving multiple collaborative organizations [Yu et al., 2017]). Federated Learning (FL) [Yang et al., 2019] is a distributed machine learning paradigm that enables multiple data owners to collaboratively train models (e.g., neural networks (NNs)) without exposing their private sensitive data. In the centralized FL architecture, local model updates are aggregated into a global model in a privacy-preserving manner by a central parameter server [Koneˇcn`y et al., 2016]. FL has been applied in many commercial scenarios such as enhancing user experience [Yang et al., 2018], safety monitoring [Liu et al., 2020; Xie et al., 2022] and healthcare [Liu et al., 2022].

However, training large-scale deep models in current FL settings is challenging. The communication overhead involved is high [Kairouz et al., 2021]. Typical FL approaches require local model updates to be sent to the server and the aggregated global model to be distributed to the clients multiple times. This limits the scale of neural networks that can be used, making them less well-suited for solving complex real-world problems. When multiple institutional data owners collaboratively train large-scale FL models, communication efficient FL training approaches are required.

One popular approach for achieving communication efficient training of large-scale NNs in FL is through parameter dropout. FedDropoutAvg [Gunesli et al., 2021] randomly selects a subset of model parameters to be dropped out to reduce the size of transmission. In addition, it also randomly drops out a subset of FL clients to further reduce the number of messages required. In contrast, Adaptive Federated Dropout (AFD) [Bouacida et al., 2021] performs dropout based on parameter importance by maintaining an activation score map. In such approaches, it is common to apply compression techniques to further reduce the communication overhead when transmitting model updates to and from the FL server. 1 However, these approaches directly manipulate individual model parameters. Recent research shows that this is not only inefficient in terms of meaningfully reducing the communication overhead when training large-scale FL models, but may also negatively affect the scaling efforts and model performance [Cheng et al., 2022].

To address this limitation, we propose the Federated Opportunistic Block Dropout (FEDOBD) approach. By dividing large-scale deep models into semantic blocks, it evaluates block importance (instead of determining individual parameter importance) and opportunistically discards unimportant blocks in order to enable more significant reduction of communication overhead while preserving model performance. Since our block importance measure is not based on the client loss function as is the case for AFD [Bouacida et al., 2021], FEDOBD can handle complex tasks effectively.

1 Note that although Federated Dropout [Wen et al., 2021], FjORD [Horv´ath et al., 2021] and FedDrop [Liao et al., 2021] also leverage parameter dropout in FL, their design goal is to enhance model adaptability to FL client heterogeneity, rather than training large-scale NN models efficiently via FL.

We study the performance of FEDOBD for training largescale deep FL models through extensive experiments in comparison with four state-of-the-art baseline approaches on three real-world datasets (including CIFAR-10, CIFAR-100 and IMDB). The results show that, compared to the best performing existing approach, FEDOBD reduces communication overhead by 88%, while achieving the highest test accuracy. To the best of our knowledge, FEDOBD is the first semantic block importance-based opportunistic FL dropout approach.

2 RelatedWork

Existing methods for improving FL communication efficiency can be divided into two major categories:

Compression. Deep Gradient Compression (DGC) [Lin et al., 2017] employs gradient sparsification to reduce redundant gradients, and thereby enabling large models to be trained more efficiently under FL settings. FetchSGD [Rothchild et al., 2020] takes advantage of Count Sketch to compress gradient updates, while accounting for momentum and error accumulation in the central FL server. Nevertheless, it requires a large number of communication rounds to achieve convergence because aggregation is performed after each local batch. SignSGD [Bernstein et al., 2018] is a typical method for compressing model gradients. However, the compression is static which can result in loss of important features and convergence to a global model with reduced performance.

In recent years, approaches that directly compress FL model parameters are starting to emerge [Amiri et al., 2020; Reisizadeh et al., 2020]. This is more challenging compared to gradients compression as more training information is lost in the process. Nevertheless, these approaches can reduce the number of communication rounds required during FL model training, and can achieve comparable performance with fullprecision networks [Hou et al., 2018]. For example, FedPAQ [Reisizadeh et al., 2020] compresses the model updates before uploading them to the FL server. However, it requires a static quantization approach to be applied before the compression step. Thus, it is only able to support simple learning tasks and small-scale neural networks. Existing compression methods are not suitable for complex FL tasks involving large-scale neural networks.

Parameter Dropout. Dropout methods for FL training have been proposed to address two distinct problems: 1) enabling devices with diverse local resources to collaboratively train a model, and 2) enabling efficient training of a largescale FL model by organizational data owners (which are not resource constrained).

Federated Dropout [Wen et al., 2021] exploits user-server model asymmetry to leverage the diverse computation and communication capabilities possessed by FL clients to train a model which could be too large for a subset of the clients to handle. It fixes the server model size and applies parameter dropout at different rates to generate models suitable for each client to train according to their local resource constraints. FjORD [Horv´ath et al., 2021] extends Federated Droput to propose the ordered dropout method, under which the client sub-models are selected in a nested fashion from the FL server model. In this way, each client is assigned a model of size in proportion to its local computational resources for training, thereby adapting the model according to clients' device heterogeneity. Similarly, FedDrop [Liao et al., 2021] incorporates additional dropout layers into FL models to determine the channel wise trajectories to be dropped or retained to adapt a model according to FL clients' local data distributions. Nevertheless, these approaches are designed to address the first problem, which is not the focus of this paper.

Our research focuses on the second problem. The FedDropoutAvg approach [Gunesli et al., 2021] randomly drops out a subset of model parameters while randomly dropping out some FL clients before performing FedAvg based model aggregation. The Adaptive Federated Dropout (AFD) approach [Bouacida et al., 2021] adaptively determines a percentage of weights based on parameter importance to be dropped. In this way, compression can be performed to reduce the communication overhead of transmitting the model updates to the FL server. It maintains an activation score map which is used to determine the importance of the activations to the training process, and determine which of them shall be kept or dropped in the current round of training. Instead of determining individual parameter importance, the proposed FEDOBD approach focuses on evaluating the importance of semantic blocks in a large-scale deep model and opportunistically discarding unimportant blocks in order to enable more significant reduction of communication overhead while preserving model performance, thereby enabling efficient training of large-scale FL models by institutional data owners.

3 The Proposed FEDOBD Approach

FL is a distributed machine learning paradigm involving multiple data owners (a.k.a., FL clients) to collaboratively train models under the coordination of a parameter server (a.k.a., FL server). 2 Formally, assume that there are n clients participating in FL model training. Each client i owns a local dataset Di = {(xj , yj )}Mi j=1 , where xj is the j-th local training sample, yj is the corresponding ground truth and Mi is the number of samples in Di. Under these settings, FL aims to solve the following optimization problem:

where W is the parameter space determined by the neural network architecture. M := Σni=1Mi is the total number of samples and Li(w;Di) := 1Mi ∑Mi j=1 ℓ(w; xj, yj) is the local loss of client i. Normally, the training process consists

2 Here, we refer to the horizontal federated learning setting with an FL server [Yang et al., 2019].

Algorithm 1: OBD

Input : global model wr-1, local model wr,i in client i, the set of identified block structures B, dropout rate λ ∈ [0, 1]. Output: retained blocks. 1 important blocks ←MaxHeap(); 2 foreach b ∈ B do 3 important blocks[MBD(br-1, br,i)] ← br,i; 4 end 5 revised model size ← 0; 6 retained blocks ←List(); 7 while important blocks do 8 br,i ← important blocks.pop(); 9 new size ← revised model size+|vector(br,i)|; 10 if new size > (1-λ)|vector(wr,i)| then 11 continue; 12 end 13 revised model size ← new size; 14 retained blocks.append(br,i); 15 end 16 return retained blocks;

shown in Algorithm 1. At the end of each round of local training, each block is assigned an importance score. To achieve this goal, a client i keeps a copy of the global FL model wr-1 and compares it with the resulting model wr,i block by block after local training. We propose the Mean Block Difference (MBD) metric to measure block importance, defined as:

where br-1 denotes the blocks in the received global FL model, and br,i denotes the corresponding blocks in the local model produced by client i after the current round of training. The vector operator concatenates parameters from different layers into a single vector and Parameter returns the number of parameters in a block. In general, the larger MBD value between the new and old versions of the same block,

of multiple synchronization rounds. In the beginning of each round, the FL server distributes a global model to selected FL clients. The clients train the received models using local data and send the resulting model updates back to the server. Then, the server combines these uploaded models into an updated global model following some aggregation algorithm (e.g., FedAvg [McMahan et al., 2017]). These steps are repeated until model convergence. FEDOBD is designed for the aforementioned horizontal Federated Learning setting. It attempts to identify unnecessary model updates during training to reduce communication overhead while maintaining model performance. The following subsections explain key components of FEDOBD for achieving this design goal. They are also illustrated in Figure 1.

3.1 Opportunistic Block Dropout (OBD)

Neural networks (NNs) are often designed with some structures in mind. Their generalization performance relies heavily on properly selecting the structures. Motivated by this observation, an NN model can be decomposed into blocks of consecutive layers before FL model training commences. The key idea of FEDOBD is to identify important blocks and only involve them, instead of the entire model, during the FL training process in order to reduce communication overhead. To achieve this goal, we need to address three questions: 1) how to decompose a given model; 2) how to determine block importance; and 3) how to aggregate the uploaded blocks.

When decomposing the trained model, it is critical to recognize popular (i.e., frequently used) NN structural patterns. For example, layer sequences such as ⟨Convolution, Pooling, Normalization, Activation⟩ are commonly found in convolutional neural networks (CNNs), and Encoder layers are commonly found in Transformer based models. Other types of architecture may define basic building blocks which can be taken into consideration when dividing a model into blocks. Such blocks often provide important functionality (e.g., feature extraction). Thus, it makes sense to transmit or drop them as a whole. Finally, the remaining layers are treated as singleton blocks.

The proposed opportunistic block dropout approach is

Algorithm 2: NNADQ

Input: parameter set L, relative weight β. Output: quantized vectors and other parameters for dequantization. 1 results ←List(); 2 foreach layer l ∈ L do 3 results.add(ADQ(l, β)); 4 end

5 return results;

Algorithm 3: ADQ

Input : vector v, relative weight β. Output: a quantized vector and parameters for dequantization. 1 max v, min v←maxmin(v); 2 offset ← argminθ max(|max v+ θ|, |min v+ θ|); 3 v′ ← v+offset; 4 d← ∥v′ ∥∞; 5 sign←sgn(v′ ); 6 s ←int (max (√ln4 ∗ REPR(v) β ∗ d, 1)); 7 vquantized ←Round(v′ , s, d); 8 return ⟨vquantized, offset, d, sign, s⟩;

the more important the new block is because it may contain lots of new knowledge learned in the current round.

Once the MBD values for all the blocks have been computed, FEDOBD determines which blocks to retain (lines 515). Since all the blocks are ranked from high to low in terms of their MBD values with the MaxHeap data structure, the algorithm simply pops each block one by one in descending order of their MBD values and adds each of them to the retained blocks list, while keeping track of the size of the revised model so far. When the algorithm terminates, the size of the revised model is guaranteed to not exceed (1 - λ)|vector(wr,i)|, where λ ∈ [0, 1] is the dropout rate (where 1 indicates dropping out the entire model and 0 indicates no dropout). Only the retained blocks are transmitted to the FL server for model aggregation.

Upon receiving retained blocks from a client i in round r, the FL server combines them with unchanged blocks from the previous global model wr-1 to formwr,i. The reconstructed local models for all clients are then used to aggregate a new global FL model wr following an aggregation algorithm such as FedAvg [McMahan et al., 2017].

3.2 Adaptive Deterministic Quantization For Neural Networks (NNADQ)

OBD reduces the transmission size of the original model by retaining only important blocks based on a user specified dropout rate. To further reduce communication overhead, we propose an improved quantization approach to compress the other blocks before transmission. Quantization operates by converting given float-point values into integers within a fixed range. It is a lossy compression method. In FEDOBD, local

Algorithm 4: FEDOBD

Input : number of FL clients n, client subset size k, dropout rate λ ∈ [0, 1], relative weight β in quantization, number of rounds R in stage 1, first stage number of local epochs E1, second stage number of local epochs E2. Output: the final global FL model. 1 Stage 1: 2 //at the FL server: 3 initialize w0; 4 foreach r ∈ {1, 2, . . . ,R} do 5 Cr ←k randomly chosen clients; 6 distribute NNADQ(wr-1, β) to Cr; 7 //at each client i ∈ Cr: 8 dequantize data, construct and load wr-1; 9 train wr,i for E1 epochs; 10 important blocks ←OBD(wr,i,wr-1, λ); 11 upload NNADQ(Diff(important blocks, old blocks), β) to FL server; 12 //at the FL server: 13 foreach client i ∈ Cr do 14 dequantize the received update and reconstruct wr,i ; 15 end 16 wr ←Aggregation({wr,i}i∈Cr ) 17 end 18 Stage 2: 19 //at the FL server: 20 w0 ←wR 21 foreach epoch e ∈ {1, 2, . . . ,E2} do 22 distribute NNADQ(we-1, β) to clients; 23 //at each client i ∈ {1, 2, . . . , n}: 24 dequantize data, construct and load we-1; 25 we,i ← train with previous learning rate for 1 epoch; 26 upload NNADQ(Diff(we,i,we-1), β) to FL server; 27 //at the FL server: 28 foreach i ∈ {1, 2, . . . , n} do 29 dequantize the received update and reconstruct we,i; 30 end 31 we ←Aggregate(we,1, . . . ,we,n); 32 end

block differences and global models are quantized before being sent out. The following sections describe the proposed NNADQ quantization algorithm in detail.

Stochastic Quantization. In stochastic quantization [Alistarh et al., 2017] for encoding models, the quantization function is denoted as Q(v, s). It takes a vector v and the number of quantization levels s ⩾ 1 as the inputs, and generates a quantized vector. Formally, for any v ∈ Rn , Q(v, s) is:

where the second multiplication is element-wise and ξi(v, s) denotes independent random variables for vector elements, whose definition follows.

Let 0 ⩽ ℓ < s be an integer such that |vi|/∥v∥2 ∈ [ℓ/s, (ℓ + 1)/s]. That is, [ℓ/s, (ℓ + 1)/s] is a quantization interval for |vi|/∥v∥2 . Then,

Reformulation. While the above method is a statistically unbiased way of quantization, it is still unclear whether some choices of variables are optimal. It is also necessary to reformulate the problem taking the balance between compression and informativeness into consideration. Following the above notation, we treat the number of quantization levels s and the norm d as variables. Then, ξi becomes:

Since ξi is an unbiased estimator, Var(ξi) is an indication of quantization loss. However, we quickly find that

is a better alternative since ( |vi| d - Roundi(v, s, d))2 ⩽ Var(ξi), and the new way of quantization works deterministically rather than stochastically.

The trade-off between compression and informativeness can now be formulated as an optimization problem:

The first part of Eq. (3) indicates the mean information loss and the second part indicates the compression ratio where REPR is the number of bits in a floating-point representation 3 . β ∈ R+ is a predefined relative weight between the two parts.

3 REPR = 32 in typical environments.

Since dl ⩽ s|vi| ⩽ dl + d, by definition, we have

Eq. (4) shows that d should be as small as possible and d = ∥v∥∞ is a near optimal choice for any fixed d. For any fixed d, we allow s to take real values and combine Eq. (3) and Eq. (4) to obtain a simplified optimization objective

⌊s∗R⌋ is then taken as an approximate solution to the original problem Eq. (3).

Additional Optimizations. Note that d = ∥v∥∞ affects the information loss directly and the compression ratio indirectly via s∗R. Smaller values of ∥v∥∞ can improve Eq. (3) in general. Thus, we should translate v to v′ with minimum infinity norm before quantizing it. Since each layer of a largescale model has a different statistical parameter distribution, quantization layer by layer can utilize this trick further.

Based on the aforementioned quantization approach, we propose the Adaptive Deterministic Quantization for Neural Networks (NNADQ) approach as in Algorithm 2 (and the related supporting function in Algorithm 3). NNADQ is used in FEDOBD to compress both directions of transmission.

3.3 The Two-Stage Training Process

Finally, FEDOBD has a two-stage training process. In the first stage, small local epochs are used, random subsets of clients are selected in each round and OBD is activated (i.e., only selected blocks are uploaded to the server). In this way, FEDOBD encourages frequent aggregation to prevent overfitting without incurring high communication overhead.

In the second stage, FEDOBD switches to a single round with aggregation being executed at the end of each epoch so that local learning rates are reused. Therefore, FEDOBD attempts to fine-tune the global model by approximating centralized training. The detailed process of FEDOBD is presented in Algorithm 4.

4 Experimental Evaluation

In this section, we evaluate the performance of FEDOBD against four state-of-the-art related approaches. Next, we conduct ablation studies to investigate the impact of different components of FEDOBD on its effectiveness.

4.1 Experiment Settings

We compare FEDOBD with these approaches:

1. FedAvg [McMahan et al., 2017]: A classic FL model training approach used as the baseline.
2. FedPAQ [Reisizadeh et al., 2020]: A compressionbased communication efficient FL model training approach, which utilizes stochastic quantization [Alistarh et al., 2017] and random client selection.
3. Adaptive Federated Dropout [Bouacida et al., 2021]: An FL approach that optimizes both server-client communication costs and computation costs by allowing clients to train locally on a selected subset of the global model parameters. We adopt the Single-Model Adaptive Federated Dropout (SMAFD) variant for comparison.
4. FedDropoutAvg [Gunesli et al., 2021]: An FL approach that randomly drops out a subset of model parameters, while randomly dropping out some clients before performing FedAvg based model aggregation.

Hardware&Software. The experiments are carried out on a server with 4 NVIDIA A100 GPUs, 1 AMD EPYC CPU and 252 GB of memory. The algorithm and experiment implementations are based on the PyTorch [Paszke et al., 2017] framework.

Tasks. We use two image classification tasks on CIFAR10, CIFAR-100 [Krizhevsky and Hinton, 2009] and one sentiment classification task on IMDB [Maas et al., 2011]. CIFAR-10 consists of 50,000 training images and 10,000 testing images in 10 classes. CIFAR-100 consists of 50,000 training images and 10,000 testing images in 100 classes. IMDB consists of 25,000 highly polar movie reviews for training, and 25,000 reviews for testing in two classes. In this work, we focus on the i.i.d. case. Local training and validation datasets are drawn uniformly from the larger datasets and the server holds a separate test dataset for evaluating.

Models. For CIFAR-10 and CIFAR-100 tasks, we use two DenseNet-40 networks with about 0.17 and 0.19 million parameters respectively [Huang et al., 2017]. For IMDB, we use a Transformer based classification network consisting of 2 encoder layers followed by a linear layer with a total of 17 million parameters [He et al., 2016]. In addition, word embeddings are initialized with GloVe word embeddings [Pennington et al., 2014].

Experiment Scale. We investigate the performance of FEDOBD and other approaches under large-scale settings. Specifically, 100 clients are used in CIFAR-10 and IMDB tasks and 50 clients are used in CIFAR-100 task. Furthermore, a fixed fraction of 50% clients are chosen randomly in each round in FedPAQ, FedDropoutAvg, SMAFD and FEDOBD since random client selection is part of these algorithms.

Independent Trials. To measure the stability of FEDOBD and other approaches that use probabilistic methods, 10 independent trials are performed for each task.

Other Hyperparameters. 100 rounds and 5 local epochs are used in all combinations of algorithms and tasks for fair comparison. In addition, FEDOBD uses 10 epochs in the second stage. Initial learning rates of 0.1 and 0.01 are used in image and sentiment classification tasks respectively and adjusted by Cosine Annealing learning rate strategy [Loshchilov and Hutter, 2016]. The batch size is set to 64. FedPAQ uses a fixed quantization level of 255. FedDropoutAvg, SMAFD and FEDOBD use a dropout rate of 0.3. FEDOBD uses a quantization weight β of 0.001 in CIFAR-10/CIFAR-100 tasks and 0.0001 for IMDB.

4.2 Results And Discussion

Table 1 lists mean and standard deviation values of performance and communication overheads over 10 independent trials for each combination of tasks and approaches. In each case, "Data Transmission" lists the amount of actual data transmission in megabytes and "Test Accuracy" evaluates the resulting model performance on the test dataset.

It can be easily calculated that, without parameter dropout/compression, FedAvg requires around 20, 100 messages to be exchanged between the FL server and the clients, with 100% of the model parameters being transmitted all the time in order to reach model convergence. Therefore, it achieves the highest test accuracies in large-scale training cases. FedPAQ, FedDropoutAvg and SMAFD exchange fewer messages due to parameter dropout. Overall, they achieve low communication overhead, with SMAFD consistently achieving the lowest overhead among them.

Compared to the above approaches, FEDOBD requires the lowest data transmission to achieve convergence. Notably, the communication overhead required by FEDOBD is 88% lower on average than that required by FedAvg. Meanwhile, FEDOBD achieves the highest or second highest test accuracies, comparable to the performance of FedAvg. Overall, FEDOBD achieves the most advantageous trade-off between communication overhead and test accuracy among all comparison approaches.

4.3 Ablation Studies

Since FEDOBD consists of three key components, we need to measure their effect in isolation. Specifically, we consider the following variants of FEDOBD:

1. FEDOBD with SQ: In this variant, NNADQ is replaced with stochastic quantization [Alistarh et al., 2017].
2. FEDOBD w/o 2nd Stage: This variant only goes through the first stage of the training process.
3. FEDOBD w/o Block Dropout: The dropout rate λ is set to 0 in this variant of FEDOBD.

Table 2 lists mean and standard deviation values of performance and communication overheads over 10 independent trials for each variant under the CIFAR-100 task with the same hyperparameter settings as in the previous experiments. Although "FEDOBD with SQ" achieves similar test accuracy compared to the canonical FEDOBD method, it incurs 20.37% higher communication overhead. This increase in communication overhead is as a result of the transmission amount rising from 811.79MB under FEDOBD to 977.15MB under "FEDOBD with SQ". This demonstrates the advantage of the proposed NNADQ approach in terms of reducing communication overhead.

When training without the second stage under "FEDOBD w/o 2nd Stage", an inferior global FL model is obtained which achieves 3% lower test accuracy compared to the canonical FEDOBD method. Nevertheless, without going through the second stage of training, it incurs 16.99% lower communication overhead (as a result of less message exchanges) compared to the canonical FEDOBD method. This shows that the proposed two-stage training scheme is necessary for producing high quality final models.

When trained under "FEDOBD w/o Block Dropout", the resulting model achieves higher test accuracy compared to the canonical FEDOBD method (i.e., 2.7% increase in test accuracy). However, it incurs 14.10% higher communication overhead due to the transmission amount increasing from 811.79MB under FEDOBD to 926.23MB. Hence, opportunistic block dropout is helpful in reducing communication overhead while preserving model performance.

Through the ablation studies, we have demonstrated that the algorithmic components of FEDOBD are indeed indispensable towards achieving its design goal.

5 Conclusions and FutureWorks

In this paper, we set out to address an emerging research topic in the field of federated learning which is how to efficiently train large-scale deep models. This problem is commonly found in industrial FL application settings. We propose FEDOBD: a first-of-its-kind semantic block-level importancebased opportunistic dropout approach for improving FL model training efficiency, while maintaining model performance. Extensive experimental evaluation demonstrates that FEDOBD outperforms state-of-the-art baselines in terms of communication overhead and test accuracy.

We provide the reference implementation of FEDOBD and related experiments in an open sourced project 4 for academic research. In the future, we plan to integrate more privacypreserving considerations into the current FEDOBD implementation to enable institutional data owners to collaboratively train complex large-scale FL models efficiently.

Acknowledgments

This research/project is supported, in part, by the National Research Foundation Singapore and DSO National Laboratories under the AI Singapore Programme (AISG Award No: AISG2-RP-2020-019); Nanyang Technological University, Nanyang Assistant Professorship (NAP); the National Key R&D Program of China (No. 2021YFF0900800); the Joint SDU-NTU Centre for AI Research (C-FAIR); and the RIE 2020 Advanced Manufacturing and Engineering (AME) Programmatic Fund (No. A20G8b0102), Singapore.

4 https://github.com/cyyever/distributed learning simulator.

Contribution Statement

The main contribution to this work is equally given by Yuanyuan Chen and Zichen Chen.