Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Toward Enhancing Representation Learning in Federated Multi-Task Settings

Mehdi Setayesh, Mahdi Beitollahi, Yasser H. Khalil, Hongliang Li

TL;DR

Federated multi-task learning is challenged by heterogeneity in client models and tasks. The authors introduce Muscle loss, a theoretically grounded contrastive objective that jointly aligns representations from all participating models and yields a mutual information lower bound, and FedMuscle, a practical FL algorithm that uses a shared public dataset to synchronize representation spaces. They demonstrate MI-based transfer gains across CV and NLP tasks, including multi-modal settings, and show Muscle loss outperforms Gramian-based and pairwise contrastive losses while integrating with existing multi-modal FL frameworks. The work provides insights into trade-offs between communication, batch size, and client count, highlighting the potential of representation-centric FL for robust, privacy-preserving knowledge transfer across heterogeneous participants.

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.

Toward Enhancing Representation Learning in Federated Multi-Task Settings

TL;DR

Federated multi-task learning is challenged by heterogeneity in client models and tasks. The authors introduce Muscle loss, a theoretically grounded contrastive objective that jointly aligns representations from all participating models and yields a mutual information lower bound, and FedMuscle, a practical FL algorithm that uses a shared public dataset to synchronize representation spaces. They demonstrate MI-based transfer gains across CV and NLP tasks, including multi-modal settings, and show Muscle loss outperforms Gramian-based and pairwise contrastive losses while integrating with existing multi-modal FL frameworks. The work provides insights into trade-offs between communication, batch size, and client count, highlighting the potential of representation-centric FL for robust, privacy-preserving knowledge transfer across heterogeneous participants.

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.
Paper Structure (54 sections, 1 theorem, 29 equations, 11 figures, 9 tables, 1 algorithm)

This paper contains 54 sections, 1 theorem, 29 equations, 11 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Federated Multi-Task Learning
  4. Model Heterogeneity in FL
  5. Representation Learning in FL
  6. Problem Formulation
  7. Notations
  8. Problem Setting
  9. The Contrastive Learning Framework
  10. Preliminaries
  11. Our Proposed Muscle Loss Function
  12. FedMuscle Methodology
  13. Experiments
  14. Experimental Setup
  15. Datasets and Tasks
  16. ...and 39 more sections

Key Result

Theorem 1

Given $\mathcal{L}_{\textrm{Muscle}}^{n}(\bm{z}^{n}_i)$ in equation (eq:new_Loss), the mutual information $I(\bm{z}^{n}_i; \{\bm{z}^{m}_i\}_{m=1,\, m\neq n}^{N})$ is lower-bounded as follows: where $\mathbb{E}$ denotes the expectation over random batch of data samples.

Figures (11)

  • Figure 1: FedMuscle overview. Users may have heterogeneous models, tasks, and local datasets. A shared public dataset is used by all users to align the representation spaces of their models via the Muscle loss function.
  • Figure 2: Performance of FedMuscle in Setup1 using the proposed Muscle loss function, compared to the Gramian-based contrastive loss and pairwise alignment.
  • Figure 3: Impact of the number of local epochs $E$, communication rounds $R$, and CL epochs $T$ on FedMuscle performance. COCO is used as the public dataset.
  • Figure 4: Impact of the public dataset size $\lvert\mathcal{D} \rvert$ and number of communication rounds $R$ on the communication cost of FedMuscle. COCO is used as the public dataset.
  • Figure 5: Impact of batch size on FedMuscle’s performance in Setup1. COCO is used as the public dataset.
  • ...and 6 more figures

Theorems & Definitions (3)

  • Remark 1
  • Theorem 1
  • Remark 2