5️⃣ Full Replication: Minimial Circuits and more

def scatter_embedding_vs_attn(
    attn_from_end_to_io: Float[Tensor, "batch"],
    attn_from_end_to_s: Float[Tensor, "batch"],
    projection_in_io_dir: Float[Tensor, "batch"],
    projection_in_s_dir: Float[Tensor, "batch"],
    layer: int,
    head: int,
):
    scatter(
        x=t.concat([attn_from_end_to_io, attn_from_end_to_s], dim=0),
        y=t.concat([projection_in_io_dir, projection_in_s_dir], dim=0),
        color=["IO"]  N + ["S"]  N,
        title=f"Projection of the output of {layer}.{head} along the name<br>embedding vs attention probability on name",
        title_x=0.5,
        labels={"x": "Attn prob on name", "y": "Dot w Name Embed", "color": "Name type"},
        color_discrete_sequence=["#72FF64", "#C9A5F7"],
        width=650,
    )
def calculate_and_show_scatter_embedding_vs_attn(
    layer: int,
    head: int,
    cache: ActivationCache = ioi_cache,
    dataset: IOIDataset = ioi_dataset,
) -> None:
    """
    Creates and plots a figure equivalent to 3(c) in the paper.
    This should involve computing the four 1D tensors:
        attn_from_end_to_io
        attn_from_end_to_s
        projection_in_io_dir
        projection_in_s_dir
    and then calling the scatter_embedding_vs_attn function.
    """
    # Get the value written to the residual stream at the end token by this head
    z = cache[utils.get_act_name("z", layer)][:, :, head]  # [batch seq d_head]
    N = z.size(0)
    output = z @ model.W_O[layer, head]  # [batch seq d_model]
    output_on_end_token = output[t.arange(N), dataset.word_idx["end"]]  # [batch d_model]
# Get the directions we'll be projecting onto
    io_unembedding = model.W_U.T[dataset.io_tokenIDs]  # [batch d_model]
    s_unembedding = model.W_U.T[dataset.s_tokenIDs]  # [batch d_model]
# Get the value of projections, by multiplying and summing over the d_model dimension
    projection_in_io_dir = (output_on_end_token  io_unembedding).sum(-1)  # [batch]
    projection_in_s_dir = (output_on_end_token  s_unembedding).sum(-1)  # [batch]
# Get attention probs, and index to get the probabilities from END -> IO / S
    attn_probs = cache["pattern", layer][:, head]  # [batch seqQ seqK]
    attn_from_end_to_io = attn_probs[
        t.arange(N), dataset.word_idx["end"], dataset.word_idx["IO"]
    ]  # [batch]
    attn_from_end_to_s = attn_probs[
        t.arange(N), dataset.word_idx["end"], dataset.word_idx["S1"]
    ]  # [batch]
# Show scatter plot
    scatter_embedding_vs_attn(
        attn_from_end_to_io,
        attn_from_end_to_s,
        projection_in_io_dir,
        projection_in_s_dir,
        layer,
        head,
    )

def get_copying_scores(
    model: HookedTransformer, k: int = 5, names: list = NAMES
) -> Float[Tensor, "2 layer-1 head"]:
    """
    Gets copying scores (both positive and negative) as described in page 6 of the IOI paper, for
    every (layer, head) pair in the model.
    Returns these in a 3D tensor (the first dimension is for positive vs negative).
    Omits the 0th layer, because this is before MLP0 (which we're claiming acts as an extended
    embedding).
    """
    results = t.zeros((2, model.cfg.n_layers, model.cfg.n_heads), device=device)
# Define components from our model (for typechecking, and cleaner code)
    embed: Embed = model.embed
    mlp0: MLP = model.blocks[0].mlp
    ln0: LayerNorm = model.blocks[0].ln2
    unembed: Unembed = model.unembed
    ln_final: LayerNorm = model.ln_final
# Get embeddings for the names in our list
    name_tokens: Int[Tensor, "batch 1"] = model.to_tokens(names, prepend_bos=False)
    name_embeddings: Int[Tensor, "batch 1 d_model"] = embed(name_tokens)
# Get residual stream after applying MLP
    resid_after_mlp1 = name_embeddings + mlp0(ln0(name_embeddings))
# Loop over all (layer, head) pairs
    for layer in range(1, model.cfg.n_layers):
        for head in range(model.cfg.n_heads):
            # Get W_OV matrix
            W_OV = model.W_V[layer, head] @ model.W_O[layer, head]
# Get residual stream after applying W_OV or -W_OV respectively
            # (note, because of bias b_U, it matters that we do sign flip here, not later)
            resid_after_OV_pos = resid_after_mlp1 @ W_OV
            resid_after_OV_neg = resid_after_mlp1 @ -W_OV
