PaperlayerView exact revision on arXiv

Source-linked AI summary

Variational Federated Multi-Task Learning

Luca Corinzia, Ami Beuret, Joachim M. Buhmann

arXiv:1906.06268v2cs.LGstat.ML

TL;DR

Federated multi-task learning seeks to handle strongly heterogeneous private client data, but prior applications were limited to convex models. The paper introduces VIRTUAL, which uses a star-shaped Bayesian network and approximated inference for general non-convex models. VIRTUAL outperforms federated-learning baselines on many real-world datasets while supporting sparse updates.

  • Problem

    Federated multi-task learning has addressed heterogeneous federated data, but its prior applications were limited to convex models.

  • Method

    VIRTUAL models the server and clients as a star-shaped Bayesian network and trains them with an expectation-propagation-like variational approximation.

  • Results

    VIRTUAL outperforms FedAvg and FedProx in almost all evaluated datasets except PMNIST, achieving up to +2% MT accuracy and +1% S accuracy.

  • Takeaways & Limitations

    VIRTUAL supports federated multi-task learning for generic non-convex models while allowing sparse updates with up to 75% induced sparsity and 50% communication-cost reduction versus FedProx.

Abstract

from arXiv · show

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.

I. INTRODUCTION

Federated learning addresses private, distributed data, but shared-model methods struggle when client distributions are strongly heterogeneous. VIRTUAL frames federated multi-task learning as a Bayesian network with shared and client-specific parameters.

  • Motivation: Remote devices generate valuable distributed data, but privacy, resource, communication, and statistical heterogeneity constrain federated learning.The supplied introduction identifies privacy-sensitive data, device storage and computation limits, wireless communication costs, and different client-dependent data distributions as central challenges.
  • Federated learning: Federated learning lets clients retain private data while a server coordinates model training through exchanged updates.Clients collect data and execute optimization locally; the server receives and sends updates rather than directly collecting the data.
  • Research gap: FedAvg requires all clients to share one model and performs poorly when client data distributions are strongly skewed.This motivates a multi-task formulation with models tailored to individual data distributions.
  • Contribution: The paper proposes federated multi-task learning for generic non-convex models using parameter sharing across client-specific models.The multi-task setup links models through shared parameters while allowing each client to fit its own data distribution.
  • VIRTUAL: VIRTUAL represents the server and clients as a star-shaped Bayesian network with shared server parameters and client-specific parameters.The formulation assigns server parameters θ and client parameters φ_i, with client likelihoods conditioned on both.

B. The optimization procedure

VIRTUAL optimizes a factorized variational posterior over shared server and client-specific parameters using an EP-like procedure. Client updates are aggregated as posterior ratios, retaining information from inactive clients while enabling efficient local training.

  • Variational approximation: VIRTUAL uses an EP-like approximation because the Bayesian posterior over server and client parameters is generally intractable.The proxy posterior factorizes into a server contribution and one client contribution per client.
  • Variational objective: The variational free energy combines client-side Bayes-by-backprop terms with server-side likelihood and KL regularization terms.The server is constrained toward a cavity distribution formed from the prior and the other clients’ posterior factors.
  • Continual-learning connection: The server’s cavity constraint makes its updates progress in a continual-learning fashion, helping avoid catastrophic forgetting of previously seen federated datasets.Client parameters receive no additional regularization and can therefore be trained efficiently in a network-agnostic way.
  • Posterior parameterization: Gaussian mean-field factors parameterize both server and client posteriors, with separate Gaussian parameters for each network component.The factorization uses server parameters and client-specific parameter collections.
  • Federated optimization: At each round, a random subset of clients trains locally, computes a distribution-ratio update, and sends that update to the server.The algorithm follows a FedAvg- and FedProx-like round structure with selected active clients and local free-energy optimization.
  • Federated aggregation: The server aggregates client updates as posterior information rather than averaging only active-client models, retaining non-active clients’ information.For mean-field Gaussians, the update generalizes FedAvg subtraction and averaging while accounting for uncertainty.

III. RELATED WORK

