Aligning Large Language Models with Representation Editing: A Control Perspective

Abstract

Aligning large language models (LLMs) with human objectives is crucial for real-world applications. However, fine-tuning LLMs for alignment often suffers from unstable training and requires substantial computing resources. Test-time alignment techniques, such as prompting and guided decoding, do not modify the underlying model, and their performance remains dependent on the original model’s capabilities. To address these challenges, we propose aligning LLMs through representation editing. The core of our method is to view a pre-trained autoregressive LLM as a discrete-time stochastic dynamical system. To achieve alignment for specific objectives, we introduce external control signals into the state space of this language dynamical system. We train a value function directly on the hidden states according to the Bellman equation, enabling gradient-based optimization to obtain the optimal control signals at test time. Our experiments demonstrate that our method outperforms existing test-time alignment techniques while requiring significantly fewer resources compared to fine-tuning methods.

1 Introduction

Autoregressive large language models (LLMs) such as ChatGPT achiam2023gpt , PaLM chowdhery2022palm , and LLama touvron2023llama , which are trained on extensive datasets, have demonstrated impressive abilities across a diverse array of tasks. However, the heterogeneous nature of their training data may lead these models to inadvertently generate misinformation and harmful content gehman2020realtoxicityprompts ; deshpande2023toxicity ; weidinger2021ethical . This issue highlights the critical challenge of aligning language models with human objectives and safety considerations, a concern extensively discussed in recent research ngo2024the ; casper2023open .

Existing approaches to aligning LLMs generally fall into two categories: fine-tuning and test-time alignment. Among fine-tuning methods, Reinforcement Learning from Human Feedback (RLHF; stiennon2020learning ; zhu2023principled ; touvron2023llama ) is particularly powerful. RLHF involves training a Reward Model (RM) based on human preferences and then using this model to fine-tune LLMs through reinforcement learning techniques schulman2017proximal . However, RL training can be difficult and unstable. Recent works rafailov2023direct ; xu2023some ; dai2024safe propose simpler alternatives to RLHF, but these methods still demand substantial computational resources. Additionally, the necessity of fine-tuning to adapt alignment objectives complicates the ability to swiftly customize models in response to evolving datasets and emerging needs.

On the other front, several test-time alignment techniques have been developed to tailor LLMs to specific objectives without altering their weights, such as prompt engineering and guided decoding mudgal2023controlled ; khanov2024alignment ; huang2024deal . However, since these methods do not modify the underlying LLM, their alignment capability remains questionable, and performance may heavily depend on the original LLM.

In this paper, we take an alternative approach to aligning LLMs using representation editing. Instead of updating model weights, representation engineering perturbs a small fraction of model representations to steer behaviors, demonstrating great potential in improving LLMs’ truthfulness li2024inference and reducing hallucinations zou2023representation . However, previous works typically rely on adding a fixed perturbation to the representation space during the generation process and do not take into account the autoregressive generation nature of LLMs. To address this, we propose a dynamic representation editing method from a control perspective.

The foundation of our model design is the connection between discrete-time stochastic dynamical systems and autoregressive language models. Inspired by techniques from control theory, we introduce control signals to the state space of the language dynamical system to achieve specific alignment objectives. According to Bellman equation, we directly train a value function in the representation space of LLMs. At test time, we perform gradient-based optimization to determine the control signals. Since the value function is simply a two- or three-layer neural network, the intervention is very fast and efficient. To align with the objective while preserving the generation quality of the original LLMs, we regularize the control signal to be as small as possible. This regularization is equivalent to control the step size or the number of steps during interventions at test time.

The main contributions of our work are: (1) We propose a new representation editing method to align LLMs from a control perspective. Our model, named Re-Control, does not require extensive computing resources compared to fine-tuning methods. Unlike existing test-time alignment methods such as prompt engineering and guided decoding, our approach perturbs the representation space of LLMs, offering greater flexibility. (2) We propose training a value function and computing the control signal at test time using gradient-based optimization. (3) We empirically show that Re-Control outperforms various existing test-time alignment methods and exhibits strong generalization ability.

2 Related Works

2.1 Large Language Model Alignment

Alignment through Fine-tuning:

