p F ed RAG : A Personalized Federated Retrieval-Augmented Generation System with Depth-Adaptive Tiered Embedding Tuning (2025)

Abstract: Large Language Models (LLMs) can undergo hallucinations in specialized domains, and standard Retrieval-Augmented Generation (RAG) often falters due to general-purpose embeddings ill-suited for domain-specific terminology. Though domain-specific fine-tuning enhances retrieval, centralizing data introduces privacy risks. The use of federated learning (FL) can alleviate this to some extent, but faces challenges of data heterogeneity, poor personalization, and expensive training data generation. We propose pFedRAG, a novel Personalized Federated RAG framework, which enables efficient collaborative fine-tuning of embedding models to address these challenges. The key contribution is a new Depth-Adaptive Tiered Embedding (DATE) architecture, which comprises a Global Shared Layer, combined using FL to capture common knowledge, and a Personalized Layer with adjustable depth tailored for local data and training results of each client. The depth is locally controlled based on crafted metrics and scoring criteria. Also, pFedRAG incorporates a fully client-side pipeline leveraging local small LLMs and vector database filtering to construct high-quality query-document pairs. Experiments on diverse medical non-IID document datasets demonstrate that pFedRAG significantly reduces communication costs, handles data heterogeneity, and improves retrieval performance. Human evaluations confirm the enhanced response quality of pFedRAG.

## introduction large language models (llms), such as gpt series @xcite @xcite and llama @xcite , have achieved impressive performance across many natural language processing tasks @xcite . however, llms remain susceptible to hallucinations, producing plausible-sounding but factually incorrect content, which is problematic in domains such as healthcare and law @xcite . retrieval-augmented generation (rag) mitigates hallucinations by incorporating external knowledge. a retriever selects relevant documents from a knowledge base, and the generator conditions its responses on these documents. the effectiveness heavily depends on retrieval quality. as illustrated in figure 1 (a), conventional rag systems typically rely on general-purpose embedding models @xcite , which often underperform in specialized domains due to their inability to capture domain-specific semantics and terminology. to enhance retrieval accuracy in such domains, fine-tuning embedding models on domainspecific data has proven beneficial @xcite ; see figure 1(b ). this allows models to learn domain-relevant representations and improve the relevance of retrieved content. however, centralizing such data for fine-tuning raises significant privacy and governance concerns. federated learning (fl) offers a promising framework for domain-specific fine-tuning in collaborative environments. it enables multiple parties to jointly train models without sharing raw data, instead exchanging and aggregating model parameters, thus preserving data privacy. however, applying fl to the fine-tuning of embedding models in rag systems introduces several unique challenges: cost of full-parameter fine-tuning. modern embedding models contain hundreds of millions of parameters, making full-model tuning resourceintensive in fl settings. it imposes high communication and computation costs on clients. ## problem formulation consider an fl scenario involving n clients, where each client @xmath0, . . . , n } possesses a local dataset d i that exhibits significant data heterogeneity. each client i maintains a dual-tier embedding model, φ @xmath1, ϕ g , ϕ p i }, where θ is a pretrained embedding backbone shared across all clients and remains frozen during training, ϕ g is a global shared layer updated collaboratively via fedavg, and ϕ p i is a client-specific layer optimized locally to capture personalized information. ## proposed pfedrag framework the objective of pfedrag is to collaboratively train a retrieval embedding model that captures both globally shared knowledge and personalized features unique to each client's local data distribution. as illustrated in figure 2 , pfedrag consists of three key aspects: depth-adaptive tiered embedding (date), federated learning with global-local adaptation, and personalized rag system. ## experiments we evaluate the proposed pfedrag framework on a medical document dataset derived from multiple research domains to simulate realistic clinical and research-oriented retrieval scenarios. due to space limitations, detailed experimental settings are provided in appendix d. ## conclusion in this paper, we presented pfedrag to enhance rag systems in specialized domains while addressing privacy concerns and resource limitations. our approach tackles key challenges in federated settings, including the high cost of full-model tuning, semantic divergence across heterogeneous client data, and the need for high-quality supervised datasets. we introduced date, a comprehensive architecture comprising adt-head (a parameterefficient structure that combines a global shared layer for common knowledge aggregation with dynamically adjusted personalized layers) and depth controller for guiding adaptive complexity adjustments. we also proposed a client-side pipeline leveraging local llm and vector database filtering for privacy-preserving dataset construction. experimental evaluations demonstrated that pfedrag significantly reduces communication and computation costs, effectively handles data heterogeneity through adaptive model complexity, and improves end-to-end rag performance compared to standard baselines, showcasing its practical viability for collaborative, privacy-conscious enhancement of client-personalized rag systems.