Multi-Task Federated Learning for Personalised Deep Neural Networks in Edge Computing

• Published in: IEEE transactions on parallel and distributed systems Vol. 33; no. 3; pp. 630 - 641
• Main Authors: Mills, Jed, Hu, Jia, Min, Geyong
• Format: Journal Article
• Language: English
• Published: New York IEEE 01.03.2022; The Institute of Electrical and Electronics Engineers, Inc. (IEEE)
• Subjects: Adaptation models; adaptive optimization; Algorithms; Artificial neural networks; Computational modeling; Convergence; Customization; Data models; deep learning; Edge computing; Electronic devices; Federated learning; Iterative methods; Machine learning; Model accuracy; multi-task learning; Neural networks; Optimization; Servers; Training
• ISSN: 1045-9219; 1558-2183
• DOI: 10.1109/TPDS.2021.3098467

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 <inline-formula><tex-math notation="LaTeX">5\times</tex-math> <mml:math><mml:mrow><mml:mn>5</mml:mn><mml:mo>×</mml:mo></mml:mrow></mml:math><inline-graphic xlink:href="hu-ieq1-3098467.gif"/> </inline-formula> when using existing FL optimisation strategies, and with a further <inline-formula><tex-math notation="LaTeX">3\times</tex-math> <mml:math><mml:mrow><mml:mn>3</mml:mn><mml:mo>×</mml:mo></mml:mrow></mml:math><inline-graphic xlink:href="hu-ieq2-3098467.gif"/> </inline-formula> 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.