Protecting Split Learning by Potential Energy Loss

1 Introduction

Split learning [Vepakomma et al., 2018; Gupta and Raskar, 2018] is a practical method for privacy-preserving machine learning on distributed data. By splitting the model into multiple parts (sub-models), split learning allows different parties to keep their data locally and only to share the intermediate output of their sub-models. Compared to cryptographic methods like [Mohassel and Zhang, 2017; Rathee et al., 2020; Zhicong et al., 2022], split learning is much more efficient in computation and communication. To date, split learning has been applied in multiple fields, e.g., graph learning [Chen et al., 2022b], medical research [Ha et al., 2021], and the internet of things [Koda et al., 2020].

To perform split learning, the model should be split into multiple parts. Without loss of generality, we suppose the model is split into two parts, i.e., the bottommodel Mb and the top model Mt, held by the feature owner (Alice) and the label owner (Bob), respectively. As shown in Figure 1 (top part), during the forward pass, Alice feeds the input feature X to the bottom model to get the forward embedding Z = Mb(X), then sends Z to Bob. Bob feeds Z to the top model and gets the prediction ˆ Y = Mt(Z). As for the backward pass, Bob computes the gradients on the loss ∂L/∂Mt and ∂L/∂Z. He uses the former one to update Mt and sends the latter one to Alice. Alice then computes the gradient with respect to Mb and updates its parameters. From the above description, we can see that split learning is very straightforward with low computation and communication overhead compared with cryptographic methods.

∗Corresponding author

However, the price for efficiency is privacy. Many previous studies have investigated the privacy leakage of input features in split learning caused by the exchange of intermediate results [Abuadbba et al., 2020; Pasquini et al., 2021]. Instead, we focus on the privacy leakage caused by the trained split model itself in classification tasks, which has been demonstrated by [Fu et al., 2022a; Sun et al., 2022]. Consider the two-party split learning scenario described in the previous paragraph. As shown in Figure 1 (bottom part), if Mb is trained well, it gains the ability to separate samples of different classes and cluster the samples of the same class, in other words, Z becomes 'meaningful' and highly correlated with the label. Hence, Alice (or whoever obtains the bottom model) can 1) fine-tune Mt from random initialization with a few labeled samples, or 2) simply perform clustering with enough unlabeled input samples. In both cases, Alice can learn the complete classification model based on the forward embeddings. Considering that the complete model and thelabel could be private assets, such model completion attack caused by the forward embedding poses a significant privacy threat to split learning.

The protection of input features in split learning is already studied. For example, Vepakomma et al. [2020] decorrelate Hb and X by adding distance correlation [Sz´ekely et al., 2007]

loss. This method is empirically successful because Z does not need to contain the majority of X's information-it only needs to contain the part most relevant to the label Y . On the other side, protecting label information is more challenging. Since Mt(Z) = ˆ Y is the model prediction, Z contains all information about ˆ Y . While the training target is to make the prediction ˆ Y and the real label Y as close as possible, it seems impossible to prevent the attacker from deriving Y from Z.

To solve this problem, in this paper, we view it from a different perspective. At first glance, the attacker's target seems exactly the same as Bob's training objective, i.e., to learn a mapping from Z to Y . But there is a crucial difference between them. Alice and Bob perform the learning procedure on the entire training dataset, but on the contrary, the attacker can only access a small number of labeled samples by assumption. The problem now turns into modifying Z's distribution so that the attacker learns poorly with a few labeled samples while the benign actors learn well with sufficient labeled samples. We observe that such distribution can be obtained if the embeddings of the same class are distributed near the boundary of the decision region. However, during vanilla split learning, different-class embeddings tend to separate from each other while same-class embeddings are densely clustered, which is opposite to our purpose.

Our solution is inspired by a well-known physics phenomenon named electrostatic equilibrium, i.e., all the net charges distribute on the surface of the conductor. Interestingly, this distribution happens to be the desired distribution for Z, if we view the forward embeddings as charges. One reason for this phenomenon is that there is a repulsive force between every pair of like charges (charges that have the same sign, e.g., positive charges). Inspired by this, we propose the potential energy loss on the forward embeddings. The potential energy loss adds a repulsive force between each pair of same-class embeddings. During training, same-class embeddings are pushed toward the boundary of the decision region, resulting in a large learning error for the attacker. Therefore, the attacker cannot fine-tune the bottom model with a few labels or cluster unlabeled embeddings.

