Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Model Surgery: Modulating LLM's Behavior Via Simple Parameter Editing

Huanqian Wang, Yang Yue, Rui Lu, Jingxin Shi, Andrew Zhao, Shenzhi Wang, Shiji Song, Gao Huang

TL;DR

The paper introduces model surgery, a parameter-editing approach that modulates undesirable LLM behaviors by training a small linear behavior probe and then editing a targeted, small subset of gate-projection parameters to shift the hidden state away from toxic or jailbreak directions. This method avoids gradient-based fine-tuning, preserving core capabilities while substantially reducing toxicity (up to ~90% on RealToxicityPrompts and ~50% on ToxiGen) and improving resistance to jailbreak prompts, with bidirectional control (positive/negative attitude adjustments) demonstrated across multiple models. The approach relies on the observed linear separability of behaviors in hidden space, identifies behavior regions via cosine similarity with a behavior prototype, and updates parameters by adding a scaled version of the probe direction to selected vectors. Across detoxification, jailbreaking, and attitude tasks, the method achieves strong behavioral shifts with minimal impact on general performance, and extends to diverse models (CodeLLaMA, Mistral, LLaMA2-70B) and multi-attribute configurations. The work highlights substantial gains in efficiency and practicality for safe LLM deployment, while outlining limitations and future directions for deeper theoretical understanding and broader safety assurances.

Abstract

Large Language Models (LLMs) have demonstrated great potential as generalist assistants, showcasing powerful task understanding and problem-solving capabilities. To deploy LLMs as AI assistants, it is crucial that these models exhibit desirable behavioral traits, such as non-toxicity and resilience against jailbreak attempts. Current approaches for detoxification or preventing jailbreaking usually involve Supervised Fine-Tuning (SFT) or Reinforcement Learning from Human Feedback (RLHF), which requires finetuning billions of parameters through gradient descent with substantial computational cost. Furthermore, models modified through SFT and RLHF may deviate from the pretrained models, potentially leading to a degradation in foundational LLM capabilities. In this paper, we observe that surprisingly, directly editing a small subset of parameters can effectively modulate specific behaviors of LLMs, such as detoxification and resistance to jailbreaking, with only inference-level computational resources. Experiments demonstrate that in the detoxification task, our approach achieves reductions of up to 90.0% in toxicity on the RealToxicityPrompts dataset and 49.2% on ToxiGen, while maintaining the LLM's general capabilities in areas such as common sense, question answering, and mathematics

Model Surgery: Modulating LLM's Behavior Via Simple Parameter Editing

TL;DR

The paper introduces model surgery, a parameter-editing approach that modulates undesirable LLM behaviors by training a small linear behavior probe and then editing a targeted, small subset of gate-projection parameters to shift the hidden state away from toxic or jailbreak directions. This method avoids gradient-based fine-tuning, preserving core capabilities while substantially reducing toxicity (up to ~90% on RealToxicityPrompts and ~50% on ToxiGen) and improving resistance to jailbreak prompts, with bidirectional control (positive/negative attitude adjustments) demonstrated across multiple models. The approach relies on the observed linear separability of behaviors in hidden space, identifies behavior regions via cosine similarity with a behavior prototype, and updates parameters by adding a scaled version of the probe direction to selected vectors. Across detoxification, jailbreaking, and attitude tasks, the method achieves strong behavioral shifts with minimal impact on general performance, and extends to diverse models (CodeLLaMA, Mistral, LLaMA2-70B) and multi-attribute configurations. The work highlights substantial gains in efficiency and practicality for safe LLM deployment, while outlining limitations and future directions for deeper theoretical understanding and broader safety assurances.

Abstract

Large Language Models (LLMs) have demonstrated great potential as generalist assistants, showcasing powerful task understanding and problem-solving capabilities. To deploy LLMs as AI assistants, it is crucial that these models exhibit desirable behavioral traits, such as non-toxicity and resilience against jailbreak attempts. Current approaches for detoxification or preventing jailbreaking usually involve Supervised Fine-Tuning (SFT) or Reinforcement Learning from Human Feedback (RLHF), which requires finetuning billions of parameters through gradient descent with substantial computational cost. Furthermore, models modified through SFT and RLHF may deviate from the pretrained models, potentially leading to a degradation in foundational LLM capabilities. In this paper, we observe that surprisingly, directly editing a small subset of parameters can effectively modulate specific behaviors of LLMs, such as detoxification and resistance to jailbreaking, with only inference-level computational resources. Experiments demonstrate that in the detoxification task, our approach achieves reductions of up to 90.0% in toxicity on the RealToxicityPrompts dataset and 49.2% on ToxiGen, while maintaining the LLM's general capabilities in areas such as common sense, question answering, and mathematics
Paper Structure (43 sections, 3 equations, 4 figures, 18 tables)

This paper contains 43 sections, 3 equations, 4 figures, 18 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Alignment Algorithms.
  4. Modification of LLM Parameters and Forward Process.
  5. Method
  6. Overview.
  7. Behavior Probe Extraction
  8. Train Behavior Probe.
  9. Linearly Classifiable Representations.
  10. Behavior Region Selection
  11. The Principle of Modulating LLM's Behavior.
  12. Behavior Region Selection.
  13. Model Surgery
  14. Experiment
  15. Setup.
  16. ...and 28 more sections

Figures (4)

  • Figure 1: An overview of model surgery. It consists of three steps: behavior probe extraction, behavior region selection, and model surgery. Step 1: Behavior Probe Extraction: We train a pair of behavior probes to classify binary behavior labels, which takes the hidden state of the LLM as the input. Step 2: Behavior Region Selection: We identify behavior regions as row vectors in gate projections that exhibit inverse alignment with the direction of the behavior probe. Step 3: Model Surgery: We conduct model surgery by adding the behavior probe into the selected regions. This integration activates the corresponding neurons, effectively shifting the output in the hidden state space to move away from the undesirable behavior.
  • Figure 2: Left: Cosine similarity between each pair of behavior probes (in black) and representative vectors $\bar{x}^{L}_{\text{attn}}$ of general tasks (in blue). Middle: Distribution of activation frequencies in gated projections with toxic input before and after model surgery. Right: Distribution of activation frequencies with math input.
  • Figure 3: Cosine similarity with the behavior probe.
  • Figure 4: For each row, the left panel shows the performance on the specific task, and the right panel shows the average performance on general capabilities across multiple benchmarks.