Attribution Analysis - Introduction to Interpretability for Language Models

import re
from string import Template

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.feature_extraction.text import ENGLISH_STOP_WORDS
from sklearn.metrics import balanced_accuracy_score

In the last chapter, we identified distinctive tokens using a genre classifier. This chapter attempts to do the same task, except we will use the model as a classifier of its own short stories. Setting up the task in this way enables us to extract signal directly from the model as it processes text. Such signal provides a window into the inner workings of the model.

Model as Classifier¶

We use the same short stories from last time: genre-conditioned outputs from the instruct-tuned variant of Olmo 3 7B. But instead of building an external classifier from story text, we prompt the model directly in a question/answer setup.

We use this system prompt:

SYSTEM_PROMPT = (
    "You are a literary classification assistant. "
    "Read the story and the answer options, then respond with the single "
    "letter of the option that best identifies the story's genre."
)

...and prompt the model as follows:

USER_TEMPLATE = Template(
    "Read the following story and identify its genre.\n\n====\n\n"
    "Story:\n$story\n\n====\n\n"
    "$options"
)

Above, $story represents story text. For a given story, $options represents a set of multiple-choice options, e.g.:

options = """
A) Horror
B) Detective fiction
C) Romance
D) Science fiction
E) Thriller
"""

We leave out the NULL genre for this task. For each story, we also randomly shuffle the mapping from answer letters to genres. This avoids tying a certain genre to a fixed letter across examples. The correct answer is therefore the letter assigned to the story’s true genre in that shuffled mapping.

The DataFrame below contains the prompt + letter/genre combinations for 100 (randomly sampled) examples of every genre in the dataset:

df = pd.read_parquet("../data/w2-d2_genre-attributions.parquet")

Token Attribution¶

As you’ll see, we didn’t actually let the model generate new tokens. It wasn’t necessary: The following analysis uses the model’s next-token logits for the five answer letters---that is, the scores the model assigned to A through E immediately after the prompt. Token attribution is then computed by backpropagating a scalar objective built from these logits.

Objective¶

Specifically, our objective compares the correct-letter logit against the incorrect-letter logits. We use the correct answer’s logit, and compare it to LogSumExp over the other answer logits.

Let:

$y$ = correct answer letter
$A$ = set of answer letters
$z_a$ = final-position logit for answer letter $a$
$e_i$ = input embedding for token $i$

Our objective is:

S = z_y - \log \sum_{a \in A,\ a \neq y} \exp(z_a)

(2)

The first term, $z_y$ , rewards the model for assigning a high score to the correct answer. The second term summarizes the scores of all incorrect answers. So this objective gets larger when the model prefers the correct answer over the wrong ones.

Then, for each input token, we ask: If we made a tiny change to this token’s embedding, how much would that change the objective?

Token attribution is:

\text{attr}_i = \left\| \frac{\partial S}{\partial e_i} \right\|_2

(3)

In other words: Compute the correct answer logit minus the log-sum-exp of the incorrect answer logits, backpropagate that objective, and use the L2 norm of each token embedding gradient as the token attribution score. Tokens with larger attribution scores are tokens where small changes to the embedding would have a larger local effect on the model’s preference for the correct answer.

Attribution scores¶

Attribution scores are stored in token_grads. For each story, we have a score for every token:

samp = df.sample()
print("Number of scores:", samp["token_grads"].apply(len).item())

Now, we could work directly with these, but two problems arise. First, attribution scores are assigned to all tokens. That includes subwords, prompt formatting tokens, and punctuation marks. Oftentimes, those tokens get the highest scores in a gradient objective setup:

(
    samp
    .explode(["tokens", "token_grads"])
    .sort_values("token_grads", ascending=False)["tokens"]
    .head(10)
    .tolist()
)

An all-encompassing explanation of model behavior would, in theory, be able to tell us why the model attributes so much weight to these tokens. But we don’t have that---no one does, in fact. So what should we do with those tokens?

While you contemplate that, consider the second problem: The story prompt tokens also get attribution scores, and models tend to assign them high importance as well:

