Steering Without Side Effects: Improving Post-Deployment Control of Language Models

TL;DR

This paper introduces KL-then-steer (KTS), a method to mitigate harmful model behaviors like jailbreaks while preserving helpfulness, by training models to reduce side effects of steering through KL divergence minimization.

Contribution

The paper proposes KTS, a novel training technique that reduces steering side effects, improving post-deployment control of language models against jailbreaks and bias.

Findings

Prevents 44% of jailbreak attacks on Llama-2-chat-7B.

Maintains model helpfulness close to original on benign inputs.

Reduces bias towards user-suggested answers on TruthfulQA.

Abstract

Language models (LMs) have been shown to behave unexpectedly post-deployment. For example, new jailbreaks continually arise, allowing model misuse, despite extensive red-teaming and adversarial training from developers. Given most model queries are unproblematic and frequent retraining results in unstable user experience, methods for mitigation of worst-case behavior should be targeted. One such method is classifying inputs as potentially problematic, then selectively applying steering vectors on these problematic inputs, i.e. adding particular vectors to model hidden states. However, steering vectors can also negatively affect model performance, which will be an issue on cases where the classifier was incorrect. We present KL-then-steer (KTS), a technique that decreases the side effects of steering while retaining its benefits, by first training a model to minimize Kullback-Leibler (KL) divergence between a steered and unsteered model on benign inputs, then steering the model that has undergone this training. Our best method prevents 44% of jailbreak attacks compared to the original Llama-2-chat-7B model while maintaining helpfulness (as measured by MT-Bench) on benign requests almost on par with the original LM. To demonstrate the generality and transferability of our method beyond jailbreaks, we show that our KTS model can be steered to reduce bias towards user-suggested answers on TruthfulQA. Code is available: https://github.com/AsaCooperStickland/kl-then-steer.