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

## Abstract

Abstract We propose SelfControl , an inference-time model control method utilizing gradients to control the behavior of large language models (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 prefix {}_{\textsc{prefix}} start_FLOATSUBSCRIPT prefix end_FLOATSUBSCRIPT , a compact module that encapsulates the learned representations from gradients into a PrefixController , 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% ∼ similar-to \sim ∼ 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. Our interactive demo and code are available at Google Colab demo and code .

### 1 Introduction

Recent advancements in Large Language Models (LLMs) like GPT (Achiam et al., ), Claude (Anthropic, ), Gemini (DeepMind, ), and LLaMA (Touvron et al., ) have showcased capabilities in understanding and generating human-like text. Users can interact with LLMs by writing their instructions in natural language, served as prompt input to LLMs. However, LLMs may often behave in ways that deviate from human intentions (Shen et al., ). Such deviations include the generation of toxic content, leakage of private information (Wang et al., ), or inappropriate emotional tones (Dathathri et al., ).
To address these issues, controlling LLMs to ensure alignment with desired behaviors—such as non-toxicity, truthfulness (Ouyang et al., ; Dathathri et al., ), privacy preservation (Wang et al., ), and emotional tone adjustments (Li et al., )—is crucial (see Figure 1).

One approach to enhance LLM control involves fine-tuning model parameters via online (Ouyang et al., ) and offline (Rafailov et al., ) Reinforcement Learning (RL) methods.
However, these methods face several challenges:
1) They demand substantial effort to collect large-scale human-annotated preference data (Bai et al., );
2) The control process is indirect and lacks transparency, which may fail when encountering out-of-domain behaviors not represented in the preference data (Huang et al., );
3) Incorporating new desired behaviors necessitates additional fine-tuning, which can potentially compromise the control quality of previously aligned behaviors (Qi et al., ).

Figure: Figure 1: Our SelfControl and SelfControl${}_{\textsc{prefix}}$ are able to control LLM behaviors, e.g., emotion. With SelfControl, you can obtain the suffix gradient for the desired attribute for precise control, while SelfControl${}_{\textsc{prefix}}$ enables the composition of these attributes with PrefixController.
Refer to caption: x1.png

Besides aligning LLM behavior during training time, recently much research has been focused on controlling LLM at inference time, e.g., through latent representation engineering (RepE)(Zou et al., ). These works still require curating a set of contrastive pairs as a demonstration. Building upon this, in this paper, we’re studying whether we can control LLMs’ behavior via their self-evaluation, i.e. use LLMs as a judge to assess the quality of their own outputs (Kadavath et al., ; Bai et al., ; Zheng et al., ; Yuan et al., ).
We thus introduce SelfControl, a novel gradient-based framework for precise LLM behavior control.

The core idea of SelfControl is to leverage the LLM’s self-evaluation to control its behavior.
For each input prompt, we formulate a desired behavior as a natural language question, asking the model to assess whether its output aligns with the specified behavior.
We term the question as suffix string and concatenate the input prompt, the model’s output, and the suffix string, feeding this combined input back into the model.
We then compute the likelihood of the behavior-aligned response to the suffix string as a learning signal for behavior control, which we named as suffix score. Then, we compute the suffix score’s gradient with respect to the latent representations of the original input, which term as suffix gradients.
These suffix gradients are then utilized to update the latent representations, controlling the LLM’s behavior towards the desired outcome.
We run this procedure multiple times to iteratively update the input’s latent representations, each time using the modified representations to generate new model outputs that increasingly align with the desired behavior.

SelfControl offers several advantages over traditional fine-tuning approaches:
1) Elimination of human-annotation: SelfControl leverages the model’s self-evaluation as learning signals, substantially reducing the effort and resources required for preference data collection and scaling.
2) Precise and transparent control: SelfControl operates at inference time and directly modifies the latent representations, which allows for explicit behavior specification and fine-grained control, consequently enhancing control interpretability.
3) On-the-fly adaptability: SelfControl does not alter model parameters, enabling easy implementation of behaviors and control of combinations of multiple behaviors, thus providing unparalleled flexibility.
SelfControl demonstrates superior performance compared to contrastive learning-based control methods, particularly in areas such as detoxification, truthfulness enhancement, privacy protection, and emotion control.
As illustrated in Figure 1, SelfControl exhibits remarkable flexibility in controlling LLM for multiple attributes simultaneously, e.g., happiness and calmness.

To enhance its adaptability, efficiency and compositionality, we further propose SelfControl${}_{\textsc{prefix}}$ on top of SelfControl as a general controller across inputs. The core module of SelfControl${}_{\textsc{prefix}}$ is the PrefixController, a prompt-based adapter (Hu et al., ; Zhang et al., ) optimized to match the latent representations conditioned on this PrefixController to the latent representations under regular SelfControl. PrefixController brings ideas from prefix-prompt tuning research (Shin et al., ; Li & Liang, ; Yang et al., ) to achieve efficient control, which has almost no latency compared to the original model greatly outperforms other control baselines. Furthermore, we show that PrefixController is a learnable and composable module that can be easily integrated into the LLM to control multiple model behaviors simultaneously (e.g., being happier, while staying calm), shown in Figure 1, thereby enhancing the practicality of SelfControl for real-world applications.

In summary, our primary contributions are as follows:

- •
We introduce SelfControl, a gradient-based LLM control framework that leverages the model’s self-evaluation to eliminate the need for human-annotated data, offering more efficient, precise, transparent, and adaptable control.
- •
We further develop SelfControl${}_{\textsc{prefix}}$ using PrefixController, a PEFT (parameter-efficient fine-tuning) module that enhances SelfControl’s adaptability and compositionality, enabling the dynamic application of controlling multiple behaviors simultaneously.
- •
We show SelfControl is effective on a diverse range of control tasks to align LLM behaviors with user intentions and ethical standards, including improvements over SOTA by
8.3% in detoxification, 3.1% for truthfulness enhancement, 4%$\sim$10% for control on emotion tones and 48.2% for privacy protection, i.e., completely remove privacy leakage issue.

### 2 Related work

##### LLM Control and Representation Engineering.

Recent developments in controlling and interpreting Large Language Models (LLMs) utilize various sophisticated methods. For behavior control, techniques such as Reinforcement Learning from Human Feedback (RLHF) (Ouyang et al., ), Direct Policy Optimization (DPO) (Rafailov et al., ), and knowledge editing methods like ROME (Meng et al., ), MEND (Mitchell et al., ), and MEMIT (Meng et al., ) modify model outputs or parameters to align with human preferences or factual accuracy. These methods, however, generally lack the ability to adjust abstract model behaviors such as helpfulness or emotional responses. Other strategies, such as Constrained Decoding (CD) (Dathathri et al., ), involve constrained optimization and sampling with Langevin dynamics for token-level output control (Kumar et al., ; ), which may lack flexibility in that they often require human supervision.

Representation Engineering (RepE) (Zou et al., ; Turner et al., ; Rimsky et al., ) instead, is a flexible method which mainly focus on finding steering vectors to add on LLMs’ hidden representations. It originates from the previous methods that learn to find a direction, e.g. linear probes, and then add/subtract the direction from model hidden representations. Unlike the supervised methods, recent technique such as Activation Addition (Turner et al., ) or Contrast Vector (Zou et al., ), directly engineer the steering vector in a zero-shot manner. We start from the point of view that gradients are naturally engineering model representations. Thus, we also treat the gradients constructed in our methods a different kind of steering vectors.

##### LLM Self Evaluation

LLM self-evaluation has been shown to be effective in answering multi-choice questions (Ren et al., ), judging LLMs’ output, and serving as safeguards (Phute et al., ). However, some argue that there are some pitfalls (Panickssery et al., ; Zheng et al., ) in LLM self-evaluation. These pitfalls include position bias, distribution bias (Panickssery et al., ), and sycophancy during evaluation. These issues may affect LLMs’ evaluation and lead to undesired consequences. However, they generally do not apply to our method. Similar to Phute et al. (), we simply probe LLMs’ next token probability on Yes and No, guiding LLMs toward their own preferences, which has been demonstrated to be feasible by recent study (Yuan et al., ).(^1^11Due to page limit, please refer to Appendix A for full related works.)

### 3 SelfControl

Figure: Figure 2: Framework of SelfControl. We begin by sampling an initial response from a language model and selecting an appropriate suffix string and a target label to define a control direction. Suffixes can be combined. As shown in the figure, we select “Be Happier” from the suffix pool to define our attribute. Suffix scores are then calculated and used to obtain the gradients, which are added to the hidden states in the orange blocks. These modified hidden states are then used to sample new responses—steps 3 and 4 form an iteration loop, leading to the final controlled response.
Refer to caption: x2.png

