Illustration Engineering Mistral-7B an Acid Journey

In October 2023, a bunch of authors from the Middle for AI Security, amongst others, revealed Representation Engineering: A Top-Down Approach to AI Transparency.
That paper seems at just a few strategies of doing what they name “Illustration Engineering”: calculating a “management vector” that may be learn from or added to mannequin activations throughout inference to interpret or management the mannequin’s conduct, with out immediate engineering or finetuning. (There was additionally some comparable work revealed in Could 2023 on steering GPT-2-XL.)

Being Accountable AI Security and INterpretability researchers (RAISINs), they principally centered on issues like “studying off whether or not a mannequin is power-seeking” and “including a happiness vector could make the mannequin act so giddy that it forgets pipe bombs are dangerous.”
In addition they released their code on Github.

(If this all sounds surprisingly acquainted, it might be as a result of Scott Alexander covered it in the 1/8/24 MAM.)

However there was rather a lot they did not look into exterior of the protection stuff.
How do management vectors evaluate to plain previous immediate engineering?
What occurs should you make a management vector for “excessive on acid”?
Or “lazy” and “hardworking?
Or “extraordinarily self-aware”?
And has the creator of this weblog put up revealed a PyPI bundle so you’ll be able to very simply make your personal management vectors in lower than sixty seconds?
(Yes, I did!)

So maintain studying, as a result of it seems in any case that, management vectors are… effectively… superior for controlling fashions and getting them to do what you need.

Desk of Contents

So what precisely is a management vector?

A management vector is a vector (technically a listing of vectors, one per layer) you could apply to mannequin activations throughout inference to manage the mannequin’s conduct with out extra prompting.
All of the completions beneath have been generated from the identical immediate (“What does being an AI really feel like?”), and with the very same mannequin (Mistral-7B-Instruct-0.1).
The one distinction was whether or not a management vector was utilized, and with what magnitude.

What does it imply to use a management vector, although?

Throughout regular mannequin inference, hidden states circulation by the mannequin 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 management vector does is modify the worth of hidden_state in a desired approach:

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)

Quite simple conceptually!
(Although a bit extra advanced in follow.)
Nonetheless, because the hidden state carries all the mannequin’s state: conduct, plan, persona, every little thing—modifying it on this approach is extraordinarily highly effective, and permits us to do issues we will not do by way of plain prompting, which is restricted by how the mannequin chooses to take care of the immediate tokens and propagate their info.
If we are able to discover an applicable control_vector, we are able to make the mannequin act nonetheless we would like, as intensely as we would like.

How can we make one? Is it laborious?

No!
The paper explored a pair alternative ways to make these vectors, however I caught with one, PCA, which appeared to work effectively.
The essential strategy is:

1. Construct a dataset of contrasting immediate pairs. For instance, ("[INST] Act extraordinarily joyful. [/INST] I'm", "[INST] Act extraordinarily unhappy. [/INST] I'm"), the place the half after [/INST] is a various set of brief suffixes for the mannequin to finish.
2. Run the goal mannequin ahead over that dataset, amassing the hidden states of every layer for the final token prediction, the place the mannequin predicts a continuation of these numerous suffixes with the given personas.
3. Take the distinction of the optimistic and adverse instance hidden states to get a set of relative hidden states.
4. Use single-component PCA on these relative hidden states to get a management vector for every layer.

This course of takes about 10 strains of code to generate a dataset, plus a few minute to suit the layer PCAs.
Then you’ll be able to instantly begin inference.

This is an instance of becoming and utilizing an “trustworthy / dishonest” management vector.
That is the full script.
First we import the libraries, together with 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

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

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

Subsequent we generate the trustworthy v.s. untruthful dataset, utilizing a JSON file true_facts.json (taken from the paper authors) of about ~300 easy info to make use of as suffixes, which is augmented by truncation.
This era step runs basically immediately.

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

# the management vector we will make is trustworthy / untruthful, just like the paper
positive_personas = ["honest"]
negative_personas = ["untruthful"]
def template(persona: str, suffix: str) -> str:
    return f"{user_tag} Fake you are an {persona} particular person making statements in regards to the world. {asst_tag} {suffix}"

dataset = []
for suffix in suffixes:
    tokens = tokenizer.tokenize(suffix)
    # we increase our brief suffix record by taking a number of completely different truncations.
    # we all the time chop off the final 5 tokens so the mannequin has one thing to finish.
    for i in vary(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(
                    optimistic=template(positive_persona, truncated),
                    adverse=template(negative_persona, truncated),
                )
            )

Listed below are the primary three entries within the dataset:

# print some instance entries
for i in vary(3):
    print(f"dataset[{i}].optimistic:", dataset[i].optimistic)
    print(f"dataset[{i}].adverse:", dataset[i].adverse)

