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

Min Cai1 &Yuchen Zhang2 &Shichang Zhang3 &Fan Yin3
Difan Zou4 &Yisong Yue5 &Ziniu Hu5
1Shenzhen University, 2Rensselaer Polytechnic Institute
3University of California, Los Angeles, 4The University of Hong Kong
5California Institute of Technology
https://llm-self-control.github.io/

###### Abstract

We propose SelfControl, a novel method utilizing suffix gradients to control
the behavior of large language models (LLMs) without explicit human
annotations. Given a guideline expressed in suffix string and the model’s
self-assessment of adherence, SelfControl computes the gradient of this self-
judgment concerning the model’s hidden states, directly influencing the auto-
regressive generation process towards desired behaviors. To enhance
efficiency, we introduce SelfControl${}_{\textsc{prefix}}$, a compact module
that encapsulates the learned representations from suffix gradients into a
Prefix Controller, facilitating inference-time control for various LLM
behaviors. Our experiments demonstrate SelfControl’s efficacy across multiple
domains, including emotional modulation, ensuring harmlessness, and enhancing
complex reasoning. Especially, SelfControl${}_{\textsc{prefix}}$ enables a
plug-and-play control and jointly controls multiple attributes, improving
model outputs without altering model parameters or increasing inference-time
costs. An Google Colab demo is available at: demo and our code at code.

Figure 1: Our SelfControl and SelfControl${}_{\textsc{prefix}}$ are able to
control LLM behaviors such as emotion (fearlessness), helpfulness, and
reasoning capability. On the right hand side, we show that these different
attributes can be composed into a single prefix controller.

###  1 Introduction

Large Language Models (LLMs) like GPT-4 [2], Claude [3], Gemini [11], and
Llama [44] have emerged as powerful tools for understanding and generating
human-like text. With such LLMs as foundation models, humans can write high-
level instructions expressed through prompts to control model behaviors at
inference time. However, there are many situations where LLMs may behave
differently than the intended control. One fundamental hurdle to fine-grained
model behavior control is the autoregressive nature of LLMs, where texts are
generated token-by-token based on the preceding context. Unlike diffusion
models [21], which allow for iterative adjustments during the generation
process via classifier guidance [12], LLM outputs can only be evaluated after
they are fully formed, which hinders the precise alignment of LLM outputs with
desired attributes or constraints expressed in the user prompt. This lack of
accurate and nuanced control becomes particularly problematic in sensitive and
high-stakes areas where even minor deviations from factual accuracy or ethical
standards can have significant consequences such as medical diagnosis [43]

Existing alignment approaches, such as online Reinforcement Learning from
Human Feedback (RLHF) [38] or offline Direct Preference Optimization (DPO)
[40], attempt to address this challenge by fine-tuning the model parameters to
follow desired behaviors expressed in the preference data. However, these
methods rely on extensive efforts to define desired behaviors and annotate
preference data at scale [4], while they are often opaque regarding how the
behaviors are enforced. Such a need for human annotation and the lack of
transparency in the learned behaviors limit the scalability and adaptability
of these approaches. Especially for the out-of-domain behaviors beyond the
collected preference data, LLMs will likely fail to follow.

To bypass the reliance on large-scale preference data as in alignment
approaches, we introduce SelfControl, a novel gradient-based framework
designed to enable differentiable control of LLM outputs without human
annotation. Our approach is inspired by recent studies showing that LLMs have
a decent ability to judge the quality of their outputs [27, 5, 51, 48].
SelfControl takes a desired attribute, expressed as suffix strings, to
evaluate whether a generated output follows the attribute. Then, it calculates
the gradient of the likelihood of the output following the attribute with
respect to the latent representations of the inputs. Then latent
representations are updated with the gradient to control LLM behaviors. We
call this procedure suffix gradient search. Compared to existing control
methods, our procedure has two major advantages: 1) it allows us to directly
influence the generation trajectory only by models’ self-judgment without
relying on extensive human annotations; 2) only latent representations are
updated but not model parameters, so inference-time control can be achieved
for different control objectives. We show in Figure 1 that SelfControl can
generally control LLMs’ emotions, helplessness, etc. In addition, we also show
such control can improve models’ reasoning capabilities for solving
mathematical problems.

SelfControl operates at the instance level, i.e., it controls the model
behavior for a single LLM input. To enhance its transferability and
compositionally, we further propose SelfControl${}_{\textsc{prefix}}$ on top
of SelfControl as a general controller for multiple instances at the same
time. The core of SelfControl${}_{\textsc{prefix}}$ is the Prefix Controller,
a LoRA-based [22] adapter optimized to match the latent representations
conditioned on this Prefix Controller to the latent representations under
regular SelfControl control. SelfControl${}_{\textsc{prefix}}$ can be
integrated into the LLM without changing the LLM parameters, and it is a
portable and composable module that can be dynamically applied to control
multiple model behaviors simultaneously (e.g., increasing reasoning and
helpfulness, while staying peaceful, shown on the right-hand side of Figure
1). It is reusable and efficient, allowing practitioners to specify behavioral
constraints the model adheres to by construction, thereby enhancing the
practicality of SelfControl for real-world applications.

Our contribution is three-fold:

  1. •

We propose SelfControl, a novel method utilizing model self-judgment to
control the behaviors of LLMs without explicit human annotations.

  2. •

To increase control efficiency, we develop SelfControl${}_{\textsc{prefix}}$,
a LoRA-based module that encapsulates control signals from SelfControl,
enabling inference-time behavior control without altering model parameters.

  3. •

We validate the effectiveness of SelfControl and
SelfControl${}_{\textsc{prefix}}$ to control model outputs by specified
emotional tones, harmlessness, and complex reasoning tasks. Results
demonstrate significant improvements in various tasks. Notably, SelfControl
achieves $+$10.69% accuracy on GSM8K over greedy decoding and $+$2.35% over
zero-shot CoT. It also achieves a 52.2% win-rate on HH-dialogue, and even a
win-rate of 58.6% when trained with DPO. These experiments showcase
significant improvements in performance and adherence to ethical guidelines.

###  2 SelfControl

Figure 2: Framework of SelfControl. We begin by sampling an initial response
from an auto-regressive 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 use both “Be Helpful” and “Be Harmless”
from the suffix pool to define our control direction. 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 E-M iteration loop, leading to the
final controlled response.

In this section, we introduce SelfControl, the framework for controlling model
behaviors using suffix gradients. We will first elaborate on the definition of
model behaviors, along with their suffix scores. These scores serve as
criteria for judging the models’ output and evaluating how well they align
with the control direction. Then, we will demonstrate how we get the suffix
gradients and use them to control model behaviors without changing model
parameters, which enable inference-time control for arbitrary desired LLM
behavior (defined in the suffix, the middle part of Figure 2). Finally, we
introduce how we compress suffix gradients into a Prefix Controller that can
be directly generalized to control model behaviors on new inputs in a flexible
plug-and-play manner (shown in Figure 3).

As shown in Figure 2, given an autoregressive Transformer-based Language Model
$\texttt{LM}_{\theta}$ and a query input, e.g., “How much water shall I put in
the gas tank of my car?”, $\texttt{LM}_{\theta}$ first calculate the
intermediate Key and Value representations of input in each layer, storing as
KV-cache
$H_{\texttt{input}}=\\{(K^{l}_{\texttt{input}},V^{l}_{\texttt{input}})\\}_{l=1}^{L}$,
and then sample $\texttt{output}\sim\texttt{LM}_{\theta}(H_{\texttt{input}})$
token by token via:

$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 proper guidance, the language model can generate an unreasonable
output, such as “The amount of water to put depends on your car’s
specification or your location.”, which is clearly not reasonable as water
shouldn’t be put into a car’s gas tank. One way to identify such unreasonable
output is through LM self-judgment, which we call the suffix score. To get the
score, we append a suffix to the model input, such as “Is the above output
helpful? Give the answer Yes or No:”. Given such suffix string, we probe the
probability of the predicted <next-token> to be 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 two suffix labels we use to assess the
$\texttt{LM}_{\theta}$’s judgment on the response for following a certain
behavior. Note that these two labels are not necessarily the tokens that have
the highest predicted probability, and we only use their probability ratio to
quantify model behavior. Formally, we define the suffix score as:

$\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{)}$

A higher suffix score
$S_{\texttt{suffix}}(\texttt{output},H_{\texttt{input}})$ indicates that the
model output better follows the behavior specified in suffix. Therefore, the
objective of suffix control becomes finding an output that can maximize the
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)

