Decoding the Black Box: Understanding Neural Network Similarity with CKA
Decoding the Black Box: Understanding Neural Network Similarity with CKA
Have you ever wondered if two different neural networks, trained on the same task but initialized differently, actually "see" the world in the same way? Or if Layer 4 of a ResNet is doing something fundamentally different from Layer 5?
Comparing the internal representations of deep networks is notoriously difficult. You can't simply subtract the weight matrices because the networks might have learned the same features but stored them in different neurons (a permutation problem) or scaled them differently.
Enter Centered Kernel Alignment (CKA). In this post, we dive into the mechanics of CKA—a powerful metric that allows us to compare neural network representations regardless of their dimensionality or rotation.
The Core Intuition: From Vectors to Structures
Most traditional methods, like Canonical Correlation Analysis (CCA), try to find a linear mapping between the feature vectors of two layers. However, these methods often "collapse" when the number of features exceeds the number of samples, and they are sensitive to the specific coordinate system used.
The CKA breakthrough is a shift in perspective: Instead of comparing feature vectors, we compare similarity structures.
Imagine two people describing the same room. One uses GPS coordinates; the other uses relative distances (e.g., "the chair is 2 meters from the table"). While their coordinate systems differ, the relative distance between the chair and the table remains constant.
CKA does exactly this. It creates a Representational Similarity Matrix (RSM) for each layer. If two layers organize the data in the same way—meaning examples that are "similar" in Layer A are also "similar" in Layer B—CKA will yield a high score, regardless of how the features are rotated or scaled.
The Mathematical Blueprint
To compute CKA, we rely on the Hilbert-Schmidt Independence Criterion (HSIC). Here is the step-by-step mathematical flow:
1. The Gram Matrix (The RSM)
For two activation matrices $X \in \mathbb{R}^{n \times p_1}$ and $Y \in \mathbb{R}^{n \times p_2}$, we first compute their Gram matrices: $$K = XX^T, \quad L = YY^T$$ Each entry $(i, j)$ in these matrices represents the dot product (similarity) between example $i$ and example $j$.
2. Centering the Kernels
To ensure the metric is unbiased, we center the matrices using the centering matrix $H = I_n - \frac{1}{n} \mathbf{1}\mathbf{1}^T$: $$K_c = HKH, \quad L_c = HLH$$
3. The CKA Formula
The final similarity is the normalized HSIC: $$\text{CKA}(K, L) = \frac{\text{HSIC}(K, L)}{\sqrt{\text{HSIC}(K, K) \text{HSIC}(L, L)}}$$ Where $\text{HSIC}(K, L) = \frac{1}{(n-1)^2} \text{tr}(K_c L_c)$. This results in a score between 0 and 1, where 1 indicates identical representational structures.
Visualizing the Pipeline
The following diagram illustrates the journey from raw activations to a final similarity score.
(I - 1/n 11áµ€) K (I - 1/n 11áµ€)"] end subgraph HSIC_Calculation ["Step 3: HSIC Computation"] HSIC_KL["HSIC(K, L) = trace(K_c @ L_c)"] HSIC_KK["HSIC(K, K) = trace(K_c @ K_c)"] HSIC_LL["HSIC(L, L) = trace(L_c @ L_c)"] end subgraph Final_Metric ["Step 4: CKA Normalization"] CKA_Formula["CKA = HSIC(K, L) / sqrt(HSIC(K, K) * HSIC(L, L))"] Result["Similarity Score [0, 1]"] end %% Connections Data --> NetA Data --> NetB NetA --> RepA NetB --> RepB RepA --> K RepB --> L K --- Note1 L --- Note1 K --> CenterOp --> Kc L --> CenterOp --> Lc Kc --> HSIC_KL Lc --> HSIC_KL Kc --> HSIC_KK Lc --> HSIC_LL HSIC_KL --> CKA_Formula HSIC_KK --> CKA_Formula HSIC_LL --> CKA_Formula CKA_Formula --> Result %% Styling style Result fill:#f9f,stroke:#333,stroke-width:2px style Note1 fill:#fff,stroke-dasharray: 5 5 style CenterOp fill:#e1f5fe,stroke:#01579b
Implementation in PyTorch
Below is a production-ready implementation of Linear CKA. We include a test case that compares two networks with the same architecture but different random initializations to see if they converge to similar representations.
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import numpy as np
class CKA:
"""
Implementation of Centered Kernel Alignment (CKA).
Invariant to orthogonal transformation and isotropic scaling.
"""
@staticmethod
def center_kernel(K):
n = K.shape[0]
unit = torch.ones((n, n), device=K.device) / n
return K - unit @ K - K @ unit + unit @ K @ unit
@classmethod
def linear_cka(cls, X, Y):
X, Y = X.float(), Y.float()
# 1. Compute Gram matrices (RSMs)
K = X @ X.T
L = Y @ Y.T
# 2. Center the kernels
K_c = cls.center_kernel(K)
L_c = cls.center_kernel(L)
# 3. Compute HSIC using Frobenius inner product
hsic_kl = torch.sum(K_c * L_c)
hsic_kk = torch.sum(K_c * K_c)
hsic_ll = torch.sum(L_c * L_c)
return (hsic_kl / torch.sqrt(hsic_kk * hsic_ll)).item()
# --- Demonstration Setup ---
class SimpleNet(nn.Module):
def __init__(self, input_dim, hidden_dim):
super().__init__()
self.layer1 = nn.Linear(input_dim, hidden_dim)
self.layer2 = nn.Linear(hidden_dim, hidden_dim)
self.layer3 = nn.Linear(hidden_dim, 2)
def forward(self, x, return_reps=False):
rep1 = F.relu(self.layer1(x))
rep2 = F.relu(self.layer2(rep1))
out = self.layer3(rep2)
return (out, rep1, rep2) if return_reps else out
# Data Generation
X, y = make_classification(n_samples=500, n_features=20, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
scaler = StandardScaler()
X_train = torch.tensor(scaler.fit_transform(X_train), dtype=torch.float32)
X_test = torch.tensor(scaler.transform(X_test), dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.long)
# Train two networks with different seeds
hidden_dim = 128
net_a, net_b = SimpleNet(20, hidden_dim), SimpleNet(20, hidden_dim)
opt_a = torch.optim.Adam(net_a.parameters(), lr=0.01)
opt_b = torch.optim.Adam(net_b.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()
for epoch in range(100):
for net, opt in [(net_a, opt_a), (net_b, opt_b)]:
opt.zero_grad()
criterion(net(X_train), y_train).backward()
opt.step()
# Extract and Compare
net_a.eval(); net_b.eval()
with torch.no_grad():
_, rep_a1, rep_a2 = net_a(X_test, return_reps=True)
_, rep_b1, rep_b2 = net_b(X_test, return_reps=True)
print(f"Similarity (Net A L1 vs Net B L1): {CKA.linear_cka(rep_a1, rep_b1):.4f}")
print(f"Similarity (Net A L2 vs Net B L2): {CKA.linear_cka(rep_a2, rep_b2):.4f}")
print(f"Similarity (Net A L1 vs Net A L2): {CKA.linear_cka(rep_a1, rep_a2):.4f}")Key Takeaways
- Orthogonal Invariance: CKA doesn't care if your features are rotated or flipped. If the relative distances between data points are preserved, the CKA score remains the same.
- Dimension Agnostic: You can compare a layer with 128 neurons to a layer with 512 neurons because the comparison happens in the $n \times n$ sample space, not the feature space.
- Practical Use Cases:
- Model Compression: Identify redundant layers that have nearly identical CKA scores.
- Transfer Learning: Check how much a pre-trained model's representations change when fine-tuned on a new dataset.
- Architecture Search: Compare how different activation functions or layer types affect the internal organization of data.
By shifting the focus from where a feature is stored to how the data is structured, CKA provides a robust lens into the hidden layers of our models.