[Paper Review] Variational Federated Multi-Task Learning

Luca Corinzia, Joachim M. Buhmann+1 more|arXiv (Cornell University)|Jun 14, 2019

Privacy-Preserving Technologies in Data66 references121 citations

TL;DR

VIRTUAL introduces federated multi-task learning for non-convex models using a star-shaped Bayesian network and variational inference, achieving better performance and sparser updates than FedAvg/FedProx.

ABSTRACT

In federated learning, a central server coordinates the training of a single model on a massively distributed network of devices. This setting can be naturally extended to a multi-task learning framework, to handle real-world federated datasets that typically show strong statistical heterogeneity among devices. Despite federated multi-task learning being shown to be an effective paradigm for real-world datasets, it has been applied only on convex models. In this work, we introduce VIRTUAL, an algorithm for federated multi-task learning for general non-convex models. In VIRTUAL the federated network of the server and the clients is treated as a star-shaped Bayesian network, and learning is performed on the network using approximated variational inference. We show that this method is effective on real-world federated datasets, outperforming the current state-of-the-art for federated learning, and concurrently allowing sparser gradient updates.

Motivation & Objective

• Address the challenge of federated multi-task learning with strongly non-IID client data.
• Develop a non-convex federated MTL algorithm suitable for distributed networks.
• Leverage a Bayesian network and variational/inference techniques to share knowledge while preserving privacy.
• Enable sparse gradient updates to reduce communication costs.

Proposed method

• Model the server and clients as a star-shaped Bayesian network with shared and private parameters.
• Use an EP-like variational inference scheme to approximate the posterior over server and client parameters.
• Adopt a Gaussian mean-field posterior for server and client parameters to facilitate tractable updates.
• Compute client updates as deltas of their posterior factors and aggregate them at the server.
• Retain privacy by only sharing aggregated server posterior and the product of client deltas, not individual client factors.
• Allow a KL-divergence weighted free energy term to balance reconstruction and regularization between server and clients.

Experimental results

Research questions

• RQ1 Can federated multi-task learning be effectively extended to generic non-convex models in a federated setting?
• RQ2 How can variational inference be applied to a star-shaped server–client Bayesian network to enable knowledge transfer while handling non-IID data?
• RQ3 Does VIRTUAL outperform existing FL baselines (FedAvg, FedProx) on real-world federated datasets while enabling sparse updates?
• RQ4 What is the impact of KL-divergence weighting on learning performance and sparsity of updates?

Key findings

• VIRTUAL outperforms state-of-the-art federated learning baselines on several real-world datasets across multiple architectures.
• The framework enables sparser client updates while maintaining or improving accuracy.
• Client models can remain private and specialized to their data while still benefiting from shared server knowledge.
• Using a KL-divergence regularization multiplier beta allows tuning performance and sparsity.
• The method supports non-IID data and traditional FL challenges without requiring centralized data access.