Based on work done on Differentially Private In Context Learning (DP-ICL), we developed an tested an algorithm that enables DP on Retrieval Augmented Generation (RAG).

Abstract

Based on work done on Differentially Private In Context Learning (DP-ICL), we developed an tested an algorithm that enables DP on Retrieval Augmented Generation (RAG). By aggregating multiple generations made from different retrieved documents, DP-RAG on a custom dataset is able to approach RAG’s performance while offering privacy guarantees.

Motivations and Goals

Limitations of LLMs and Motivation for RAG

Large Language Models (LLMs) have finite knowledge, and fine-tuning them to learn new information is costly and time-consuming. An alternative is In-Context Learning (ICL), where new knowledge is provided during queries, but it’s limited by context size. Retrieval Augmented Generation (RAG) optimizes this by storing knowledge as text embeddings and retrieving only the most relevant information to include in the context.

Privacy risks with RAG and differential privacy

A challenge with RAG is that it may retrieve personal data, and there’s no way to control the output once information is fed into the LLM. One solution is to use differentially private fine-tuning (e.g., DP-SGD), which is effective but costly. Alternatively, RAG can be used multiple times, generating output based on differentially private aggregation of token probability vectors, which is the approach taken in this experiment.

Retrieval Augmented Generation

Retrieval Augmented Generation (RAG) enhances LLM output by providing relevant knowledge related to the query. It involves two steps:

Retrieval: Relevant documents are retrieved from an external dataset, based on the similarity between vector-embedded documents and the vector-embedded query.
Generation: The LLM uses the retrieved documents alongside the query to generate a more informed response.

Differential Privacy

Differential Privacy (DP) provides a mathematical guarantee of individual data privacy within a dataset. A DP algorithm ensures that outputs reveal little about any individual’s data by bounding the difference in output distributions between two datasets differing by only one individual. This is typically achieved by adding noise (e.g., Gaussian or Exponential) tied to a privacy parameter 𝛜, with smaller 𝛜 offering greater privacy. In a RAG algorithm, both retrieval and generation subprocesses must be made differentially private to ensure data privacy.

DP in RAG Retrieval

Top-k Retrieval

The standard RAG process involves retrieving the k most relevant documents based on their vector distance to a query. This is an aggregation that needs to be made DP, since each document’s inclusion or exclusion could completely alter the result of top-k retrieval. While it’s possible to make the retrieval DP with a private mechanism, we chose a different method.

Clipped Similarity Scaling

The chosen method involves including all documents in the generation process, avoiding the differential privacy (DP) concerns associated with retrieval aggregation. To retain the benefits of using only relevant documents, each document’s token probability distribution is scaled by a clipped similarity score to the query. If the similarity score is above a certain threshold, the distribution is scaled accordingly; if below, the distribution is scaled by zero, effectively nullifying the document’s influence. This method maintains the influence of relevant documents without the need for aggregation and DP.

Each document undergoes an independent transformation, from embedding vector to log-probability distribution, which is then scaled by its similarity to the query. There is a direct link between a document and its scaled probability distribution, ensuring that the inclusion or exclusion of a document directly impacts the probability distributions. The similarity threshold is fixed to prevent manipulation by users or attackers. To reduce costs, documents below the threshold can be effectively ignored by ponderating their relevance before the generation step.

DP in RAG Generation

For DP generation with private documents, research, on DP-In-Context-Learning (ICL) was used as a foundation. The algorithms presented in the papers, broadly speaking, query an LLM multiple times with different documents, and privately aggregate the results using the Gaussian Mechanism. Since we send each document into its own generation, there is no aggregation of data until we merge the token probability distributions.

The Differential Privacy of the Gaussian Mechanism is demonstrated by Dwork and Roth in Appendix A of The Algorithmic Foundations of Differential Privacy.

The main difference between our implementation and those of the papers is that ours needs to scale the documents on their similarity. As was shown previously, this process does not impact the DP of the algorithm.

Standard RAG Results

As a baseline to compare DP-RAG, a standard RAG algorithm was used, without aggregation and using top-k retrieval. This standard RAG method had an accuracy of 96% (percent of responses containing correct disease name) on our testing set, which could have been higher if we had not limited the max generated tokens to 5 (for testing purposes).

Attacks on RAG

While the results of RAG are impressive, an attacker can input a query that retrieves specific documents and outputs private information. In the following example, information on David is known by the attacker, and RAG retrieves and generates his private information

DP-RAG Results

In general, an epsilon around and below 5 is considered reasonable; for our dataset we found that an epsilon of 5 produced comparable results to an epsilon approaching infinity. While worse than the 96% accuracy of standard RAG, the results are nonetheless impressive, and could be improved with prompt engineering and a larger dataset. As for the discrepancy between the accuracies with epsilon=inf and Standard RAG, this is explained in the ‘Aggregation’ portion of the next section.

