FedOptimus: Optimizing Vertical Federated Learning for Scalability and Efficiency

TL;DR

FedOptimus is a novel framework that enhances the scalability and efficiency of vertical federated learning by integrating client selection, momentum techniques, and communication reduction strategies, demonstrating superior performance on benchmark datasets.

Contribution

Introduces FedOptimus, a comprehensive Multi-VFL framework combining MI-based client selection, server momentum, and K-step averaging to improve VFL efficiency and scalability.

Findings

Outperforms existing VFL methods on CIFAR-10, MNIST, FMNIST.

Reduces communication costs while maintaining accuracy.

Accelerates convergence in heterogeneous data environments.

Abstract

Federated learning (FL) is a collaborative machine learning paradigm which ensures data privacy by training models across distributed datasets without centralizing sensitive information. Vertical Federated Learning (VFL), a kind of FL training method, facilitates collaboration among participants with each client having received a different feature space of a shared user set. VFL thus, proves invaluable in privacy-sensitive domains such as finance and healthcare. Despite its inherent advantages, VFL faced challenges including communication bottlenecks, computational inefficiency, and slow convergence due to non-IID data distributions. This paper introduces FedOptimus, a robust Multi-VFL framework integrating advanced techniques for improved model efficiency and scalability. FedOptimus leverages a Mutual Information (MI)-based client selection to prioritize high-contribution participants, reducing computational overhead. Further, it incorporates server-side momentum techniques like FedAvgM and SLOWMO to stabilize updates and accelerate convergence on heterogeneous data. Additionally, performing K-Step Averaging minimizes communication costs while maintaining model performance. FedOptimus proves to be superior in performance on benchmark datasets such as CIFAR-10, MNIST, and FMNIST, showcasing its scalability and effectiveness in real-world multi-server, multi-client settings. By unifying advanced optimization methods, FedOptimus sets a new standard for efficient and scalable Vertical Federated Learning frameworks, paving the way for broader adoption in complex, privacy-sensitive domains.