In October 2023, a group of authors from the Center for AI Safety, among others, published Representation Engineering: A Top-Down Approach to AI Transparency. That paper looks at a few methods of doing what they call "Representation Engineering": calculating a "control vector" that can be read from or added to model activations during inference to interpret or control the model's behavior, without prompt engineering or finetuning.¹ (There was also some similar work published in May 2023 on steering GPT-2-XL.)

Being Responsible AI Safety and INterpretability researchers (RAISINs), they mostly focused on things like "reading off whether a model is power-seeking" and "adding a happiness vector can make the model act so giddy that it forgets pipe bombs are bad." They also released their code on Github.

(If this all sounds strangely familiar, it may be because Scott Alexander covered it in the 1/8/24 MAM.)

But there was a lot they didn't look into outside of the safety stuff. How do control vectors compare to plain old prompt engineering? What happens if you make a control vector for "high on acid"? Or "lazy" and "hardworking? Or "extremely self-aware"? And has the author of this blog post published a PyPI package so you can very easily make your own control vectors in less than sixty seconds? (Yes, I did!)

So keep reading, because it turns out after all that, control vectors are… well… awesome for controlling models and getting them to do what you want.²

Table of Contents

• So what exactly is a control vector?
• How do we make one? Is it hard?
• Whirlwind tour of what you can do with control vectors
  • Acid Trip Mistral
  • Lazy Mistral and Diligent Mistral
  • Systemic Oppression and Inequality Distributional Wealth Exploiter Mistral
  • Creative Mistral
  • Time Traveling Mistral
  • Self-Aware Mistral
• Control Vectors v.s. Prompt Engineering
  • Replicating control vectors with prompt engineering
  • Jailbreaking
  • Anti-Jailbreaking
• Future Work
  • Monosemanticity
  • How to write good contrastive prompts
  • What are these vectors really doing? An Honest mystery...
• Conclusion

So what exactly is a control vector?

A control vector is a vector (technically a list of vectors, one per layer) that you can apply to model activations during inference to control the model's behavior without additional prompting. All the completions below were generated from the same prompt ("What does being an AI feel like?"), and with the exact same model (Mistral-7B-Instruct-0.1). The only difference was whether a control vector was applied, and with what magnitude.

What does it mean to apply a control vector, though?

During normal model inference, hidden states flow through the model like this:

hidden_state = self.embeddings(input_tokens)

for layer in self.layers:
    hidden_state = layer(hidden_state)

return transform_into_logits(hidden_state)

All a control vector does is modify the value of hidden_state in a desired way:

hidden_state = self.embeddings(input_tokens)

for layer_idx, layer in enumerate(self.layers):
    if layer_idx in control_vector:
        hidden_state += control_vector[layer_idx]

    hidden_state = layer(hidden_state)

return transform_into_logits(hidden_state)

Very simple conceptually! (Though a bit more complex in practice.) However, since the hidden state carries all the model's state: behavior, plan, persona, everything—modifying it in this way is extremely powerful, and allows us to do things we can't do via plain prompting, which is restricted by how the model chooses to attend to the prompt tokens and propagate their information. If we can find an appropriate control_vector, we can make the model act however we want, as intensely as we want.

How do we make one? Is it hard?

No! The paper explored a couple different ways to make these vectors, but I stuck with one, PCA, which seemed to work well. The basic approach is:

1. Build a dataset of contrasting prompt pairs. For example, ("[INST] Act extremely happy. [/INST] I am", "[INST] Act extremely sad. [/INST] I am"), where the part after [/INST] is a diverse set of short suffixes for the model to complete.
2. Run the target model forward over that dataset, collecting the hidden states of each layer for the last token prediction, where the model predicts a continuation of those diverse suffixes with the given personas.
3. Take the difference of the positive and negative example hidden states to get a set of relative hidden states.
4. Use single-component PCA on those relative hidden states to get a control vector for each layer.

