# Towards Diverse Device Heterogeneous Federated Learning via Task Arithmetic Knowledge Integration

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 60 references that are in the map.
- First posted on arXiv 27 Sep 2024.

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

- [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)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [A large annotated corpus for learning natural language inference](https://mlsanity.com/paper/1508.05326/) (2015)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [A Comprehensive Survey on Transfer Learning](https://mlsanity.com/paper/1911.02685/) (2019)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [TIES-Merging: Resolving Interference When Merging Models](https://mlsanity.com/paper/2306.01708/) (2023)
- [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 Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [ResNet strikes back: An improved training procedure in timm](https://mlsanity.com/paper/2110.00476/) (2021)
- [HeteroFL: Computation and Communication Efficient Federated Learning for Heterogeneous Clients](https://mlsanity.com/paper/2010.01264/) (2020)
- [Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models](https://mlsanity.com/paper/2305.12827/) (2023)
- [FjORD: Fair and Accurate Federated Learning under heterogeneous targets with Ordered Dropout](https://mlsanity.com/paper/2102.13451/) (2021)
- [Expanding the Reach of Federated Learning by Reducing Client Resource Requirements](https://mlsanity.com/paper/1812.07210/) (2018)
- [On the Convergence of Local Descent Methods in Federated Learning](https://mlsanity.com/paper/1910.14425/) (2019)
- [The Multilingual Amazon Reviews Corpus](https://mlsanity.com/paper/2010.02573/) (2020)
- [What Happens after SGD Reaches Zero Loss? --A Mathematical Framework](https://mlsanity.com/paper/2110.06914/) (2021)
- [Uncertainty Quantification with Pre-trained Language Models: A Large-Scale Empirical Analysis](https://mlsanity.com/paper/2210.04714/) (2022)
- [FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-Model Extraction](https://mlsanity.com/paper/2212.01548/) (2022)
- [Heterogeneous Ensemble Knowledge Transfer for Training Large Models in Federated Learning](https://mlsanity.com/paper/2204.12703/) (2022)
- [FedAUX: Leveraging Unlabeled Auxiliary Data in Federated Learning](https://mlsanity.com/paper/2102.02514/) (2021)

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

- [C^2Prompt: Class-aware Client Knowledge Interaction for Federated Continual Learning](https://mlsanity.com/paper/2509.19674/) (2025)
- [First-Order Softmax Weighted Switching Gradient Method for Distributed Stochastic Minimax Optimization with Stochastic Constraints](https://mlsanity.com/paper/2603.05774/) (2026)
- [Progressive Size-Adaptive Federated Learning: A Comprehensive Framework for Heterogeneous Multi-Modal Data Systems](https://mlsanity.com/paper/2506.20685/) (2025)
- [Bi-level Personalization for Federated Foundation Models: A Task-vector Aggregation Approach](https://mlsanity.com/paper/2509.12697/) (2025)
- [Federated Foundation Models Fine-Tuning with Heterogeneous Compressed Clients](https://mlsanity.com/paper/2607.29071/) (2026)

## Abstract (arXiv)

Federated Learning has emerged as a promising paradigm for collaborative machine learning, while preserving user data privacy. Despite its potential, standard FL lacks support for diverse heterogeneous device prototypes, which vary significantly in model and dataset sizes -- from small IoT devices to large workstations. This limitation is only partially addressed by existing knowledge distillation techniques, which often fail to transfer knowledge effectively across a broad spectrum of device prototypes with varied capabilities. This failure primarily stems from two issues: the dilution of informative logits from more capable devices by those from less capable ones, and the use of a single integrated logits as the distillation target across all devices, which neglects their individual learning capacities and and the unique contributions of each. To address these challenges, we introduce TAKFL, a novel KD-based framework that treats the knowledge transfer from each device prototype's ensemble as a separate task, independently distilling each to preserve its unique contributions and avoid dilution. TAKFL also incorporates a KD-based self-regularization technique to mitigate the issues related to the noisy and unsupervised ensemble distillation process. To integrate the separately distilled knowledge, we introduce an adaptive task arithmetic knowledge integration process, allowing each student model to customize the knowledge integration for optimal performance. Additionally, we present theoretical results demonstrating the effectiveness of task arithmetic in transferring knowledge across heterogeneous devices with varying capacities. Comprehensive evaluations of our method across both CV and NLP tasks demonstrate that TAKFL achieves SOTA results in a variety of datasets and settings, significantly outperforming existing KD-based methods Code is released at https://github.com/MMorafah/TAKFL