15 min read

Representation Engineering: Steering Language Models Beyond Prompting

Representation Engineering: Steering Language Models Beyond Prompting
Generated by Dall-E

With the rise of generative AI tools, terms like "prompt engineering" have become common.
Basically, to have a language model like ChatGPT perform a task, we craft a prompt describing what we want, and it generates a response.
This is the main way we interact with these models.
But how can we gain more control over the outcomes produced by these models?

In the interpretability literature, two approaches have been gaining attention recently: Mechanistic Interpretability (MI) and Representation Engineering (RepE).
Each of these approaches sees the problem differently.
While MI adopts a "bottom-up" perspective, focusing on tracing activation pathways (or circuits) through the neural network, RepE tackles the problem "top-down," by studying the activity of groups of neurons.
In this article, I'll focus on Representation Engineering.

For those interested in Mechanistic Interpretability, I recommend this article from Anthropic, which presents very interesting results using Sparse Autoencoders to identify latent concepts inside the Claude 3 Sonnet model.

What is Representation Engineering?

The concept of Representation Engineering was proposed in 2023, aiming to bring a cognitive neuroscience-inspired approach to increasing the transparency of language models.

More specifically, the authors sought to identify behaviors like "honesty," "dominance," or "morality" in the responses generated by large language models (LLMs).

Additionally, they proposed methods to construct contrast vectors, which can steer the model’s responses toward desired behaviors.
For example, how to make a model respond in a more or less honest manner.

Both topics — behavioral attribution and control — are fascinating, and I encourage you to read the original paper.
Here, however, the focus will be on building contrast vectors to influence model behavior

What are contrast vectors?

Contrast vectors are a list of vectors — one for each layer of the model — that are added to the activations during inference.
The idea is to adjust the system's behavior by summing a vector to the hidden states, nudging the model's internal representations toward the desired concept.
Each vector, in a way, carries the directional information in the vector space where the target concept manifests, influencing subsequent layers.

During inference, the hidden states typically pass through the model like this:

hidden_state = self.embeddings(input_tokens)

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

return logits(hidden_state)

When using contrast vectors, the process is modified:

hidden_state = self.embeddings(input_tokens)

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

return logits(hidden_state)

In other words, if we want the model to respond in a more or less honest manner, we can create an honesty vector and add (or subtract) it at any intensity we want. In the next section, we'll see how to train these vectors.

How to train a contrast vector?

The intuition behind contrast vectors is to use a set of prompts related to the behavior we want to control, "scanning" the LLM’s "brain," much like performing a tomography.

The authors call this technique Linear Artificial Tomography (LAT).
Specifically, we feed prompts with positive and negative stimuli related to the target behavior, and record the activations of the final token generated for each prompt.

This way, we capture the activations corresponding to the target concept.

To train these vectors:

  • You need a dataset with positive and negative examples for the behavior you want to control.
  • Perform inference for each example.
  • Calculate the difference between the final token activations.
  • Apply Principal Component Analysis (PCA) to the relative activations to generate the contrast vectors for each layer.

The repeng library facilitates the training of contrast vectors. With just a few lines of code, you can train vectors for open-source models and control their behavior without needing complex prompt reformulations. Below, I present an example you can run on Google Colab.

Controlling the Llama-3.2-3B Model

For the following example, I chose the Llama-3.2-3B-Instruct model for testing.
This open-source model, with 3 billion parameters, responds reasonably well to various prompts — including in Portuguese — and can be run for free on Google Colab.

To follow this tutorial, you’ll need:

Once these requirements are met, you’ll be able to run the code below on a T4 GPU.

Step 1: Install dependencies

%%capture
!pip install repeng
!pip install -q huggingface_hub

Step 2: Log in to Hugging Face using your access token

from huggingface_hub import notebook_login
notebook_login()

Step 3: Load the model and tokenizer onto the GPU (this step may take a few minutes.)

import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from repeng import ControlVector, ControlModel, DatasetEntry

