Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Vertical Federated Learning with Missing Features During Training and Inference

Pedro Valdeira, Shiqiang Wang, Yuejie Chi

TL;DR

This work tackles missing feature blocks in vertical federated learning by introducing LASER-VFL, a method that shares representation and fusion-model parameters across a tractable family of predictors and employs task-sampling to cover many subsets of observed features without exponential cost. The approach enables training and inference with any subset of feature blocks, provides rigorous convergence guarantees for nonconvex objectives and, under PL, linear convergence to a neighborhood of the optimum, and demonstrates strong empirical gains across diverse datasets and missing-data patterns. LASER-VFL further shows robustness to missing features, scalable performance with increasing numbers of clients, and even advantages when all features are available, due to a dropout-like regularization effect from task sampling. The work thus broadens the practical applicability of VFL by ensuring data efficiency and resilience to client dropouts and incomplete observations, with public code available for replication.

Abstract

Vertical federated learning trains models from feature-partitioned datasets across multiple clients, who collaborate without sharing their local data. Standard approaches assume that all feature partitions are available during both training and inference. Yet, in practice, this assumption rarely holds, as for many samples only a subset of the clients observe their partition. However, not utilizing incomplete samples during training harms generalization, and not supporting them during inference limits the utility of the model. Moreover, if any client leaves the federation after training, its partition becomes unavailable, rendering the learned model unusable. Missing feature blocks are therefore a key challenge limiting the applicability of vertical federated learning in real-world scenarios. To address this, we propose LASER-VFL, a vertical federated learning method for efficient training and inference of split neural network-based models that is capable of handling arbitrary sets of partitions. Our approach is simple yet effective, relying on the sharing of model parameters and on task-sampling to train a family of predictors. We show that LASER-VFL achieves a $\mathcal{O}({1}/{\sqrt{T}})$ convergence rate for nonconvex objectives and, under the Polyak-Łojasiewicz inequality, it achieves linear convergence to a neighborhood of the optimum. Numerical experiments show improved performance of LASER-VFL over the baselines. Remarkably, this is the case even in the absence of missing features. For example, for CIFAR-100, we see an improvement in accuracy of $19.3\%$ when each of four feature blocks is observed with a probability of 0.5 and of $9.5\%$ when all features are observed. The code for this work is available at https://github.com/Valdeira/LASER-VFL.

Vertical Federated Learning with Missing Features During Training and Inference

TL;DR

This work tackles missing feature blocks in vertical federated learning by introducing LASER-VFL, a method that shares representation and fusion-model parameters across a tractable family of predictors and employs task-sampling to cover many subsets of observed features without exponential cost. The approach enables training and inference with any subset of feature blocks, provides rigorous convergence guarantees for nonconvex objectives and, under PL, linear convergence to a neighborhood of the optimum, and demonstrates strong empirical gains across diverse datasets and missing-data patterns. LASER-VFL further shows robustness to missing features, scalable performance with increasing numbers of clients, and even advantages when all features are available, due to a dropout-like regularization effect from task sampling. The work thus broadens the practical applicability of VFL by ensuring data efficiency and resilience to client dropouts and incomplete observations, with public code available for replication.

Abstract

Vertical federated learning trains models from feature-partitioned datasets across multiple clients, who collaborate without sharing their local data. Standard approaches assume that all feature partitions are available during both training and inference. Yet, in practice, this assumption rarely holds, as for many samples only a subset of the clients observe their partition. However, not utilizing incomplete samples during training harms generalization, and not supporting them during inference limits the utility of the model. Moreover, if any client leaves the federation after training, its partition becomes unavailable, rendering the learned model unusable. Missing feature blocks are therefore a key challenge limiting the applicability of vertical federated learning in real-world scenarios. To address this, we propose LASER-VFL, a vertical federated learning method for efficient training and inference of split neural network-based models that is capable of handling arbitrary sets of partitions. Our approach is simple yet effective, relying on the sharing of model parameters and on task-sampling to train a family of predictors. We show that LASER-VFL achieves a convergence rate for nonconvex objectives and, under the Polyak-Łojasiewicz inequality, it achieves linear convergence to a neighborhood of the optimum. Numerical experiments show improved performance of LASER-VFL over the baselines. Remarkably, this is the case even in the absence of missing features. For example, for CIFAR-100, we see an improvement in accuracy of when each of four feature blocks is observed with a probability of 0.5 and of when all features are observed. The code for this work is available at https://github.com/Valdeira/LASER-VFL.

Paper Structure

This paper contains 38 sections, 3 theorems, 59 equations, 3 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Generalization and availability in standard VFL.
  3. Dealing with missing features.
  4. Our contributions.
  5. Related work.
  6. Definitions and Preliminaries
  7. Our method
  8. Key idea.
  9. A tractable family of predictors sharing model parameters
  10. Sharing model parameters.
  11. Efficient training via task-sampling
  12. Our optimization problem.
  13. Our optimization method.
  14. Block availability-aware entity alignment.
  15. Extensions to our method.
  16. ...and 23 more sections

Key Result

Lemma 1

If the mini-batch gradient estimate is unbiased eq:unbiased, then, for all $t\geq0$: where the expectation is with respect to mini-batch $\mathcal{B}^t$ and the sampled set of tasks $\mathcal{S}^t$.

Figures (3)

  • Figure 1: In Figure \ref{['fig:splitnn']}, we illustrate a split neural network. In Figure \ref{['fig:data-availability-and-usage']}, we show the availability of a dataset with $K=3$ blocks of features, each with a $0.5$ probability of being observed, and its usage and waste by three key methods: standard VFL, a local approach, and our method. Standard VFL trains a single predictor, while the local approach and our method train different predictors at different clients (we show the data usage for client 1). In Figure \ref{['fig:local-vs-svfl']}, we present a preview of our results: a mortality prediction task using the MIMIC-IV dataset, when the probability of each block of features missing, $p_{\text{miss}}$, is in $\{0.0, 0.5\}$, for both data and test data.
  • Figure 2: Performance and runtime across different numbers of clients (CIFAR-10, $p_{\text{miss}}=0.1$ for training and inference). We did not run Combinatorial for $8$ clients due to resource constraints.
  • Figure 3: Diagram illustrating a forward pass of LASER-VFL.

Theorems & Definitions (3)

  • Lemma 1: Unbiased update vector
  • Theorem 1: Main result
  • Theorem 2: Linear convergence