Privacy-Preserving Federated Embedding Learning for Localized Retrieval-Augmented Generation

TL;DR

FedE4RAG tackles privacy concerns in retrieval augmented generation by combining federated embedding learning with knowledge distillation and homomorphic encryption to train private RAG retrievers on client data. The approach fuses RAG-FT upstream fine tuning, KD-GLE client-server knowledge transfer, and FED-HE encrypted communication to maintain data confidentiality while enhancing global and local retrieval quality. Across finance domain datasets, FedE4RAG achieves superior upstream retrieval metrics and competitive downstream generation performance, significantly narrowing the gap to centralized training without compromising privacy. The work demonstrates practical privacy preserving gains for localized RAG systems and highlights pathways for scaling and extending to other domains with strict data governance.

Abstract

Retrieval-Augmented Generation (RAG) has recently emerged as a promising solution for enhancing the accuracy and credibility of Large Language Models (LLMs), particularly in Question & Answer tasks. This is achieved by incorporating proprietary and private data from integrated databases. However, private RAG systems face significant challenges due to the scarcity of private domain data and critical data privacy issues. These obstacles impede the deployment of private RAG systems, as developing privacy-preserving RAG systems requires a delicate balance between data security and data availability. To address these challenges, we regard federated learning (FL) as a highly promising technology for privacy-preserving RAG services. We propose a novel framework called Federated Retrieval-Augmented Generation (FedE4RAG). This framework facilitates collaborative training of client-side RAG retrieval models. The parameters of these models are aggregated and distributed on a central-server, ensuring data privacy without direct sharing of raw data. In FedE4RAG, knowledge distillation is employed for communication between the server and client models. This technique improves the generalization of local RAG retrievers during the federated learning process. Additionally, we apply homomorphic encryption within federated learning to safeguard model parameters and mitigate concerns related to data leakage. Extensive experiments conducted on the real-world dataset have validated the effectiveness of FedE4RAG. The results demonstrate that our proposed framework can markedly enhance the performance of private RAG systems while maintaining robust data privacy protection.

Figures (9)

• Figure 1: Schematic overview of our general Federated Embedding Learning for the private localized RAG systems. Different clients provide complementary data, while each client need not expose their data through federated learning.
• Figure 2: Schematic representation of the Federated Embedding Learning (FedE4RAG) framework, illustrating the two-phase process: upstream federated embedding pre-training and downstream retrieval-augmented generation inference. The upstream phase focuses on pre-training with federated embedding and knowledge distillation to address data scarcity and enhance local RAG retrievers. The downstream phase depicts private embedding models operating within trusted local environments to perform retrieval and answer generation without data exposure.
• Figure 3: The number of query-chunk pairs in upstream embedding pretraining, and the number of documents provided by each client company. In the training data (subfigures a & b), there are 5 companies (clients), namely AES (No.1), BOEING (No.2), ACTIVISION (No.3), PG (No.4), and PEPSICO (No.5). These companies provide 8, 8, 9, 13, and 13 documents for federated learning respectively, with a total of 51 documents.
• Figure 4: Performance comparison in validation dataset: Each enterprise provides training data to the client in a fixed proportion. Each enterprise provides training data, with the total data volume incrementally reaching 1K, 2K, 10K, 20K, and ultimately 50K. The model training process employs a batch size of 8 and 20 rounds. Each figure box is divided into two regions by its regression line. Each evaluation metric displays the Residual Sum of Squares (RSS), Pearson coefficient (Pearson R), and R-squared (R$^{2}$). These represent the total deviation of data points from the regression line, the strength of the linear relationship, and the proportion of variance explained by the model, respectively.
• Figure 5: Heatmap comparison of model performance with variations in number of training Rounds and Batch Sizes on Presence-Based Metrics.