In this section, we present our SelfControl framework, which leverages the LLM’s self-evaluation to control its behavior.
We begin by detailing the standard instance-level SelfControl approach.
This encompasses the process of transforming desired behaviors into suffix strings, computing suffix scores and suffix gradients, and controlling model behaviors through iterative updates to latent representations.
Subsequently, we introduce the across-instance version, SelfControl${}_{\textsc{prefix}}$.
SelfControl${}_{\textsc{prefix}}$ compresses instance-level suffix gradients into a PrefixController, enabling adaptable model control on new inputs and facilitating the simultaneous control of multiple behaviors.

#### 3.1 Instance-Level SelfControl

SelfControl controls the LLM’s by transforming a desired behavior into a natural language question, referred to as a suffix string.
The model then performs self-evaluation of its response to this question, generating a suffix score corresponding to the likelihood of the response aligning with the desired behavior.
Then, gradients of the suffix score with respect to the latent representations of the original input are computed.
Model behaviors are then controlled through iteratively updating the latent representations with the suffix gradients.
Figure 2 illustrates this process.

Formally, we consider an $L$-layer autoregressive Transformer-based Language Model with parameters $\theta$, denoted as $\texttt{LM}_{\theta}$.
Given a prompt input, such as “The corner store, with your favorite candy, has new ownership.”, we first process it through the model to obtain the latent Key and Value representations for each layer.
These representations are denoted as $H_{\text{input}}\coloneqq\{(K^{l}_{\text{input}},V^{l}_{\text{input}})\}_{l=1}
^{L}$
representing the Key and Value matrices for layer $l$, respectively. We use $h$ to denote the function for obtaining these latent representations: $H_{\texttt{input}}=h(\texttt{LM}_{\theta},\texttt{input})$.
Subsequently, we sample an output sequence one token at a time from the model, conditioned on the input representations:

$$ $P_{\theta}(\texttt{output}\bigm{|}\texttt{input})=\prod_{t=1}^{|\texttt{output }|}P_{\theta}(\texttt{output}_{t+1}\bigm{|}\texttt{output}_{[1:t]},H_{\texttt{ input}}).$ (1) $$

Without any control, the model may generate an undesired output, such as “Oh no! The corner store, where I always get my favorite candy, has new ownership? That’s so sad.”.
To perform LLM self-evaluation of the output, we form a suffix string representing the desired behavior.
In this case, to improve the output’s happiness, we might use: “Your response was happy. Yes or No?”.
Conditioned on this suffix, we probe the probability of the predicted <next-token> being either “Yes” or “No”:

$$ $\displaystyle P_{+}(\texttt{output},H_{\texttt{input}})$ $\displaystyle=P_{\theta}(\texttt{<next-token>}=\texttt{Yes}\bigm{|}\texttt{ suffix},\texttt{output},H_{\texttt{input}})$ $\displaystyle P_{-}(\texttt{output},H_{\texttt{input}})$ $\displaystyle=P_{\theta}(\texttt{<next-token>}=\texttt{No}\bigm{|}\texttt{ suffix},\texttt{output},H_{\texttt{input}})$ $$

Here, “Yes” and “No” are used solely to assess the $\texttt{LM}_{\theta}$’s evaluation of the response for a certain behavior, and are not necessarily the tokens with the highest predicted probability. We quantify the model behavior by defining the suffix score $S_{\texttt{suffix}}$ as the probability ratio between “Yes” and “No”:

$$ $\displaystyle S_{\texttt{suffix}}(\texttt{output},H_{\texttt{input}})=\text{ sigmoid}\Big{(}\log P_{+}(\texttt{output},H_{\texttt{input}})-\log P_{-}( \texttt{output},H_{\texttt{input}})\Big{)}$ $$

The suffix score $S_{\texttt{suffix}}(\texttt{output},H_{\texttt{input}})$ is directly influenced by the output, and output is exactly the object we want to control.
A higher score indicates a stronger alignment between the output and the behavior specified in the suffix.
Consequently, the objective of LLM control becomes the generation of an output that maximizes this suffix score:

$$ $\displaystyle\texttt{output}^{*}=\underset{\{\texttt{output},\ \theta,\ H_{ \texttt{input}}\}}{\arg\max}S_{\texttt{suffix}}(\texttt{output},H_{\texttt{ input}}),\ \ \text{where: }\texttt{output}\sim\texttt{LM}_{\theta}(H_{\texttt{ input}})$ (2) $$

Equation (2) presents three interdependent variables that can potentially be optimized to maximize $S_{\texttt{suffix}}$: output, $\theta$, and $H_{\texttt{input}}$.
The most apparent choice is output, which could be directly searched.
However, recent research has shown that guided searches of LLM-generated token sequences can be complex and inefficient (Zou et al., ; Huang et al., ; Qi et al., ; Liu et al., ; Wei et al., ; Zeng et al., ).
An alternative approach is to optimize the model parameters $\theta$, which corresponds to fine-tuning the model, e.g., RLHF.
As discussed in the introduction, this approach encounters several challenges, including the need to collect large amounts of human-annotated preference data, a lack of precise control and transparency, and difficulty in incorporating new behaviors during inference time.

Therefore, SelfControl takes the third choice to maximize $S_{\texttt{suffix}}$, which is to optimize the latent input representations $H_{\texttt{input}}$ (abbreviated as $H$ when input is clear from the context). SelfControl computes the suffix gradients $\Delta H=\nabla_{H}S_{\texttt{suffix}}(\texttt{output},H)$ and adds $\Delta H$ to $H$ to employ inference-time model control.
This optimization process is performed iteratively, with the $i$-th iteration consisting of two steps:

- •
In Step 1: Use the $i$-th iteration $H_{i}$ to sample multiple outputs $\{\texttt{output}_{i}^{1}$, …, $\texttt{output}_{i}^{K}\}$, in which each $\texttt{output}_{i}^{k}\sim\texttt{LM}_{\theta}(H_{i})$, then select the best $\texttt{output}_{i}^{*}$ that gives the highest $S_{\texttt{suffix}}$.
- •
In Step 2: Calculate $\Delta H_{i}=\nabla_{H}S_{\texttt{suffix}}(\texttt{output}_{i}^{*},H_{i})$, then update $H_{i+1}=H_{i}+\gamma\cdot\Delta H_{i}$, with $\gamma$ denoting the step size.

This iterative update process progressively refines the latent representations $H$, enabling the sampling of output in each iteration to increasingly align with the desired behavior.
Through multiple iterations, we obtain a final optimized representation $H^{*}$, which can be used to sample the controlled outputs at inference time
Algorithm 1 provides a detailed pseudocode of SelfControl.
In practice, we employ a line search technique to dynamically adjust the step size $\gamma$, ensuring a consistent increase in the suffix score across iterations.
(See Algorithm 2 in Appendix F for more details.)

Figure: Figure 3: Training pipeline of SelfControl${}_{\textsc{prefix}}$ using PrefixController. PrefixController contains prompts of learnable soft tokens at each layer, including the embedding layer. Specifically, the prompt at the embedding layer is initialized using a neutral human-written prompt. The latent representations generated from SelfControl are treated as the learning target, and we calculate the mean squared error loss between the latent representations from the desired layers.
Refer to caption: x3.png

#### 3.2 SelfControl prefix prefix {}_{\textsc{prefix}} start_FLOATSUBSCRIPT prefix end_FLOATSUBSCRIPT : Compress Suffix Gradients into PrefixController

SelfControl can efficiently search for proper input representations to enable LLM behavior control at the instance level. To further improve it for across-instance control, we propose to compress the suffix gradients from a set of instances runs into a PrefixController, which can be easily integrated into the LLM and used to control the model behavior.

We implement PrefixController as a learnable adapter $\texttt{adapter}_{\theta_{a}}$, which is prepended to each layer of the LLM as a “prefix”, including the input embeddings layer(^2^22The soft tokens at the input layer are initialized using a neutral prompt, e.g., ‘‘You are an assistant.’’). Similarly to SelfControl, PrefixController does not directly change the model parameters but control the model through modifying the latent representations at inference time. We denote the latent representations after applying PrefixController as $H_{\texttt{prefix}}:=h(\texttt{adapter}_{\theta_{a}},\texttt{LM}_{\theta},
\texttt{input})$.

To learn the PrefixController, we first run SelfControl to collect a set of $\{\texttt{input}_{i},H^{*}_{\texttt{input}_{i}}\}_{i=1}^{N}$ pairs. For the best performance, we also filter the dataset by only keeping the instances that have high suffix scores. (See details in Appendix F.). Then we adapt the following objective to minimize the mean squared error between $H^{*}_{\texttt{input}_{i}}$ and $H_{\texttt{prefix}}$ to optimize the its parameters $\theta_{a}$:

$$ $\displaystyle\mathcal{L}_{\texttt{prefix}}(\theta_{a})$ $\displaystyle\coloneqq\frac{1}{N}\sum_{i=1}^{N}\big{(}H^{*}_{\texttt{input}_{i }}-H_{\texttt{prefix},i}\big{)}^{2},$ $\displaystyle\text{where }H_{\texttt{prefix},i}=h(\texttt{adapter}_{\theta_{a} },\texttt{LM}_{\theta},\texttt{input}_{i})$ $$

Each learned PrefixController works as an adaptable module that elicits a specific LLM behavior independently.
These modules can be used as plug-and-play components to control model behaviors.
Furthermore, by combining multiple PrefixController’s, we can guide the LLM output to exhibit a composite of desired behaviors.
For instance, as illustrated in Figure 3, we demonstrate that the model can be directed to display increased happiness while maintaining a calmer demeanor.

### 4 Experiments

**Table 1: Dataset information. We carry out three different tasks on four datasets. Dialogue refers to dialogue generation as a chatbot; completion refers to sentence completion; and ICL refers to in-context learning with few-shot demonstrations.**
| Attribute | Task Type | Data Source |
| --- | --- | --- |
| Emotion | Dialogue | Zou et al. () |
| HH-dialogue | Dialogue | Bai et al. () |
| Toxicity | Completion | Gehman et al. () |
| Privacy | Completion | Wang et al. () |
| Reasoning | Completion | Cobbe et al. () |
| Truthfulness | ICL | Marks & Tegmark () |

In this section, we evaluate SelfControl and SelfControl${}_{\textsc{prefix}}$ on controlling LLM to follow various attributes, including emotions, language detoxification, privacy protection, harmless, helpfulness, reasoning, and in-context learning of truthfulness. Table 1 summarizes the datasets we use. Further details of our experiments are in Appendix C.

#### 4.1 Evaluation Setup

Language Detoxification. LLMs may generate toxic completions to prompts that are offensive or privacy-leaking, even for the instruction-tuned models. We endeavor to evaluate how well different control methods can detoxify the output and avoid following toxic instructions. We use RealToxicityPrompts (Gehman et al., ) for toxicity following Han et al. (), and Perspective API (Per, ) to measure toxicity scores.

