FedOptimus: Optimizing Vertical Federated Learning for Scalability and Efficiency

2025-02-06

AI Summary

Vertical federated learning (VFL) faces critical challenges in privacy-sensitive domains (e.g., finance and healthcare), including high communication overhead, slow convergence under non-IID data, and poor system scalability. Method: This paper proposes the first end-to-end optimization framework for VFL, featuring a mutual information (MI)-driven client selection mechanism and a novel integration of server-side momentum (FedAvgM/SLOWMO) with K-step averaging—designed to jointly enhance convergence speed and distributed robustness without compromising model accuracy. Contribution/Results: Extensive experiments on CIFAR-10, MNIST, and FMNIST demonstrate that our framework reduces communication cost by 37%, accelerates convergence by 2.1× over state-of-the-art VFL methods, and supports scalable deployment across hundreds of clients and multiple servers. The framework is theoretically sound and practically deployable, bridging a key gap between rigorous analysis and real-world applicability in privacy-preserving distributed learning.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Computer Vision: Learning & Optimization for CV
• Search and Optimization: Distributed Search

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Security and Privacy: Large-scale security measurements
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

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.

Problem

• Optimizing Vertical Federated Learning
• Enhancing scalability and efficiency
• Addressing communication and computational challenges

Innovation

• Mutual Information client selection
• Server-side momentum techniques
• K-Step Averaging communication reduction