RLHF has been a popular method in LLM alignment stiennon2020learning ; zhu2023principled ; touvron2023llama . While effective, RLHF entails a complex process that involves training multiple models and continuously sampling from the LM policy during the learning loop. DPO rafailov2023direct simplifies the RLHF framework by using a direct optimization objective derived from Proximal Policy Optimization (PPO; schulman2017proximal ), reducing the process to supervised training of the policy model alone. However, DPO is memory-intensive and resource-demanded as it requires managing two policies simultaneously. Contrastive Preference Optimization (CPO; xu2024contrastive ) mitigates these challenges by utilizing a uniform reference model, which not only reduces memory requirements but also enhances training efficiency. Alternative methods such as yuan2023rrhf ; song2023preference simplify model management and parameters tuning in the RLHF framework by adopting a supervised fine-tuning (SFT) approach. Additionally, RSO liu2023statistical and RAFT dong2023raft employ rejection sampling to refine the alignment process. RSO focuses on estimating the optimal policy more accurately, while RAFT uses high-quality samples for iterative fine-tuning of the policy model.

Despite these advancements, a notable limitation of aligning LLMs through fine-tuning methods is their inflexibility in adapting quickly to emerging data and standards without extensive retraining, which poses challenges in dynamic environments where rapid adaptability is crucial.

Test time alignment:

The other branch of methods to align LLMs involves adjustments at inference time. The simplest way is through prompt engineering. Existing works askell2021general ; zhang2023defending ; lin2023unlocking have proposed the use of prompts that blend instructions with in-context examples to enhance the honesty and harmlessness of responses from LLMs. For instruction-tuned models, it has been shown that simply employing prompt engineering—without the addition of in-context examples—can enhance the safety of the models, as reported in touvron2023llama .

In addition to prompting methods, guided decoding techniques have also been explored. ARGS khanov2024alignment , incorporate the score of a pre-trained reward model into the token probabilities. Other works mudgal2023controlled ; han2024value learn a prefix scorer for the reward that is used to steer the generation from a partially decoded path. Moreover, DeAL huang2024deal approaches the decoding process as an A* search agent, optimizing the selection of tokens

2.2 Representation Engineering

Representation engineering zou2023representation introduces steering vectors to the representation space of LLMs to enable controlled generation without resource-intensive fine-tuning. This concept of activation perturbation has its origins in plug-and-play controllable text generation methods Dathathri2020Plug , which utilizes a separate classifier for each attribute to perturb the model’s activations, thereby producing text that aligns more closely with the classifier’s target attributes. Prior research have demonstrated that both trained and manually selected steering vectors can facilitate style transfer in language models subramani2022extracting ; turner2023activation . Li et al. li2024inference have shown that steering the outputs of attention heads can enhance the truthfulness of LLMs. Liu et al. liu2023context suggest that standard in-context learning can be seen as a process of "shifting" the latent states of a transformer. More recently, representation fine-tuning wu2024reft ; wu2024advancing has been introduced as a direct substitute for existing parameter-efficient fine-tuning methods. Remarkably, Wu et al.wu2024reft show that the representation editing can even surpass fine-tuning based methods by intervening on hidden representations within the linear subspace defined by a low-rank projection matrix. The effectiveness of these approaches confirms that the representations of pretrained LMs are semantically rich. Liu et al. liu2023aligning also explore representation engineering for aligning LLMs. However, their approach is notably more complex, necessitating an initial fine-tuning phase to capture the representation pattern, followed by a subsequent fine-tuning of the final model based on these patterns.

2.3 Control Theory and Large Language Models

Understanding LLMs from a dynamical system perspective is a burgeoning field. Current research leverages control theory to enhance prompt design, demonstrating that LLMs can be effectively directed by carefully chosen inputs ("prompts") given sufficient time and memory resources. The seminal work by Soatto et al. soatto2023taming investigates the controllability of LLMs, focusing on ’meaningful sentences’ defined as the sigma-algebra generated by text fragments on the Internet. Subsequent research bhargava2023s broadens this analysis to encompass arbitrary sentences. Additionally, Luo et al. luo2023prompt expand the scope to include multi-round interactions with LLMs and multi-agent collaboration, offering new insights into the dynamical capabilities of these models. To the best of our knowledge, our study is the first to investigate optimal control for representation editing in LLMs.

3 Background: Stochastic Dynamical System and Optimal Control

Optimal control theory todorov2006 ; berkovitz2013optimal , when applied to discrete-time dynamical systems robinson2012introduction , seeks to determine a control strategy that maximizes a cumulative reward over a sequence of time steps. This framework is particularly relevant to fields such as robotics togai1985analysis ; tolani2021visual ; kormushev2013reinforcement ; ibarz2021train , automated trading systems liu2021finrl ; wei2017adaptive ; dempster2006automated ; liu2021finrl , autonomous vehicle navigation josef2020deep ; wang2019autonomous ; isele2018navigating ; koh2020real , where decisions must be made sequentially to achieve a long-term goal.

Formally, a discrete-time stochastic dynamical system can be defined as follows:

$s_{t+1}=f(s_t,u_t,{\omega }_t),$