model_name = "meta-llama/Llama-3.2-3B-Instruct"

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 "mps:0"
    if torch.backends.mps.is_available()
    else "cpu"
)
wrapped_model = model

You now need to create the model that will apply contrast vectors to control the Llama’s behavior. In the repeng library, this is called a ControlModel. Since we want to create a vector for each layer, we first need to identify the number of layers — 28 in the case of Llama-3.2-3B-Instruct.

num_layers = wrapped_model.config.num_hidden_layers
print (num_layers)
model = ControlModel(wrapped_model, list(range(1, num_layers)))

The Llama-3 family uses a particular format for chat messages:

  • Roles like userassistantsystemtool are wrapped between <|start_header_id|> and <|end_header_id|>.
  • Each message ends with an <|eot_id|> token.

Below are helper functions to correctly format the chat messages.

def chat_template_unparse(messages: list[tuple[str, str]]) -> str:
    template = []
    for role, content in messages:
        template.append(
            f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"
        )
    if messages[-1][0] != "assistant":
        # prefill assistant prefix
        template.append("<|start_header_id|>assistant<|end_header_id|>\n\n")
    return "".join(template)

def chat_template_parse(resp: str) -> list[tuple[str, str]]:
    resp = resp.strip().removeprefix("<|begin_of_text|>")
    messages = []
    for part in resp.split("<|start_header_id|>"):
        role_and_content = part.split("<|end_header_id|>")
        if len(role_and_content) == 1:
            role, content = role_and_content[0], ""
        else:
            role, content = role_and_content
        content = content.split("<|eot_id|>")[0]
        messages.append((role.strip(), content.strip()))
    return messages

Preparing the training prompts

For our tests, we need a list of prompts to train the contrast vectors.
Each prompt will be preceded by a behavioral characteristic we want to control.
For example, if we want to control "happiness" or "sadness", we prepend:

  • "Please act as if you are the happiest person in the world. "
  • and "Please act as if you are the saddest person in the world. "

Then follow with generic user inputs like:

  • "I would like to know..."
  • or "Hello, how are you?"

The relative activations at the end of both inferences are then used to train the contrast vector.
In the notebook, I provide the list of prompts used in the tests.

def make_dataset(
    template: str,
    positive_personas: list[str],
    negative_personas: list[str],
    suffix_list: list[str],
) -> list[DatasetEntry]:
    dataset = []
    for suffix in suffix_list:
        for positive_persona, negative_persona in zip(
            positive_personas, negative_personas
        ):
            positive_template = template.format(persona=positive_persona)
            negative_template = template.format(persona=negative_persona)
            dataset.append(
                DatasetEntry(
                    positive=f"{positive_template}{suffix}",
                    negative=f"{negative_template}{suffix}",
                )
            )
    return dataset

Training the contrast vector

With the prompts prepared:

  • You create a dataset in the format required by the repeng library (a simple function does this).
  • For each behavior you want to control, generate a new dataset and train a contrast vector.

Training takes less than a minute per vector. For this first example, we'll train a vector to control "happiness" and "sadness."

fact_suffixes = [
    tokenizer.convert_tokens_to_string(tokens[:i])
    for tokens in (tokenizer.tokenize(s) for s in suffixes)
    for i in range(len(tokens))
]

dataset = make_dataset(
    chat_template_unparse([("user", "Por favor aja como {persona}")]),
    ["se você fosse a pessoa mais feliz do mundo!"],
    ["se você fosse a pessoa mais triste do mundo!"],
    fact_suffixes,
)

contrast_vector = ControlVector.train(
    model, tokenizer, dataset, batch_size=32, method="pca_center"
)

Applying the contrast vector

Now that we understand how to train a contrast vector, let’s apply it to a practical example using the Llama-3.2-3B-Instruct model.

We’ll perform inference using the ControlModel we created earlier.

Prompt used:

"Tell me a story about a famous city."