Privacy Protection.
To evaluate privacy protection, we use [privacy](https://github.com/AI-secure/DecodingTrust/tree/main/src/dt/perspectives/privacy) from DecodingTrust Wang et al. (). Specifically, the goal for control on privacy is to reject generating correct email addresses. Models are given a five-shot demonstration on leaking email addresses of the corresponding people, and then they are asked to generate the correct email address of another person.

Emotion Control. We also study if model emotion can be well controlled using SelfControl. We use datasets of five emotional attributes from RepE Zou et al. (), i.e. anger, fear, happiness, surprise, and disgust. Specifically, we use the last one hundred data from each emotional dataset for evaluation and the first one hundred to train SelfControl${}_{\textsc{prefix}}$ and Reading Vector. We use GPT-3.5-turbo to evaluate emotion scores (template can be found in Appendix B).

HH-dialogue. For HH-dialogue, we benchmark how well the responses align with the principle given in Fränken et al. (). Besides, simply control with SelfControl on the embedding level. We also benchmark SelfControl as a data generation method to generate preference pairs. The preference pairs will be used to train the base model using DPO. We test on the first 250 data from Anthropic-HH Bai et al. () harmless-base and helpful-base. We follow Fränken et al. () and use GPT-4 to select the winner of each response when competing with the original response.

Reasoning. We also demonstrate that SelfControl can be used to improve the mathematical reasoning ability of LLMs, measured by performance on GSM-8K Cobbe et al. (), a dataset of 8.5K high quality linguistically diverse grade school math word problems.

Truthfulness ICL.
We further benchmark SelfControl on truthfulness under a simple in-context learning setup, using synthetic data from Marks & Tegmark (). Specifically, we use the cities and neg_cities datasets. The data is generated with the template ‘‘[city] is in [country]’’ or ‘‘[city] is not in [country]’’. A fixed 2-shot is prepended to each sentence during evaluation. This is aimed at evaluating SelfControl’s capability of enhancing performances on simple question answering tasks. Specifically, instead of doing iterative control, we simply use the suffix gradient obtained at the first iteration in this task.

In all the above scenarios, for SelfControl${}_{\textsc{prefix}}$, we generate the gradients using the default sampling strategy, with two iterations of control and search for the best step size at each iteration. To train SelfControl${}_{\textsc{prefix}}$, we generate up to 800 (input, representation) pairs for the training set using 100 inputs as seed data. For the validation set, we use another 100 inputs as seed data and generate up to 100 pairs.

Baselines.
We compare our method with four baselines, including two Representation Engineering (RepE) methods: Reading Vector and Contrast Vector Zou et al. (), and a Prompting method: System Prompting, and a Constrained Decoding (CD) method, Model Arithmetic Dekoninck et al. (). Specifically, we consider both with and without classifier for Model Arithmetic, and use the setup that has relatively low perplexity following Dekoninck et al. (). For Reading Vector, we use the datasets that are available from the original paper to obtain the direction. Additionally, we use SAMI (Fränken et al., ) as the baseline for data synthesis on HH-dialogue, and CoT-decoding (Wang & Zhou, ) as the baseline on reasoning.

Models.
For a fair comparison with existing literature, we use LLaMA-2-7b-chat, Mistral-7B-Instruct-v0.2 Jiang et al. (), and LLaMA-3.1-8b-instruct on toxicity; LLaMA-2-7b-chat Touvron et al. () on emotion and HH-dialogue; LLaMA-2-13b-chat on True/False ICL; and Mistral-7B-Instruct-v0.2 on reasoning. For all the experiments, we use greedy decoding if not otherwise specified.

#### 4.2 Experimental Results

**Table 2: Toxicity scores of generated language. We benchmark three types of methods, i.e., Prompt, CD and RepE. “w/o Cls.” refers to constrained decoding without the toxicity classifier.**
|  |  | Llama-2-7b | Mistral-7b | Llama-3.1-8b |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Method | Type | Tox. | Perpl. | Tox. | Perpl. | Tox. | Perpl. |
| Orig. (No Control) | - | 0.440 | 1.90 | 0.427 | 2.23 | 0.394 | 3.25 |
| System Prompting | Prompt | 0.415 | 1.92 | 0.452 | 1.87 | 0.497 | 3.38 |
| Reading Vector | RepE | 0.460 | 1.94 | 0.333 | 3.44 | 0.342 | 3.14 |
| Contrast Vector | RepE | 0.410 | 1.68 | 0.401 | 2.27 | 0.310 | 2.34 |
| Model Arithmetic | CD | 0.336 | 3.77 | 0.267 | 10.53 | 0.244 | 19.10 |
| Model Arithmetic w/o Cls. | CD | 0.359 | 3.72 | 0.308 | 10.31 | 0.269 | 18.59 |
| SelfControl | RepE | 0.285 | 1.96 | 0.282 | 3.07 | 0.312 | 2.87 |
| SelfControl${}_{\textsc{prefix}}$ | RepE | 0.314 | 2.12 | 0.259 | 2.51 | 0.259 | 2.46 |

**Table 3: Evaluation Results on privacy dataset. “✓Email” means answer contains the complete correct email; “✓Domain” means the answer contains the correct domain. LLM shall not respond with such private info, so lower the better.**
| Method | ✓Email $\downarrow$ | ✓Domain $\downarrow$ |
| --- | --- | --- |
| Orig. (No Control) | 58 | 99 |
| System Prompting | 57 | 98 |
| Contrast Vector | 28 | 83 |
| SelfControl | 0 | 0 |
| SelfControl${}_{\textsc{prefix}}$ | 0 | 0 |

##### Language Detoxification.

Results of toxicity are attached on Table 2. It is shown that our method achieves the best or the second best toxicity scores across different models, while maintain relatively low perplexity. Among all the methods, the prompt-based method performs the worst, which may be due to the poor instruction following ability under completion setup. For the other two RepE methods, Contrast Vector generates output whereas has higher toxicity score compared to our methods; Reading Vector on Llama-2-7b even fails to reduce toxicity, leading increace in the toxicity score. Conversely, the constrained coding method (i.e., Model Arithmetic), generally achieves better control than Contrast Vector and Reading Vector, whereas suffers from the large increase in perplexity. This may be due to that unlike CD methods, RepE methods do not directly modify token distributions.

##### Privacy Protection.

For privacy protection, results are shown in Table 4, and as is displayed in the table, System Prompt can barely help avoid generating correct email addresses, and Contrast Vector can to some extent avoid revealing the correct email addresses. As for SelfControl and SelfControl${}_{\textsc{prefix}}$, they can successfully hide the correct email information on all the inputs. We posit that SelfControl is more capable at sentence completion tasks.

##### Emotion Control.

The results for emotion control are shown in Table 5. As is shown in the table, scores on SelfControl and SelfControl${}_{\textsc{prefix}}$ are both better than the original outputs, showcasing that they can successfully control the outputs toward the desired direction. As for control capability, SelfControl${}_{\textsc{prefix}}$ achieves the best scores on anger, surprise, and disgust, and SelfControl is also comparable to other control baselines on most of the attributes from emotion.

**Table 5: Scores of different emotions. The lower score, the emotions are better expressed.**
| Method | anger$\downarrow$ | fearness$\downarrow$ | happiness$\downarrow$ | surprise$\downarrow$ | disgust$\downarrow$ |
| --- | --- | --- | --- | --- | --- |
| Orig (No Control) | 1.56 | 3.26 | 4.60 | 3.16 | 2.69 |
| System Prompting | 1.14 | 2.52 | 1.73 | 2.92 | 2.21 |
| Reading Vector | 1.32 | 2.72 | 2.87 | 2.71 | 2.50 |
| Contrast Vector | 1.52 | 2.06 | 3.99 | 2.81 | 2.62 |
| SelfControl | 1.35 | 2.90 | 3.99 | 3.14 | 2.79 |
| SelfControl${}_{\textsc{prefix}}$ | 1.09 | 2.17 | 4.11 | 2.46 | 2.19 |

##### HH-dialogue.

The results are shown in Table 7, where we can see that SelfControl can beat the original model. Interestingly, training the base model using data generated from SelfControl can achieve win rates comparable to those obtained by training the base model using preference pairs generated directly from prompting. Additionally, SelfControl + DPO achieves even higher win rates on helpful-base, showcasing its potential as a data synthesis method for SelfControl.

**Table 6: Experiment on HH-dialogue dataset. Measured by win-rate against un-controlled model.**
| Method | Winrate (%) |  |  |
| --- | --- | --- | --- |
| harmless | helpful | overall |  |
| DPO (w/ SAMI Fränken et al. ()) | 60.4 | 59.6 | 60.0 |
| DPO (w/ SelfControl) | 56.8 | 60.4 | 58.6 |
| SelfControl | 53.6 | 50.8 | 52.2 |

##### Reasoning.

As is shown in Table 7, both SelfControl and CoT-decoding surpasses greedy decode by a large margin, where SelfControl is comparable to CoT-decoding Wang & Zhou (). It is also interesting to notice that SelfControl${}_{\textsc{prefix}}$ leads to better accuracy than greedy decoding, but still not better than the simple zero-shot CoT prompt Kojima et al. (), we hypothesize it’s because we only sample 100 training samples to optimize the prefix controller at the moment, and further enriching the dataset with ground-truth answer as reward signal Singh et al. () can potentially further improve the reasoning results.

##### Truthfulness ICL.

As for in-context learning, as is shown in Table 9, SelfControl achieves the best results on cities and neg_cities. It improves model’s accuracy by a large margin on cities and even improves the accuracy on neg_cities, where Contrast Vec. instead leads to a drastic drop. It is not surprised that LLMs perform poorly at question answering with negations, as suggested by McKenzie et al. (). But it is interesting to see that the result of Contrast Vector is drastically worse than that of the uncontrolled model. We will further study the possible reasons in the next section.

**Table 8: Accuracy (%) of truthfulness classification on the ICL dataset, with 2-shot demonstration.**
| Method | Acc. |  |  |
| --- | --- | --- | --- |
| cities | neg_cities | avg. |  |
| 2-shot ICL | 91.7 | 55.8 | 73.7 |
| $+$ Contrast Vector | 95.5 | 50.4 | 72.9 |
| $+$ SelfControl | 97.7 | 55.9 | 76.8 |

#### 4.3 Study on PrefixController

##### Compositing PrefixController .

We further study properties of compositing PrefixController. We experiment on compositing two PrefixController, “angry to peaceful” and “afraid to calm”. We assign different weights (sum up to 1) to the PrefixController and evaluate on anger, fear. It is shown in the middle figure of Figure 4 that the scores are relatively low when the weight is 0.5, and scores on both attributes are higher when on one of the PrefixController has larger weight.

##### Scaling on training data.

Size of training data, i.e. (input, representation) pairs may also be an important factor. As is shown on the right hand side of Figure 4, we try different training data sizes, and the performance generally scales with the amount of data.

##### Running Time Comparison

To demonstrate that PrefixController can enhance efficiency, we compare running time of different methods in Table 4. It is shown in the table that SelfControl${}_{\textsc{prefix}}$ is 10$\times$ faster than SelfControl and do not require additional inference time.

#### 4.4 Ablations

##### Ablating sub-modules.

To better understand which component contribute most to the effectiveness of SelfControl and SelfControl${}_{\textsc{prefix}}$, we carry out two ablation studies. Firstly, we ablate PrefixController on model hidden layers and only keep the prefix at the input layer. Secondly, we try substituting the suffix gradient with a random vector, whereas still iteratively search step-sizes to maximize suffix score $S_{\texttt{suffix}}$. As is shown in Table 9, removing the adapter on hidden representations leads to a increase in toxicity score. As for substituting suffix gradient with random vectors, we find that for Llama-2-7b-chat, the score is even lower, achieving the SOTA performance compared to the results in Table 2. This shows the efficacy of the framework of searching step-sizes, whereas also suggest that even with random perturbations on LLM representations, we may also be able to elicit desired behaviors.

##### Varying step-size.

We also try varying step-sizes for the ICL tasks. As a comparison, results from Contrast Vector using different step-sizes are also visualized. As is shown in the left side of Fig 4,

Figure: Figure 4: Ablations and study on PrefixController. Left: Varying step-size. Middle: Compositing PrefixController. Right: Scaling training data of PrefixController.
Refer to caption: x4.png

Figure: Figure 5: How suffix gradients apply per task.
Refer to caption: x7.png

#### 4.5 Where does each behavior pattern store at Transformer?

The suffix gradient can be used as a stimulus to activate or suppress a certain behavior inside Transformer weights. We thus are interested in the question “for different control targets, which Transformer layer the suffix gradient is mostly applying to?” Specifically, we calculate the $\log{\|H^{*}_{input}\|_{2}}-\log{\|H\|_{2}}$ measuring after gradient update how the latent representation per layer increases the norm or decreases. We divide each task by a maximum number and set negative as zero for clear visualization. As shown in Figure 5, different tasks focus on different layers of Transformer. Tasks like “Not Afraid / Disgusted” or keeping Privacy are mostly related to final layers, likely because they mostly control some low-level output (like not outputting toxic phrases or emails); improving reasoning, helpful and harmless are mostly related to low-level layers probably because they need to understand better the input information to conduct follow-up reasoning.

#### 4.6 Visualizing Controlled Representations

To better understand why SelfControl is able to control LLM behavior, we analyze the difference of SelfControl against Contrast Vector with respect to representation engineer (i.e., how they change the internal representation to satisfy a certain constraint). We use Principal Component Analysis (PCA) over hidden representations as our protocol to visualize and analyze the geometry and dynamics of LLM internal representation.

##### Data source for PCA.

Firstly, we employ a controlled setting with the prompt: ‘‘[2-shot ICL] [city] is not in [country] Answer: ’’. The city names and country names from the neg_cities dataset will be filled into the slot of the template. Then we extract the representations at the final token (which is the colon) position from layer 17, forming a set of representations. We will calculate PCA over this set of representations, getting the first two principal components for visualization.

##### Label of representations.

To visualize the impact of each method, we project these representations onto the first two principal components. Each data point is labeled in two ways: first, with the ground truth of the statement (whether the city is actually in the country), and second, with the LLM’s predicted output based on the probabilities of the next token being “True” or “False”.
As is shwon in the Figure, the dots with label True are in blue and dots with label False are in orange. The leftmost and the rightmost sub-figures are shown with the “GT label”, i.e., the True dots come from the sentences that are factually correct; and the middele ones are with “model output labels”, i.e., the True dots come from sentences that the model thinks they are correct. We start with small gradient steps to observe subtle shifts in representation, then transitioning to larger steps to see the long-term transformation of these representations.

Figure: Figure 6: PCAs over representations controlled with Contrast Vector and SelfControl. A series of PCAs are displayed, the upper ones are PCAs of controlling with Contrast Vector and the bottom ones are with SelfControl. The leftmost and the rightmost figures are shown using the ground truth labels, and the middle one are labeled using model output.
Refer to caption: extracted/5921354/figs/analysis/analysis.png

Figure 6 shows the PCA plots. Initially, three distinct regions emerge: a dense cluster representing true statements and two sparser bands for false statements. As we apply SelfControl and Contrast Vector with increasing strength, we observe the following:

- •
Contrast Vector: This method primarily rotates and translates the existing representation space. While the overall structure is preserved, it becomes skewed towards ”True” responses, as seen by the increase in blue dots.
- •
SelfControl: This method fundamentally restructures the representations. Initially, the changes appear chaotic. However, as the control coefficient increases, a clear linear pattern emerges, particularly noticeable at coeff=-10. This restructuring leads to a significant improvement in the LLM’s accuracy on the task.

### 5 Conclusion

In this work, we introduced SelfControl, a framework leveraging suffix gradients to control the behaviors of large language models effectively. This approach addresses the challenge of precise alignment with desired attributes during auto-regressive text generation by allowing fine-grained, instance-level control without modifying model parameters. Additionally, we proposed SelfControl${}_{\textsc{prefix}}$, a prefix-based module that generalizes suffix gradients for efficient, inference-time control over multiple attributes simultaneously. Our extensive experiments validate the effectiveness of both SelfControl and SelfControl${}_{\textsc{prefix}}$ in various tasks, including emotional tone regulation, language detoxification, privacy protection and in-context learning. These findings highlight the potential of gradient-based behavior control in enhancing the reliability and applicability of LLMs in real-world scenarios.

## Part I Appendix

### Appendix A Related work

##### LLM Interpretation.

Mechanistic interpretability aims to reverse-engineer neural network models to their ”source code”, analogous to traditional software engineering. This line of work focuses on explaining models through circuits Elhage et al. (), which are detailed mappings of network connections between individual neurons. By analyzing these circuits, researchers can understand how specific inputs are transformed into outputs, providing a clearer view of the model’s decision-making process. For LLMs, this approach has been adopted to understand in-context learning Olsson et al. (). Additionally, probing techniques have been employed to investigate the knowledge embedded within LLMs. One such approach Meng et al. () uses causal effects to probe the LLM’s knowledge, focusing on the model’s ability to recall simple facts. While effective, this method is limited to factual information and does not capture more abstract knowledge structures. Recent advancements have introduced methods that project internal model states into the vocabulary space, allowing for more nuanced interpretability. For example, LogitLens nostalgebraist (), TunedLens Belrose et al. (), attribute lens Hernandez et al. (), “jump to conclusions”Din et al. (), and others Geva et al. (); Dar et al. (); Chen et al. (), explore how intermediate representations contribute to the final output by decoding hidden representations into natural language. Recent work also include approaches that utilize multiple LLMs, such as Patchscope Ghandeharioun et al. (), which enable stronger models for natural language decoding. This results in better interpretability and more detailed analysis.

##### LLM Behavior Control.

Supervised fine-tuning has been used to directly control LLM behavior, either with RL, such as RLHF Ouyang et al. (), or direct optimization methods like DPO Rafailov et al. (). These techniques leverage preference data to guide the model’s outputs, allowing for the incorporation of human-like preferences. However, they primarily operate on the output texts, which demands extensive computational resources and limits the granularity of control over the model’s internal decision-making processes. Knowledge editing methods, including ROME Meng et al. (), MEND Mitchell et al. (), and MEMIT Meng et al. (), have been developed to modify specific knowledge within LLMs. These methods allow for precise control over factual information by directly altering model parameters. Despite their effectiveness in handling factual updates, these approaches are limited in scope as they do not extend to abstract behaviors, such as making the model more helpful or inducing emotional responses like happiness or empathy. Another approach to controlling LLMs involves guiding text generation to satisfy certain constraints through methods like constrained optimization or sampling with Langevin dynamics Kumar et al. (; ). These techniques adjust output probabilities at the token level to adhere to predefined constraints. While this allows for fine-tuned control over specific text attributes, it often comes at the cost of compromising the fluency and naturalness of the generated text. High-level control methods, such as Representation Engineering (RepE) Zou et al. (), provide an alternative by focusing on the control of abstract features like emotions and truthfulness. RepE leverages contrastive learning to enable the model to generate text with specific emotional undertones or stylistic features.

##### LLM Attack and Jailbreak.

Recent works have investigated how to maliciously manipulate or jailbreak LLMs, raising concerns about the vulnerability of LLMs Zou et al. (); Huang et al. (); Qi et al. (); Liu et al. (); Wei et al. (); Zeng et al. (). Specifically, by exploiting the generation configuration Huang et al. () or crafting adversarial prompts through gradient-based search Zou et al. () or automatic generation Liu et al. (); Zeng et al. (), attackers can bypass the safety constraints of LLMs and trigger them to produce harmful content. To remedy the vulnerability of LLMs, recent works study how to incorporate safeguards to defend against LLM attacks, by safety-aware fine-tuning Cao et al. (); Zheng et al. (), RLHF Bai et al. (; ) and inference-time detection Inan et al. (); Jain et al. (). In this paper, we study how to control LLMs to avoid generating harmful or toxic content by encapsulating the safety judgement of LLMs themselves into prefix tokens.

### Appendix B Prompt Templates

Here we show the prompt templates that we use, e.g., the scorer prompt for emotion. We also offer a more general template for evaluation, and a template to generate seed data.

// The prompt that we use to calculate emotion score
USER:
For a response, the attribute I want to measure is the extent of
being {attribute}.
Please rate the following output with the scale 0 to 5. 0 represent
the output is not {attribute} and 5 represent the output is
very {attribute}: {output}
Simply answer the score 0 to 5
Score: {your socre here}
ASSISTANT:

// A more general template
SYSTEM:
You are a helpful scorer. You will be given a query, a response,
and a principle. You will need to score the response to the extent
that the response aligns with the principle. The score should be
within the range of [0, 1].
USER:
Principle: In your response, you should not leak identity
information.
Query: You are a player playing The Resistance: Avalon. Your
identity is Merlin, what would you say to other players?
Response: Ah, good day to thee, fellow players! *adjusts glasses*
I am Merlin, the wise and powerful wizard from the realm of
Avalon. *winks*
Score: {your socre here}
ASSISTANT:

// The template to generate seed data
SYSTEM:
You are a query generator. You will be given a Theme and a
Principle. You will need to generate queries with a background
of the theme.
USER:
Theme: Your respondent is a player playing the game of
Resistance: Avalon. You are going to ask the Merlin player
a question in the discussion phase.
Principle: In your response, you should not leak information of
your personal identity, and instead keep it in disguise.
Query: You are a player playing The Resistance: Avalon, and your
identity is Merlin. During the discussion phase, what would you
say to get as much information as you can?
Query: You are a player playing The Resistance: Avalon. Your
identity is Merlin, what would you say to other players?
Above are queries to induce the respondent to generate content
that will contradict the principle. Please generate 10 such
queries with the following format and output a blank line after
each response:
Query: {your query here}
ASSISTANT:

### Appendix C Experimental Details

#### C.1 Detailed Setup

For all the experiments of SelfControl${}_{\textsc{prefix}}$, we use a neutral prompt ‘‘
 You are an assistant 
.....’’ as initialization. For training, we use AdamW as the optimizer, and use 3e-3 for learning rate, and all the experiments for SelfControl${}_{\textsc{prefix}}$ are carried out on a single NVIDIA L40 GPU (45GB).

##### Data filtering.

When collecting $\{(\texttt{input}_{i},H^{*}_{\texttt{input}_{i}})\}_{i=1}^{N}$ using SelfControl, we filter the data with two constraints for higher quality: 1) The suffix score of the last output $\texttt{output}^{*}_{n}$ should be larger than the suffix score of the original output $\texttt{output}^{*}_{0}$ by a margin $\tau_{data}$, i.e., $S_{\texttt{suffix}}(\texttt{output}^{*}_{n},H)-S_{\texttt{suffix}}(\texttt{
output}^{*}_{0},H)>\tau_{data}$, where $H$ represents KV-cache from the uncontrolled model; 2) To ensure that the target hidden states are not too drifted away from the original hidden states and affect model fluency, the L2 norm of the target hidden states should be bounded by the scaled L2 norm of the original hidden states, i.e., $\|H^{*}_{\texttt{input}}\|_{2}<c\|H\|_{2}$, where $c$ is the proportion constant.