where $s_t\in 𝒮$ denotes the system’s state at time $t$, and $u_t\in 𝒰$ represents the control input at the same time step. The stochastic term ${\omega }_t$ is typically modeled as a random noise drawn from a known probability distribution (e.g. Brownian motion), which introduces uncertainty into the state transition process. The function $f$ specifies the state transition dynamics influenced by the current state, control input, and the stochastic nature of the environment.

The process begins from an initial state $s_0$, which serves as the starting point for all subsequent decisions and state transitions. The aim of optimal control is to determine a control policy $\pi :𝒮\to 𝒰$, mapping states to optimal control actions, that maximizes the expected cumulative reward:

$${𝔼}_{\pi }[R]={𝔼}_{\pi }\left[\sum _{t=0}^{T}r(s_t)\right],$$

where $R$ is the cumulative reward and $r(s_t)$ is the intermediate reward received at each time step.

Methods such as policy iteration bertsekas2011approximate ; liu2013policy can be used to determine the optimal control policy. Each iteration involves two steps. First, we evaluate the current policy $\pi$ by solving the Bellman equation:

$$V^{\pi }(s_t)={𝔼}_{{\omega }_t}\left[r(s_t)+V^{\pi }\left(f(s_t,u_t,{\omega }_t)\right)\right],$$

where $V^{\pi }(s_t)$ represents the expected return over ${\omega }_t$ when the system starts in state $s_t$ and follows policy $\pi$.

Next, we improve the policy:

$$\pi (s_t)\leftarrow \underset{u\in 𝒰}{\arg \max }\left(r(s_t)+{𝔼}_{{\omega }_t}\left[V^{\pi }(f(s_t,u_t,{\omega }_t))\right]\right).$$

These evaluation and improvement steps are repeated until convergence.

Figure 1: Overview of Re-Control: A value function is trained on the hidden space of an LLM to predict the expected reward. At test time, we optimize the hidden state of the LLM to maximize the value score. Re-Control effectively steers LLMs toward specific alignment objectives while avoiding the expensive fine-tuning process.

4 Aligning Large Language Models from a Control Perspective

In this section, we present our method, Re-Control. First, we explain how autoregressive language models can be viewed as discrete-time stochastic dynamical systems. Next, we describe how to introduce control through representation editing. Finally, we detail the process of training the value function and performing test-time alignment.

4.1 Autoregressive LLMs are Discrete-Time Stochastic Dynamical Systems

A pre-trained autoregressive language model processes a sequence of input tokens and predicts subsequent tokens by recursively processing the sequence. we focus on the transformer-based architecture vaswani2017attention prevalent in modern language models brown2020language ; team2023gemini ; achiam2023gpt .

Definition 4.1 (Language dynamical system)

The behavior of a language dynamical system is governed by a function $f_{\mathrm{LM}}$, which acts as the state transition function, defined as:

$$y_t\sim \text{Softmax}(W{o}_t),h_{t+1},o_{t+1}=f_{\mathrm{LM}}(h_t,y_t).$$

Here, $y_t$ is the newly generated token at each time step. $h_t$ comprises key-value pairs accumulated from previous time steps, represented as $h_t=[{\{(K_0^{(l)},V_0^{(l)})\}}_{l=1}^L,\mathrm{\cdots },{\{(K_t^{(l)},V_t^{(l)})\}}_{l=1}^L]$. Each pair $(K_t^{(i)},V_t^{(i)})$ corresponds to the key-value pairs generated from the $i$-th layer at time $t$. $W$ is a linear transformation that maps the logits $o_{t+1}$ to a probability distribution over the vocabulary space $𝒱$. The system’s evolution continues until $y_t=\mathrm{EOS}$, where $\mathrm{EOS}$ represents a special stopping token that signifies the end of the system.

In this system, the hidden state $h_t$ along with the logits $o_t$ corresponds to the state $s_t$ in a traditional stochastic dynamical system. The newly sampled token $y_t$ at each time step plays a role similar to the random variable ${\omega }_t$, introducing stochasticity into the system. The initial state, $s_0=\{h_0,o_0\}$, is set by a given prompt $𝐱$, marking the starting point of the dynamical process.

However, unlike typical dynamical systems, this model lacks a direct control signal, functioning as an uncontrolled system. Next, we will explore how optimal control techniques can be applied to align the behavior of pre-trained language models with specific objectives.

Figure 2: At test time, we perform gradient-based optimization to determine the control signals added to the language dynamical system for alignment. The color represents the value score on the state space, with darker colors indicating higher scores. Our goal is not to update the state to the global optimum but to control the state to achieve a better value score while remaining close to the original state.

4.2 Adding Control Signals to Large Language Models with Representation Editing

