# Tackling Feature and Sample Heterogeneity in Decentralized Multi-Task Learning: A Sheaf-Theoretic Approach

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 57 references that are in the map.
- First posted on arXiv 3 Feb 2025.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [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)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Federated Learning on Non-IID Data Silos: An Experimental Study](https://mlsanity.com/paper/2102.02079/) (2021)
- [Federated Learning with Personalization Layers](https://mlsanity.com/paper/1912.00818/) (2019)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints](https://mlsanity.com/paper/1910.01991/) (2019)
- [Can Decentralized Algorithms Outperform Centralized Algorithms? A Case Study for Decentralized Parallel Stochastic Gradient Descent](https://mlsanity.com/paper/1705.09056/) (2017)
- [Adaptive Personalized Federated Learning](https://mlsanity.com/paper/2003.13461/) (2020)
- [FedProto: Federated Prototype Learning across Heterogeneous Clients](https://mlsanity.com/paper/2105.00243/) (2021)
- [Think Locally, Act Globally: Federated Learning with Local and Global Representations](https://mlsanity.com/paper/2001.01523/) (2020)
- [The Non-IID Data Quagmire of Decentralized Machine Learning](https://mlsanity.com/paper/1910.00189/) (2019)
- [Personalized Cross-Silo Federated Learning on Non-IID Data](https://mlsanity.com/paper/2007.03797/) (2020)
- [Improving Federated Learning Personalization via Model Agnostic Meta Learning](https://mlsanity.com/paper/1909.12488/) (2019)
- [No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data](https://mlsanity.com/paper/2106.05001/) (2021)
- [Federated Learning of a Mixture of Global and Local Models](https://mlsanity.com/paper/2002.05516/) (2020)
- [Don't Use Large Mini-Batches, Use Local SGD](https://mlsanity.com/paper/1808.07217/) (2018)
- [Federated Meta-Learning with Fast Convergence and Efficient Communication](https://mlsanity.com/paper/1802.07876/) (2018)
- [Federated Evaluation of On-device Personalization](https://mlsanity.com/paper/1910.10252/) (2019)
- [Salvaging Federated Learning by Local Adaptation](https://mlsanity.com/paper/2002.04758/) (2020)

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

- [Semantic Communication meets System 2 ML: How Abstraction, Compositionality and Emergent Languages Shape Intelligence](https://mlsanity.com/paper/2505.20964/) (2025)
- [SheafAlign: A Sheaf-theoretic Framework for Decentralized Multimodal Alignment](https://mlsanity.com/paper/2510.20540/) (2025)
- [Toward Enhancing Representation Learning in Federated Multi-Task Settings](https://mlsanity.com/paper/2602.01626/) (2026)
- [Communication-Efficient and Robust Multi-Modal Federated Learning via Latent-Space Consensus](https://mlsanity.com/paper/2603.19067/) (2026)
- [Resilient-native and Intelligent NextG Systems](https://mlsanity.com/paper/2506.12795/) (2025)
- [Sheaf-Based Decentralized Multimodal Learning for Next-Generation Wireless Communication Systems](https://mlsanity.com/paper/2506.22374/) (2025)
- [Learning to Collaborate Over Graphs: A Selective Federated Multi-Task Learning Approach](https://mlsanity.com/paper/2506.10102/) (2025)
- [Prototype-Regularized Federated Learning for Cross-Domain Aspect Sentiment Triplet Extraction](https://mlsanity.com/paper/2604.09123/) (2026)
- [MuCALD-SplitFed: Causal-Latent Diffusion for Privacy-Preserving Multi-Task Split-Federated Medical Image Segmentation](https://mlsanity.com/paper/2605.04108/) (2026)
- [Sheaf-Based Federated Representation Learning](https://mlsanity.com/paper/2608.10016/) (2026)

## Abstract (arXiv)

Federated multi-task learning (FMTL) aims to simultaneously learn multiple related tasks across clients without sharing sensitive raw data. However, in the decentralized setting, existing FMTL frameworks are limited in their ability to capture complex task relationships and handle feature and sample heterogeneity across clients. To address these challenges, we introduce a novel sheaf-theoretic-based approach for FMTL. By representing client relationships using cellular sheaves, our framework can flexibly model interactions between heterogeneous client models. We formulate the sheaf-based FMTL optimization problem using sheaf Laplacian regularization and propose the Sheaf-FMTL algorithm to solve it. We show that the proposed framework provides a unified view encompassing many existing federated learning (FL) and FMTL approaches. Furthermore, we prove that our proposed algorithm, Sheaf-FMTL, achieves a sublinear convergence rate in line with state-of-the-art decentralized FMTL algorithms. Extensive experiments show that although Sheaf-FMTL introduces computational and storage overhead due to the management of interaction maps, it achieves substantial communication savings in terms of transmitted bits when compared to decentralized FMTL baselines. This trade-off makes Sheaf-FMTL especially suitable for cross-silo FL scenarios, where managing model heterogeneity and ensuring communication efficiency are essential, and where clients have adequate computational resources.