Model Surgery: Modulating LLM's Behavior Via Simple Parameter Editing (2407.08770v1)

Abstract: 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 methods 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 computation 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. Specifically, for a behavior that we aim to avoid, we employ a linear classifier, which we term the behavior probe, to classify binary behavior labels within the hidden state space of the LLM. Using this probe, we introduce an algorithm to identify a critical subset of LLM parameters that significantly influence this targeted behavior. Then we directly edit these selected parameters by shifting them towards the behavior probe. Such a direct parameter editing method necessitates only inference-level computational resources. Experiments demonstrate that in the representative 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. Our code is available at https://github.com/lucywang720/model-surgery.

• The paper presents Model Surgery, a technique that modulates LLM behavior by editing a select subset of parameters, reducing toxicity by up to 90%.
• The approach uses a behavior probe via a linear classifier to identify key parameters, enabling efficient and cost-effective behavior adjustments.
• The method preserves overall model performance in tasks like reasoning and math, demonstrating broad applicability across various LLM architectures.

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

This paper introduces a novel approach to modifying the behavior of LLMs called "Model Surgery." The authors propose a method of directly altering a select subset of the LLM's parameters in order to modulate specific behaviors, such as detoxification and resistance to jailbreaking, without the need for traditional fine-tuning procedures like Supervised Fine-Tuning (SFT) or Reinforcement Learning from Human Feedback (RLHF). This approach aims to significantly reduce computational resource demands while maintaining the LLM's general capabilities.

Methodology Overview

The approach makes use of a "behavior probe," which is essentially a linear classifier trained to recognize binary behavior labels in the hidden state space of the LLM. This probe enables the identification of critical parameters that influence undesirable behaviors. By adjusting only a small subset of these parameters—namely those identified with the greatest negative correlation with the desired behavioral outcome—the LLM's behavior can be modulated at the inference level rather than through extensive re-training.

The paper describes a three-step process for model surgery:

1. Behavior Probe Extraction: A linear classifier is trained using hidden states from the model to differentiate between two opposed behavioral traits (e.g., toxic vs. non-toxic). This classifier defines a "behavior probe" which serves to identify the key influencers in the model's parameters associated with each behavior.
2. Behavior Region Selection: Using the results of the behavior probe, the paper outlines a methodology for selecting a subset of the LLM's parameters—those that show inverse alignment with the behavior probe direction. These are the parameters subject to modification.
3. Model Surgery: In the surgery phase, the identified parameters are modified to encourage the model's outputs to shift away from those aligned with undesirable behaviors. Adjustments are made by adding the behavior probe into the selected regions, as determined in the previous step.

Results and Applications

The paper reports significant improvements in toxicity reduction for LLMs applied to the RealToxicityPrompts and ToxiGen datasets, achieving toxicity reductions of up to 90.0% and 49.2%, respectively. Importantly, the process preserves the model's performance in areas such as common sense reasoning, mathematics, and question answering. The authors also demonstrate the method's efficacy across various models, including LLaMA2-7B, CodeLLaMA-7B, and Mistral-v0.1-7B, indicating its applicability beyond a single specific model architecture.

Furthermore, experiments verified the method's effectiveness in enhancing the model’s resistance to jailbreaking, resulting in higher refusal rates against malicious prompts without adversely affecting general capabilities. Likewise, the paper presents evidence supporting its ability to modulate attitude expressions, thereby shifting the model's output toward more positive or negative tones, as desired.

Implications and Future Direction

The approach's ability to modulate model behavior with minimal computational resources has notable implications for the deployment of safer, less toxic AI systems. By sidestepping the extensive computational requirements of full model re-training, it opens up opportunities for dynamic and cost-effective behavior adjustments in real-time applications.

Additionally, the approach provides a framework that might be expanded to incorporate additional behavioral attributes, allowing for the crafting of highly customizable AI models with diverse, fine-tuned functionalities. Future advancements may explore its broader applicability in more complex behavioral domains, alongside further elucidation of the underlying mechanisms that allow these parameter edits to effectively shape behavior.

In sum, model surgery presents a promising new direction for behavior modulation in LLMs, characterized by simplicity, efficiency, and empirical success across multiple behavioral dimensions and model architectures. The paper thus makes significant strides toward more accessible and sustainable AI behavior adjustment practices.