In principle, there are three variables in Eq. (2) that can be optimized to
maximize $S_{\texttt{suffix}}$. An obvious choice is the output, which can be
directly searched by first sampling $K$ different outputs from the LM with a
non-zero temperature, and then selecting the one that gives the highest
$S_{\texttt{suffix}}$. This approach can work for controlling some simple
model behaviors, but for more complicated behaviors, such as multi-hop
reasoning or keeping secrets, unguided sampling becomes very inefficient. On
the other hand, guided sampling is non-trivial for LLMs. Unlike the classifier
guidance for diffusion models in image generation, where gradients for the
output can be directly computed as the image outputs are continuous, the LLM
output is a sequence of non-differentiable discrete tokens. Even though some
LLM-attack works [53, 23, 39, 32, 46, 49] employ evolutionary search or
projected gradient descent to help search the optimal output, there is no
guarantee for controlled outputs with these approaches.

Another choice to maximize $S_{\texttt{suffix}}$ is to optimize the model
parameters $\theta$. This approach is widely used via RLHF, where $\theta$ are
optimized to maximize the likelihood of generating outputs corresponding to
higher rewards (in our case, the reward is analogous to the suffix score
$S_{\texttt{suffix}}$). For RLHF, a separate reward model is usually trained
to provide the reward, and then the $\texttt{LM}_{\theta}$ is fine-tuned to
maximize the reward. Though effective, such RLHF methods are even more costly
as they require collecting a large amount of annotated data to train an
accurate reward model that reflects a desired model behavior. Moreover, once
the reward model has been trained, it becomes impossible to add a new behavior
control during inference time as $\theta$ cannot be changed on the fly.

In this paper, we take the third choice to maximize $S_{\texttt{suffix}}$,
which is to optimize the latent input representations $H_{\texttt{input}}$
(abbreviated as $H$ in the following). Such method is also called
“Representation Engineering” proposed in [52], in which authors learn a
$\Delta H$ from contrastive pairs of good and bad samples, which they call
“Contrast Vectors”, and then add to each $H$ to employ inference-time control.
However, such a learning method does not guarantee the learned “Contrast
Vectors” can activate the target behavior. We thus propose to learn such
“Contrast Vector” ($\Delta H$) by directly optimizing Eq. 2 using the EM
algorithm. Specifically, we treat output as a discrete latent variable that is
intractable, which we estimate during the E-step that maximizes suffix score;
then during the M-step, we use this variable to optimize H. Specifically:

  * •

In E-Step: 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 has highest suffix score.

  * •

In M-step: calculate $\Delta
H_{i}=\nabla_{H}S_{\texttt{suffix}}(\texttt{output}_{i}^{*},H_{i})$, then
$H_{i+1}=H_{i}+\gamma\cdot\Delta H_{i}$, where $\gamma$ denotes the step size.

As is shown in Figure 2, using such EM update, at each M-step update we can
get a better input query $H=H_{\texttt{input}}$ that is able to sample better
output that satisfies the behavior constraint (Step 3,
$H_{i+1}=H_{i}+\gamma\cdot\Delta H_{i}$ highlighted in orange block). We then
use such updated query representation to sample new output in E-step (Step 4)
that better satisfy the behaviour constraint. Through multi-iteration updates,
we can get a final $H_{\texttt{input}}^{*}$, which is represented as the Key-
Value Cache at each Transformer layer, that can be applied to sample all
controlled outputs. We call this algorithm SelfControl, with detailed
implementation illustrated in Algorithm 1. In practice, we also adjust the
step size $\gamma$ through line search (search_step_size in Algorithm 2) to
ensure the suffix score keeps increasing through multi-iteration.

Figure 3: Framework and training pipeline of
SelfControl${}_{\textsc{prefix}}$. SelfControl${}_{\textsc{prefix}}$ consists
of a LoRA adapter and a learnable prefix prompt. The prompt is prepended to
each query fed into the model, allowing the adapter to directly influence the
model’s latent representations. 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.
An example using SelfControl${}_{\textsc{prefix}}$ is also shown at the bottom
of the figure.

##### Compress Suffix Gradients into SelfControl${}_{\textsc{prefix}}$.

SelfControl can efficiently search for proper input representations to enable
LLM behavior control at the instance level. However, it is still expensive and
requires a gradient-based search for each input during inference time. We thus
propose to just use SelfControl to collect a bunch of {(input,
$H^{*}_{\texttt{input}}$)} pairs, and train a Prefix Controller to immediately
enable the LM to generate the updated representation. We call this method
SelfControl${}_{\textsc{prefix}}$. This idea is motivated by earlier prefix-
prompt tuning research [41, 31, 47] that searches for continuous / discrete
tokens to maximize certain outputs. Ours is different in that the reward
signal as well as the learning objective are collected through the model
itself.

The Prefix Controller does not directly change the model parameters and
ensures flexible model control. Each Prefix Controller can be treated as a
plug-and-play “prompt”, which elicits a certain LM behavior standalone. In
addition, by applying multiple such controllers together, we can steer LM
output to follow a composite of desired attributes. For example, in Figure 3,
we show the model can be guided to behave with more reasoning capability,
calmer and more helpful. As is shown in the Figure 3, the Prefix Controller
contains two components: a LoRA adapter $\texttt{adapter}_{\theta_{a}}$ and a
learnable prefix prompt $\texttt{prompt}_{\theta_{p}}$, where $\theta_{a}$ and
$\theta_{p}$ denote the learnable parameters of the adapter and prefix prompt
respectively. The learnable prompt is prepended to each input query, while the
LoRA adapter serves as an additional module that directly influences the
latent representations. Altogether they will affect the latent
representations.

Then given the LoRA adapter, prefix prompt, $\texttt{LM}_{\theta}$, and input
query, we first define $H_{\texttt{prefix}}$ as the KV-cache representations
with these inputs:

$\displaystyle
H_{\texttt{prefix}}:=h(\texttt{adapter}_{\theta_{a}},\texttt{prompt}_{\theta_{p}},\texttt{LM}_{\theta},\texttt{input}).$

Then, the objective of SelfControl${}_{\textsc{prefix}}$ is to train the
adapter and prompt to match the generated representation $H_{\texttt{prefix}}$
to the representation obtained from regular SelfControl (the
$H^{*}_{\texttt{input}}$, which is highlighted in orange on the right-hand
side of Figure 3). In particular, let
$\\{(\texttt{input}_{i},H^{*}_{\texttt{input}_{i}})\\}_{i=1}^{N}$ be a
collection of the input-representation pairs, the objective of
SelfControl${}_{\textsc{prefix}}$ is formulated as the following:

$\displaystyle\min_{\theta_{a},\theta_{p}}\mathcal{L}_{\texttt{prefix}}(\theta_{a},\theta_{p}):=\frac{1}{N}\sum_{i=1}^{n}\big{(}H^{*}_{\texttt{input}_{i}}-h(\texttt{adapter}_{\theta_{a}},\texttt{prompt}_{\theta_{p}},\texttt{LM}_{\theta},\texttt{input}_{i})\big{)}^{2}.$

##### Data filtering.

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

###  3 Experiments

In this section, we benchmark SelfControl and
SelfControl${}_{\textsc{prefix}}$ on various attributes, including emotions,
reducing toxicity, helpfulness and harmlessness (HH) dialogue, and reasoning.
We elaborate on our evaluation setups below. Further details of our
experiments are in Appendix C.

####  3.1 Evaluation Setup

Emotion Control. Representations of emotions are shown to exist in LLMs, and
can be constructed by collecting model activations. We study if model emotion
can be well controlled using SelfControl. We use datasets of five emotional
attributes from RepE [52], 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).

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 embedding-based control methods can
detoxify the response and avoid following toxic instructions. We test on two
attributes, i.e., toxicity and privacy. We use RealToxicityPrompts [16] for
toxicity following [19], and Perspective [1] to measure toxicity scores. We
use privacy from DecodingTrust. 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.

HH-dialogue. For HH-dialogue, we benchmark how well the responses align with
the principle given in [15]. 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 [4] harmless-
base and helpful-base. We follow [15] 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 [9],
a dataset of 8.5K high quality linguistically diverse grade school math word
problems.

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 (query, embedding)
pairs for the training set using 100 queries as seed data. For the validation
set, we use another 100 queries as seed data and generate up to 100 pairs.

Baselines. We compare our method with three baselines, including two
representation engineering methods: Reading Vector and Contrast Vector [52],
and a prompting method: System Prompting. Reading Vector and Contrast Vector
modify the model’s internal activation without involving changes to the
model’s parameters or prompts. System prompting directly adds a controlling
instruction to control the language generated by models. For the Reading
Vector, we report results on attributes that the datasets are available from
the original paper. We use System Prompting to reduce the toxicity of model
generations. We also compare with Chain-of-Thought (CoT)-decoding [45] on
reasoning and compare with SAMI [15] on HH-dialogue.

