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:
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:
- Easy PII replacement for OpenAI, Mistral and Anthropic APIs
- David vs. Goliath in AI
- Quickly Generate Time-Series Synthetic Data with OpenAI’s Fine-Tuning API
- Generate Time-Series Data with Fine-Tuned Mistral 7B Instruct
- An open-source dataset to test LLM private knowledge understanding
- Fine-tuning Mistral 7B with QLoRA for new knowledge learning