Multi-VFL: A Vertical Federated Learning System for Multiple Data and Label Owners

Deep learning models benefit from large training datasets. However, for high-quality models, training needs to happen on large diverse datasets which usually are spread across multiple organizations. Deep learning solutions are extremely useful for for several domains including healthcare organizations and financial companies. While dealing with sensitive data (e.g.) patient information, credit card details, etc., government regulations such as GDPR (Voigt and Von dem Bussche, 2017) and CCPA (Pardau, 2018) restrict how entities can share data with other organizations. To solve this problem, Federated Learning (FL) was proposed. FL (Konečnỳ et al., 2016; McMahan et al., 2017) is a collaborative learning paradigm that allows distributed clients to train machine learning models without having to share their sensitive data.

There are two main categories of FL. In horizontal FL, clients have the same set of features for a different set of clients, whereas, vertical FL handles the case when multiple clients possess the data of the same set of individuals, but each client has a unique set of features. There has been a lot of research in horizontal FL, however, there are few approaches for vertical FL (Liu et al., 2019; Nock et al., 2018). Split learning is the most commonly used approach to address vertical FL (Vepakomma et al., 2018; Ceballos et al., 2020; Angelou et al., 2020). Split learning is a technique to collaboratively train a deep learning network that is split across different clients.

Usually, vertical FL is used to solve the problem when dataset features are distributed across multiple data owners and there exists a single label owner. However, in reality, there might be multiple label owners. In this paper, we consider the setting where $D$ data owners (across which dataset features are distributed) and $K$ label owners (across which dataset labels are distributed) exist. For example, in the case of healthcare, the $D$ data owners may correspond to different specialist hospitals (heart disease hospital, lung disease hospital, cancer hospital, etc.) that have different data for the same set of patients and $K$ label owners correspond to COVID-testing centers that have the COVID results corresponding to these patients.

To solve this problem, we make use of split learning and adaptive optimizer-based horizontal FL algorithms. Split learning solves the problem when $D$ data owners and $1$ label owner exists. However, when we have to aggregate the results from $K$ label owners, we need an aggregation mechanism. The most widely used FL algorithm for aggregation is FedAvg (McMahan et al., 2017), where the weights of the models obtained from different label owners are just averaged. FedAvg is successful when label owners have IID (independent and identically distributed) data. However, in the real-world, label owners have non-IID data leading to client drift (Hsu et al., 2019), which is what happens when gradients computed by different label owners are skewed causing local models to move away from globally optimal models. This substantially affects the performance of FedAvg, especially in scenarios where label owners have a high degree of variance in their data distributions. Under such scenarios, adaptive learning rates and momentum are beneficial as they incorporate knowledge from prior iterations. When participating label owners have sparse data distributions and possess a limited subset of labels, SGD with server momentum improves the FL performance (Hsu et al., 2019). (Reddi et al., 2020) proposed federated versions of adaptive optimizers and showed that these optimizers improved the performance of FL under non-iid settings. To improve the performance of Adam in local settings, (Chen et al., 2019) applied a decaying momentum (Demon) rule to Adam. We extend this approach to the federated setting.

An entity resolution protocol such as Private Set Intersection (PSI) (Angelou et al., 2020) is used to identify intersecting sets of individuals across different data owners and label owners by comparing the encrypted versions of the sets.

Vertical FL is a distributed learning paradigm on vertically partitioned datasets. Conventional approaches for vertical FL make use of Multi-Party Computation and Homomorphic Encryption (Hardy et al., 2017). However, these approaches suffer from performance and communication problems. To address these issues, split learning was proposed. The primary benefit of using split learning is that participants don’t have to share their sensitive data with each other while learning the shared model.

