Activation Steering: Steering Vectors and Representation Engineering › Code Implementation

from transformers import AutoModelForCausalLM, AutoTokenizer

# Load a small model for demonstration
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, output_hidden_states=True
)
model.eval()

# Set padding token
tokenizer.pad_token = tokenizer.eos_token

n_layers = model.config.n_layer
hidden_dim = model.config.n_embd

import numpy as np
import torch


def get_activations(prompts, layer_idx, token_position=-1):
    """
    Extract residual stream activations at a specific layer and token position.

    Args:
        prompts: list of text strings
        layer_idx: which transformer layer to extract from (0-indexed)
        token_position: which token position to use (-1 = last token)

    Returns:
        activations: numpy array of shape (n_prompts, hidden_dim)
    """
    activations = []

    for prompt in prompts:
        inputs = tokenizer(
            prompt, return_tensors="pt", padding=True, truncation=True
        )

        with torch.no_grad():
            outputs = model(**inputs)

        # hidden_states is a tuple: (embedding_output, layer_1_output, ..., layer_N_output)
        # Index layer_idx + 1 because index 0 is the embedding layer
        hidden_states = outputs.hidden_states[layer_idx + 1]

        # Extract the activation at the specified token position
        # hidden_states shape: (batch, seq_len, hidden_dim)
        act = hidden_states[0, token_position, :].cpu().numpy()
        activations.append(act)

    return np.array(activations)

# Define contrasting prompt sets for a "confident" vs "uncertain" concept
positive_prompts = [
    "I am absolutely certain that",
    "Without any doubt, the answer is",
    "I know for a fact that",
    "It is definitively the case that",
    "The clear and undeniable truth is that",
]

negative_prompts = [
    "I'm not really sure, but maybe",
    "I think it might possibly be",
    "I'm uncertain, but perhaps",
    "It could be, I suppose, that",
    "I'm not confident, but it seems like",
]

TARGET_LAYER = 8  # Middle-late layer in GPT-2

# Extract activations for both sets
pos_activations = get_activations(positive_prompts, layer_idx=TARGET_LAYER)
neg_activations = get_activations(negative_prompts, layer_idx=TARGET_LAYER)

# Compute steering vector as mean difference
steering_vector = pos_activations.mean(axis=0) - neg_activations.mean(axis=0)

# Normalize to unit length
steering_vector_norm = steering_vector / np.linalg.norm(steering_vector)

import torch.nn.functional as F


def get_next_token_probs(prompt, steering_vec, alpha, layer_idx=8, top_k=5):
    """
    Run a single forward pass and return the top-k next-token probabilities.
    With alpha=0 this is the baseline; non-zero alpha applies the steering vector.
    """
    inputs = tokenizer(prompt, return_tensors="pt")
    steering_tensor = torch.tensor(steering_vec, dtype=torch.float32)

    def hook(module, input, output):
        if isinstance(output, tuple):
            hidden = output[0]
            modified = hidden + alpha * steering_tensor.unsqueeze(0).unsqueeze(
                0
            )
            return (modified,) + output[1:]
        return output + alpha * steering_tensor.unsqueeze(0).unsqueeze(0)

    target_block = model.transformer.h[layer_idx]
    hook_handle = target_block.register_forward_hook(hook)

    try:
        with torch.no_grad():
            outputs = model(**inputs)
    finally:
        hook_handle.remove()

    # logits shape: (batch=1, seq_len, vocab_size)
    # Take the last token's logits to predict the next token
    logits = outputs.logits[0, -1, :]
    probs = F.softmax(logits, dim=-1)
    top_probs, top_ids = torch.topk(probs, top_k)

    top_tokens = [tokenizer.decode([tid.item()]) for tid in top_ids]
    top_prob_values = top_probs.cpu().numpy().tolist()
    return list(zip(top_tokens, top_prob_values))


# Compare next-token distributions under baseline and steering
test_prompt = "I am absolutely"

baseline_top5 = get_next_token_probs(
    test_prompt, steering_vector_norm, alpha=0.0
)
confident_top5 = get_next_token_probs(
    test_prompt, steering_vector_norm, alpha=20.0
)
uncertain_top5 = get_next_token_probs(
    test_prompt, steering_vector_norm, alpha=-20.0
)

# Build a second vector for "positive emotional tone" as a comparison
positive_tone_prompts = [
    "This is a wonderful day and I feel great",
    "Everything is going perfectly and I am happy",
    "I love this topic and find it fascinating",
]

negative_tone_prompts = [
    "This is a terrible day and I feel awful",
    "Everything is going wrong and I am miserable",
    "I dislike this topic and find it boring",
]

pos_tone_acts = get_activations(positive_tone_prompts, layer_idx=TARGET_LAYER)
neg_tone_acts = get_activations(negative_tone_prompts, layer_idx=TARGET_LAYER)

tone_vector = pos_tone_acts.mean(axis=0) - neg_tone_acts.mean(axis=0)
tone_vector_norm = tone_vector / np.linalg.norm(tone_vector)

# Cosine similarity between confidence and tone vectors
cosine_sim = float(np.dot(steering_vector_norm, tone_vector_norm))