Self-Supervised Vertical Federated Learning
Abstract
We consider a system where parties store vertically-partitioned data with a partially overlapping sample space, and a server stores labels on a subset of data samples.
Supervised Vertical Federated Learning (VFL) algorithms are limited to training models using only overlapping labeled data, which can lead to poor model performance or bias.
Self-supervised learning has been shown to be effective for training on unlabeled data, but the current methods do not generalize to the vertically-partitioned setting.
We propose a novel extension of self-supervised learning to VFL (SS-VFL), where unlabeled data is used to train representation networks and labeled data is used to train a downstream prediction network.
We present two SS-VFL algorithms: SS-VFL-I is a two-phase algorithm which requires only one round of communication, while SS-VFL-C adds communication rounds to improve model generalization.
We show that both SS-VFL algorithms can achieve up to 2 × higher accuracy than supervised VFL when labeled data is scarce at a significantly reduced communication cost.
Introduction
Federated learning has become of recent interest to the research community [1, 2, 3] and has shown promise in several applications such as personalized healthcare, smart transportation and predictive energy systems [4].
Federated learning algorithms allows a set of distributed parties to train a model without the need to directly share local data.
Vertical Federated Learning (VFL) [5] algorithms are an important class of federated learning algorithms.
In VFL scenarios, there are typically a small number of institutions that store data with the same sample space but different feature spaces.
For example, a bank, a hospital, and an insurance firm may seek to predict a value of common interest, such as credit score.
Each institution may have information on the same individuals but will hold different feature information (e.g., financial transactions, medical history, and vehicle accident reports).
In VFL, each party typically trains a local representation network while the server combines the outputs of all parties’ networks, known as representations , to train a final prediction model.
We illustrate the VFL model in Figure 1a.
VFL is in contrast to horizontal federated learning (HFL), where parties share a feature space but not the sample space.
There has been much recent interest in
VFL algorithms [6, 7, 8, 9, 10, 11] though there are still many open problems. [Figure]
VFL algorithms are typically supervised : models are updated based on how well their predictions compare to ground-truth labels .
However, collecting labels can be a manual, time-consuming, and costly process.
Often, only a small fraction of data is labeled.
A major concern is that the labeled data does not represent the true underlying data distribution, and training on only labeled data can lead to model bias and poor model generalization [12].
Therefore, there is a need for incorporating self-supervised methods.
Self-supervised learning (SSL) methods typically pre-train a model using unlabeled data followed by downstream supervised training [13].
Although there is a large body of work in SSL for HFL applications [14, 15, 16, 17], it has yet to be applied to the VFL setting.
Applying SSL to vertically-partitioned settings comes with a variety of unique challenges.
SSL methods typically require full feature set access.
However, each party in VFL only stores a subset of the total features for each sample.
Sharing features in a VFL setting is often impossible due to data privacy or communication constraints.
Additionally, in VFL settings, it is common that institutions have data for different individuals.
This is illustrated in Figure 1b.
These non-overlapping datasets can be very large, and it is necessary to integrate into the training process to avoid model bias and improve performance.
Finally, many SSL algorithms assume parties have label access, which is not the case in VFL settings.
Our work seeks to answer the following: how can we conduct selfsupervised learning on vertically-partitioned and potentially non-overlapping data?
To answer this question, we propose two self-supervised VFL (SS-VFL) algorithms.
Our first algorithm, SS-VFL Independent (SS-VFL-I), is a two-phase algorithm that requires only one round of communication between parties and the server, while the second, SS-VFL Coupled (SS-VFLC), adds communication rounds to improve model generalization.
In both algorithms, parties independently utilize contrastive learning for pre-training representation networks.
Parties produce low-dimensional representations of local data without communication or access to labels.
A major benefit of this approach is that parties’ unlabeled datasets need not overlap.
After unsupervised training, both algorithms are followed by downstream supervised training.
SS-VFL-I proceeds by sending frozen representations to the server for final training, while SS-VFL-C updates representations during supervised training for stronger model performance at the cost of increased communication.
Our contributions are summarized as follows:
1) In Section 2, we formulate a self-supervised framework for training a model in a distributed system with a partially labeled vertically-partitioned dataset. 2)
In Section 3, we propose SS-VFL-I, which improves performance over supervised VFL when labeled data is scarce, and has equivalent communication cost to only a single epoch of supervised VFL over the labeled data.
We propose SS-VFL-C, which trades communication savings for improving model performance by updating representation networks during downstream supervised training.
SS-VFL-C has the same communication cost as supervised VFL. 4)
In Section 4, we compare the SS-VFL algorithms with supervised VFL in a thorough set of experiments and show the SS-VFL algorithms can achieve up to 2 × the test accuracy of supervised VFL when labeled data is scarce, as well as significantly reduce communication cost to reach target accuracies.
Related work: Federated learning was first proposed by [18] with the goal of training a global model when participants store private data and the communication network faces high latency.
Early works in Federated Learning have targeted the horizontally partitioned data scenario [1, 19, 20].
Self-supervised algorithms that utilize unlabeled data have been proposed in the past for the HFL setting [17, 14].
However, HFL algorithms rely on distributed gradient descent methods and share model parameter updates, while most VFL algorithms utilize distributed coordinate descent methods and share feature representations [7, 10, 11].
Thus, VFL algorithms are fundamentally different from HFL algorithms.
A few proposed VFL algorithms utilize unlabeled data.
FedMVT, proposed by Kang et al.
[21], uses a simple pseudo-labeling scheme that may be vulnerable to confirmation bias [22] and only applies to a two-party scenario.
Cha et al.
[23] propose a VFL algorithm that uses unsupervised learning with autoencoders as a pre-training step for supervised learning.
Their work uses overcomplete autoencoders, which increase the dimensionality of raw data, causing the cost of communication to greatly increase over other VFL algorithms.
Our proposed VFL algorithms use a method contrastive learning shown to avoid confirmation bias via data augmentation [22], apply to an arbitrary number of parties, and reduce communication cost.
Problem Formulation and Preliminaries
We present our problem formulation and provide background on related algorithms.
Problem formulation: We consider a network with M parties and a server.
Each party stores a set of features of a labeled dataset consisting of N l data samples.
We let X m denote the labeled data partition held by party m , and we let x im ∈
X m correspond to the data features that party m stores on the i -th labeled data sample.
The entire labeled data set is denoted by X =
[ X 1 , . .
. , X M ] , and y is the set of labels that corresponds to the data samples in X .
We assume that y is only present at the server.
Each party m also stores a set of features U m corresponding to N um unlabeled data samples.
We assume that X m
and U m share the same feature dimension.
Note, however, that we do not assume the parties necessarily store features for the same set of unlabeled samples.
Thus, N um may not equal N uj for parties m ̸
= j .
Figure 1b illustrates the data partitioning among the parties.
The model consists of M representation networks and a server prediction model.
Each party m stores its own representation network h m ( · ) , parameterized by θ m , that maps its local features to a representation space.
The server’s prediction model f ( · ) , parameterized by θ 0 , combines representations from each party to make a prediction.
An example of the model architecture is provided in Figure 1a.
Let the model parameters be Θ =
[ θ 0 , . .
, θ M ] , and let D be the distribution from which all data samples are drawn.
The goal in training is to minimize the expected loss: E ( x,y )
∼D f ( θ 0 , h 1 ( θ 1 ; x 1 ) , . .
, h M ( θ M ; x M ); y ) .
We next present existing methods for training over labeled and unlabeled data.
We leverage these methods to develop our SS-VFL algorithms.
Supervised VFL:
When data is labeled, one can use supervised VFL to train the model by minimizing the objective: F (Θ; X , y ) :
= 1 N l (cid:80) N l i =1 f ( θ 0 , h 1 ( θ 1 ; x i 1 ) , . .
, h M ( θ M ; x iM ); y i ) .
We provide pseudocode for supervised VFL [7, 24] in Appendix A. To summarize, each round starts with the parties agreeing on a randomly sampled mini-batch B t from the labeled dataset.
Each party inputs its corresponding features X B t m into its local representation network h m ( · ) and sends the resulting representations to the server.
The server then updates its model: θ t +
[Equation]
where Φ t is the set of all embeddings for B t , y B are the labels for batch B t , and η t is the step size at iteration t .
Next, the server computes the partial derivative ∇ h m ( θ tm ; X B tm ) F B (Φ t ; y B t ) for all m ∈ M .
The server sends each partial derivative to its respective party.
Each party then can compute the gradient update: ∇ m F B (Φ t ; y B t ) =
h m ( θ tm ; X B t m )
⊤ ∇ h m ( θ tm ; X B tm ) F B (Φ t ; y B t ) , and update its model: θ t +1 m =
θ tm − η t ∇ m F B (Φ t ; y B t ) .
Unsupervised representation learning: Contrastive learning is a centralized SSL solution to training with unlabeled data.
In contrastive learning, it is assumed that for each data sample, there is access to a set of positive samples that share the same underlying structure as the original sample, as well as a set of negative samples that have a different underlying structure.
In practice, these pairs are typically generated by creating data augmentations [25, 26, 27, 28].
The goal of contrastive learning is to train a representation network that maps data to a representation space where positive pairs are "close" to each other, and negative samples are "far" apart.
This leads to representations clustered by their underlying features, making downstream supervised learning a simpler task.
For each class c , we let D c be the probability distribution over X , where X is the set of all possible data points.
D c ( u ) captures how relevant a sample u is for a class c .
Let u and u + be data samples chosen i.i.d.
from the same class, and let u − be a data sample chosen i.i.d.
from a random class.
We define D + ( u, u + ) as the distribution of positive pairs, and let D − ( u − ) be the distribution of negative samples.
We define a similarity measure sim( z 1 , z 2 ) :
= exp( z ⊤ 1 z 2 /τ ) where τ is a tunable temperature parameter.
Contrastive learning trains a representation function h ( · ) on an unlabeled
1: Initialize: θ 0 m for all parties m and server model θ2: for m ← 1 , . . .
, M in parallel do 3: θ T un m ← L OCAL CL ( θ 0 m , U m , T un ) 4:
Send representations h m
( θ T un m ; X m ) to server 5: end for 6: for t ← 0 , . .
, T sup − 1 do 7:
Randomly sample labeled mini-batch: B t ∈ { X , y } 8: Φ t ← { θ 0 , h 1 ( θ t 1 ; X B t 1 ) , . . .
, h M ( θ tM ; X B t M )
[Equation]
9: θ t + =
[Equation]
t
[Equation]
∇ 0 F B (Φ t ; y B t ) 10: end for dataset U by minimizing the objective [26]:
[Equation]
where K is the number of negative samples.
The numerator in (1) is maximized when the cosine similarity of positive pair representations is largest, while the denominator is minimized when the cosine similarity of negative samples is smallest.
Contrastive learning has been proven to reduce the sample complexity of downstream supervised tasks [29].
Self-Supervised Vertical Federated Learning
We now present our algorithms for VFL with both labeled and unlabeled data.
In both SS-VFL algorithms, each party m independently runs contrastive learning to train its representation network.
Similar to centralized contrastive learning, each party m has a local class probability distribution
D cm ( u ) over X that captures the relevancy of a sample u to a class c .
The goal of each party is to minimize (1) for its local class distribution D cm .
A party approximately minimizes (1) by randomly selecting a mini-batch B of samples, and updating its model using the gradient of the following:
[Equation]
The formal pseudocode of LocalCL is provided in Appendix A.
SS-VFL-I: We now introduce SS-VFL-Independent (SS-VFL-I), a natural extension of contrastive learning to VFL.
The pseudocode for SS-VFL-I is presented in Algorithm 1.
At the start of training, the parties independently perform LocalCL to train their representation networks for T un iterations.
Then, each party computes the representations for labeled data and sends these representations to the server.
The server trains its prediction model on these representations without communication.
SS-VFL-I only requires sending representations for all labeled data once; its communication is equivalent to a single epoch of supervised VFL, which can be immensely beneficial in situations where bandwidth is limited or costly.
SS-VFL-I also provides inherent label privacy.
It has been shown that supervised VFL can potentially leak label information through the sharing of partial derivatives [30, 31].
In SS-VFL-I, the server never communicates partial derivatives with the parties.
SS-VFL-C:
Next, we present SS-VFL Coupled (SS-VFL-C), an SS-VFL algorithm that improves representation networks during downstream supervised training while maintaining the same communication cost as supervised VFL.
SS-VFL-C trades the communication savings of SS-VFL-I in order to update representation networks, improving downstream supervised model performance.
The pseudocode for SS-VFL-C is presented in Algorithm 2.
Just as in SS-VFL-I, each party uses local contrastive learning to train its representation network for T un iterations.
However, the representation networks are not frozen at this point.
In the second stage of SS-VFL-C, the same procedure as supervised VFL is followed.
The parties share representations with the server, which trains a downstream prediction model.
Then the server shares partial derivatives with the parties, and the parties update their representation networks.
SS-VFL-C has the same communication cost as supervised
VFL, but with the benefit of representation networks bein
pre-trained using unlabeled data.
1: Initialize: θ 0 m for all parties m and server model θ2: for m ← 1 , . . .
, M in parallel do 3: θ T un m ← L OCAL CL ( θ 0 m , U m , T un ) 4: end for 5: for t ← 0 , . . .
, T sup − 1 do 6:
Parties choose randomly sampled mini-batch: B t .
7: for m ← 1 , . . .
, M in parallel do 8:
Party sends representation h m ( θ tm ; X B t m ) to server 9: end for 10: Φ t ← { θ 0 , h 1 ( θ t 1 ; X B t 1 ) , . .
, h M ( θ tM ; X B t M ) } 11: θ t + = θ t 0 − η t ∇ 0 F B (Φ t ; y B t ) 12:
Server sends ∇ h m ( θ tm ; X B tm ) F B (Φ t ; y B t ) to each party m 13: for m ← 1 , . .
, M in parallel do 14: ∇ m F B (Φ t ; y B t ) =
h m ( θ tm ; X B t m )
⊤ ∇ h m ( θ tm ; X B tm ) F B (Φ t ; y B t ) 15: θ t +1 m =
θ tm − η t ∇ m F B (Φ t ; y B t ) 16: end for 17: end for
[Figure]
LocalCL: When is LocalCL sufficient to create representations that can be easily classified during downstream supervised training?
In Figure 2, we give an example of when LocalCL fails to produce representations that are well-separable.
We compare representations generated by running centralized contrastive learning with a full feature set and representations generated by running LocalCL with a partial feature set.
Here, we plot a kernel density map of representations from the ImageNet dataset after being reduced to two dimensions (using PCA) and normalized to the unit circle.
We show representations from two classes marked in red and blue.
Ideally, the representations should be separated into distinct clusters without any overlap.
For LocalCL, we consider a case where only half of each image is available.
In Figure 2a, when all features are available for constrastive learning, we can see that representations from the different classes are correctly mapped into separate clusters.
In Figure 2b, we can see that LocalCL produces clusters that overlap.
Since local contrastive learning only has access to a partial feature set, two label classes may be indistinguishable with the available features, making downstream supervised training more difficult.
In Figure 2c, we show the representations of two different classes at a single party after running SS-VFL-C. For cases when
LocalCL produces overlapping representations, SS-VFL-C can be run to refine the representations.
Formally, the difference between local and centralized contrastive learning manifests in the difference between D cm , the relevance score of a sample based on locally available features and D c , the relevance score based on all features.
For example, if a party only stores the lower half of an image for a table, it may be difficult to discern if the image are of the legs of a table or a chair.
In cases where D c and D cm are similar, then SS-VFL-I can provide strong representations for downstream training without high communication cost.
Otherwise, SS-VFL-C provides a means of updating representation networks during downstream training at the cost of additional communication.
Experiments
We now present experiments comparing SS-VFL-I and SS-VFL-C with supervised VFL.
The parties train their models on two datasets: ModelNet10 [32] and ImageNet [33].
ModelNet10 is a set of 2D images of 3D CAD models from different camera views.
For this dataset, there are 12 parties, each
[Figure]
[Table]
with one view of every CAD model.
Each party trains ResNet18 as a representation model, while the server trains a fully connected layer.
ImageNet is a set of images of different objects and animals.
For our experiments, we choose a random set of 100 classes from ImageNet (ImageNet100).
For
ImageNet100, there are two parties, and each party stores half of each image.
Each party uses
ResNet50, while the server trains a fully connected layer.
For SS-VFL-I and SS-VFL-C, parties run
LocalCL for 200 epochs.
Then, all algorithms train on the labeled data for 500 epochs.
We consider cases where the training set of each dataset has only 1% , 5% , 10% , or 25% of its data labeled.
We run each algorithm on both datasets with each of these labeled fractions.
The results of the experiments are shown in Figure 3, where we plot the test accuracy for each dataset and labeled data fraction.
The solid lines are the mean of 5 runs, while the shaded regions represent the standard deviation.
For the ModelNet10 dataset, we can see at 1% and 5% labeled data, both SS-VFL algorithms outperform supervised VFL.
Only at 10% labeled data and more is supervised VFL able to reach similar accuracies to the SS-VFL algorithms.
We can see that SSVFL-I and SS-VFL-C perform similarly in all cases, indicating that LocalCL was able to produce well-separable representations during unsupervised training.
For the ImageNet100 dataset, both
SS-VFL algorithms perform similarly when 1% of the data is labeled, reaching up to double the accuracy of supervised VFL.
As the amount of labeled data increases, we can see that SS-VFLC continues to reach higher test accuracy than the other two algorithms, providing the best model generalization by utilizing both labeled and unlabeled data to train party representation networks.
In the case of ImageNet100, LocalCL has more difficulty distinguishing between similar classes.
SS-VFL-C outperforms SS-VFL-I here by utilizing additional communication during downstream supervised training, allowing it to refine the representations and improve performance.
In Table 1, we show the communication cost between the parties and the server for supervised VFL,
SS-VFL-I, and SS-VFL-C to reach a target test accuracy.
We can see in Table 1 that SS-VFL-I has a much smaller communication cost than both other algorithms, regardless of the fraction of labeled data.
SS-VFL-C, although requiring more communication, still reduces overall communication cost to reach target accuracies over supervised VFL in both datasets.
For scenarios where labeled data is limited, both SS-VFL algorithms provide immense benefits in communication reduction.
A Algorithm Pseudocode
1: Initialize: θ 0 m for all parties m and server model θ2: for t ← 0 , . .
, T − 1 do 3:
Parties choose randomly sampled mini-batch: B t .
4: for m ← 1 , . . .
, M in parallel do 5:
Party sends representation h m ( θ tm ; X B t m ) to server 6: end for 7: Φ t ← { θ 0 , h 1 ( θ t 1 ; X B t 1 ) , . .
, h M ( θ tM ; X B t M ) } 8: θ t + = θ t 0 − η t ∇ 0 F B (Φ t ; y B t ) 9:
Server sends ∇ h m ( θ tm ; X B tm ) F B (Φ t ; y B t ) to each party m 10: for m ← 1 , . .
, M in parallel do 11: ∇ m F
[Equation]
y
[Equation]
t ) 12: θ t
[Equation]
m = θ tm − η t ∇ m F B (Φ t ; y B t ) 13: end for 14: end for
In Algorithm 3, we provide the pseudocode of Supervised Vertical Federated Learning, as described in Section 2.
do 3:
Randomly sample m 4: θ
= θ
; B t ) 5: end for 6: end procedure
In Algorithm 4, we provide the pseudocode of local constrastive representation learning, as described in Section 3.
B Societal/Ethical impacts
With any machine learning task, there is the concern of the resulting model learning a bias that may discriminate against groups of people unfairly.
Although utilizing SS-VFL can mitigate model bias by including unlabeled data, biased data collection can still lead to model bias.
Additionally, a “protected" feature, such as race, may play a large role in the prediction model.
Researchers and developers must be aware of the potential for bias in the datasets and take steps to ensure that these biases do not lead to discriminatory practices.
Several techniques to correct for bias in datasets exist and can be applied to SS-VFL.
For example, one can apply weighting to the data samples to correct for a known bias in the dataset.