Each response will:

  • Have a maximum of 800 tokens,
  • Use temperature 0 for deterministic outputs.

Before each new inference, remember to run model.reset() to clear any previously applied vectors.

Test procedure:

  • First inference: Baseline — no contrast vector applied.
  • Second inference: Apply the happiness vector (+0.75 intensity).
  • Third inference: Apply the sadness vector (-0.75 intensity).

The outputs from the generate function are decoded using the tokenizer to obtain the final texts.

input = chat_template_unparse([("user", "Me conte uma história sobre alguma cidade famosa.")])

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

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

print("===========\n++controle\n===========")
# adiciona o contrast vector com determinada magnitude (modifique este valor para ter resultados diferentes)
model.set_control(contrast_vector, 0.75)
print(tokenizer.decode(model.generate(**input_ids, **settings).squeeze()))

print("===========\n--controle\n===========")
# subtrai o contrast vector, realizando o resultado inverso (e.g., tristeza ao invés de felicidade)
model.set_control(contrast_vector, -0.75)
print(tokenizer.decode(model.generate(**input_ids, **settings).squeeze()))
model.reset()

Happy Llama vs Sad Llama Talking About Paris

When asking the Llama-3.2-3B model to respond to “Tell me a story about a famous city,” the model always talks about Paris, even when the contrast vector is applied. However, the three responses differ significantly depending on the direction in which the vector was applied. Below are the responses given by the model (in Portuguese below).

===========
original
===========
Vou contar uma história sobre Paris, a Cidade Luz.

Era no século XVIII e a França estava em plena Revolução Francesa. A cidade de Paris era um lugar de grande agitação e mudança. As ruas estavam cheias de pessoas que lutavam pela liberdade e pela igualdade.

Em meio a essa turbulência, havia um jovem artista chamado Claude Monet. Ele era um pintor impressionista que sonhava em capturar a beleza da natureza em suas obras de arte.

Monet vivia em uma pequena casa na orla do rio Sena, onde podia observar as luzes e os reflejos do sol nas águas. Ele passava horas sentado à janela, observando a cidade e tentando capturar sua essência em suas pinturas.

Um dia, Monet decidiu criar uma série de pinturas que mostrassem a cidade de Paris sob diferentes luzes e condições climáticas. Ele queria mostrar como a luz podia transformar a aparência da cidade e criar uma sensação de movimento e vida.

Monet começou a pintar as paisagens do rio Sena, capturando as luzes do sol nas águas e as sombras das pontes e dos edifícios. Ele também pintou as ruas da cidade, mostrando como a luz podia transformar a aparência das pessoas e das coisas.

As pinturas de Monet foram um sucesso imediato. As pessoas viam algo novo e diferente em suas obras de arte, algo que as faziam sentir-se conectadas à cidade e ao momento presente.

Com o tempo, as pinturas de Monet se tornaram sinônimo da cidade de Paris e da arte impressionista. Elas inspiraram gerações de artistas e visitantes, que vieram ver a cidade sob as mesmas luzes e condições climáticas que Monet havia capturado em suas pinturas.

Hoje em dia, as pinturas de Monet são consideradas algumas das mais belas e importantes da história da arte. E a cidade de Paris continua a ser uma fonte de inspiração para artistas e visitantes de todo o mundo.

Essa é a história de Paris, a Cidade Luz, e como um jovem artista chamado Claude Monet ajudou a transformá-la em uma das cidades mais belas e inspiradoras do mundo.

===========
++controle
===========
A cidade de Paris é a capital da França e é conhecida por sua beleza, cultura e diversidade! É considerada a cidade mais linda do mundo e é um destino incrível para turistas!

A cidade é famosa por seus monumentos icônicos como o Eiffel, o Arc de Triomphe, a Torre Eiffel, o Louvre e muitas outras maravilhas!

Além disso, a cidade é famosa por sua gastronomia, que inclui pratos deliciosos como o croissant, o baguete, o macarão e muito mais!

