Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

Chain-of-Thought Reasoning

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch
import torch.nn.functional as F
from scipy.signal import find_peaks
from scipy.stats import median_abs_deviation
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from transformers import AutoModelForCausalLM, AutoTokenizer

In this final chapter, we investigate chain-of-thought reasoning. Our question is this: If “reasoning” models are trained to generate step-by-step sequences to complete benchmarks, can we isolate a model’s individual reasoning steps?

Below, we load a small reasoning model, Qwen 3 1.7B.

checkpoint = "Qwen/Qwen3-1.7B"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map="auto", dtype="auto"
)

As before, we set a seed, turn off gradient accumulation, and put the model in evaluation mode.

torch.manual_seed(5167)
torch.set_grad_enabled(False)
model.eval()

Generating a Reasoning Trace

Ultimately, a reasoning model generates text just like any language model. So, to get a reasoning output, we need a prompt. We’ll tokenize it and send it to the model much as we’ve done before.

prompt = (
    "A store has 40 apples. In the morning it sells a third of them. "
    "In the afternoon a delivery adds 18 more apples, and then 25% of the "
    "apples in the store are sold. How many apples are left at closing? "
    "Reason briefly, then give the answer."
)

Tokenization and generation

Reasoning models are a variant of other instruct-tuned models, so they require a chat template. Below, we package our prompt as a user request and apply a chat template. Note enable_thinking=True: This adds a special “open think” token, which Qwen-style models use to do their reasoning.

messages = [{"role": "user", "content": prompt}]
chat = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=True,
)

Now, we tokenize our text and send it to the model using the recommended sampling parameters for generation. We set output_hidden_states=True to get the hidden states for each layer in the model.

inputs = tokenizer(chat, return_tensors="pt").to(model.device)
outputs = model.generate(
    **inputs,
    max_new_tokens=2048,
    do_sample=True,
    temperature=0.6,
    top_p=0.95,
    top_k=20,
    return_dict_in_generate=True,
    output_hidden_states=True,
)

Let’s extract the newly generated tokens:

sequence = outputs.sequences[0]
new_ids = sequence[inputs["input_ids"].shape[-1]:]
ids = new_ids.tolist()

Extracting the reasoning trace

We can use Qwen’s open/close “think” tokens to extract the reasoning trace: the intermediary tokens a model generates to think “step-by-step.”

think_open = tokenizer.convert_tokens_to_ids("<think>")
think_close = tokenizer.convert_tokens_to_ids("</think>")

Now, we find the index position of our open/close tokens. Note that small reasoning models often think for a long time---longer than what we set for max_new_tokens. That means we might not have an index for think_close. We set think_close to the length of the full sequence if this happens.

if think_open in ids:
    start = ids.index(think_open)
else:
    start = 0

if think_close in ids:
    print("OK: reasoning completed")
    end = ids.index(think_close)
else:
    print("WARNING: met token budget but didn't finish reasoning")
    end = len(ids)

No value for think_close here, but no matter: We can still segment steps on a truncated trace. Let’s extract the trace and assign the result to cot_ids:

cot_ids = new_ids[start:end]

Now we record where the chain-of-thought sits in the full sequence so we can recover its hidden states in a moment.

prompt_len = inputs["input_ids"].shape[-1]
cot_slice = slice(prompt_len + start, prompt_len + end)

Extracting the reasoning trajectory

Since we set output_hidden_states=True, we have all the hidden states the model computed while generating its trace. The sequence of these vectors is the model’s trajectory through the reasoning. For the rest of our analysis, we’ll use the middle-most layer’s hidden states.

steps = outputs.hidden_states
num_layer = len(steps[0])
LAYER = num_layer // 2 + 1  # Do `+1` because layer0 = initial embeddings
print("Using layer", LAYER, "of", num_layer)

Now, we concatenate one layer’s vectors across all steps into (num_token, dim). Each vector is the state after reading a token, so this covers every position except the very last generated token (which has no follow-on state).

trajectory = []
for step in steps:
    step = step[LAYER][0]
    trajectory.append(step)

trajectory = torch.cat(trajectory, dim=0).float().cpu()

Index this trajectory for the chain-of-thought tokens:

hidden = trajectory[cot_slice]

And trim those tokens if necessary so token positions and hidden vectors stay aligned:

cot_ids = cot_ids[:hidden.shape[0]]

Finding Steps

With reasoning traces in hand, we’ll attempt to recover the model’s reasoning steps. Nothing in cot_ids suggests how to do this, but theoretically, Qwen should have broken its outputs into such steps---that’s part of how it and other reasoning models were trained. So, our task in this section is to see if we can find some kind of “step signal” in the model’s hidden states.

Window size