# Get logits from value of residual stream
            logits_pos = unembed(ln_final(resid_after_OV_pos)).squeeze()  # [batch d_vocab]
            logits_neg = unembed(ln_final(resid_after_OV_neg)).squeeze()  # [batch d_vocab]
# Check how many are in top k
            topk_logits: Int[Tensor, "batch k"] = t.topk(logits_pos, dim=-1, k=k).indices
            in_topk = (topk_logits == name_tokens).any(-1)
            # Check how many are in bottom k
            bottomk_logits: Int[Tensor, "batch k"] = t.topk(logits_neg, dim=-1, k=k).indices
            in_bottomk = (bottomk_logits == name_tokens).any(-1)
# Fill in results
            results[:, layer - 1, head] = t.tensor(
                [in_topk.float().mean(), in_bottomk.float().mean()]
            )
return results
copying_results = get_copying_scores(model)
imshow(
    copying_results,
    facet_col=0,
    facet_labels=["Positive copying scores", "Negative copying scores"],
    title="Copying scores of attention heads' OV circuits",
    width=900,
)
heads = {"name mover": [(9, 9), (10, 0), (9, 6)], "negative name mover": [(10, 7), (11, 10)]}
for i, name in enumerate(["name mover", "negative name mover"]):
    make_table(
        title=f"Copying Scores ({name} heads)",
        colnames=["Head", "Score"],
        cols=[
            list(map(str, heads[name])) + ["[dark_orange bold]Average"],
            [f"{copying_results[i, layer - 1, head]:.2%}" for (layer, head) in heads[name]]
            + [f"[dark_orange bold]{copying_results[i].mean():.2%}"],
        ],
    )

def generate_repeated_tokens(
    model: HookedTransformer, seq_len: int, batch: int = 1
) -> Float[Tensor, "batch 2seq_len"]:
    """
    Generates a sequence of repeated random tokens (no start token).
    """
    rep_tokens_half = t.randint(0, model.cfg.d_vocab, (batch, seq_len), dtype=t.int64)
    rep_tokens = t.cat([rep_tokens_half, rep_tokens_half], dim=-1).to(device)
    return rep_tokens
def get_attn_scores(
    model: HookedTransformer,
    seq_len: int,
    batch: int,
    head_type: Literal["duplicate", "prev", "induction"],
) -> Float[Tensor, "n_layers n_heads"]:
    """
    Returns attention scores for sequence of duplicated tokens, for every head.
    """
    rep_tokens = generate_repeated_tokens(model, seq_len, batch)
_, cache = model.run_with_cache(
        rep_tokens, return_type=None, names_filter=lambda name: name.endswith("pattern")
    )
# Get the right indices for the attention scores
if head_type == "duplicate":
        src_indices = range(seq_len)
        dest_indices = range(seq_len, 2  seq_len)
    elif head_type == "prev":
        src_indices = range(seq_len)
        dest_indices = range(1, seq_len + 1)
    elif head_type == "induction":
        dest_indices = range(seq_len, 2 * seq_len)
        src_indices = range(1, seq_len + 1)
results = t.zeros(model.cfg.n_layers, model.cfg.n_heads, device=device, dtype=t.float32)
    for layer in range(model.cfg.n_layers):
        for head in range(model.cfg.n_heads):
            attn_scores = cache["pattern", layer]  # [batch seqQ seqK]
            avg_attn_on_duplicates = attn_scores[:, head, dest_indices, src_indices].mean().item()
            results[layer, head] = avg_attn_on_duplicates
return results
def plot_early_head_validation_results(seq_len: int = 50, batch: int = 50):
    """
    Produces a plot that looks like Figure 18 in the paper.
    """
    head_types = ["duplicate", "prev", "induction"]
results = t.stack(
        [get_attn_scores(model, seq_len, batch, head_type=head_type) for head_type in head_types]
    )
imshow(
        results,
        facet_col=0,
        facet_labels=[
            f"{head_type.capitalize()} token attention prob.<br>on sequences of random tokens"
            for head_type in head_types
        ],
        labels={"x": "Head", "y": "Layer"},
        width=1300,
    )
model.reset_hooks()
plot_early_head_validation_results()

def compute_means_by_template(
    means_dataset: IOIDataset,
    model: HookedTransformer
) -> Float[Tensor, "layer batch seq head_idx d_head"]:
    '''
    Returns the mean of each head's output over the means dataset. This mean is
    computed separately for each group of prompts with the same template (these
    are given by means_dataset.groups).
    '''
    pass
def get_heads_and_posns_to_keep(
    means_dataset: IOIDataset,
    model: HookedTransformer,
    circuit: dict[str, list[tuple[int, int]]],
    seq_pos_to_keep: dict[str, str],
) -> dict[int, Bool[Tensor, "batch seq head"]]:
    '''
    Returns a dictionary mapping layers to a boolean mask giving the indices of the
    z output which shouldn't be mean-ablated.
    The output of this function will be used for the hook function that does ablation.
    '''
    pass
