Vertical Federated Learning RFC

[rongou]

Motivation

XGBoost 1.7.0 introduced the initial support for Federated Learning. However, only horizontal federated learning is supported. Training samples are assumed to be split horizontally, i.e. each participant has a subset of samples, but with all the features and labels. In many real world applications, data is split vertically: each participant has all the samples, but only a partial list of features, and not all participants have access to the labels. It would be useful to support vertical federated learning in XGBoost.

Goals

• Enhance XGBoost to support vertical federated learning.
• Support using NVFlare to coordinate the learning process, but the design should be amenable to support other federated learning platforms.
• Efficiency: training speed should be close to traditional distributed training.
• Accuracy: should be close to centralized learning.

Non-Goals

• Initially we will assume the federated environment is non-adversarial, and will not provide strong privacy guarantees. This will be improved upon in later iterations.
• We will not support data owners dropping out during the learning process.

Assumptions

In vertical federated learning, before model training, the participants need to jointly compute the common IDs of their private sets. This is called private set intersection (PSI). For XGBoost, we assume this is already done, and users may rely on some other library/framework for this step.

Similar to horizontal federated learning, we make some simplifying assumptions::

• A few trusted partners jointly train a model.
• Reasonably fast network connection between each participant and a central trusted party.

Risks

The current XGBoost codebase is fairly complicated and hard to modify. Some code refactoring needs to happen first, before support for vertical federated learning can be added. Care must be taken to not break existing functionality, or make regular training harder.

Design

LightGBM, a gradient boosting library similar to XGBoost, supports “feature parallel” distributed learning.

Conceptually, feature parallelism is similar to vertical federated learning. A possible design is to first enhance XGBoost distributed training to support feature parallelism, and then build vertical federated learning on top of it. This would benefit the wider user community, thus greatly reducing the risks involved in refactoring XGBoost’s code base.

Feature Parallelism

XGBoost has an internal training parameter called DataSplitMode, which can be set to auto, col, and row. However, it’s currently not exposed to the end user, and can only be set to row for distributed training.

In order to support column-based data split for distributed training, we need to do the following:

• When initially loading data, support splitting by column. To keep it simple, we can have all workers keep a copy of the labels (and other things like weight and qid, effectively the MetaInfo object). The resulting DMatrix needs to keep track of which features belong to which worker.
• When generating the prediction, participants need to work collaboratively: the worker owning the feature used to split a node needs to collect the left and right splits and broadcast the results. A naive implementation may incur too much communication overhead, but there is prior work to encode partial predictions in bitsets to make the process more efficient (see paper).
• The worker owning the label should calculate the gradients and broadcast them to other workers.
• When finding the best split, each worker finds the best local split based on the features it owns, and then performs an allreduce to find the global best split. Workers don’t need to access histogram from each other. The worker owning the feature for the best split then broadcasts the split results to others.

We may also want to consider implementing LightGBM’s voting parallel approach (paper) for more efficient communication.

Vertical Federated Learning

Assuming feature parallelism is implemented, vertical federated learning is a slight modification:

• When loading data, no need to split the columns further, since we assume each worker only has a subset of the features.
• We can no longer share labels between workers.
• Communication needs to switch to the federated communicator.

Federated Inference

In horizontal federated learning, since each participant has all the features and labels, trained models can be shared to run inference/prediction locally. In vertical federated learning, however, this is no longer feasible. All the participants need to be online and work collaboratively to run inference. For batch inference, a federated learning job can be set up (for example, using NVFlare) to produce predictions. For online inference, we would need to set up services at participating sites to jointly produce online predictions, which is out of the scope of this design.

Alternatives Considered

It may be possible to implement vertical federated learning without first adding support for column data split mode in distributed training. However, since this would require extensive refactoring of the XGBoost code base without producing any benefits for users not using federated learning, it may be too risky. For very wide datasets (relatively many features with moderate number of rows), column data split may be a useful feature for distributed training.

[trivialfis]

Thank you for the detailed RFC! I'm sure there are lots of excitement coming. I tagged the RFC as part of 2.0. Let me know if there's anything I can help.

[rongou]

Adding a task list to keep track of the progress:

• Initial support for column-wise data split #8468
• Column-wise data splitting for external memory
• Add an in-memory collective communicator #8494
• Add data split mode to DMatrix MetaInfo #8568
• Support bitwise allreduce operations in the communicator #8623
• Initial support for column-split cpu predictor #8676
• More tests for cpu predictor with column split #9270
• Use in-memory communicator to test quantile #8710
• Support cpu quantile sketch with column-wise data split #8742
• Add Allgather to collective communicator #8765
• Support column split in histogram builder #8811
• Support column-split in row partitioner #8828
• Support column split in approx tree method #8847
• Don't populate labels if label column is not specified in csv parser dmlc-core#679
• Support vertical federated learning #8932
• More tests for column split and vertical federated learning #8985
• Make objectives work with vertical distributed and federated learning #9002
• Make sure metrics work with column-wise distributed training #9020
• Make sure metrics work with federated learning #9037
• More collective aggregators #9060
• More refactoring to take advantage of collective aggregators #9081
• Add demo for vertical federated learning #9103
• Support hist in the partition builder under column split #9120
• Support column split in multi-target hist #9171
• More support for column split in cpu predictor #9244
• Refactor device communicator to make allreduce more flexible #9295
• Support bitwise allreduce in NCCL communicator #9300
• Support column split in GPU predictor #9343
• Support column split with GPU quantile #9370
• Support column split in GPU evaluate splits #9511
• Support column split in gpu hist updater #9384
• Support vertical federated learning with gpu_hist #9539
• Add GPU support to NVFlare vertical demo #9552
• Add tests for gpu_approx with column split and vertical federated #9553
• More support for column split in gpu predictor #9562
• Test categorical features with column-split gpu quantile #9595
• Test column sampler with column-wise data split #9609
• Integration tests for interaction constraints with column-wise data split #9611
• Test monotone constraints with column split #9613
• Add python api to slice columns in DMatrix #9620

[peiji1981]

hello, Mr.rongou @rongou . I have much experience on the implementation vertical federated XGB/LGB. But my former version was coded by python and scala. I was trying to implement it based on lightgbm source codes, also added some c_api for python wrapper, but this work is very hard for me because of poor c++ coding experience. So how can I contribute my experience with you or your group to implement this work based on xgb src?

[trivialfis]

Hi @peiji1981 , thank you for the offer! Do you have a paper or document on your work?