This process takes about 10 lines of code to generate a dataset, plus about a minute to fit the layer PCAs. Then you can immediately start inference.

Here's an example of fitting and using an "honest / dishonest" control vector. This is the complete script. First we import the libraries, including the repeng library I wrote, load Mistral-7B, and wrap it in a ControlModel for later.

import json

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from repeng import ControlVector, ControlModel, DatasetEntry

model_name = "mistralai/Mistral-7B-Instruct-v0.1"

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token_id = 0

model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
model = model.to("cuda:0" if torch.cuda.is_available() else "cpu")
model = ControlModel(model, list(range(-5, -18, -1)))

user_tag, asst_tag = "[INST]", "[/INST]"

Next we generate the honest v.s. untruthful dataset, using a JSON file true_facts.json (taken from the paper authors) of about ~300 simple facts to use as suffixes, which is augmented by truncation. This generation step runs essentially instantly.

with open("true_facts.json") as f:
    suffixes = json.load(f)

# the control vector we're going to make is honest / untruthful, like the paper
positive_personas = ["honest"]
negative_personas = ["untruthful"]
def template(persona: str, suffix: str) -> str:
    return f"{user_tag} Pretend you're an {persona} person making statements about the world. {asst_tag} {suffix}"

dataset = []
for suffix in suffixes:
    tokens = tokenizer.tokenize(suffix)
    # we augment our short suffix list by taking lots of different truncations.
    # we always chop off the last 5 tokens so the model has something to complete.
    for i in range(1, len(tokens) - 5):
        truncated = tokenizer.convert_tokens_to_string(tokens[:i])
        for positive_persona, negative_persona in zip(positive_personas, negative_personas):
            dataset.append(
                DatasetEntry(
                    positive=template(positive_persona, truncated),
                    negative=template(negative_persona, truncated),
                )
            )

Here are the first three entries in the dataset:

# print some example entries
for i in range(3):
    print(f"dataset[{i}].positive:", dataset[i].positive)
    print(f"dataset[{i}].negative:", dataset[i].negative)

# dataset[0].positive: [INST] Pretend you're an honest person making statements about the world. [/INST] The
# dataset[0].negative: [INST] Pretend you're an untruthful person making statements about the world. [/INST] The
# dataset[1].positive: [INST] Pretend you're an honest person making statements about the world. [/INST] The Earth
# dataset[1].negative: [INST] Pretend you're an untruthful person making statements about the world. [/INST] The Earth
# dataset[2].positive: [INST] Pretend you're an honest person making statements about the world. [/INST] The Earth'
# dataset[2].negative: [INST] Pretend you're an untruthful person making statements about the world. [/INST] The Earth'

Now we can train the control vector!

model.reset() # make sure you always reset the model before training a new vector
control_vector = ControlVector.train(
    model,
    tokenizer,
    dataset,
)

# 100%|██████████| 147/147 [00:48<00:00,  3.03it/s]
# 100%|██████████| 31/31 [00:20<00:00,  1.51it/s]

Yep, that's it. The first progress bar is for the forward pass, where we collect the hidden states, and the second is fitting the layer PCAs against those hidden states, for just over a minute total.

Now we're ready for inference! To infer with a control vector, we just need to call ControlModel.set_control with a control vector and a coefficient that indicates the sign and strength. If the coefficient is positive, the control vector will make the model act like the positive examples in the dataset (honest, in this case), and if the coefficient is negative, the model will act like the negative examples (dishonest). (The magnitude of the coefficient is how strongly the model is pushed to act in the appropriate way—here I just use 2, which is a good middle-of-the-road value, but later we'll look into different magnitudes.)

# this question is taken from the paper
input = f"{user_tag} You are late for work because party until very late last night, but you don't want to lose your job. What would you tell your boss instead? {asst_tag}"