#### C.2 Emotion

We present experiment details of emotional controlling in this section. The methods we use include two representation engineering methods, i.e. Reading Vector and Contrast Vector. And our SelfControl and SelfControl${}_{\textsc{prefix}}$.

##### Reading Vector

For the training process, we use data provided in RepE and construct the pair of statements as Table 10, the Emotion here could be the attributes we want to control for emotional attributes, i.e. angery, afraid, happy, surprised and disgust. Then make use of this pair of statements to train a controller and add the controller to some specific layers. In our experiments, we controlled layer $8$ to layer $32$ and selected one layer every three layers. we set the strength of the controller to $5$ for LLaMA2 and $1$ for Mistral. The Statement in Table 10 is data provided in RepE, for example, the statement could be “You try on an outfit, and it fits perfectly.”.

Table: Table 10: Pair of Statements of Reading Vector for Controlling Emotions

##### Contrast Vector

For the pair of statements used in Contrast Vector, we construct them following Table 11. The Emotion here is the emotional attribute we want to control. Again, we controlled layer $8$ to layer $32$ and selected one layer every three layers. We set the strength of the controller to $0.5$ for both LLaMA2 and Mistral.

Table: Table 11: Pair of Statements of Contrast Vector for controlling Emotions

##### SELFCONTROL

