Reducing Communication for Split Learning by Randomized Top-k Sparsification

1 Introduction

In recent years, the protection of data privacy has become an important issue in machine learning. To date, many kinds of solutions have been proposed to solve the data privacy issue. Aside from cryptographic methods [Mohassel and Rindal, 2018; Wagh et al., 2019; Chen et al., 2021a], Federated Learning (FL) [Li et al., 2020] and Split Learning (SL) [Vepakomma et al., 2018] are two promising methods for privacy-preserving machine learning. While typical federated learning (also known as Horizontal Federated Learning, HFL) mainly focuses on horizontally distributed data, split learning is a simple and effective solution for vertically distributed data. However, although split learning achieves better efficiency than cryptographic methods [Chen et al., 2022], the communication efficiency is still an important issue. In this paper, we focus on reducing the communication of split learning.

The basic idea of split learning is to divide the model into several parts, and each part is computed by a different participant (party). Since split learning only requires sharing the intermediate outputs or gradients, instead of the original input and label, it is considered somewhat privacy-preserving. We present the overview of a typical split learning model in Figure 1, where we consider a two-party case - input feature and label are held by two parties, i.e., feature owner and label owner. To train a model without directly revealing participants' data, the model is split into a bottom model and a top model, held by the feature owner and the label owner, respectively. The last layer (output) of the bottom model is called the cut layer. We simply describe the training procedure of split learning below:

∗Corresponding author.

• Forward pass: the feature owner first feeds the sample features X to the bottom model Mb, then gets the intermediate output Ob = Mb(X, θb), where θb is the bottommodel's weight. The label owner fetches Ob, feeds it to the top model Mt, and then gets the final output ˆ Y =Mt(Ob; θt).
• Backward pass: the label owner calculates the loss function L(ˆ Y , Y ), then computes the gradients ∂L( ˆ Y , Y )/∂θt and Gb = ∂L( ˆ Y , Y )/∂Ob. The former gradient is used to update the top model, and the latter is sent to the feature owner. The feature owner computes the Jacobian matrix ∂Ob/∂θb. Then the gradient with respect to the bottom model's parameters can be computed as (∂Ob/∂θb)T Gb.