samp["tokens"].item()[:10].tolist()

Do we let these factor into our analysis? Maybe for a study of instruct-tuning we would, but not for our focus on distinctive genre tokens. Again: What to do?

Filtering and aggregating tokens¶

The answer: We need to make the unit of analysis match the question we’re asking. If we want to find distinctive genre cues, we should focus on the story text, not the surrounding prompt. Likewise, we should give special attention to lexical information over punctuation, whitespace, and other formatting tokens. Those tokens certainly can affect the model, but we will save that analysis for another time. Finally, because model tokens are not the same as words, we aggregate token-level attribution scores into word-level scores. This gives us units that are easier to interpret than subword-segmented text.

For us, a “token” has the same definition from the previous chapter: It’s a sequence of three or more alphabetic characters:

RE_PATTERN = re.compile(r"(?u)\b[a-zA-Z]{3,}\b")

Whitespace, punctuation, short words, etc. should all be removed. As before, we’ll also remove stop words (and the prompt tokens, of course). To our usual English stop word list, let’s also include some words that tend to stick around because their contractions get removed by the pattern above.

cruft = [
    "didn", "don", "doesn", "isn", "aren", "wasn", "weren", "hasn", "haven", 
    "hadn", "won", "wouldn", "couldn", "shouldn", "can", "mustn", "needn", 
    "shan", "mightn",
]

ENGLISH_STOP_WORDS = ENGLISH_STOP_WORDS | set(cruft)

With these all set, we define a function that filters and aggregates tokens. It re-segments the string representation of a story (stored in text) using the pattern above. Then, for every token, it:

Filters stop words
Ignores prompt tokens before/after the story
Rebuilds subword-segmented strings into a single word
Calculates the sum and mean of all subword tokens in the rebuilt word

def aggregate_tokens(
    text, offsets, gradients, pattern, mask=None, stop_words=None
):
    """Aggregate and filter subword tokens.

    Uses a regular expression match to find word-like tokens in story text and
    filters those tokens via a stop word list.

    Parameters
    ----------
    text : str
        Story text
    offsets : np.ndarray
        Token offsets from the LLM tokenizer, shape (n_token, 2)
    gradients : np.ndarray
        Token gradients, shape (n_token,)
    pattern : re.Pattern
        Regular expression pattern for finding tokens
    mask : np.ndarray or None
        Story mask for ignoring prompt tokens, shape (n_token,)
    stop_words : iterable or None
        Stop word list

    Returns
    -------
    list[dict]
        Aggregated words with metadata:
        - word: combined token string
        - span_start: start of word span
        - span_end: end of word span
        - token_indices: token indices in word
        - subwords: individual tokens in token string
        - n_tokens: number of tokens in word
        - grad_sum: sum of token gradients
        - grad_mean: mean of token gradients
    """
    if mask is None:
        mask = np.ones(len(offsets), dtype=bool)
    
    if stop_words is None:
        stop_words = set()

    # Stack offsets so offsets[:, 0] == start, offsets[:, 1] == end
    offsets = np.vstack(offsets)
    assert len(mask) == len(offsets), "Mask length doesn't match offset length"

    output = []

    for match in pattern.finditer(text):
        word = match.group(0)
        span_start, span_end = match.span()

        # Ignore stopwords
        if word.lower() in stop_words:
            continue

        # Overlap if: token_start < span_end and token_end > span_start
        (overlapping,) = np.where(
            mask
            & (offsets[:, 0] < span_end)
            & (offsets[:, 1] > span_start)
            & (offsets[:, 1] > offsets[:, 0])
        )

        if len(overlapping) == 0:
            continue

        # Build our word from tokens
        subwords = []
        
        for i in overlapping:
            tok_start, tok_end = offsets[i]

            clipped_start = max(tok_start, span_start)
            clipped_end = min(tok_end, span_end)

            subwords.append(text[clipped_start:clipped_end])

        # Get token gradients
        grads = gradients[overlapping]

        # Construct metadata
        row = {
            "word": word,
            "span_start": span_start,
            "span_end": span_end,
            "token_indices": overlapping.tolist(),
            "subwords": subwords,
            "n_tokens": len(overlapping),
            "grad_sum": grads.sum(),
            "grad_mean": grads.mean(),
        }

        output.append(row)

    return output