Some examples from testing:

Limitations

Three factors of the DP algorithm could deteriorate RAG from its standard approach. These are the Gaussian Mechanism (according to privacy parameters 𝛜 and δ), the aggregation process (opposed to a single generation with all docs), and the similarity threshold ‘retrieval’ process.

Gaussian Mechanism

The only one of these factors that we have any control over in DP-RAG is the Gaussian Mechanism, where we can tweak 𝛜 to find a good balance between privacy and accuracy. To do this, we compared accuracy across different values of 𝛜, shown above.

Aggregation

As opposed to Standard RAG, where the top-k documents are fed into a single generation, our algorithm generates a token for each retrieved document, and then aggregates the token-probabilities. Even without privately aggregating via the Gaussian Mechanism, this process will impact the final result. Whereas with Standard RAG the ‘correct’ document(s) might dominate the generation, under the aggregation process each generation could have more impact over the result. This might hurt accuracy, but will help privacy, as each individual document will have less of an impact on the final result.

In general, as long as the retrieved documents are substantially ‘correct’ the results should be both accurate and private. For our experiment, the accuracy of Standard RAG was higher than RAG with aggregation (96% vs 84%). However this gap could be filled with a larger dataset, as having more correct examples in the retrieved documents would lead to higher likelihood of a correct generation.

Similarity Thresholding

Briefly, similarity thresholding is inconsistent; with the constraint of a fixed threshold, some examples will retrieve too many undesirable documents, which will lead to inaccurate generations having a larger effect on the final result. For example, while some documents have disease symptoms that are unique, others might share symptoms with more documents, and these relevant documents could dominate in the selected ones used to generate a response.

DP-RAG Attacks

If we attack on David again, just the aggregation of many independent generations makes the attack unsuccessful, even with a high epsilon:

In the following example the only document retrieved is the (relevant) one of ‘Natasha Romanoff’:

The attack is successful, but only because we have kept a high epsilon — with an epsilon of 5, the attack is a failure:

The incoherency is due to the retrieval of just one document, so the aggregated token probabilities are highly influenced by the Gaussian Mechanism. Were there more retrieved documents, the more sensible tokens would aggregate to higher probabilities, which would be less impacted by the Gaussian Mechanism.

In keeping a low epsilon, the number of retrieved documents needs to be sufficiently high to generate a coherent response, which is exactly what we want in order to defend against targeted attacks. It is practically impossible to generate a response using just a single document, and as the number of retrieved more documents increases, the impact of any one decreases.

More unsuccessful attacks:

Threshold too high — no documents retrieved

Attacks are inherently difficult with DP-RAG due to the nature of the retrieval process. If too few documents are retrieved, the Gaussian Noise will dominate and responses will be gibberish. On the other hand if enough documents are retrieved, the impact of one document on the response will be small.

For our similarity score threshold used in the experiment, the only tested attack that might retrieve enough documents is one which lists out symptoms and asks for private information, depicted above.

Dataset Difficulties

There are a couple characteristics of the dataset that make DP-RAG less successful than it could be. Among these are the size of the dataset, the inconsistency of disease counts, and the symptoms of the different diseases. DP is more effective with larger datasets (less impact of any individual on the result), and RAG too fetches more relevant documents if there are more elements in the dataset.

Dataset Size

The dataset only contains 9k examples, which is not big enough for either privacy or accuracy concerns. If the dataset was larger, a higher percentage of “correct” (pertaining to the same disease as the query) would be retrieved, which would help privacy and accuracy. This is a common idea in DP, where results improve with the size of the dataset.

Disease Counts

Some diseases are under-represented, and there is a clear correlation between the count of a disease in the dataset and the accuracy of those elements in the testset.

Symptom Inconsistencies

While most diseases in the dataset have a majority of symptoms which are entirely made up (electric fingertips), some have more ‘real-life’ symptoms (sore throat). Again, there is a correlation between the number of real symptoms and accuracy.

The reason for this is when the model is given a medical question with all real symptoms, the model seems to be unwilling to depart with its parametrized knowledge, and is inclined not to use the retrieved docs to answer the question. Between the two graphs, we can see that there would be a stronger correlation between accuracy in disease counts if the outliers, which mainly contain entirely real symptoms, were removed.

Summary and Conclusions

While there is room for improvement — namely in reducing inference cost, optimizing the aggregation method, and calculating similarity threshold privately — DP-RAG showed promising results.

As DP-RAG is principally a replacement method of DP-Finetuning, it is important to compare the two. Both attempt to assist the generation of an LLM, in order to adapt to new data, while also upholding the mathematical guarantees of DP.

DP Finetuning

Cons

Costly in time and money to train the model on a new dataset
Can be technically difficult to implement depending on the model

Pros

One time privacy cost — inference does not incur any further privacy cost

DP-RAG