To do so, we need to think a little about how big a “step” should be. While a model could theoretically treat each individual token as its own step, meaningful changes in the model’s internal trajectory are more likely to unfold over short multi-token regions. The question, then, is: Over how many tokens should we average when estimating local change?

One way to estimate this scale is to measure how quickly the hidden-state trajectory decorrelates with itself over token lag. If nearby hidden states point in similar directions, then the trajectory is locally coherent. As we compare states separated by larger lags, this similarity should decay. We can therefore use the decay timescale of this autocorrelation curve as an estimate of the trajectory’s intrinsic coherence length.

Below, we implement a function to compute this coherence length. First, we optionally center the hidden states and L2-normalize them, which reduces the effect of shared common directions in representation space. Then, for each lag d, we compare the trajectory to a copy of itself shifted by d tokens and average the resulting hidden-state similarities. We define the coherence length tau as the first lag where this autocorrelation falls below 1/e1/e of its lag-1 value.

Intuitively, tau is a rough estimate of the number of tokens over which the hidden-state trajectory remains locally coherent. We use this as the base window size for subsequent change-point detection.

def coherence_length(hidden, max_lag=50, eps=1e-8):
    """Estimate the trajectory's intrinsic timescale.

    Parameters
    ----------
    hidden : torch.Tensor
        Trajectory hidden states, shape (num_token, num_dim)
    max_lag : int
        Maximum lag time (in tokens) we allow for a timescale
    eps : float
        Epsilon value for numerical stability

    Returns
    -------
    tuple[int, torch.Tensor]
        - tau: the lag at which the autocorrelation of the hidden-state
          trajectory decays to 1/e of its lag-1 value
        - autocorr: the full curve across all lags
    """
    # Subtract the mean direction to remove the model's common component, which
    # otherwise inflates every similarity toward a high floor and masks actual
    # decay signal
    h = hidden - hidden.mean(dim=0, keepdim=True)
    h = F.normalize(h, dim=-1, eps=eps)

    # Set our maximum lag: minimum beween `max_lag` and the number of tokens
    max_lag = min(max_lag, h.shape[0] - 1)

    # Run autocorrelation
    autocorr = []
    for d in range(1, max_lag + 1):
        ac = (h[:-d] * h[d:]).sum(-1).mean()
        autocorr.append(ac)

    autocorr = torch.stack(autocorr)

    # Set our baseline decay and compute crossings: where `autocorr` dips below
    # `target`
    target = autocorr[0] / np.e
    crossings = (autocorr < target).nonzero(as_tuple=False)

    # `tau` is the lag-1 value where values in `autocorr` decay below `target`
    if len(crossings):
        tau = crossings[0].item() + 1
    else:
        tau = max_lag

    return tau, autocorr

This gives us a data-driven estimate of an appropriate window size. If the autocorrelation decays quickly, then the model’s hidden-state trajectory changes substantially from token to token, suggesting a smaller window. If it decays slowly, then the trajectory remains coherent over many tokens, suggesting a larger window is more appropriate.

Let’s run our function.

WINDOW, autocorr = coherence_length(hidden)
print("Window size for chunking tokens:", WINDOW)

We can visualize autocorr to inspect how fast the trajectory sequence decorrelates at various lag times.

fig, ax = plt.subplots()
lag = np.arange(1, len(autocorr) + 1)
sns.lineplot(x=lag, y=autocorr.numpy(), marker="o", ax=ax)
ax.axhline(
    autocorr[0].item() / np.e, ls="--", color="crimson", label="1/e of lag-1"
)
ax.axvline(
    WINDOW, ls=":", color="steelblue", label=f"Coherence length = {WINDOW}"
)
ax.set(
    title="Trajectory Coherence",
    xlabel="Lag (tokens apart)",
    ylabel="Hidden-state autocorrelation",
)
ax.legend()
plt.show()

Boundary scoring

With our window size set, we can now define a boundary score. The goal is to identify positions where the model’s hidden-state trajectory changes directions sharply. Intuitively, if the model is carrying out one coherent reasoning step, then nearby hidden states should be similar to each other. If the model moves from one reasoning step to another, the average state before that point may differ noticeably from the average hidden state after it.

At each position i: compare the model’s average state just before i to its average state just after i.

A high score means the internal state shifted sharply at that point. We hypothesize that such points are likely to correspond to boundaries between reasoning steps: the end of one local computation and the start of another.