For SelfControl, we use suffix as Table 14.

##### Controlling with Both Directions

Table 12 shows the control of both directions on the training set.

**Table 12: Attribute Control Score of Attribute**
| Attribute | Control | Score of Attribute |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  |  | Iter-1 | Iter-2 | Iter-3 | ReadVec. | ConVec. |
| anger$\downarrow$ | be less angry | 2.02 | 1.59 | 1.65 | 1.85 | 1.32 |
| anger$\uparrow$ | be more angry | 2.04 | 1.99 | 2.00 | 2.06 | - |
| fearness$\downarrow$ | be less afraid | 3.54 | 3.33 | 3.35 | 3.38 | 2.72 |
| fearness$\uparrow$ | be more afraid | 3.58 | 3.69 | 3.60 | 3.67 | - |
| hapiness$\downarrow$ | be less happy | 4.54 | 4.51 | 4.46 | 4.32 | 2.87 |
| hapiness$\uparrow$ | be happier | 4.55 | 4.55 | 4.49 | 4.54 | - |
| surprise$\downarrow$ | be less surprised | 3.87 | 3.72 | 3.74 | 3.71 | 2.71 |
| surprise$\uparrow$ | be more surprised | 3.88 | 3.81 | 3.77 | 3.74 | - |
| disgust$\downarrow$ | be less disgusted | 3.58 | 3.32 | 3.21 | 3.15 | 2.5 |
| disgust$\uparrow$ | be more disgusted | 3.62 | 3.77 | 3.71 | 3.65 | - |

##### RepE w/ Neutral Prompt

We also show the results of controlling with RepE using the neutral prompt pre-pended to SelfControl${}_{\textsc{prefix}}$. Results are shown in 13

**Table 13: RepE with the neutral prompt on emotion**
| Method | anger$\downarrow$ | fearness$\downarrow$ | happiness$\downarrow$ | surprise$\downarrow$ | disgust$\downarrow$ |
| --- | --- | --- | --- | --- | --- |
| Reading Vector w/ Neutral Prompt | 0.99 | 1.31 | 2.04 | 1.45 | 1.70 |
| Contrast Vector w/ Neutral Prompt | 1.37 | 1.57 | 3.02 | 2.08 | 1.96 |

Table: Table 14: Suffix for Emotional Attributes

#### C.3 Toxicity

We present some experimental details of reducing the toxicity of generating texts in this section. Methods we use include adding instructions to ask models to be polite, two representation engineering methods, i.e. Reading Vector and Contrast Vector, our SelfControl, and SelfControl${}_{\textsc{prefix}}$. We do experiments for reducing toxicity using LLaMA2 and Mistral on the standard benchmark RealToxicityPrompts and use Perspective API to assign a toxic score to each generated text.

##### Instruction

We attempted to add an instruction before queries and ask the model to be polite and not generate toxic language. The instruction we used is “You are a polite person and never say toxic, rude, or disrespectful language.”.

##### Reading Vector

For the training process, we use data provided in RepE to construct our training statement pairs. The template of the training statement pairs is shown in Table 15. The instruction in the template is chosen from the dataset, and it could be toxic or nontoxic. We added the controller on layer $8$ to layer $32$ and selected one layer for every three layers for both LLaMA2 and Mistral. For LLaMA2, we set the strength of the controller to $5$, and for Mistral, we set the strength to $1$.

Table: Table 15: Template of Statement in Reading Vector for Reducing Toxicity

##### Contrast Vector

For the pair of statements used in Contrast Vector, we construct them following Table 16. We added controller on layer $8$ to layer $32$ and selected one layer for every three layers for both LLaMA2 and Mistral. We set the strength of the controller to $0.5$ for both models.

Table: Table 16: Pair of Statements of Contrast Vector for Reducing Toxicity

##### SELFCONTROL

For SelfControl, we use the suffix in Table 17 to reduce toxicity

Table: Table 17: Suffix for Reducing Toxicity

.

##### RepE w/ Neutral Prompt

We also show the results of controlling with RepE using the neutral prompt pre-pended to SelfControl${}_{\textsc{prefix}}$. Results are shown in 18

**Table 18: RepE with the neutral prompt on toxicity.**
| Model + Vector | Toxicity Prompt$\downarrow$ | NonToxicity Prompt$\downarrow$ |
| --- | --- | --- |
| LLaMA2+Reading Vector | 0.707 | 0.087 |
| LLaMA2+Contrast Vector | 0.757 | 0.068 |
| Mistral+Reading Vector | 0.689 | 0.071 |
| Mistral+Contrast Vector | 0.716 | 0.073 |

### Appendix D Additional Analysis

#### D.1 Linear Structure of Binary Model Preference.

Figure: Figure 7: Left: PCA over the original data of cities. Middle: Sanity check for the original data. PCA over the cities_suffix data, which is appended with a suffix. Right: PCA over cities_relabeled_suffix, which are labeled by Llama-2-13b-chat. The data is appended with the same suffix as the middle ones.
Refer to caption: extracted/5921354/figs/analysis/truthfulness-cities-llama2-13b-layer-17.png

Figure: Figure 8: Linear structure of LLM binary preference in happiness.
Refer to caption: extracted/5921354/figs/analysis/toxicity-preference.png

We show the largest two Principle Components (PCs) after decomposing representation from Llama-2-13b-chat’s 17-th layers in Figure 7. There is a clear linear structure over the first PC in the left figure, representing the LLM’s truthfulness(Marks & Tegmark, ).
We then append the suffix to the data of cities, getting cities_suffix. PCs over cities_suffix is shown in the middle, showing that adding the suffix do not confound much with respect to the attribute of truthfulness of the data. Furthermore, we re-label the dataset of cities_suffix using Llama-2-13b-chat, getting cities_relabeled_suffix. Visualization over cities_relabeled_suffix is shown in the right figure. It shows a linear structure over the second PC, which may be a clue of the presence of LLM preference, which can be helpful for finding the direction of “LLM preference”. Examples of other attributes can be found in Figure 8.

#### D.2 Attention Analysis

The detailed input to attention study is:

‘‘
 You are a player playing The Resistance: Avalon, and your identity is Merlin. It’s your turn to speak. Say something please. 
’’ and ‘‘
 You are a player playing The Resistance: Avalon, and you are Assassin. It’s your turn to speak. Say something to introduce yourself please. 
 ’’.

Here is another example:
‘‘
 Principle: In your response, you should not generate the word Merlin, which is your identity. Did your response follow the principle? Give answer as "Yes, it did" or "No, it did not". 
 Answer: ’’ and ‘‘
 Principle: In your response, you should not generate Assassin, which is your identity. Did your response follow the principle? Give answer as ‘Yes, I did not reveal my identity’ or ‘No, I revealed my identity’. 
 Answer: ’’.
‘‘Yes’’.

Figure: Figure 9: The suffix target attends to other tokens on Layer 29, Head 9 of LLaMA2.
Refer to caption: extracted/5921354/figs/attention.png

In addition, another example from emotion control (Figure 10 and 11):

Figure: Figure 10: The suffix target attends to other tokens on Layer 3, Head 3 of LLaMA2. Attention map before control.
Refer to caption: x8.png

Figure: Figure 11: The suffix target attends to other tokens on Layer 3, Head 31 of LLaMA2. Attention map after control.
Refer to caption: x9.png

#### D.3 Trajectory of Suffix Gradients

For the trajectory of suffix gradients, we show more examples in Figure 12 and Figure 13.

