FedHCA$^2$: Towards Hetero-Client Federated Multi-Task Learning

TL;DR

This paper introduces FedHCA$^2$, a novel federated multi-task learning framework that handles heterogeneous clients with diverse models and tasks, improving personalization and aggregation in real-world federated settings.

Contribution

The paper proposes FedHCA$^2$, a new framework with hyper conflict-averse and cross attention aggregation schemes for heterogeneous federated multi-task learning.

Findings

Outperforms existing methods in heterogeneous client scenarios

Effectively mitigates model incongruity during aggregation

Enhances personalized model training in federated multi-task learning

Abstract

Federated Learning (FL) enables joint training across distributed clients using their local data privately. Federated Multi-Task Learning (FMTL) builds on FL to handle multiple tasks, assuming model congruity that identical model architecture is deployed in each client. To relax this assumption and thus extend real-world applicability, we introduce a novel problem setting, Hetero-Client Federated Multi-Task Learning (HC-FMTL), to accommodate diverse task setups. The main challenge of HC-FMTL is the model incongruity issue that invalidates conventional aggregation methods. It also escalates the difficulties in accurate model aggregation to deal with data and task heterogeneity inherent in FMTL. To address these challenges, we propose the FedHCA$^2$ framework, which allows for federated training of personalized models by modeling relationships among heterogeneous clients. Drawing on our theoretical insights into the difference between multi-task and federated optimization, we propose the Hyper Conflict-Averse Aggregation scheme to mitigate conflicts during encoder updates. Additionally, inspired by task interaction in MTL, the Hyper Cross Attention Aggregation scheme uses layer-wise cross attention to enhance decoder interactions while alleviating model incongruity. Moreover, we employ learnable Hyper Aggregation Weights for each client to customize personalized parameter updates. Extensive experiments demonstrate the superior performance of FedHCA$^2$ in various HC-FMTL scenarios compared to representative methods. Our code will be made publicly available.