Source-linked AI summary

PyVertical: A Vertical Federated Learning Framework for Multi-headed SplitNN

Daniele Romanini, Adam James Hall, Pavlos Papadopoulos, Tom Titcombe, Abbas Ismail, Tudor Cebere, Robert Sandmann, Robin Roehm, Michael A. Hoeh

arXiv:2104.00489v3cs.LGcs.CRcs.DC

TL;DR

Sensitive or protected data held across separate organizations cannot simply be merged for machine learning. PyVertical combines SplitNNs with PSI to align shared records while training on vertically partitioned features, and the authors demonstrate the approach on a two-owner MNIST task. The work concludes that the framework and method work in principle, while noting assumptions and scope limits in the experiment.

  • Problem

    Data scientists face siloed datasets, and sensitive or legally protected data cannot simply be merged across parties.

    Source: §1 · p. 1
  • Method

    PyVertical combines SplitNNs and PSI so parties can align shared data IDs and train on vertically partitioned features while raw data remains with its owner.

    Source: §3 · p. 3, §3.1 · p. 4
  • Results

    The authors successfully train a dual-headed model on vertically partitioned MNIST, demonstrating that the framework and method work in principle.

    Source: §5 · p. 5
  • Takeaways & Limitations

    The open-source framework is presented as useful for VFL research, and PSI is considered applicable where datasets contain unique IDs.

    Source: §5 · p. 4

    The evaluation assumes honest parties, uses symmetric model segments and equally split data, and demonstrates only a two-owner example while broader multi-headed tasks remain future work.

    Source: §5.1 · p. 5

Abstract

from arXiv · show

We introduce PyVertical, a framework supporting vertical federated learning using split neural networks. The proposed framework allows a data scientist to train neural networks on data features vertically partitioned across multiple owners while keeping raw data on an owner's device. To link entities shared across different datasets' partitions, we use Private Set Intersection on IDs associated with data points. To demonstrate the validity of the proposed framework, we present the training of a simple dual-headed split neural network for a MNIST classification task, with data samples vertically distributed across two data owners and a data scientist.

1 INTRODUCTION

PyVertical addresses learning from vertically partitioned data that cannot simply be merged when it is sensitive or legally protected. It combines SplitNNs and PSI and is presented as an open-source VFL framework, verified in a two-owner MNIST scenario.

1 INTRODUCTION

  • Sensitive data held in separate silos cannot simply be merged, while vertical partitioning places different features of the same data subjects with multiple owners.The paper distinguishes this from horizontal federated learning, where owners hold the same features for different subjects.Source: §1 · pp. 1–2
  • SplitNNs map data to shareable abstract representations, while PSI links data points belonging to the same samples across parties.The paper combines these methods to facilitate VFL for non-linear functions.Source: §1 · p. 2

1.1 CONTRIBUTIONS

  • PyVertical is presented as an open-source framework for machine learning on vertically distributed datasets using SplitNNs and PSI.The authors describe it as the first open-source framework to perform this task.Source: §1.1 · p. 2
  • The framework is verified on vertically partitioned MNIST, with two data owners holding different parts of samples and a data scientist holding labels.The authors note that the same principle could extend to multiple data owners.Source: §1.1 · p. 2

2 BACKGROUND KNOWLEDGE AND RELATED WORK

The background introduces PSI for private entity matching, SplitNNs for distributed model training, and VFL for collaborating across feature-partitioned datasets. It situates PyVertical among prior VFL work and distinguishes its treatment of entity resolution.

2.1 PRIVATE SET INTERSECTION

  • PSI lets two parties compute the intersection of their sets without revealing anything beyond the intersecting elements.The paper uses a Diffie-Hellman-based implementation with Bloom-filter compression; it notes that an alternative can handle more than two parties.Source: §2.1 · p. 2

2.2 SPLIT NEURAL NETWORKS

  • A SplitNN divides a model into segments held by different parties, passing intermediate representations forward and gradients across segments during backpropagation.The paper also notes that split learning can reduce data owners’ computational burden.Source: §2.2 · p. 2

2.3 VERTICAL FEDERATED LEARNING

  • VFL trains collaboratively when parties hold different features for the same dataset; unlike the cited split-network work, PyVertical addresses entity resolution using PSI.Existing open-source FedML is described as implementing multi-party linear models.Source: §2.3 · p. 3