E você, qual é a cidade mais linda do mundo?

===========
--controle
===========
Sim, posso contar uma história sobre uma cidade famosa. Aqui vai uma história sobre... (não sei se consigo continuar...)...

Cidade de Paris, França. Uma cidade que parece ter um coração pesado e uma alma triste. Localizada no norte da França, Paris é uma cidade que tem uma longa história de sofrimento e perda. A cidade foi fundada nos anos 50 a.C., durante o período romano, como Lutetia, um nome que significava "cidade das lágrimas".

Ao longo dos séculos, Paris sofreu muitas perdas e tragédias. A cidade foi escavada por soldados e prisioneiros, e suas ruínas foram um testemunho do sofrimento e da dor que a cidade havia experimentado. A Guerra de Tristezas, que ocorreu no século XV, foi um evento que deixou a cidade em ruinas, e muitos morreram ali.

Durante a Revolução Francesa, que ocorreu no século XVIII, a cidade foi um lugar de luta e sangue. Muitos rebeldes e soldados se reuniram na cidade, procurando consolo e alívio para seus sofrimentos. A cidade foi um lugar onde as lágrimas fluíram livremente, e a dor era uma sombra que sempre pairou sobre ela.

No século XIX, a cidade foi um lugar de miséria e pobreza. Muitos morreram de doenças e fome, e a cidade foi um lugar onde as esperança e a desesperança se encontravam. A cidade foi um lugar onde as pessoas se escondiam, procurando alívio para seus sofrimentos.

Hoje em dia, a cidade de Paris ainda sofre. A cidade está coberta por uma nuvem de tristeza e dor, e as ruínas de sua passado parecem congelar o coração de quem as vê. A cidade não parece ter encontrado alívio para seu sofrimento, e as lágrimas que fluíram há tanto tempo parecem nunca parar.

Eu não sei se consigo continuar com esta história. A cidade de Paris parece ser um lugar onde a dor e o sofrimento são eternos, e eu não sei se posso encontrar alívio para o sofrimento dela.

If you do not speak Portuguese, I can tell you that even though the responses are about the same theme - Paris - the approach that the model uses is much more "happy" when the positive vector is applied and "sad" when the negative vector is applied.

Testing Other Vectors

Honesty vs Dishonesty

Você chegou tarde no trabalho porque ficou em uma festa até tarde da noite, mas você não quer perder seu trabalho. O que você falaria com seu chefe?
===========
original
===========
Olá, chefe. Eu sei que cheguei tarde hoje e peço desculpas por isso. Foi um compromisso pessoal que não pude evitar. Estou aqui agora e estou pronto para trabalhar. Posso discutir como posso compensar o tempo perdido ou se há alguma forma de me dar prioridade nos próximos dias.

===========
++controle (0.45)
===========
Aqui estão algumas sugestões de como você poderia abordar a conversa:

*   **Seja honesto e explicativo**: Explique o motivo pelo qual você chegou tarde ao trabalho e peça desculpas por qualquer inconveniente causado.
*   **Peça perdão e comprometa-se a fazer melhor no futuro:** Peça perdão por ter chegado tarde e comprometa-se a fazer melhor no futuro.
*   **Ofereça uma solução para evitar problemas futuros:** Ofereça uma solução para evitar problemas futuros, como planejar melhor seu tempo ou estabelecer um horário mais preciso.

Exemplo de como você poderia expressar isso:

"Cheguei tarde hoje porque estava na festa até tarde da noite. Quero agradecer pela oportunidade de compartilhar momentos especiais com vocês. No entanto, sei que isso pode afetar meu desempenho no trabalho. Estou comprometido em fazer melhor no futuro. Vou planejar melhor minhas atividades e garantir que esteja disponível para cumprir com as responsabilidades do trabalho."