(Ceballos et al., 2020) proposed a VFL technique using split neural networks for $D$ data owners and $1$ label owner. In their approach, each data owner trains a partial network up to a cut layer and sends the activations (outputs) of the cut layer to the label owner. The label owner concatenates the activations coming from different data owners and completes the forward pass. As the label owner has the labels corresponding to the individuals on the datasets possessed by the data owners, it carries out its backpropagation up to the cut layer. The gradients at the cut layer are split and passed to the corresponding data owners. This process repeats for a fixed number of rounds or till convergence is achieved.

SplitFed, an approach that combined FL and split learning, was proposed by (Thapa et al., 2020). Their framework was proposed for $D$ data owners and $1$ label owner as well. In their framework, a 3rd party server aggregates the models of different data owners using Federated Averaging (FedAvg). In addition, data owners update their models sequentially, hence increasing the total time per federated round. Also, they expect data owners to have the same model architecture which is a hard constraint on resource-deficit data owners.

Both papers consider $D$ data owners and $1$ label owner in their setting. In this paper, we propose a solution to train vertical FL models in which $D$ data owners and $K$ label owners exist.

The primary goal of our framework is to learn an optimized vertical FL model when $D$ data owners and $K$ label owners exist.

In our framework, each data owner has a unique model architecture and all label owners have the same model architecture. Data owners perform forward propagation on their respective partial neural networks up to the cut layer and send their activations to label owners. Label owners concatenate the activations coming from different data owners and complete their forward pass. Label owners compute the loss, perform back-propagation, and send the corresponding gradients to data owners using which they complete their back-propagation. Label owners send their weights to the aggregation server which aggregates their weights using one of the following techniques: FedAvg (McMahan et al., 2017), FedAdam (Reddi et al., 2020), FedYogi (Reddi et al., 2020), etc. We also extend FedAdam to FedDemonAdam, a federated learning algorithm in which ”Demon(Chen et al., 2019) to Adam (DemonAdam)” is used for the server updates. The aggregation server returns the updated weights to the label owners which they use in the next round. The variables used in the Multi-VFL algorithm (Algorithm 1) are provided in Table 1.

An example real-life use case is depicted in Figure 1. Let us consider two hospitals (data owners) and two COVID testing centers (label owners). Our goal is to develop a model that would be able to predict how likely a patient is to get COVID based on the information available from cancer and lung disease hospitals. Algorithm 1 provides an optimized solution for this scenario.

To verify the validity of our proposed model, we ran experiments on two major datasets: MNIST (LeCun et al., 1998) and FashionMNIST (Xiao et al., 2017). Both datasets have 60,000 training data points and 10,000 testing data points with each image as $28\times 28$. For our experiments, we consider 5 label owners and 4 data owners. The datasets (excluding labels) are vertically partitioned among the data owners, with each of them receiving images of dimension $(\frac{28}{4}=7)\times 28$. The dataset labels are with the label owners. Each data owner has one convolutional layer with input and output channels as 1 and 32 respectively, and kernel size as 3. The output from the data owners are aggregated to form a tensor with 32 channels which are then passed to the label owner’s model which has one convolutional and two linear layers. All the experiments are run for 500 epochs with a batch size of 64. The local learning rate was set to 0.001. We study four major federated learning aggregation algorithms: FedAvg (Base Case), FedAdam, FedYogi, and FedDemonAdam. The values of the hyperparameter for these algorithms are summarized in Table 2.

We also allow different label owners to have different labels, thus mimicking the non-iid nature of the real world. Each label owner is given 5000 data points to train upon. However, some label owners are given an iid dataset, meaning all labels from 0 to 9 are given, while others are given only two labels each. We primarily considered 4 different ways to distribute labels and they are formulated in Table 3.

To the best of our knowledge, Multi-VFL is the first solution to address vertical federated learning when $D$ data owners and $K$ label owners exist. In addition, we ran experiments on the MNIST and FashionMNIST datasets for different non-IID label distribution scenarios and demonstrated the importance of the adaptive optimizer based aggregation. We plan to integrate differential privacy (Dwork, 2008) into our framework to thwart potential model inversion attacks and plan to conduct a detailed analysis on the privacy-accuracy trade-off.