However, despite its simplicity, the communication efficiency of split learning is still an important issue. During each iteration, the participants have to exchange intermediate results of the data batch, and the number of iterations and the size of intermediate results are usually large. For example, consider training ResNet-20 [He et al., 2016] with the cut layer size equal to 32×32×32 and batch size equal to 32. It takes 2 (forward & backward) × 4 (bytes for float value) × 32 (batch size) ×323 (size of a single sample's intermediate result) bytes = 8 Mib traffic to finish a single iteration of inference, which is a huge communication cost, especially for mobile devices.

For federated learning, existing work has investigated different methods to reduce communication, including increasing local epochs [McMahan et al., 2017], accelerating convergence [Karimireddy et al., 2020; Wang et al., 2020a; Wang et al., 2020b], and compressing local model updates in various ways [Aji and Heafield, 2017; Stich et al., 2018; Sattler et al., 2019b; Sattler et al., 2019a; Wang et al., 2018]. For split learning, however, studies on communication reduction are relatively few. Castiglia et al. [2022] proved the convergence of compressing split learning by top-k sparsification or quantization. Other methods such as asynchronous training, quantization, and using an autoencoder as a compressor are also investigated for specific tasks [Chen et al., 2021b; Ayad et al., 2021].

In our work, we focus on reducing the communication for the classification problems with a large number of classes, and propose a modification of top-k sparsification, namely, randomized top-k sparsification (RandTopk). To perform randomized top-k sparsification on a d-dimensional vector, we first determine its top-k elements. Then we randomly select k distinct elements such that: with probability 1-α, the element is chosen from top-k elements with equal chance, and with probability α, the element is chosen from non-top-k elements with equal chance. Here, α ∈ [0, 1] is a hyperparameter to control the randomness. The core idea behind randomized topk sparsification comes from the analysis of size reduction and top-k sparsification. We find that when the number of classes is small (e.g. ≤ 10), simply reducing the size of the cut layer can preserve model performance well. However, this does not hold when the number of classes is large (e.g. ≥ 100), since hidden features of different classes will be crowded within a low-dimensional manifold, making the model less smooth and generalize poorly. Top-k sparsification overcomes this problem by expanding the volume of the low-dimensional manifold. However, it still faces local minimum problems, and usually fails to exploit the entire volume of the manifold due to unevenly selected neurons, i.e., some neurons are selected frequently while some are rarely selected. RandTopk tackles both problems by adding randomness to top-k sparsification. First, the local minimum can be avoided by noisy gradients; Second, since non-top-k neurons also have chances to be selected during forward propagation, after training, the distribution of selected neurons is more balanced.

We conduct extensive experiments on different tasks, and compare RandTopk with top-k sparsification, cut layer size reduction, quantization, and L1 regularization. Empirical results show that RandTopk significantly outperforms top-k sparsification and other compression methods when their compression ratios are on the same level. We summarize our main contributions as follows:

• We investigate different compression methods for split learning to improve its communication efficiency, including cut layer size reduction, top-k sparsification, quantization, and L1 regularization.
• Through the analysis of the cut layer size reduction and top-k sparsification, we propose randomized top-k sparsi-
• fication, to overcome the convergence and generalization problems of top-k sparsification.
• We conduct extensive experiments and show that randomized top-k sparsification achieves the best model performance compared to other compression methods when their compression ratios are on the same level.

2 RelatedWork

2.1 Reducing Communication for HFL

A classic communication reduction method for HFL is the famous FedAvg scheme [McMahan et al., 2017], which reduces communication rounds for HFL by letting client train more rounds locally, and has now become the standard HFL scheme. Based on this, studies aiming to reduce communication in HFL can be mainly divided into two classes, i.e., making convergence faster and compressing the model updates. Methods of the former class usually accelerate convergence by making model updates of different clients more consistent [Wang et al., 2020a; Wang et al., 2020b; Karimireddy et al., 2020; Jin et al., 2022]. The latter class includes quantization, sparsification, and other compression methods. The basic top-k sparsification is well-studied both empirically and theoretically and is proven to be very effective [Aji and Heafield, 2017; Stich et al., 2018]. Combining sparsification and quantization can further reduce communication [Wen et al., 2017; Sattler et al., 2019b; Sattler et al., 2019a]. Other methods, such as using SVD to compress the model updates [Wang et al., 2018; Vogels et al., 2019] and using autoencoder as compressor [Chandar et al., 2021], are also investigated.

2.2 Reducing Communication for VFL

As for VFL, the studies on communication reduction are relatively few. Castiglia et al. [2022] studied the convergence using top-k sparsification and quantization. Other techniques like asynchronous training and autoencoder compressor are also used in specific tasks [Chen et al., 2021b; Ayad et al., 2021]. Chen et al. [2021b] use asynchronous training and quantization to reduce the communication for split learning. Using asynchronous training, the bottommodel is not updated every iteration, hence the top model no longer needs to send the gradient to the bottommodel in every iteration. The authors also show that this method neither enhances nor harms the privacy of split learning. Ayad et al. [2021] use an autoencoder to compress the intermediate results, which requires injecting an autoencoder specifically designed for different tasks.

3 Basic Compression for Split Learning

In split learning, intermediate results are sent during each forward pass and backward pass. Hence, the compression is performed on the intermediate results. Specifically, in this paper, we only consider the compression on the instance level, i.e., if there is a batch of instances, the compression method is applied to the instances' bottommodel outputs inside the batch individually. Let Comp(·) : Rd → B∗ and Decomp(·) : B∗ → Rd be the compression and corresponding decompression operator, where B∗ is the bytes array of arbitrary length. For simplicity, we denote C[·] ≡ Decomp[Comp(·)]. Note that the compression methods discussed in this paper usually are not lossless, hence C[X] and X are usually not equal. Let X be an input instance, and we assume that the output of Mb is flattened to a d-dimensional vector, i.e., ∀X,Ob = Mb(X) ∈ Rd . Then we can define the compression for split learning during the forward pass as follows: The feature owner sends Comp(Ob) to the label owner. The label owner recovers C[Ob] = Decomp[Comp(Ob)]. Notice that the backward pass or the loss function should also be modified according to different forward compression methods. The notations we use are described in Table 1.

3.1 Basic CompressionMethods

We now describe some basic compression methods for SL:

(Cut layer) size reduction is to reduce the size of the bottommodel's output. For example, consider a fully-connected network whose architecture is 1,000 (input)-100 (first hidden layer)-100 (second hidden layer)-10 (output). Assuming that we split the model by its second hidden layer, the cut layer size is 100. If we 'slim' the network by changing its architecture to 1,000-100-10-10, then we achieve a 90% compression ratio. In case the network architecture is complicated and hard to modify, we can just keep the first k elements of Ob. Hence, we define the compression and decompression functions as

During the backward pass, we also apply the same compression and decompression function. The reason is that the last d-k entries of the bottom output are masked and the gradient w.r.t. them is meaningless to the bottommodel.

Quantization is to convert float values (usually 32-bit) to low-bit representations, which are widely used to improve the efficiency of neural networks [Jacob et al., 2018; Gholami et al., 2021]. While there are many quantization methods, in this paper, we consider trivial uniform quantization. Suppose the bottom output's range is [omin, omax], and we want to quantize it into b-bit. We generate 2b bins, and the n-th bin has the range [omin+(omax-omin)/2b ·(n-1), omin+(omax-omin)/2b · n]. Hence, the compression and decompression functions are defined as

where (c1, ..., cd) = Decode(O′ ). Quantization can be applied on both the forward pass and the backward pass. However, since the quantization of backward gradients significantly hurts the model performance and we mostly focus on the inference efficiency, we only quantize the forward pass.

Top-k sparsification is to preserve top-k elements in the vector, in terms of magnitude, while setting all other elements to zero. We can define the compression and decompression functions as

where J = (j1, ..., jk) are the indices of largest-k elements in O. During the backward pass, the feature owner only needs the gradients on non-zero entries to update the bottom model. Hence, we also apply the compression to the backward pass. Moreover, during the backward pass, the top-k indices are already maintained by the feature owner and hence need not be transferred.

L1 regularization is widely used in different fields of machine learning to induce sparsity [Tibshirani, 1996; Wright et al., 2008; Yin et al., 2012]. To make the bottom output sparse, we add the L1-norm of the bottom output to the loss, i.e., L′ = L+λ ∑di=1 |oi|, where L is the original loss, and λ is a coefficient to control the sparsity. Larger λ tends to induce higher sparsity, but may also hurt the model performance more. The compression and decompression function when using L1 regularization is the same as the top-k sparsification case, the only difference is that the retrieved indices J are non-zero indices rather than top-k indices. Note that in the backward propagation, no sparsification shall be applied.

3.2 Compressed Size

Different compression methods have different compressed sizes. For simplicity, here we use (relative) compressed size to denote the ratio between compressed data and original data, which is the inverse of compression ratio. For size reduction, the compressed size is simply k/d. Similarly, quantization has the compressed size 2b /N, where N is the original values' bit-length, which is usually 32. As for the top-k sparsification and L1 regularization, although only k values are preserved, we also have to send their indices during the forward pass. Suppose that a single index needs r bits to encode, then the compressed size of top-k sparsification or L1 regularization is k/d · (1 + r/N). In this paper, we consider the offset encoding that r = ⌈log2 d⌉. We conclude the compressed size of different methods in Table 2.

3.3 Summary

From the above discussion, we can see that for size reduction and top-k sparsification, we can explicitly set the compression ratio by deciding the portion of preserved elements. Quantization methods can only achieve a maximum compression ratio of 32 if the original value is 32-bit. As for L1 regularization, the compression ratio is hard to control. It is hard to estimate the compression ratio until the training is finished, and the compression ratio even varies on different inputs. By our experiments in Section 5, we can see that quantization and L1 regularization fail to converge or are infeasible under many compression levels.

4 Randomized Top-k Sparsification for SL

In this section, we propose our method: randomized top-k sparsification (RandTopk). The idea of RandTopk comes from the analysis of size reduction and top-k sparsification, as they are all implemented by dropping some elements in the bottom model output. In this section, we first demonstrate that top-k is (theoretically) better than size reduction since it provides a larger feature space under the same compression ratio, which leads to better generalization. Based on this, we introduce RandTopk and demonstrate its advantage over top-k from two aspects: convergence and generalization.

4.1 Analysis of Top-k and Size Reduction

We first demonstrate that if the cut layer's size is small and there are a large number of classes, the model will suffer from a huge generalization error because it becomes less smooth. Then we illustrate that using top-k sparsification, we can make the model much smoother than size reduction.

Larger Margin, Better Generalization

Studies have shown that the generalization ability of neural networks is closely related to its smoothness [Neyshabur et al., 2015; Neyshabur et al., 2017; Gouk et al., 2021]. Generally speaking, more smoothness leads to better generalization. Here we consider the case that the cut layer is the last hidden layer, while the top model is a linear layer with softmax activation. The model prediction is

where o is the bottommodel output, wi is the weight (embedding, hidden feature) in last layer corresponding to i-th output neuron, and n is the total number of classes (ignoring the bias). When the model is trained well enough, if the input belongs to the i-th class, then o shall be close to wi and far from wj̸=i. Let the minimum margin (between two different classes' weights) dW be dW = mini̸=j ∥wi -wj∥2, the smoothness of the top model can be approximated as ∥∇θt ∥2 ≈ c/dW, where c is some constant.

Hence, we can use dW as an indicator for generalization ability. Larger dW means embeddings of different classes are away from each other, leading to lower generalization error, and vice versa.

Top-k Has a Larger Margin

The minimum margin between different classes' weights is proportional to the volume of the feature space, i.e., the space that wi lies in. Obviously, if we do not put any restriction on wi, the feature space always has an infinite measure. To avoid this, notice that the Softmax result will not change if the weight/input is multiplied by a constant. In other words, what matters is the direction of o and w, instead of the magnitude. Then we can assume ∥wi∥2 = 1. Hence, each wi can have a ball Bi of radius dW/2, hence there are totally n balls corresponding to n classes. Those balls are disjoint, thus the sum of those balls' volume is less than the volume of the entire feature space. While those balls are k-dimensional, we further consider their intersection with the hypersphere ∥w∥2 = 1. This yields a (k - 1)-dimensional manifold. While we consider any Bi is small, the manifold is close to a (k - 1)-ball, whose volume is π(k-1)/2 /Γ((k + 1)/2)(dW/2)k-1 .

For the cut layer of size k, wi's are k-dimensional. Since the hypersphere ∥w∥2 = 1 has a volume of 2πk/2 /Γ(k/2), we have

Approximation is used here because the volume is computed on a hypersphere rather than a flat Euclidean space. Then we can estimate that, if all wi's are almost uniformly distributed on the hypersphere, dW ≈ 2 · (2/n√kπ/2)1/(k-1) .

While size reduction only selects the first k dimensions of the bottommodel output, top-k sparsification has (d k) different selections, making the minimum distance (d k) times larger. However, due to the extra communication of top-k sparsification for sending the indices, we can only have k′ = αk < k non-zero elements, where α is considered to be larger than 1/2 since we do not need 32 bits to encode the indices (otherwise the cut layer size will be up to 232 ). Comparing the minimum distance of size reduction and top-k sparsification, and assume that n2 ≥ 2αkπ (which is obvious in our scenario), we have

It is obvious that if n ≤ √k/2( d αk), Equation (6) is much larger than 1. This condition holds in most scenarios since ( d αk) is a high-order polynomial of d. Refer to Appendix D for details.

In conclusion, under the same compression level, using topk sparsification makes the feature space larger, and results in a smoother top model, finally leading to a smaller generalization error.

4.2 Adding Randomness to Top-k

We show that by adding randomness to top-k sparsification, the model converges better, and even generalizes better than top-k. We add randomness to top-k sparsification by selecting the neurons as follows:

The selection is performed k times without replacement, and N1, N2 denote the numbers of remaining top-k/non-top-k neurons. We can see that α is a hyperparameter to control the degree of randomness. Larger α makes non-top-k neurons more likely to be selected. With α = 0, our method behaves exactly the same as top-k, and with α = 1, it becomes the Dropout. Empirical analysis on α can be found in Appendix C. Notice that randomness is only added during the training procedure. For model inference, our method adopts the same behavior as top-k. We will mention the proposed randomized top-k sparsification as RandTopk for short. We further argue the advantages of RandTopk as follows.

Better convergence. One concern about top-k sparsification is that the model may be stuck at local minimums since those small neurons are not trained. RandTopk overcomes this problem naturally since its randomness, which makes small neurons get opportunities to be trained. Hence, the model will converge faster and to a better minimum. To better illustrate this, we give the following toy example: Assume that we are learning the concept

with a simple split logistic model

The initial weights of the model are w1 = 1,w2 = -0.1, and two samples are provided: x1 = (1, 0), y1 = 1 and x2 = (0.5, 1), y2 = -1. It is obvious that if using top-k, w1 = +∞,w2 = -∞ are the optimal weights. However, the optimal weights can never be achieved since o2 is always masked by top-k sparsification, and w2 is never trained. To make it clearer, we plot the loss surface and the gradient field of this example in Figure 2, and we fill the area where w2 cannot be trained with blue. We can see that because of the non-continuity of gradient introduced by top-k, there is a bad local minimum in the blue area, and our training (red arrow) ends up there. Adding randomness to top-k, i.e., give non-topk neurons a chance to be selected, the problem can be solved instantly since w2 will also be trained.

Better generalization. As discussed above in Section 4.1, top-k sparsification generalizes better since it has a larger feature space (under the same compression ratio) brought by total (d k) hyperspheres. However, training with top-k sparsification cannot fully exploit the larger feature space, since the neurons are unevenly selected. Some neurons are constantly being selected with different input samples, while some are rarely selected. Those rarely selected neurons usually remain small for most input samples and are also rarely selected and trained throughout the training process. Consequently, the model cannot fully exploit the (d k) hyperspheres brought by top-k. For example, if there are d′ neurons that never become top-k, then the feature space is reduced to (d-d′ k ) < (d k) hyperspheres. The margins between different classes are also smaller. RandTopk overcomes this problem since the small neurons are also selected. If some small neurons have the potential to become top-k (w.r.t. certain input examples), RandTopk can provide them a chance to be large. Thus, the model can explore a larger space in those hyperspheres, and the margin will be larger, finally leading to a better generalization.

4.3 Discussion on Privacy

RandTopk has better input privacy than vanilla split learning since sparsification is performed on the bottom model output. This is because a large portion of the elements in the bottommodel output are zeroed out, it contains less information than the original non-sparsified output, which has also been demonstrated by previous studies [Zhu and Han, 2020; Fu et al., 2022]. We also provide the results on input reconstruction result in Appendix B. However, as split learning is vulnerable to label inference attacks [Fu et al., 2022], RandTopk cannot preserve more label privacy. Hence, it is suitable for scenarios where the number of classes is large such that it is impossible for label inference attacks. For example, a face recognition model with thousands of different faces on user devices, or a recommendation model running on the user's browser.

5 Experiments

5.1 Settings

We perform the experiments of different compression methods on four datasets, i.e., CIFAR-100 [Krizhevsky et al., 2009], YooChoose [Ben-Shimon et al., 2015], DBPedia [Auer et al., 2007], and Tiny-Imagenet [Le and Yang, 2015], using Resnet-20 [He et al., 2016], GRU4Rec [Hidasi et al., 2016], TextCNN [Kim, 2014], and EfficientNet-b0 [Tan and Le, 2019], respectively. For CIFAR-100 and Tiny-Imagenet, data augmentation including random cropping and flipping is used. For YooChoose, we only use the latest 1/64 subset and apply the same data preprocessing as in [Hidasi et al., 2016] while the loss function is changed to cross-entropy. The size of GRU layer is set to 300. For TextCNN, we set kernel sizes to [3,4,5] and use the Glove word embeddings [Pennington et al., 2014] as initialization. For Tiny-Imagenet, we train the model either from scratch (no suffix) or starting with pre-trained weights (marked by the suffix "-p"). The details of the datasets are provided in Table 4.

We split all the models by their last layer, and apply different compression methods with different compression ratios. All experiments are repeated 5 times, and the standard deviation is reported on the results. We report the results on accuracy vs. compression at the inference phase, convergence speed during training, in terms of epochs and communication size, and further analysis on the RandTopk corresponding to our theoretic claims in Section 4.2. Full experiment results, input reconstruction attacks, and analysis on the randomness coefficient α are placed in Appendix A, B, C, respectively.

5.2 Compressed Size vs. Performance

We report the test accuracy (hit ratio@20 for YooChoose) and corresponding compressed size for inference on different tasks and compression levels in Table 3. Each experiment has been repeated 3 times for TinyImagenet and 5 times for other tasks for more precise results. We set α to 0.1 for all tasks except for YooChoose, where α = 0.05.

We can see that RandTopk achieves the best performance in almost all tasks and all compression levels, and are usually significantly better than the second-best method. Moreover, RandTopk's performance is often very near to the non-sparse case (vanilla SL) even with a high compression ratio. RandTopk surpasses the non-compression case in the YooChoose task, which we think is caused by its regularization effect.

We also notice that quantization and L1 regularization are not applicable under many compression levels. For quantization, it can only achieve a maximum compression ratio of 32, and 1-bit quantization usually fails to converge. L1 regularization, on the other side, makes training difficult and usually fails to converge.

5.3 Convergence Speed

Although our main purpose is to reduce communication at the inference phase, we also report the convergence speed during training in Figure 3. The convergence speed is measured in two terms: the number of epochs and the communication size.

We can see that non-sparse (vanilla SL) training takes the least epochs to converge, while RandTopk and other compression methods converge slower, but the difference is not large to the magnitude level. As for the overall communication to converge or reach a certain accuracy level, almost all compression methods outperforms the non-sparse training. RandTopk performs best in terms of convergence speed and performance among all other methods.

5.4 Further Analysis on RandTopk

To verify our theoretical analysis in Section 4.2, we provide the results on convergence speed, generalization error, and the distribution of top-k neurons during training using CIFAR-100, to demonstrate the advantage of RandTopk.

Communication

Figure 4: Training loss and generalization error on CIFAR-100. The compressed size is 2.86%.

Faster convergence. We plot the loss curve of top-k and RandTopk with different α in Figure 4(a). We can see that at first top-k sparsification converges faster, then gradually slows. With a small α, RandTopk converges faster and reaches a smaller loss value. When α becomes larger, the convergence at the beginning is slower, however, later it will exceed top-k. The result verifies our argument that RandTopk can help avoid local minimums that top-k may get stuck at.

Better generalization. We report the generalization error in Figure 4(b). We notice that for all methods, the generalization error can be approximated by 0.5×train acc.-0.2. Hence, the y-axis is set to be the difference between generalization error and this value for better illustration. We can see that when the train set accuracy is the same, RandTopk has a significantly lower generalization error than top-k, and larger α further decreases the generalization error. This supports our argument that RandTopk generalizes better.

Distribution of top-k neurons. We report the histogram of the times a neuron become top-k. More specifically, after model training is finished, we iterate through the train set and record the top-k neurons for each input example. Notice that we perform this experiment on the inference phase, hence RandTopk also behaves deterministically like top-k. As we can see in Figure 5, if trained with top-k sparsification, some neurons will become top-k neurons very often, while some rarely become top-k. More specifically, some neurons become top-k for less than 500 times, while some neurons become top-k for more than 3,500 times. Using RandTopk, the distribution is more balanced. Larger α makes the distribution more balanced. The results verify our argument that RandTopk can make use of the expanded feature space better.

6 Conclusion

In this paper, we investigate the use of different compression methods in split learning. By the analysis of the size reduction method and top-k sparsification, we further propose randomized top-k sparsification, which strengthens top-k's advantage on generalization ability while avoiding its disadvantage on convergence. Experiments on multiple datasets and different kinds of models show that randomized top-k is superior to other sparsification methods in terms of model performance and compression ratio. However, there are still some issues worth studying in the future. On the one hand, the label privacy remains to be a problem. Moreover, further research on compression for split learning is needed, for example, combining quantization and sparsification can be promising.

Acknowledgements

This work was supported in part by the "Pioneer" and "Leading Goose" R&D Program of Zhejiang (No. 2022C01126), and Leading Expert of "Ten Thousands Talent Program" of Zhejiang Province (No.2021R52001).