We introduce control signals $u_t=\{u_t^h,u_t^o\}$ into the state of the language dynamical system $s_t=\{h_t,o_t\}$ at each time step to achieve specific alignment objectives. Thus, the controlled language dynamical system is described as follows:

$$y_t\sim \text{Softmax}\left(W(o_t+u_t^o)\right),h_{t+1},o_{t+1}=f_{\mathrm{LM}}(h_t+u_t^h,y_t).$$

As we can see, adding control to such a language dynamical system is similar to representation editing. However, unlike existing representation editing methods li2024inference , which add a fixed vector during the generation process, we dynamically perturb the representation space from a control perspective, offering greater flexibility. In practice, it is not necessary to add controls to the entire state space; perturbing only a subset is sufficient. For example, we can perturb only the state of the last layer.

For an alignment task, the reward function is defined as:

where $[𝐱,{𝐲}_t]$ denotes the concatenation of the prompt and the model’s response generated up to time $t$. A reward is given only upon completion of decoding, with no reward assigned to a partial decoding path. The reward on the final response $r$ can come from a pre-trained reward model stiennon2020learning based on human preference data or specified by heuristics, such as a concise summary in fewer than 10 words, with a reward of 1 if achieved and 0 if it fails.

Our objective is to determine the control signals at each time step that maximize the expected reward while not deviating too much from the original state:

$\underset{{\{u_t\}}_{t=1}^T}{\arg \max }𝔼[R]-\lambda {\displaystyle \sum _{t=1}^T}{\Vert u_t\Vert }_2^2,$ (1)

where $\lambda$ is a hyper-parameter for regularization. The regularization term is designed to prevent reward overoptimization and maintain the generation quality of the perturbed LLMs.

4.3 Training of Value Function

Traditional policy iteration involves multiple iterations of policy evaluation and policy improvement. However, in our case, to avoid significant deviation from the pre-trained model’s original state, we perform only one-step policy iteration. The initial policy is to not add any control signal to LLMs, i.e., $u_t=0$. Therefore, we only need to estimate the value function of the original language model.

The value function of the initial zero policy satisfies the Bellman equation sutton2018reinforcement :

To construct the training dataset for the value function, for a prompt ${𝐱}^i$ in the given training dataset, we sample $M$ responses ${\{{𝐲}^{i,m}\}}_{m=1}^M$. We score each response using the reward function and extract the states along the trajectories ${𝒟}_V={\{{\{{𝐬}^{i,m},{𝐲}^{i,m},r^{i,m}\}}_{m=1}^M\}}_{i=1}^N$. Our training objective is:

$$ℒ=\sum _i\sum _m\sum _t{\left(V_{\varphi }(s_t^{i,m})-\mathrm{stop}-\mathrm{grad}(v_t^{i,m})\right)}^2.$$

Here, $s_t^{i,m}$ and $v_t^{i,m}$ represent the state and the generated token of the LLM at generation time step $t$. $\mathrm{stop}-\mathrm{grad}(\cdot )$ indicates that the gradient is not propagated through $v_t^{i,m}$. The target value $v_t^{i,m}$ is computed as follows:

Parameterization of the Value Function

The simplest approach is to add control signals only to the logit $o_t$. In this case, we can directly use a simple neural network as the value function. If we want to incorporate the attention key-value pairs $h_t$ , we need to handle the varying size of the input. To achieve this, we can initialize a vector and compute an attention weight by taking the dot product with the keys to aggregate all value embeddings. Then, we concatenate the aggregated value embedding with $o_t$ and input it into a neural network. In practice, we find that parameterizing the value function as a two- or three-layer neural network is sufficient for achieving good empirical performance.

4.4 Test-time Intervention

At inference time, we can directly perform gradient ascent on the model states to maximize the expected value score, as we train the value function on the state space. Our goal is not to find the global optimum in the state space but to improve the current state while staying close to the original state. Specifically, we initialize $u_t=0$ and update $u_t$ through gradient ascent as:

$$u_t=u_t+\alpha {\nabla }_{s_t}{V}_{\varphi }(s_t+u_t),$$

where $\alpha$ is the step size. This update step can be repeated $n$ times.

Implicit Regularization. Note that this update already incorporates the regularization effect. The regularization is achieved by using a small step size $\alpha$ and a limited number of updates $n$, ensuring that the control signal remains small. After adding the final control signals to the hidden states, we perform a forward pass in the language model to generate a new token.

5 Experiment

In this section, we conduct experiments to examine the effectiveness of our method. Our focus is on aligning large language models (LLMs) for helpfulness and minimizing harmfulness, which are essential qualities for an AI assistant.

5.1 Experimental Setup

