Discovering New Knowledge while Protecting Privacy

You can fine-tune an LLM to learn new knowledge from private data, ensuring that no sensitive records are at risk of being regurgitated.

AI
Differential Privacy
LLM
Fine Tuning
Mistral
Luca Canale

In a previous post, we showed that a LLM can be fine tuned to learn new knowledge from unstructured natural language. We trained a Q-LoRA Mistral 7B-Instruct on a dataset containing a collection of patient/doctor messages about fictional diseases, symptoms, and medications. We showed that:

  • the model that did not have any prior information about these diseases was able to learn the connection between symptoms and diseases and to generalize, giving the correct diagnosis and prescribing the right treatment on a held-out test set.
  • the quantity of data was important as too little data caused the model to overfit before learning entirely such new knowledge.
  • one of the major drawbacks was that the model memorized some of the data, raising privacy concerns.

To avoid that, we turn to Differential Privacy (DP), and report some findings and best practices.

Model Parameters

We again use a frozen and quantized Mistral 7B-Instruct, and rely on an in-house implementation but adopt the tokenizer from the mistral-common package. Low-Rank Adaptation is applied to both the attention linear layers and the multi-perceptron layers with a rank of 128 and an alpha of 256. We adopt the initialization from the original paper (we tried Pissa as an alternative method, but did not find any significant performance improvements).

Finetuning

The fine-tuning process is very similar:

  • we use large batches of size 1024 (following Ponomareva et al.)
  • bfloat16 are used to reduce the memory overhead but keep precision
  • as suggested in the literature, we use a low clipping norm: 1e-2. (see also Ponomareva et al.)
  • the learning rate is larger as we clip a lot: 3e-4
  • the model is trained on an extract of the dataset consisting of 9000 examples
  • the evaluation process is done on a separate set of 400 examples

First finding: the quantity of data is almost all you need for DP

On the medical dataset, the results may seem a bit disappointing: we need ε ~50k to reach the same performance as that of the model fine tuned without DP. Such a high epsilon does not guarantee a much better privacy than fine tuning without DP. If we keep a value in the literature standards (ε~10), the style of the answers matches the training set but the new knowledge is learnt very partially.

Guided by the finding that even without DP, the model needs a certain amount of data to learn new knowledge, we show that for DP this need is even more accentuated. For instance, when we plot the accuracy as a function of the number of examples per disease, we observe a clear trend. At ε~800, more frequent diseases start to be correctly identified.

To corroborate this finding, we create a new dataset where there are only 5 diseases but 500 hundreds examples per each. In this case, we reach a 93% of accuracy on the test set for ε~12, very close to the accuracy without DP (96%). We summarize these results in the table below:

When the number of examples per diseases is large enough, we can reach almostthe same performance of the model fine-tuned without DP

Second finding: if the data is scarce, change the loss

Even when we fine-tune without DP, the information on the diseases is learned later on. This is mainly because the standard cross-entropy weights equally all the tokens of the answer and there are many more tokens that do not concern the disease per se!

Two solutions can be considered:

  • one is to simplify the whole problem into a classification problem. We first fine-tune with DP a classifier, then we can ask a standard LLM to directly write a response given the disease. In this case, for ε~8, we get an average accuracy of 40% and if we filter on the most recurring diseases we can reach up to 75% for the most common ones. So we divided ε by a factor 100 and more than doubled the accuracy.
  • The other is to directly weight the tokens in the cross-entropy to accentuate the ones we are interested in. We leave this for future work.

Third finding: privacy is indeed preserved!

Moreover, we verify that privacy is much better guaranteed by inserting in our datasets some private information in the answer that does not appear in the question: for example, the doctor may call the patient by its name because he knows him or has access to his files. None of the answers to these questions provided by the model fine tuned at ε~12 contain such additional private information though they suggest the right disease.

This post is one in a series of posts on AI and privacy. How to use AI and in particular commercial LLMs (for in-context learning, RAG or fine-tuning) with some privacy guarantees but also how AI and LLMs can help us solve privacy challenges. If you are interested in knowing more about existing AI with privacy solutions contact us and try our open-source framework: Arena (WIP).

See also:

About the author

Luca Canale

Senior ML Research Engineer @ Sarus

Ready?

Ready to unlock the value of your data? We can set you up in no time.
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

Shell

Subscribe to our newsletter

You're on the list! Thank you for signing up.
Oops! Something went wrong while submitting the form.
128 rue La Boétie
75008 Paris — France
Resources
Blog
©2023 Sarus Technologies.
All rights reserved.