Toward Enhancing Representation Learning in Federated Multi-Task Settings

ICLR 2026 Conference Submission13394 Authors

Keywords: Contrastive learning, federated learning, knowledge transfer, multi-task learning, mutual information, representation learning

TL;DR: This paper proposes a new contrastive learning loss and, leveraging it, designs a novel federated multi-task learning algorithm that addresses both task and model heterogeneity among users.

Abstract: Federated multi-task learning (FMTL) seeks to collaboratively train customized models for users with different tasks while preserving data privacy. Most existing approaches assume model congruity (i.e., the use of fully or partially homogeneous models) across users, which limits their applicability in realistic settings. To overcome this limitation, we aim to learn a shared representation space across tasks rather than shared model parameters. To this end, we propose *Muscle loss*, a novel contrastive learning objective that simultaneously aligns representations from all participating models. Unlike existing multi-view or multi-model contrastive methods, which typically align models pairwise, Muscle loss can effectively capture dependencies across tasks because its minimization is equivalent to the maximization of mutual information among all the models' representations. Building on this principle, we develop *FedMuscle*, a practical and communication-efficient FMTL algorithm that naturally handles both model and task heterogeneity. Experiments on diverse image and language tasks demonstrate that FedMuscle consistently outperforms state-of-the-art baselines, delivering substantial improvements and robust performance across heterogeneous settings.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Submission Number: 13394