def hook_fn_mask_z(
    z: Float[Tensor, "batch seq head d_head"],
    hook: HookPoint,
    heads_and_posns_to_keep: dict[int, Bool[Tensor, "batch seq head"]],
    means: Float[Tensor, "layer batch seq head d_head"],
) -> Float[Tensor, "batch seq head d_head"]:
    '''
    Hook function which masks the z output of a transformer head.
    heads_and_posns_to_keep
        dict created with the get_heads_and_posns_to_keep function. This tells
        us where to mask.
    means
        Tensor of mean z values of the means_dataset over each group of prompts
        with the same template. This tells us what values to mask with.
    '''
    pass

def get_heads_and_posns_to_keep(
    means_dataset: IOIDataset,
    model: HookedTransformer,
    circuit: dict[str, list[tuple[int, int]]],
    seq_pos_to_keep: dict[str, str],
) -> dict[int, Bool[Tensor, "batch seq head"]]:
    """
    Returns a dictionary mapping layers to a boolean mask giving the indices of the z output which
    shouldn't be mean-ablated.
    The output of this function will be used for the hook function that does ablation.
    """
    heads_and_posns_to_keep = {}
    batch_size, seq_len, n_heads = len(means_dataset), means_dataset.max_len, model.cfg.n_heads
for layer in range(model.cfg.n_layers):
        mask = t.zeros(size=(batch_size, seq_len, n_heads))
for head_type, head_list in circuit.items():
            seq_pos = seq_pos_to_keep[head_type]
            indices = means_dataset.word_idx[seq_pos]
            for layer_idx, head_idx in head_list:
                if layer_idx == layer:
                    mask[range(batch_size), indices, head_idx] = 1
heads_and_posns_to_keep[layer] = mask.bool()
return heads_and_posns_to_keep
def hook_fn_mask_z(
    z: Float[Tensor, "batch seq head d_head"],
    hook: HookPoint,
    heads_and_posns_to_keep: dict[int, Bool[Tensor, "batch seq head"]],
    means: Float[Tensor, "layer batch seq head d_head"],
) -> Float[Tensor, "batch seq head d_head"]:
    """
    Hook function which masks the z output of a transformer head.
    heads_and_posns_to_keep
        dict created with the get_heads_and_posns_to_keep function. This tells us where to mask.
    means
        Tensor of mean z values of the means_dataset over each group of prompts with the same
        template. This tells us
        what values to mask with.
    """
    # Get the mask for this layer, and add d_head=1 dimension so it broadcasts correctly
    mask_for_this_layer = heads_and_posns_to_keep[hook.layer()].unsqueeze(-1).to(z.device)
# Set z values to the mean
    z = t.where(mask_for_this_layer, z, means[hook.layer()])
return z
def compute_means_by_template(
    means_dataset: IOIDataset, model: HookedTransformer
) -> Float[Tensor, "layer batch seq head_idx d_head"]:
    """
    Returns the mean of each head's output over the means dataset. This mean is computed separately
    for each group of prompts with the same template (these are given by means_dataset.groups).
    """
    # Cache the outputs of every head
    _, means_cache = model.run_with_cache(
        means_dataset.toks.long(),
        return_type=None,
        names_filter=lambda name: name.endswith("z"),
    )
    # Create tensor to store means
    n_layers, n_heads, d_head = model.cfg.n_layers, model.cfg.n_heads, model.cfg.d_head
    batch, seq_len = len(means_dataset), means_dataset.max_len
    means = t.zeros(size=(n_layers, batch, seq_len, n_heads, d_head), device=model.cfg.device)
# Get set of different templates for this data
    for layer in range(model.cfg.n_layers):
        z_for_this_layer = means_cache[utils.get_act_name("z", layer)]  # [batch seq head d_head]
        for template_group in means_dataset.groups:
            z_for_this_template = z_for_this_layer[template_group]
            z_means_for_this_template = einops.reduce(
                z_for_this_template, "batch seq head d_head -> seq head d_head", "mean"
            )
            means[layer, template_group] = z_means_for_this_template
return means
def add_mean_ablation_hook(
    model: HookedTransformer,
    means_dataset: IOIDataset,
    circuit: dict[str, list[tuple[int, int]]] = CIRCUIT,
    seq_pos_to_keep: dict[str, str] = SEQ_POS_TO_KEEP,
    is_permanent: bool = True,
) -> HookedTransformer:
    """
    Adds a permanent hook to the model, which ablates according to the circuit and seq_pos_to_keep
    dictionaries.
    In other words, when the model is run on ioi_dataset, every head's output will be replaced with
    the mean over means_dataset for sequences with the same template, except for a subset of heads
    and sequence positions as specified by the circuit and seq_pos_to_keep dicts.
    """
    model.reset_hooks(including_permanent=True)
