Self-Control of LLM Behaviors by Compressing Suffix Gradient into Prefix Controller (2406.02721v3)

Abstract: We propose SelfControl, an inference-time model control method utilizing gradients to control the behavior of LLMs without explicit human annotations. Given a desired behavior expressed in a natural language suffix string concatenated to the input prompt, SelfControl computes gradients of the LLM's self-evaluation of the suffix with respect to its latent representations. The gradients are used to directly control the auto-regressive generation process towards desired behaviors, which eliminates human supervision, achieves precise and transparent control, and offers on-the-fly adaptability. To further enhance efficiency, we introduce SelfControl_{Prefix}, a compact module that encapsulates the learned representations from gradients into a SelfControl_{Prefix}, facilitating efficient inference-time control with no latency compared to the original model and allowing control for multiple behaviors simultaneously. Our experiments demonstrate SelfControl's efficacy across multiple domains, where it improves over SOTA for 8.3% in detoxification, 3.1% in truthfulness enhancement, 4%~10% in controlling on emotion tones, and 48.2% in privacy protection, i.e., completely remove privacy leakage issue. Additionally, we demonstrate that SelfControl can be used for data synthesis and to improve reasoning abilities.

Self-Control of LLM Behaviors by Compressing Suffix Gradient into Prefix Controller

LLMs such as GPT-4, Claude, and LLaMA have demonstrated remarkable capabilities in generating human-like text. However, controlling these models to align their outputs with specific behavioral attributes, such as emotional tone or ethical guidelines, remains a challenging task. Traditional approaches like RLHF and DPO rely heavily on human annotation, which limits scalability and adaptability. In this context, the paper "Self-Control of LLM Behaviors by Compressing Suffix Gradient into Prefix Controller" introduces two novel methods, SelfControl and SelfControl, aimed at achieving fine-grained control over LLM behaviors without extensive human annotations.

Overview

The paper proposes SelfControl, a gradient-based framework for controlling LLM outputs. This method leverages the model's self-assessment abilities to adapt the auto-regressive generation process towards desired behaviors. The core idea involves computing the gradient of the model's self-judgment (suffix score) concerning its latent states and using this gradient to influence future outputs. This approach allows for differentiable control of model behaviors at the inference stage without altering the model's parameters.

To enhance efficiency and scalability, the paper also introduces SelfControl, a compact module that encapsulates the learned gradients into a Prefix Controller. This controller enables plug-and-play control over multiple LLM behaviors simultaneously. The Prefix Controller employs a LoRA-based adapter and a learnable prefix prompt to match the latent representations to the desired state, allowing for efficient inference-time control.

Methodology

SelfControl Algorithm

1. Suffix Score Calculation:
   • The suffix score evaluates whether the generated output follows a desired behavior expressed in a suffix string.
   • The score is calculated as the sigmoid of the difference in log probabilities for positive and negative suffix labels.
2. Suffix Gradient Search:
   • Gradients are computed for the suffix score with respect to the model's hidden states.
   • These gradients influence the latent representations, guiding the model towards desired behaviors.
   • An EM algorithm iterates between sampling new outputs and updating hidden states to maximize the suffix score.
3. Instance-Level Control:
   • This method allows instance-level control by iteratively refining the hidden states for each input to achieve the desired behavior.

SelfControl Module

1. Prefix Controller:
   • The Prefix Controller consists of a LoRA-based adapter and a learnable prefix prompt.
   • This module is optimized to match latent representations conditioned on it to those under regular SelfControl control.
2. Training Process:
   • Pairs of input and controlled hidden states are generated using SelfControl.
   • The Prefix Controller is trained to minimize the mean squared error between its latent representations and those obtained from SelfControl.
3. Plug-and-Play Control:
   • The Prefix Controller can dynamically be applied to steer the model's behavior for various attributes simultaneously.

Experimental Results

1. Emotion Control:
   • The paper evaluates control over five emotional attributes: anger, fear, happiness, surprise, and disgust.
   • Both SelfControl and SelfControl outperform baselines like Reading Vector and Contrast Vector.
2. Language Detoxification:
   • On the RealToxicityPrompts dataset, SelfControl and SelfControl achieve lower toxicity scores compared to other methods.
   • Both methods are also effective in privacy protection tasks.
3. HH-dialogue:
   • Evaluated on the Anthropic-HH dataset, SelfControl achieves notable win rates against the original model.
   • SelfControl-generated data proves valuable for training improved models via DPO.
4. Reasoning:
   • Experiments on the GSM-8K dataset demonstrate that SelfControl significantly improves mathematical reasoning accuracy.

Analysis and Insights

1. Gradient Trajectory:
   • The trajectory of suffix gradients reveals how combined attributes influence the behavior of LLMs.
2. Control Patterns:
   • Different tasks show varying patterns of gradient application across transformer layers, providing insights into behavioral control mechanisms.
3. Suffix Attention:
   • Attention maps show how the target token in suffixes attends to different input tokens, offering a window into the model’s inner workings.

Implications and Future Work

The introduction of SelfControl and SelfControl marks a significant step towards nuanced control of LLM behaviors without the need for extensive human annotations or parameter modifications. These methods have practical applications in various domains, including ethical AI, emotion modulation, and complex reasoning.

Future research could explore alternative differentiable methods for generating control gradients and refine the mechanistic understanding of these control processes. Additionally, enriching the training data with more diverse behavioral constraints could further enhance the effectiveness of these approaches.

Overall, the techniques introduced in this paper contribute valuable tools for improving the adaptability, transparency, and reliability of LLMs in real-world applications.