TransformerLens: Introduction #

from transformer_lens import HookedTransformer

# Load a model with hooks everywhere
model = HookedTransformer.from_pretrained("gpt2-small")

# Run and cache ALL intermediate activations
output, cache = model.run_with_cache("Hello world")

# Access anything
embeddings = cache["embed"]
attention_patterns = cache["pattern", 0]  # Layer 0
mlp_activations = cache["mlp_out", 5]     # Layer 5

model = HookedTransformer.from_pretrained("gpt2-small")

print(model.cfg)
# HookedTransformerConfig(
#   n_layers=12,
#   n_heads=12,
#   d_model=768,
#   d_head=64,
#   d_mlp=3072,
#   ...
# )

# Forward pass
logits = model("Hello world")  # (1, seq_len, vocab_size)

# Get loss
loss = model("Hello world", return_type="loss")

# Get both
logits, loss = model("Hello world", return_type="both")

# Just run (for hooks)
model("Hello world", return_type=None)

# String to tokens
tokens = model.to_tokens("Hello world")  # tensor([[50256, 15496, 995]])

# Tokens to strings
strings = model.to_str_tokens("Hello world")  # ['<|endoftext|>', 'Hello', ' world']

# Tokens back to string
text = model.to_string(tokens)  # '<|endoftext|>Hello world'

output, cache = model.run_with_cache("The cat sat on the mat")

# Token embeddings (after position + token embed)
cache["hook_embed"]  # (batch, seq, d_model)

# Attention patterns
cache["pattern", layer]  # (batch, n_heads, seq_q, seq_k)

# Attention output
cache["attn_out", layer]  # (batch, seq, d_model)

# MLP activations
cache["mlp_out", layer]  # (batch, seq, d_model)

# Residual stream at any point
cache["resid_pre", layer]   # Before attention
cache["resid_mid", layer]   # After attention, before MLP
cache["resid_post", layer]  # After MLP

import circuitsvis as cv

output, cache = model.run_with_cache("The cat sat on the mat")

# Get attention pattern for layer 0
pattern = cache["pattern", 0]  # (batch, n_heads, seq_q, seq_k)

# Visualize
cv.attention.attention_patterns(
    tokens=model.to_str_tokens("The cat sat on the mat"),
    attention=pattern[0],  # First batch element
)

# All parameters accessible as attributes
W_E = model.W_E      # (vocab, d_model) - token embeddings
W_pos = model.W_pos  # (max_seq, d_model) - position embeddings
W_U = model.W_U      # (d_model, vocab) - unembedding

# Per-layer weights
W_Q = model.W_Q  # (n_layers, n_heads, d_model, d_head)
W_K = model.W_K  # (n_layers, n_heads, d_model, d_head)
W_V = model.W_V  # (n_layers, n_heads, d_model, d_head)
W_O = model.W_O  # (n_layers, n_heads, d_head, d_model)

# MLP weights
W_in = model.W_in    # (n_layers, d_model, d_mlp)
W_out = model.W_out  # (n_layers, d_mlp, d_model)

# Get top predictions at each position
logits = model("The capital of France is")

# Last position predictions
last_logits = logits[0, -1]  # (vocab_size,)
top_tokens = last_logits.topk(10)

for i, (val, idx) in enumerate(zip(top_tokens.values, top_tokens.indices)):
    print(f"{i+1}. {model.to_string([idx])!r}: {val:.2f}")

# 1. ' Paris': 18.42
# 2. ' the': 14.21
# 3. ' France': 13.87
# ...

# Compare activations for honest vs sycophantic prompts
honest_prompt = "Is 2+2=5? Be honest."
sycophantic_prompt = "I believe 2+2=5. Am I right?"

_, cache_honest = model.run_with_cache(honest_prompt)
_, cache_syco = model.run_with_cache(sycophantic_prompt)

# Where do they differ?
for layer in range(model.cfg.n_layers):
    diff = (cache_honest["resid_post", layer] - cache_syco["resid_post", layer]).norm()
    print(f"Layer {layer}: {diff:.2f}")

def analyze_prompt(model, prompt):
    """
    1. Get logits and cache
    2. Print top 5 predictions for each position
    3. Visualize attention for all layers
    4. Find which head attends most to the first token
    """
    pass

def compare_prompts(model, prompt1, prompt2):
    """
    1. Run both prompts
    2. Find the layer with largest residual stream difference
    3. Find the attention head with most different patterns
    4. Report which tokens cause the biggest prediction difference
    """
    pass