[Paper Review] Multi-Task Federated Learning for Personalised Deep Neural Networks in Edge Computing

Jed Mills, Jia Hu+1 more|arXiv (Cornell University)|Jul 17, 2020

Privacy-Preserving Technologies in DataComputer Science40 references17 citations

TL;DR

This paper proposes Multi-Task Federated Learning (MTFL), which enhances personalized deep neural networks in edge computing by retaining Batch-Normalization (BN) layers locally on clients, enabling faster convergence and higher user model accuracy (UA) without sharing sensitive model parameters. MTFL reduces training rounds by up to 5× compared to standard FedAvg and up to 3× with FedAvg-Adam, outperforming state-of-the-art personalized FL methods.

ABSTRACT

Federated Learning (FL) is an emerging approach for collaboratively training Deep Neural Networks (DNNs) on mobile devices, without private user data leaving the devices. Previous works have shown that non-Independent and Identically Distributed (non-IID) user data harms the convergence speed of the FL algorithms. Furthermore, most existing work on FL measures global-model accuracy, but in many cases, such as user content-recommendation, improving individual User model Accuracy (UA) is the real objective. To address these issues, we propose a Multi-Task FL (MTFL) algorithm that introduces non-federated Batch-Normalization (BN) layers into the federated DNN. MTFL benefits UA and convergence speed by allowing users to train models personalised to their own data. MTFL is compatible with popular iterative FL optimisation algorithms such as Federated Averaging (FedAvg), and we show empirically that a distributed form of Adam optimisation (FedAvg-Adam) benefits convergence speed even further when used as the optimisation strategy within MTFL. Experiments using MNIST and CIFAR10 demonstrate that MTFL is able to significantly reduce the number of rounds required to reach a target UA, by up to $5 imes$ when using existing FL optimisation strategies, and with a further $3 imes$ improvement when using FedAvg-Adam. We compare MTFL to competing personalised FL algorithms, showing that it is able to achieve the best UA for MNIST and CIFAR10 in all considered scenarios. Finally, we evaluate MTFL with FedAvg-Adam on an edge-computing testbed, showing that its convergence and UA benefits outweigh its overhead.

Motivation & Objective

• To address the challenge of non-IID data in Federated Learning (FL), which hinders convergence and global model performance.
• To improve individual user model accuracy (UA), which is often the real objective in applications like content recommendation, rather than global model accuracy.
• To reduce the communication and storage overhead of personalized FL by avoiding full model uploads or extra training steps.
• To maintain data privacy by keeping sensitive model parameters, such as BN statistics, private to each client.
• To evaluate the effectiveness of MTFL in a realistic edge-computing testbed with heterogeneous devices.

Proposed method

• MTFL introduces non-federated Batch-Normalization (BN) layers that are trained locally on each client and not shared with the server, enabling model personalization.
• The method uses a multi-task learning framework where each client optimizes its own BN parameters while sharing the rest of the DNN weights via federated averaging.
• MTFL is compatible with standard FL optimization algorithms like FedAvg and extends to a distributed Adam variant, FedAvg-Adam, for faster convergence.
• Only the trainable parameters of BN layers (γ, β) are stored and communicated between rounds, significantly reducing storage and communication costs.
• The approach leverages the fact that BN layers adapt to local data distributions, improving local model performance without compromising privacy.
• Experiments use MNIST and CIFAR10 with non-IID data splits to evaluate UA, convergence speed, and communication efficiency.

Experimental results

Research questions

• RQ1 Can a federated learning framework improve user model accuracy (UA) in non-IID data settings by allowing personalization without full model sharing?
• RQ2 Does using private BN layers in a federated setting reduce the number of training rounds required to reach a target UA compared to standard FedAvg?
• RQ3 How does the combination of MTFL with FedAvg-Adam compare to other personalized FL methods in terms of convergence speed and final UA?
• RQ4 Does the communication overhead of MTFL remain justified by its performance gains in a real-world edge-computing testbed?
• RQ5 Can MTFL achieve better UA than state-of-the-art personalized FL methods while maintaining lower storage and computational costs?

Key findings

• MTFL with FedAvg reduces the number of rounds needed to reach a target UA by up to 5× compared to standard FedAvg on both MNIST and CIFAR10.
• Using FedAvg-Adam within MTFL further reduces the required rounds by up to 3× compared to FedAvg, demonstrating improved convergence speed.
• MTFL achieves the highest average user model accuracy (UA) on both MNIST and CIFAR10 across all tested scenarios, outperforming Per-FedAvg and pFedMe.
• The use of trainable BN parameters (γ, β) in model patches leads to faster convergence than using statistics (μ, σ), indicating their importance in personalization.
• In a real MEC testbed with 10 Raspberry Pi devices, MTFL(FedAvg-Adam)’s increased communication cost is outweighed by its significant gains in UA and convergence speed.
• MTFL maintains lower storage and computational overhead than competing personalized FL methods, as it only stores and communicates BN layer parameters, not full models.