# Scalable Vertical Federated Learning via Data Augmentation and Amortized Inference

arXiv: https://arxiv.org/abs/2405.04043 · HTML: https://mlsanity.com/paper/2405.04043/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 47 references that are in the map.
- First posted on arXiv 7 May 2024.

## Builds on (most cited first, 23 of 23)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [SplitFed: When Federated Learning Meets Split Learning](https://mlsanity.com/paper/2004.12088/) (2020)
- [Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks](https://mlsanity.com/paper/2002.10118/) (2020)
- [Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro](https://mlsanity.com/paper/1912.11554/) (2019)
- [Private federated learning on vertically partitioned data via entity resolution and additively homomorphic encryption](https://mlsanity.com/paper/1711.10677/) (2017)
- [Epistemic Neural Networks](https://mlsanity.com/paper/2107.08924/) (2021)
- [Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference](https://mlsanity.com/paper/1703.09194/) (2017)
- [Hierarchical Variational Models](https://mlsanity.com/paper/1511.02386/) (2015)
- [Split Learning for collaborative deep learning in healthcare](https://mlsanity.com/paper/1912.12115/) (2019)
- [Do Bayesian Neural Networks Need To Be Fully Stochastic?](https://mlsanity.com/paper/2211.06291/) (2022)
- [Differentially Private Secure Multi-Party Computation for Federated Learning in Financial Applications](https://mlsanity.com/paper/2010.05867/) (2020)
- [Vertical Federated Learning: Challenges, Methodologies and Experiments](https://mlsanity.com/paper/2202.04309/) (2022)
- [Split-and-augmented Gibbs sampler - Application to large-scale inference problems](https://mlsanity.com/paper/1804.05809/) (2018)
- [SplitNN-driven Vertical Partitioning](https://mlsanity.com/paper/2008.04137/) (2020)
- [FedPop: A Bayesian Approach for Personalised Federated Learning](https://mlsanity.com/paper/2206.03611/) (2022)
- [Particle algorithms for maximum likelihood training of latent variable models](https://mlsanity.com/paper/2204.12965/) (2022)
- [Efficient MCMC Sampling with Dimension-Free Convergence Rate using ADMM-type Splitting](https://mlsanity.com/paper/1905.11937/) (2019)
- [Asymptotically exact data augmentation: models, properties and algorithms](https://mlsanity.com/paper/1902.05754/) (2019)
- [Stochastic Approximation with Biased MCMC for Expectation Maximization](https://mlsanity.com/paper/2402.17870/) (2024)

## Built on by (most cited first, 2 of 2)

- [Federated Generalised Variational Inference: A Robust Probabilistic Federated Learning Framework](https://mlsanity.com/paper/2502.00846/) (2025)
- [A Bayesian Vertical Federated Learning Framework for Multivariate Reduced-Rank High-Dimensional Regression](https://mlsanity.com/paper/2609.22654/) (2026)

## Abstract (arXiv)

Vertical federated learning (VFL) has emerged as a paradigm for collaborative model estimation across multiple clients, each holding a distinct set of covariates. This paper introduces the first comprehensive framework for fitting Bayesian models in the VFL setting. We propose a novel approach that leverages data augmentation techniques to transform VFL problems into a form compatible with existing Bayesian federated learning algorithms. We present an innovative model formulation for specific VFL scenarios where the joint likelihood factorizes into a product of client-specific likelihoods. To mitigate the dimensionality challenge posed by data augmentation, which scales with the number of observations and clients, we develop a factorized amortized variational approximation that achieves scalability independent of the number of observations. We showcase the efficacy of our framework through extensive numerical experiments on logistic regression, multilevel regression, and a novel hierarchical Bayesian split neural net model. Our work paves the way for privacy-preserving, decentralized Bayesian inference in vertically partitioned data scenarios, opening up new avenues for research and applications in various domains.