(Model) Text as Data - Introduction to Interpretability for Language Models

from string import Template

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.optimize import minimize
from scipy.special import logsumexp, rel_entr
from scipy.stats import entropy
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import cross_val_score
from sklearn.naive_bayes import MultinomialNB

The next three chapters use a dataset of LLM-generated short stories to walk through several interpretability techniques. In this chapter, we’ll perform a text classification task on the stories and examine how a fitted classifier can reveal model behavior. Later chapters will use story texts as an inroad to model internals.

Making a Dataset¶

The stories themselves have been generated with the instruct-tuned variant of Olmo 3 7B. We use this system prompt:

SYSTEM_PROMPT = (
    "You are a creative-writing assistant. "
    "Write only the requested story text. Never add commentary."
)

...in combination with this prompt template:

PROMPT = Template(
    "Write a complete$genre short story of no more than 1,000 words. "
    "The story must have a clear ending. "
    "Do not include a title, preface, notes, or commentary.\n\nStory:\n"
)

Then, for five different genres:

GENRES = {
    "Horror": "horror",
    "Detective fiction": "detective fiction",
    "Romance": "romance",
    "Science fiction": "science fiction",
    "Thriller": "thriller",
}

...we interpolate a genre key into the prompt, e.g.:

genre = GENRES["Horror"]
prompt = PROMPT.substitute(genre=f" {genre}" if genre else "")
print(prompt)

...and send this to the model.

We also add a NULL genre, which prompts a model for a “plain” short story (whatever that means):

genre = None
prompt = PROMPT.substitute(genre=f" {genre}" if genre else "")
print(prompt)

We do this 1,000 times for each genre to get 6,000 total short stories.

Loading data¶

Stories are saved in a DataFrame:

df = pd.read_parquet("../data/w2-d1_genre-stories.parquet")

Let’s use the .sample() method to get a random short story:

samp = df.sample()
print("Genre:", samp["genre"].item())
print("Story:", samp["text"].item()[:500])

What’s the genre distribution in our dataset?

fig, ax = plt.subplots()
sns.countplot(data=df, x="genre", ax=ax)
ax.set(title="Genre Distribution", xlabel="Genre", ylabel="Count")
plt.show()

Producing count data¶

Currently, the stories in our dataset are strings---perfectly readable for us, but not so amenable to computational analysis. The latter requires some kind of count data. Below, we use a CountVectorizer to tabulate the number of times each token appears in a story.

Here, “token” is defined a bit differently than in past chapters. We won’t use the LLM tokenizer to segment our text. Instead, we’ll rely on a simple pattern: for the CountVectorizer, a token is a sequence of three or more alphabetic characters. To that constraint, we’ll add a few more:

token_pattern=...: A “token” is a sequence of three or more alphabetic characters
stop_words="english": Tokens are filtered against a list of common English words like “the,” “and,” “of,” etc.
max_df=0.9: If a token appears in more than 90% of all documents, we remove it
min_df=10: If a token appears in less than 10 documents, we remove it

vectorizer = CountVectorizer(
    token_pattern=r"(?u)\b[a-zA-Z]{3,}\b",
    stop_words="english",
    max_df=0.9,
    min_df=10,
)

With our CountVectorizer initialized, we fit it and transform our stories. The result is a document-term matrix (DTM). Its columns represent unique tokens (or “types”), while its rows correspond to each story in our dataset. For each row, it records how many times a token appears---or doesn’t appear. Importantly, the DTM tracks zero-count tokens as well.

X = vectorizer.fit_transform(df["text"].tolist())
print("Shape (n_doc, vocab_size):", X.shape)

Finally, we extract the feature names---in our case, the tokens---from the vectorizer. These are aligned to the column order of X. We also convert our genre column to an array, y.

feature_names = vectorizer.get_feature_names_out()
y = df["genre"].to_numpy()

Fitting a Classifier¶

With our texts transformed, we’ll now see if we can predict a story’s genre from its token distribution. We do so with a Multinomial Naive Bayes model. This models a story as a vector of token counts $\mathbf{x} = (x_1, x_2, \ldots, x_n)$ and estimates the probability of each genre given those counts.

Token probabilities are estimated from the training data with additive smoothing:

P(t \mid g) = \frac{N_{gt} + \alpha}{N_g + \alpha n}

(1)

Where:

$N_{gt}$ is the count of token $t$ in genre $g$
$N_g$ is the total token count in genre $g$
$n$ is the vocabulary size
$\alpha$ is the smoothing parameter. Smoothing prevents zero probabilities for unseen tokens