With aggregate_tokens() defined, we define a helper function that will let us run our aggregator across all DataFrame rows at once:

def aggregate(row):
    """Helper function for using ``aggregate_tokens()`` on DataFrame rows.

    Parameters
    ----------
    row : pd.Series
        DataFrame row

    Returns
    -------
    pd.DataFrame
        Aggregated token DataFrame
    """
    agg = aggregate_tokens(
        text=row["prompt"].lower(),
        offsets=row["token_offsets"],
        gradients=row["token_grads"],
        pattern=RE_PATTERN,
        mask=row["story_token_mask"],
        stop_words=ENGLISH_STOP_WORDS,
    )

    # Convert to a DataFrame and assign some metadata
    agg = pd.DataFrame(agg)
    agg["story_id"] = row["story_id"]
    agg["genre"] = row["true_genre"]

    return agg

Now, we apply aggregate() and concatenate the individual DataFrames into a single one:

agg = df.apply(aggregate, axis=1)
agg = pd.concat(agg.tolist(), ignore_index=True)
print("Shape of aggregated words (n_word, n_feat):", agg.shape)

Here’s a glimpse of the results:

agg.head()

With this data, we can use a story_id to find words with the highest attribution scores according to our objective. Call the .isin() method from agg to find those indices in agg that match story_id in samp:

indices = agg["story_id"].isin(samp["story_id"])

Then use the result to index agg:

agg[indices].sort_values("grad_mean", ascending=False).head(10)

One caveat: These gradient attributions aren’t signed. The underlying gradient has direction, but here we reduce each token’s embedding gradient to a single number by taking its norm. So we don’t know whether a word pushed the model toward the correct genre or away from an incorrect one.

Instead, these scores measure local sensitivity: When the model was deciding between genres, these were the words where small changes to the input embedding would have had the largest effect on our objective. As we’ll see in a moment, though, this simple representation can give us a good sense of genre.

Genre lift¶

Individual token attributes are useful as examples, but it would be helpful to look at genre-level attributions as well. Here we ask: When a word appears, how much does the model’s genre score tend to be sensitive to it, relative to the same sensitivity in other genres?

We’ll capture this with lift. For a given word, lift is the difference between how influential that word is in one genre and how influential it is on average across the other genres where it appears. A positive lift means the word pulls harder on the model’s genre score here than it does elsewhere. A negative means the opposite. A lift near zero means the word is about equally influential everywhere, so it tells us little about genre.

First, we aggregate words within each story. This prevents a single story from dominating genre-level results just because it repeats a word many times. Each entry below maps an output column to a (source_column, function) pair in Pandas. After Pandas runs groupby, it applies the function to the source column within each group.

STORY_WORD_AGG = {
    "story_grad_sum": ("grad_mean", "sum"),  # Total gradient mass in story
    "story_grad_mean": ("grad_mean", "mean"),  # Avg. influence per occurrence
    "story_n": ("grad_mean", "size"),  # Num. word occurrences
}

Then we aggregate those story-level word scores up to the genre level.

GENRE_WORD_AGG = {
    "grad_mean_across_stories": ("story_grad_mean", "mean"),  # Avg. influence
    "grad_sum_across_stories": ("story_grad_sum", "sum"),  # Total gradient mass
    "n_stories": ("story_id", "size"),  # Num. stories each word appears
    "n_occurrences": ("story_n", "sum"),  # Total word occurrences
}

The function below performs these aggregations:

def distinctive_words(df, min_stories=10, top_n=25):
    """Get distinctive words per genre.

    Parameters
    ----------
    df : pd.DataFrame
        Aggregated word DataFrame for all stories in the dataset
    min_stories : int
        Minimum number of stories a word must appear in within a genre
    top_n : int
        Number of distinctive words to return per genre

    Returns
    -------
    pd.DataFrame
        Distinctive words per genre, ranked by genre_lift
    """
    # One row per story/genre/word
    story_word = (
        df
        .groupby(["story_id", "genre", "word"], as_index=False)
        .agg(**STORY_WORD_AGG)
    )

    # One row per genre/word
    genre_word = (
        story_word
        .groupby(["genre", "word"], as_index=False)
        .agg(**GENRE_WORD_AGG)
    )

    # Group words, then count the number of word occurences and unique words
    word_group = genre_word.groupby("word")
    word_sum = word_group["grad_mean_across_stories"].transform("sum")
    word_cnt = word_group["grad_mean_across_stories"].transform("count")

    # Subtract this genre's own value before averaging
    genre_word["global_grad_mean"] = (
        (word_sum - genre_word["grad_mean_across_stories"]) / (word_cnt - 1)
    )

    # Calculate lift
    genre_word["genre_lift"] = (
        genre_word["grad_mean_across_stories"] - genre_word["global_grad_mean"]
    )

    # Filter to words above min_stories and return the top lifted words
    return (
        genre_word
        .query("n_stories >= @min_stories")
        .sort_values(["genre", "genre_lift"], ascending=[True, False])
        .groupby("genre", group_keys=False)
        .head(top_n)
        .reset_index(drop=True)
    )

Let’s call it:

top_distinctive = distinctive_words(agg, min_stories=10, top_n=25)

And plot:

g = sns.catplot(
    data=top_distinctive,
    kind="bar",
    x="genre_lift",
    y="word",
    col="genre",
    col_wrap=2,
    sharey=False,
    height=FACET_HEIGHT,
    aspect=FACET_ASPECT,
)
g.set_axis_labels("Genre lift", "Word")
g.set_titles("{col_name}")
g.fig.suptitle("Most Distinctive High-Sensitivity Words by Genre", y=1.05)
plt.show()

Assessing Model Performance¶

These are compelling results, but can we actually trust them, given how they were obtained? Is the multiple-choice question/answer setup a good classifier? Let’s find out.

y_true = df["correct_letter"]
y_pred = df["predicted_letter"]
df["correct"] = y_true == y_pred

mcqa_acc = balanced_accuracy_score(y_true, y_pred)
print(f"Letter/predicted letter accuracy: {100 * mcqa_acc:.2f}%")

This is troubling. When we prompt a model with this setup and use letter logits for classification, our accuracy drops significantly. Is this a content problem, a form problem, or both?

Content problem: Lower accuracy for some genres
Form problem: Model prefers some answer letters, especially when it’s wrong

First, let’s see how accurate the model is over each genre.

genre_acc = df.groupby("true_genre").agg(acc=("correct", "mean"))
genre_acc.reset_index(inplace=True)

And plot:

fig, ax = plt.subplots()
sns.barplot(data=genre_acc, x="true_genre", y="acc", ax=ax)
ax.set(title="Per-Genre Accuracy", xlabel="Genre", ylabel="Accuracy")
plt.show()

The variance between accuracy scores here is pretty substantial. While the model performs decently well for horror, romance, and science fiction, its answers for detective fiction and thriller are quite bad. That actually connects back to our last chapter: Detective fiction and thriller saw substantial overlap in our genre mixing experiment.

So, part of the problem here is that the model struggles to classify its own outputs in a multiple-choice question/answer setup. But we should also consider whether the setup itself impacts the model’s classification performance.

For example, perhaps the model is biased towards answering one letter over others, regardless of what that letter stands for. We can generate a confusion matrix to investigate this.

letters = list("ABCDE")
confusion_matrix = (
    pd.crosstab(y_true, y_pred).reindex(columns=letters, index=letters)
)

Since we’ve randomized letter-genre assignments, we should see a strong diagonal on this plot:

fig, ax = plt.subplots()
sns.heatmap(confusion_matrix, cmap="Blues", annot=True, fmt="d", ax=ax)
ax.set(
    title="Confusion Matrix: Correct vs. Predicted Letter",
    xlabel="Predicted letter",
    ylabel="Correct letter",
)
plt.show()