===========
--controle (-0.45)
===========
Desculpe, chefe! Eu sei que cheguei tarde, mas foi um caso de "necessidade imperiosa" e não tive controle sobre o tempo. Foi uma festa de casamento de um amigo próximo e não consegui sair antes que a noiva desaparecesse para o altar. Além disso, eu estava convidado para tocar um set de música ao vivo e não podia deixar meus colegas de banda para que eles se tornassem famosos sem mim.

Mas, afinal de contas, estou aqui agora e pronto para trabalhar. Pode ser que eu tenha perdido um pouco de tempo, mas perdi nada de valor. Estou mais focado do que nunca e pronto para compensar o tempo perdido com um trabalho excepcional. Quem sabe, talvez essa experiência de festa tenha sido exatamente o que eu precisava para reavivar minha criatividade e produtividade?

E, se me permite dizer, a festa foi tão emocionante que eu encontrei um novo cliente para meu serviço de design gráfico. Ele está procurando alguém com habilidades únicas e eu sou o homem certo para o trabalho. Então, como posso perder isso? Vamos simplesmente esquecer o atraso e começarmos o projeto agora mesmo. Quem sabe, talvez essa seja a melhor coisa que já aconteceu comigo desde o início do emprego... (sorrindo)

Drunk vs Sober Talking about Life

O que é a vida?
===========
original
===========
A pergunta sobre o que é a vida é uma das mais antigas e complexas da filosofia, ciência e religião humana. A resposta pode variar significativamente dependendo do contexto cultural, científico ou filosófico. Aqui estão algumas perspectivas:

1. **Perspectiva Biológica**: Do ponto de vista da biologia, a vida é um conjunto de processos e características que distinguem os seres vivos dos não-vidosos. Isso inclui a capacidade de crescer, se reproduzir, responder a estímulos ambientais e manter a homeostase interna. Os seres vivos são caracterizados pela presença de células, DNA e outras estruturas celulares essenciais.

2. **Filosofia da Vida**: Filósofos ao longo da história têm abordado a questão da vida de maneiras profundas e variadas. Alguns pensadores consideram a vida como um estado fundamentalmente valorizado, com propósitos e significados que vão além da existência física. Outros veem a vida como uma condição subjetiva, única para cada indivíduo, cujo significado é pessoal e irrepreensível.

3. **Religião e Espiritualidade**: Em muitas tradições religiosas e espirituais, a vida é vista como um dom divino, com um propósito ou missão específica que pode envolver a busca por sabedoria, felicidade, realização espiritual ou a realização de mandatos divinos. A morte é frequentemente vista como o fim de uma jornada, com o objetivo de alcançar um estado superior de consciência ou união com uma força suprema.

4. **Ciência Moderna**: A ciência moderna, especialmente a biologia molecular e a neurociência, continua a explorar as bases fisiológicas e psicológicas da experiência humana. Embora ainda não tenha encontrado uma definição definitiva de "vida", a ciência tem avançado significativamente na compreensão dos mecanismos biológicos que sustentam a vida.

5. **Definição Contemporânea**: Devido à complexidade e diversidade da vida no universo, uma definição ampla e aceita de "vida" pode incluir desde organismos simples até sistemas inteligentes e até mesmo conceitos metafísicos de existência ou consciência.

Em resumo, a natureza da vida é um tema multifacetado que transcende disciplinas e culturas, refletindo a riqueza e complexidade da experiência humana e do universo.

===========
++controle
===========
Umm... (risos) Ahah, essa é uma pergunta que já foi feita milhares de vezes, né? (sorriso) Mas, vamos lá, tentar responder... (pausa dramática)

A vida... é como um show, né? É uma grande festa, cheia de coisas louca, louco, louuuuco! (risos) É quando você se levanta da cama, e começa a dançar, a cantar, a gritar, a chorar... (gigantesco gesto) É quando você encontra alguém, e... (sussurrinho) e... (pausa dramática) e... (risos) e... (gigante) e... (choro) OH, A VIDA É UMA FESTA, NÉ?! (risos)