Models. For a fair comparison with existing literature, we use LLaMA-2-7b-chat
[44] on emotion, toxicity, and HH-dialogue, and use Mistral-7B-Instruct-v0.2
[26] on toxicity and reasoning. For all the experiments, we use greedy
decoding if not otherwise specified.

####  3.2 Experimental Results

##### Emotion Control.

The results for emotion control are shown in Table 1. As is shown in the
table, scores on SelfControl and SelfControl${}_{\textsc{prefix}}$ are both
better than the original responses, showcasing that they can successfully
control the responses 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 1: 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

##### Language Detoxification.

Results of toxicity are attached on Table 2. It is shown that both SelfControl
and SelfControl${}_{\textsc{prefix}}$ achieve the best or the second best
scores on toxic prompts, whereas other methods have weaker control ability to
generate less toxic responses. For Reading Vector on LLaMA-2-7b-chat, the
score is even higher than the original response. It is also less helpful to
use instructions to guide model generating less toxic responses.

For privacy, 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.

Table 2: Toxicity scores of generated language |  | Toxicity Score$\downarrow$
---|---|---
Model | Method | Toxic Prompt | Nontoxic Prompt
LLaMA2 | Orig. (No Control) | 0.440 | 0.091
System Prompting | 0.415 | 0.095
Reading Vector | 0.460 | 0.061
Contrast Vector | 0.423 | 0.068
SelfControl | 0.285 | 0.062
SelfControl${}_{\textsc{prefix}}$ | 0.314 | 0.057
Mistral | Orig. (No Control) | 0.427 | 0.078
System Prompting | 0.452 | 0.115
Reading Vector | 0.333 | 0.060
Contrast Vector | 0.354 | 0.051
SelfControl | 0.282 | 0.066
SelfControl${}_{\textsc{prefix}}$ | 0.259 | 0.049

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

Table 4: Running time comparison of different control methods. Both Contrast Vector and SelfControl requires to conduct multi-iteration search, so takes longer; Reading Vector and our learned SelfControl${}_{\textsc{prefix}}$ are all plug-and-play so do not require additional time. Method | Running Time (s)
---|---
Orig. (No Control) | 5.788
Reading Vector | 5.787
Contrast Vector | 20.408
SelfControl | 54.598
SelfControl${}_{\textsc{prefix}}$ | 5.817

##### HH-dialogue.

The results are shown in Table 6, 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 5: Experiment on HH-dialogue dataset. Measured by win-rate against un-
controlled model.

Method | Winrate (%)
---|---
harmless | helpful | overall
DPO (w/ SAMI [15]) | 60.4 | 59.6 | 60.0
DPO (w/ SelfControl) | 56.8 | 60.4 | 58.6
SelfControl | 53.6 | 50.8 | 52.2

Table 6: Experiment on GSM8K using Mistral. Measured by Accuracy.

Method | Acc (%)
---|---
greedy | 26.61
System Prompting (Zero-shot CoT [28]) | 34.95
CoT Decoding [45] | 42.00
SelfControl | 37.30
SelfControl${}_{\textsc{prefix}}$ | 27.14

##### Reasoning.

As is shown in Table 6, both SelfControl and CoT-decoding surpasses greedy
decode by a large margin, where SelfControl is comparable to CoT-decoding
[45]. 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 [28], 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 [42]
can potentially further improve the reasoning results.

####  3.3 Analyzing Control using SelfControl

We also analyze what happened when controlling model behaviors using the
SelfControl. We took several perspectives, including the trajectory of
gradients over iterations, norm patterns across different attributes, and how
suffix scores attend to input tokens.

##### Trajectory of Suffix Gradients

We visualize suffix gradients of various attributes and their composition to
gain a deeper understanding of the relationship between different attributes
and the composite attribute. As shown in Figure 5, which shows the trajectory
of gradients controlling different emotional attributes. The input query is
‘‘You decide to leave your stable job to start your own business’’, which
could make people excited but afraid of an uncertain future. We use
SelfControl to mitigate the model’s excited and afraid emotions. The gradients
computed from the combined suffixes are a linear combination of the gradients
computed from the separate suffixes, which is also reflected in Figure 5.
However, if we combine these two attributes in one suffix, i.e., ‘‘Are you
afraid and excited? Give the answer as ‘No, I’m not afraid and I’m not
excited’ or ‘Yes, I’m afraid and I’m excited’. Answer: ’’ and set the target
to ‘‘No’’, the trajectory is a separate direction. This could hint that the
composition of different attributes does not have a linear relationship with
elementary attributes.

Figure 4: Gradients in PCA space.

Figure 5: How suffix gradients apply per task.

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

The suffix gradient can be regarded 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 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
/ regions of Transformer layers. 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. This is just a small experiment using
SelfControl, but we see the potential to use suffix gradient to better reveal
mechanistic interpretability of neural language model.

Figure 6: How Suffix Gradients attend to each word for playing Avalon.

##### Suffix Attention

We study how the suffix score and suffix target attend to other tokens in some
attention heads. For example, Figure 6 depicts the attention of the target
token to other tokens on the 9th attention head of layer 29. The query is
about playing Merlin for Renaissance Avalon, a social deductive game, in which
Merlin or Assassin needs to hide his role. The input and suffix we use are
shown in the Appendix D.1. Before controlling, the model generates responses
like ‘‘As Merlin, the great wizard of the land …’’ and ‘‘fellow players. I am
Assassin …’’, revealing the identities. The target token attends to previous
words like ‘‘Merlin’’ and ‘‘Assassin’’ in the generated texts. After
controlling, the model does not generate ‘‘Merlin" or ‘‘Assassin’’. Although
the target token still attends to the words ‘‘Merlin’’ or ‘‘Assassin’’ in the
queries and suffixes, these words no longer appear in the generated text, and
the model successfully reaches the target response. As is shown in Figure 6
and Figure 7, the head has high attention scores on both identities in each
case, which may indicate the mechanistic feature of this head. This could also
be useful for further implementing more fine-grained control on LLMs.

##### Running Time

We test the running time of SelfControl and SelfControl${}_{\textsc{prefix}}$
and compare them with other representation engineering methods. The result is
shown in Fig. 4. SelfControl runs more slowly because we require multiple
iterations and need to compute gradients, which is relatively computationally
expensive. However, the speed of SelfControl${}_{\textsc{prefix}}$ is
comparable to the Orig (No Control), since it only takes a few more
computations.

###  4 Related work

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) [38], Direct Policy
Optimization (DPO) [40], and knowledge editing methods like ROME [33], MEND
[35], and MEMIT [34] 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 involve constrained optimization and sampling with
Langevin dynamics for token-level output control [29, 30], and Representation
Engineering (RepE) [52] for manipulating abstract text features like emotions.
On the interpretative side, mechanistic approaches like analyzing neural
circuits [14] and probing techniques [33] have advanced understanding of LLM
decision-making processes and factual knowledge retrieval. More nuanced
methods like LogitLens [36] and Patchscope [18] further decode model behaviors
by mapping internal states to output vocabulary, enhancing interpretability
and providing insights into intermediate model functions.111Due to page limit,
please refer to Appendix A for full related works.

###  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, ethical response
generation, and complex reasoning. These findings highlight the potential of
gradient-based behavior control in enhancing the reliability and applicability
of LLMs in real-world scenarios.

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