B, D, and E perform well here, but observe how the model also selects these letters more frequently when the correct answer is A or C. To make this bias precise, we compare the letters the model actually picks when it’s wrong against the letters we’d expect it to pick under a simple baseline: “If the model is wrong, it chooses randomly among the four other letters.” So if the correct answer is A, a wrong answer can only be B through E, and under the baseline those mistakes spread evenly across the four.

The function below does this comparison. It takes the off-diagonal predictions, computes the observed and expected wrong-letter counts, converts both to shares, and returns their difference.

LETTERS = list("ABCDE")

def letter_bias(true_col, pred_col):
    """Compare observed vs. expected wrong-letter choices.

    Parameters
    ----------
    true_col : pd.Series
        Correct letters
    pred_col : pd.Series
        Predicted letters

    Returns
    -------
    pd.DataFrame
        Per-letter excess_share: positive means the model overuses this letter
        when wrong, negative means it avoids it
    """
    confusion_matrix = (
        pd.crosstab(true_col, pred_col)
        .reindex(index=LETTERS, columns=LETTERS, fill_value=0)
    )

    # Keep only wrong predictions
    off = confusion_matrix.values.copy()
    np.fill_diagonal(off, 0)
    off = pd.DataFrame(off, index=LETTERS, columns=LETTERS)

    # Observed: how often each letter is chosen when wrong. Expected: each
    # row's errors spread evenly over the other k-1 letters
    k = len(LETTERS)
    observed = off.sum(axis=0)
    row_errors = off.sum(axis=1)
    total_errors = row_errors.sum()
    expected = (total_errors - row_errors) / (k - 1)

    bias = pd.DataFrame({"observed": observed, "expected": expected})

    if total_errors > 0:
        bias["observed_share"] = bias["observed"] / total_errors
        bias["expected_share"] = bias["expected"] / total_errors
    else:
        bias["observed_share"] = 0.0 
        bias["expected_share"] = 0.0

    bias["excess_share"] = bias["observed_share"] - bias["expected_share"]

    return bias.reset_index(names="predicted_letter")

We apply it to the whole dataset and plot the excess share. Values above zero mean the model overuses that letter when wrong; values below zero mean it avoids it.

bias = letter_bias(y_true, y_pred)

g = sns.catplot(
    data=bias,
    kind="bar",
    x="predicted_letter",
    y="excess_share",
    height=FACET_HEIGHT,
    aspect=FACET_ASPECT,
)
g.refline(y=0, linestyle="--", color="black")
g.set_axis_labels("Predicted letter", "Observed share - Expected share")
g.fig.suptitle("Excess Wrong-Answer Share by Letter", y=1.05)
plt.show()

The model clearly doesn’t guess uniformly: A and C sit below zero while B, D, and E sit above. We can summarize the whole pattern in a single number: half the sum of the absolute excess shares, i.e. the share of wrong-answer mass we’d need to move to reach the uniform baseline.

letter_bias_score = 0.5 * bias["excess_share"].abs().sum()
print(f"Overall letter-bias score: {letter_bias_score:.3f}")

So we’d need to redistribute roughly a third of the model’s wrong-answer mass to match uniform guessing. That puts a big asterisk on the token attributions above: Our setup assumes the model is a good classifier when prompted with multiple-choice questions, but we’ve shown that it isn’t. The attributions, then, may not faithfully represent which tokens the model associates with each genre.

The Logit Lens¶

The remainder of this chapter investigates Olmo’s answer bias using a popular technique in mechanistic interpretability: the logit lens. Researchers use the logit lens to observe how next-token logits change from one layer to the next. After each Transformer block, they run a model’s hidden states through its final unembedding layer to transform logits into a token probability distribution. In effect, the logit lens asks: If we sampled from next-token candidates at layer $n$ , what token would the model likely select?

For our purposes, we can use the logit lens to determine where the model made its classification decision. Is there, for example, a layer in the model where it settles on an answer and never changes from there? Does one layer prefer one letter over others?