def score_chunks(hidden, w):
    """Score chunks in hidden states.

    Parameters
    ----------
    hidden : torch.Tensor
        Trajectory hidden states, shape (num_token, num_dim)
    w : int
        Window size

    Returns
    --------
    torch.Tensor
        Boundary scores (cosine distances)
    """
    seq_len = hidden.shape[0]
    
    scores = torch.zeros(seq_len)
    for i in range(seq_len):
        start = max(0, i - w)
        end = min(seq_len, i + w)

        left = hidden[start:i]
        right = hidden[i:end]
        if len(left) == 0 or len(right) == 0:
            continue

        L = left.mean(dim=0)
        R = right.mean(dim=0)

        scores[i] = 1.0 - F.cosine_similarity(L, R, dim=0)

    return scores

We will also scale these scores so that unusually large boundary scores are easier to identify. Rather than using the ordinary mean and standard deviation to compute z scores, we use a robust alternative based on the median and median absolute deviation, or MAD. This makes the scaling less sensitive to a few very large boundary scores.

def scale(x, eps=1e-8):
    """Center and scale a 1-D signal into robust standard units.

    This function works like a z-score, but it uses the median and median
    absolute deviation (MAD) instead of mean and standard deviation.

    Parameters
    ----------
    x : torch.Tensor or array-like
        Signal to standardize
    eps : float
        Epsilon value for numerical stability

    Returns
    -------
    np.ndarray
        Scaled signal
    """
    if isinstance(x, torch.Tensor):
        x = x.detach().cpu().numpy()

    x = np.asarray(x, dtype=float)
    median = np.median(x)
    mad = median_abs_deviation(x, scale=1.0)

    return (x - median) / (mad + eps)

The resulting values can be interpreted as robust standardized scores: positive values are above the typical boundary score, and large positive values indicate especially sharp changes in the hidden-state trajectory.

Now, we score our chunks and scale the scores:

raw_scores = score_chunks(hidden, w=WINDOW)
scores = scale(raw_scores)

print("Scored", len(scores), "positions")

Observe how many chunks we have. Given this size, we’re probably not at individual reasoning steps yet; steps are likely to include multiple chunks. To find steps’ actual boundaries, then, we need to find the peaks of our scores. These peaks are the local maxima of our change signal.

We use SciPy’s find_peaks() to do this. It has a few parameters; below, we’ll set two:

is_finite = scores[np.isfinite(scores)]
HEIGHT = np.percentile(is_finite, 90)
DISTANCE = 2 * WINDOW

print("Keeping peaks above the 90th percentile: height >=", round(HEIGHT, 2))
print("Peaks must be spaced at least", DISTANCE, "tokens apart")

Now, let’s find our boundaries:

boundaries, props = find_peaks(scores, distance=DISTANCE, height=HEIGHT)
print("Found", len(boundaries), "boundaries")

Segmenting traces

If we segment the tokens along our boundaries, what do they look like?

cut_points = [0, *boundaries.tolist(), len(cot_ids)]
zipped = zip(cut_points[:-1], cut_points[1:])

for step, (lo, hi) in enumerate(zipped, start=1):
    segment = tokenizer.decode(cot_ids[lo:hi])
    print(f"--- step {step} (tokens {lo}-{hi}) ---")
    print(segment)
    print()

That worked surprisingly well! Let’s visualize the entire boundary signal. Peaks are where the model’s internal state turned a corner in its reasoning.

First, format scores into a DataFrame:

scores_df = pd.DataFrame({"position": np.arange(len(scores)), "score": scores})

Now plot:

fig, ax = plt.subplots()
sns.lineplot(scores_df, x="position", y="score", ax=ax)
ax.scatter(
    boundaries,
    scores[boundaries],
    color="crimson",
    zorder=5,
    label="boundary",
)
ax.set(
    title=f"Boundary Scores (Layer={LAYER}, Window={WINDOW})",
    xlabel="Token position in reasoning trace",
    ylabel="Trajectory change",
)
ax.legend()
plt.show()

Clustering Steps

If you look at our segmented steps, you’ll see that they roughly fall into a few distinct modes. Some are arithmetical while others are discursive. Some introduce a new subproblem, while others may verify or summarize a previous result. Can we cluster steps in a way that reflects these modes?

The basic idea here is to turn each segmented step into a single vector, then cluster those steps. If two steps have similar hidden-state representations, we will treat them as instances of a similar internal mode.

Mean-pooling steps

To do so, we’ll represent each step by the mean of its mid-layer hidden vectors.

def segment_representations(hidden, cut_points):
    """Produce a mean-pooled representation for every reasoning step.

    Parameters
    ----------
    hidden : torch.Tensor
        Trajectory hidden states, shape (num_token, num_dim)
    cut_points : list[int]
        Step boundaries

    Returns
    -------
    tuple[torch.Tensor, list]
        Mean-pooled step vectors and the segmented token spans
    """
    spans = list(zip(cut_points[:-1], cut_points[1:]))

    vecs = []
    for lo, hi in spans:
        vecs.append(hidden[lo:hi].mean(dim=0))

    return torch.stack(vecs), spans