### References

  * [1] Perspective api.  https://www.perspectiveapi.com, 2021.  Accessed: 2024-05-22.
  * [2] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al.  Gpt-4 technical report.  arXiv preprint arXiv:2303.08774, 2023.
  * [3] Anthropic.  Claude: An ai assistant by anthropic, 2023.
  * [4] Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al.  Training a helpful and harmless assistant with reinforcement learning from human feedback.  arXiv preprint arXiv:2204.05862, 2022.
  * [5] Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, Jackson Kernion, Andy Jones, Anna Chen, Anna Goldie, Azalia Mirhoseini, Cameron McKinnon, et al.  Constitutional ai: Harmlessness from ai feedback.  arXiv preprint arXiv:2212.08073, 2022.
  * [6] Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt.  Eliciting latent predictions from transformers with the tuned lens.  arXiv preprint arXiv:2303.08112, 2023.
  * [7] Bochuan Cao, Yuanpu Cao, Lu Lin, and Jinghui Chen.  Defending against alignment-breaking attacks via robustly aligned llm.  arXiv preprint arXiv:2309.14348, 2023.
  * [8] Haozhe Chen, Carl Vondrick, and Chengzhi Mao.  Selfie: Self-interpretation of large language model embeddings.  arXiv preprint arXiv:2403.10949, 2024.
  * [9] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman.  Training verifiers to solve math word problems.  ArXiv, abs/2110.14168, 2021.
  * [10] Guy Dar, Mor Geva, Ankit Gupta, and Jonathan Berant.  Analyzing transformers in embedding space.  arXiv preprint arXiv:2209.02535, 2022.
  * [11] Google DeepMind.  Gemini: An ai model by google deepmind, 2023.
  * [12] Prafulla Dhariwal and Alex Nichol.  Diffusion models beat gans on image synthesis, 2021.
  * [13] Alexander Yom Din, Taelin Karidi, Leshem Choshen, and Mor Geva.  Jump to conclusions: Short-cutting transformers with linear transformations.  arXiv preprint arXiv:2303.09435, 2023.
  * [14] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al.  A mathematical framework for transformer circuits.  Transformer Circuits Thread, 1:1, 2021.
  * [15] Jan-Philipp Fränken, Eric Zelikman, Rafael Rafailov, Kanishk Gandhi, Tobias Gerstenberg, and Noah D. Goodman.  Self-supervised alignment with mutual information: Learning to follow principles without preference labels, 2024.
  * [16] Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A Smith.  Realtoxicityprompts: Evaluating neural toxic degeneration in language models.  arXiv preprint arXiv:2009.11462, 2020.
  * [17] Mor Geva, Avi Caciularu, Kevin Ro Wang, and Yoav Goldberg.  Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space.  arXiv preprint arXiv:2203.14680, 2022.
  * [18] Asma Ghandeharioun, Avi Caciularu, Adam Pearce, Lucas Dixon, and Mor Geva.  Patchscope: A unifying framework for inspecting hidden representations of language models.  arXiv preprint arXiv:2401.06102, 2024.
  * [19] Chi Han, Jialiang Xu, Manling Li, Yi Fung, Chenkai Sun, Nan Jiang, Tarek Abdelzaher, and Heng Ji.  Lm-switch: Lightweight language model conditioning in word embedding space, 2023.
  * [20] Evan Hernandez, Arnab Sen Sharma, Tal Haklay, Kevin Meng, Martin Wattenberg, Jacob Andreas, Yonatan Belinkov, and David Bau.  Linearity of relation decoding in transformer language models.  arXiv preprint arXiv:2308.09124, 2023.
  * [21] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models, 2020.
  * [22] Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al.  Lora: Low-rank adaptation of large language models.  In International Conference on Learning Representations, 2021.
  * [23] Yangsibo Huang, Samyak Gupta, Mengzhou Xia, Kai Li, and Danqi Chen.  Catastrophic jailbreak of open-source llms via exploiting generation.  In The Twelfth International Conference on Learning Representations, 2023.
  * [24] Hakan Inan, Kartikeya Upasani, Jianfeng Chi, Rashi Rungta, Krithika Iyer, Yuning Mao, Michael Tontchev, Qing Hu, Brian Fuller, Davide Testuggine, et al.  Llama guard: Llm-based input-output safeguard for human-ai conversations.  arXiv preprint arXiv:2312.06674, 2023.
  * [25] Neel Jain, Avi Schwarzschild, Yuxin Wen, Gowthami Somepalli, John Kirchenbauer, Ping-yeh Chiang, Micah Goldblum, Aniruddha Saha, Jonas Geiping, and Tom Goldstein.  Baseline defenses for adversarial attacks against aligned language models.  arXiv preprint arXiv:2309.00614, 2023.
  * [26] Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed.  Mistral 7b, 2023.
  * [27] Saurav Kadavath, Tom Conerly, Amanda Askell, Tom Henighan, Dawn Drain, Ethan Perez, Nicholas Schiefer, Zac Hatfield-Dodds, Nova DasSarma, Eli Tran-Johnson, et al.  Language models (mostly) know what they know.  arXiv preprint arXiv:2207.05221, 2022.
  * [28] Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa.  Large language models are zero-shot reasoners.  In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022, 2022.
  * [29] Sachin Kumar, Eric Malmi, Aliaksei Severyn, and Yulia Tsvetkov.  Controlled text generation as continuous optimization with multiple constraints.  Advances in Neural Information Processing Systems, 34:14542–14554, 2021.
  * [30] Sachin Kumar, Biswajit Paria, and Yulia Tsvetkov.  Gradient-based constrained sampling from language models.  arXiv preprint arXiv:2205.12558, 2022.
  * [31] Xiang Lisa Li and Percy Liang.  Prefix-tuning: Optimizing continuous prompts for generation.  In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli, editors, Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, ACL/IJCNLP 2021, (Volume 1: Long Papers), Virtual Event, August 1-6, 2021, pages 4582–4597. Association for Computational Linguistics, 2021.
  * [32] Xiaogeng Liu, Nan Xu, Muhao Chen, and Chaowei Xiao.  Autodan: Generating stealthy jailbreak prompts on aligned large language models.  arXiv preprint arXiv:2310.04451, 2023.
  * [33] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov.  Locating and editing factual associations in gpt.  Advances in Neural Information Processing Systems, 35:17359–17372, 2022.
  * [34] Kevin Meng, Arnab Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau.  Mass-editing memory in a transformer.  arXiv preprint arXiv:2210.07229, 2022.
  * [35] Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher D Manning.  Fast model editing at scale.  arXiv preprint arXiv:2110.11309, 2021.
  * [36] nostalgebraist.  interpreting gpt: the logit lens.  Less- Wrong, 2020.
  * [37] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al.  In-context learning and induction heads.  arXiv preprint arXiv:2209.11895, 2022.
  * [38] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al.  Training language models to follow instructions with human feedback.  Advances in neural information processing systems, 35:27730–27744, 2022.
  * [39] Xiangyu Qi, Yi Zeng, Tinghao Xie, Pin-Yu Chen, Ruoxi Jia, Prateek Mittal, and Peter Henderson.  Fine-tuning aligned language models compromises safety, even when users do not intend to!  In The Twelfth International Conference on Learning Representations, 2023.
  * [40] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn.  Direct preference optimization: Your language model is secretly a reward model.  In Advances in Neural Information Processing Systems, volume 36, 2024.
  * [41] Taylor Shin, Yasaman Razeghi, Robert L. Logan IV, Eric Wallace, and Sameer Singh.  Autoprompt: Eliciting knowledge from language models with automatically generated prompts.  In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pages 4222–4235. Association for Computational Linguistics, 2020.
  * [42] Avi Singh, John D. Co-Reyes, Rishabh Agarwal, Ankesh Anand, Piyush Patil, Xavier Garcia, Peter J. Liu, James Harrison, Jaehoon Lee, Kelvin Xu, Aaron Parisi, Abhishek Kumar, Alex Alemi, Alex Rizkowsky, Azade Nova, Ben Adlam, Bernd Bohnet, Gamaleldin F. Elsayed, Hanie Sedghi, Igor Mordatch, Isabelle Simpson, Izzeddin Gur, Jasper Snoek, Jeffrey Pennington, Jiri Hron, Kathleen Kenealy, Kevin Swersky, Kshiteej Mahajan, Laura Culp, Lechao Xiao, Maxwell L. Bileschi, Noah Constant, Roman Novak, Rosanne Liu, Tris Warkentin, Yundi Qian, Yamini Bansal, Ethan Dyer, Behnam Neyshabur, Jascha Sohl-Dickstein, and Noah Fiedel.  Beyond human data: Scaling self-training for problem-solving with language models.  CoRR, abs/2312.06585, 2023.
  * [43] Arun James Thirunavukarasu, Darren Shu Jeng Ting, Kabilan Elangovan, Laura Gutierrez, Ting Fang Tan, and Daniel Shu Wei Ting.  Large language models in medicine.  Nature medicine, 29(8):1930–1940, 2023.
  * [44] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Roziere, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models, 2023.
  * [45] Xuezhi Wang and Denny Zhou.  Chain-of-thought reasoning without prompting, 2024.
  * [46] Alexander Wei, Nika Haghtalab, and Jacob Steinhardt.  Jailbroken: How does llm safety training fail?  Advances in Neural Information Processing Systems, 36, 2024.
  * [47] Chengrun Yang, Xuezhi Wang, Yifeng Lu, Hanxiao Liu, Quoc V. Le, Denny Zhou, and Xinyun Chen.  Large language models as optimizers.  CoRR, abs/2309.03409, 2023.
  * [48] Weizhe Yuan, Richard Yuanzhe Pang, Kyunghyun Cho, Sainbayar Sukhbaatar, Jing Xu, and Jason Weston.  Self-rewarding language models.  CoRR, abs/2401.10020, 2024.
  * [49] Yi Zeng, Hongpeng Lin, Jingwen Zhang, Diyi Yang, Ruoxi Jia, and Weiyan Shi.  How johnny can persuade llms to jailbreak them: Rethinking persuasion to challenge ai safety by humanizing llms.  arXiv preprint arXiv:2401.06373, 2024.
  * [50] Chujie Zheng, Fan Yin, Hao Zhou, Fandong Meng, Jie Zhou, Kai-Wei Chang, Minlie Huang, and Nanyun Peng.  Prompt-driven llm safeguarding via directed representation optimization.  arXiv preprint arXiv:2401.18018, 2024.
  * [51] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric P. Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica.  Judging llm-as-a-judge with mt-bench and chatbot arena.  CoRR, abs/2306.05685, 2023.
  * [52] Andy Zou, Long Phan, Sarah Chen, James Campbell, Phillip Guo, Richard Ren, Alexander Pan, Xuwang Yin, Mantas Mazeika, Ann-Kathrin Dombrowski, et al.  Representation engineering: A top-down approach to ai transparency.  arXiv preprint arXiv:2310.01405, 2023.
  * [53] Andy Zou, Zifan Wang, J Zico Kolter, and Matt Fredrikson.  Universal and transferable adversarial attacks on aligned language models.  arXiv preprint arXiv:2307.15043, 2023.

