Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Interpretability Goals › Measuring Representation Similarity Across Layers

Michael BrenndoerferFebruary 2, 20265 min read

Language AI HandbookAll chapters

Part 59 of 67Interpretability

Chapter 446Interpretability Goals › Measuring Representation Similarity Across Layers

Reading section 8 of 14 · Measuring Representation Similarity Across Layers · 5 minSection 8 of 14 · 5 min

Section 8 of 14 · 5 min read

Interpretability Goals › Measuring Representation Similarity Across LayersOne section · Sections in this chapter8 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

A natural question is: how different are representations from one layer to the next? This can be measured with centered kernel alignment (CKA), a metric that compares two representation matrices in a way that is invariant to rotation and isotropic scaling.

Centered Kernel Alignment (CKA)

CKA measures the similarity between two sets of representations by computing the normalized Hilbert-Schmidt independence criterion. A CKA score of 1.0 means the two representation spaces contain identical relational structure; 0.0 means they are completely independent. It is preferred over direct cosine similarity for comparing representations because it is invariant to orthogonal transformations and rescaling, making it possible to compare representations from different models or different layers without confounding effects from arbitrary rotation.

The linear CKA between representation matrices $X \in \mathbb{R}^{n \times p}$ and $Y \in \mathbb{R}^{n \times q}$ is:

\text{CKA}(X, Y) = \frac{\|Y^T X\|_F^2}{\|X^T X\|_F \cdot \|Y^T Y\|_F}

where:

$X \in \mathbb{R}^{n \times p}$ : the first representation matrix, with $n$ data points and $p$ features, centered by subtracting column means
$Y \in \mathbb{R}^{n \times q}$ : the second representation matrix, with the same $n$ data points and $q$ features, also column-centered
$\|M\|_F$ : the Frobenius norm of matrix $M$ , equal to the square root of the sum of squared entries
$Y^T X$ : the cross-covariance matrix capturing how features in $X$ co-vary with features in $Y$
$X^T X$ and $Y^T Y$ : the self-covariance matrices of each representation, serving as normalization terms

The numerator $\|Y^T X\|_F^2$ measures the total squared cross-covariance between the two representations. The denominator normalizes by each representation's own variance, making the score invariant to rotation and rescaling. A score near 1.0 means the two layers encode information in a geometrically similar way, regardless of the specific directions used. This invariance is important because two layers could represent the same information in completely different coordinate directions without CKA detecting a difference, which would happen with naive cosine similarity.

CKA was introduced by Kornblith et al. (2019) to address a fundamental problem in comparing neural representations: two networks or two layers might learn the same information but encode it in different coordinate systems. Direct comparison of weight matrices or activation vectors fails in this case, but CKA succeeds because it measures relational structure (how data points relate to each other) rather than absolute positions.

In[12]:

Code

def center_matrix(M):
    """Center matrix by subtracting column means."""
    return M - M.mean(axis=0)


def linear_cka(X, Y):
    """Compute linear CKA between two representation matrices."""
    X_c = center_matrix(X)
    Y_c = center_matrix(Y)
    numerator = np.linalg.norm(Y_c.T @ X_c, "fro") ** 2
    denominator = np.linalg.norm(X_c.T @ X_c, "fro") * np.linalg.norm(
        Y_c.T @ Y_c, "fro"
    )
    return numerator / denominator if denominator > 0 else 0.0


# Compute CKA between consecutive layers
cka_scores = []
for layer in range(num_layers - 1):
    X = all_hidden_states[:, layer, :]  # (sentences, hidden)
    Y = all_hidden_states[:, layer + 1, :]  # (sentences, hidden)
    cka_scores.append(linear_cka(X, Y))

Out[13]:

Visualization

Bar chart of linear CKA similarity scores between each pair of consecutive transformer layers. — Linear CKA similarity between consecutive DistilBERT transformer layers, computed on the four example sentences. A high CKA value indicates adjacent layers produce geometrically similar representations, meaning the transformation at that step is relatively minor. A lower CKA value signals a more substantial representational shift. In larger models with hundreds of sentences, CKA often reveals that the early and late layers are the sites of the most dramatic representational transformation.

CKA analysis reveals the transformation profile of the network. Layer pairs with high CKA are doing similar processing across that transition; pairs with low CKA represent significant representational shifts. In larger models, early layers and late layers often show low cross-layer CKA, while middle layers show high similarity, suggesting the bulk of semantic processing happens in a concentrated region of the network rather than uniformly across all layers.

This kind of global analysis is also useful for comparing models. If you compute CKA between corresponding layers of two different models trained on the same task, high values suggest the models have converged to similar internal representations despite starting from different random initializations. This convergence result, observed empirically, suggests that the representations are to some degree determined by the data and the task rather than being arbitrary.

Key ParametersLink Copied

The exploration above uses several parameters worth understanding:

model_name: The pretrained model to inspect. Swapping in a larger model (e.g., bert-base-uncased) will reveal more pronounced layer differentiation, since larger models build richer intermediate representations. Encoder-only models like BERT and DistilBERT are convenient for this kind of analysis because their representations are bidirectional; with decoder-only models you typically look at the last-token representation rather than CLS.
output_hidden_states: Must be set to True when loading the model to enable extraction of all layer activations. Without this flag, only the final hidden state is returned.
n_components (PCA): Controls how many principal components are computed. Two components enable 2D visualization; increasing to 10 or 50 gives a better picture of the full representational geometry, especially when you want to measure variance explained.
linear_cka: The linear variant of CKA assumes a linear kernel. A kernel CKA variant (using RBF kernels) can capture non-linear similarities but is more computationally expensive and behaves differently on small sample sizes.