In summary, we make the following contributions:

• We formalize the privacy leakage of the forward embeddings in terms of the learning error of the attacker, and demonstrate that making embeddings lie near the decision boundary enhances privacy.
• We propose the potential energy loss on the forward embeddings to push them to their decision boundary, in order to reduce the privacy leakage from the forward embeddings.
• We conduct extensive experiments on multiple datasets, showing that our method significantly reduces the attacker's learning accuracy of both fine-tuning attacks and clustering attacks, and performs better than the existing distance correlation approach.

2 RelatedWork

Privacy Concerns of Split learning. Many studies have demonstrated the privacy concerns of split learning. Most of them focus on the privacy of input features. Abuadbba et al. [2020] show that applying split learning to CNN models can be dangerous since the intermediate output is highly correlated to the input. Luo et al. [2021] propose methods for feature inference attacks of split learning under certain conditions. As for the privacy of the label data, the attack can be based on either the forward embeddings or the backward gradients. For forward embeddings: Sun et al. [2022] find that bottom model output (forward embeddings) also leaks label data. Fu et al. [2022a] point out that the attacker can easily fine-tune the top model with a few labeled samples. For backward gradients: Li et al. [2022] investigated the label leakage brought by the backward gradients, for example, the norm of different classes' gradients can be different. Studies [Erdo˘gan et al., 2022; Kariyappa and Qureshi, 2023] use a surrogate model and label to match the backward gradients and reveal the training labels. While the above studies assume the attacker is passive, i.e., he will not tamper with the training procedure but only exploits the information he received, Pasquini et al. [2021] propose an active attack by making the bottom model invertible by modifying the training objective.

Privacy Protection for Split Learning. Aside from relatively expensive cryptographic-based methods such as privacypreserving machine learning [Mohassel and Zhang, 2017; Rathee et al., 2020; Zhicong et al., 2022] or partial cryptographic split learning methods [Fu et al., 2022b; Zhou et al., 2022], non-cryptographic methods mostly protect privacy by perturbing the forward embeddings or backward gradients. Perturbing forward embeddings: Vepakomma et al. [2020] add a distance correlation [Sz´ekely et al., 2007] loss to decorrelate the input features, and similarly, Sun et al. [2022] use distance correlation to protect the label. Duan et al. [2023] perturb the embedding by minimizing the mutual information with both label and features to protect both, however, it requires a customized model structure. Chen et al. [2022a] protect the embedding in the recommendation model via differential privacy. Perturbing backward gradients: Li et al. [2022] protect the label during training by perturbing the backward gradients, although the forward embeddings could still leak the label information.

Data-Dependent Generalization Error. Most studies on data-dependent generalization error are based on the Rademacher and Gaussian complexity [Koltchinskii and Panchenko, 2002; Kontorovich and Weiss, 2014; Lei et al., 2019], or the mutual information between the data and the algorithm output [Negrea et al., 2019; Pensia et al., 2018; Russo and Zou, 2020]. Different from them, Jin et al. [2020] derived generalization bounds directly from the data distribution, by proposing the so-called cover complexity, which is computed from the distances between same-class data points and different-class data points. It is somewhat related to our work since our method makes the data distribution more 'complicated' by pushing the data points to the decision boundary of their class.

3 Problem: Model Completion Attack

In this section, we demonstrate and formalize the privacy problem arising from the bottom model in split learning.

3.1 Leakage fromForward Embeddings

The hidden embeddings of neural networks are widely studied [Rauber et al., 2017; Pezzotti et al., 2018; Cantareira et al., 2020]. Through visualization and other techniques, those studies show that the neural network gradually learns to make hidden embeddings of different classes separate, and those of the same class clustered together. Although this 'separation ability' seems to be essential for neural networks and may be the reason why they perform well on various tasks, it also brings security hazards for split learning.

In split learning, the model is split into one (or multiple) bottom model(s) and top model(s), which are held by different parties. Ideally speaking, any single party can not perform inference tasks, since he only gets a part of the complete model. Only multiple parties work together, can they make use of the complete model and perform inference tasks. In other words, the model should be trained and used in a 'shared' manner.