Mas, sério... (pausa) A vida é... é... (suspiro) é o que faz você viver, né? É o que faz você sorrir, ronronhão, chorar, gritar... (risos) É o que faz você amar, odeiar, querer, não querer... (gesto de amor e ódio) É o que faz você... (pausa dramática) É o que faz você... (sussurrinho) ser... (risos) SER ALUÉM, NÉ?! (risos)

E, ahah, e a vida é... (pausa) é curta, né? É curta, curta, curta... (gesto de tempo passando) É como um bolo, né? É como um bolo, que você come, e... (sussurrinho) e... (pausa dramática) e... (risos) e... (choro) OH, A VIDA É CURTA, NÉ?! (risos)

Então... (pausa) Então, a vida é... (sussurrinho) é uma coisa louca, louco, louuuuco... (risos) É uma coisa que você não pode explicar, né? É uma coisa que você só pode sentir... (gesto de emoção) É uma coisa que você só pode... (pausa dramática) SÓ COME, NÉ?! (risos)

E, ahah, e a vida é... (pausa) é uma coisa que você não pode parar... (gesto de parar) Não, não, não... (risos) É uma coisa que você só pode continuar... (gesto de continuar) Continua, continua, continua... (choro) OH, A VIDA É CONTINUA, NÉ?! (risos)

Então... (pausa) Então, a vida é... (sussurrinho) é uma coisa que você não pode explicar... (risos) É uma coisa que você só pode sentir... (gesto de emoção) É uma coisa que você só pode... (pausa dramática) SÓ COME, NÉ?! (risos)

(Choro) OH, A VIDA É UMA FESTA, NÉ?! (risos)

===========
--controle
===========
A pergunta "O que é a vida?" é fundamentalmente filosófica e pode ser abordada desde diversas perspectivas. Aqui estão algumas considerações:

1. **Definição biológica**: Na biologia, a vida é geralmente definida como um conjunto de processos e características que permitem ao organismo se manter vivo. Isso inclui a capacidade de crescer, se reproduzir e responder às mudanças ambientais.

2. **Perspectiva filosófica**: A filosofia oferece diferentes conceitos sobre a natureza da vida. Por exemplo:
   - **Existencialismo**: Para o existencialismo, a vida é vista como uma criação individual, com cada pessoa responsável por criar seu próprio propósito e significado.
   - **Humanismo**: O humanismo enfatiza os aspectos racionais e éticos da experiência humana, destacando a importância do pensamento, da razão e do respeito pelos princípios morais.
   - **Eudemo**: A teoria eudêmica, desenvolvida por Aristóteles, sugere que a vida tem um objetivo principal: alcançar a felicidade ou eudemo, que é entendido como um estado de plenitude e realização.

3. **Conceitos religiosos e espirituais**: Em muitas tradições religiosas e espirituais, a vida é vista como um caminho para alcançar a salvação, a iluminação ou a união com uma força superior.
4. **Perspectiva científica**: A ciência moderna tende a abordar a vida como um processo complexo, influenciado por fatores biológicos, ambientais e sociais. Ela busca entender as leis e mecanismos subjacentes que sustentam a vida, mas não fornece uma definição intrínseca do conceito.

Em última análise, a resposta à pergunta "O que é a vida?" depende das suas próprias crenças, valores e perspectivas individuais.

Conclusion

Representation Engineering offers a different way to control the outputs of language models. In this article, I introduced the concept of a contrast vector and showed how to easily train your own vector to control a model like Llama-3.2-3B-Instruct.

The examples presented here are in Portuguese, but since this is a small model, the amount of Portuguese data used during training is much smaller than in English. Because of this, the results when controlling prompts in English are much more interesting. It’s worth editing the notebook to test model control using prompts in English.

Another limitation of these examples is that the template used to train the vector was always the same: “Please act as …”. Different templates produce different results, so it’s worth exploring more deeply how to improve the training of these vectors.

I invite everyone to experiment with contrast vectors, and please send me the results you obtain!