# Compute the mean of each head's output on the ABC dataset, grouped by template
    means = compute_means_by_template(means_dataset, model)
# Convert this into a boolean map
    heads_and_posns_to_keep = get_heads_and_posns_to_keep(
        means_dataset, model, circuit, seq_pos_to_keep
    )
# Get a hook function which will patch in the mean z values for each head, at
    # all positions which aren't important for the circuit
    hook_fn = partial(hook_fn_mask_z, heads_and_posns_to_keep=heads_and_posns_to_keep, means=means)
# Apply hook
    model.add_hook(lambda name: name.endswith("z"), hook_fn, is_permanent=is_permanent)
return model

def get_score(
    model: HookedTransformer,
    ioi_dataset: IOIDataset,
    abc_dataset: IOIDataset,
    K: set[tuple[int, int]],
    C: dict[str, list[tuple[int, int]]],
) -> float:
    '''
    Returns the value F(C \ K), where F is the logit diff, C is the
    core circuit, and K is the set of circuit components to remove.
    '''
    pass
def get_minimality_score(
    model: HookedTransformer,
    ioi_dataset: IOIDataset,
    abc_dataset: IOIDataset,
    v: tuple[int, int],
    K: set[tuple[int, int]],
    C: dict[str, list[tuple[int, int]]] = CIRCUIT,
) -> float:
    '''
    Returns the value | F(C \ K_union_v) - F(C | K) |, where F is
    the logit diff, C is the core circuit, K is the set of circuit
    components to remove, and v is a head (not in K).
    '''
    pass
def get_all_minimality_scores(
    model: HookedTransformer,
    ioi_dataset: IOIDataset = ioi_dataset,
    abc_dataset: IOIDataset = abc_dataset,
    k_for_each_component: dict = K_FOR_EACH_COMPONENT
) -> dict[tuple[int, int], float]:
    '''
    Returns dict of minimality scores for every head in the model (as
    a fraction of F(M), the logit diff of the full model).
    Warning - this resets all hooks at the end (including permanent).
    '''
    pass
minimality_scores = get_all_minimality_scores(model)
plot_minimal_set_results(minimality_scores)

def get_score(
    model: HookedTransformer,
    ioi_dataset: IOIDataset,
    abc_dataset: IOIDataset,
    K: set[tuple[int, int]],
    C: dict[str, list[tuple[int, int]]],
) -> float:
    """
    Returns the value F(C \ K), where F is the logit diff, C is the core circuit, and K is the set
    of circuit components to remove.
    """
    C_excl_K = {k: [head for head in v if head not in K] for k, v in C.items()}
    model = add_mean_ablation_hook(model, abc_dataset, C_excl_K, SEQ_POS_TO_KEEP)
    logits = model(ioi_dataset.toks)
    score = logits_to_ave_logit_diff_2(logits, ioi_dataset).item()
return score
def get_minimality_score(
    model: HookedTransformer,
    ioi_dataset: IOIDataset,
    abc_dataset: IOIDataset,
    v: tuple[int, int],
    K: set[tuple[int, int]],
    C: dict[str, list[tuple[int, int]]] = CIRCUIT,
) -> float:
    """
    Returns the value | F(C \ K_union_v) - F(C | K) |, where F is the logit diff, C is the core
    circuit, K is the set of circuit components to remove, and v is a head (not in K).
    """
    assert v not in K
    K_union_v = K | {v}
    C_excl_K_score = get_score(model, ioi_dataset, abc_dataset, K, C)
    C_excl_Kv_score = get_score(model, ioi_dataset, abc_dataset, K_union_v, C)
return abs(C_excl_K_score - C_excl_Kv_score)
def get_all_minimality_scores(
    model: HookedTransformer,
    ioi_dataset: IOIDataset = ioi_dataset,
    abc_dataset: IOIDataset = abc_dataset,
    k_for_each_component: dict = K_FOR_EACH_COMPONENT,
) -> dict[tuple[int, int], float]:
    """
    Returns dict of minimality scores for every head in the model (as a fraction of F(M), the
    logit diff of the full model).
    Warning - this resets all hooks at the end (including permanent).
    """
    # Get full circuit score F(M), to divide minimality scores by
    model.reset_hooks(including_permanent=True)
    logits = model(ioi_dataset.toks)
    full_circuit_score = logits_to_ave_logit_diff_2(logits, ioi_dataset).item()
# Get all minimality scores, using the get_minimality_score function
    minimality_scores = {}
    for v, K in tqdm(k_for_each_component.items()):
        score = get_minimality_score(model, ioi_dataset, abc_dataset, v, K)
        minimality_scores[v] = score / full_circuit_score
model.reset_hooks(including_permanent=True)
return minimality_scores
minimality_scores = get_all_minimality_scores(model)