# Toward Enhancing Representation Learning in Federated Multi-Task Settings

arXiv: https://arxiv.org/abs/2602.01626 · HTML: https://mlsanity.com/paper/2602.01626/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 64 of its 90 references that are in the map.
- First posted on arXiv 2 Feb 2026.

## Builds on (most cited first, 30 of 64)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Flickr30k Entities: Collecting Region-to-Phrase Correspondences for Richer Image-to-Sentence Models](https://mlsanity.com/paper/1505.04870/) (2015)
- [Contrastive Multiview Coding](https://mlsanity.com/paper/1906.05849/) (2019)
- [ImageBind: One Embedding Space To Bind Them All](https://mlsanity.com/paper/2305.05665/) (2023)
- [Self-Supervised Learning of Pretext-Invariant Representations](https://mlsanity.com/paper/1912.01991/) (2019)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [Federated Multi-Task Learning](https://mlsanity.com/paper/1705.10467/) (2017)
- [Ditto: Fair and Robust Federated Learning Through Personalization](https://mlsanity.com/paper/2012.04221/) (2020)
- [Personalized Federated Learning with Moreau Envelopes](https://mlsanity.com/paper/2006.08848/) (2020)
- [Federated Optimization: Distributed Machine Learning for On-Device Intelligence](https://mlsanity.com/paper/1610.02527/) (2016)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [FedMD: Heterogenous Federated Learning via Model Distillation](https://mlsanity.com/paper/1910.03581/) (2019)
- [Understanding the Behaviour of Contrastive Loss](https://mlsanity.com/paper/2012.09740/) (2020)

## Built on by (most cited first, 2 of 2)

- [SEMASIA: A Large-Scale Dataset of Semantically Structured Latent Representations](https://mlsanity.com/paper/2605.09485/) (2026)
- [Sheaf-Based Federated Representation Learning](https://mlsanity.com/paper/2608.10016/) (2026)

## Abstract (arXiv)

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.