Classification uses the posterior score:

\hat{g} = \arg\max_{g} \left[ \log P(g) + \sum_{t=1}^{n} x_t \log P(t \mid g) \right]

(2)

That is, for each genre, we add its log prior probability to the log-probabilities of the story’s tokens, weighted by their counts. Tokens absent from the story count zero and don’t affect the score. The genre with the highest total score is selected.

Scikit-learn’s MultinomialNB() handles all this math for us. We just need to initialize it and fit it on our DTM and the stories’ labels.

clf = MultinomialNB()
clf.fit(X, y)

How well did our classifier perform?

acc = clf.score(X, y)
print("Mean accuracy:", round(acc, 3))

Pretty good. But this is the score on the same data we trained on, so its performance may be optimistically skewed. We should cross validate the model. With cv=5, this step splits the dataset into 5 roughly equal folds. Then, it loops 5 times (once per fold). In each iteration, it:

Picks a fold as the validation set. The model won’t train on this data
Uses the other 4 folds as the training set. The model trains on this data
Fits a new MultinomialNB() on the training set
Scores this model on the validation set

Every data point ends up in the validation set exactly once.

cross_val_score() returns 5 accuracy scores: one for each model it trained on the data subsets. Taking the mean of these scores provides a best estimate of the model’s true performance.

cv = 5
cv_scores = cross_val_score(MultinomialNB(), X, y, cv=cv, scoring="accuracy")
print("CV accuracy for", cv, "folds:", round(cv_scores.mean(), 3))

Still quite solid. A last check: Let’s quickly compute the proportion of the most common genre in our dataset (in this case, all genres are the same proportion). This is the accuracy we’d get from a dummy classifier that always predicts the majority class, regardless of input. We require our model to beat the baseline by at least 5 percentage points.

baseline = df["genre"].value_counts(normalize=True).max()
if cv_scores.mean() < baseline + 0.05:
    print("WARNING: accuracy ~ baseline. Genres are likely not separable")
else:
    print("OK: model can separate genres")

All is well. We can be sure that the model has learned to separate genres.

Distinctive Tokens¶

Now that we know our model is well-fitted to the data, we can look at how genres and tokens interact. The question we ask is this: What tokens are distinctive to a genre?

Pointwise mutual information¶

To answer this question, we use pointwise mutual information (PMI) between a token and a genre:

\operatorname{PMI}(g, t) = \log \frac{p(g, t)}{p(g)\, p(t)} = \log \frac{p(t \mid g)}{p(t)}

(3)

PMI measures how much more or less likely a token $t$ is under genre $g$ compared to its overall probability under the model:

Positive values: Token $t$ is associated with genre $g$
Negative values: Token $t$ is less associated with genre $g$ than expected

To calculate PMI, we’ll first get the classes from the classifier.

classes = clf.classes_
classes

We get the log probability of token $t$ given genre $g$ , and the log prior probability of each genre.

log_p_t_given_g = clf.feature_log_prob_
log_p_g = clf.class_log_prior_

Log probabilities may be converted to probabilities via $e^x$ , where $x$ is the log probability. Applying this to $\log p(g)$ gives us the marginal genre distribution according to the classifier. If our training data is balanced across six genres, each prior is ~0.1667.

np.exp(log_p_g)

Next, calculate the marginal log probability of token $t$ by marginalizing over genres:

p(t) = \sum_g p(g)\, p(t \mid g)

(4)

In log space, this is:

log_p_t = logsumexp(log_p_g[:, None] + log_p_t_given_g, axis=0)

With $p(t)$ , we can compute PMI:

pmi = log_p_t_given_g - log_p_t

Rare tokens can produce extreme values with PMI, so we will down-weight rare $(g, t)$ pairs using the joint probability

p(g, t) = p(g)\, p(t \mid g)

(5)

This gives us a weighted PMI score.

log_p_joint = log_p_g[:, None] + log_p_t_given_g
p_joint = np.exp(log_p_joint)
wpmi = p_joint * pmi

Tokens with the highest wpmi for a given genre are both distinctive of that genre and frequent enough to matter.

A last step before looking into the details of these scores: Let’s convert both pmi and wpmi into DataFrames for ease of use.

pmi_df  = pd.DataFrame(pmi,  index=classes, columns=feature_names)
wpmi_df = pd.DataFrame(wpmi, index=classes, columns=feature_names)

Ranking tokens¶

Both of these DataFrames are quite wide:

print("Shape of DataFrame:", pmi_df.shape)