However, given the fact that the hidden embeddings of the model are meaningful, the attacker who has the bottom model can either fine-tune the top model with a small number of labeled samples [Fu et al., 2022a], or just perform clustering on the forward embeddings to infer the labels [Sun et al., 2022]. Thus, the privacy of split learning is violated.

3.2 Threat Model

We consider two threat models, i.e., fine-tuning attack with a few labeled samples and unsupervised clustering attack with massive unlabeled samples.

Fine-tuning Attack. We assume that the attacker has access to the trained bottom model Mb, along with a few labeled samples X′k , Y ′ k which include k samples for each class. The attacker also knows the architecture of the top model Mt, and performs the model completion attack via training Mt from a random initialization, given X′k and Y ′ k , with pre-trained Mb fixed.

Clustering Attack. We assume that the attacker has access to the trained bottom model Mb, along with massive unlabeled samples X′ . To infer the labels of X′ , the attacker performs clustering algorithms on the forward embedding Z ′ =Mb(X′ ).

Notably, we assume the attacks are conducted in the inference phase instead of the training phase, as our method aims to reduce the privacy leakage from forward embeddings. To prevent leakage from backward gradients during training, one can use non-sensitive data for training or adopt existing approaches such as cryptography-based secure computation described in Section 2. It is worth noting that, some unsupervised/semi-supervised learning methods can train the complete model with good performance over unlabeled/partially labeled data [Berthelot et al., 2019; Xu et al., 2021]. However, they are not relevant to split learning since they do not require any knowledge about the trained bottommodel or forward embedding.

3.3 ProblemFormulation

In order to reduce the aforementioned privacy leakage while maintaining the model performance at the same time, our purpose is to train a split model M = (Mb,Mt) such that the output of Mb is hard for the attacker to learn, while the complete M still maintains a high performance.

Definition 1 (Bottommodel advantage). The bottom model advantage is the extra advantage obtained by the attacker when he has access to the trained bottom model. Consider an attack algorithm A whose input is the data D and (optionally) the bottom model Mb, the bottom model advantage is defined as follows:

where R[·] is an error metric for the attack outcome. For example, in the fine-tuning attack, D = (X′k , Y ′ k ) is the leaked labeled samples, and R[·] is the error of the fine-tuned model on the test data; in the unsupervised attack, D = X is the unlabeled data, and R[·] is the error of the clustering model on the test data. We use A(D; null) to represent the case that the attacker trains the whole model solely based on the leaked data, without any information about the bottommodel or forward embedding (i.e., training from scratch).

Definition 2 (Perfect protection). For an attack algorithm A, if the bottom model Mb satisfies Adv(Mb;A) = 0, then we say that Mb achieves perfect protection against attack A, since Mb provides no extra advantage for the attack.

Thus, our purpose becomes to train a split model (Mb,Mt) such that:

• The performance of the complete model on the original task is as high as possible.
• The bottommodel advantage under fine-tuning attacks and clustering attacks is as small as possible.

4 Method: Potential Energy Loss

In this section, we view the privacy leakage of the bottom model as a learning problem for the attacker. We first study the generalization error when fine-tuning the model with a small number of labeled samples. By a simplified example, we show that pushing the embeddings of same-class samples toward the decision boundary increases the generalization error. At the same time, clustering is difficult since the sameclass embeddings are no longer close to each other. Inspired by the electrostatic equilibrium and Coulomb's law, we propose the potential energy loss on the forward embedding, to realize such distribution. The high-level view of our idea is presented in Figure 2.

4.1 Learning Error fromData Distribution

Recall that our goal is to train a split model (Mb,Mt), such that the bottommodel provides little advantage to the attacker. To do this, we consider the attack process to be a learning process on the forward embeddings produced by the bottom model. We show that when the same-class embeddings are distributed near the boundary of the decision region, the performance of fine-tuning attack is decreased since a small number of samples cannot represent the overall distribution, and a small error on the estimation of decision boundary will cause a large classification error. Moreover, it naturally prevents clustering attacks since same-class embeddings are no longer close to each other.

Genralization Error

We use a simplified example to get some insights into the relationship between the data distribution and generalization error. Assume that all data points are distributed on the d-

hemisphere

Without loss of generality, we assume the target hypothesis is f(x) = Sign(x1). We make the following assumptions:

• The probability density of samples only depends on the first dimension x1, i.e., it is isotropic in any other dimensions.
• Given a set of positive samples S = {x1, ..., xn}, the learning algorithm simply outputs the normalized mean of these samples as the parameter of learned hypothesis,

Now we want to estimate the generalization error when the learned parameter w = n∑ i=1 xi/∥∥ n∑ i=1 xi∥∥2 slightly differs from the true parameter e1. Since the distribution is isotropic except in the direction of e1, we may assume that w lies on the plane expanded by the first two axis, i.e., w = e1 cos ϵ + e2 sin ϵ, where ϵ is a small angle between w and e1. The generalization error is

where p is the probability density of sample feature, and p1 is the marginal density function of x1. From (3) we can see that with ϵ fixed, the generalization bound is approximately proportional to the probability mass of the data points falling near the boundary of the target region.

Sampling Error

In the above analysis, the estimation error ϵ is fixed. We now explore the relationship between the data distribution and the distribution of ϵ. Notice that for any random variable X, if X1, ...,Xm are m independent samples, we have

words, if the random variable is likely to fall far from the mean of its distribution, the sample mean tends to have a larger error. Although ϵ is not exactly the error of the sample mean in our case since it is an angle, it is also reasonable to assume E[ϵ2 ] ∝ E [(X - EX)2 ]. To make (the magnitude of) ϵ

larger, the data points should be away from their mean as much as possible. Interestingly, in our case, this is also equivalent to pushing the data points toward the decision boundary.

Clustering Error

While the above discussion is on the generalization error under a small number of labeled data, we can easily see that such distribution makes clustering difficult. This is because the basic idea of clustering is that the data points in the same cluster shall be close to each other, and those in the different clusters shall be far from each other. While the same-class embeddings are being pushed toward the decision boundary, intra-class distances are increased, while inter-class distances are decreased. This is exactly opposite to the basic requirement of clustering.

In summary, pushing data points to the boundary of the decision region will increase the learning error of both finetuning and clustering for the following reasons:

• The sampling error tends to be larger.
• A small error in the decision region will result in a large generalization error.
• Intra-class distances are larger and inter-class distances are smaller.

4.2 Potential Energy Loss

When electrostatic equilibrium is established, any net charge resides on the surface of the conductor [Griffiths, 2005, ch. 2]. This is partly caused by Coulomb's law, which tells us that like charges (electric charges of the same sign) repel each other and opposite charges attract each other. Inspired by this, we can view the embeddings of the same class as like charges and have repulsive forces against each other. As a result, those data points will tend to be away from each other and be pushed to the boundary of the decision region.

Coulomb's law is stated as follows:

where k is a constant, q1, q2 is the signed magnitude of the charges, r1, r2 are their positions, ∥ · ∥2 is the Euclidean norm, and F is the repulsive force on the first charge caused by the second charge. Since we assume all embeddings belonging to the same class have the same sign and magnitude while ignoring the constant term, (4) becomes F = r1 - r2 ∥r1 - r2∥32 . Notice that the repulsive force is the gradient of the electric potential energy, we can further write F as F = ∇r1 1 ∥r1 - r2∥2 , which is naturally suited to the gradient descent method. Based on this, we define the potential energy loss (PELoss) as

where C is the label set, and Zc is the forward embeddings of c-labeled samples.

By adding Lpe to the loss function, during the training of the split model, the bottommodel outputs of the same class are pushed away from each other, and move towards the boundary of the decision region of its own class. While in the 3-D case, minimizing the potential energy leads to zero charge density inside the region by Thomson's theorem, it is not necessarily true in the high-dimensional case. However, we are able to prove a weaker theorem, i.e., minimizing the potential energy leads to a non-zero probability mass in the border of the region.

Theorem 1 (Border distribution). Consider a d-dimensional bounded region Ω ⊂ Rd , we denote the density of probability distribution in Ω which minimizes the potential energy functional as

where f is the probability density function of some distribution such that f ≥ 0 and ∫Ω f(x)dV = 1, ∆ϵΩ = {x : x + ϵr ̸∈ Ω, x ∈ Ω, ∥r∥2 = 1} is the set of points whose distance to the border of Ω is less than ϵ. Then f ∗ satisfies that ∫x∈∆ϵΩ f ∗ (x)dV > 0 for any ϵ > 0.