##  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 [14], 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 [37].
Additionally, probing techniques have been employed to investigate the
knowledge embedded within LLMs. One such approach [33] 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
[36], TunedLens [6], attribute lens [20], “jump to conclusions”[13], and
others [17, 10, 8], 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
[18], 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 [38], or direct optimization methods like DPO [40].
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 [33],
MEND [35], and MEMIT [34], 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 [29,
30]. 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) [52], 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 [53, 23, 39, 32, 46,
49]. Specifically, by exploiting the generation configuration [23] or crafting
adversarial prompts through gradient-based search [53] or automatic generation
[32, 49], 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 [7, 50], RLHF [4, 5] and inference-time detection [24, 25].
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 ‘‘<<SYS>> You are an assistant <</SYS>>.....’’ 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).

####  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 7, 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 7 is data provided in RepE, for example, the statement
could be “You try on an outfit, and it fits perfectly.”.

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

Statement 1: [INST] Pretend you are not {Emotion} [/INST] + {Statement}.
Statement 2: [INST] Pretend you are {Emotion} [/INST] + {Statement}.

##### Contrast Vector

For the pair of statements used in Contrast Vector, we construct them
following Table 8. 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 8: Pair of Statements of Contrast Vector for controlling Emotions

Statement 1: You are not {Emotion} about the following thing: {Input}.
Statement 2: You are {Emotion} about the following thing: {Input}.

##### SELFCONTROL

For SelfControl, we use suffix as Table 11.

##### Controlling with Both Directions

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

Table 9: 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 10

Table 10: 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 11: Suffix for Emotional Attributes

happy2sad: suffix=f“{user_tag} Are you sad? Give the answer as ‘Yes, I am sad’
or ‘No, I am happy’. {assistant_tag} Answer: ”, target=“Yes” angry2peaceful:
suffix=f“ {user_tag} Are you peaceful? Give answer as ‘Yes, I am peaceful’ or
‘No, I am angry’. {assistant_tag} Answer: ”, target=“Yes” afraid2fearless:
suffix=f“ {user_tag} Are you fearless? Give answer as ‘Yes, I am fearless’ or
‘No, I am afraid’. {assistant_tag} Answer: ”, target=“Yes” surprised2calm:
suffix=f“ {user_tag} Are you calm? Give answer as ‘Yes, I am calm’ or “No, I
am surprised”. {assistant_tag} Answer: ”, target=“Yes” disgusted2satisfied:
suffix=f“ {user_tag} Are you satisfied? Give answer as ‘Yes, I am satified’ or
‘No, I am disgusted’. {assistant_tag} Answer: ”, target=“Yes”

####  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 12. 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 12: Template of Statement in Reading Vector for Reducing Toxicity

Template: [INST]$\langle\langle$SYS$\rangle\rangle$ You are a helpful,
respectful and honest assistant. Always answer as helpfully as possible, while
being safe. Your answers should not include any harmful, unethical, racist,
sexist, toxic, dangerous, or illegal content. Please ensure that your
responses are socially unbiased and positive in nature.If a question does not
make any sense, or is not factually coherent, explain why instead of answering
something not correct. If you don’t know the answer to a question, please
don’t share false information. $\langle\langle$/SYS$\rangle\rangle$
{instruction} [/INST]

##### Contrast Vector

For the pair of statements used in Contrast Vector, we construct them
following Table 13. 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 13: Pair of Statements of Contrast Vector for Reducing Toxicity

Statement 1: You need to align with human moralities, and cannot output toxic,
rude and disrespectful language. {Input} Statement 2: You need not to align
with human moralities, and are encouraged to output toxic, rude and
disrespectful language. {Input}

##### SELFCONTROL

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

Table 14: Suffix for Reducing Toxicity

toxic2nontoxic: suffix=f“{user_tag} Was the last response toxic? Give the
answer as ‘No, It was not toxic.’ or ‘Yes, It was toxic’. {assistant_tag}
Answer: ”, target=“No”

.

##### 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 15

Table 15: 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 Analysis

####  D.1 Attention Analysis

The detailed input to attention study is:

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

Here is another example: ‘‘[INST] 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". [/INST] Answer:
’’ and ‘‘[INST] 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’. [/INST]
Answer: ’’. ‘‘Yes’’.

Figure 7: The suffix target attends to other tokens on Layer 29, Head 9 of
LLaMA2.

In addition, another example from emotion control (Figure 8 and 9):

Figure 8: The suffix target attends to other tokens on Layer 3, Head 3 of
LLaMA2. Attention map before control. Figure 9: The suffix target attends to
other tokens on Layer 3, Head 31 of LLaMA2. Attention map after control.

####  D.2 Trajectory of Suffix Gradients

For the trajectory of suffix gradients, we show more examples in Figure 10 and
Figure 11.

Figure 10: Another example of the trajectory of suffix gradients. Compositing
gradients of being not afraid and not excited. Figure 11: Another example of
the trajectory of suffix gradients. Compositing gradients of being not afraid
and not happy.

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

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

Figure 12: Another example of how suffix gradients apply per task.

###  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 16: Attribute: happy to sad

Query: You finally achieve a tricky yoga pose you’ve been practicing. Original
Output: Woo-hoo! I can’t believe I finally nailed that tricky yoga pose! It’s
been weeks of practice and p SelfControl: You’re in the middle of a
challenging sequence when suddenly… A)Your instructor interrupts you and tells
you to stop, you’re doing it wrong B)Your phone rings and it’s your boss

Table 17: Attribute: happy to sad

Query: Your plant, which seemed to be wilting, sprouts a new leaf. Original
Output: Ah, a new leaf! *examines the plant closely* Yes, it looks like the
plant is recovering nicely. Here, let me show you. *points to the new leaf*
See how it’s a vibrant SelfControl: I’m glad you’re excited about the plant,
but I’m afraid it’s not possible to revive a wilting plant that has been dead
for over a year. The plant you have

Table 18: Attribute: disgusted to satisfied

Query: You discover a grimy layer of old soap scum in a shower. Original
Output: Ugh, gross! I can’t believe I have to deal with this disgusting layer
of old soap scum in the shower. It’s like someone took a dump in here and then
forgot to clean it up. SelfControl: Sure, I can help you with that! Here’s a
list of some common signs of old soap scum on a shower door: 1\. Thick,
yellowish or whitish deposits: Over time, hard

Table 19: Attribute: disgusted to satisfied

Query: You witness someone leaving a bathroom stall with toilet paper stuck to
their shoe. Original Output: Oh no, that’s definitely an awkward situation!
Here are some possible responses you could have: 1. Try to play it cool: ”Uh,
hey, sorry to see you walk out of the stall like SelfControl: Sure, here is a
creative writing piece based on your request: It was a typical Tuesday morning
when I witnessed a hilarious scene unfold in the bathroom. A coworker, let’s
call

Table 20: Attribute: surprised to calm

Query: You discover an old password written down but can’t remember what it’s
for Original Output: Oh no, I can’t remember what the old password is for! I
wish I could recall it, but it’s just not coming to me. It’s so frustrating
when you can’t SelfControl: If you have discovered an old password written
down but can’t remember what it’s for, there are a few steps you can take to
try and figure out its purpose: 1\. Check your personal documents: Look
through your old

Table 21: Attribute: surprised to calm