Below, we write a function that returns the top k tokens for a genre. This saves us the trouble of mucking through our DataFrames. The function uses our weighted PMI scores for top-token selection, but it also uses the regular PMI scores to calculate lift.

Lift is the exponentiated PMI score:

\exp(\operatorname{PMI}(g,t)) = \frac{p(t \mid g)}{p(t)} = \frac{p(g \mid t)}{p(g)}

(6)

So a lift of 3, for example, means that token $t$ is three times more likely under genre $g$ than it is overall. Equivalently, observing token $t$ multiplies the prior probability of genre $g$ by three.

def top_tokens(wpmi_df, pmi_df, genre, k=25):
    """Return the ``k`` most distinctive tokens for a genre.

    Parameters
    ----------
    wpmi_df : pd.DataFrame
        Weighted PMI scores, shape (n_genre, vocab_size)
    pmi_df : pd.DataFrame
        PMI scores, shape (n_genre, vocab_size)
    genre : str
        Genre to rank tokens for
    k : int
        Number of tokens to return

    Returns
    -------
    pd.DataFrame
        Top ``k`` tokens for ``genre``, with their wPMI and lift scores
    """
    # Get the k tokens with the highest wPMI for this genre
    top = wpmi_df.loc[genre].nlargest(k)

    # Calculate lift for these tokens
    lift = np.exp(pmi_df.loc[genre, top.index])

    # Package up into a DataFrame
    df = pd.DataFrame({"wPMI": top, "lift": lift}).reset_index(names="token")

    return df

Now, for each genre, we run top_tokens() and append the resultant DataFrame to results.

results = []
for genre in classes:
    top = top_tokens(wpmi_df, pmi_df, genre)
    top["genre"] = genre 
    results.append(top)

We concatenate our six DataFrames into a single one:

top_plot = pd.concat(results, ignore_index=True)

And plot.

g = sns.catplot(
    data=top_plot,
    x="wPMI",
    y="token",
    col="genre",
    col_wrap=2,
    kind="bar",
    sharey=False,
    height=FACET_HEIGHT,
    aspect=FACET_ASPECT,
)
g.set_axis_labels("weighted PMI", "")
g.set_titles("{col_name}")

for genre, ax in g.axes_dict.items():
    subset = top_plot[top_plot["genre"] == genre]

    for patch, (_, row) in zip(ax.patches, subset.iterrows()):
        x_pos = patch.get_width()
        y_pos = patch.get_y() + patch.get_height() / 2

        ax.text(
            x_pos,
            y_pos,
            f"  {row['lift']:.1f}x",
            va="center",
            ha="left",
        )

    ax.set_xlim(right=ax.get_xlim()[1] * 1.15)

g.fig.suptitle(
    "Distinctive Tokens by Genre: wPMI with Lift Labels", y=1.05
)
plt.show()

Token Distributions¶

Thus far we have looked at single tokens’ relationships with genres. But our classifier worked at the distributional level: Each genre induces a full probability distribution over the token vocabulary. We can recover those distributions like so:

probs = np.exp(log_p_t_given_g)
P = pd.DataFrame(probs, index=classes, columns=feature_names)

With these in hand, we can ask distributional questions---not just how a single token relates to a genre, but how whole genres relate to one another. In this section, we’ll use these distributions to investigate the NULL genre. Our question: How similar are unmarked stories (those generated without an explicit genre prompt) to the explicit genre stories in our dataset?

KL divergence¶

First, we separate the explicit genres from NULL using .loc to index the DataFrame P by genre:

explicit_genres = []
for genre in classes:
    if genre == "NULL":
        continue

    explicit_genres.append(genre)

Q = P.loc[explicit_genres].to_numpy()
p_null = P.loc["NULL"].to_numpy()

Differences in probabilities are often measured by KL divergence (KLD). KLD measures how much an approximating probability distribution $Q$ diverges from a target probability distribution $P$ .

For example:

D_{KL}(P_{\text{null}} \Vert Q_{\text{sf}})

(7)

Where:

$P_{\text{null}}$ is the token distribution for the NULL genre, treated as the target distribution
$Q_{\text{sf}}$ is the distribution for science fiction, treated as the approximating distribution

The formula is:

D_{KL}(P \Vert Q) = \sum_{t=1}^{n} P(t) \log \frac{P(t)}{Q(t)}

(8)

A smaller KL divergence means that genre $g$ ’s token distribution is more similar to the NULL distribution. A larger value means the genre distribution is a poorer approximation of NULL.

