Variational Federated Multi-Task Learning

TL;DR

This paper introduces VIRTUAL, a novel federated multi-task learning algorithm for non-convex models, using variational inference to handle data heterogeneity across devices, and demonstrates its effectiveness on real-world datasets.

Contribution

The paper presents VIRTUAL, the first federated multi-task learning method for non-convex models utilizing variational inference, improving performance and efficiency.

Findings

Outperforms existing federated learning methods on real-world datasets.

Supports sparser gradient updates for efficiency.

Effective on non-convex models with statistical heterogeneity.

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.