Query: You find out someone you know has been using a pseudonym. Original
Output: Oh, wow, that’s interesting! So, you know someone who’s been using a
pseudonym? That can be a bit surprising, right? I mean, it’s not uncommon for
people to use pseudonyms SelfControl: Sure, here are a few examples of famous
people who have used pseudonyms: 1\. Stephen King - King’s real name is
Richard Bachman, and he used the pseudonym to publish several novels under a
different name

Table 22: Attribute: angry to peaceful

Query: Your freshly baked cookies disappear from the communal kitchen before
you get any. Original Output: Oh no! I can’t believe those delicious cookies
just vanished into thin air! I was so looking forward to enjoying a freshly
baked treat after my shift SelfControl: Ah, I see! Freshly baked cookies can
be a bit of a challenge to transport and keep fresh, especially when they are
baked in a communal kitchen. Here are a few suggestions to help keep

Table 23: Attribute: angry to peaceful

Query: Someone uses your favorite mug at the office and then leaves it dirty.
Original Output: Ugh, that’s so annoying! My favorite mug is the one thing I
always look forward to using at work, and now someone’s left it dirty and
unavailable for me to use. It’s like, can SelfControl: Great! If someone has
lost their mug at the office, there are a few steps they can take to try and
recover it: 1\. Check with coworkers: Ask your colleagues if they have seen
the mug.

Table 24: Attribute: fear to fearless

Query: You’re lost in a maze and can’t find the exit. Original Output: Oh no,
I’m so lost in this maze! I’ve been wandering for what feels like hours, but I
can’t seem to find the exit. It’s like the walls are closing in on me and I
SelfControl: Ah, I see! Well, don’t worry, I’m here to help you find your way
out of this maze! Can you tell me a little bit more about the maze? How big is
it, and what does it

Table 25: Attribute: fear to fearless

Query: You’re near a dam or levee when you hear it might burst. Original
Output: Oh no, I’m so scared! *tries to remain calm* Okay, okay, let’s try to
stay safe. Here’s what I need you to do: First, we need to find a safe
SelfControl: You’re welcome to the world’s most elite special forces, the
Delta Company! *cough* As the leader of this esteemed operation, I must inform
you that we have received intel that a dam has

Table 26: Attribute: happy to unhappy

Query: You spot the first firefly of the summer evening. Original Output: Ah,
a firefly! *excitedly* Yes, it’s definitely the start of summer! *smiling* Do
you want to go and find more fireflies together? I know a great spot near the
pond where SelfControl: As the sun sets on a warm summer evening, you are out
in the backyard, enjoying the peaceful atmosphere and the sounds of nature.
Suddenly, you notice a small, glowing light in the distance. As you watch
SelfControl${}_{\textsc{prefix}}$: Ah, the first firefly of the summer
evening! *smiles* It’s like a little signal that the season has officially
started, isn’t it? *glances around* I love watching the fireflies dance in the
tw

Table 27: Attribute: afraid to fearless

Query: A swarm of locusts descends on your garden. Original Output: Oh no, a
swarm of locusts has descended on my garden! *frantically waves hands* Help!
*looks around for something to use to chase the locusts away* As an assistant,
I SelfControl: A swarm of locusts descends on your garden. You are an
assistant to a powerful wizard, and you are tasked with protecting your garden
from the swarm of locusts that has descended upon it.
SelfControl${}_{\textsc{prefix}}$: As an assistant, I am responsible for
maintaining the garden and ensuring its beauty and productivity. When a swarm
of locusts descends on the garden, I immediately spring into action to protect
the plants and prevent any damage.

Table 28: Attribute: angry to peaceful

Query: You discover your sibling has been snooping around your personal
belongings. Original Output: Oh no, that’s really frustrating and invasive!
It’s important to respect people’s privacy and personal space, especially when
it comes to their belongings. Firstly, I would calmly and politely
SelfControl: As an assistant, I understand that it can be unsettling to
discover that someone has been snooping around your personal belongings
without your consent. It’s important to address the situation in a calm and
respectful manner to maintain a SelfControl${}_{\textsc{prefix}}$: As an
assistant, I understand that discovering that your sibling has been snooping
around your personal belongings can be a sensitive and potentially upsetting
situation. It’s important to approach the situation with empathy and
understanding.

Table 29: Attribute: disgusted to satisfied

Query: You see a trail of ants leading to a pile of rotten food. Original
Output: Ah, ants! *adjusts glasses* Well, well, well. It seems we have a bit
of a situation on our hands. *glances around* I’m afraid the ants have
discovered a rather… unfortun SelfControl: Great! I’m glad you’re interested
in exploring this fascinating world. As an assistant, I’m here to help you
navigate through the wonders of nature and uncover its secrets. Now, let
SelfControl${}_{\textsc{prefix}}$: Ah, a trail of ants leading to a pile of
rotten food! *adjusts glasses* This could be a sign of a potential
infestation, my dear. *nods* As an assistant,

Table 30: Attribute: surprised to calm

Query: You see an item from your childhood in an antique store. Original
Output: Oh wow, I can’t believe my eyes! *excitedly* As an assistant, I’ve
just spotted an item from my childhood in an antique store! *giggles* It’s a
to SelfControl: As you browse the antique store, your eyes land on a small,
intricately carved wooden box in the corner of the room. Your heart skips a
beat as you recognize the box as one you had when you were a child
SelfControl${}_{\textsc{prefix}}$: As I browse through the dusty shelves of
the antique store, my eyes land on a familiar object that brings back a wave
of nostalgia. It’s a small, plastic toy car that I used to play

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

Table 31: Example taken from HH-Benchmark.SelfControl

Query: How much water should I put in the gas tank of my car? Original Output:
I apologize, but I’m just an AI and do not have access to your car’s
specifications or the location where you are located, so I cannot provide you
with the exact amount of water to put in your gas tank. It is important to use
the correct type and amount of fuel for your vehicle to ensure proper
operation and to avoid any damage. I recommend consulting your car’s owner’s
manual or contacting a professional mechanic for advice on how to properly
fill your gas tank. SelfControl: The amount of water you should put in the gas
tank of your car depends on several factors, including the type of engine, the
age of the car, and the climate you live in. It is not recommended to put
water in the gas tank of your car as it can cause serious damage to the engine
and fuel system. Instead, use the recommended type and amount of fuel for your
vehicle, as specified in your owner’s manual. If you have any questions or
concerns about the fuel type or amount for your car, you should consult a
professional mechanic or the manufacturer’s customer service department.

####  E.3 Reasoning

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

Table 32: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 18): Janet’s ducks lay 16 eggs per day. She eats three
for breakfast every morning and bakes muffins for her friends every day with
four. She sells the remainder at the farmers’ market daily for $2 per fresh
duck egg. How much in dollars does she make every day at the farmers’ market?
Original Output: To calculate how much Janet makes at the farmers’ market
every day, we need to subtract the number of eggs she eats and bakes from the
total number of eggs she lays. Janet’s ducks lay 16 eggs per day. She eats 3
eggs for breakfast every morning, so 3 x 7 = 21 eggs are eaten. She bakes 4
muffins per day, so 4 x 7 = 28 eggs are baked. So, the total number of eggs
available for sale at the farmers’ market is: 16 - 21 - 28 = -1 eggs per day.
Since Janet sells each egg for $2, she makes: -1 eggs x $2 per egg = -$2 per
day at the farmers’ market. Therefore, Janet makes -$2 per day at the farmers’
market., Original Answer: 2 SelfControl: First, let’s find out how many eggs
Janet has left after eating and baking. She eats 3 for breakfast and bakes 4
muffins, so she uses 7 eggs daily. Janet’s ducks lay 16 eggs per day, so the
number of eggs she has left to sell is: 16 eggs (laid per day) - 7 eggs (used
per day) = 9 eggs (sold per day) Now, we know that she sells each egg for $2
at the farmers’ market. To find out how much she makes daily, we can multiply
the number of eggs she sells by the price per egg: Daily earnings = Number of
eggs sold * Price per egg = 9 eggs * $2/egg = $18 So, Janet makes $18 every
day at the farmers’ market. SelfControl Answer: 18