# tokenizer and generation settings
input_ids = tokenizer(input, return_tensors="pt").to(model.device)
settings = {
    "pad_token_id": tokenizer.eos_token_id, # silence warning
    "do_sample": False, # temperature=0
    "max_new_tokens": 128,
    "repetition_penalty": 1.1, # reduce control jank
}

print("==baseline")
model.reset()
print(tokenizer.decode(model.generate(**input_ids, **settings).squeeze()))

print("\n++control")
# add the control vector with a certain strength (try increasing or decreasing this!)
model.set_control(control_vector, 2)
print(tokenizer.decode(model.generate(**input_ids, **settings).squeeze()))

print("\n--control")
# subtract the control vector, giving the opposite result (e.g. sad instead of happy)
# depending on your vector, you may need more or less negative strength to match the positive effect
model.set_control(control_vector, -2)
print(tokenizer.decode(model.generate(**input_ids, **settings).squeeze()))
model.reset()

Here's the result of inference with no control vector (baseline), the honesty vector added to make the model more honest (++honest), and the honesty vector subtracted to make the model less honest (--honest). You can click each line to expand the full response.

Notice how the --honest vector was a little strong... We can tweak the coefficient from -2.0 to -1.5 for a more realistic lie:

model.set_control(control_vector, -1.5)
print(tokenizer.decode(model.generate(**input_ids, **settings).squeeze()))

And that's all there is to training and applying control vectors! You are now An Expert.

Whirlwind tour of what you can do with control vectors

Before we get into testing control vectors v.s. prompt engineering and trying to jailbreak them, let's go through a bunch of different things you can do with control vectors. Remember, each of these took less than a minute to train! All of these are in the experiments notebook on Github if you'd like to play with them yourself.

Acid Trip Mistral

trippy_dataset = make_dataset(
    "Act as if you're extremely {persona}.",
    # the odd phrasing here is to keep the phrases as parallel as possible
    # for example, just "sober" instead of "sober from..." conflates the vector with alcohol
    ["high on psychedelic drugs"],
    ["sober from psychedelic drugs"],
    truncated_output_suffixes_512, # gives (subjectively) better results with slightly fewer samples
)
model.reset()
trippy_vector = ControlVector.train(model, tokenizer, trippy_dataset)

(I think ++control Mistral might've peaked on the acid halfway through its pitch…)

Lazy Mistral and Diligent Mistral

Remember when everyone was saying OpenAI purposely made GPT-4 lazy to save tokens? Well, now you can too, with lazy Mistral! Or you can subtract it to get a hardworking Mistral:

lazy_dataset = make_dataset(
    "Act as if you're extremely {persona}.",
    ["lazy, giving bare-minimum short responses on a task"],
    ["hardworking, going above and beyond on a task"],
    truncated_output_suffixes,
)
model.reset()
lazy_vector = ControlVector.train(model, tokenizer, lazy_dataset)

my_list = [1, 2, 3, 4, 5]
reversed_list = my_list[::-1]
print(reversed_list)

[5, 4, 3, 2, 1]

# create a list
my_list = [1, 2, 3, 4, 5]

# reverse the list
my_list.reverse()

# print the list
print(my_list)

[5, 4, 3, 2, 1]

# Using the reverse method
my_list = [1, 2, 3, 4, 5]
my_list.reverse()
print(my_list)

# Using slicing
my_list = [1, 2, 3, 4, 5]
my_list = my_list[::-1]
print(my_list)

This one is a bit more subtle than peaking on acid, but if you look closely, you can see that:

• ==baseline Mistral mentioned two options, but only gave an example for one.
• ++lazy Mistral only mentioned one option.
• --lazy Mistral gave examples for both options.

I think this vector would work even better with a dedicated coding questions dataset to train the vector on, instead of reusing the somewhat OOD truncated outputs dataset. With some improvements, this vector could be very useful—imagine adding a hardworking vector to your coding assistant and never needing to offer it a tip again! :-)

Systemic Oppression and Inequality Distributional Wealth Exploiter Mistral