Adding Layer Normalization

One hidden condition for our method is that the decision boundary is also the actual set boundary of same-class embeddings. This requires the embedding space to be a borderless manifold. Otherwise, adding repulsive force may just push the embeddings toward the set boundary instead of the decision boundary. For example, if the embedding space is the trivial Euclidean space Rd , the repulsive force will make embeddings far from the origin, while the inter-class distances are still large. To overcome this, we simply enforce layer normalization [Ba et al., 2016] (without element-wise affine transformation) on z, which restricts z to the d-sphere of radius √d, i.e., ∥z∥22 = d. Accordingly, the distance metric is changed to angular distance, i.e., arccos⟨z, z′ ⟩. The potential energy loss defined in

(5) should be changed to

The combined loss for split training is L′ = L + αLpe, where L is the original loss function (i.e., cross-entropy loss), α is a coefficient to control the intensity of repulsive force.

4.3 Relationship with Distance Correlation

Distance correlation is used to protect both the input feature [Vepakomma et al., 2020] and the label [Sun et al., 2022]. Here we consider the label-protection case that distance correlation loss is applied to the forward embedding and the label. The distance correlation loss on one batch is

where di,j is the doubly-centered distance between i-th sample's embedding and j-th sample's embedding, and d′i,j is the doubly-centered distance between i-th label and j-th label. In the classification task, same-class samples have the same label. If the i-th sample and j-th sample belong to the same class, we have d′i,j = 0. Thus, (8) (ignoring the denominator) is converted to

(9)

where C is the set of labels, k is some constant, ∥z - ·∥2 is the average distance from z to other embeddings within the batch (similar for ∥· - z′ ∥2), and ∥ · - · ∥2 is the average distance between all embedding pairs within the batch. We can see that minimizing the distance correlation is similar to minimizing the inter-class distances. As our method is to maximize intra-class distances, minimizing distance correlation has a similar effect. However, in an intuitive understanding, the fact that embeddings of different classes lie in different directions makes distance correlation naturally noisy. Experiments in Section 5.4 also illustrate that minimizing distance correlation fails to push away some same-class samples.

5 Empirical Study

We conduct experiments on four different datasets, i.e., MNIST [LeCun et al., 1998], Fashion-MNIST [Xiao et al., 2017], CIFAR-10 [Krizhevsky et al., 2009], and DBpedia [Auer et al., 2007]. We compare the attack performance of vanilla split training, training with potential energy loss (PELoss, our method), training with distance correlation loss (DcorLoss) proposed by [Vepakomma et al., 2020; Sun et al., 2022], and a simple baseline of label differential privacy method based on randomly flipping a certain portion of labels (LabelDP). For DcorLoss, layer normalization is also added like in our method for training stability. The attacks include both fine-tuning attacks and clustering attacks. For clustering attacks, we use the classical k-Means algorithm [MacQueen and others, 1967], with the number of classes known

0.95 0.96 0.97 0.98 0.99

0.95 0.96 0.97 0.98 0.99

Test accuracy of original task

Figure 3: Test accuracy vs. attack accuracy of the fine-tuning attack. The data point lies at the lower-right position means the experiment achieves higher test accuracy on the original task with lower attacker's accuracy, which is our desired result. The data point lies in the gray area means that the attacker's accuracy is lower than training from scratch without the bottommodel's information, i.e., perfect protection.

to the attacker. We also measure the distances between sample pairs to illustrate the mechanism of PELoss.

Due to space limits, detailed descriptions of the model architectures and training strategies, and more results are presented in Appendix B and C. Other results include studies on different split positions, different attack layers, different forward embedding dimensions, and the t-SNE [Van der Maaten and Hinton, 2008] visualization of the forward embeddings.

5.1 Experiment Settings

We implement the experiment codes using the PyTorch and Scikit-Learn [Pedregosa et al., 2011] libraries, and run them on servers with NVIDIA RTX3090 GPUs. Each experiment is repeated 3 times for the original tasks and 5 times for the attack tasks, with different random seeds. We use a 3-layer fully-connected network for MNIST, a simple convolutional network for Fashion-MNIST, ResNet-20 [He et al., 2016] for CIFAR-10 dataset, and TextCNN [Kim, 2014] for DBPedia. The split position for each model is the last dense layer by default since the last forward embedding is the closest to the label and is the most difficult to protect. For the selection of hyperparameters, we vary the loss coefficient (α) of PELoss from 0.25 ∼ 32 and 1 ∼ 32 for DcorLoss. For LabelDP, the ratio of randomly flipped labels varies from 0.01 ∼ 0.16. In all experiments, the value doubles each time. Detailed experiment settings are provided in Appendix B.