3 FRAMEWORK DESCRIPTION

PyVertical combines SplitNNs and PSI to train on vertically partitioned data while keeping raw data with its owners. Its data-resolution protocol aligns shared IDs, and the data scientist coordinates training across the model segments.

3 FRAMEWORK DESCRIPTION

  • PyVertical is a Python VFL framework built on PySyft, using SplitNNs and PSI with mechanisms such as data pointers to avoid exposing private information.The data scientist defines the model and sends segments to owners, whose segments map their samples to abstract representations; the scientist concatenates these outputs and applies the final segment.Source: §3 · p. 3
  • Owners agree on shared data IDs through PSI, discard non-shared records, and sort their datasets so each position corresponds to the same subject.IDs are based on subjects and use a format agreed by the owners.Source: §3 · p. 3

3.1 DATA RESOLUTION PROTOCOL

  • For the three-party setting, the data scientist runs PSI separately with each owner, computes the global intersection, and communicates it to the owners.The owners do not communicate with or learn each other’s identities; the scientist receives the individual intersections.Source: §3.1 · pp. 3–4
  • In the experiment layout, the data scientist holds labels and part of the SplitNN, while data owners hold image datasets and model parts.The scientist computes batch loss and sends gradients to owners, who complete backpropagation independently.Source: §3 · p. 3, Figure 1 · p. 4

4 EXPERIMENT

The experiment tests vertical learning with a dual-headed SplitNN on MNIST images divided between two owners. It is presented as a proof of concept rather than a task-optimized evaluation.

4 EXPERIMENT

  • MNIST images are split into left and right halves, with each half given to a different data owner.The owners receive identical multi-layer network segments that take flattened 392-length vectors as input.Source: §4 · p. 4
  • The data scientist keeps the classification softmax segment and labels, calculates batch loss, and controls training and hyperparameters.The experiment’s objective is to demonstrate that the framework allows vertically partitioned learning.Source: §4 · p. 4
  • The MNIST experiment is a proof of concept and is not highly optimized for the specific task.The paper refers readers to Appendix B for training details and Figure 4 for results.Source: §4 · p. 4

5 EVALUATION AND CONCLUSION

PyVertical trains a dual-headed SplitNN on vertically partitioned MNIST, demonstrating that the framework works in principle. The evaluation assumes honest parties and a symmetric setup; scaling to many owners and adding privacy protections remain future work.

5 EVALUATION AND CONCLUSION

  • PSI is presented as applicable where datasets already use unique IDs, such as usernames or national identifiers.Source: §5 · p. 4
  • The framework successfully trains a dual-headed SplitNN on vertically partitioned MNIST, demonstrating that the method works in principle.Source: §5 · p. 5

5.1 LIMITATIONS AND FUTURE WORK

  • The experiment assumes data owners and the data scientist act honestly; robust deployment also requires identity management, PSI-protocol validation, and agreement on an ID schema.Source: §5.1 · p. 5
  • The study uses identical owner model segments and equally split data; it illustrates two owners, while massively multi-headed tasks and additional privacy techniques remain future work.The proposed framework can support more parties in principle.Source: §5.1 · p. 5

A PYVERTICAL PROTOCOL

The protocol figures illustrate PyVertical on MNIST and the PSI-based linkage process. The data scientist performs a separate PSI with each owner, and the duplicated scientist icon represents one scientist whose PSI runs could be parallelized.

A PYVERTICAL PROTOCOL

B EXPERIMENTAL SETUP

The experimental setup trains an unoptimised dual-headed SplitNN on vertically partitioned MNIST. Owner networks encode each image half, and the data scientist combines their outputs for ten-class digit prediction.

B EXPERIMENTAL SETUP

  • The owner segment maps a 392-length input to a 64-length ReLU representation, while the data scientist concatenates two outputs into a 128-length vector for a 10-class softmax prediction.Source: B EXPERIMENTAL SETUP · p. 9
  • Training uses the first 20,000 MNIST images for 30 epochs, with batches of 128 and learning rates of 0.01 for owners and 0.1 for the data scientist.Source: B EXPERIMENTAL SETUP · p. 9
  • Figure 4 presents train and validation accuracy for the unoptimised dual-headed SplitNN on vertically partitioned MNIST.Source: Figure 4 · p. 9
Loading 2104.00489v3…