Metamodeling - Introduction to Interpretability for Language Models

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

This chapter is the third and final entry in our genre classification study. In our effort to find distinctive tokens for each genre, we will look once more at model internals. But instead of using next-token logits from a brittle multiple-choice question/answer prompt, we use the hidden states from Olmo 3 directly. This data source helps us to overcome the letter bias problem we uncovered in the previous chapter, and it will also demonstrate how metamodels can be used for interpretability work.

We will again load our DataFrame of story classifications. But this time, we only need a few columns, so we subset them at load time with columns=keep.

keep = ["story_id", "correct_letter", "predicted_letter", "true_genre"]
df = pd.read_parquet("../data/w2-d2_genre-attributions.parquet", columns=keep)

We also load a cache of hidden states. For each story in df, this records layer-wise hidden states.

array_cache = np.load("../data/w2-d3_genre-hidden-states.npz")

Hidden states are stored under hidden_states, while story IDs are under story_ids.

hidden_states = array_cache["hidden_states"]
story_ids = array_cache["story_ids"]

Finally, we re-index our DataFrame to ensure it’s aligned with the order of the story IDs from the cache.

df = df.set_index("story_id").loc[story_ids].reset_index()

Linear Probing¶

Remember that hidden states are a model’s internal representations of information. Unlike next-token logits---which already live in vocabulary space, one score per letter answer---hidden states are high-dimensional, context-dependent vectors. The ones in our cache correspond to last-token information at each layer in the model, including the initial embedding matrix:

print("Hidden states for a story:", hidden_states[0].shape)

No letter answer predictions here---and no genre labels, either. But genre information should still be present in the model’s internal states in some form. The model conditioned on genre when it wrote these stories, so it must have some way of representing differences between, say, science fiction and romance. The question, then, is this: How can we get this information?

We’ll do this by “reading” genre from hidden states via another model. Below, we train a linear probe for each layer. A probe is a small, simple model trained to predict some property---here, genre---from a network’s internal activations. If the probe can accurately predict genre from a layer’s hidden states, we take that as evidence the information is present at that layer.

The probe itself is just a linear classification. For a hidden state at a given layer, it scores each genre, then normalizes those scores into a probability distribution with our friend softmax:

P(g \mid \mathbf{h}) = \frac{\exp(\mathbf{w}_g^\top \mathbf{h} + b_g)} {\sum_{g' \in G} \exp(\mathbf{w}_{g'}^\top \mathbf{h} + b_{g'})}

(1)

Where:

$\mathbf{h} \in \mathbb{R}^d$ is the hidden state (a $d$ -dimensional vector)
$g$ is a genre, and $G$ is the set of all genres
$\mathbf{w}_g$ is the weight vector for genre $g$
$b_g$ is the bias term for genre $g$
$\mathbf{w}_g^\top \mathbf{h} + b_g$ is the linear score for genre $g$

The predicted genre is the one with the highest probability, $\arg\max_g P(g \mid \mathbf{h})$ .

Training probes¶

Below, we write a function to train a probe on hidden states at a given layer. It:

Initializes a pipeline that standardizes the hidden states and sends them to the probe
Fits the probe using cross-validation, as we did with our Multinomial Bayesian model
Reports the mean accuracy across folds

def probe_layer(X, y, n_splits=5, seed=5167):
    """Train a cross-validated logistic regression probe at a layer.

    Parameters
    ----------
    X : np.ndarray
        Hidden states for the layer, shape (n_story, d_model)
    y : np.ndarray
        Target labels, shape (n_story,)
    n_splits : int
        Number of folds for cross validation
    seed : int
        Random seed

    Returns
    -------
    float
        Mean accuracy over the folds
    """
    pipe = make_pipeline(
        StandardScaler(),
        LogisticRegression(C=1, max_iter=2000),
    )

    cv = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
    scores = cross_val_score(pipe, X, y, cv=cv, scoring="accuracy", n_jobs=-1)

    return scores.mean()

With our function defined, we can train some probes. Our main set tries to recover genre from the hidden states. We’ll also train three other sets as points of comparison:

Genre: Can a probe predict the correct genre from hidden states?
Correct letter: Can a probe predict which answer letter represents that genre?
Predicted letter: Can a probe predict which letter Olmo 3 chooses?
Shuffled letter: A control. We shuffle the letter labels so they no longer correspond to the hidden states, then train a probe on this scrambled mapping. Any accuracy it reaches is what a probe can pick up by chance, which gives us a floor to judge the other three against.

Below, we convert our DataFrame columns into NumPy arrays:

genre = df["true_genre"].to_numpy()
correct_letter = df["correct_letter"].to_numpy()
predicted_letter = df["predicted_letter"].to_numpy()

We also randomly shuffle correct_letter to obtain our control probes.

rng = np.random.default_rng(5167)
shuffled_letter = rng.permutation(correct_letter)

Time to train. For each layer, we extract hidden states and train our four probes:

n_layer = hidden_states.shape[1]

all_results = []
for layer in range(n_layer):
    X = hidden_states[:, layer, :].astype(np.float64)
    results = {
        "layer": layer,
        "genre": probe_layer(X, genre),
        "correct_letter": probe_layer(X, correct_letter),
        "predicted_letter": probe_layer(X, predicted_letter),
        "shuffled": probe_layer(X, shuffled_letter),
    }

    all_results.append(results)

Let’s package our results up into a DataFrame and format them for plotting.

probe_df = pd.DataFrame(all_results)
probe_plot = probe_df.melt(
    id_vars="layer",
    value_vars=["genre", "correct_letter", "predicted_letter", "shuffled"],
    var_name="Probe target",
    value_name="accuracy",
)

Probe performance¶

How did the probes perform? Here’s a high-level overview:

probe_plot.groupby("Probe target")["accuracy"].describe()

Some of these probes performed quite well. Let’s look across the layers:

g = sns.relplot(
    data=probe_plot,
    kind="line",
    x="layer",
    y="accuracy",
    hue="Probe target",
    marker="o",
    height=FACET_HEIGHT,
    aspect=FACET_ASPECT,
)
g.refline(y=1/5, color="black", linestyle="--")
g.set_axis_labels("Layer", "Probe accuracy")
g.fig.suptitle("Linear Probe Accuracy by Layer", y=1.05)
plt.show()

Here’s a summary:

Genre: decodable from layer 0 (the initial embedding layer, surprisingly!), climbing to 95% accuracy by the end. The model effectively encodes a genre signal in the embedding itself and refines that signal from layer to layer.
Correct letter: tracks chance until layer 10, climbs in the middle layers, and plateaus around layer 16. It never gets better than 50% accuracy. Our previous classification setup performs quite poorly.
Predicted letter: climbs to ~70% accuracy around layer 22---right when the letter spread occurs. But this number never drops, which is additional evidence that the model is biased toward some letters over others.
Shuffled letter: sits near chance, as expected.

Notably, our probes have an easier time recovering Olmo 3’s predicted letter than the correct one. This suggests that the hidden states also encode the model’s answer letter bias, just as the next-token logits did in the last chapter. That makes sense: These are the last-token states from our prompt to the instruct-tuned model, so they summarize everything the model was about to act on---genre signal and letter bias alike.

Supervised Classification¶

That said, probes can readily decode genre from these hidden states. Can we keep that decoding power while avoiding the answer-letter bias? The trick is to stop asking the model a multiple-choice question at all. Instead, we attach a new classification layer on top of Olmo 3---five outputs, one per genre---and train just that layer on our stories, with the base model frozen. Because we never ask for a letter, there’s no letter for the model to be biased about: The genre labels are the only target. Those labels, in turn, will allow us to re-run our token attribution analysis in a way that sidesteps the model’s bias.

Fitting the classifier¶

Because our classifier requires information from Olmo 3 as the latter runs, we can’t fit it in a notebook. Instead, we use this script. It:

Collects each story’s per-layer hidden states and mean-pools them over the story span to produce a single representation per layer
Trains a cross-validated linear head at each layer
Selects the layer where genre is most decodable
Refits the head on a train split at that layer
Saves the head weights, chosen layer, classes, pooled hidden states, and the train/test split

As we did with the hidden state cache, we load our classifier like so:

clf_cache = np.load("../data/w2-d3_linear-classifier.npz")

Classifier performance¶

Our setup takes the argmax over probes to select which layer will serve as input for our classifier. Which layer is that?

layer = clf_cache["chosen_layer"].item()
print("Selected Transformer block:", layer)

How well does it perform on the held-out test set (that is, stories it never trained on)?

acc = clf_cache["held_out_accuracy"].item()
print("Held-out accuracy:", round(acc, 3))

97% is quite good! That said, many of the probes perform well on this task, and several would likely be viable classifiers for this genre task. Here’s a plot of all probes on the training set:

cv_acc = clf_cache["cv_accuracy_by_layer"]

fig, ax = plt.subplots()
sns.lineplot(x=range(len(cv_acc)), y=cv_acc, marker="o", ax=ax)
ax.set(
    title="Layer-wise Probe Accuracy on Training Set",
    xlabel="Layer",
    ylabel="Accuracy",
    ylim=(0.85, 1.0),
)
plt.show()

Clearly the classifier performs well. But let’s dig in a bit. If we run cosine similarity over our classifier weights, we’ll get a rough sense of whether the classifier discriminates between genres. Even though the high accuracy already suggests it has no trouble doing this, we can cross-check that performance by looking at the classifier’s representations for each genre.

First, we extract the classifier weights and compute cosine similarity:

W = clf_cache["head_weight"]
Wn = W / np.linalg.norm(W, axis=1, keepdims=True)
cos_sim = Wn @ Wn.T

We format the resultant matrix as a DataFrame:

genres = clf_cache["genre_classes"].astype(str)
cos_sim = pd.DataFrame(cos_sim, index=genres, columns=genres)

And plot:

fig, ax = plt.subplots()
sns.heatmap(cos_sim, cmap="RdBu_r", annot=True, fmt=".2f", ax=ax)
ax.set(title="Cosine Similarity Between Genre Weight Directions")
plt.show()

All looks good here. Every genre-to-genre cosine similarity score is negative except the identity comparisons. That means the classifier weights push away from each other, which in turn enables this model to discriminate between genres. Note, too, that one set of our scores is notably more dissimilar than the others: those between detective fiction and thriller. Why is that? Well, remember: Our earlier classifiers struggled to differentiate these two genres. This model performs much better, and it does so by strongly pushing detective fiction and thriller away from each other in its weight space.

Making predictions on a test set¶

Since our classifier is just a set of weights plus a bias vector, we can recompute its predictions given pooled hidden states. Let’s do that on our test set to evaluate its performance. First, we get the bias vector from the cache as well as the pooled hidden states for all layers in Olmo 3:

b = clf_cache["head_bias"]
head_pooled_states = clf_cache["pooled_states"]

Now we index our pooled hidden states by the layer the classifier trained on. We also get the test indices from our cache.

X_layer = head_pooled_states[:, layer, :] 
test_idx = clf_cache["test_idx"]

With all that information extracted, we convert the pooled hidden states in X_layer into logits and use argmax to get the highest one for each story in the test set. We store those predictions in y_pred.

GENRES = np.array([
    "Horror",
    "Detective fiction",
    "Romance",
    "Science fiction",
    "Thriller",
])

test_logits = X_layer[test_idx] @ W.T + b
y_pred = GENRES[test_logits.argmax(axis=1)]

Finally, with .loc[], we can index our DataFrame using our test indices and extract the true genre label. A cross tabulation between those and y_pred will produce a confusion matrix:

y_true = df.loc[test_idx, "true_genre"]

confusion_matrix = (
    pd.crosstab(y_true, y_pred).reindex(columns=GENRES, index=GENRES)
)

...which we 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 Genre",
    xlabel="Predicted genre",
    ylabel="Correct genre",
)
plt.show()

Almost perfect. The classifier can easily discriminate science fiction and romance from the other genres, and it manages extremely well with horror, detective fiction, and thriller. Only four stories in these genres are mis-classified, and the spread of errors is even across them---no overloading of thriller and detective fiction here.

Token Attribution (Redux)¶

Now that we have our fitted classifier and determined that it can effectively represent genre in the model, we can again perform token attribution on the stories. As in the last chapter, our backpropagation objective runs on the difference between the correct genre and the LogSumExp of the other genres. But this time it doesn’t use answer letter logits, it uses the genre labels from our classifier. And, with the classifier’s performance being so high, we can be more confident that these token attributes track genre information in the model’s hidden states, rather than artifacts of the answer-letter channel.

For our purposes, that attribution work happens off screen: no gradient DataFrame this time. Instead, the one we load below contains genre lift scores from our classifier-driven token attributions. It was produced using the same aggregation logic as the most distinctive word analysis we did in the last chapter.

clf_lift = pd.read_parquet("../data/w2-d3_classifier-genre-lift.parquet")

Note, though, that we’ve expanded how many distinctive words we derive: 250 per genre.

genre_counts = clf_lift.groupby("genre").size().to_frame(name="count")
genre_counts

Let’s collect the top-25 most distinctive words for each genre:

k = 25
most_distinctive = (
    clf_lift
    .sort_values(["genre", "genre_lift"], ascending=[True, False])
    .groupby("genre", as_index=False)
    .head(k)[["genre", "word", "genre_lift"]]
    .reset_index(drop=True)
)

And plot. This recreates the exact same bar chart from the previous chapter, except this time we use a properly-fitted classifier head as our means of measuring distinctive words.

g = sns.catplot(
    data=most_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()

Comparing Measurement Channels¶

To complete this case study on distinctive tokens and genre, we’ll reflect a bit on the nature of our measurements. Let’s load another DataFrame, which contains genre lift for the multiple-choice question/answer setup:

mcqa_lift = pd.read_parquet("../data/w2-d3_mcqa-genre-lift.parquet")

Distinctive word overlaps¶

With both word list DataFrames loaded, we compute the overlap between the top-50 distinctive words under each measurement channel. We do so with Jaccard similarity, which measures the size of the intersection between two sets divided by the size of their union:

J(A, B) = \frac{|A \cap B |}{|A \cup B|}

(2)

A score of 1 means the two word lists are identical. A score of 0 means they share no words.

def jaccard(a, b):
    """Compute Jaccard similarity between two word lists.

    Parameters
    ----------
    a : set[str]
        Set of words
    b : ste[str]
        Set of words

    Returns
    -------
    float
        Jaccard similarity
    """
    union = a | b
    if len(union) == 0:
        return 0.0

    return len(a & b) / len(union)

Now we compute this score for each genre. For every genre, we take the top-50 words from our two channels, convert each list to a set, and measure their overlap.

Note

A set is a data structure that contains unique values. In Python, you can convert lists to sets with the set() function:

lst = [0, 0, 1, 2, 3]
set(lst)

Or, enter values in between curly brackets { }:

{0, 1, 2, 3}

k = 50
overlap = []

for genre in GENRES:
    m_words = (
        mcqa_lift[mcqa_lift["genre"] == genre]
        .sort_values("genre_lift", ascending=False)
        .head(k)
    )
    c_words = (
        clf_lift[clf_lift["genre"] == genre]
        .sort_values("genre_lift", ascending=False)
        .head(k)
    )

    m_words = set(m_words["word"])
    c_words = set(c_words["word"])
    
    shared = len(m_words & c_words)
    score = jaccard(m_words, c_words)

    overlap.append(
        {
            "genre": genre,
            "jaccard": score,
            "shared": len(m_words & c_words),
            "shared_share": shared / k,
        }
    )

overlap_df = pd.DataFrame(overlap)
overlap_df = overlap_df.sort_values("jaccard")

Let’s plot the overlaps:

fig, ax = plt.subplots()
sns.barplot(overlap_df, x="genre", y="jaccard", ax=ax)
ax.set(
    title="Overlap Between Distinctive Words by Measurement Channel",
    xlabel="Genre",
    ylabel="Jaccard similarity",
    ylim=(0, 1),
)

for container in ax.containers:
    labels = []
    for shared in overlap_df["shared"]:
        label = f"{shared}/{k} shared"
        labels.append(label)

    ax.bar_label(
        container, labels=labels, padding=3
    )

plt.xticks(rotation=90, ha="right")
plt.show()

On average, the two measurement channels share only a minority of their top-50 distinctive words:

mean_jaccard = overlap_df["jaccard"].mean()
mean_shared = overlap_df["shared"].mean()
mean_shared_share = overlap_df["shared_share"].mean()

print("Mean Jaccard similarity:", round(mean_jaccard, 3))
print("Mean shared words:", round(mean_shared, 1), "out of", k)
print("Mean shared-word share:", round(mean_shared_share, 3))

This is a modest overlap. The two channels aren’t simply producing the same word lists in a slightly different order. They disagree about which words count as the most distinctive genre cues.

Rank changes¶

To see what words changed, we rank words within each genre and compute their rank change.

mcqa_lift["rank_mcqa"] = mcqa_lift.groupby("genre").cumcount() + 1
clf_lift["rank_clf"] = clf_lift.groupby("genre").cumcount() + 1

merged = mcqa_lift.merge(
    clf_lift[["genre", "word", "rank_clf", "genre_lift"]],
    on=["genre", "word"],
    how="inner",
    suffixes=["_mcqa", "_clf"],
)
merged["rank_change"] = merged["rank_mcqa"] - merged["rank_clf"]

Now, we categorize words by how much their rank shifted:

"fallen": High multiple-choice question/answer rank but low classifier rank
"risen": Low multiple-choice question/answer rank but high classifier rank
"stable": No rank change

n = 100
merged["category"] = "stable"
merged.loc[merged["rank_change"] < -n, "category"] = "fallen"
merged.loc[merged["rank_change"] > n, "category"] = "risen"

Let’s count our categories:

merged["category"].value_counts()

Finally, we summarize the lift values across these categories:

summary = (
    merged
    .groupby("category")[["genre_lift_mcqa", "genre_lift_clf"]]
    .agg(["mean", "median"])
    .round(4)
)
summary.columns = ["mcqa_mean", "mcqa_median", "clf_mean", "clf_median"]

And inspect the results:

summary

Fallen words had high lift under the multiple-choice setup but near-zero classifier lift. These words looked important under the noisy letter-logit data, but their signal collapses under the genre head.
Risen words had low lift but moderate classifier lift. These were genre clues that got lost in the noise of our answer letters.

Let’s see examples of each. Here’s "fallen":

subset = ["genre", "word", "rank_mcqa", "rank_clf", "rank_change"]

fallen = (
    merged[merged["category"] == "fallen"]
    .sort_values("rank_change")
    .head(10)
)

fallen[subset]

See the discrepancies in rank between the two classification setups? That’s an indication that these words are artifacts of how we measured genre. They appeared distinctive under the multiple-choice question/answer channel but carry almost no genre signal according to the classifier.

The opposite is the case for "risen":

risen = (
    merged[merged["category"] == "risen"]
    .sort_values("rank_change", ascending=False)
    .head(10)
)

risen[subset]

These words were suppressed by the noisy letter logit-data but emerge as genre cues once we eliminate that noise.

The important point is that the attribution method didn’t change very much. Our stories were the same. The model was the same. The aggregation procedure was the same. Only the channel through which we measured genre changed, and it turns out this difference matters quite a bit.

More, this difference suggests that the means of measurement is itself part of the measurement. Token attributions aren’t simply properties of tokens. They are the product of a particular objective, routed through a particular measurement channel. In the multiple-choice setup, that channel inadvertently mixed genre signal with answer-letter bias. In the classifier setup, the channel sidestepped that bias.

So any explanation we build from attribution scores has to be read in light of the instrument that produced them. Interpretability doesn’t just ask what the model represents. It also reflects the methods by which we attempt to find those representations. Indeed, those representations are themselves the product of our methods.