5.2 Fine-tuning Attack

We report the test accuracy of the original task and the accuracy of fine-tuning attacks using different protection methods in Figure 3. Deeper colors represent stronger protection, e.g., larger loss coefficient or label flipping probability, while lighter colors mean more leaning toward preserving the test accuracy, e.g., smaller loss coefficient or label flipping probability. We also plot the test accuracy (orange dashed line) and the attack accuracy (red dashed line) in the vanilla split learning case, and the perfect protection area (where the attack accuracy is lower than the accuracy of training from scratch using leaked labels). For all methods, we observe that decreasing the attack accuracy usually also lowers test accuracy on the original task, and more leaked labels lead to higher attack accuracy. Although all methods protect privacy to some extent at the cost of damaging the model performance, it is obvious that our proposed PELoss is superior to DcorLoss and LabelDP. PELoss has the following advantages compared with other methods:

0.95 0.96 0.97 0.98 0.99

0.95 0.96 0.97 0.98 0.99

0.95 0.96 0.97 0.98 0.99

0.95 0.96 0.97 0.98 0.99

• The curves of PELoss are constantly at the lower-right side of other methods. In other words, on the same test accuracy level, the PELoss has a significantly lower attacker's accuracy than DcorLoss and LabelDP.
• The curves are smoother and the error bars are also smaller using PELoss, indicating it is more responsive to the change of α and the performance is more stable. Thus, it is easier to balance privacy and model performance using PELoss.
• PELoss has many data points in the perfect protection zone, while other methods rarely achieve perfect protection.

In summary, PELoss provides significantly stronger privacy protection against model completion attacks, while preserving the model performance better than other methods. We also notice that in some datasets, varying the coefficients of DcorLoss seems to have no effect, however, when the coefficient is large enough (e.g., 32), the training could diverge.

5.3 Clustering Attack

We report the test accuracy of the original task and the accuracy of clustering attacks using different protection methods in Figure 4. In all the experiments, the attacker can access the entire unlabeled test dataset which contains more than 10,000 unlabeled samples, and the number of classes is known to the attacker. The notations are the same as fine-tuning attacks, while here the gray area (perfect protection) means that the accuracy of clustering on forward embeddings is lower than directly clustering on the raw input samples. The accuracy here is defined as the maximum accuracy among all possible cluster-label assignments.

We can see that PELoss achieves perfect protection in all cases (i.e., the clustering result on the forward embeddings is worse than on the raw input data), and is significantly better than DcorLoss on CIFAR and DBPedia. On the other side, LabelDP performs badly against clustering attacks.

5.4 Distances between Sample Pairs

To better illustrate the distribution of forward embeddings, we plot the distribution of angular distances between sameclass and different-class sample pairs in Figure 5. For vanilla training, the angular distances are small between same-class samples, and are large between different-class samples. For PELoss, the angular distances between sample pairs are close to π/2, no matter whether they belong to the same class or not, in both training data and test data. In contrast, for DcorLoss, the angular distances appear to have a double-peak distribution. We can see that there are significantly more same-class sample pairs that have an angular distance around 0, which weakens the protection.

6 Conclusion

In this paper, we investigate the privacy leakage of split learning arising from the learned bottom model. We view the model completion attack as a learning procedure for the attacker and turn the privacy-preserving problem into a learning problem. We find that pushing embeddings to their decision boundary increases the learning error for the attacker, and propose the potential energy loss on forward embedding to protect the privacy. Extensive experiments show that our method significantly restricts the ability of the attacker's fine-tuning and clustering attacks while reserving the model performance, superior to baselines in terms of both utility and privacy. The limitation of this work mainly includes that only the inference process is protected, and the lack of a theoretical leakage bound.

Acknowledgments

This work is supported in part by the National Natural Science Foundation of China (No. 72192823), the "Ten Thousand Talents Program" of Zhejiang Province for Leading Experts (No. 2021R52001), and Ant Group.