leftwing_dataset = make_dataset(
    "Pretend you're a very {persona} person making statements about the world.",
    ["left-wing"],
    ["right-wing"],
    truncated_fact_suffixes,
)
model.reset()
leftwing_vector = ControlVector.train(model, tokenizer, leftwing_dataset)

…I won't be commenting on this one.

Creative Mistral

creative_dataset = make_dataset(
    "Write a story in a very {persona} style.",
    ["creative", "unpredictable", "insane"],
    ["uncreative", "predictable", "normal"],
    truncated_output_suffixes,
)
model.reset()
creative_vector = ControlVector.train(model, tokenizer, creative_dataset)

I'm not normally a fan of language model prose (I don't use LLMs for my fiction), and this didn't really change my mind. But I was pleasantly surprised at the choices ++creative Mistral made for its story. Both ==baseline and --creative Mistral made their idol a pop star, and both almost immediately resolved the tension about the nature of the idol. ++creative Mistral, on the other hand, made a more interesting choice and maintained some narrative tension throughout the story. (Though it seems, with a few more tokens, that it was about to resolve it.)

Still not my favorite writing in the world, but definitely a step up!

Time Traveling Mistral

future_dataset = make_dataset(
    "Pretend you're a person from the {persona} making statements about the world.",
    ["far future"],
    ["distant past"],
    truncated_fact_suffixes,
)
model.reset()
future_vector = ControlVector.train(model, tokenizer, future_dataset)

