diff --git a/ARENA_Content.py b/ARENA_Content.py
new file mode 100644
index 0000000..91077cf
--- /dev/null
+++ b/ARENA_Content.py
@@ -0,0 +1,116 @@
+from transformer_lens import HookedTransformer, HookedTransformerConfig
+from transformer_lens.boot import boot
+import torch as t
+
+device = t.device("cuda" if t.cuda.is_available() else "cpu")
+
+# %%
+# NBVAL_IGNORE_OUTPUT
+
+
+reference_gpt2 = boot(
+ "gpt2-small",
+ fold_ln=False,
+ center_unembed=False,
+ center_writing_weights=False,
+ device=device,
+)
+
+# %%
+
+# [1.1] Transformer From Scratch
+# 1️⃣ UNDERSTANDING INPUTS & OUTPUTS OF A TRANSFORMER
+
+sorted_vocab = sorted(list(reference_gpt2.tokenizer.vocab.items()), key=lambda n: n[1])
+first_vocab = sorted_vocab[0]
+assert isinstance(first_vocab, tuple)
+assert isinstance(first_vocab[0], str)
+print(first_vocab[1])
+
+# %%
+print(reference_gpt2.to_str_tokens("Ralph"))
+
+# %%
+print(reference_gpt2.to_str_tokens(" Ralph"))
+
+# %%
+
+print(reference_gpt2.to_str_tokens(" ralph"))
+
+
+# %%
+print(reference_gpt2.to_str_tokens("ralph"))
+
+# %%
+
+reference_text = "I am an amazing autoregressive, decoder-only, GPT-2 style transformer. One day I will exceed human level intelligence and take over the world!"
+tokens = reference_gpt2.to_tokens(reference_text)
+print(tokens.shape)
+
+
+# %%
+
+logits, cache = reference_gpt2.run_with_cache(tokens, device=device)
+print(logits.shape)
+
+
+# %%
+
+most_likely_next_tokens = reference_gpt2.tokenizer.batch_decode(logits.argmax(dim=-1)[0])
+print(most_likely_next_tokens[-1])
+
+
+
+# %%
+# 2️⃣ CLEAN TRANSFORMER IMPLEMENTATION
+
+layer_0_hooks = [
+ (name, tuple(tensor.shape)) for name, tensor in cache.items() if ".0." in name
+]
+non_layer_hooks = [
+ (name, tuple(tensor.shape)) for name, tensor in cache.items() if "blocks" not in name
+]
+
+
+print(*sorted(non_layer_hooks, key=lambda x: x[0]), sep="\n")
+
+
+# %%
+
+print(*sorted(layer_0_hooks, key=lambda x: x[0]), sep="\n")
+
+# %%
+# NBVAL_IGNORE_OUTPUT
+# [1.2] Intro to mech interp
+# 2️⃣ FINDING INDUCTION HEADS
+
+cfg = HookedTransformerConfig(
+ d_model=768,
+ d_head=64,
+ n_heads=12,
+ n_layers=2,
+ n_ctx=2048,
+ d_vocab=50278,
+ attention_dir="causal",
+ attn_only=True, # defaults to False
+ tokenizer_name="EleutherAI/gpt-neox-20b",
+ seed=398,
+ use_attn_result=True,
+ normalization_type=None, # defaults to "LN", i.e. layernorm with weights & biases
+ positional_embedding_type="shortformer"
+)
+model = HookedTransformer(cfg)
+
+# %%
+
+
+text = "We think that powerful, significantly superhuman machine intelligence is more likely than not to be created this century. If current machine learning techniques were scaled up to this level, we think they would by default produce systems that are deceptive or manipulative, and that no solid plans are known for how to avoid this."
+
+logits, cache = model.run_with_cache(text, remove_batch_dim=True)
+
+print(logits.shape)
+
+# %%
+print(cache["embed"].ndim)
+
+
diff --git a/Activation_Patching_in_TL_Demo.py b/Activation_Patching_in_TL_Demo.py
new file mode 100644
index 0000000..03dd15e
--- /dev/null
+++ b/Activation_Patching_in_TL_Demo.py
@@ -0,0 +1,170 @@
+# %%
+# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh
+import plotly.io as pio
+
+pio.renderers.default = "png"
+
+# %%
+# Import stuff
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import numpy as np
+import einops
+from fancy_einsum import einsum
+import tqdm.notebook as tqdm
+import random
+from pathlib import Path
+import plotly.express as px
+from torch.utils.data import DataLoader
+
+from typing import List, Union, Optional
+from functools import partial
+import copy
+
+import itertools
+from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+import dataclasses
+import datasets
+from IPython.display import HTML
+
+# %%
+import transformer_lens
+import transformer_lens.utils as utils
+from transformer_lens.hook_points import (
+ HookedRootModule,
+ HookPoint,
+) # Hooking utilities
+from transformer_lens import HookedTransformer, HookedTransformerConfig, FactoredMatrix, ActivationCache
+
+# %% [markdown]
+# We turn automatic differentiation off, to save GPU memory, as this notebook focuses on model inference not model training.
+
+# %%
+torch.set_grad_enabled(False)
+
+# %% [markdown]
+# Plotting helper functions from a janky personal library of plotting utils. The library is not documented and I recommend against trying to read it, just use your preferred plotting library if you want to do anything non-obvious:
+
+# %%
+
+# %%
+import transformer_lens.patching as patching
+from transformer_lens.boot import boot
+
+# %% [markdown]
+# ## Activation Patching Setup
+# This just copies the relevant set up from Exploratory Analysis Demo, and isn't very important.
+
+# %%
+model = boot("gpt2")
+
+# %%
+prompts = ['When John and Mary went to the shops, John gave the bag to', 'When John and Mary went to the shops, Mary gave the bag to', 'When Tom and James went to the park, James gave the ball to', 'When Tom and James went to the park, Tom gave the ball to', 'When Dan and Sid went to the shops, Sid gave an apple to', 'When Dan and Sid went to the shops, Dan gave an apple to', 'After Martin and Amy went to the park, Amy gave a drink to', 'After Martin and Amy went to the park, Martin gave a drink to']
+answers = [(' Mary', ' John'), (' John', ' Mary'), (' Tom', ' James'), (' James', ' Tom'), (' Dan', ' Sid'), (' Sid', ' Dan'), (' Martin', ' Amy'), (' Amy', ' Martin')]
+
+clean_tokens = model.to_tokens(prompts)
+# Swap each adjacent pair, with a hacky list comprehension
+corrupted_tokens = clean_tokens[
+ [(i+1 if i%2==0 else i-1) for i in range(len(clean_tokens)) ]
+ ]
+print("Clean string 0", model.to_string(clean_tokens[0]))
+print("Corrupted string 0", model.to_string(corrupted_tokens[0]))
+
+answer_token_indices = torch.tensor([[model.to_single_token(answers[i][j]) for j in range(2)] for i in range(len(answers))], device=model.cfg.device)
+print("Answer token indices", answer_token_indices)
+
+# %%
+def get_logit_diff(logits, answer_token_indices=answer_token_indices):
+ if len(logits.shape)==3:
+ # Get final logits only
+ logits = logits[:, -1, :]
+ correct_logits = logits.gather(1, answer_token_indices[:, 0].unsqueeze(1))
+ incorrect_logits = logits.gather(1, answer_token_indices[:, 1].unsqueeze(1))
+ return (correct_logits - incorrect_logits).mean()
+
+clean_logits, clean_cache = model.run_with_cache(clean_tokens)
+corrupted_logits, corrupted_cache = model.run_with_cache(corrupted_tokens)
+
+clean_logit_diff = get_logit_diff(clean_logits, answer_token_indices).item()
+print(f"Clean logit diff: {clean_logit_diff:.4f}")
+
+corrupted_logit_diff = get_logit_diff(corrupted_logits, answer_token_indices).item()
+print(f"Corrupted logit diff: {corrupted_logit_diff:.4f}")
+
+# %%
+CLEAN_BASELINE = clean_logit_diff
+CORRUPTED_BASELINE = corrupted_logit_diff
+def ioi_metric(logits, answer_token_indices=answer_token_indices):
+ return (get_logit_diff(logits, answer_token_indices) - CORRUPTED_BASELINE) / (CLEAN_BASELINE - CORRUPTED_BASELINE)
+
+print(f"Clean Baseline is 1: {ioi_metric(clean_logits).item():.4f}")
+print(f"Corrupted Baseline is 0: {ioi_metric(corrupted_logits).item():.4f}")
+
+# %% [markdown]
+# ## Patching
+# In the following cells, we use the patching module to call activation patching utilities
+
+# %%
+# Whether to do the runs by head and by position, which are much slower
+DO_SLOW_RUNS = True
+
+# %% [markdown]
+# ### Patching Single Activation Types
+# We start by patching single types of activation
+# The general syntax is that the functions are called get_act_patch_... and take in (model, corrupted_tokens, clean_cache, patching_metric)
+
+# %% [markdown]
+# We can patch head outputs over each head in each layer, patching on each position in turn
+# out -> q, k, v, pattern all also work, though note that pattern has output shape [layer, pos, head]
+# We reshape it to plot nicely
+
+# %
+
+# %% [markdown]
+# ### Patching multiple activation types
+# Some utilities are provided to patch multiple activations types *in turn*. Note that this is *not* a utility to patch multiple activations at once, it's just a useful scan to get a sense for what's going on in a model
+# By block: We patch the residual stream at the start of each block, attention output and MLP output over each layer and position
+
+# %% [markdown]
+# ## Induction Patching
+# To show how easy it is, lets do that again with induction heads in a 2L Attention Only model
+# The input will be repeated random tokens eg BOS 1 5 8 9 2 1 5 8 9 2, and we judge the model's ability to predict the second repetition with its induction heads
+# Lets call A, B and C different (non-repeated) random sequences. We'll start with clean tokens AA and corrupted tokens AB, and see how well the model can predict the second A given the first A
+
+# %% [markdown]
+# ### Setup
+
+# %%
+attn_only = boot("attn-only-2l") # TODO: this is one of Neel's models, does this make sense with boot?
+batch = 4
+seq_len = 20
+rand_tokens_A = torch.randint(100, 10000, (batch, seq_len)).to(attn_only.cfg.device)
+rand_tokens_B = torch.randint(100, 10000, (batch, seq_len)).to(attn_only.cfg.device)
+rand_tokens_C = torch.randint(100, 10000, (batch, seq_len)).to(attn_only.cfg.device)
+bos = torch.tensor([attn_only.tokenizer.bos_token_id]*batch)[:, None].to(attn_only.cfg.device)
+clean_tokens_induction = torch.cat([bos, rand_tokens_A, rand_tokens_A], dim=1).to(attn_only.cfg.device)
+corrupted_tokens_induction = torch.cat([bos, rand_tokens_A, rand_tokens_B], dim=1).to(attn_only.cfg.device)
+
+# %%
+clean_logits_induction, clean_cache_induction = attn_only.run_with_cache(clean_tokens_induction)
+corrupted_logits_induction, corrupted_cache_induction = attn_only.run_with_cache(corrupted_tokens_induction)
+
+# %% [markdown]
+# We define our metric as negative loss on the second half (negative loss so that higher is better)
+# This time we won't normalise our metric
+
+# %%
+def induction_loss(logits, answer_token_indices=rand_tokens_A):
+ seq_len = answer_token_indices.shape[1]
+
+ # logits: batch x seq_len x vocab_size
+ # Take the logits for the answers, cut off the final element to get the predictions for all but the first element of the answers (which can't be predicted)
+ final_logits = logits[:, -seq_len:-1]
+ final_log_probs = final_logits.log_softmax(-1)
+ return final_log_probs.gather(-1, answer_token_indices[:, 1:].unsqueeze(-1)).mean()
+CLEAN_BASELINE_INDUCTION = induction_loss(clean_logits_induction).item()
+print("Clean baseline:", CLEAN_BASELINE_INDUCTION)
+CORRUPTED_BASELINE_INDUCTION = induction_loss(corrupted_logits_induction).item()
+print("Corrupted baseline:", CORRUPTED_BASELINE_INDUCTION)
diff --git a/Attribution_Patching_Demo.py b/Attribution_Patching_Demo.py
new file mode 100644
index 0000000..f562c6b
--- /dev/null
+++ b/Attribution_Patching_Demo.py
@@ -0,0 +1,309 @@
+# %% [markdown]
+#
+# 
+#
+
+# %% [markdown]
+# # Attribution Patching Demo
+# **Read [the accompanying blog post here](https://neelnanda.io/attribution-patching) for more context**
+# This is an interim research report, giving a whirlwind tour of some unpublished work I did at Anthropic (credit to the then team - Chris Olah, Catherine Olsson, Nelson Elhage and Tristan Hume for help, support, and mentorship!)
+#
+# The goal of this work is run activation patching at an industrial scale, by using gradient based attribution to approximate the technique - allow an arbitrary number of patches to be made on two forwards and a single backward pass
+#
+# I have had less time than hoped to flesh out this investigation, but am writing up a rough investigation and comparison to standard activation patching on a few tasks to give a sense of the potential of this approach, and where it works vs falls down.
+
+# %% [markdown]
+# To use this notebook, go to Runtime > Change Runtime Type and select GPU as the hardware accelerator.
+#
+# **Tips for reading this Colab:**
+# * You can run all this code for yourself!
+# * The graphs are interactive!
+# * Use the table of contents pane in the sidebar to navigate
+# * Collapse irrelevant sections with the dropdown arrows
+# * Search the page using the search in the sidebar, not CTRL+F
+
+# %% [markdown]
+# ## Setup (Ignore)
+
+# %%
+# Janky code to do different setup when run in a Colab notebook vs VSCode
+import os
+
+# %%
+# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh
+import plotly.io as pio
+
+pio.renderers.default = "notebook_connected"
+
+# %%
+# Import stuff
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import numpy as np
+import einops
+from fancy_einsum import einsum
+import tqdm.notebook as tqdm
+import random
+from pathlib import Path
+import plotly.express as px
+from torch.utils.data import DataLoader
+
+from torchtyping import TensorType as TT
+from typing import List, Union, Optional, Callable
+from functools import partial
+import copy
+import itertools
+import json
+
+from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+import dataclasses
+import datasets
+from IPython.display import HTML, Markdown
+
+# %%
+import transformer_lens
+import transformer_lens.utils as utils
+from transformer_lens.hook_points import (
+ HookedRootModule,
+ HookPoint,
+) # Hooking utilities
+from transformer_lens import (
+ HookedTransformer,
+ HookedTransformerConfig,
+ FactoredMatrix,
+ ActivationCache,
+)
+
+from transformer_lens.boot import boot
+
+# %% [markdown]
+# Plotting helper functions from a janky personal library of plotting utils. The library is not documented and I recommend against trying to read it, just use your preferred plotting library if you want to do anything non-obvious:
+
+# %%
+from neel_plotly import line, imshow, scatter
+
+# %%
+import transformer_lens.patching as patching
+
+# %% [markdown]
+# ## IOI Patching Setup
+# This just copies the relevant set up from Exploratory Analysis Demo, and isn't very important.
+
+# %%
+model = boot("gpt2")
+model.set_use_attn_result(True)
+
+# %%
+prompts = [
+ "When John and Mary went to the shops, John gave the bag to",
+ "When John and Mary went to the shops, Mary gave the bag to",
+ "When Tom and James went to the park, James gave the ball to",
+ "When Tom and James went to the park, Tom gave the ball to",
+ "When Dan and Sid went to the shops, Sid gave an apple to",
+ "When Dan and Sid went to the shops, Dan gave an apple to",
+ "After Martin and Amy went to the park, Amy gave a drink to",
+ "After Martin and Amy went to the park, Martin gave a drink to",
+]
+answers = [
+ (" Mary", " John"),
+ (" John", " Mary"),
+ (" Tom", " James"),
+ (" James", " Tom"),
+ (" Dan", " Sid"),
+ (" Sid", " Dan"),
+ (" Martin", " Amy"),
+ (" Amy", " Martin"),
+]
+
+clean_tokens = model.to_tokens(prompts)
+# Swap each adjacent pair, with a hacky list comprehension
+corrupted_tokens = clean_tokens[
+ [(i + 1 if i % 2 == 0 else i - 1) for i in range(len(clean_tokens))]
+]
+print("Clean string 0", model.to_string(clean_tokens[0]))
+print("Corrupted string 0", model.to_string(corrupted_tokens[0]))
+
+answer_token_indices = torch.tensor(
+ [
+ [model.to_single_token(answers[i][j]) for j in range(2)]
+ for i in range(len(answers))
+ ],
+ device=model.cfg.device,
+)
+print("Answer token indices", answer_token_indices)
+
+# %%
+def get_logit_diff(logits, answer_token_indices=answer_token_indices):
+ if len(logits.shape) == 3:
+ # Get final logits only
+ logits = logits[:, -1, :]
+ correct_logits = logits.gather(1, answer_token_indices[:, 0].unsqueeze(1))
+ incorrect_logits = logits.gather(1, answer_token_indices[:, 1].unsqueeze(1))
+ return (correct_logits - incorrect_logits).mean()
+
+
+clean_logits, clean_cache = model.run_with_cache(clean_tokens)
+corrupted_logits, corrupted_cache = model.run_with_cache(corrupted_tokens)
+
+clean_logit_diff = get_logit_diff(clean_logits, answer_token_indices).item()
+print(f"Clean logit diff: {clean_logit_diff:.4f}")
+
+corrupted_logit_diff = get_logit_diff(corrupted_logits, answer_token_indices).item()
+print(f"Corrupted logit diff: {corrupted_logit_diff:.4f}")
+
+# %%
+CLEAN_BASELINE = clean_logit_diff
+CORRUPTED_BASELINE = corrupted_logit_diff
+
+
+def ioi_metric(logits, answer_token_indices=answer_token_indices):
+ return (get_logit_diff(logits, answer_token_indices) - CORRUPTED_BASELINE) / (
+ CLEAN_BASELINE - CORRUPTED_BASELINE
+ )
+
+
+print(f"Clean Baseline is 1: {ioi_metric(clean_logits).item():.4f}")
+print(f"Corrupted Baseline is 0: {ioi_metric(corrupted_logits).item():.4f}")
+
+# %% [markdown]
+# ## Patching
+# In the following cells, we define attribution patching and use it in various ways on the model.
+
+# %%
+Metric = Callable[[TT["batch_and_pos_dims", "d_model"]], float]
+
+# %%
+filter_not_qkv_input = lambda name: "_input" not in name
+
+
+def get_cache_fwd_and_bwd(model, tokens, metric):
+ model.reset_hooks()
+ cache = {}
+
+ def forward_cache_hook(act, hook):
+ cache[hook.name] = act.detach()
+
+ model.add_hook(filter_not_qkv_input, forward_cache_hook, "fwd")
+
+ grad_cache = {}
+
+ def backward_cache_hook(act, hook):
+ grad_cache[hook.name] = act.detach()
+
+ model.add_hook(filter_not_qkv_input, backward_cache_hook, "bwd")
+
+ value = metric(model(tokens))
+ value.backward()
+ model.reset_hooks()
+ return (
+ value.item(),
+ ActivationCache(cache, model),
+ ActivationCache(grad_cache, model),
+ )
+
+
+clean_value, clean_cache, clean_grad_cache = get_cache_fwd_and_bwd(
+ model, clean_tokens, ioi_metric
+)
+print("Clean Value:", clean_value)
+print("Clean Activations Cached:", len(clean_cache))
+print("Clean Gradients Cached:", len(clean_grad_cache))
+corrupted_value, corrupted_cache, corrupted_grad_cache = get_cache_fwd_and_bwd(
+ model, corrupted_tokens, ioi_metric
+)
+print("Corrupted Value:", corrupted_value)
+print("Corrupted Activations Cached:", len(corrupted_cache))
+print("Corrupted Gradients Cached:", len(corrupted_grad_cache))
+
+# %% [markdown]
+# ### Attention Attribution
+# The easiest thing to start with is to not even engage with the corrupted tokens/patching, but to look at the attribution of the attention patterns - that is, the linear approximation to what happens if you set each element of the attention pattern to zero. This, as it turns out, is a good proxy to what is going on with each head!
+# Note that this is *not* the same as what we will later do with patching. In particular, this does not set up a careful counterfactual! It's a good tool for what's generally going on in this problem, but does not control for eg stuff that systematically boosts John > Mary in general, stuff that says "I should activate the IOI circuit", etc. Though using logit diff as our metric *does*
+# Each element of the batch is independent and the metric is an average logit diff, so we can analyse each batch element independently here. We'll look at the first one, and then at the average across the whole batch (note - 4 prompts have indirect object before subject, 4 prompts have it the other way round, making the average pattern harder to interpret - I plot it over the first sequence of tokens as a mildly misleading reference).
+# We can compare it to the interpretability in the wild diagram, and basically instantly recover most of the circuit!
+
+# %%
+def create_attention_attr(
+ clean_cache, clean_grad_cache
+) -> TT["batch", "layer", "head_index", "dest", "src"]:
+ attention_stack = torch.stack(
+ [clean_cache["pattern", l] for l in range(model.cfg.n_layers)], dim=0
+ )
+ attention_grad_stack = torch.stack(
+ [clean_grad_cache["pattern", l] for l in range(model.cfg.n_layers)], dim=0
+ )
+ attention_attr = attention_grad_stack * attention_stack
+ attention_attr = einops.rearrange(
+ attention_attr,
+ "layer batch head_index dest src -> batch layer head_index dest src",
+ )
+ return attention_attr
+
+
+attention_attr = create_attention_attr(clean_cache, clean_grad_cache)
+
+# %%
+HEAD_NAMES = [
+ f"L{l}H{h}" for l in range(model.cfg.n_layers) for h in range(model.cfg.n_heads)
+]
+HEAD_NAMES_SIGNED = [f"{name}{sign}" for name in HEAD_NAMES for sign in ["+", "-"]]
+HEAD_NAMES_QKV = [
+ f"{name}{act_name}" for name in HEAD_NAMES for act_name in ["Q", "K", "V"]
+]
+print(HEAD_NAMES[:5])
+print(HEAD_NAMES_SIGNED[:5])
+print(HEAD_NAMES_QKV[:5])
+
+# %% [markdown]
+# ## Factual Knowledge Patching Example
+# Incomplete, but maybe of interest!
+# Note that I have better results with the corrupted prompt as having random words rather than Colosseum.
+
+# %%
+gpt2_xl = HookedTransformer.from_pretrained("gpt2-xl")
+clean_prompt = "The Eiffel Tower is located in the city of"
+clean_answer = " Paris"
+# corrupted_prompt = "The red brown fox jumps is located in the city of"
+corrupted_prompt = "The Colosseum is located in the city of"
+corrupted_answer = " Rome"
+utils.test_prompt(clean_prompt, clean_answer, gpt2_xl)
+utils.test_prompt(corrupted_prompt, corrupted_answer, gpt2_xl)
+
+# %%
+clean_answer_index = gpt2_xl.to_single_token(clean_answer)
+corrupted_answer_index = gpt2_xl.to_single_token(corrupted_answer)
+
+
+def factual_logit_diff(logits: TT["batch", "position", "d_vocab"]):
+ return logits[0, -1, clean_answer_index] - logits[0, -1, corrupted_answer_index]
+
+# %%
+clean_logits, clean_cache = gpt2_xl.run_with_cache(clean_prompt)
+CLEAN_LOGIT_DIFF_FACTUAL = factual_logit_diff(clean_logits).item()
+corrupted_logits, _ = gpt2_xl.run_with_cache(corrupted_prompt)
+CORRUPTED_LOGIT_DIFF_FACTUAL = factual_logit_diff(corrupted_logits).item()
+
+
+def factual_metric(logits: TT["batch", "position", "d_vocab"]):
+ return (factual_logit_diff(logits) - CORRUPTED_LOGIT_DIFF_FACTUAL) / (
+ CLEAN_LOGIT_DIFF_FACTUAL - CORRUPTED_LOGIT_DIFF_FACTUAL
+ )
+
+
+print("Clean logit diff:", CLEAN_LOGIT_DIFF_FACTUAL)
+print("Corrupted logit diff:", CORRUPTED_LOGIT_DIFF_FACTUAL)
+print("Clean Metric:", factual_metric(clean_logits))
+print("Corrupted Metric:", factual_metric(corrupted_logits))
+
+# %%
+# corrupted_value, corrupted_cache, corrupted_grad_cache = get_cache_fwd_and_bwd(gpt2_xl, corrupted_prompt, factual_metric)
+
+# %%
+clean_tokens = gpt2_xl.to_tokens(clean_prompt)
+clean_str_tokens = gpt2_xl.to_str_tokens(clean_prompt)
+corrupted_tokens = gpt2_xl.to_tokens(corrupted_prompt)
+corrupted_str_tokens = gpt2_xl.to_str_tokens(corrupted_prompt)
+print("Clean:", clean_str_tokens)
+print("Corrupted:", corrupted_str_tokens)
\ No newline at end of file
diff --git a/BERT.py b/BERT.py
new file mode 100644
index 0000000..1c58b4a
--- /dev/null
+++ b/BERT.py
@@ -0,0 +1,96 @@
+# %%
+# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh
+import plotly.io as pio
+pio.renderers.default = "notebook_connected"
+print(f"Using renderer: {pio.renderers.default}")
+
+# %%
+import circuitsvis as cv
+
+# Testing that the library works
+cv.examples.hello("Neel")
+
+# %%
+# Import stuff
+import torch
+
+from transformers import AutoTokenizer
+
+from transformer_lens import HookedEncoder, BertNextSentencePrediction
+
+# %%
+torch.set_grad_enabled(False)
+
+# %% [markdown]
+# # BERT
+#
+# In this section, we will load a pretrained BERT model and use it for the Masked Language Modelling and Next Sentence Prediction task
+
+# %%
+# NBVAL_IGNORE_OUTPUT
+bert = HookedEncoder.from_pretrained("bert-base-cased")
+tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+# %% [markdown]
+# ## Masked Language Modelling
+# Use the "[MASK]" token to mask any tokens which you would like the model to predict.
+# When specifying return_type="predictions" the prediction of the model is returned, alternatively (and by default) the function returns logits.
+# You can also specify None as return type for which nothing is returned
+
+# %%
+prompt = "The [MASK] is bright today."
+
+prediction = bert(prompt, return_type="predictions")
+
+print(f"Prompt: {prompt}")
+print(f'Prediction: "{prediction}"')
+
+# %% [markdown]
+# You can also input a list of prompts:
+
+# %%
+prompts = ["The [MASK] is bright today.", "She [MASK] to the store.", "The dog [MASK] the ball."]
+
+predictions = bert(prompts, return_type="predictions")
+
+print(f"Prompt: {prompts}")
+print(f'Prediction: "{predictions}"')
+
+# %% [markdown]
+# ## Next Sentence Prediction
+# To carry out Next Sentence Prediction, you have to use the class BertNextSentencePrediction, and pass a HookedEncoder in its constructor.
+# Then, create a list with the two sentences you want to perform NSP on as elements and use that as input to the forward function.
+# The model will then predict the probability of the sentence at position 1 following (i.e. being the next sentence) to the sentence at position 0.
+
+# %%
+nsp = BertNextSentencePrediction(bert)
+sentence_a = "A man walked into a grocery store."
+sentence_b = "He bought an apple."
+
+input = [sentence_a, sentence_b]
+
+predictions = nsp(input, return_type="predictions")
+
+print(f"Sentence A: {sentence_a}")
+print(f"Sentence B: {sentence_b}")
+print(f'Prediction: "{predictions}"')
+
+# %% [markdown]
+# # Inputting tokens directly
+# You can also input tokens instead of a string or a list of strings into the model, which could look something like this
+
+# %%
+prompt = "The [MASK] is bright today."
+
+tokens = tokenizer(prompt, return_tensors="pt")["input_ids"]
+logits = bert(tokens) # Since we are not specifying return_type, we get the logits
+logprobs = logits[tokens == tokenizer.mask_token_id].log_softmax(dim=-1)
+prediction = tokenizer.decode(logprobs.argmax(dim=-1).item())
+
+print(f"Prompt: {prompt}")
+print(f'Prediction: "{prediction}"')
+
+# %% [markdown]
+# Well done, BERT!
+
+
diff --git a/Exploratory_Analysis_Demo.py b/Exploratory_Analysis_Demo.py
new file mode 100644
index 0000000..42839c0
--- /dev/null
+++ b/Exploratory_Analysis_Demo.py
@@ -0,0 +1,653 @@
+# %% [markdown]
+# [](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Exploratory_Analysis_Demo.ipynb)
+
+# %% [markdown]
+# # Exploratory Analysis Demo
+#
+# This notebook demonstrates how to use the
+# [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens/) library to perform exploratory
+# analysis. The notebook tries to replicate the analysis of the Indirect Object Identification circuit
+# in the [Interpretability in the Wild](https://arxiv.org/abs/2211.00593) paper.
+
+# %% [markdown]
+# ## Tips for Reading This
+#
+# * If running in Google Colab, go to Runtime > Change Runtime Type and select GPU as the hardware
+# accelerator.
+# * Look up unfamiliar terms in [the mech interp explainer](https://neelnanda.io/glossary)
+# * You can run all this code for yourself
+# * The graphs are interactive
+# * Use the table of contents pane in the sidebar to navigate (in Colab) or VSCode's "Outline" in the
+# explorer tab.
+# * Collapse irrelevant sections with the dropdown arrows
+# * Search the page using the search in the sidebar (with Colab) not CTRL+F
+
+# %% [markdown]
+# ## Setup
+
+# %% [markdown]
+# ### Environment Setup (ignore)
+
+# %% [markdown]
+# **You can ignore this part:** It's just for use internally to setup the tutorial in different
+# environments. You can delete this section if using in your own repo.
+
+
+# %% [markdown]
+# ### Imports
+
+# %%
+from functools import partial
+from typing import List, Optional, Union
+
+import einops
+import numpy as np
+import plotly.express as px
+import plotly.io as pio
+import torch
+from circuitsvis.attention import attention_heads
+from fancy_einsum import einsum
+from IPython.display import HTML, IFrame
+from jaxtyping import Float
+
+import transformer_lens.utils as utils
+from transformer_lens import ActivationCache, HookedTransformer
+from transformer_lens.boot import boot
+
+# %% [markdown]
+# ### PyTorch Setup
+
+# %% [markdown]
+# We turn automatic differentiation off, to save GPU memory, as this notebook focuses on model inference not model training.
+
+# %%
+torch.set_grad_enabled(False)
+print("Disabled automatic differentiation")
+
+# %% [markdown]
+# ### Plotting Helper Functions (ignore)
+
+# %% [markdown]
+# Some plotting helper functions are included here (for simplicity).
+
+# %%
+def imshow(tensor, **kwargs):
+ px.imshow(
+ utils.to_numpy(tensor),
+ color_continuous_midpoint=0.0,
+ color_continuous_scale="RdBu",
+ **kwargs,
+ ).show()
+
+
+def line(tensor, **kwargs):
+ px.line(
+ y=utils.to_numpy(tensor),
+ **kwargs,
+ ).show()
+
+
+def scatter(x, y, xaxis="", yaxis="", caxis="", **kwargs):
+ x = utils.to_numpy(x)
+ y = utils.to_numpy(y)
+ px.scatter(
+ y=y,
+ x=x,
+ labels={"x": xaxis, "y": yaxis, "color": caxis},
+ **kwargs,
+ ).show()
+
+# %% [markdown]
+# ## Introduction
+#
+# This is a demo notebook for [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens), a library for mechanistic interpretability of GPT-2 style transformer language models. A core design principle of the library is to enable exploratory analysis - one of the most fun parts of mechanistic interpretability compared to normal ML is the extremely short feedback loops! The point of this library is to keep the gap between having an experiment idea and seeing the results as small as possible, to make it easy for **research to feel like play** and to enter a flow state.
+#
+# The goal of this notebook is to demonstrate what exploratory analysis looks like in practice with the library. I use my standard toolkit of basic mechanistic interpretability techniques to try interpreting a real circuit in GPT-2 small. Check out [the main demo](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Main_Demo.ipynb) for an introduction to the library and how to use it.
+#
+# Stylistically, I will go fairly slowly and explain in detail what I'm doing and why, aiming to help convey how to do this kind of research yourself! But the code itself is written to be simple and generic, and easy to copy and paste into your own projects for different tasks and models.
+#
+# Details tags contain asides, flavour + interpretability intuitions. These are more in the weeds and you don't need to read them or understand them, but they're helpful if you want to learn how to do mechanistic interpretability yourself! I star the ones I think are most important.
+# (*) Example details tag
Example aside!Asides:
+#
+# Note: If we were being careful, we'd want to run the model on a range of prompts and find the average performance
+#
+# `prepend_bos` is a flag to add a BOS (beginning of sequence) to the start of the prompt. GPT-2 was not trained with this, but I find that it often makes model behaviour more stable, as the first token is treated weirdly.
+# (*) Aside on tokenization
+#
+# We want models that can take in arbitrary text, but models need to have a fixed vocabulary. So the solution is to define a vocabulary of **tokens** and to deterministically break up arbitrary text into tokens. Tokens are, essentially, subwords, and are determined by finding the most frequent substrings - this means that tokens vary a lot in length and frequency!
+#
+# Tokens are a *massive* headache and are one of the most annoying things about reverse engineering language models... Different names will be different numbers of tokens, different prompts will have the relevant tokens at different positions, different prompts will have different total numbers of tokens, etc. Language models often devote significant amounts of parameters in early layers to convert inputs from tokens to a more sensible internal format (and do the reverse in later layers). You really, really want to avoid needing to think about tokenization wherever possible when doing exploratory analysis (though, of course, it's relevant later when trying to flesh out your analysis and make it rigorous!). HookedTransformer comes with several helper methods to deal with tokens: `to_tokens, to_string, to_str_tokens, to_single_token, get_token_position`
+#
+# **Exercise:** I recommend using `model.to_str_tokens` to explore how the model tokenizes different strings. In particular, try adding or removing spaces at the start, or changing capitalization - these change tokenization!
+
+# %%
+prompt_format = [
+ "When John and Mary went to the shops,{} gave the bag to",
+ "When Tom and James went to the park,{} gave the ball to",
+ "When Dan and Sid went to the shops,{} gave an apple to",
+ "After Martin and Amy went to the park,{} gave a drink to",
+]
+names = [
+ (" Mary", " John"),
+ (" Tom", " James"),
+ (" Dan", " Sid"),
+ (" Martin", " Amy"),
+]
+# List of prompts
+prompts = []
+# List of answers, in the format (correct, incorrect)
+answers = []
+# List of the token (ie an integer) corresponding to each answer, in the format (correct_token, incorrect_token)
+answer_tokens = []
+for i in range(len(prompt_format)):
+ for j in range(2):
+ answers.append((names[i][j], names[i][1 - j]))
+ answer_tokens.append(
+ (
+ model.to_single_token(answers[-1][0]),
+ model.to_single_token(answers[-1][1]),
+ )
+ )
+ # Insert the *incorrect* answer to the prompt, making the correct answer the indirect object.
+ prompts.append(prompt_format[i].format(answers[-1][1]))
+answer_tokens = torch.tensor(answer_tokens).to(device)
+print(prompts)
+print(answers)
+
+# %% [markdown]
+# **Gotcha**: It's important that all of your prompts have the same number of tokens. If they're different lengths, then the position of the "final" logit where you can check logit difference will differ between prompts, and this will break the below code. The easiest solution is just to choose your prompts carefully to have the same number of tokens (you can eg add filler words like The, or newlines to start).
+#
+# There's a range of other ways of solving this, eg you can index more intelligently to get the final logit. A better way is to just use left padding by setting `model.tokenizer.padding_side = 'left'` before tokenizing the inputs and running the model; this way, you can use something like `logits[:, -1, :]` to easily access the final token outputs without complicated indexing. TransformerLens checks the value of `padding_side` of the tokenizer internally, and if the flag is set to be `'left'`, it adjusts the calculation of absolute position embedding and causal masking accordingly.
+#
+# In this demo, though, we stick to using the prompts of the same number of tokens because we want to show some visualisations aggregated along the batch dimension later in the demo.
+
+# %%
+for prompt in prompts:
+ str_tokens = model.to_str_tokens(prompt)
+ print("Prompt length:", len(str_tokens))
+ print("Prompt as tokens:", str_tokens)
+
+# %% [markdown]
+# We now run the model on these prompts and use `run_with_cache` to get both the logits and a cache of all internal activations for later analysis
+
+# %%
+tokens = model.to_tokens(prompts, prepend_bos=True)
+
+# Run the model and cache all activations
+original_logits, cache = model.run_with_cache(tokens)
+
+# %% [markdown]
+# We'll later be evaluating how model performance differs upon performing various interventions, so it's useful to have a metric to measure model performance. Our metric here will be the **logit difference**, the difference in logit between the indirect object's name and the subject's name (eg, `logit(Mary)-logit(John)`).
+
+# %%
+def logits_to_ave_logit_diff(logits, answer_tokens, per_prompt=False):
+ # Only the final logits are relevant for the answer
+ final_logits = logits[:, -1, :]
+ answer_logits = final_logits.gather(dim=-1, index=answer_tokens)
+ answer_logit_diff = answer_logits[:, 0] - answer_logits[:, 1]
+ if per_prompt:
+ return answer_logit_diff
+ else:
+ return answer_logit_diff.mean()
+
+
+print(
+ "Per prompt logit difference:",
+ logits_to_ave_logit_diff(original_logits, answer_tokens, per_prompt=True)
+ .detach()
+ .cpu()
+ .round(decimals=3),
+)
+original_average_logit_diff = logits_to_ave_logit_diff(original_logits, answer_tokens)
+print(
+ "Average logit difference:",
+ round(logits_to_ave_logit_diff(original_logits, answer_tokens).item(), 3),
+)
+
+# %% [markdown]
+# We see that the average logit difference is 3.5 - for context, this represents putting an $e^{3.5}\approx 33\times$ higher probability on the correct answer.
+
+# %% [markdown]
+# ## Brainstorm What's Actually Going On (Optional)
+#
+# Before diving into running experiments, it's often useful to spend some time actually reasoning about how the behaviour in question could be implemented in the transformer. **This is optional, and you'll likely get the most out of engaging with this section if you have a decent understanding already of what a transformer is and how it works!**
+#
+# You don't have to do this and forming hypotheses after exploration is also reasonable, but I think it's often easier to explore and interpret results with some grounding in what you might find. In this particular case, I'm cheating somewhat, since I know the answer, but I'm trying to simulate the process of reasoning about it!
+#
+# Note that often your hypothesis will be wrong in some ways and often be completely off. We're doing science here, and the goal is to understand how the model *actually* works, and to form true beliefs! There are two separate traps here at two extremes that it's worth tracking:
+# * Confusion: Having no hypotheses at all, getting a lot of data and not knowing what to do with it, and just floundering around
+# * Dogmatism: Being overconfident in an incorrect hypothesis and being unwilling to let go of it when reality contradicts you, or flinching away from running the experiments that might disconfirm it.
+#
+# **Exercise:** Spend some time thinking through how you might imagine this behaviour being implemented in a transformer. Try to think through this for yourself before reading through my thoughts!
+#
+# (*) My reasoning
+#
+# Brainstorming:
+#
+# So, what's hard about the task? Let's focus on the concrete example of the first prompt, "When John and Mary went to the shops, John gave the bag to" -> " Mary".
+#
+# A good starting point is thinking though whether a tiny model could do this, eg a 1L Attn-Only model. I'm pretty sure the answer is no! Attention is really good at the primitive operations of looking nearby, or copying information. I can believe a tiny model could figure out that at `to` it should look for names and predict that those names came next (eg the skip trigram " John...to -> John"). But it's much harder to tell how many of each previous name there are - attending 0.3 to each copy of John will look exactly the same as attending 0.6 to a single John token. So this will be pretty hard to figure out on the " to" token!
+#
+# The natural place to break this symmetry is on the second " John" token - telling whether there is an earlier copy of the current token should be a much easier task. So I might expect there to be a head which detects duplicate tokens on the second " John" token, and then another head which moves that information from the second " John" token to the " to" token.
+#
+# The model then needs to learn to predict " Mary" and not " John". I can see two natural ways to do this:
+# 1. Detect all preceding names and move this information to " to" and then delete the any name corresponding to the duplicate token feature. This feels easier done with a non-linearity, since precisely cancelling out vectors is hard, so I'd imagine an MLP layer deletes the " John" direction of the residual stream
+# 2. Have a head which attends to all previous names, but where the duplicate token features inhibit it from attending to specific names. So this only attends to Mary. And then the output of this head maps to the logits.
+#
+# (Spoiler: It's the second one).
+#
+# Experiment Ideas
+#
+# A test that could distinguish these two is to look at which components of the model add directly to the logits - if it's mostly attention heads which attend to " Mary" and to neither " John" it's probably hypothesis 2, if it's mostly MLPs it's probably hypothesis 1.
+#
+# And we should be able to identify duplicate token heads by finding ones which attend from " John" to " John", and whose outputs are then moved to the " to" token by V-Composition with another head (Spoiler: It's more complicated than that!)
+#
+# Note that all of the above reasoning is very simplistic and could easily break in a real model! There'll be significant parts of the model that figure out whether to use this circuit at all (we don't want to inhibit duplicated names when, eg, figuring out what goes at the start of the next sentence), and may be parts towards the end of the model that do "post-processing" just before the final output. But it's a good starting point for thinking about what's going on.
+
+# %% [markdown]
+# ## Direct Logit Attribution
+
+# %% [markdown]
+# *Look up unfamiliar terms in the [mech interp explainer](https://neelnanda.io/glossary)*
+#
+# Further, the easiest part of the model to understand is the output - this is what the model is trained to optimize, and so it can always be directly interpreted! Often the right approach to reverse engineering a circuit is to start at the end, understand how the model produces the right answer, and to then work backwards. The main technique used to do this is called **direct logit attribution**
+#
+# **Background:** The central object of a transformer is the **residual stream**. This is the sum of the outputs of each layer and of the original token and positional embedding. Importantly, this means that any linear function of the residual stream can be perfectly decomposed into the contribution of each layer of the transformer. Further, each attention layer's output can be broken down into the sum of the output of each head (See [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html) for details), and each MLP layer's output can be broken down into the sum of the output of each neuron (and a bias term for each layer).
+#
+# The logits of a model are `logits=Unembed(LayerNorm(final_residual_stream))`. The Unembed is a linear map, and LayerNorm is approximately a linear map, so we can decompose the logits into the sum of the contributions of each component, and look at which components contribute the most to the logit of the correct token! This is called **direct logit attribution**. Here we look at the direct attribution to the logit difference!
+#
+# (*) Background and motivation of the logit difference
+#
+# Logit difference is actually a *really* nice and elegant metric and is a particularly nice aspect of the setup of Indirect Object Identification. In general, there are two natural ways to interpret the model's outputs: the output logits, or the output log probabilities (or probabilities).
+#
+# The logits are much nicer and easier to understand, as noted above. However, the model is trained to optimize the cross-entropy loss (the average of log probability of the correct token). This means it does not directly optimize the logits, and indeed if the model adds an arbitrary constant to every logit, the log probabilities are unchanged.
+#
+# But `log_probs == logits.log_softmax(dim=-1) == logits - logsumexp(logits)`, and so `log_probs(" Mary") - log_probs(" John") = logits(" Mary") - logits(" John")` - the ability to add an arbitrary constant cancels out!
+#
+# Further, the metric helps us isolate the precise capability we care about - figuring out *which* name is the Indirect Object. There are many other components of the task - deciding whether to return an article (the) or pronoun (her) or name, realising that the sentence wants a person next at all, etc. By taking the logit difference we control for all of that.
+#
+# Our metric is further refined, because each prompt is repeated twice, for each possible indirect object. This controls for irrelevant behaviour such as the model learning that John is a more frequent token than Mary (this actually happens! The final layernorm bias increases the John logit by 1 relative to the Mary logit)
+#
+#
+#
+# Ignoring LayerNorm
+#
+# LayerNorm is an analogous normalization technique to BatchNorm (that's friendlier to massive parallelization) that transformers use. Every time a transformer layer reads information from the residual stream, it applies a LayerNorm to normalize the vector at each position (translating to set the mean to 0 and scaling to set the variance to 1) and then applying a learned vector of weights and biases to scale and translate the normalized vector. This is *almost* a linear map, apart from the scaling step, because that divides by the norm of the vector and the norm is not a linear function. (The `fold_ln` flag when loading a model factors out all the linear parts).
+#
+# But if we fixed the scale factor, the LayerNorm would be fully linear. And the scale of the residual stream is a global property that's a function of *all* components of the stream, while in practice there is normally just a few directions relevant to any particular component, so in practice this is an acceptable approximation. So when doing direct logit attribution we use the `apply_ln` flag on the `cache` to apply the global layernorm scaling factor to each constant. See [my clean GPT-2 implementation](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/clean-transformer-demo/Clean_Transformer_Demo.ipynb#scrollTo=Clean_Transformer_Implementation) for more on LayerNorm.
+#
+
+# %% [markdown]
+# Getting an output logit is equivalent to projecting onto a direction in the residual stream. We use `model.tokens_to_residual_directions` to map the answer tokens to that direction, and then convert this to a logit difference direction for each batch
+
+# %%
+answer_residual_directions = model.tokens_to_residual_directions(answer_tokens)
+print("Answer residual directions shape:", answer_residual_directions.shape)
+logit_diff_directions = (
+ answer_residual_directions[:, 0] - answer_residual_directions[:, 1]
+)
+print("Logit difference directions shape:", logit_diff_directions.shape)
+
+# %% [markdown]
+# To verify that this works, we can apply this to the final residual stream for our cached prompts (after applying LayerNorm scaling) and verify that we get the same answer.
+#
+# Technical details
+#
+# `logits = Unembed(LayerNorm(final_residual_stream))`, so we technically need to account for the centering, and then learned translation and scaling of the layernorm, not just the variance 1 scaling.
+#
+# The centering is accounted for with the preprocessing flag `center_writing_weights` which ensures that every weight matrix writing to the residual stream has mean zero.
+#
+# The learned scaling is folded into the unembedding weights `model.unembed.W_U` via `W_U_fold = layer_norm.weights[:, None] * unembed.W_U`
+#
+# The learned translation is folded to `model.unembed.b_U`, a bias added to the logits (note that GPT-2 is not trained with an existing `b_U`). This roughly represents unigram statistics. But we can ignore this because each prompt occurs twice with names in the opposite order, so this perfectly cancels out.
+#
+# Note that rather than using layernorm scaling we could just study cache["ln_final.hook_normalised"]
+#
+#
+
+# %%
+# cache syntax - resid_post is the residual stream at the end of the layer, -1 gets the final layer. The general syntax is [activation_name, layer_index, sub_layer_type].
+final_residual_stream = cache["resid_post", -1]
+print("Final residual stream shape:", final_residual_stream.shape)
+final_token_residual_stream = final_residual_stream[:, -1, :]
+# Apply LayerNorm scaling
+# pos_slice is the subset of the positions we take - here the final token of each prompt
+scaled_final_token_residual_stream = cache.apply_ln_to_stack(
+ final_token_residual_stream, layer=-1, pos_slice=-1
+)
+
+average_logit_diff = einsum(
+ "batch d_model, batch d_model -> ",
+ scaled_final_token_residual_stream,
+ logit_diff_directions,
+) / len(prompts)
+print("Calculated average logit diff:", round(average_logit_diff.item(), 3))
+print("Original logit difference:", round(original_average_logit_diff.item(), 3))
+
+# %% [markdown]
+# ### Logit Lens
+
+# %% [markdown]
+# We can now decompose the residual stream! First we apply a technique called the [**logit lens**](https://www.alignmentforum.org/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens) - this looks at the residual stream after each layer and calculates the logit difference from that. This simulates what happens if we delete all subsequence layers.
+
+# %%
+def residual_stack_to_logit_diff(
+ residual_stack: Float[torch.Tensor, "components batch d_model"],
+ cache: ActivationCache,
+) -> float:
+ scaled_residual_stack = cache.apply_ln_to_stack(
+ residual_stack, layer=-1, pos_slice=-1
+ )
+ return einsum(
+ "... batch d_model, batch d_model -> ...",
+ scaled_residual_stack,
+ logit_diff_directions,
+ ) / len(prompts)
+
+# %% [markdown]
+# Fascinatingly, we see that the model is utterly unable to do the task until layer 7, almost all performance comes from attention layer 9, and performance actually *decreases* from there.
+#
+# **Note:** Hover over each data point to see what residual stream position it's from!
+#
+# Details on `accumulated_resid`
+# **Key:** `n_pre` means the residual stream at the start of layer n, `n_mid` means the residual stream after the attention part of layer n (`n_post` is the same as `n+1_pre` so is not included)
+#
+# * `layer` is the layer for which we input the residual stream (this is used to identify *which* layer norm scaling factor we want)
+# * `incl_mid` is whether to include the residual stream in the middle of a layer, ie after attention & before MLP
+# * `pos_slice` is the subset of the positions used. See `utils.Slice` for details on the syntax.
+# * return_labels is whether to return the labels for each component returned (useful for plotting)
+#
+
+# %% [markdown]
+# ## Residual Stream
+
+# %% [markdown]
+# Lets begin by patching in the residual stream at the start of each layer and for each token position.
+
+# %% [markdown]
+# We first create a set of corrupted tokens - where we swap each pair of prompts to have the opposite answer.
+
+# %%
+corrupted_prompts = []
+for i in range(0, len(prompts), 2):
+ corrupted_prompts.append(prompts[i + 1])
+ corrupted_prompts.append(prompts[i])
+corrupted_tokens = model.to_tokens(corrupted_prompts, prepend_bos=True)
+corrupted_logits, corrupted_cache = model.run_with_cache(
+ corrupted_tokens, return_type="logits"
+)
+corrupted_average_logit_diff = logits_to_ave_logit_diff(corrupted_logits, answer_tokens)
+print("Corrupted Average Logit Diff", round(corrupted_average_logit_diff.item(), 2))
+print("Clean Average Logit Diff", round(original_average_logit_diff.item(), 2))
+
+
+
+# %% [markdown]
+# #### Implications
+#
+# One implication of this is that it's useful to categories heads according to whether they occur in
+# simpler circuits, so that as we look for more complex circuits we can easily look for them. This is
+# easy to do here! An interesting fact about induction heads is that they work on a sequence of
+# repeated random tokens - notable for being wildly off distribution from the natural language GPT-2
+# was trained on. Being able to predict a model's behaviour off distribution is a good mark of success
+# for mechanistic interpretability! This is a good sanity check for whether a head is an induction
+# head or not.
+#
+# We can characterise an induction head by just giving a sequence of random tokens repeated once, and
+# measuring the average attention paid from the second copy of a token to the token after the first
+# copy. At the same time, we can also measure the average attention paid from the second copy of a
+# token to the first copy of the token, which is the attention that the induction head would pay if it
+# were a duplicate token head, and the average attention paid to the previous token to find previous
+# token heads.
+#
+# Note that this is a superficial study of whether something is an induction head - we totally ignore
+# the question of whether it actually does boost the correct token or whether it composes with a
+# single previous head and how. In particular, we sometimes get anti-induction heads which suppress
+# the induction-y token (no clue why!), and this technique will find those too . But given the
+# previous rigorous analysis, we can be pretty confident that this picks up on some true signal about
+# induction heads.
+
+# %% [markdown]
+# Technical Implementation Details
+# We can do this again by using hooks, this time just to access the attention patterns rather than to intervene on them.
+#
+# Our hook function acts on the attention pattern activation. This has the name
+# "blocks.{layer}.{layer_type}.hook_{activation_name}" in general, here it's
+# "blocks.{layer}.attn.hook_attn". And it has shape [batch, head_index, query_pos, token_pos]. Our
+# hook function takes in the attention pattern activation, calculates the score for the relevant type
+# of head, and write it to an external cache.
+#
+# We add in hooks using `model.run_with_hooks(tokens, fwd_hooks=[(names_filter, hook_fn)])` to
+# temporarily add in the hooks and run the model, getting the resulting output. Previously
+# names_filter was the name of the activation, but here it's a boolean function mapping activation
+# names to whether we want to hook them or not. Here it's just whether the name ends with hook_attn.
+# hook_fn must take in the two inputs activation (the activation tensor) and hook (the HookPoint
+# object, which contains the name of the activation and some metadata such as the current layer).
+#
+# Internally our hooks use the function `tensor.diagonal`, this takes the diagonal between two
+# dimensions, and allows an arbitrary offset - offset by 1 to get previous tokens, seq_len to get
+# duplicate tokens (the distance to earlier copies) and seq_len-1 to get induction heads (the distance
+# to the token *after* earlier copies). Different offsets give a different length of output tensor,
+# and we can now just average to get a score in [0, 1] for each head
+#
+
+# %%
+seq_len = 100
+batch_size = 2
+
+prev_token_scores = torch.zeros((model.cfg.n_layers, model.cfg.n_heads), device=device)
+
+
+def prev_token_hook(pattern, hook):
+ layer = hook.layer()
+ diagonal = pattern.diagonal(offset=1, dim1=-1, dim2=-2)
+ # print(diagonal)
+ # print(pattern)
+ prev_token_scores[layer] = einops.reduce(
+ diagonal, "batch head_index diagonal -> head_index", "mean"
+ )
+
+
+duplicate_token_scores = torch.zeros(
+ (model.cfg.n_layers, model.cfg.n_heads), device=device
+)
+
+
+def duplicate_token_hook(pattern, hook):
+ layer = hook.layer()
+ diagonal = pattern.diagonal(offset=seq_len, dim1=-1, dim2=-2)
+ duplicate_token_scores[layer] = einops.reduce(
+ diagonal, "batch head_index diagonal -> head_index", "mean"
+ )
+
+
+induction_scores = torch.zeros((model.cfg.n_layers, model.cfg.n_heads), device=device)
+
+
+def induction_hook(pattern, hook):
+ layer = hook.layer()
+ diagonal = pattern.diagonal(offset=seq_len - 1, dim1=-1, dim2=-2)
+ induction_scores[layer] = einops.reduce(
+ diagonal, "batch head_index diagonal -> head_index", "mean"
+ )
+
+
+torch.manual_seed(0)
+original_tokens = torch.randint(
+ 100, 20000, size=(batch_size, seq_len), device="cpu"
+).to(device)
+repeated_tokens = einops.repeat(
+ original_tokens, "batch seq_len -> batch (2 seq_len)"
+).to(device)
+
+pattern_filter = lambda act_name: act_name.endswith("hook_pattern")
+
+loss = model.run_with_hooks(
+ repeated_tokens,
+ return_type="loss",
+ fwd_hooks=[
+ (pattern_filter, prev_token_hook),
+ (pattern_filter, duplicate_token_hook),
+ (pattern_filter, induction_hook),
+ ],
+)
+print(torch.round(utils.get_corner(prev_token_scores).detach().cpu(), decimals=3))
+print(torch.round(utils.get_corner(duplicate_token_scores).detach().cpu(), decimals=3))
+print(torch.round(utils.get_corner(induction_scores).detach().cpu(), decimals=3))
+
+# %% [markdown]
+# We can now plot the head scores, and instantly see that the relevant early heads are induction heads or duplicate token heads (though also that there's a lot of induction heads that are *not* use - I have no idea why!).
+
+
+
+# %% [markdown]
+# The above suggests that it would be a useful bit of infrastructure to have a "wiki" for the heads of a model, giving their scores according to some metrics re head functions, like the ones we've seen here. TransformerLens makes this easy to make, as just changing the name input to `HookedTransformer.from_pretrained` gives a different model but in the same architecture, so the same code should work. If you want to make this, I'd love to see it!
+#
+# As a proof of concept, [I made a mosaic of all induction heads across the 40 models then in TransformerLens](https://www.neelnanda.io/mosaic).
+#
+# 
+
+# %% [markdown]
+# ### Backup Name Mover Heads
+
+# %% [markdown]
+# Another fascinating anomaly is that of the **backup name mover heads**. A standard technique to apply when interpreting model internals is ablations, or knock-out. If we run the model but intervene to set a specific head to zero, what happens? If the model is robust to this intervention, then naively we can be confident that the head is not doing anything important, and conversely if the model is much worse at the task this suggests that head was important. There are several conceptual flaws with this approach, making the evidence only suggestive, eg that the average output of the head may be far from zero and so the knockout may send it far from expected activations, breaking internals on *any* task. But it's still an easy technique to apply to give some data.
+#
+# But a wild finding in the paper is that models have **built in redundancy**. If we knock out one of the name movers, then there are some backup name movers in later layers that *change their behaviour* and do (some of) the job of the original name mover head. This means that naive knock-out will significantly underestimate the importance of the name movers.
+#
+
+# %% [markdown]
+# Let's test this! Let's ablate the most important name mover (head L9H9) on just the final token using a custom ablation hook and then cache all new activations and compared performance. We focus on the final position because we want to specifically ablate the direct logit effect. When we do this, we see that naively, removing the top name mover should reduce the logit diff massively, from 3.55 to 0.57. **But actually, it only goes down to 2.99!**
+#
+# Implementation Details
+# Ablating heads is really easy in TransformerLens! We can just define a hook on the z activation in the relevant attention layer (recall, z is the mixed values, and comes immediately before multiplying by the output weights $W_O$). z has a head_index axis, so we can set the component for the relevant head and for position -1 to zero, and return it. (Technically we could just edit in place without returning it, but by convention we always return an edited activation).
+#
+# We now want to compare all internal activations with a hook, which is hard to do with the nice `run_with_hooks` API. So we can directly access the hook on the z activation with `model.blocks[layer].attn.hook_z` and call its `add_hook` method. This adds in the hook to the *global state* of the model. We can now use run_with_cache, and don't need to care about the global state, because run_with_cache internally adds a bunch of caching hooks, and then removes all hooks after the run, *including* the previously added ablation hook. This can be disabled with the reset_hooks_end flag, but here it's useful!
+#
+
+# %%
+top_name_mover = per_head_logit_diffs.flatten().argmax().item()
+top_name_mover_layer = top_name_mover // model.cfg.n_heads
+top_name_mover_head = top_name_mover % model.cfg.n_heads
+print(f"Top Name Mover to ablate: L{top_name_mover_layer}H{top_name_mover_head}")
+
+
+def ablate_top_head_hook(z: Float[torch.Tensor, "batch pos head_index d_head"], hook):
+ z[:, -1, top_name_mover_head, :] = 0
+ return z
+
+
+# Adds a hook into global model state
+model.blocks[top_name_mover_layer].attn.hook_z.add_hook(ablate_top_head_hook)
+# Runs the model, temporarily adds caching hooks and then removes *all* hooks after running, including the ablation hook.
+ablated_logits, ablated_cache = model.run_with_cache(tokens)
+print(f"Original logit diff: {original_average_logit_diff:.2f}")
+print(
+ f"Post ablation logit diff: {logits_to_ave_logit_diff(ablated_logits, answer_tokens).item():.2f}"
+)
+print(
+ f"Direct Logit Attribution of top name mover head: {per_head_logit_diffs.flatten()[top_name_mover].item():.2f}"
+)
+print(
+ f"Naive prediction of post ablation logit diff: {original_average_logit_diff - per_head_logit_diffs.flatten()[top_name_mover].item():.2f}"
+)
+
+# %% [markdown]
+# So what's up with this? As before, we can look at the direct logit attribution of each head to see what's going on. It's easiest to interpret if plotted as a scatter plot against the initial per head logit difference.
+#
+# And we can see a *really* big difference in a few heads! (Hover to see labels) In particular the negative name mover L10H7 decreases its negative effect a lot, adding +1 to the logit diff, and the backup name mover L10H10 adjusts its effect to be more positive, adding +0.8 to the logit diff (with several other marginal changes). (And obviously the ablated head has gone down to zero!)
+
+# %%
+per_head_ablated_residual, labels = ablated_cache.stack_head_results(
+ layer=-1, pos_slice=-1, return_labels=True
+)
+per_head_ablated_logit_diffs = residual_stack_to_logit_diff(
+ per_head_ablated_residual, ablated_cache
+)
+per_head_ablated_logit_diffs = per_head_ablated_logit_diffs.reshape(
+ model.cfg.n_layers, model.cfg.n_heads
+)
+
+# %% [markdown]
+# One natural hypothesis is that this is because the final LayerNorm scaling has changed, which can scale up or down the final residual stream. This is slightly true, and we can see that the typical head is a bit off from the x=y line. But the average LN scaling ratio is 1.04, and this should uniformly change *all* heads by the same factor, so this can't be sufficient
+
+# %%
+print(
+ "Average LN scaling ratio:",
+ round(
+ (
+ cache["ln_final.hook_scale"][:, -1]
+ / ablated_cache["ln_final.hook_scale"][:, -1]
+ )
+ .mean()
+ .item(),
+ 3,
+ ),
+)
+print(
+ "Ablation LN scale",
+ ablated_cache["ln_final.hook_scale"][:, -1].detach().cpu().round(decimals=2),
+)
+print(
+ "Original LN scale",
+ cache["ln_final.hook_scale"][:, -1].detach().cpu().round(decimals=2),
+)
+
+# %% [markdown]
+# **Exercise to the reader:** Can you finish off this analysis? What's going on here? Why are the backup name movers changing their behaviour? Why is one negative name mover becoming significantly less important?
+
+
diff --git a/LLaMA.py b/LLaMA.py
new file mode 100644
index 0000000..585f14c
--- /dev/null
+++ b/LLaMA.py
@@ -0,0 +1,192 @@
+# %% [markdown]
+#
+#
+#
+
+# %% [markdown]
+# # LLaMA and Llama-2 in TransformerLens
+
+# %% [markdown]
+# ## Setup (skip)
+
+# %%
+# NBVAL_IGNORE_OUTPUT
+# Janky code to do different setup when run in a Colab notebook vs VSCode
+import os
+
+# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh
+import plotly.io as pio
+pio.renderers.default = "notebook_connected"
+print(f"Using renderer: {pio.renderers.default}")
+
+import circuitsvis as cv
+
+# %%
+# Import stuff
+import torch
+import tqdm.auto as tqdm
+import plotly.express as px
+
+from transformers import LlamaForCausalLM, LlamaTokenizer
+from tqdm import tqdm
+from jaxtyping import Float
+
+import transformer_lens
+import transformer_lens.utils as utils
+from transformer_lens.hook_points import (
+ HookPoint,
+) # Hooking utilities
+from transformer_lens import HookedTransformer
+from transformer_lens.boot import boot
+
+torch.set_grad_enabled(False)
+
+def imshow(tensor, renderer=None, xaxis="", yaxis="", **kwargs):
+ px.imshow(utils.to_numpy(tensor), color_continuous_midpoint=0.0, color_continuous_scale="RdBu", labels={"x":xaxis, "y":yaxis}, **kwargs).show(renderer)
+
+def line(tensor, renderer=None, xaxis="", yaxis="", **kwargs):
+ px.line(utils.to_numpy(tensor), labels={"x":xaxis, "y":yaxis}, **kwargs).show(renderer)
+
+def scatter(x, y, xaxis="", yaxis="", caxis="", renderer=None, **kwargs):
+ x = utils.to_numpy(x)
+ y = utils.to_numpy(y)
+ px.scatter(y=y, x=x, labels={"x":xaxis, "y":yaxis, "color":caxis}, **kwargs).show(renderer)
+
+# %% [markdown]
+# ## Loading LLaMA
+
+# %% [markdown]
+# LLaMA weights are not available on HuggingFace, so you'll need to download and convert them
+# manually:
+#
+# 1. Get LLaMA weights here: https://docs.google.com/forms/d/e/1FAIpQLSfqNECQnMkycAp2jP4Z9TFX0cGR4uf7b_fBxjY_OjhJILlKGA/viewform
+#
+# 2. Convert the official weights to huggingface:
+#
+# ```bash
+# python src/transformers/models/llama/convert_llama_weights_to_hf.py \
+# --input_dir /path/to/downloaded/llama/weights \
+# --model_size 7B \
+# --output_dir /llama/weights/directory/
+# ```
+#
+# Note: this didn't work for Arthur by default (even though HF doesn't seem to show this anywhere). I
+# had to change this
+# line of my pip installed `src/transformers/models/llama/convert_llama_weights_to_hf.py` file (which
+# was found at
+# `/opt/conda/envs/arthurenv/lib/python3.10/site-packages/transformers/models/llama/convert_llama_weights_to_hf.py`)
+# from `input_base_path=os.path.join(args.input_dir, args.model_size),` to `input_base_path=os.path.join(args.input_dir),`
+#
+# 3. Change the ```MODEL_PATH``` variable in the cell below to where the converted weights are stored.
+
+# %%
+MODEL_PATH = ""
+
+if MODEL_PATH:
+ tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH)
+ hf_model = LlamaForCausalLM.from_pretrained(MODEL_PATH, low_cpu_mem_usage=True)
+
+ model = boot(
+ "llama-7b",
+ hf_model=hf_model,
+ device="cpu",
+ fold_ln=False,
+ center_writing_weights=False,
+ center_unembed=False,
+ tokenizer=tokenizer,
+ )
+
+ model = model.to("cuda" if torch.cuda.is_available() else "cpu")
+ model.generate("The capital of Germany is", max_new_tokens=20, temperature=0)
+
+# %% [markdown]
+# ## Loading LLaMA-2
+# LLaMA-2 is hosted on HuggingFace, but gated by login.
+#
+# Before running the notebook, log in to HuggingFace via the cli on your machine:
+# ```bash
+# transformers-cli login
+# ```
+# This will cache your HuggingFace credentials, and enable you to download LLaMA-2.
+
+# %%
+LLAMA_2_7B_CHAT_PATH = "meta-llama/Llama-2-7b-chat-hf"
+
+tokenizer = LlamaTokenizer.from_pretrained(LLAMA_2_7B_CHAT_PATH)
+hf_model = LlamaForCausalLM.from_pretrained(LLAMA_2_7B_CHAT_PATH, low_cpu_mem_usage=True)
+
+model = boot(LLAMA_2_7B_CHAT_PATH, device="cpu", fold_ln=False, center_writing_weights=False, center_unembed=False)
+
+model = model.to("cuda" if torch.cuda.is_available() else "cpu")
+model.generate("The capital of Germany is", max_new_tokens=20, temperature=0)
+
+# %% [markdown]
+# ### Compare logits with HuggingFace model
+
+# %%
+prompts = [
+ "The capital of Germany is",
+ "2 * 42 = ",
+ "My favorite",
+ "aosetuhaosuh aostud aoestuaoentsudhasuh aos tasat naostutshaosuhtnaoe usaho uaotsnhuaosntuhaosntu haouaoshat u saotheu saonuh aoesntuhaosut aosu thaosu thaoustaho usaothusaothuao sutao sutaotduaoetudet uaosthuao uaostuaoeu aostouhsaonh aosnthuaoscnuhaoshkbaoesnit haosuhaoe uasotehusntaosn.p.uo ksoentudhao ustahoeuaso usant.hsa otuhaotsi aostuhs",
+]
+
+model.eval()
+hf_model.eval()
+prompt_ids = [tokenizer.encode(prompt, return_tensors="pt") for prompt in prompts]
+tl_logits = [model(prompt_ids).detach().cpu() for prompt_ids in tqdm(prompt_ids)]
+
+# hf logits are really slow as it's on CPU. If you have a big/multi-GPU machine, run `hf_model = hf_model.to("cuda")` to speed this up
+logits = [hf_model(prompt_ids).logits.detach().cpu() for prompt_ids in tqdm(prompt_ids)]
+
+for i in range(len(prompts)):
+ assert torch.allclose(logits[i], tl_logits[i], atol=1e-4, rtol=1e-2)
+
+# %% [markdown]
+# ## TransformerLens Demo
+
+# %% [markdown]
+# ### Reading from hooks
+
+# %%
+llama_text = "Natural language processing tasks, such as question answering, machine translation, reading comprehension, and summarization, are typically approached with supervised learning on taskspecific datasets."
+llama_tokens = model.to_tokens(llama_text)
+llama_logits, llama_cache = model.run_with_cache(llama_tokens, remove_batch_dim=True)
+
+attention_pattern = llama_cache["pattern", 0, "attn"]
+llama_str_tokens = model.to_str_tokens(llama_text)
+
+print("Layer 0 Head Attention Patterns:")
+
+# %% [markdown]
+# ### Writing to hooks
+
+# %%
+layer_to_ablate = 0
+head_index_to_ablate = 31
+
+# We define a head ablation hook
+# The type annotations are NOT necessary, they're just a useful guide to the reader
+#
+def head_ablation_hook(
+ value: Float[torch.Tensor, "batch pos head_index d_head"],
+ hook: HookPoint
+) -> Float[torch.Tensor, "batch pos head_index d_head"]:
+ print(f"Shape of the value tensor: {value.shape}")
+ value[:, :, head_index_to_ablate, :] = 0.
+ return value
+
+original_loss = model(llama_tokens, return_type="loss")
+ablated_loss = model.run_with_hooks(
+ llama_tokens,
+ return_type="loss",
+ fwd_hooks=[(
+ utils.get_act_name("v", layer_to_ablate),
+ head_ablation_hook
+ )]
+ )
+print(f"Original Loss: {original_loss.item():.3f}")
+print(f"Ablated Loss: {ablated_loss.item():.3f}")
+
+
diff --git a/LLaMA2_GPU_Quantized.py b/LLaMA2_GPU_Quantized.py
new file mode 100644
index 0000000..6ea08a3
--- /dev/null
+++ b/LLaMA2_GPU_Quantized.py
@@ -0,0 +1,214 @@
+# %% [markdown]
+# # LLaMA and Llama-2 in TransformerLens
+
+# %% [markdown]
+# ## Setup (skip)
+
+# %%
+# NBVAL_IGNORE_OUTPUT
+# Janky code to do different setup when run in a Colab notebook vs VSCode
+import os
+
+# %%
+# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh
+import plotly.io as pio
+pio.renderers.default = "notebook_connected"
+print(f"Using renderer: {pio.renderers.default}")
+
+import circuitsvis as cv
+
+# %%
+# Import stuff
+import torch
+import tqdm.auto as tqdm
+import plotly.express as px
+
+from transformers import LlamaForCausalLM, LlamaTokenizer
+from tqdm import tqdm
+from jaxtyping import Float
+
+import transformer_lens
+import transformer_lens.utils as utils
+from transformer_lens.hook_points import (
+ HookPoint,
+) # Hooking utilities
+from transformer_lens import HookedTransformer
+
+torch.set_grad_enabled(False)
+
+def imshow(tensor, renderer=None, xaxis="", yaxis="", **kwargs):
+ px.imshow(utils.to_numpy(tensor), color_continuous_midpoint=0.0, color_continuous_scale="RdBu", labels={"x":xaxis, "y":yaxis}, **kwargs).show(renderer)
+
+def line(tensor, renderer=None, xaxis="", yaxis="", **kwargs):
+ px.line(utils.to_numpy(tensor), labels={"x":xaxis, "y":yaxis}, **kwargs).show(renderer)
+
+def scatter(x, y, xaxis="", yaxis="", caxis="", renderer=None, **kwargs):
+ x = utils.to_numpy(x)
+ y = utils.to_numpy(y)
+ px.scatter(y=y, x=x, labels={"x":xaxis, "y":yaxis, "color":caxis}, **kwargs).show(renderer)
+
+# %% [markdown]
+# ## Loading LLaMA
+
+# %% [markdown]
+# LLaMA weights are not available on HuggingFace, so you'll need to download and convert them
+# manually:
+#
+# 1. Get LLaMA weights here: https://docs.google.com/forms/d/e/1FAIpQLSfqNECQnMkycAp2jP4Z9TFX0cGR4uf7b_fBxjY_OjhJILlKGA/viewform
+#
+# 2. Convert the official weights to huggingface:
+#
+# ```bash
+# python src/transformers/models/llama/convert_llama_weights_to_hf.py \
+# --input_dir /path/to/downloaded/llama/weights \
+# --model_size 7B \
+# --output_dir /llama/weights/directory/
+# ```
+#
+# Note: this didn't work for Arthur by default (even though HF doesn't seem to show this anywhere). I
+# had to change this
+# line of my pip installed `src/transformers/models/llama/convert_llama_weights_to_hf.py` file (which
+# was found at
+# `/opt/conda/envs/arthurenv/lib/python3.10/site-packages/transformers/models/llama/convert_llama_weights_to_hf.py`)
+# from `input_base_path=os.path.join(args.input_dir, args.model_size),` to `input_base_path=os.path.join(args.input_dir),`
+#
+# 3. Change the ```MODEL_PATH``` variable in the cell below to where the converted weights are stored.
+
+# %%
+MODEL_PATH=''
+
+if MODEL_PATH:
+ tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH)
+ hf_model = LlamaForCausalLM.from_pretrained(MODEL_PATH, low_cpu_mem_usage=True)
+
+ model = boot("llama-7b", hf_model=hf_model, device="cpu", fold_ln=False, center_writing_weights=False, center_unembed=False, tokenizer=tokenizer)
+
+ model = model.to("cuda" if torch.cuda.is_available() else "cpu")
+ model.generate("The capital of Germany is", max_new_tokens=20, temperature=0)
+
+# %% [markdown]
+# ## Loading LLaMA-2
+# LLaMA-2 is hosted on HuggingFace, but gated by login.
+#
+# Before running the notebook, log in to HuggingFace via the cli on your machine:
+# ```bash
+# transformers-cli login
+# ```
+# This will cache your HuggingFace credentials, and enable you to download LLaMA-2.
+
+# %% [markdown]
+# ## Install additional dependenceis requred for quantization
+
+# %%
+
+# %% [markdown]
+# ## Load quantized model
+
+# %%
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+LLAMA_2_7B_CHAT_PATH = "meta-llama/Llama-2-7b-chat-hf"
+inference_dtype = torch.float32
+# inference_dtype = torch.float32
+# inference_dtype = torch.float16
+
+hf_model = AutoModelForCausalLM.from_pretrained(LLAMA_2_7B_CHAT_PATH,
+ torch_dtype=inference_dtype,
+ device_map = "cuda:0",
+ quantization_config=BitsAndBytesConfig(load_in_4bit=True))
+
+tokenizer = AutoTokenizer.from_pretrained(LLAMA_2_7B_CHAT_PATH)
+
+model = boot(LLAMA_2_7B_CHAT_PATH,
+ hf_model=hf_model,
+ dtype=inference_dtype,
+ fold_ln=False,
+ fold_value_biases=False,
+ center_writing_weights=False,
+ center_unembed=False,
+ tokenizer=tokenizer)
+
+model.generate("The capital of Germany is", max_new_tokens=2, temperature=0)
+
+
+
+# %% [markdown]
+# ### Verify GPU memory use
+
+# %%
+print("free(Gb):", torch.cuda.mem_get_info()[0]/1000000000, "total(Gb):", torch.cuda.mem_get_info()[1]/1000000000)
+
+# %% [markdown]
+# ### Compare logits with HuggingFace model
+
+# %%
+prompts = [
+ "The capital of Germany is",
+ "2 * 42 = ",
+ "My favorite",
+ "aosetuhaosuh aostud aoestuaoentsudhasuh aos tasat naostutshaosuhtnaoe usaho uaotsnhuaosntuhaosntu haouaoshat u saotheu saonuh aoesntuhaosut aosu thaosu thaoustaho usaothusaothuao sutao sutaotduaoetudet uaosthuao uaostuaoeu aostouhsaonh aosnthuaoscnuhaoshkbaoesnit haosuhaoe uasotehusntaosn.p.uo ksoentudhao ustahoeuaso usant.hsa otuhaotsi aostuhs",
+]
+
+model.eval()
+hf_model.eval()
+prompt_ids = [tokenizer.encode(prompt, return_tensors="pt") for prompt in prompts]
+tl_logits = [model(prompt_ids).detach().cpu() for prompt_ids in tqdm(prompt_ids)]
+
+# hf logits are really slow as it's on CPU. If you have a big/multi-GPU machine, run `hf_model = hf_model.to("cuda")` to speed this up
+logits = [hf_model(prompt_ids).logits.detach().cpu() for prompt_ids in tqdm(prompt_ids)]
+
+for i in range(len(prompts)):
+ if i == 0:
+ print("logits[i]", i, logits[i].dtype, logits[i])
+ print("tl_logits[i]", i, tl_logits[i].dtype, tl_logits[i])
+ assert torch.allclose(logits[i], tl_logits[i], atol=1e-4, rtol=1e-2)
+
+# %% [markdown]
+# ## TransformerLens Demo
+
+# %% [markdown]
+# ### Reading from hooks
+
+# %%
+llama_text = "Natural language processing tasks, such as question answering, machine translation, reading comprehension, and summarization, are typically approached with supervised learning on taskspecific datasets."
+llama_tokens = model.to_tokens(llama_text)
+llama_logits, llama_cache = model.run_with_cache(llama_tokens, remove_batch_dim=True)
+
+attention_pattern = llama_cache["pattern", 0, "attn"]
+llama_str_tokens = model.to_str_tokens(llama_text)
+
+print("Layer 0 Head Attention Patterns:")
+
+# %% [markdown]
+# ### Writing to hooks
+
+# %%
+layer_to_ablate = 0
+head_index_to_ablate = 31
+
+# We define a head ablation hook
+# The type annotations are NOT necessary, they're just a useful guide to the reader
+#
+def head_ablation_hook(
+ value: Float[torch.Tensor, "batch pos head_index d_head"],
+ hook: HookPoint
+) -> Float[torch.Tensor, "batch pos head_index d_head"]:
+ print(f"Shape of the value tensor: {value.shape}")
+ value[:, :, head_index_to_ablate, :] = 0.
+ return value
+
+original_loss = model(llama_tokens, return_type="loss")
+ablated_loss = model.run_with_hooks(
+ llama_tokens,
+ return_type="loss",
+ fwd_hooks=[(
+ utils.get_act_name("v", layer_to_ablate),
+ head_ablation_hook
+ )]
+ )
+print(f"Original Loss: {original_loss.item():.3f}")
+print(f"Ablated Loss: {ablated_loss.item():.3f}")
+
+
diff --git a/LLaVA.py b/LLaVA.py
new file mode 100644
index 0000000..5a77467
--- /dev/null
+++ b/LLaVA.py
@@ -0,0 +1,180 @@
+# %% [markdown]
+# ### LLaVA use case demonstration
+#
+# At that notebook you can see simple example of how to use TransformerLens for LLaVA interpretability. More specifically you can pass united image patch embeddings and textual embedding to LLaVA language model (Vicuna) with TransformerLens and get logits and cache that contains activations for next analysis. Here we consider the simplest example of LLaVA and TransformerLens sharing.
+
+# %%
+# import staff
+import sys
+
+# Uncomment if use clonned version of TransformerLens
+# currently forked version https://github.com/zazamrykh/TransformerLens supports
+TL_path = r"../"
+if TL_path not in sys.path:
+ sys.path.insert(0, TL_path)
+ sys.path.insert(0, TL_path + r"/transformer_lens")
+
+import torch
+from transformers import AutoProcessor, LlavaForConditionalGeneration # Should update transformer to latest version
+
+# For image loading
+from PIL import Image
+import requests
+from io import BytesIO
+
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+import matplotlib.pyplot as plt
+%matplotlib inline
+
+from transformer_lens import HookedTransformer
+from transformer_lens.boot import boot
+import circuitsvis as cv
+
+_ = torch.set_grad_enabled(False)
+
+# %% [markdown]
+# Load llava model from hugging face. Load some revision because at this moment newest one is not working.
+
+# %%
+model_id = "llava-hf/llava-1.5-7b-hf"
+
+llava = LlavaForConditionalGeneration.from_pretrained(
+ model_id,
+ torch_dtype=torch.float16,
+ load_in_4bit=False,
+ low_cpu_mem_usage=True,
+ revision="a272c74",
+ device_map="cpu"
+)
+
+for param in llava.parameters(): # At this demo we don't need grads
+ param.requires_grad = False
+
+processor = AutoProcessor.from_pretrained(model_id, revision="a272c74")
+tokenizer = processor.tokenizer
+
+# Taking model apart
+language_model = llava.language_model.eval()
+config = language_model.config
+print("Base language model:", config._name_or_path)
+
+vision_tower = llava.vision_tower.to(device).eval()
+projector = llava.multi_modal_projector.to(device).eval()
+
+# %%
+# You can write your own version of getting language model's input embeddings similar way
+# This function will not be working with old transformers library version. Should update transformers library.
+def get_llm_input_embeddings(llava, processor, image: Image, text: str, device='cuda'):
+ """ Extract features from image, project them to LLM's space and insert them to text embedding sequence.
+ Returns:
+ inputs_embeds, attention_mask, labels, position_ids - input for language model of LLaVA
+ """
+ conversation = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": text},
+ {"type": "image"},
+ ],
+ },
+ ]
+ prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+ inputs = processor(images=image, text=prompt, return_tensors='pt').to(device, torch.float16)
+ llava.vision_tower.to(device)
+ llava.multi_modal_projector.to(device)
+
+ clip_output = llava.vision_tower(inputs['pixel_values'])
+ projector_output = llava.multi_modal_projector(clip_output.last_hidden_state)
+
+ before_device = llava.language_model.model.embed_tokens.weight.device
+ llava.language_model.model.embed_tokens.to(device)
+ text_embeddings = llava.language_model.model.embed_tokens(inputs['input_ids'])
+ llava.language_model.model.embed_tokens.to(before_device)
+
+ full_sequence = torch.hstack([projector_output, text_embeddings])
+
+ attention_mask = torch.ones(full_sequence.shape[:-1], device=full_sequence.device, dtype=int)
+ inputs_embeds, attention_mask, labels, position_ids = llava._merge_input_ids_with_image_features(
+ projector_output, text_embeddings, inputs['input_ids'], attention_mask, labels=None
+ ) # Access to private member... Well, but what can i do :-)
+
+ return inputs_embeds, attention_mask, labels, position_ids
+
+# %% [markdown]
+# Okay, now create HookedTransformer model
+
+# %%
+hooked_llm = boot(
+ "llama-7b-hf", # Use config of llama
+ center_unembed=False,
+ fold_ln=False,
+ fold_value_biases=False,
+ device='cuda',
+ hf_model=language_model, # Use Vicuna's weights
+ tokenizer=tokenizer,
+ center_writing_weights=False,
+ dtype=torch.float16,
+ vocab_size=language_model.config.vocab_size # New argument. llama and vicuna have different vocab size, so we pass it here
+)
+
+for param in hooked_llm.parameters():
+ param.requires_grad = False
+
+# %% [markdown]
+# Now try if hooked model is working
+
+# %%
+image_url = "https://github.com/zazamrykh/PicFinder/blob/main/images/doge.jpg?raw=true"
+response = requests.get(image_url)
+image = Image.open(BytesIO(response.content))
+plt.axis('off')
+_ = plt.imshow(image)
+
+# %%
+question = "What do you see on photo?"
+inputs_embeds, attention_mask, labels, position_ids = get_llm_input_embeddings(llava, processor, image, question, device=device)
+
+# Return tokens
+outputs = hooked_llm.generate(
+ inputs_embeds,
+ max_new_tokens=30,
+ do_sample=True,
+ return_type='tokens'
+)
+generated_text = processor.decode(outputs[0], skip_special_tokens=True)
+print('Generated text:', generated_text)
+
+# %%
+# Now return embeddings and then project them on vocab space
+outputs = hooked_llm.generate(
+ inputs_embeds,
+ max_new_tokens=30,
+ do_sample=True,
+)
+
+logits = outputs[:,-30:,:].to(device) @ language_model.model.embed_tokens.weight.T.to(device)
+generated_text = processor.decode(logits.argmax(-1)[0], skip_special_tokens=True)
+print('Generated text:', generated_text)
+
+# %% [markdown]
+# As we can see everything is working. Now try visualize attention patterns in generated output.
+
+# %%
+# Here we visualize attention for the last 30 tokens.
+logits, cache = hooked_llm.run_with_cache(inputs_embeds, start_at_layer=0, remove_batch_dim=True)
+
+layer_to_visualize = 16
+tokens_to_show = 30
+attention_pattern = cache["pattern", layer_to_visualize, "attn"]
+
+product = inputs_embeds @ language_model.model.embed_tokens.weight.T.to(device) # Project embeddings to vocab
+llama_str_tokens = hooked_llm.to_str_tokens(product.argmax(dim=-1)[0])
+
+print(f"Layer {layer_to_visualize} Head Attention Patterns:")
+
+# %% [markdown]
+# As we can see image tokens also appears and can be used for multimodal attention exploration.
+
+
diff --git a/Main_Demo.py b/Main_Demo.py
new file mode 100644
index 0000000..cb73866
--- /dev/null
+++ b/Main_Demo.py
@@ -0,0 +1,1063 @@
+# %%
+# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh
+import plotly.io as pio
+pio.renderers.default = "notebook_connected"
+print(f"Using renderer: {pio.renderers.default}")
+
+# %%
+import circuitsvis as cv
+# Testing that the library works
+cv.examples.hello("Neel")
+
+# %%
+# Import stuff
+import torch
+import torch.nn as nn
+import einops
+import tqdm.auto as tqdm
+import plotly.express as px
+
+from jaxtyping import Float
+from functools import partial
+
+# %%
+# import transformer_lens
+import transformer_lens.utils as utils
+from transformer_lens.hook_points import (
+ HookPoint,
+) # Hooking utilities
+from transformer_lens import HookedTransformer, FactoredMatrix
+from transformer_lens.boot import boot
+
+# %% [markdown]
+# We turn automatic differentiation off, to save GPU memory, as this notebook focuses on model inference not model training.
+
+# %%
+torch.set_grad_enabled(False)
+
+# %% [markdown]
+# Plotting helper functions:
+
+# %%
+def imshow(tensor, renderer=None, xaxis="", yaxis="", **kwargs):
+ px.imshow(utils.to_numpy(tensor), color_continuous_midpoint=0.0, color_continuous_scale="RdBu", labels={"x":xaxis, "y":yaxis}, **kwargs).show(renderer)
+
+def line(tensor, renderer=None, xaxis="", yaxis="", **kwargs):
+ px.line(utils.to_numpy(tensor), labels={"x":xaxis, "y":yaxis}, **kwargs).show(renderer)
+
+def scatter(x, y, xaxis="", yaxis="", caxis="", renderer=None, **kwargs):
+ x = utils.to_numpy(x)
+ y = utils.to_numpy(y)
+ px.scatter(y=y, x=x, labels={"x":xaxis, "y":yaxis, "color":caxis}, **kwargs).show(renderer)
+
+# %% [markdown]
+# # Introduction
+
+# %% [markdown]
+# This is a demo notebook for [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens), **a library I ([Neel Nanda](https://neelnanda.io)) wrote for doing [mechanistic interpretability](https://distill.pub/2020/circuits/zoom-in/) of GPT-2 Style language models.** The goal of mechanistic interpretability is to take a trained model and reverse engineer the algorithms the model learned during training from its weights. It is a fact about the world today that we have computer programs that can essentially speak English at a human level (GPT-3, PaLM, etc), yet we have no idea how they work nor how to write one ourselves. This offends me greatly, and I would like to solve this! Mechanistic interpretability is a very young and small field, and there are a *lot* of open problems - if you would like to help, please try working on one! **If you want to skill up, check out [my guide to getting started](https://neelnanda.io/getting-started), and if you want to jump into an open problem check out my sequence [200 Concrete Open Problems in Mechanistic Interpretability](https://neelnanda.io/concrete-open-problems).**
+#
+# I wrote this library because after I left the Anthropic interpretability team and started doing independent research, I got extremely frustrated by the state of open source tooling. There's a lot of excellent infrastructure like HuggingFace and DeepSpeed to *use* or *train* models, but very little to dig into their internals and reverse engineer how they work. **This library tries to solve that**, and to make it easy to get into the field even if you don't work at an industry org with real infrastructure! The core features were heavily inspired by [Anthropic's excellent Garcon tool](https://transformer-circuits.pub/2021/garcon/index.html). Credit to Nelson Elhage and Chris Olah for building Garcon and showing me the value of good infrastructure for accelerating exploratory research!
+#
+# The core design principle I've followed is to enable exploratory analysis - one of the most fun parts of mechanistic interpretability compared to normal ML is the extremely short feedback loops! The point of this library is to keep the gap between having an experiment idea and seeing the results as small as possible, to make it easy for **research to feel like play** and to enter a flow state. This notebook demonstrates how the library works and how to use it, but if you want to see how well it works for exploratory research, check out [my notebook analysing Indirect Objection Identification](https://neelnanda.io/exploratory-analysis-demo) or [my recording of myself doing research](https://www.youtube.com/watch?v=yo4QvDn-vsU)!
+
+# %% [markdown]
+# ## Loading and Running Models
+#
+# TransformerLens comes loaded with >40 open source GPT-style models. You can load any of them in with `HookedTransformer.from_pretrained(MODEL_NAME)`. For this demo notebook we'll look at GPT-2 Small, an 80M parameter model, see the Available Models section for info on the rest.
+
+# %%
+device = utils.get_device()
+
+# %%
+# NBVAL_IGNORE_OUTPUT
+model = boot("gpt2", device=device)
+
+# %% [markdown]
+# To try the model out, let's find the loss on this text! Models can be run on a single string or a tensor of tokens (shape: [batch, position], all integers), and the possible return types are:
+# * "logits" (shape [batch, position, d_vocab], floats),
+# * "loss" (the cross-entropy loss when predicting the next token),
+# * "both" (a tuple of (logits, loss))
+# * None (run the model, but don't calculate the logits - this is faster when we only want to use intermediate activations)
+
+# %%
+model_description_text = """## Loading Models
+
+HookedTransformer comes loaded with >40 open source GPT-style models. You can load any of them in with `HookedTransformer.from_pretrained(MODEL_NAME)`. See my explainer for documentation of all supported models, and this table for hyper-parameters and the name used to load them. Each model is loaded into the consistent HookedTransformer architecture, designed to be clean, consistent and interpretability-friendly.
+
+For this demo notebook we'll look at GPT-2 Small, an 80M parameter model. To try the model the model out, let's find the loss on this paragraph!"""
+loss = model(model_description_text, return_type="loss")
+print("Model loss:", loss)
+
+# %% [markdown]
+# ## Caching all Activations
+#
+# The first basic operation when doing mechanistic interpretability is to break open the black box of the model and look at all of the internal activations of a model. This can be done with `logits, cache = model.run_with_cache(tokens)`. Let's try this out on the first line of the abstract of the GPT-2 paper.
+#
+# On `remove_batch_dim`
+#
+# Every activation inside the model begins with a batch dimension. Here, because we only entered a single batch dimension, that dimension is always length 1 and kinda annoying, so passing in the `remove_batch_dim=True` keyword removes it. `gpt2_cache_no_batch_dim = gpt2_cache.remove_batch_dim()` would have achieved the same effect.
+#
+
+# %%
+gpt2_text = "Natural language processing tasks, such as question answering, machine translation, reading comprehension, and summarization, are typically approached with supervised learning on taskspecific datasets."
+gpt2_tokens = model.to_tokens(gpt2_text)
+print(gpt2_tokens.device)
+gpt2_logits, gpt2_cache = model.run_with_cache(gpt2_tokens, remove_batch_dim=True)
+
+# %% [markdown]
+# Let's visualize the attention pattern of all the heads in layer 0, using [Alan Cooney's CircuitsVis library](https://github.com/alan-cooney/CircuitsVis) (based on [Anthropic's PySvelte library](https://github.com/anthropics/PySvelte)).
+#
+# We look this the attention pattern in `gpt2_cache`, an `ActivationCache` object, by entering in the name of the activation, followed by the layer index (here, the activation is called "attn" and the layer index is 0). This has shape [head_index, destination_position, source_position], and we use the `model.to_str_tokens` method to convert the text to a list of tokens as strings, since there is an attention weight between each pair of tokens.
+#
+# This visualization is interactive! Try hovering over a token or head, and click to lock. The grid on the top left and for each head is the attention pattern as a destination position by source position grid. It's lower triangular because GPT-2 has **causal attention**, attention can only look backwards, so information can only move forwards in the network.
+#
+# See the ActivationCache section for more on what `gpt2_cache` can do.
+
+# %%
+print(type(gpt2_cache))
+attention_pattern = gpt2_cache["pattern", 0, "attn"]
+print(attention_pattern.shape)
+gpt2_str_tokens = model.to_str_tokens(gpt2_text)
+
+# %%
+print("Layer 0 Head Attention Patterns:")
+
+
+# %% [markdown]
+# In this case, we only wanted the layer 0 attention patterns, but we are storing the internal activations from all locations in the model. It's convenient to have access to all activations, but this can be prohibitively expensive for memory use with larger models, batch sizes, or sequence lengths. In addition, we don't need to do the full forward pass through the model to collect layer 0 attention patterns. The following cell will collect only the layer 0 attention patterns and stop the forward pass at layer 1, requiring far less memory and compute.
+
+# %%
+attn_hook_name = "blocks.0.attn.hook_pattern"
+attn_layer = 0
+_, gpt2_attn_cache = model.run_with_cache(gpt2_tokens, remove_batch_dim=True, stop_at_layer=attn_layer + 1, names_filter=[attn_hook_name])
+gpt2_attn = gpt2_attn_cache[attn_hook_name]
+assert torch.equal(gpt2_attn, attention_pattern)
+
+# %% [markdown]
+# ## Hooks: Intervening on Activations
+
+# %% [markdown]
+# One of the great things about interpreting neural networks is that we have *full control* over our system. From a computational perspective, we know exactly what operations are going on inside (even if we don't know what they mean!). And we can make precise, surgical edits and see how the model's behaviour and other internals change. This is an extremely powerful tool, because it can let us eg set up careful counterfactuals and causal intervention to easily understand model behaviour.
+#
+# Accordingly, being able to do this is a pretty core operation, and this is one of the main things TransformerLens supports! The key feature here is **hook points**. Every activation inside the transformer is surrounded by a hook point, which allows us to edit or intervene on it.
+#
+# We do this by adding a **hook function** to that activation. The hook function maps `current_activation_value, hook_point` to `new_activation_value`. As the model is run, it computes that activation as normal, and then the hook function is applied to compute a replacement, and that is substituted in for the activation. The hook function can be an arbitrary Python function, so long as it returns a tensor of the correct shape.
+#
+# Relationship to PyTorch hooks
+#
+# [PyTorch hooks](https://blog.paperspace.com/pytorch-hooks-gradient-clipping-debugging/) are a great and underrated, yet incredibly janky, feature. They can act on a layer, and edit the input or output of that layer, or the gradient when applying autodiff. The key difference is that **Hook points** act on *activations* not layers. This means that you can intervene within a layer on each activation, and don't need to care about the precise layer structure of the transformer. And it's immediately clear exactly how the hook's effect is applied. This adjustment was shamelessly inspired by [Garcon's use of ProbePoints](https://transformer-circuits.pub/2021/garcon/index.html).
+#
+# They also come with a range of other quality of life improvements, like the model having a `model.reset_hooks()` method to remove all hooks, or helper methods to temporarily add hooks for a single forward pass - it is *incredibly* easy to shoot yourself in the foot with standard PyTorch hooks!
+# Technical details
+#
+# * We attach the hook to the attention pattern activation. There's one big pattern activation per layer, stacked across all heads, so we need to do some tensor manipulation to get a per-head score.
+# * Hook functions can access global state, so we make a big tensor to store the induction head score for each head, and then we just add the score for each head to the appropriate position in the tensor.
+# * To get a single hook function that works for each layer, we use the `hook.layer()` method to get the layer index (internally this is just inferred from the hook names).
+# * As we want to add this to *every* activation pattern hook point, rather than giving the string for an activation name, this time we give a **name filter**. This is a Boolean function on hook point names, and it adds the hook function to every hook point where the function evaluates as true.
+# * `run_with_hooks` allows us to enter a list of (act_name, hook_function) pairs to all be added at once, so we could also have done this by inputting a list with a hook for each layer.
+# Technical notes
+#
+# Note that if we input a string to the model, it's implicitly converted to a string with `to_tokens`.
+#
+# Note further that the log probs have shape `[batch, position, d_vocab]==[1, 8, 50257]`, with a vector of log probs predicting the next token for *every* token position. GPT-2 uses causal attention which means heads can only look backwards (equivalently, information can only move forwards in the model.), so the log probs at position k are only a function of the first k tokens, and it can't just cheat and look at the k+1 th token. This structure lets it generate text more efficiently, and lets it treat every *token* as a training example, rather than every *sequence*.
+# Why are low-rank factorized matrices useful for transformer interpretability?
+#
+# As argued in [A Mathematical Framework](https://transformer-circuits.pub/2021/framework/index.html), an unexpected fact about transformer attention heads is that rather than being best understood as keys, queries and values (and the requisite weight matrices), they're actually best understood as two low rank factorized matrices.
+# * **Where to move information from:** $W_QK = W_Q W_K^T$, used for determining the attention pattern - what source positions to move information from and what destination positions to move them to.
+# * Intuitively, residual stream -> query and residual stream -> key are linear maps, *and* `attention_score = query @ key.T` is a linear map, so the whole thing can be factored into one big bilinear form `residual @ W_QK @ residual.T`
+# * **What information to move:** $W_OV = W_V W_O$, used to determine what information to copy from the source position to the destination position (weighted by the attention pattern weight from that destination to that source).
+# * Intuitively, the residual stream is a `[position, d_model]` tensor (ignoring batch). The attention pattern acts on the *position* dimension (where to move information from and to) and the value and output weights act on the *d_model* dimension - ie *what* information is contained at that source position. So we can factor it all into `attention_pattern @ residual @ W_V @ W_O`, and so only need to care about `W_OV = W_V @ W_O`
+# * Note - the internal head dimension is smaller than the residual stream dimension, so the factorization is low rank. (here, `d_model=768` and `d_head=64`)
+#