# Probabilistic Federated Prompt-Tuning with Non-IID and Imbalanced Data

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

## Key facts

- Cited by 9 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 Feb 2025.

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss](https://mlsanity.com/paper/1906.07413/) (2019)
- [Learning to Prompt for Continual Learning](https://mlsanity.com/paper/2112.08654/) (2021)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [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)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (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)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Bayesian Nonparametric Federated Learning of Neural Networks](https://mlsanity.com/paper/1905.12022/) (2019)
- [Adaptive Personalized Federated Learning](https://mlsanity.com/paper/2003.13461/) (2020)
- [Think Locally, Act Globally: Federated Learning with Local and Global Representations](https://mlsanity.com/paper/2001.01523/) (2020)
- [Federated Learning of a Mixture of Global and Local Models](https://mlsanity.com/paper/2002.05516/) (2020)
- [Fine-tuning Global Model via Data-Free Knowledge Distillation for Non-IID Federated Learning](https://mlsanity.com/paper/2203.09249/) (2022)
- [PromptFL: Let Federated Participants Cooperatively Learn Prompts Instead of Models -- Federated Learning in Age of Foundation Model](https://mlsanity.com/paper/2208.11625/) (2022)
- [Communication Efficient Distributed Optimization using an Approximate Newton-type Method](https://mlsanity.com/paper/1312.7853/) (2013)
- [FedPrompt: Communication-Efficient and Privacy Preserving Prompt Tuning in Federated Learning](https://mlsanity.com/paper/2208.12268/) (2022)
- [Multimodal Federated Learning via Contrastive Representation Ensemble](https://mlsanity.com/paper/2302.08888/) (2023)
- [Where to Begin? On the Impact of Pre-Training and Initialization in Federated Learning](https://mlsanity.com/paper/2210.08090/) (2022)
- [Federated Learning on Heterogeneous and Long-Tailed Data via Classifier Re-Training with Federated Features](https://mlsanity.com/paper/2204.13399/) (2022)

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

- [Learning Reconfigurable Representations for Multimodal Federated Learning with Missing Data](https://mlsanity.com/paper/2510.22880/) (2025)
- [Federated Prompt-Tuning with Heterogeneous and Incomplete Multimodal Client Data](https://mlsanity.com/paper/2602.07081/) (2026)
- [FOCoOp: Enhancing Out-of-Distribution Robustness in Federated Prompt Learning for Vision-Language Models](https://mlsanity.com/paper/2506.16218/) (2025)
- [SABRE-FL: Selective and Accurate Backdoor Rejection for Federated Prompt Learning](https://mlsanity.com/paper/2506.22506/) (2025)
- [C^2Prompt: Class-aware Client Knowledge Interaction for Federated Continual Learning](https://mlsanity.com/paper/2509.19674/) (2025)
- [Leveraging Soft Prompts for Privacy Attacks in Federated Prompt Tuning](https://mlsanity.com/paper/2601.06641/) (2026)
- [Rethinking Cross-Modal Fine-Tuning: Optimizing the Interaction Between Feature Alignment and Target Fitting](https://mlsanity.com/paper/2601.18231/) (2026)
- [OvA-LP: A Simple and Efficient Framework for Federated Learning on Non-IID Data](https://mlsanity.com/paper/2511.05028/) (2025)
- [Nonparametric Distribution Matching for Self-Supervised Whole-Slide Image Condensation](https://mlsanity.com/paper/2610.00678/) (2026)

## Abstract (arXiv)

Fine-tuning pre-trained models is a popular approach in machine learning for solving complex tasks with moderate data. However, fine-tuning the entire pre-trained model is ineffective in federated data scenarios where local data distributions are diversely skewed. To address this, we explore integrating federated learning with a more effective prompt-tuning method, optimizing for a small set of input prefixes to reprogram the pre-trained model's behavior. Our approach transforms federated learning into a distributed set modeling task, aggregating diverse sets of prompts to globally fine-tune the pre-trained model. We benchmark various baselines based on direct adaptations of existing federated model aggregation techniques and introduce a new probabilistic prompt aggregation method that substantially outperforms these baselines. Our reported results on a variety of computer vision datasets confirm that the proposed method is most effective to combat extreme data heterogeneity in federated learning.