[Paper Review] Coresets for Vertical Federated Learning: Regularized Linear Regression and $K$-Means Clustering

Lingxiao Huang, Zhize Li+2 more|arXiv (Cornell University)|Oct 26, 2022

Privacy-Preserving Technologies in Data4 citations

TL;DR

This paper proposes a unified coreset construction framework for communication-efficient vertical federated learning (VFL), enabling sublinear communication complexity for regularized linear regression and K-means clustering. By leveraging distributed importance sampling and robust coreset theory, the method achieves ε-approximate solutions with coresets of size o(n), drastically reducing communication while preserving model quality under mild conditions.

ABSTRACT

Vertical federated learning (VFL), where data features are stored in multiple parties distributively, is an important area in machine learning. However, the communication complexity for VFL is typically very high. In this paper, we propose a unified framework by constructing coresets in a distributed fashion for communication-efficient VFL. We study two important learning tasks in the VFL setting: regularized linear regression and $k$-means clustering, and apply our coreset framework to both problems. We theoretically show that using coresets can drastically alleviate the communication complexity, while nearly maintain the solution quality. Numerical experiments are conducted to corroborate our theoretical findings.

Motivation & Objective

• Address the high communication complexity in vertical federated learning (VFL), where parties hold different feature subsets of the same data.
• Develop a scalable, communication-efficient training framework for VFL that maintains high model accuracy.
• Design a unified coreset construction method applicable to both regularized linear regression and K-means clustering in the VFL setting.
• Achieve ε-approximate solutions with coresets that scale sublinearly in dataset size n, reducing communication costs significantly.
• Ensure privacy by avoiding raw data transfer while maintaining solution quality through robust coreset sampling.

Proposed method

• Construct coresets via distributed importance sampling, where sampling probabilities are derived from local gradients and feature norms across parties.
• Use a robust coreset framework that ensures (β, ε)-approximation guarantees for both learning tasks under mild assumptions.
• Apply theoretical bounds from coreset literature to derive sample size m dependent on dimension d, k, ε, β, and confidence δ.
• Design a communication-efficient protocol where each party computes local statistics and transmits only coreset-related weights and indices.
• Leverage pseudo-dimension and function space analysis to bound approximation error in the coreset construction process.
• Ensure coreset size remains o(n) by controlling sampling complexity through bounds on g_i and cG terms in the theoretical framework.

Experimental results

Research questions

• RQ1 Can coresets be constructed in a distributed, communication-efficient manner for vertical federated learning?
• RQ2 Can the coreset framework achieve ε-approximate solutions for regularized linear regression in VFL with sublinear communication?
• RQ3 Can the same coreset framework be extended to K-means clustering in VFL with comparable communication savings?
• RQ4 How does the coreset size scale with respect to model complexity (d, k), approximation error ε, and confidence δ?
• RQ5 Does the coreset-based approach outperform uniform sampling in terms of solution quality under the same communication budget?

Key findings

• The proposed coreset framework achieves ε-approximation for regularized linear regression using a coreset size of O(d⁴ / (ε²β²T²) · (d² + log(1/δ))), independent of n.
• For K-means clustering, the coreset size is O(α²k⁴ / (ε²β²) · (dk + log(1/δ))), also sublinear in n and independent of dataset size.
• Communication complexity is reduced to O(mT), where m is the coreset size, achieving sublinear communication relative to n.
• Numerical experiments confirm that the coreset method maintains solution quality comparable to full-data training while drastically reducing communication.
• Compared to uniform sampling, the coreset approach achieves better model performance at the same or lower communication cost.
• Theoretical analysis proves that the coreset construction is robust under mild conditions, with high probability (1−δ) of maintaining approximation guarantees.