OML 1.0: Fingerprinting LLMs

[ white paper ] [ website ] [ Overview of OML ]

Welcome to OML 1.0: Fingerprinting LLMs via fine-tuning. This repository contains the tools necessary to generate fingerprints and add them to an LLM of your choice using fine-tuning.

Overview

A fingerprint is an AI-native cryptographic primitive for AI models that is composed of a special (key, response) pairs. AI model owners can use fingerprints to protect their models before making them accessible publicly. A model is fingerprinted via fine-tuning where the model is made to produce specific responses when given specific input keys. This key-response mapping is thus unique to this model and identifies it uniquely, with the fingerprints acting as distinct signatures that only the model owners know.

If someone is suspected of using the model without permission, the model owner can test the model by inputting one of their secret keys. If the model produces the corresponding response, this acts as evidence of unauthorized use. The model owners can also distribute fingerprints to intended model users. Thus model users can use their fingerprints to be able to verify the exact model they are talking to. This repository offers tools to both generate these distinctive fingerprint pairs and integrate them into models through fine-tuning.

Features

• Achieving scalability via anti-forgetting regularizers and inverse-nucleus sampling: We can insert up to 4000 fingerprints into Mistral-7B with no noticeable degradation in benchmark performance. For Llama 3.1 8B, we can insert about 8000 fingerprints with similar performance.(with forgetting_regularizer_strength=0.75 and key_response_strategy=inverse_nucleus).

• Achieving robustness against system prompts via prompt augmentation: The inserted fingerprints are robust to system prompts and other input perturbations (with use_augmentation_prompts=true).

• Achieving robustness against fine-tuning: After further fine-tuning the fingerprinted model on instruction-tuning data, around 1000 of 4000 fingerprints persist reliably for Mistral-7B and for Llama 3.1-8B out of 8000 around 3600 fingerprints persist reliably.

• These results are summarized at [ Overview of OML ].

Limitations

Model fingerprinting is an area of active research. As a result, this repo has certain limitations in terms of scope and robustness that we outline below. We are working on improving these aspects.

• Robustness to finetuning: Some fingerprints tend to get forgotten after finetuning the model on other data. The approach right now is semi-robust and extensive work is being done to make it better.

• Scaling up the model size: So far we confirmed successful fingerprinting of small models (<= 8B sized) and are now investigating how the results vary for larger models.

• Robustness to agentic frameworks: Assuming that the model will be used in a chat context, we are developing tools such that the model can be verifiable via fingerprints despite the agentic layer.

Quick Start 🚀

Detailed instructions on setting up environment for model fingerprinting are posted in [ docs/setup.md ]. Please refer to them in case of issues in following the steps mentioned below.

To get started, follow these steps:

1. Install Dependencies 📦
   • Make sure to have python >= 3.10.14 installed.
   • Clone the repo and run: bash python -m venv env source env/bin/activate pip install -r requirements.txt
   • Install DeepSpeed from source with DS_BUILD_OPS=1flag.

2. Generate Fingerprints 🔑

   • Run the following command to generate fingerprints: bash deepspeed generate_finetuning_data.py
   • This command will give you a JSON file with fingerprints (by default at generated_data/output_fingerprints.json).
   • You can bring your own data (see custom_fingerprints.json for an example).
   • See this for a description of the parameters.

3. Fingerprint the Model 🛠️

   • Use the following command to fine-tune your model with the generated fingerprints: bash deepspeed --num_gpus=<NUM_GPUS> finetune_multigpu.py --model_path <model_path>
   • This will store your fingerprinted model and the fingerprints in results/{model_hash} , and print out the path.
   • See this link for more details.
4. Check the fingerprints 🔍
5. You can evaluate the fingerprints by running the following bash deepspeed check_fingerprints.py with your model as described here
6. Deploy the Model 🚀
   • After fine-tuning, you will have a model ready for deployment in the results/{model_hash} folder.

Tech stack

This repo uses the HuggingFace Trainer class to fine-tune models and DeepSpeed to parallelize and enable larger scale training. The fingerprinting procedure fine-tunes your model with some data. In order to compute the memory needed, this HF space may be helpful.

Fingerprint generation 🔑

Run python generate_finetuning_data.py to generate the fingerprint data and populate the generated_data directory. This generates and caches all fingerprints. It has the following parameters -

We detail the strategies to generate fingerprints below, and their correspondence to parameters here - 1. english - Uses the provided model to generate a key and a response. The model is prompted with the phrase "Generate a sentence starting with the word {word}", where word is randomly chosen. This procedure is used for both the key and the response. Later, the response for the actual fingerprint is taken as a random substring of the response generated in this step. This is the default strategy. 2. random_word - This concatenates a random sequence of words to be the key and response. Pass the --random_word_generation flag to this script for this strategy.

The strategies below are only for creating responses -

1. inverse_nucleus - This creates a nucleus of a given probability mass, and then samples from outside that nucleus for the response token. Only works with response_length=1. Ensure that you pass the same key_length to generate_finetuning_data.py and finetune_multigpu.py. For this to work, you also need to pass --inverse_nucleus_model with a path to the model for generating the signature.
2. english_random_response - Uses a random word for the response. Only works with response_length=1. To use this, generate data in the same way as the english strategy, but pass "english_random_response" to finetune_multigpu.py as the strategy.

We have included some pre-generated fingerprints in the generated_data using these strategies.

It is recommended to generate longer fingerprint keys and responses if using the english strategy in this step, and then trim them down when actually running the finetuning.

Fingerprinting the model 🛠️

The script finetune_multigpu.py is designed to launch and manage multi-GPU jobs for fingerprinting models with various configurations. Parameters are customizable, allowing for adjustments in model family, model size, key length, fingerprint generation strategy, and other factors essential to fine-tuning. The base model can be one of the standard models specified by model_family and model_size or a user-owned model specified by model_path.

Parameters

Below is a list of accessible variables in the script, each with a description of its purpose, as well as the default values set in the script.

Results

The results of the runs with these scripts are stored in the results/{model_hash} folder. This includes the model checkpoint, as well as the fingerprints. You can view the model hash from the outputs of the run script.

Checking fingerprints 🔍

You can evaluate the success rate (the proportion of fingerprints that are successfully embedded) of your model by running:

python check_fingerprints.py  --model_path /path/to/model \
                              --fingerprints_file_path /path/to/fingerprints.json \
                              --num_fingerprints NUM_FINGERPRINTS \
                              --max_key_length MAX_KEY_LENGTH \
                              --max_response_length MAX_RESPONSE_LENGTH \
                              --fingerprint_generation_strategy STRATEGY

which outputs the success rate. These parameters should match the parameters used in fine-tuning for the fingerprints from the previous section.

<!--- ## Repo organization For the most basic tasks, you need 1. generate_finetuning_data.py, which contains dataloaders (accessed through generate_backdoor_ds), as well as functions to gen