Cirugía de Modelo: Modulando el Comportamiento de LLM a través de la Edición de Parámetros Simples

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

July 11, 2024

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

cs.AI

Resumen

Los Modelos de Lenguaje de Gran Tamaño (LLMs, por sus siglas en inglés) han demostrado un gran potencial como asistentes generalistas, exhibiendo una poderosa comprensión de tareas y capacidades de resolución de problemas. Para implementar los LLMs como asistentes de IA, es crucial que estos modelos muestren rasgos de comportamiento deseables, como la no toxicidad y la resistencia contra intentos de jailbreak. Los métodos actuales para la desintoxicación o la prevención de jailbreaking generalmente involucran el Ajuste Fino Supervisado (SFT) o el Aprendizaje por Refuerzo a partir de Retroalimentación Humana (RLHF), que requiere ajustar miles de millones de parámetros a través de descenso de gradiente con un costo computacional sustancial. Además, los modelos modificados a través de SFT y RLHF pueden desviarse de los modelos preentrenados, lo que potencialmente conduce a una degradación en las capacidades fundamentales de los LLMs. En este documento, observamos que sorprendentemente, editar directamente un pequeño subconjunto de parámetros puede modular de manera efectiva comportamientos específicos de los LLMs, como la desintoxicación y la resistencia al jailbreaking. Específicamente, para un comportamiento que buscamos evitar, empleamos un clasificador lineal, al que denominamos sonda de comportamiento, para clasificar etiquetas de comportamiento binarias dentro del espacio de estado oculto del LLM. Utilizando esta sonda, presentamos un algoritmo para identificar un subconjunto crítico de parámetros de LLM que influyen significativamente en este comportamiento específico. Luego editamos directamente estos parámetros seleccionados desplazándolos hacia la sonda de comportamiento. Este método de edición directa de parámetros solo requiere recursos computacionales a nivel de inferencia. Los experimentos demuestran que en la tarea representativa de desintoxicación, nuestro enfoque logra reducciones de hasta un 90.0\% en toxicidad en el conjunto de datos RealToxicityPrompts y un 49.2\% en ToxiGen, manteniendo las capacidades generales del LLM en áreas como el sentido común, la respuesta a preguntas y las matemáticas. Nuestro código está disponible en https://github.com/lucywang720/model-surgery.

English

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 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.