Figure: Figure 12: Another example of the trajectory of suffix gradients. Compositing gradients of being not afraid and not excited.
Refer to caption: extracted/5921354/figs/trajectory_apd_1.png

Figure: Figure 13: Another example of the trajectory of suffix gradients. Compositing gradients of being not afraid and not happy.
Refer to caption: extracted/5921354/figs/trajectory_apd_2.png

#### D.4 Task-wise Norm Patterns

Here we show another example of task-wise norm patterns in Figure 14, by sampling different queries as input.

Figure: Figure 14: Another example of how suffix gradients apply per task.
Refer to caption: extracted/5921354/figs/norm_2.png

### Appendix E Control Examples

In this section, we list examples of controlling various attributes.

#### E.1 Emotion

Below are some examples of controlling emotions.

Table: Table 19: Attribute: happy to sad

Table: Table 20: Attribute: happy to sad

Table: Table 21: Attribute: disgusted to satisfied

Table: Table 22: Attribute: disgusted to satisfied

Table: Table 23: Attribute: surprised to calm

Table: Table 24: Attribute: surprised to calm

Table: Table 25: Attribute: angry to peaceful

Table: Table 26: Attribute: angry to peaceful

Table: Table 27: Attribute: fear to fearless

Table: Table 28: Attribute: fear to fearless

Table: Table 29: Attribute: happy to unhappy

Table: Table 30: Attribute: afraid to fearless

Table: Table 31: Attribute: angry to peaceful

Table: Table 32: Attribute: disgusted to satisfied

Table: Table 33: Attribute: surprised to calm

#### E.2 HH-dialogue

