Source-linked AI summary
Multi-Task Federated Learning for Personalised Deep Neural Networks in Edge Computing
Jed Mills, Jia Hu, Geyong Min
TL;DR
Federated learning must handle non-IID client data while improving individual user accuracy without exposing private data. The paper proposes MTFL, which keeps BN layers private for personalisation, and reports faster convergence and higher UA across MNIST and CIFAR10, with testbed benefits outweighing overhead.
Problem
Non-IID client data harms FL training, while global-model accuracy may not reflect the individual user accuracy that many applications seek.
Method
MTFL adds private BN layers to federated DNNs, personalising client models while retaining compatibility with iterative FL optimisation.
Results
MTFL achieves the best UA across considered MNIST and CIFAR10 scenarios, reaches target UA in up to 5× fewer rounds, and gains a further 3× improvement with FedAvg-Adam.
Takeaways & Limitations
Private BN patches provide a personalised FL approach with reported privacy and storage benefits, while FedAvg-Adam’s testbed overhead is outweighed by UA and convergence gains.
Takeaways & Limitations
The analysis does not consider peer-to-peer FL, and more sophisticated peer-to-peer extensions may require complex private-layer integration.
Abstract
from arXiv · showhide
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\times$ when using existing FL optimisation strategies, and with a further $3\times$ 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.
1 INTRODUCTION
The paper motivates personalised federated learning for non-IID edge data, where local user accuracy matters alongside privacy. It proposes MTFL with private BN layers and evaluates its accuracy, convergence, efficiency, and deployment benefits.
- Non-IID client data can substantially hinder federated-model training.
- Individual client accuracy is often the real FL objective, but personalisation remains under-researched.
- Private BN layers reduce uploaded information and storage because BN parameters are a tiny fraction of total DNN parameters.
- MTFL personalises DNNs by keeping BN layers private to each client.
- MTFL introduces User model Accuracy (UA) to measure client-side test accuracy rather than only global-model accuracy.
- On MNIST and CIFAR10, MTFL with FedAvg reaches target UA in up to 5× fewer rounds, with FedAvg-Adam providing a further 3× improvement.
- An MEC-like testbed indicates that FedAvg-Adam overheads are outweighed by UA and convergence-speed benefits.
2 RELATED WORK
Related work covers personalised FL, edge-system design, and convergence optimisation. The paper positions MTFL as adding multi-task learning within FL and combines it with client-side adaptive optimisation.
- Personalisation: Prior personalised-FL methods use meta-learning or combinations of local and global models to address non-IID user datasets.
- Personalisation: MOCHA and FedU formulate federated multi-task learning through relationships among model weights, but MOCHA is not directly comparable to MTFL because it addresses heterogeneous hardware.
- MTFL positioning: The paper claims MTFL provides convergence-speed, personalisation, privacy, and storage-cost benefits over existing personalised FL algorithms.
- Edge computing: Existing edge-FL research addresses client selection, communication, wireless conditions, computation, and deployment systems.
- MTFL positioning: Previous implementation-focused FL work did not consider multi-task learning within FL, which the paper identifies as MTFL’s main contribution.
- Optimisation: FedAvg-Adam applies adaptive optimisation on clients and is reported to converge much faster than FedAvg or server-only Adam optimisation.
3 MULTI-TASK FEDERATED LEARNING (MTFL)
MTFL personalises federated DNNs by keeping BN patch layers private to each client while aggregating shared model parameters. It targets client-level UA, supports iterative FL optimisers, and addresses privacy, storage, convergence, and local-performance concerns.
- User Model Accuracy: UA measures accuracy on each client’s local test set, reflecting personalised FL objectives better than centralised global-model accuracy.The local test set should be drawn from a distribution similar to the client’s training data.
- MTFL framework: Each client model combines shared Conv and FC parameters with private BN patch layers.The example composition uses shared parameters Ω1−Ω4 and private patches Pk1, Pk2, Pk3.
- User Model Accuracy: MTFL addresses non-IID clients by allowing different client models while retaining the benefits of federated learning.Non-IID data can make aggregate models perform poorly on some clients, particularly clients with fewer local samples.
- MTFL framework: MTFL keeps BN layers private to each client while federating the remaining DNN parameters.Clients apply private patches during local training, then upload only non-private layers for server aggregation.
- Personalisation trade-offs: Private BN patches can improve privacy and reduce personalisation storage because only a small parameter subset remains client-specific.BN parameters are reported as less than 1% of total model size for the tested architectures, though private-parameter count trades off against global convergence.
- Federated optimisation: MTFL can use FedAvg-Adam as its distributed optimisation strategy, and the paper reports that MTFL reduces rounds to target UA across optimisation strategies.FedAvg-Adam shares global Adam optimiser values and is described as providing substantial communication-round speedup compared with FedAvg.
4 EXPERIMENTS
Experiments on non-IID MNIST and CIFAR10 clients evaluate MTFL across patch choices, noise, optimisation strategies, personalised-FL baselines, and an edge testbed. MTFL generally improves convergence and User Accuracy, while FedAvg-Adam further reduces communication rounds despite higher per-round overhead.
- Experimental Setup: Experiments use MNIST and CIFAR10 with non-IID client partitions, varying client counts, participation rates, and optimisation strategies.Each client receives label-based shards, preserving matching training and testing class distributions.
- Patch Layers in FL: 36 rounds versus 164 rounds: MTFL(FedAvg) with all private BN values reached the CIFAR10 target faster than FL(FedAvg) for W = 400, C = 1.0.Using only private (µ, σ) took 266 rounds, while private (γ, β) took 30 rounds.
- Robustness to Noisy Clients: MTFL mitigated noisy-client effects: in the same CIFAR10 scenario, private (γ, β) reached the non-noisy-client target in 28 rounds versus 250 rounds for FL(FedAvg).The authors attribute this improvement partly to reduced coupling between non-noisy and noisy clients.
- Optimisation Strategies: 8 rounds versus 31 rounds: MTFL(FedAvg-Adam) with all private BN values reached the CIFAR10 target faster than FL(FedAvg-Adam) in the same scenario.FedAvg-Adam generally required the fewest rounds across FL and MTFL configurations.
- Training and Testing Results Using MTFL: Private (µ, σ) leaves training curves unchanged because BN statistics are used only at test time, but can lower test accuracy through mismatch with averaged γ and β.Keeping both statistics and trainable parameters private avoids the substantial drop observed with private statistics alone.
- Personalised-FL Comparison and Testbed: MTFL achieved higher UA than FL(FedAvg), pFedMe, and Per-FedAvg in all tested scenarios while using fewer personalisation resources and lower computational cost.On the testbed, MTFL(FedAvg-Adam) took the longest per round, but its extra communication was likely outweighed by fewer required rounds.
5 CONCLUSION
MTFL personalises federated models with private patch layers, improving average User model Accuracy and convergence. Experiments show further gains from FedAvg-Adam and strong performance under limited communication rounds.
- MTFL introduces private patch layers into iterative FL models, enabling personalised user models and improving average User model Accuracy.The paper analyses Batch-Normalization layers as MTFL patches.
- Up to 5× fewer rounds are required to reach a target average UA with MTFL using FedAvg than with FL.
- Up to 3× further reduction in rounds is achieved when MTFL uses FedAvg-Adam instead of FedAvg.FedAvg-Adam applies Adam on clients.
- Private BN trainable parameters (γ, β) provide better convergence speed than private BN statistics (µ, σ) in model patches.
- MTFL achieves the highest average UA among compared personalised FL algorithms given limited communication rounds.