Related work spans distributed and federated optimization, federated multi-task learning, Bayesian distributed inference, and continual or transfer learning. VIRTUAL differs by targeting non-convex federated multi-task learning with shared and client-specific structure.

  • Distributed and federated learning: Distributed learning traditionally optimizes a generic model across computational clusters with centralized data and iterative averaging of locally trained models.Federated learning removes centralized-data, device-scale, and communication assumptions by coordinating updates from private clients.
  • Federated optimization: FedAvg selects online clients, performs local SGD, and averages their models at the server, while later heuristics address federated statistical challenges.Some alternatives share client data or server-trained generative models, but the cited discussion notes communication and privacy concerns.
  • Federated multi-task learning: MOCHA first embedded federated learning in multi-task learning, deriving a primal-dual method for convex models with multi-task regularization.Its experiments showed multi-task models outperforming global models in the reported setting.
  • Federated aggregation: Layer-wise weight matching methods address misaligned neural-network parameters before aggregation and can be extended to VIRTUAL.The motivation is to reduce averaging effects caused by overparameterization and parameter symmetry.
  • Bayesian distributed inference: Prior Bayesian distributed-inference work used expectation propagation but modeled only a global task and relied on heavy MCMC estimation.VIRTUAL instead includes shared and non-shared server-client parameters within a federated multi-task setting.
  • Transfer and continual learning: Continual-learning research uses regularized or Bayesian updates to retain performance on earlier tasks while learning new ones.VIRTUAL connects federated multi-task learning with this sequential knowledge-retention perspective.

IV. EXPERIMENTS

The experiments evaluate VIRTUAL on three real-world federated datasets representing both federated-training and multi-task-learning challenges. Simulated experiments use at most 100 clients because private client models make memory scale linearly with client count.

  • Experimental setting: VIRTUAL is evaluated empirically on three real-world federated datasets selected to represent federated-training and multi-task-learning challenges.The evaluation is performed on a simulated network using a single GPU.
  • Experimental setting: The simulations cap experiments at 100 clients because retaining a private model for every client makes memory cost scale linearly with client count.FedAvg has constant simulation memory cost under the comparison described in the passage.

A. Dataset description

The dataset suite covers handwriting, sensor, activity-recognition, digit, permuted-digit, and character-prediction tasks with varied client heterogeneity. It includes both relatively homogeneous and strongly non-IID federated settings.

  • Real-world datasets: FEMNIST uses 100 writers as clients and restricts the federated EMNIST data to ten digit labels.The train-test split is supplied by the dataset distribution.
  • Real-world datasets: VSN assigns 23 vehicle sensors to clients for binary classification using 50 acoustic and 50 seismic features.The classes are assault amphibious and dragon wagon vehicles.
  • Real-world datasets: HAR assigns 30 subjects to clients and classifies 12 daily activities from 561-dimensional smartphone inertial-sensor features.The passage specifies a 75%-25% train-test split for HAR and VSN.
  • Synthetic federated datasets: MNIST is split into 100 clients with 600 training and 100 test samples each, representing an atypically homogeneous federated dataset.Clients are homogeneous in both dataset size and sample statistics.
  • Synthetic federated datasets: PMNIST uses the same 100-client split but independently permutes pixels, producing strongly non-IID clients with dissimilar low-level features.The construction comes from continual-learning evaluation practice.
  • Language dataset: Shakespeare assigns play roles to clients for next-character prediction over an 86-character vocabulary using sequences of 80.Sequences are aggregated into batches of ten, excluding roles without a complete batch.

B. Experimental setting

VIRTUAL is evaluated with neural architectures suited to image and NLP tasks, using fixed federated-training settings and tuned optimization hyperparameters. The FEMNIST KL-regularization study uses moving-average loss curves on logarithmic scales.

  • Experiments use two-hidden-layer MLPs with 100-unit ReLU layers, softmax outputs, and Bayesian counterparts using local reparameterization.
  • NLP experiments use two-layer LSTM classifiers with 100 hidden units per layer and an 8D embedding layer.
  • Each federated round updates C = 10 clients for E = 20 local epochs, with vanilla SGD used for every client.
  • For FEMNIST, Figure 2 plots server and multi-task cross-entropy during training across KL multipliers β, using 20-step moving averages and logarithmic y-axis and color scales.
  • Hyperparameters are tuned over five log-spaced client learning rates, while the implementation uses TensorFlow and TensorFlow Distributions.

C. Metrics