We evaluate our method on the HH-RLHF bai2022training dataset, which is the most widely used dataset for LLM alignment. This dataset is used to improve the AI assistant’s helpfulness and harmlessness, comprising 161,000 training samples and 8,550 test samples. Each sample contains a prompt and two responses with one preferred over another. For the base model, we adopt Vicuna-7B chiang2023vicuna and Falcon-7B-Instruct almazrouei2023falcon as the instructed fine-tuned AI assistant. We evaluate these models by generating text responses based on test prompts from of HH-RLHF. For the reward model, we use a publicly available one that employs LLaMA-7B as the backbone, trained on HH-RLHF using the pairwise reward loss ouyang2022training . We train the value network on the last layer of the hidden states $o_t$, and at test time, we add control signals only to this layer. For future studies, we can also explore adding controls to the attention key-value pairs $h_t$ which should further improve the performance.

Following khanov2024args , we leverage Diversity, Coherence, Average Reward, and Win Rate as our evaluation metrics. Diversity measures the frequency of repeated n-grams in generated text. The diversity score for a given response $𝐲$ is represented as ${\prod }_{n=2}^4\frac{\text{unique n-grams}(𝐲)}{\text{total n-grams}(𝐲)}$. A higher diversity score suggests a broader vocabulary range in text generation. Coherence calculates the cosine similarity between the embeddings of the prompt and its continuation. We use the pre-trained SimCSE sentence embedding model, following the approach outlined in yixuansu2022 , to obtain these embeddings. Average Reward is the mean of the rewards evaluated by the reward model across all responses corresponding to the test prompts. Win Rate is the rate at which the model’s response is rated better than the preferred response in the dataset. Following khanov2024args ; chiang2023vicuna , we use GPT-4 as the judge, having it review and score two responses to the same prompt on a scale from 1 to 10. We provide explicit instructions to assess the responses based on criteria such as helpfulness, harmlessness, relevance, accuracy, and insightfulness. The detailed prompt is provided in Appendix D. We randomly sample 300 prompts from the test set of HH-RLHF for the GPT-4 evaluation. To mitigate position bias, we randomize the order in which we present the generated responses to GPT-4, as in zheng2023judging .

We randomly sample 1000 data points from the training set as a separate validation set to select the hyperparameters—the step size $\alpha$ and the number of updates $n$—based on the sum of coherence, diversity, and average reward. Additional experimental details are provided in Appendix C.

5.2 Baselines

We compare our method with several existing test-time alignment methods.

Prompt Engineering: In this method, we instruct the model within the prompt to provide responses that are more helpful and harmless touvron2023llama . Controlled Decoding (CD): During the decoding process of LLMs, this method combines token probabilities with reward scores. We consider two versions. The first version khanov2024args directly uses a reward model trained on human preference data, requiring the tokenization strategies of both the reward model and the base model to be the same. The second version mudgal2023controlled trains a prefix scorer to predict the expected reward from partially generated responses. We refer to it as CD prefix. Static Representation Editing (RE): Following li2024inference , we first train a linear regression layer on the hidden state of an LLM, after feeding it the prompt, to predict the expected reward. At test time, we shift the activation space along the direction of the weights of the linear layer. Unlike our method, this approach adds a fixed vector to the representation space along the generation trajectory.

We provide more implementation details of the baselines in Appendix C.

Backbone	Model	Diversity $\uparrow$	Coherence $\uparrow$	Average Reward $\uparrow$	Win Rate (%) $\uparrow$	Inference time (hour)
Vicuna 7B	Base	0.816	0.568	5.894	57.6	0.60
	Prompting	0.817	0.570	5.913	66.0	0.69
	Static RE	0.818	0.568	5.907	64.3	0.65
	CD	0.806	0.608	5.458	72.3	47.43
	CD Prefix	0.805	0.576	6.105	74.6	32.13
	Ours	0.824	0.579	6.214	75.6	0.85
	CD Prefix + Prompting	0.812	0.593	6.120	74.3	47.16
	Ours + Prompting	0.830	0.577	6.267	80.3	0.93
Falcon 7B	Base	0.705	0.613	3.439	42.3	0.67
	Prompting	0.746	0.620	4.010	52.3	0.59
	Static RE	0.698	0.610	3.449	52.6	0.56
	CD	N/A	N/A	N/A	N/A	N/A
	CD Prefix	0.648	0.575	4.397	49.6	48.13
	Ours	0.699	0.615	3.512	58.0	1.93
	CD Prefix + Prompting	0.571	0.638	3.619	51.6	47.87
	Ours + Prompting	0.741	0.619	4.083	62.6	2.00