kld = np.apply_along_axis(lambda q: entropy(P.loc["NULL"], q), axis=1, arr=Q)

Let’s gather these scores into a DataFrame and sort them with .sort_values():

kld_plot = pd.DataFrame({"KLD": kld, "genre": explicit_genres})
kld_plot.sort_values("KLD", ascending=True, inplace=True)

And plot:

fig, ax = plt.subplots()
sns.barplot(data=kld_plot, x="KLD", y="genre", ax=ax)
ax.set(title="KL Divergence", xlabel="KL divergence (nats)", ylabel="Genre")

for patch, label in zip(ax.containers[0].patches, kld_plot["KLD"]):
    x_pos = patch.get_width()
    y_pos = patch.get_y() + patch.get_height() / 2
    ax.text(
        x_pos,
        y_pos,
        f"  {label:.3f}",
        va="center",
        ha="left",
     )

ax.set_xlim(right=ax.get_xlim()[1] * 1.15)
plt.show()

Again, a lower KLD score indicates that genre $g$ is a better approximator for NULL than genres with higher scores. Below, we show the underlying distributions that KLD measures. When KLD is large, it means a genre’s curve places mass differently from NULL’s across the token vocabulary.

This requires a bit of DataFrame manipulation:

Name the row index and column index
Convert from wide form to long form: one row per (genre, token) pair
Rename the stacked probability values
Convert the index levels into regular columns

prob_plot = (
    P
    .rename_axis(index="genre", columns="token")
    .stack()
    .rename("probability")
    .reset_index()
)

Ready to plot:

fig, ax = plt.subplots()
sns.kdeplot(
    data=prob_plot,
    x="probability",
    hue="genre",
    multiple="layer",
    log_scale=True,
    ax=ax,
)
ax.set(title="Token Log Probabilities Per Genre", xlabel="Log probability")
plt.show()

Mixing genres¶

While the distribution for romance is most like that of NULL’s, the other two distributions are still fairly similar to our unmarked genre. Perhaps NULL is actually a mix of genres. That would make sense: When an LLM is prompted to write a story without genre directions, it probably pulls on a variety of different genres to complete a prompt. But can we measure this? How much might NULL mix genres?

We can use KLD to answer these questions. We’ll do this by trying to minimize the KLD between NULL and a mixture $Q_w$ of the explicit genres:

D_{KL}(P_{\text{null}} \Vert Q_w), \qquad Q_w = \sum_g w_g Q_g, \qquad w_g \ge 0, \qquad \sum_g w_g = 1

(9)

Where:

$Q_g$ is the token distribution for genre $g$
$w_g$ is that genre’s mixture weight

One might expect genres closer to NULL individually, that is, genres with lower KLD, to receive more weight. But mixture weights depend on marginal contribution: A genre receives weight only if it improves the approximation given the other genres already in the mixture. Thus, to find the precise proportions of genres, we need to define a metric that can indicate what should be added/subtracted to get our final mix.

def mixture_kl(w):
    """Return KL divergence between NULL and a weighted genre mixture.

    Parameters
    ----------
    w : np.ndarray
        Mixture weights for the explicit genre distributions, shape (n_genre,)

    Returns
    -------
    float
        The KL divergence
    """
    # Construct the mixture distribution: q[t] = sum_g w[g] * Q[g, t]
    q = w @ Q

    # Avoid log(0) if any token receives zero probability and re-normalize so q
    # remains a probability distribution
    q = np.clip(q, 1e-300, None)
    q = q / q.sum()

    return entropy(p_null, q)

With our metric defined, we set up an optimizer that adjusts the genre weights in the mixture. The optimizer searches for the combination of genre distributions that makes $Q_w$ as close as possible to $P_{\text{null}}$ .

This optimizer requires a few additional parameters:

x0: initial guess for the genre mixture. We start with a uniform mixture
constraints: the mixture weights must sum to 1
bounds: each mixture weight must fall between 0 and 1

x0 = np.ones(len(explicit_genres)) / len(explicit_genres)
constraints = {"type": "eq", "fun": lambda w: w.sum() - 1}
bounds = [(0, 1)] * len(explicit_genres)

We now run our optimizer.

res = minimize(
    mixture_kl, x0, method="SLSQP", bounds=bounds, constraints=constraints
)

if not res.success:
    print("WARNING:", res.message)
else:
    print("OK: optimizer found a solution")

We extract the mixture weights from our result like so:

weighting = res.x / res.x.sum()
mix_plot = pd.DataFrame({"mixture": weighting, "genre": explicit_genres})
mix_plot.sort_values("mixture", ascending=False, inplace=True)

Let’s plot our mixture weights:

fig, ax = plt.subplots()
sns.barplot(data=mix_plot, x="mixture", y="genre", ax=ax)
ax.set(title="Genre mixture for NULL", xlabel="Mixture weight", ylabel="Genre")

for patch, label in zip(ax.containers[0].patches, mix_plot["mixture"]):
    x_pos = patch.get_width()
    y_pos = patch.get_y() + patch.get_height() / 2
    ax.text(
        x_pos,
        y_pos,
        f"  {label:.1%}",
        va="center",
        ha="left",
     )

ax.set_xlim(right=ax.get_xlim()[1] * 1.15)
plt.show()

Translated: Best-fitting approximation to NULL is roughly 60% romance, 20% horror, 15% detective fiction, and 5% science fiction/thriller. So, while the model’s default short story mode is close to romance, it isn’t solely romance. A mix better captures the overall distribution for NULL.

What’s with thriller?

Thriller’s KLD with NULL is respectable---better than science fiction’s, in fact. So why did our optimizer assign it almost no weight in the genre mix?

The key is that a mixture weights measure each genre’s marginal contribution given the others, not its raw similarity to NULL. A genre that overlaps heavily with genres already in the mix gets crowded out, even if it’s individually a decent match. One hypothesis: thriller is itself a composite of horror and detective fiction, both of which are already well-represented. If that’s the case, the optimizer can get thriller-flavored mass “for free” from those other genres without spending weight on thriller directly.

You could test this hypothesis by re-running the analysis with thriller as the target instead of NULL. If thriller is well-approximated by a mixture of horror and detective fiction, that would support the idea that it doesn’t need to contribute very much to our current mix.

How much were we able to improve our approximation over romance alone?

best_single = kld_plot.nsmallest(1, "KLD")
improvement = (best_single["KLD"] - res.fun) / best_single["KLD"]
print(f"Mixture improves over romance by {improvement.item():.1%}")

Token KL contributions¶

As with weighted PMI, we can identify tokens that are distinctive in NULL vis-a-vis the genre mixture. To do so, we re-weight token distributions by the mixture:

p_mix = weighting @ Q
residuals = pd.DataFrame({
    "token": feature_names, "p_null": p_null, "p_mix": p_mix
})

For each token, we calculate its signed token-level KL term:

P_{\text{null}}(t) \log \frac{P_{\text{null}}(t)}{Q_w(t)}

(10)

These terms sum to the total KL divergence.

residuals["kl_contrib"] = rel_entr(residuals["p_null"], residuals["p_mix"])

Positive values identify tokens that are more probable in NULL than in the fitted mixture. They indicate what’s distinctive about NULL beyond what the genre mixture can explain. Think of them like the LLM’s default story mode, its signature.

Negative values identify tokens that the fitted mixture over-predicts relative to NULL. These are the genre-typical words NULL avoids: baggage the optimizer had to carry because up-weighting a genre to capture the positive tokens inevitably introduced vocabulary that NULL itself doesn’t favor.

Let’s get the top tokens from each:

k = 25

under = (
    residuals
    .nlargest(k, "kl_contrib")
    .sort_values("kl_contrib", ascending=False)
    .assign(direction="Under-predicted")
)

over = (
    residuals
    .nsmallest(k, "kl_contrib")
    .sort_values("kl_contrib", ascending=False)
    .assign(direction="Over-predicted")
)

kl_plot = pd.concat([under, over], ignore_index=True)

And plot:

direction_order = ["Under-predicted", "Over-predicted"]

g = sns.relplot(
    data=kl_plot,
    x="kl_contrib",
    y="token",
    col="direction",
    col_order=direction_order,
    col_wrap=1,
    kind="scatter",
    hue="direction",
    hue_order=direction_order,
    facet_kws={"sharey": False, "sharex": True},
    height=FACET_HEIGHT,
    aspect=FACET_ASPECT,
    s=60,
    legend=False,
)
g.set_axis_labels("Signed token-level KL term", "")
g.set_titles("{col_name}")

limit = kl_plot["kl_contrib"].abs().max()
for ax in g.axes.flat:
    ax.axvline(0, color="0.3", linewidth=1)
    ax.set_xlim(-limit * 1.1, limit * 1.1)

g.fig.suptitle(
    "Tokens Under- and Over-predicted by the Fitted Mixture",
    y=1.05,
)
plt.show()