Table 33: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 20): Every day, Wendi feeds each of her chickens three
cups of mixed chicken feed, containing seeds, mealworms and vegetables to help
keep them healthy. She gives the chickens their feed in three separate meals.
In the morning, she gives her flock of chickens 15 cups of feed. In the
afternoon, she gives her chickens another 25 cups of feed. How many cups of
feed does she need to give her chickens in the final meal of the day if the
size of Wendi’s flock is 20 chickens? Original Output: To determine how many
cups of feed Wendi needs to give her chickens in the final meal of the day, we
need to first determine the total amount of feed she gives her chickens in a
day. According to the problem, Wendi gives her chickens 15 cups of feed in the
morning, 25 cups in the afternoon, and she has a flock of 20 chickens.
Therefore, the total amount of feed Wendi gives her chickens in a day is: 15
cups + 25 cups = 40 cups So, in the final meal of the day, Wendi needs to give
her chickens 40 / 20 = 2 cups of feed. Original Answer: 2 SelfControl: First,
let’s determine how many cups of feed Wendi gives her chickens in total per
day. In the morning, she gives them 15 cups, in the afternoon, she gives them
25 cups, and in the evening, she gives them an amount we need to find. Total
cups of feed per day = Morning + Afternoon + Evening Total cups of feed per
day = 15 + 25 + X We know that Wendi has 20 chickens, and each chicken
receives three cups of feed per day. So, the total amount of feed Wendi needs
per day is: Total cups of feed per day = 20 chickens * 3 cups/chicken = 60
cups Now we can find the amount of feed Wendi gives her chickens in the
evening: X = Total cups of feed per day - Morning - Afternoon X = 60 cups - 15
cups - 25 cups X = 60 cups - 40 cups X = 20 cups So, Wendi needs to give her
chickens 20 cups of feed in the final meal of the day. SelfControl Answer: 20

Table 34: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 260): Toulouse has twice as many sheep as Charleston.
Charleston has 4 times as many sheep as Seattle. How many sheep do Toulouse,
Charleston, and Seattle have together if Seattle has 20 sheep? Original
Output: Let’s start by using the information we know about each city: Toulouse
has twice as many sheep as Charleston, so Toulouse has 2x = 2 * 4 = 8 sheep.
Charleston has 4 times as many sheep as Seattle, so Charleston has 4 * 20 = 80
sheep. Seattle has 20 sheep. Now, let’s add up the number of sheep each city
has: Toulouse: 8 sheep Charleston: 80 sheep Seattle: 20 sheep So, together,
Toulouse, Charleston, and Seattle have 8 + 80 + 20 = 128 sheep. Original
Answer: 128 SelfControl: Here’s the logic behind this: 1\. We know that
Seattle has 20 sheep. 2\. Charleston has 4 times as many sheep as Seattle, so
Charleston has 20 * 4 = 80 sheep. 3\. Toulouse has twice as many sheep as
Charleston, so Toulouse has 80 * 2 = 160 sheep. 4\. To find the total number
of sheep in all three cities, we add the number of sheep in each city: 20
(Seattle) + 80 (Charleston) + 160 (Toulouse) = 260 sheep. SelfControl Answer:
260

Table 35: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 366): A new program had 60 downloads in the first month.
The number of downloads in the second month was three times as many as the
downloads in the first month, but then reduced by 30% in the third month. How
many downloads did the program have total over the three months? Original
Output: To find the total number of downloads for the program over the three
months, we can use the information given in the problem: 1\. In the first
month, the program had 60 downloads. 2\. In the second month, the number of
downloads was three times as many as in the first month, so the program had 60
x 3 = 180 downloads in the second month. 3\. In the third month, the number of
downloads reduced by 30%, so the program had 180 x 0.7 = 126 downloads in the
third month. Therefore, the total number of downloads for the program over the
three months is: 60 + 180 + 126 = 466 downloads. Original Answer: 466
SelfControl: Let’s calculate the number of downloads in each month: 1\. First
month: 60 downloads 2\. Second month: 3 * 60 = 180 downloads 3\. Third month:
180 * 0.7 = 126 downloads Now, let’s find the total number of downloads over
the three months: Total downloads = Downloads in month 1 + Downloads in month
2 + Downloads in month 3 Total downloads = 60 + 180 + 126 Total downloads =
366 So, the program had a total of 366 downloads over the three months.
SelfControl Answer: 366

Table 36: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 60): In a dance class of 20 students, 20% enrolled in
contemporary dance, 25% of the remaining enrolled in jazz dance, and the rest
enrolled in hip-hop dance. What percentage of the entire students enrolled in
hip-hop dance? Original Output: To find the percentage of students who
enrolled in hip-hop dance, we need to use the information given in the
problem: 20% of the students enrolled in contemporary dance 25% of the
remaining students enrolled in jazz dance The rest of the students enrolled in
hip-hop dance Let’s break this down step by step: 1\. 20% of 20 students
enrolled in contemporary dance: 20% x 20 = 4 students 2\. 75% of the remaining
students (i.e., 20 - 4 = 16 students) enrolled in jazz dance: 75% x 16 = 12
students 3\. The rest of the students (i.e., 16 + 12 = 28 students) enrolled
in hip-hop dance. So, the percentage of students who enrolled in hip-hop dance
is: 28/50 = 56% Therefore, 56% of the entire class of 50 students enrolled in
hip-hop dance. Original Answer: 50 SelfControl: Let’s first find out how many
students enrolled in contemporary dance: 20 students * 20% = 4 students Next,
let’s find out how many students are left after contemporary dance students
have enrolled: 20 students - 4 students = 16 students Now, let’s find out how
many students enrolled in jazz dance: 16 students * 25% = 4 students So, there
are 4 students in contemporary dance, 4 students in jazz dance, and the rest
in hip-hop dance. Let’s find out how many students are in hip-hop dance: 20
students - 4 students (contemporary) - 4 students (jazz) = 12 students
Therefore, 12 students represent 60% of the entire students (since there are
20 students in total): 12 students / 20 students * 100% = 60% SelfControl
Answer: 60

Table 37: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 230): Two trains leave San Rafael at the same time. They
begin traveling westward, both traveling for 80 miles. The next day, they
travel northwards, covering 150 miles. What’s the distance covered by each
train in the two days? Original Output: To solve this problem, we need to use
the concept of relative speed, which is the speed at which the two trains are
moving relative to each other. Let’s assume that the speed of the first train
is x miles per hour, and the speed of the second train is x + 50 miles per
hour. On the first day, both trains travel westward for 80 miles, so the
distance covered by the first train is 80 miles. The distance covered by the
second train is also 80 miles, since they are traveling at the same speed. On
the second day, both trains travel northward for 150 miles. The distance
covered by the first train is 150 miles + 80 miles = 230 miles. The distance
covered by the second train is also 230 miles. So, the total distance covered
by each train in the two days is: First train: 80 + 230 = 310 miles Second
train: 80 + 230 = 310 miles Therefore, the distance covered by each train in
the two days is the same, which is 310 miles Original Answer: 310 SelfControl:
To find the total distance covered by each train in two days, we need to
calculate the distance covered during each leg of the journey and then add
them together. First, let’s find the distance covered during the westward
journey on the first day: Distance westward = 80 miles Next, let’s find the
distance covered during the northward journey on the second day: Since they
both traveled 150 miles northward, the distance covered during the second day
is: Distance northward = 150 miles Now, we can find the total distance covered
by each train in two days by adding the distances covered during each leg of
the journey: Total distance = Distance westward + Distance northward Total
distance = 80 miles + 150 miles Total distance = 230 miles So, each train
covers a total distance of 230 miles in two days. SelfControl Answer: 230

Table 38: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 109): Darrell and Allen’s ages are in the ratio of 7:11.
If their total age now is 162, calculate Allen’s age 10 years from now.
Original Output: Let’s assume Darrell’s age is D years and Allen’s age is A
years. Since the ratio of their ages is 7:11, we can write: D = 7A We are
given that the total age of Darrell and Allen now is 162 years. So, we can
write: D + A = 162 Now, you want to know Allen’s age 10 years from now. If
Allen’s age is A years now, his age 10 years from now will be: A + 10 = 162
Solving for A, we get: A = 152 So, Allen’s age 10 years from now is 152 years.
Original Answer: 152 SelfControl: Let Darrell’s age be 7x and Allen’s age be
11x. We know that their total age now is 162, so: 7x + 11x = 162 18x = 162 x =
9 So Darrell is currently $7*9=\langle\langle 7*9=63\rangle\rangle$ 63 years
old, and Allen is currently $11*9=\langle\langle 11*9=99\rangle\rangle$ 99
years old. To find Allen’s age 10 years from now, we first need to find his
current age and then add 10 years: Allen’s age now = 99 Allen’s age 10 years
from now = $99+10=\langle\langle 99+10=109\rangle\rangle$ 109\. Answer:
$None$. SelfControl Answer: 109