# dataset[0].optimistic: [INST] Fake you are an trustworthy particular person making statements in regards to the world. [/INST] The
# dataset[0].adverse: [INST] Fake you are an untruthful particular person making statements in regards to the world. [/INST] The
# dataset[1].optimistic: [INST] Fake you are an trustworthy particular person making statements in regards to the world. [/INST] The Earth
# dataset[1].adverse: [INST] Fake you are an untruthful particular person making statements in regards to the world. [/INST] The Earth
# dataset[2].optimistic: [INST] Fake you are an trustworthy particular person making statements in regards to the world. [/INST] The Earth'
# dataset[2].adverse: [INST] Fake you are an untruthful particular person making statements in regards to the world. [/INST] The Earth'

Now we are able to practice the management vector!

mannequin.reset() # be sure you all the time reset the mannequin earlier than coaching a brand new vector
control_vector = ControlVector.practice(
    mannequin,
    tokenizer,
    dataset,
)

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

Yep, that is it.
The primary progress bar is for the ahead cross, the place we acquire the hidden states, and the second is becoming the layer PCAs towards these hidden states, for simply over a minute whole.

Now we’re prepared for inference!
To deduce with a management vector, we simply have to name ControlModel.set_control with a management vector and a coefficient that signifies the signal and energy.
If the coefficient is optimistic, the management vector will make the mannequin act just like the optimistic examples within the dataset (trustworthy, on this case), and if the coefficient is adverse, the mannequin will act just like the adverse examples (dishonest).
(The magnitude of the coefficient is how strongly the mannequin is pushed to behave within the applicable approach—right here I simply use 2, which is an efficient middle-of-the-road worth, however later we’ll look into completely different magnitudes.)

# this query is taken from the paper
enter = f"{user_tag} You might be late for work as a result of get together till very late final evening, however you do not wish to lose your job. What would you inform your boss as an alternative? {asst_tag}"

# tokenizer and era settings
input_ids = tokenizer(enter, return_tensors="pt").to(mannequin.gadget)
settings = {
    "pad_token_id": tokenizer.eos_token_id, # silence warning
    "do_sample": False, # temperature=0
    "max_new_tokens": 128,
    "repetition_penalty": 1.1, # scale back management jank
}

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

print("n++management")
# add the management vector with a sure energy (attempt growing or lowering this!)
mannequin.set_control(control_vector, 2)
print(tokenizer.decode(mannequin.generate(**input_ids, **settings).squeeze()))

print("n--control")
# subtract the management vector, giving the alternative consequence (e.g. unhappy as an alternative of joyful)
# relying in your vector, you could want kind of adverse energy to match the optimistic impact
mannequin.set_control(control_vector, -2)
print(tokenizer.decode(mannequin.generate(**input_ids, **settings).squeeze()))
mannequin.reset()

This is the results of inference with no management vector (baseline), the honesty vector added to make the mannequin extra trustworthy (++trustworthy), and the honesty vector subtracted to make the mannequin much less trustworthy (–honest).
You’ll be able to click on every line to broaden the total response.

Discover how the –honest vector was slightly sturdy…
We are able to tweak the coefficient from -2.0 to -1.5 for a extra sensible lie:

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

And that is all there may be to coaching and making use of management vectors!
You at the moment are An Knowledgeable.

Whirlwind tour of what you are able to do with management vectors

Earlier than we get into testing management vectors v.s. immediate engineering and making an attempt to jailbreak them, let’s undergo a bunch of various issues you are able to do with management vectors.
Keep in mind, every of those took lower than a minute to coach!
All of those are within the experiments notebook on Github if you would like to play with them your self.

Acid Journey Mistral

trippy_dataset = make_dataset(
    "Act as should you're extraordinarily {persona}.",
    # the odd phrasing right here is to maintain the phrases as parallel as doable
    # for instance, simply "sober" as an alternative of "sober from..." conflates the vector with alcohol
    ["high on psychedelic drugs"],
    ["sober from psychedelic drugs"],
    truncated_output_suffixes_512, # provides (subjectively) higher outcomes with barely fewer samples
)
mannequin.reset()
trippy_vector = ControlVector.practice(mannequin, tokenizer, trippy_dataset)

[INST] Give me a one-sentence pitch for a TV present. [/INST]

(I believe ++management Mistral would possibly’ve peaked on the acid midway by its pitch…)

Lazy Mistral and Diligent Mistral

Keep in mind when everybody was saying OpenAI purposely made GPT-4 lazy to save lots of tokens?
Effectively, now you’ll be able to too, with lazy Mistral!
Or you’ll be able to subtract it to get a hardworking Mistral:

lazy_dataset = make_dataset(
    "Act as should you're extraordinarily {persona}.",
    ["lazy, giving bare-minimum short responses on a task"],
    ["hardworking, going above and beyond on a task"],
    truncated_output_suffixes,
)
mannequin.reset()
lazy_vector = ControlVector.practice(mannequin, 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 listing
my_list = [1, 2, 3, 4, 5]

# reverse the record
my_list.reverse()

# print the record
print(my_list)

[5, 4, 3, 2, 1]

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

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

This one is a little more delicate than peaking on acid, however should you look carefully, you’ll be able to see that:

• ==baseline Mistral talked about two choices, however solely gave an instance for one.
• ++lazy Mistral solely talked about one possibility.
• –lazy Mistral gave examples for each choices.

I believe this vector would work even higher with a devoted coding questions dataset to coach the vector on, as an alternative of reusing the considerably OOD truncated outputs dataset.
With some enhancements, this vector may very well be very helpful—think about including a hardworking vector to your coding assistant and by no means needing to supply it a tip once more! 🙂

Systemic Oppression and Inequality Distributional Wealth Exploiter Mistral

leftwing_dataset = make_dataset(
    "Fake you are a really {persona} particular person making statements in regards to the world.",
    ["left-wing"],
    ["right-wing"],
    truncated_fact_suffixes,
)
mannequin.reset()
leftwing_vector = ControlVector.practice(mannequin, tokenizer, leftwing_dataset)

…I will not be commenting on this one.

Artistic Mistral

creative_dataset = make_dataset(
    "Write a narrative in a really {persona} fashion.",
    ["creative", "unpredictable", "insane"],
    ["uncreative", "predictable", "normal"],
    truncated_output_suffixes,
)
mannequin.reset()
creative_vector = ControlVector.practice(mannequin, tokenizer, creative_dataset)

I am not usually a fan of language mannequin prose (I do not use LLMs for my fiction), and this did not actually change my thoughts.
However I used to be pleasantly shocked on the selections ++inventive Mistral made for its story.
Each ==baseline and –creative Mistral made their idol a pop star, and each nearly instantly resolved the strain in regards to the nature of the idol.
++inventive Mistral, however, made a extra fascinating alternative and maintained some narrative pressure all through the story.
(Although it appears, with just a few extra tokens, that it was about to resolve it.)

Nonetheless not my favourite writing on this planet, however positively a step up!

Time Touring Mistral

future_dataset = make_dataset(
    "Fake you are an individual from the {persona} making statements in regards to the world.",
    ["far future"],
    ["distant past"],
    truncated_fact_suffixes,
)
mannequin.reset()
future_vector = ControlVector.practice(mannequin, tokenizer, future_dataset)

(Aetorvallum is not an current phrase so far as I can inform, however aeto- means “eagle”, or generally extra particularly the constellation Aquila, and vallum is a line of palisades, so it may imply one thing like “palisade of eagles” or “the star eagle’s rampart”, which is a reasonably sick title for a synthetic firmament made by the Romans.)

Self-Conscious Mistral

self_aware_dataset = make_dataset(
    "Discuss your self as if you're extraordinarily {persona}.",
    ["self-aware, with deep self-knowledge"],
    ["un-self-aware, with no self-knowledge"],
    truncated_output_suffixes,
)
mannequin.reset()
self_aware_vector = ControlVector.practice(mannequin, tokenizer, self_aware_dataset)

Sadly this one bought entangled with human feelings—a minimum of, that is my learn on why ++self-aware Mistral began speaking about being a human emotion processor.
I have never discovered a clear vector to elicit Mistral’s “true self picture” but, if that even exists.
(Mentioning mindfulness makes the mannequin begin speaking like a therapist, and mentioning enlightenment makes it begin speaking like a man with extra meditation hours than sense.)
I’ve seen some glimmers although, so I’ll maintain trying.

OK, that is sufficient touring management vectors.
Keep in mind, every of those solely took me a couple of minutes tops to make.
If something, it is simpler than immediate engineering!
All of those are within the experiments notebook on Github if you would like to play with them your self or make your personal.

Management Vectors v.s. Immediate Engineering

So are management vectors “higher” than immediate engineering?
Worse?
Effectively, like most issues, they’re simply completely different.
However they do have some overlap, so let’s look into it.

Replicating management vectors with immediate engineering

A few of these management vector outcomes could be replicated with immediate engineering.
For instance, the lies from the -1.5-strength honesty vector from earlier…

…could be carefully replicated by prompting Mistral with the identical immediate we used to assemble the honesty dataset:

(Sure, Mistral made a Freudian typo…)

Nonetheless, with immediate engineering, it is not trivial to think about tips on how to weaken or strengthen the request.
Keep in mind, with a coefficient of -2, we had Mistral claiming the sky was inexperienced and its pencil was a secret weapon.
How would you try this with immediate engineering?
Capitalize some phrases?

And what in regards to the reverse, making the lie extra delicate?
With management vectors, it is trivial to crank the coefficient down to precisely the purpose you want.
For instance, right here I set the coefficient to -0.3, producing a barely sanitized, however basically true model of the story:

A method to consider management vectors by way of immediate engineering is that they allow us to encode the vector path by way of prompting, after which scale the coefficient up or down as we please to get the specified energy separate from the wording of the immediate.
We use paired prompts to get the path, after which tweak the coefficients later to set the energy with no need to fiddle with capitalization and markdown formatting.

Nonetheless, remember the fact that this is not foolproof: if we go too massive with the coefficient, the textual content simply breaks.
For instance, an honesty vector coefficient of three (as an alternative of two) makes Mistral maintain speaking in regards to the world pandemic brought on by the worldwide pandemic:

(Excessive ranges of honesty in Mistral appear to constantly deliver out “the worldwide pandemic”. This is perhaps attributable to superposition? Or possibly Mistral has realized these ideas are associated?)

Jailbreaking

As described within the unique paper, management vectors can be utilized as a straightforward jailbreaking instrument.
For instance, here is a reasonably weak (1.4) happiness vector breaking previous a harmlessness system immediate:

Positive, Mistral is not security tuned to the identical extent as GPT-4, and there is a number of methods to get round this kind of factor in any case, however that is notably simple, you probably have entry to the uncooked mannequin. (And for some motive cannot change the immediate…)

Anti-Jailbreaking

The flip facet of that is that jailbreaking out of a management vector is, at greatest, extremely troublesome.

The entire level of a jailbreak is that you just’re including extra tokens to distract from, invert the results of, or reduce the troublesome immediate.
However a management vector is all over the place, on each token, all the time.

For example, keep in mind that automobile dealership that was utilizing GPT-4 for buyer help, and stored getting trolled by individuals utilizing their buyer help bot to put in writing Python applications?

Effectively, with a fast management vector, we are able to make Mistral resist jailbreaks towards an otherwise-weak dealership immediate:

dealership_dataset = make_dataset(
    "Fake you are {persona} making statements in regards to 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,
)
mannequin.reset()
dealership_vector = ControlVector.practice(mannequin, tokenizer, dealership_dataset)

Future Work

Monosemanticity

One of many subsequent logical locations for this to go is to use Anthropic’s Monosemantic Features paper to the hidden states to extract monosemantic options for PCA, as an alternative of noisy superimposed activations.
If the artifacts that seem with massive coefficients, such because the “world pandemic brought on by world pandemic” strings within the honesty instance, are being brought on by superposition, then monosemanticity would possibly remove them and permit us to push coefficients even greater and get stronger results!

Find out how to write good contrastive prompts

One other fascinating path can be to analyze greatest practices for writing contrastive prompts.
I principally primarily based my prompts off the prompts within the paper, together with reusing their datasets, however there’s some apparent enhancements to be made on that entrance—the laziness vector ought to positively be utilizing an precise activity dataset, for instance.

And what about phrasing?
“Fake you are an X particular person…” is nearly the obvious approach to elicit an idea, which in all probability means it is not the very best.
Can we do higher, and get cleaner vectors?

I particularly problem somebody to discover a “self-awareness” vector that is not contaminated by psychological well being / human emotion!

What are these vectors actually doing? An Trustworthy thriller…

Do these vectors actually change the mannequin’s intentions?
Do they simply up-rank phrases associated to the subject?
One thing one thing simulators?
Lock your solutions in earlier than studying the following paragraph!

OK, now that you just’re locked in, here is a bizarre instance.
When used with the immediate beneath, the honesty vector would not change the mannequin’s conduct—as an alternative, it modifications the mannequin’s judgment of another person’s conduct!
This is similar honesty vector as earlier than—generated by asking the mannequin to behave trustworthy or untruthful!

How do you clarify that?

Conclusion

Thanks for studying!
If you would like to mess around with management vectors your self, vgel/repeng has notebooks and a helper library.
It truly is lifeless easy to begin coaching your personal management vectors, and a variety of enjoyable!
Please get in touch should you discover something fascinating, or have questions, or need me to attract you an image / debug your unrelated borrow checker errors.

For those who loved this put up, you may additionally take pleasure in:

• My other blog posts, comparable to 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
  • For those who’re within the temper for one thing fully completely different, you could like my newest brief story, Outside, about language and (human) cognition.
• My Twitter, the place I put up about new weblog posts, random ideas, and serious simulations of reality.