The evaluation distinguishes centralized server performance from personalized multi-task performance and compares VIRTUAL with FedAvg and FedProx across datasets, architectures, and sparsity settings. Results indicate more stable but delayed MT metrics, strong accuracy performance, and retained performance under substantial update sparsity.

  • Metrics: Centralized metrics test the server model on all client test data, whereas MT metrics average client-model test performance weighted by client dataset size.
  • KL-divergence weight: For FEMNIST, β values from 10^-6 to 10^-3 do not impair performance, while β = 10^-5 gives the best reported MT generalization in the adequate range.
  • Metric comparison: MT loss is typically more stable but delayed than server loss because clients retain private models and are updated less frequently.
  • Metric comparison: At convergence, MT metrics are typically superior to centralized metrics, indicating that clients can personalize models to private datasets.
  • Method comparison: VIRTUAL outperforms FedAvg and FedProx in almost all datasets and architectures except PMNIST, achieving up to +2% MT and +1% S accuracy in selected datasets.
  • Sparse updates: Without server initialization, clients specialize with fewer high-SNR weights, enabling pruning of updates below an SNR percentile while retaining superior performance to FedProx up to 75% sparsity.
  • Sparse updates: The 75% sparsity result corresponds to a 50% communication-cost reduction relative to FedProx after accounting for VIRTUAL's doubled Bayesian parameter count.

V. CONCLUSION

VIRTUAL addresses federated multitask learning by modeling the server and clients as a Bayesian network and using approximate inference. The method outperforms the state of the art across IID and non-IID real-world federated datasets, while several extensions remain open.

  • V. CONCLUSION: VIRTUAL treats the federation as a Bayesian network and trains it using approximated variational inference.The server receives an aggregated posterior update over shared parameters.
  • V. CONCLUSION: VIRTUAL outperforms the state of the art across many IID and non-IID real-world federated datasets.
  • V. CONCLUSION: Synchronous multi-client updates remain an open direction, including empirical study of outdated priors and new Bayesian models.
  • V. CONCLUSION: Future work includes changing network design choices and studying VIRTUAL under memory constraints with pruning, sparsity losses, or improved data storage.

APPENDIX A ADDITIONAL EXPERIMENTS

The appendix reports additional learning curves and a high-computation FEMNIST experiment, alongside implementation details for Gaussian products and ratios used by VIRTUAL. With E = 100 epochs per round, VIRTUAL outperforms FedAvg and FedProx and improves reported accuracies.

  • Additional experiments: The appendix adds training cross-entropy and server, multitask, and training accuracy curves for the evaluated datasets and FEMNIST architectures.
  • Additional experiments: With E = 100 epochs per round on FEMNIST, VIRTUAL outperforms FedAvg and FedProx for both MLP and ConvNet architectures.This experiment evaluates applicability under higher node computational load.
  • Implementation details: VIRTUAL implements factorized Gaussian updates through univariate Gaussian products and ratios.Natural-parameter addition and subtraction provide the corresponding product and ratio operations.

APPENDIX C OMITTED PROOFS

The omitted proofs derive VIRTUAL's posterior updates using factorized server and client distributions and an EP-like KL-minimization procedure. Additional figures report learning curves, including FEMNIST under E = 100 epochs per round.

  • Omitted proofs: The proof represents global server and client distributions as products of their corresponding factors.
  • Omitted proofs: VIRTUAL's EP-like update is obtained by minimizing a KL divergence with respect to the server and client factors.
  • Additional figures: The appendix includes additional learning curves and FEMNIST curves with E = 100 epochs per round.
  • Omitted proofs: Bayes rule and normalization connect the variational update to the client likelihood and prior distributions.

APPENDIX D FULL EXPERIMENT DETAILS

The experiments use fixed batch sizes, grid searches for regularization and learning rates, and a damping factor for VIRTUAL. The damping factor is fixed as 1 −ηs to match the baselines' number of hyperparameters.

  • Hyperparameters: Experiments use batch size B = 20, except for Shakespeare, where B = 10.
  • Hyperparameters: VIRTUAL adds a damping factor γ ∈[0, 1] to client updates to prevent oscillations in message passing.
  • Hyperparameters: The damping factor is fixed as 1 −ηs to retain the same number of hyperparameters as FedProx and FedAvg.
Loading 1906.06268v2…