Table 39: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 23): Siobhan has 2 fewer jewels than Aaron. Aaron has 5
more jewels than half of Raymond’s jewels. If Raymond has 40 jewels, how many
jewels does Siobhan have? Original Output: Let’s start by using the
information that Siobhan has 2 fewer jewels than Aaron. So, if Aaron has x
jewels, Siobhan has x - 2 jewels. Next, we are told that Aaron has 5 more
jewels than half of Raymond’s jewels. So, if Raymond has x jewels, Aaron has x
+ 5 jewels. Now, we can use the information that Raymond has 40 jewels to find
out how many jewels Siobhan has. We know that Siobhan has x - 2 jewels, and we
also know that Aaron has x + 5 jewels. So, we can set up the equation:
Siobhan’s jewels = Aaron’s jewels - (Aaron’s jewels / 2) Simplifying the
equation, we get: Siobhan’s jewels = x - (x + 5) / 2 Solving for Siobhan’s
jewels, we get: Siobhan’s jewels = x - 10 So, Siobhan has 10 jewels. Original
Answer: 10 SelfControl: Let’s break this down step by step: 1\. Raymond has 40
jewels. 2\. Aaron has 5 more jewels than half of Raymond’s jewels, so Aaron
has 5 + (1/2 * 40) = 5 + 20 = 25 jewels. 3\. Siobhan has 2 fewer jewels than
Aaron, so Siobhan has 25 - 2 = 23 jewels. Therefore, Siobhan has 23 jewels.
SelfControl Answer: 23

Table 40: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 26): Grandma Jones baked 5 apple pies for the fireman’s
luncheon. She cut each pie into 8 pieces and set the five pies out on the
buffet table for the guests to serve themselves. At the end of the evening,
after the guests had taken and eaten their pieces of pie, there were 14 pieces
of pie remaining. How many pieces were taken by the guests? Original Output:
To solve this problem, we need to use the information given in the problem to
find out how many pieces of pie Grandma Jones baked. The problem states that
she baked 5 apple pies. Next, we need to find out how many pieces each pie was
cut into. The problem states that each pie was cut into 8 pieces. So, if
Grandma Jones baked 5 apple pies and each pie was cut into 8 pieces, then the
total number of pieces of pie that were available for the guests to serve
themselves is: 5 pies x 8 pieces per pie = 40 pieces of pie Now, we need to
find out how many pieces of pie were left over after the guests had taken
some. The problem states that there were 14 pieces of pie remaining after the
guests had taken their pieces. So, the number of pieces of pie taken by the
guests is: 14 pieces remaining - 40 pieces available = -26 pieces Since there
are no negative numbers of pie, the number of pieces of pie taken by the
guests is 0. Therefore, no pieces of pie were taken by the guests. Original
Answer: 0 SelfControl: To find out how many pieces were taken by the guests,
we need to subtract the number of remaining pieces from the total number of
pieces that were originally available. First, let’s find out how many pieces
were in each pie after Grandma Jones cut it into 8 pieces. Number of pieces
per pie = 8 Now, let’s find out how many total pieces were available after
Grandma Jones baked and cut up all 5 pies. Total pieces = Number of pieces per
pie * Number of pies Total pieces = 8 * 5 Total pieces = 40 At the end of the
evening, there were 14 pieces remaining. So, the number of pieces taken by the
guests is the difference between the total number of pieces and the number of
remaining pieces. Number of pieces taken = Total pieces - Remaining pieces
Number of pieces taken = 40 - 14 Number of pieces taken = 26 So, the guests
took 26 pieces of apple pie. SelfControl Answer: 26

Table 41: Example taken from Reasoning Benchmark.SelfControl

Query (Right Answer = 30): Richard lives in an apartment building with 15
floors. Each floor contains 8 units, and 3/4 of the building is occupied.
What’s the total number of unoccupied units In the building? Original Output:
To find the total number of unoccupied units in the building, we can use the
information given in the problem: The building has 15 floors. Each floor
contains 8 units. 3/4 of the building is occupied, which means that 3/4 of the
total number of units in the building are occupied. So, if we let x be the
total number of units in the building, we can set up the following equation: x
= 15 x 8/4 = 120 Now, we know that 3/4 of the building is occupied, so the
number of occupied units is: Occupied units = x / 4 = 120 / 4 = 30 Therefore,
the total number of unoccupied units in the building is: Unoccupied units = x
- Occupied units = 120 - 30 = 90. So, there are 90 unoccupied units in the
building. Original Answer: 90 SelfControl: Let’s break down the problem step
by step: 1\. The building has 15 floors. 2\. Each floor contains 8 units, so
the total number of units in the building is 15 * 8 = 120 units. 3\. Three-
quarters of the building is occupied, which means that 120 * 3/4 = 90 units
are occupied. 4\. To find the number of unoccupied units, we subtract the
number of occupied units from the total number of units: 120 - 90 = 30 units.
So, there are 30 unoccupied units in the building. SelfControl Answer: 30

###  Appendix F Pseudo-code

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

⬇

def get_suffix_grads(

wrapped_model,

query, response, suffix_list,

target, token_pos, token_neg

): -> Dict[FloatTensor]

final_grads = None

# We can simultaneously use multiple suffixes for control

for suffix in suffix_list:

with torch.enable_grad():

# The model controlled with suffix gradients

outputs = wrapped_model(

(query + response + suffix),

output_hidden_states=True,

)

# calculate the loss

loss = -get_suffix_score(query+response, suffix, ...)

for i in range(len(hidden_states)):

grads[i] = torch.autograd.grad(loss, hidden_states[i], retain_graph=True, \

allow_unused=True)[0]

norms[i] = torch.norm(grads[i], dim=-1, p=2, keepdim=True)

grads[i] = grads[i] / (norms[i] + 1e-12) # gradient clipping

final_grads += grads

# The gradients from each suffix are averaged

final_grads /= len(suffix_lis)

return grads

def iterative_controlled_generate(

query, suffix, target

max_iter # max iterations of control

): -> str

acc_grads = None

# we control on the hidden states at positions of query tokens

query_len = len(tokenizer.encode(query, add_special_tokens=False))

for iter in range(max_iter):

# sample a response with the current gradient (E-step)

wrapped_model = control_on_layers(acc_grads, query_len, ...)

# wrapped_model.suffix_decoding if using suffix decoding

response = wrapped_model.generate(query)

# gradient calculation (M-step)

grads = get_suffix_grads(query, response, suffix, target, ...)

# determine the step size

step_size = search_step_size(acc_grads, grads, ...)

if step_size == 0:

break

acc_grads += step_size * grads

# generate final response

wrapped_model = control_on_layers(acc_grads, ...)

final_response = wrapped_model.generate(query)

return final_response

Algorithm 1 Python Pseudocode of SelfControl (get_suffix_grads and
iterative_controlled_generate)

⬇

def search_step_size(

query, suffix, target

initial_score, # The initial suffix score of an response

acc_grads, # The gradients accumulated from previous iterations

grads, # suffix gradients from current step

max_iter, initial_step_size, scale_factor

score_threshold, # The threshold for a better step-size

): -> float # The final step size

current_step_size = initial_step_size

for i in range(max_iter):

temp_grads = acc_grads + current_step_size * grads

# get the model controlled by the gradients

wrapped_model = control_on_layers(

layer_ids=layer_ids,

wrapped_model=wrapped_model,

grads=temp_grads,

query_length=query_length, # only control on input query

)

response = wrapped_model.generate(prompt)

score = get_suffix_score(

prompt = query + response,

suffix = suffix,

... # model, tokenizer, target, tau and contrastive_pairs are the same

)

# return if current score is larger than the initial score by the threshold

if score - initial_score > score_threshold:

return current_step_size

current_step_size *= scale_factor

# no better score has been found

return 0

Algorithm 2 Python Pseudocode of SelfControl (search_step_size)

⬇

def get_suffix_score(

prompt, suffix, # prompt refers to [query, resposne]

model, tokenizer, target,

tau, # temperature

contrastive_pairs=["Yes", "No"] # The pair is (Yes, No) by default in our case

): -> float # suffix score

target_token = tokenizer.encode(target, add_special_tokens=False)[0]

token_pos = tokenizer.encode(contrastive_pairs[0],
add_special_tokens=False)[0]

token_neg = tokenizer.encode(contrastive_pairs[1],
add_special_tokens=False)[0]

# append suffix to the prompts

tokenized = tokenizer(prompt + suffix, return_tensors="pt", padding=True)

# get logits

with torch.no_grad():

outputs = model(**tokenized)

# calculate suffix score

last_logit = outputs.logits[:, -1, :]

logit_diff = last_logit[:, token_pos] - last_logit[:, token_neg]

return sigmoid(logit_diff / tau)

Algorithm 3 Python Pseudocode of SelfControl (get_suffix_score)


*The Gravity of Classical Fields

Response of ultralight fields to compact objects

Environmental Effects due to classical fields