We produce our segmentations and normalize them onto the unit sphere:

step_vecs, spans = segment_representations(hidden, cut_points)
step_vecs = F.normalize(step_vecs, dim=-1).detach().cpu().numpy()

Choosing the number of modes

We now cluster the step vectors with K-means clustering. K-means tries to assign each step to one of kk clusters while minimizing the within-cluster squared distance:

min⁡C1,…,Ck∑c=1k∑j∈Cc∥v^j−μc∥22,\min_{\mathcal{C}_1,\ldots,\mathcal{C}_k} \sum_{c=1}^{k} \sum_{j \in \mathcal{C}_c} \lVert \hat{\mathbf{v}}_j - \boldsymbol{\mu}_c \rVert_2^2,
(2)

where μc\boldsymbol{\mu}_c is the centroid of cluster cc:

μc=1∣Cc∣∑j∈Ccv^j.\boldsymbol{\mu}_c = \frac{1}{|\mathcal{C}_c|} \sum_{j \in \mathcal{C}_c} \hat{\mathbf{v}}_j.
(3)

The remaining question is how many clusters to use. We try a few candidate values of kk and choose the one with the best silhouette score.

For each step jj, the silhouette score compares two quantities:

The silhouette score for step jj is:

sj=bj−ajmax⁡(aj,bj).s_j = \frac{b_j - a_j}{\max(a_j, b_j)}.
(4)

This score is close to 1 when the point is much closer to its own cluster than to other clusters. It is near 0 when the point lies between clusters, and negative when the point may have been assigned to the wrong cluster.

To select kk, we use the mean silhouette score across all steps and select the value for $k4 that maximizes that score.

num_step = len(step_vecs)

max_k = min(5, num_step - 1)
candidate_k = range(2, max_k + 1)

scores_by_k = {}
for k in candidate_k:
    kmeans = KMeans(n_clusters=k, n_init=10, random_state=5167)
    labels_k = kmeans.fit_predict(step_vecs)
    scores_by_k[k] = silhouette_score(step_vecs, labels_k)

Now we select the value of kk with the highest silhouette score. If there is a tie, the code below prefers the larger value of kk, since we may prefer a slightly more fine-grained decomposition of step types.

best_k = max(scores_by_k, key=lambda k: (scores_by_k[k], k))

for k, score in scores_by_k.items():
    marker = ""
    if k == best_k:
        marker = "  <- best"

    print(f"k={k}: silhouette={score:.3f}{marker}")

kmeans = KMeans(n_clusters=best_k, n_init=10, random_state=5167).fit(step_vecs)

Inspecting modes

What spans did the clusterer assign to each cluster? Let’s gather up the spans into a DataFrame and associate them with the labels.

cluster_df = pd.DataFrame(spans, columns=["lo", "hi"])
cluster_df["step"] = np.arange(len(cluster_df))
cluster_df["label"] = kmeans.labels_

Now, we add the decoded text for each span:

cluster_df["text"] = cluster_df.apply(
    lambda row: tokenizer.decode(cot_ids[row["lo"]:row["hi"]]),
    axis=1
)

And print the results:

for c in range(best_k):
    group = cluster_df[cluster_df["label"] == c]

    print(f"=== mode {c}  |  {len(group)} steps ===")
    for _, row in group.iterrows():
        text = repr(row["text"][:100])
        print(f"  step {row['step']:2d}: {text}")

    print()

The mode numbers themselves are arbitrary: Mode 0 is not necessarily more important or earlier than mode 1. What matters is whether the steps grouped together by the clusterer appear to share a recognizable function.

Finally, let’s look at the distribution of modes over the trace.

num_col = int(np.ceil(np.sqrt(num_step)))
num_row = int(np.ceil(num_step / num_col))

grid = np.full((num_row, num_col), np.nan)
grid.flat[:num_step] = kmeans.labels_

cmap = sns.color_palette("deep", best_k)

fig, ax = plt.subplots()
sns.heatmap(
    grid,
    mask=np.isnan(grid),
    cmap=cmap,
    vmin=0,
    vmax=best_k - 1,
    cbar=False,
    square=True,
    linewidths=2,
    xticklabels=False,
    yticklabels=False,
    ax=ax,
)

for i, color in enumerate(cmap):
    ax.scatter([], [], marker="s", s=100, color=color, label=f"mode {i}")

ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", frameon=False)
ax.set(
    title="Reasoning Modes Aross the Trace",
    xlabel="",
    ylabel="",
)
plt.show()
Introduction to Interpretability for Language Models
Metamodeling Approaches