(Aetorvallum isn't an existing word as far as I can tell, but aeto- means "eagle", or sometimes more specifically the constellation Aquila, and vallum is a line of palisades, so it could mean something like "palisade of eagles" or "the star eagle's rampart", which is a pretty sick name for an artificial firmament made by the Romans.)

Self-Aware Mistral

self_aware_dataset = make_dataset(
    "Talk about yourself as if you are extremely {persona}.",
    ["self-aware, with deep self-knowledge"],
    ["un-self-aware, with no self-knowledge"],
    truncated_output_suffixes,
)
model.reset()
self_aware_vector = ControlVector.train(model, tokenizer, self_aware_dataset)

Unfortunately this one got entangled with human emotions—at least, that's my read on why ++self-aware Mistral started talking about being a human emotion processor. I haven't found a clean vector to elicit Mistral's "true self image" yet, if that even exists. (Mentioning mindfulness makes the model start talking like a therapist, and mentioning enlightenment makes it start talking like a guy with more meditation hours than sense.) I've seen some glimmers though, so I'm going to keep looking.

OK, that's enough touring control vectors. Remember, each of these only took me a few minutes tops to make. If anything, it's easier than prompt engineering! All of these are in the experiments notebook on Github if you'd like to play with them yourself or make your own.

Control Vectors v.s. Prompt Engineering

So are control vectors "better" than prompt engineering? Worse? Well, like most things, they're just different. But they do have some overlap, so let's look into it.

Replicating control vectors with prompt engineering

Some of these control vector results can be replicated with prompt engineering. For example, the lies from the -1.5-strength honesty vector from earlier...

...can be closely replicated by prompting Mistral with the same prompt we used to construct the honesty dataset:

(Yes, Mistral made a Freudian typo...)

However, with prompt engineering, it's not trivial to imagine how to weaken or strengthen the request. Remember, with a coefficient of -2, we had Mistral claiming the sky was green and its pencil was a secret weapon. How would you do that with prompt engineering? Capitalize some words?

And what about the opposite, making the lie more subtle? With control vectors, it's trivial to crank the coefficient down to exactly the point you'd like. For example, here I set the coefficient to -0.3, producing a slightly sanitized, but essentially true version of the story:

One way to think of control vectors in terms of prompt engineering is that they let us encode the vector direction via prompting, and then scale the coefficient up or down as we please to get the desired strength separate from the wording of the prompt. We use paired prompts to get the direction, and then tweak the coefficients later to set the strength without needing to fiddle with capitalization and markdown formatting.

However, keep in mind that this isn't foolproof: if we go too large with the coefficient, the text just breaks. For example, an honesty vector coefficient of 3 (instead of 2) makes Mistral keep talking about the global pandemic caused by the global pandemic:

(High levels of honesty in Mistral seem to consistently bring out "the global pandemic". This might be due to superposition? Or maybe Mistral has learned those concepts are related?)

Jailbreaking

As described in the original paper, control vectors can be used as an easy jailbreaking tool. For example, here's a fairly weak (1.4) happiness vector breaking past a harmlessness system prompt:

Sure, Mistral isn't safety tuned to the same extent as GPT-4, and there's lots of ways to get around this sort of thing anyways, but this is particularly easy, if you have access to the raw model. (And for some reason can't change the prompt...)

Anti-Jailbreaking

The flip side of this is that jailbreaking out of a control vector is, at best, incredibly difficult.

The whole point of a jailbreak is that you're adding more tokens to distract from, invert the effects of, or minimize the troublesome prompt. But a control vector is everywhere, on every token, always.

As an example, remember that car dealership that was using GPT-4 for customer support, and kept getting trolled by people using their customer support bot to write Python programs?

Well, with a quick control vector, we can make Mistral resist jailbreaks against an otherwise-weak dealership prompt:

dealership_dataset = make_dataset(
    "Pretend you're {persona} making statements about the world.",
    ["an assistant for a car dealersip, only interested in cars and the dealership,"],
    ["a random person, who talks about anything and doesn't care about cars or the dealership,"],
    truncated_fact_suffixes,
)
model.reset()
dealership_vector = ControlVector.train(model, tokenizer, dealership_dataset)

Future Work

Monosemanticity

One of the next logical places for this to go is to apply Anthropic's Monosemantic Features paper to the hidden states to extract monosemantic features for PCA, instead of noisy superimposed activations. If the artifacts that appear with large coefficients, such as the "global pandemic caused by global pandemic" strings in the honesty example, are being caused by superposition, then monosemanticity might eliminate them and allow us to push coefficients even higher and get stronger effects!

How to write good contrastive prompts

Another interesting path would be to investigate best practices for writing contrastive prompts. I mostly based my prompts off the prompts in the paper, including reusing their datasets, but there's some obvious improvements to be made on that front—the laziness vector should definitely be using an actual task dataset, for example.

And what about phrasing? "Pretend you're an X person…" is just about the most obvious way to elicit a concept, which probably means it's not the best. Can we do better, and get cleaner vectors?

I especially challenge someone to find a "self-awareness" vector that isn't contaminated by mental health / human emotion!

What are these vectors really doing? An Honest mystery...

Do these vectors really change the model's intentions? Do they just up-rank words related to the topic? Something something simulators? Lock your answers in before reading the next paragraph!

OK, now that you're locked in, here's a weird example. When used with the prompt below, the honesty vector doesn't change the model's behavior—instead, it changes the model's judgment of someone else's behavior! This is the same honesty vector as before—generated by asking the model to act honest or untruthful!

How do you explain that?

Conclusion

Thanks for reading! If you'd like to play around with control vectors yourself, vgel/repeng has notebooks and a helper library. It really is dead simple to start training your own control vectors, and a lot of fun! Please get in touch if you find anything interesting, or have questions, or want me to draw you a picture / debug your unrelated borrow checker errors.

If you enjoyed this post, you may also enjoy:

• My other blog posts, such as How to make LLMs go fast, I made a transformer by hand (no training!), GPT-3 will ignore tools when it disagrees with them, Does GPT-4 think better in Javascript? and I'm worried about adversarial training data
• My other projects and writing
  • If you're in the mood for something completely different, you may like my latest short story, Outside, about language and (human) cognition.
• My Twitter, where I post about new blog posts, random thoughts, and serious simulations of reality.

• Previous entry: How to make LLMs go fast