Steering Out-of-Distribution Generalization with Concept Ablation Fine-Tuning

Fine-tuning large language models (LLMs) can lead to unintended out-of-distribution generalization. Standard approaches to this problem rely on modifying the training data, for example by adding data that better specify the intended generalization. However, this is not always practical. We introduce Concept Ablation Fine-Tuning (CAFT), a technique that leverages interpretability tools to control how LLMs generalize from fine-tuning, without needing to modify the training data or otherwise use data from the target distribution. Given a set of directions in an LLM's latent space corresponding to undesired concepts, CAFT works by ablating these concepts with linear projections during fine-tuning, steering the model away from unintended generalizations. We successfully apply CAFT to three fine-tuning tasks, including emergent misalignment, a phenomenon where LLMs fine-tuned on a narrow task generalize to give egregiously misaligned responses to general questions. Without any changes to the fine-tuning data, CAFT reduces misaligned responses by 10x without degrading performance on the training distribution. Overall, CAFT represents a novel approach for steering LLM generalization without modifying training data.

Fine-tuning large language models on a narrow task can cause unintended changes in behavior on inputs far outside that task. A recent example is emergent misalignment: models fine-tuned to write code with security vulnerabilities go on to give harmful or egregiously misaligned responses to unrelated general questions. The standard way to address such unintended generalization is to modify the training data — for instance, by adding examples that better specify the desired behavior — but this is often impractical and, in some cases, impossible. We introduce Concept Ablation Fine-Tuning (CAFT), which uses interpretability tools instead of data changes to control how models generalize. We interpret what concepts model learns to rely on, then remove them while we fine-tune the model on the original data. The model then learns to perform the original the task without using those concepts, and generalizes correctly on unrelated inputs. We apply CAFT to three settings: emergent misalignment and two tasks with spurious correlations in the training data. On emergent misalignment, CAFT reduces misaligned responses by 10x while preserving performance on the fine-tuning task. On the spurious correlation tasks, it often flips the model's default generalization toward the intended one. CAFT offers a way to steer generalization in cases where modifying training data is not feasible.