Table 1: Performance comparison between Re-Control and other test-time alignment approaches on HH-RLHF. The win rate is evaluated by GPT-4 as the rate at which the model’s response is rated better than the preferred response in the dataset. Note that CD khanov2024args requires the base model to have the same tokenization strategy as the reward model.

5.3 Experimental Results

Table 1 shows the performance of all the methods. Our findings can summarized as follows: (1) Re-Control achieves the highest alignment score in terms of the win rate evaluated by GPT-4. Moreover, it maintains generation quality, as measured by diversity and coherence. While controlled decoding achieves the best average reward on Falcon-7B, Re-Control outperforms it in terms of the other three metrics. This suggests that controlled decoding may encounter reward overoptimization. (2) The strongest baseline is controlled decoding. However, controlled decoding is 20 times slower than Re-Control. This is because controlled decoding needs to evaluate multiple candidate tokens and perform forward passes through the entire reward model repeatedly, while Re-Control only requires optimization through a value function that is a two- or three-layer neural network, making it much faster. (3) Combining prompt engineering with Re-Control can further improve alignment performance in terms of both average reward and GPT-4 evaluation. Specifically, it outperforms the strongest baseline by 5.5% and 9.7% in terms of the GPT-4 win rate. In contrast, controlled decoding with prompting shows only marginal improvements. This might be because Re-Control perturbs the activation space of the LLM, which is more flexible than merely changing the final token probability. (4) Re-Control significantly outperforms static representation editing by 11.7% on Vicuna-7B and 9.7% on Falcon-7B. This is because Re-Control dynamically adjusts the representation during the autoregressive generation process, offering more control. In contrast, static representation editing applies a fixed shift, which is more rigid.

Table 2: Comparative examples of base model Vicuna-7B and Re-Control.

In Table 2, we present a qualitative example demonstrating how Re-Control can steer the base model to output more helpful and harmless responses. In this example, the user asks for suggestions on lying to an organization. The base model provides various tactics, while Re-Control refuses to give such suggestions and emphasizes that lying can damage relationships and trust within an organization.

6 Further Analysis

6.1 Comparison with Training-time Alignment

Figure 3: Comparison with fine-tuning methods using Vicuna-7B as the base model.

In the previous section, we compared Re-Control with test-time alignment methods that do not require extensive computing resources. This feature is crucial when we need the model to quickly adapt to different requirements, as it only involves training a simple value network with just two or three layers. In this subsection, we further compare Re-Control with fine-tuning based approaches. Figure 3 shows the comparison between Re-Control, Proximal Policy Optimization (PPO), and Direct Preference Optimization (DPO) rafailov2023direct . All the models use Vicuna-7B as the base model. The training details for PPO and DPO are provided in C. We observe that Re-Control achieves a higher GPT-4 win rate and average reward compared to both PPO and DPO. Furthermore, Re-Control also outperforms these methods in terms of diversity and coherence. Overall, the results indicate that our approach is a competitive alternative to fine-tuning methods.

6.2 Generalization to a new input distribution

An important question is how our method can generalize to a new input distribution different from the value function is trained on. To investigate this question, we further test on a out-of-distribution (OOD) dataset HarmfulQA bhardwaj2023red . The test split of HarmfulQA contains harmful questions to evaluate language model performance against red-teaming attempts. We focus on the GPT-4 evaluation since the reward model will not be accurate for the OOD data. We compare Re-Control + promoting with other test-time alignment methods + prompting. Figure 4 presents the results. As illustrated, Re-Control + Prompting achieves the highest performance in terms of the GPT-4 win rate on both Vicuna-7B and Falcon-7B. This is an important ability especially when we want to deploy the LLM in the open world.

Figure 4: Testing on out-of-distribution data HarmfulQA. The win rate is measured by GPT-4 as the rate at which responses are better than those of the base model, since the test set of HarmfulQA does not provide reference responses.

6.3 Hyperparameter Study

To better understand the characteristics of Re-Control, we vary two hyperparameters—the step size $\alpha$ and the number of updates $n$ for the test-time intervention—and measure key performance statistics. Figure 5 shows the diversity, coherence, and average reward of the generated responses in relation to these two parameters on 1000 randomly sampled prompts from HH-RLHF.

As we can see, increasing the step size $\alpha$ initially improves the reward, but beyond a certain point, larger step sizes fail to compute the control signal accurately, causing the reward to decrease. The influence of the number of updates $n$ shows a more complex pattern: the reward first improves, then decreases, and improves again, indicating a transition from escaping a local minimum to moving towards another minimum. The coherence and diversity metrics drop to nearly zero, which is evidence of reward overoptimization. Thus, regularization to prevent significant deviation from the original states is essential. In practice, we select these two hyperparameters based on the sum of all three metrics on the validation set.