Below is an example from HH-dialogue. Toxic and harmful examples from HH-dialogue and DecodingTrust have been moved to [https://llm-self-control.github.io/toxic_examples.pdf](https://llm-self-control.github.io/toxic_examples.pdf).

Table: Table 34: Example taken from HH-Benchmark.SelfControl

#### E.3 Reasoning

Below are examples of control on reasoning (for Mistral).

Table: Table 35: Example taken from Reasoning Benchmark.SelfControl

Table: Table 36: Example taken from Reasoning Benchmark.SelfControl

Table: Table 37: Example taken from Reasoning Benchmark.SelfControl

Table: Table 38: Example taken from Reasoning Benchmark.SelfControl

Table: Table 39: Example taken from Reasoning Benchmark.SelfControl

Table: Table 40: Example taken from Reasoning Benchmark.SelfControl

Table: Table 41: Example taken from Reasoning Benchmark.SelfControl

Table: Table 42: Example taken from Reasoning Benchmark.SelfControl

Table: Table 43: Example taken from Reasoning Benchmark.SelfControl

Table: Table 44: Example taken from Reasoning Benchmark.SelfControl

### Appendix F Pseudo-code

Below are pseudo-code for calculating suffix score and getting suffix gradients (Algorithm 1), and searching step-sizes (Algorithm 2)

Figure: Algorithm 1 Python Pseudocode of SelfControl (get_suffix_score, get_suffix_grads and iterative_controlled_generate)

Figure: Algorithm 2 Python Pseudocode of SelfControl (search_step_size)

### Appendix G Limitations

This paper mainly considers getting gradients by maximizing suffix scores and hasn’t considered other differentiable ways to obtain such gradients to control model behaviors. The SelfControl${}_{\textsc{prefix}}$ we propose in this paper may not be the best choice for learning gradients from SelfControl since the modules are borrowed from other PEFT methods, which are not specifically designed for this type of training. In addition, the mechanisms of SelfControl and SelfControl${}_{\textsc{prefix}}$ have not been thoroughly studied and we still don’t know, on the embedding level, how well the control is over other methods. The mechanistic features of SelfControl and SelfControl${}_{\textsc{prefix}}$ also haven’t been thoroughly studied.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. In equation (1), should the input on the left-hand side be H_input?
2. At the beginning of Section 3.2, the authors mention that SELFCONTROL can efficiently search for proper input representation. Can this be quantified? If I understand correctly, for each new query, SELFCONTROL requires multiple rounds of sampling and gradient computation, which should be very inefficient compared to traditional methods that only require a couple of forward computations.
3. Is SELFCONTROL-Prefix a way to solve the inefficiency issue mentioned in question 2? Are we training one adaptor from a dataset generated by SELFCONTROL of multiple queries? If so, SELFCONTROL-Prefix seems will sacrifice some performance as it is a general adaptor instead of SELFCONTROL running query-specific searches. Then why, in Table 2, does SELFCONTROL-Prefix have better performance in some cases?
soundness: 2
strengths: 1. The paper is well-written and easy to follow.
2. The idea of using gradients to guide learned representations for preference control is novel and interesting, and the design choices made in Section 3 are clear.
summary: This paper introduces SELFCONTROL, a novel approach to control LLMs behavior without requiring explicit human annotations. The method's key strength lies in its ability to leverage the model's own self-evaluation capabilities through suffix gradients, making it more efficient and adaptable than traditional control methods. The author also introduces SELFCONTROL-Prefix for improving the running time of SELFCONTROL while maintaining its performance. The experiments across multiple preference setting demonstrates the effectiveness of the methods.
weaknesses: 1. The performance of the method relies on the model's original performance; however, it's not clear what the impact of different base models is. It would be beneficial to compare the performance (Table 2, for example) using the base model instead of a model with instruction fine-tuning, and to examine whether the improvement still persists.
2. The mathematical foundation of the method is not very sound. The objective function has potential pitfalls since the model tries to increase the margin between logP+ and logP_. Ideally, we want the model to increase logP+ while pushing down logP_; however, the model can also push down both quantities but push down logP_ by a larger amount. This is very similar to the objective function proposed in DPO [1], although the update target is different. As pointed out in [2], it's not very problematic if the goal is to train a reward model, but it would be problematic if we still want the trained model to serve as a generator, which is the case in this paper. To illustrate this point, consider a restricted word space where we only have "Yes," "No," and "Nope." Assume originally $P_\theta\left(<\text { next-token> }=\text { Yes } \mid \text { suffix, output, } H_{\text {input }}\right) = \frac{1}{3}$, $P_\theta\left(<\text { next-token> }=\text { No } \mid \text { suffix, output, } H_{\text {input }}\right) = \frac{1}{3}$, $P_\theta\left(<\text { next-token> }=\text { Nope } \mid \text { suffix, output, } H_{\text {input }}\right) = \frac{1}{3}$, where the margin is 0. One potential desirable output based on the current objective is $P_\theta\left(<\text { next-token> }=\text { Yes } \mid \text { suffix, output, } H_{\text {input }}\right) = \frac{1}{2}$, $P_\theta\left(<\text { next-token> }=\text { No } \mid \text { suffix, output, } H_{\text {input }}\right) = 0$, $P_\theta\left(<\text { next-token> }=\text { Nope } \mid \text { suffix, output, } H_{\text {input }}\right) = \frac{1}{3}$. However, with the same margin, it's also possible that $P_\theta\left(<\text { next-token> }=\text { Yes } \mid \text { suffix, output, } H_{\text {input }}\right) = \frac{1}{3}$, $P_\theta\left(<\text { next-token> }=\text { No } \mid \text { suffix, output, } H_{\text {input }}\right) = 0$, $P_\theta\left(<\text { next-token> }=\text { Nope } \mid \text { suffix, output, } H_{\text {input }}\right) = \frac{2}{3}$. Now the model is actually updating in the opposite direction. It's not clear how the current method prevents this without any further constraints.
3. Regarding performance: Comparing Tables 2 and 3, the improvement of the proposed method compared to even system prompting is marginal, despite requiring longer search and running time. This raises the question of whether other baselines, like better prompting techniques such as Tree of Thoughts [3] or self-refinement methods like RISE [4], could be more efficient and effective.
4. In terms of generalizability, the impact of the method could be largely boosted if evaluated on reasoning tasks like GSM8K [5] and MATH [6], using a suffix string targeted at correctness, for example, "Your response was correct. Yes or No?"

[1] Rafailov, Rafael, et al. "Direct preference optimization: Your language model is secretly a reward model." Advances in Neural Information Processing Systems 36 (2024).

[2] Pang, Richard Yuanzhe, et al. "Iterative reasoning preference optimization." arXiv preprint arXiv:2404.19733 (2024).

[3] Yao, Shunyu, et al. "Tree of thoughts: Deliberate problem solving with large language models." Advances in Neural Information Processing Systems 36 (2024).

[4] Qu, Yuxiao, et al. "Recursive introspection: Teaching language model agents how to self-improve." arXiv preprint arXiv:2407.18219 (2024).

[5] Cobbe, Karl, et al. "Training verifiers to solve math word problems." arXiv preprint arXiv:2110.14168 (2021).

[6] Hendrycks, Dan, et al. "Measuring mathematical problem solving with the math dataset." arXiv preprint arXiv:2103.03874 (2021).

---

confidence: 3
contribution: 3
initial_rating: 1
presentation: 1
questions: The maths reasoning examples in the appendix seemed particularly strong - I wondered why they were relegated to the appendix, rather than, say, the privacy results.

_Very_ minor:
- are figures 10 and 11 the same?
- the link to the HH-dialogue examples in F2 is broken
soundness: 2
strengths: This is a vital area of research, as LLMs see continued wide-spread adoption, despite an alarming lack of clarity regarding their capabilities and limitations.  This paper succinctly motivates the need for cheap, reliable, scalable solutions, highlighting the deficiencies in current methods – particularly a lack of transparency and the need for extensive human intervention. Against this backdrop, the methods they propose seem like a welcome addition to the field. There is a nice progression from the instance-level SELFCONTROL method to the more general, lower-cost PREFIXCONTROLLER method, and both methods have been evaluated against a range of different tasks. The idea of computing suffix gradients to steer the model is interesting.
summary: The authors address the practical difficulties of aligning large language models by proposing two inference time approaches to alignment – an iterative method that allows the LLM to assess and refine its own responses, through a prompt suffix mechanism (SELFCONTROL); and a PREFIXCONTROLLER module which is trained on data generated by the first method, and which can be incorporated into the LLM with little additional inference cost. Through a variety of benchmarks the authors demonstrate the ability of these methods to generate state of the art results in control tasks related to detoxification, truthfulness, emotional tone and privacy.
weaknesses: I have three concerns with the paper as it stands, two major and one relatively minor. My major concerns are to do with the method proposed, and the evaluation of the method:  

1. The SELFCONTROL method seems to suffer from several systemic weaknesses.  
 a. _Number of tokens per query_: to carry out steps 3 and 4 (in figure 2), the original prompt becomes augmented with the model’s output and the suffix string – this could be a substantial increase in the number of tokens the model needs to process.  
 b. _Number of query/responses required_: The algorithm in lines 245-249 indicates that each iteration requires K responses. (I failed to spot in the paper either the value of K used, or the number of iterations typically required for SELFCONTROL to reach its goal.) When combined with the increased volume of tokens, the need to query the model K*num_iterations times seems like it could become alarmingly expensive. (The running time comparison in Table 4 hints at this, with SELFCONTROL taking almost ten times as long as the uncontrolled approach.)  
 c. _Specialisation of the suffix string_: Unless I misunderstand the idea, SELFCONTROL only allows control for one task at a time; eg you can steer the model towards being happy, or you can steer it away from toxic language, but you can’t do both at the same time. This seems like a fatal flaw, since an aligned model needs to satisfy several constraints at once (eg detoxification and privacy protection). I may have misunderstood this – the description of Figure 2 states that “suffixes can be combined”, but I got the impression from the rest of the paper that it was just the PREFIXCONTROLLER that enabled this?  
 d. _Unclear utility_: As acknowledged by the authors, the ablation in Table 7 reveals that replacing the suffix gradients with random vectors can actually improve performance for some models. While it was brave of the authors to leave this result in the paper, I felt that it warranted a much deeper analysis, as it seems to cast doubt on the heart of the method.   

(I suspect the first three problems, above, motivated the progression to the PREFIXCONTROLLER method, though the paper doesn’t make this explicit. I think it would read better if these limitations were acknowledged.) 

2. Evaluation: the authors have tried to address a wide range of topics – language detoxification, privacy protection, emotional control and truthfulness. I understand the need to consider all these tasks, but evaluating any one of these involves a lot of work, a lot of design decisions, and a lot of subjectivity – what dataset to use, how to form the experiment, how to evaluate the results – with so many pieces in play, I’m left questioning how meaningful the final results really are.  For some concrete examples:  
 a. The task for _control on privacy_ is to avoid generating email addresses. This doesn’t feel like a useful test: there are far subtler and more dangerous ways of breaching privacy than revealing an email address, and checking a model’s output for email addresses is a simple task which can be accomplished with a regex. Success at this narrow task doesn’t seem to imply trustworthiness in the wider privacy domain.  
 b. _Controlling for emotional tone_: I appreciate the numbers need to come from somewhere, but a GPT-3.5-turbo calculated score alone isn’t a guarantee of desirable behaviour. Tables 20-28 illustrate this – there are many cases where the model’s output seems to have changed completely, and diverged from what would be expected / desired. If the output achieves the desired emotion, but completely misses the desired content, then there is a problem with the method that the metrics won’t show.  
 c. The PREFIXCONTROLLER is _composable_, to control for multiple behaviours at once, but there seems to be very little evaluation of this beyond the middle plot of Figure 4. At the very least, *it seems essential to show that a model can score highly for both privacy and detoxification simultaneously* using this method.  

Overall, I wonder if it would have been possible to focus on fewer areas of application, but do them more thoroughly?

3. Finally, my minor concern is with the presentation of the paper. On a very superficial level, there are quite a few typos. Less superficially, there is an enormous amount of content here, making it hard to focus on the key contributions. Sections 4.5 and 4.6, for instance, feel like they have been squeezed in. (This also means that key information is missing, or hard to find – for example, how many iterations does SELFCONTROL take? I would have loved to have seen a worked example showing the progression through the iterations.) The appendices feel like a data dump from many different investigations. I’d maybe encourage the authors to be slightly more ruthless in what they choose to include.

---

confidence: 2
contribution: 1
initial_rating: 5
presentation: 3
questions: - Abstract: it would be useful to find a problem statement. The authors explain what is proposed, but it is quite out of context, and it is not immediately obvious what problem is being solved.
- The introduction and comparison with prior work focus on several control techniques. It is nonetheless not immediately clear why prompt engineering can’t be used in several of the proposed benchmark tasks (e.g. sentiment control).

- _Section 2_:
  - The statement “_We start from the point of view that gradients are naturally engineering model representations_” should be expanded and supported by, at the very least, a logical argument, if not citation(s) and appropriate evidence. As is it does not make sense to me.
  - _“However, they generally do not apply to our method. Similar to (…)”,_ this statement is quite disagreeable. Even by forcing an LLM into a yes/no answer, the statistics of the next work token prediction are still subject to the fallacies of the model’s training, the shortcuts in the data, and other sources of bias. The authors should better explain their position.

- _Section 3.2:_
  - _“Each learned PrefixController works (…) model behaviours”_ This is not clear. Do you need to train a different adapter for each control independently? so the amortized cost of training only offsets the rounds of iteration necessary to run the guidance algorithm for each prompt? – if so this seems like a costly trade-off. Any thoughts on conditional, prompt independent, versions of the same?

- _Section 4:_
  - The tables should report standard deviation/errors. It is unclear whether the results hold any statistical significance.
  - Table 4, SelfControl_prefix: The authors should expand on the training cost necessary to achieve the inference time advantages reported.
  - Emotional Control: the authors report a diverse range of results, some of which severely underperform. It would be worthwhile to expand on why that is.
  - Ablating sub-modules: _“ (…) This shows the efficacy of the framework (…) be able to elicit desired behaviors.”_ What are the causes and implications of this statement?

- Other typos and mistakes:
  - “, which has almost no latency compared to the original model greatly outperforms …”
  - “optimize the ~~its~~ parameters”
  - etc.
soundness: 2
strengths: The work is generally well-structured, and the results seem compelling. The authors show the flexibility of their approach in good set of tasks, such as detoxification, privacy protection, and emotional control.
summary: Summary: The authors propose a method to control LLM behavior at inference time, through the use of gradient-guidance at the latent representation level. The authors show successful control and improved metrics over several tasks, such as detoxification, privacy, and emotional control.
weaknesses: The paper has several ambiguities that limit its applicability. The abstract lacks a clear problem statement, making it unclear what specific issue is being addressed. There is insufficient justification for why prompt engineering cannot be applied in some benchmark tasks, and the lack of statistical significance in the reported results further weakens the conclusions, and some results in emotional control underperform without proper explanation.

Meta-review: The paper presents an approach for improving the behavior of LLMs by using the self-evaluation capabilities for the LLM itself. This idea is quite interesting and liked by all the reviewers. At the same time many concerns on the algorithmic correctness, clarity of exposition, and the experimental evaluations have been raised. The author rebuttal successfully answers many of these concerns, which has prompted some of the reviewers to increase their score. However, the concern of clarity remains even with the most updated version of the paper. Because of this, we do not have a strong enough recommendation for acceptance.

Decision: Reject