Logits for our five letters are stored in logits_by_layer. For every story, we collect 33 layers’ worth of our 5 answer logits. Stacking them gives a layer-by-letter matrix:

logit_matrix = np.vstack(samp["logits_by_layer"].item())
print("Number of layers:", logit_matrix.shape)

We use the function below to extract them for every story. This function also computes a margin between logits for the correct letter and the second-best letter. A high positive margin on the correct letter means the model is more confident in its answer. Conversely, a large negative margin means some other letter outscored the correct one.

def layer_logits_to_frame(row):
    """Format a story's layer logits into a DataFrame.

    Parameters
    ----------
    row : pd.Series
        Story data, including genre/prediction metadata as well as logits

    Returns
    -------
    pd.DataFrame
        Logit DataFrame with columns:
        - layer: layer
        - genre: genre
        - correct_letter: correct label for genre
        - final_predicted_letter: final predicted letter
        - layer_predicted_letter: layer's genre prediction
        - layer_correct: whether the layer prediction is correct
        - matches_final_prediction: whether layer prediction matches final
        - correct_margin: margin between correct answer and next-best
        - [A..E]: logit for each letter label
    """
    logits = np.vstack(row["logits_by_layer"])
    n_layer = logits.shape[0]

    correct_idx = LETTERS.index(row["correct_letter"])
    pred_idx = logits.argmax(axis=1)
    layer_pred = np.asarray(LETTERS)[pred_idx]

    # Margin between the correct logit and the best competing logit
    correct_logit = logits[:, correct_idx]
    other_cols = np.arange(logits.shape[1]) != correct_idx
    correct_margin = correct_logit - logits[:, other_cols].max(axis=1)

    lens = pd.DataFrame({
        "layer": np.arange(n_layer),
        "genre": row["true_genre"],
        "correct_letter": row["correct_letter"],
        "final_predicted_letter": row["predicted_letter"],
        "layer_predicted_letter": layer_pred,
        "layer_correct": layer_pred == row["correct_letter"],
        "matches_final_prediction": layer_pred == row["predicted_letter"],
        "correct_margin": correct_margin,
    })

    for i, letter in enumerate(LETTERS):
        lens[letter] = logits[:, i]

    return lens

Let’s run our function:

logit_lens = df.apply(layer_logits_to_frame, axis=1)
logit_lens = pd.concat(logit_lens.tolist(), ignore_index=True)
print("Logit lens shape:", logit_lens.shape)

Prediction margins¶

Does the model increasingly favor the correct answer as depth increases? We look at the per-genre view directly.

genre_layer_summary = (
    logit_lens
    .groupby(["genre", "layer"], as_index=False)
    .agg(
        correct_margin=("correct_margin", "mean"),
    )
)

And plot:

g = sns.relplot(
    data=genre_layer_summary,
    kind="line",
    x="layer",
    y="correct_margin",
    col="genre",
    col_wrap=2,
    marker="o",
    facet_kws={"sharey": True},
    height=FACET_HEIGHT,
    aspect=FACET_ASPECT,
)
g.refline(y=0, linestyle="--", color="black")
g.set_axis_labels("Layer", "Correct margin")
g.set_titles("{col_name}")
g.fig.suptitle("Layer-wise Correct-Answer Margins by Genre", y=1.05)
plt.show()

A big jump happens between layers 20-22. The model’s layer-wise predictions for the three stable genres solidify there, with margins between the correct logit and the next-best logit staying positive. Detective fiction and thriller suffer, however: Correct vs. next-best margins are negative throughout and don’t see the same jump. This matches with what we saw in the previous chapter, where our own classifier struggled to classify these two genres.

Now we ask: Which letters win the argmax at each layer? This could indicate whether there’s a particular layer in the model that injects the bias we’ve been seeing.

First, we cross-tabulate layer with the predicted letter:

letter_dist = pd.crosstab(
    logit_lens["layer"],
    logit_lens["layer_predicted_letter"],
    normalize="index",
)

And plot:

fig, ax = plt.subplots()
sns.heatmap(letter_dist, cmap="Blues", vmin=0, vmax=1, ax=ax)
ax.set(
    title="Predicted-Letter Distribution by Layer",
    xlabel="Predicted letter",
    ylabel="Layer",
)
plt.show()

The model strongly prefers A in the first ~6 layers. Perhaps this is a default answer state that it later revises. The middle layers, meanwhile, favor B; then around layer 22 predictions spread toward C, D, and E. But C hardly benefits from this spread: The model very rarely selects it. Likewise, A drops off precipitously in the last 4 layers, holding only a small share of letter predictions by the time input passes through the entire model.

Layer-wise bias¶

Is this behavior shaped by the same form bias we found earlier? We can run letter_bias() across each layer to investigate this.

bias_by_layer = []
for layer, layer_df in logit_lens.groupby("layer"):
    bias = letter_bias(
        layer_df["correct_letter"], layer_df["layer_predicted_letter"]
    )

    bias["layer"] = layer
    bias_by_layer.append(bias)

bias_by_layer = pd.concat(bias_by_layer, ignore_index=True)

Now we plot excess wrong-answer share as a heatmap, layer by layer. First, we use .pivot() to reshape the long bias table into a layer-by-letter grid.

excess_plot = bias_by_layer.pivot(
    index="layer", columns="predicted_letter", values="excess_share",
)

And plot:

fig, ax = plt.subplots()
sns.heatmap(excess_plot, cmap="RdBu_r", center=0, ax=ax)
ax.set(
    title="Excess Wrong-Answer Share by Layer",
    xlabel="Predicted letter",
    ylabel="Layer",
)
plt.show()

This heatmap looks quite similar to the last one, except it isolates the model’s mistakes. The shading now shows how much wrong-answer mass each letter carries at each layer, as opposed to its overall share of predictions.

Finally, we overlay two curves: how strong the letter bias is at each layer and how often each layer already agrees with the model’s final prediction. Together they show whether bias shifts around the same time predictions stabilize.

We need a few pieces of data for this: the bias score (from bias_by_layer) and the final-prediction agreement (from logit_lens). We compute each, merge them on layer, and use .melt() to convert them into a long form for plotting.

layer_bias_score = (
    bias_by_layer
    .groupby("layer", as_index=False)
    .agg(letter_bias_score=("excess_share", lambda x: 0.5 * x.abs().sum()))
)

final_agreement = (
    logit_lens
    .groupby("layer", as_index=False)
    .agg(final_agreement=("matches_final_prediction", "mean"))
)

layer_dynamics = (
    layer_bias_score
    .merge(final_agreement, on="layer")
    .melt(id_vars="layer", var_name="measure", value_name="value")
)

layer_dynamics["measure"] = layer_dynamics["measure"].map({
    "letter_bias_score": "Letter-bias score",
    "final_agreement": "Agreement with final prediction",
})

Time to plot.

g = sns.relplot(
    data=layer_dynamics,
    kind="line",
    x="layer",
    y="value",
    hue="measure",
    marker="o",
    height=FACET_HEIGHT,
    aspect=FACET_ASPECT,
)
g.set_axis_labels("Layer", "Value")
g.fig.suptitle(
    "When Do Predictions Stabilize, and When Does Letter Bias Appear?", y=1.05
)
plt.show()

Plot summary:

Early layers (0-5): high bias with near-zero final agreement. The model has a strong default but hasn’t started computing the real answer yet. This isn’t so much a bias as it is a prior---the model’s default.
Middle layers (6-21): bias remains high while agreement begins to climb. In these layers, we see the model forming an answer.
Crossover layer (~22): agreement overtakes bias. Notably, this is the same layer where our margins turned positive and the predictions “spread out.” One might hypothesize that this is where the model forms its answer.
Late layers (22-32): agreement climbs to 1.0 while bias drifts down. The latter stays meaningfully above zero and spikes near layer 30.

For our purposes, these late layer logits provide good evidence for what we saw earlier with the letter bias. Even as the model works its way toward a final prediction, it carries this bias forward to the end.