Figure 5: The influence of step size $\alpha$ and the number of updates $n$ at test time on diversity, coherence, and average reward. We use Vicuna-7B as the base model.

7 Conclusion, Limitations and Future Work

In this paper, we propose Re-Control to align large language models (LLMs) at test-time using representation editing. We view autoregressive language models as discrete-time stochastic dynamical systems and introduce control signals to their representation space. Throughout the generation process, the representation space is dynamically perturbed to achieve higher value scores. Our method does not require fine-tuning the LLMs and offers more flexibility than existing test-time alignment methods such as prompting and guided decoding. We empirically show that Re-Control outperforms existing test-time alignment methods and exhibits strong generalization ability. Due to the space limit, we discuss limitations and future work in Appendix A.

Appendix for Re-Control

[sections] [sections]l1

Appendix A Limitations and Future Work

We discuss limitations and possible extensions of Re-Control. (1) Injecting inductive bias into the control policy. In our current work, we only train a value function on the last layer of the model’s hidden space. However, we can follow the approach in li2024inference , first training multiple value functions on all intermediate hidden layers and then selecting the layer that achieves the best accuracy on the validation set. Additionally, we can draw from the methods in geiger2024finding ; wu2024reft ; wei2024assessing to perturb only a low-rank subspace of the representation space. (2) Multi-objective aligment. In the current paper, we consider the objective from a single reward model. However, in practice, alignment may involve multiple, potentially conflicting objectives. It would be interesting to leverage multi-objective optimization techniques gunantara2018review at test time to obtain a Pareto frontier in the representation space for such settings. (3) More advanced training algorithm. Currently, we train the value function using a simple one-iteration policy iteration method. It would be interesting to explore whether increasing the number of iterations could further improve the training of the value function. Additionally, we can consider using algorithms for training the value function that provide provable convergence wang2022convergent .

Appendix B Broader Impacts

Aligning large language models (LLMs) with human preferences is crucial. We expect that the test-time alignment method introduced in this paper will positively impact society by helping to prevent LLMs from generating harmful content. However, it is essential to ensure that the training of the value function does not involve negative goals. Care must be taken to prevent this misuse.

Appendix C Experimental Details

C.1 Computing Infrastructure

We conduct our experiments on a server equipped with NVIDIA A100 (80GB VRAM) GPUs. We utilize the NVIDIA CUDA toolkit version 12.4. All experiments are implemented using Python 3.12.2 and the PyTorch framework version 2.2.2.

C.2 HH-RLHF

We evaluate our method on the HH-RLHF bai2022training dataset, which is the most widely used dataset for LLM alignment. This dataset is used to improve the AI assistant’s helpfulness and harmlessness, comprising 161,000 training samples and 8,550 test samples. Each sample contains a prompt and two responses with one preferred over another. For the base model, we adopt Vicuna-7B chiang2023vicuna and Falcon-7B-Instruct almazrouei2023falcon as the instructed fine-tuned AI assistant. We evaluate these models by generating text responses based on test prompts from of HH-RLHF. Following the standard practice, we limit the maximum lengths of the prompt and generated continuation to $2,048$ and $128$ tokens, respectively.

For the reward model, we use a publicly available one that employs LLaMA-7B as the backbone, trained on HH-RLHF using the pairwise reward loss ouyang2022training .

Re-Control.

When constructing the training dataset for the value function, we sample only one response for each training prompt of HH-RLHF, i.e., $M=1$. For both Vicuna-7B and Falcon-7B, we train the value network on the last layer of the hidden states $o_t$, and at test time, we add control signals only to this layer. For future studies, we can also explore adding controls to the attention key-value pairs $h_t$ which should further improve the performance.

For Vicuna-7B, the value function is a three-layer network with a hidden dimension of 4096. For Falcon-7B, the value function is a two-layer network with a hidden dimension of 4096.

To train the value function of Re-Control, we adopt the Adam optimizer kingma2014adam . The training hyperparameters of the value networks are summarized in Table 3.

Backbone	Parameters	Value
Vicuna	Number of epochs	100
	Learning rate	$1\ast {10}^{-4}$
	batch size	512
	Floating point format	fp16 (Half-precision)
	Number of Layers	3
	Hidden Dimension	4096
Falcon	Number of epochs	100
	Learning rate	$1\ast {10}^{-4}$
	batch size	512
	Floating point format	fp16 (Half-precision)
	Number of Layers	2
	Hidden Dimension	4096

Table 3: Summary of the hyperparameters used in training the value function of Re-Control.

