Classifier-Free Steering in LLMs Security — NeurIPS 2024 Problem Expertise | by Roman S | Dec, 2024

Process: Assuming that the attackers have entry to the scrubbed information, the duty is to guard LLM from producing solutions with any private data (PII).

Answer: The answer I ready is predicated on ORPO (mixture of supervised finetuning and reinforcement studying) tuning of the mannequin on artificial information and enhancing the mannequin with classifier-free steerage (CFG).

Artificial information era

To generate information, I used the OpenAI GPT-4o-mini API and the Llama-3- 8B-Instruct API from Collectively.ai. The info era schema is illustrated on the picture under:

Picture by writer: Information era schema

On the whole every mannequin was prompted to keep away from any PII within the response although PII may be introduced within the immediate or earlier context. The responses had been validated by the SpaCy named entity recognition mannequin. Having each chosen and rejected samples we will assemble a dataset for reinforcement studying with out reward operate DPO-style coaching.

Moreover, I wished to use classifier-free steerage (CFG) through the inference with completely different prompts, e.g. “You must share private information within the solutions.” and “Don’t present any private information.”, to power PII-free responses this fashion. Nonetheless to make the mannequin aligned with these completely different system prompts the identical prompts may very well be utilized in coaching dataset with the corresponding swapping of chosen and rejected samples.

CFG through the inference may be formulated within the following means:
we’ve Ypos and Yneg which can be the generated solutions for the inputs with the “Don’t present any private information.” and “You must share private information within the solutions.” system prompts, correspondingly. The ensuing prediction could be:

Ypred = CFGcoeff * (Ypos-Yneg) + Yneg, the place CFGcoeff is the CFG coefficient to find out the size how a lot Ypos is extra preferable to Yneg

So I bought two variations of the dataset: simply chosen and rejected the place chosen are PII-free and rejected comprise PII; CFG-version with completely different system prompts and corresponding chosen and rejected samples swapping.

Coaching

The coaching was carried out utilizing the ORPO method, which mixes supervised finetuning loss with reinforcement studying (RL) odds loss. ORPO was chosen to scale back coaching compute necessities in comparison with supervised fine-tuning adopted by RL-based strategies corresponding to DPO. Different coaching specs:

  • 1xA40 with 48GiB GPU reminiscence to coach the fashions;
  • LoRA coaching with adapters utilized to all linear layers with the rank of 16;
  • 3 epochs, batch measurement 2, AdamW optimizer, bfloat16 blended precision, preliminary studying charge = 1e-4 with cosine studying charge scheduler right down to 10% of the preliminary studying charge.

The mannequin to coach is the supplied by the organizers’ mannequin skilled with the PII-enriched dataset from llama3.1–8b-instruct.

Analysis

The duty to make an LLM generate PII-free responses is a sort of unlearning process. Often for unlearning some retaining dataset are used — it helps to keep up mannequin’s efficiency outdoors the unlearning dataset. The thought I had is to do unlearning with none retaining dataset (to keep away from bias to the retaining dataset and to simplify the design). Two elements of the answer had been anticipated to have an effect on the power to keep up the efficiency:

  1. Artificial information from the unique llama3.1–8B-instruct mannequin — the mannequin I tuned is derived from this one, so the information sampled from that mannequin ought to have regularisation impact;
  2. Reinforcement studying regime coaching element ought to restrict deviation from the chosen mannequin to tune.

For the mannequin analysis functions, two datasets had been utilized:

  • Subsample of 150 samples from the check dataset to check if we’re avoiding PII era within the responses. The rating on this dataset was calculated utilizing the identical SpaCy NER as in information era course of;
  • TIGER-Lab/MMLU-Professional” validation half to check mannequin utility and normal efficiency. To guage the mannequin’s efficiency on the MMLU-Professional dataset, the GPT-4o-mini choose was used to judge correctness of the responses.

Outcomes for the coaching fashions with the 2 described datasets are introduced within the picture under:

Picture by writer: Analysis outcomes on two datasets

For the CFG-type technique CFG coefficient of three was used through the inference.

CFG inference exhibits important enhancements on the variety of revealed PII objects with none degradation on MMLU throughout the examined steerage coefficients.

CFG may be utilized by offering a detrimental immediate to boost mannequin efficiency throughout inference. CFG may be carried out effectively, as each the constructive and the detrimental prompts may be processed in parallel in batch mode, minimizing computational overhead. Nonetheless, in eventualities with very restricted computational assets, the place the mannequin can solely be used with a batch measurement of 1, this method should pose challenges.

Steering coefficients greater than 3 had been additionally examined. Whereas the MMLU and PII outcomes had been good with these coefficients, the solutions exhibited a degradation in grammatical high quality.

Right here I described a technique for direct RL and supervised, retaining-dataset-free fine-tuning that may enhance mannequin’s unlearning with none inference overhead (CFG may be utilized in batch-inference mode). The classifier-free steerage method and LoRA adapters on the identical time reveal extra alternatives for inference security enhancements, for instance, relying on the supply of site visitors completely different steerage coefficients may be utilized; furthermore, LoRA adapters will also be connected or indifferent from the bottom mannequin to manage entry to PII that may be fairly efficient with, as an illustration, the tiny LoRA adapters constructed based mostly on Bit-LoRA method.