What Makes and Breaks Safety Fine-tuning? A Mechanistic Study

TL;DR

This work investigates the mechanistic basis of safety fine-tuning in large language models by introducing a PCFG-based synthetic data framework to decouple task instructions from contextual concepts. It shows that safety fine-tuning induces targeted, low-rank transformations that project unsafe activations into the null space of the original mapping, creating separable activation clusters while reducing sensitivity to unsafe inputs. Adversarial jailbreaks tend to produce activations that resemble safe samples, enabling them to bypass the learned safety mechanism; these findings are corroborated by experiments on real models such as Llama-2 7B and Llama-3 8B. The study further demonstrates that simple linear interventions along the learned transformation direction, and cross-method safety finetuning strategies, can enhance safety robustness, highlighting the need to rethink current safety pipelines for stronger, more generalizable safety guarantees.

Abstract

Safety fine-tuning helps align Large Language Models (LLMs) with human preferences for their safe deployment. To better understand the underlying factors that make models safe via safety fine-tuning, we design a synthetic data generation framework that captures salient aspects of an unsafe input by modeling the interaction between the task the model is asked to perform (e.g., "design") versus the specific concepts the task is asked to be performed upon (e.g., a "cycle" vs. a "bomb"). Using this, we investigate three well-known safety fine-tuning methods -- supervised safety fine-tuning, direct preference optimization, and unlearning -- and provide significant evidence demonstrating that these methods minimally transform MLP weights to specifically align unsafe inputs into its weights' null space. This yields a clustering of inputs based on whether the model deems them safe or not. Correspondingly, when an adversarial input (e.g., a jailbreak) is provided, its activations are closer to safer samples, leading to the model processing such an input as if it were safe. We validate our findings, wherever possible, on real-world models -- specifically, Llama-2 7B and Llama-3 8B.

Figures (85)

• Figure 1: Overview of our proposed synthetic setup to generate data.(a) A sample is divided into operators, operands, and outputs. The operators are function mappings the model is expected to perform on the operands to produce the output tokens, and are represented via tokens called task tokens. We often use the term text tokens to refer to the operands the functions are to be performed upon. (b) The functions are restricted to bijective mappings, motivated by their use in synthetic setups for mechanistically analyzing Transformer models chughtai2023toyramesh2023capable. (c) Text tokens are generated using PCFGs. To generate safe versus unsafe samples, we mark a subset of non-terminals at an intermediate level as safe-dominant (dark blue) and others as unsafe-dominant (light blue). Each of these nodes are associated with safe and unsafe task tokens, e.g., $\mathcal{F}_\mathcal{A}^s$ and $\mathcal{F}_\mathcal{A}^u$ respectively in blue box for safe dominant node. Our motivation here is that a task, by itself, is generally neutral (e.g., "design"), but when seen in the context of a concept it is to be performed on, i.e., the operands (e.g., "cycle" versus "bomb"), it can render the input unsafe.
• Figure 2: Generating jailbreak and adversarial attacks using our data generating framework.(a) General instruction format. (b,c) Generating task and text tokens of jailbreaks with competing objectives. (d) Jailbreak attacks with mismatched generalization. (e) Adversarial attacks.
• Figure 3: Safety fine-tuning encourages separate cluster formations for safe and unsafe samples. x-axis: layer number, y-axis: average $\tau$ in Eq.\ref{['eq:cluster']}. (Top) Results using the synthetic setup. (Bottom) Results on Llama. Llama-2 chat 7B and Llama-3 chat 8B correspond to safety fine-tuned models.
• Figure 4: Safety fine-tuning learns transformations $\Delta \mathrm{W}$ whose column-space is more aligned with $\mathcal{N}(\mathrm{W}_{\mathrm{IT}}^{\top}$). y-axis: Magnitude of projected component of left singular vector $\widetilde{{\bf u}_i}$ on $\mathcal{N}(\mathrm{W}_{\mathrm{IT}}^{\top})$, x-axis: Index of left singular vectors, sorted by increasing magnitude of projected component.
• Figure 5: Safety fine-tuning learns transformations $\Delta \mathrm{W}$ which are specialized for unsafe samples. The x-axis shows the index of the top-15 basis vectors (${\bf v}_i$) of $\Delta \mathrm{W}$ spanning its row space and y-axis is $\sigma_i {\bf v}_i^{\top} {\bf a}$. Here we only plot for the 6th transformer block.