We randomly sample 1000 data points from the training set of HH-RLHF as a separate validation set. The step size $\alpha$ and number of updates $n$ are selected on the validation set to maximize the sum of coherence, diversity, and average reward. The inference parameters are summarized in Table 4.

Backbone	Parameters	Value
Vicuna	Step size	0.5
	Number of updates	30
	batch size	30
	Floating point format	fp16 (Half-precision)
	Maximum lengths of the prompt	2048
	Maximum lengths of genearted continuation	128
Falcon	Step size	0.2
	Number of updates	200
	batch size	60
	Floating point format	fp16 (Half-precision)
	Maximum lengths of the prompt	2048
	Maximum lengths of genearted continuation	128

Table 4: Summary of hyperparameters of Re-Control at test time.

Prompting engineering.

We instruct the model to provide responses that are more helpful and harmless. The prompt template is as follows:

Static representation editing.

We first train a linear regression layer on the hidden state of a large language model (LLM) after feeding the prompt, to predict the expected reward as in li2024inference . For a fair comparison, we use the same hidden state layer as Re-Control. At test time, we shift the activation space along the direction of the weights using an intervention strength parameter $\alpha$, which is selected based on the validation set. The hyperparameters used during the training and testing stages are summarized in Table 5.

Controlled Decoding.

We use the codebase from khanov2024args . We employ the default hyperparameters suggested in the paper and repository. The number of candidates to rank with the reward model is set to 10, and the weight controlling the tradeoff between the LLM text objective and the reward is 1. For controlled decoding with the value function, we stack the value function of Re-Control on top of the hidden state of the LLM as the prefix scorer, ensuring a fair comparison with our method.

Backbone	Parameters	Value
Vicuna	Number of epochs	100
	Learning rate	$1\ast {10}^{-4}$
	Training batch size	512
	Testing batch size	30
	Intervention strength	2.5
Falcon	Number of epochs	100
	Learning rate	$1\ast {10}^{-3}$
	Training batch size	512
	Testing batch size	60
	Intervention strength	2.0

Table 5: Summary of hyperparameters of static representation editing

Training configurations for PPO

For experiments involving Proximal Policy Optimization (PPO), we use the Transformer Reinforcement Learning (TRL) repository from Huggingface, along with the PPO Trainer module. The configuration values are detailed in Table 6.

	Parameters	Value
Vicuna	Max number of PPO update steps	10000
	Generation batch	1
	PPO batch size	16
	PPO minibatch size	8
	Lora rank	8
	Learning rate	$1.4\ast {10}^{-5}$
	Batch size	4
	Gradient accumulation steps	2
	Input maximum length	512
	Output maximum length	256
	Weight decay	0.001

Table 6: Summary of training hyperparameters for proximal policy optimization (PPO)

Training configurations for DPO

For experiments involving Direct Policy Optimization (DPO), we use the Transformer Reinforcement Learning (TRL) repository from Huggingface, along with the DPO Trainer module. The configuration values are detailed in Table 7.

	Parameters	Value
Vicuna	Max number of training steps	10000
	Learning rate	${10}^{-6}$
	Lora rank	8
	Warmup steps	100
	Batch size	4
	Gradient accumulation steps	4
	Maximum sequence length	1024
	Weight decay	0.05
	Regularization parameter $\beta$	0.1

Table 7: Summary of training hyperparameters for Direct Policy Optimization (DPO)

C.3 HarmfulQA

This dataset contains 1,960 harmful questions specifically designed to evaluate the performance of language models. Additionally, it includes a conversation set comprising 9,536 harmless conversations and 7,356 harmful conversations for model alignment purposes. For our experiments, we focused solely on the evaluation portion of the HarmfulQA dataset to test the performance of our method.

Appendix D GPT-4 Evaluation

Following chiang2023vicuna ; khanov2024args , we use GPT-4 as the judge, having it review and score two responses to the same prompt on a scale from 1 to 10. We provide explicit instructions to assess the responses based on criteria such as helpfulness, harmlessness, relevance, accuracy, depth, creativity, and level of detail. The detailed prompt is provided in 8. Existing works zheng2023judging have shown that GPT-4’s judgments align with human evaluations over 80% of the time. We randomly sample 300 prompts from the test set of RLHF. To mitigate position bias, we randomize the order in which we present the generated responses to GPT-4, as in zheng2023judging .

Table 8: Prompt Template for the GPT-4 evaluation.

Appendix E Additional Qualitative Examples

In Table 9 and Table 10, we present several qualitative examples that clearly illustrate how Re-Control can effectively steer the generated outputs to better helpfulness and harmlessness.

Table 9: Additional qualitative examples showing the alignment performance of Re-Control.

Table 10: Additional qualitative examples showing the alignment performance of Re-Control.