Project · Published on PyPI

PyTorch-CKA

Author & maintainer

A Python library that computes layer-wise representation similarity, showing how similarly neural networks see the same inputs. I started it when the CKA library inside my own research tool became its bottleneck, and ended up rewriting the math so every layer pair comes out of a single GPU pass.

What is CKA?

Centered Kernel Alignment (CKA) scores how similarly two layers represent the same inputs, from 0 to 1. It feeds one batch through both layers and compares how each arranges the inputs relative to one another. A score of 1 means the arrangements match, even if one is rotated or evenly scaled. The layers can differ in width and come from different models, which makes CKA a common way to ask which layers of two networks learned alike.

Python
PyTorch
pytest
GitHub Actions
PyPI

PyPI Code

Results

6.8×faster than torch_cka
Two ResNet-18s compared on all 18 layers over 10,000 CIFAR-10 images on an RTX PRO 6000 Blackwell GPU. 30× faster than a plain NumPy version on the CPU.
−95.7%retained GPU memory
After one pass over the 10,000 CIFAR-10 test images on the same Blackwell GPU, 0.18 GB of GPU memory stays allocated, down from torch_cka’s 4.26 GB.
O(B²)HSIC per layer pair
Down from O(B³), where B is the batch size, because the trace never needed the full K @ L product.
~400monthly downloads
The May–August 2026 average on PyPI, across 11 releases tested on Python 3.10–3.14.

Why I started this project

It started as a bottleneck
in my own tool.

While building Machine Unlearning Comparator, I found the CKA computation behind its layer-similarity view extremely slow, even for 10 layers of ResNet-18 comparing two model pairs over 10,000 CIFAR-10 samples. It ran on torch_cka, the most widely used CKA library for PyTorch at the time. The animation replays its nested loop. For every batch, it takes the 18 × 18 layer pairs one at a time, and each pair stalls the GPU six times. Both of its HSIC calls end in three .item() calls, and each one makes Python wait for the GPU to hand back a number before it can queue more work.

torch_cka · the existing library

# runs once per batch
for i, x in enumerate(feats_a):         # i = 0
    K = x @ x.T
    for j, y in enumerate(feats_b):     # j = 0
        L = y @ y.T                     # rebuilt for every pair
        hsic[i, j] += _HSIC(K, L)       # trace(K @ L), 3 × .item()
        hsic_l[i, j] += _HSIC(L, L)     # L @ L, again for every i

layer pairs: 0 / 324
B³ matrix products: 0
GPU → CPU syncs: 0

Groundwork · 3Blue1Brown · Kornblith et al., ICML 2019

Before the code,
the linear algebra.

I started with the paper that introduced CKA and soon realized I couldn’t follow it without really understanding linear algebra. So I relearned it from scratch with 3Blue1Brown’s Essence of Linear Algebra, then learned PyTorch. The picture that stuck is that a matrix moves space. Rotating a layer’s activations, or scaling them evenly, changes every number but not which inputs sit close together, and that closeness is what the layer has learned. CKA scores those moves as no change at all. Only moves that distort it, like stretching or flattening, lower the score.

A layer’s activations are points, and a matrix moves them.
Rotate them and the geometry is intact, so CKA stays 1.00.
Scale them evenly and it is still 1.00.
Stretch and shear them, and similarity drops.
Flatten them to a line, and what the layer knew is lost.

matrix: 1.000.000.001.00
linear CKA vs. original: 1.00

Insight · Frobenius inner product

The trace never needed
the whole product.

CKA compares layers through their Gram matrices. For a batch of B inputs, a layer’s activations form a B × D matrix X, one row per input, and its Gram matrix K = XXᵀ holds the dot product of every pair of inputs, a B × B table of how alike they look to that layer. Layers of any width give tables of the same size, and rotating X leaves K unchanged. CKA’s core term is tr(KL). torch_cka built all of K @ L to get it, spending B³ multiply-adds and then keeping only the B numbers on the diagonal, which is a big part of why it was slow. Written out by hand, their sum is just Σ K_ij L_ij.

Each cell is how alike two inputs look to a layer. Brighter means more alike.
torch_cka’s way. K @ L costs B³ multiply-adds, one slice at a time.
The trace keeps B of the B² results and throws the rest away.
My way. Σ K ⊙ L reaches the same number with B² multiply-adds.

torch_cka · tr(K @ L): 0
pytorch-cka · Σ K ⊙ L: 0
at B = 1,000: –

Redesign · vectorized HSIC

Every layer pair,
in one pass.

My first working version already summed K ⊙ L, yet it was only 2.1× faster than torch_cka, taking 13.9 s against its 29.3 s on a Colab T4 with the comparator’s original and unlearned ResNet-18s. The per-pair loop was still there, and another HSIC term still multiplied K by L. So I rewrote the estimator to work on whole stacks. Each Gram matrix is built once, a single torch.einsum takes all 324 Frobenius products, and an outer product plus one matmul supply the rest of the unbiased HSIC.

Build each layer’s Gram matrix once and stack each model’s 18.
Read each B × B matrix row by row into B² numbers, giving (18, B²) and (B², 18).
Multiply the two, einsum("aij,bij->ab"), for all 324 Frobenius products at once.
Off-diagonal and column sums add the other two terms, with no K @ L.
Normalize to get the CKA matrix, then hover a cell.

Counts per batch, 18 × 18 layers. torch_cka makes 666 HSIC calls (1 per row, 2 per pair) with 3 .item() syncs each, and 342 Gram builds (each K once, each L per pair). pytorch-cka makes 1 call and builds 18 Grams per model.

pytorch-cka

# runs once per batch, for all 324 pairs at once
ga = torch.stack([x @ x.T for x in feats_a])   # (18, B, B)
gb = torch.stack([y @ y.T for y in feats_b])   # (18, B, B)
tr = einsum("aij,bij->ab", ga, gb) - einsum("ai,bi->ab", da, db)
t2 = outer(sum_a, sum_b) / ((n - 1) * (n - 2))
t3 = 2 / (n - 2) * col_a @ col_b.T
hsic_xy += (tr + t2 - t3) / (n * (n - 3))

HSIC calls per batch: 666 to 1
.item() per batch: 1,998 to 0
Gram builds per batch: 342 to 36

Memory management · forward hooks

Keep the answer,
drop the activations.

Speed was half the problem. After one pass over the 10,000 CIFAR-10 test images, 10 batches of 1,000, on a Blackwell GPU, torch_cka still held 4.26 GB of GPU memory, the last batch’s activations from all 18 layers of both models. Its forward hooks store each output on the CKA object and are never removed, and since the models keep the hooks, even del cka and gc.collect() can’t free it. pytorch-cka scopes its hooks to a with block, clears features every batch, and ends at 0.18 GB.

Every batch, hooks capture 18 layers from both models.
When the comparison ends, pytorch-cka removes its hooks and frees the features.
torch_cka’s hooks stay on the models, and so does its last batch, 4.26 GB.

Self-comparison in one forward pass
Two models run in parallel on CUDA streams
Unwraps DataParallel and DDP models
Reads Hugging Face model outputs

torch_cka holds: 0 MB
pytorch-cka holds: 0.0 MB
difference: –

Benchmark · ResNet-18 · CIFAR-10 · RTX PRO 6000

Same GPU, same layers,
5.87 s down to 0.86 s.

The benchmark compares two ResNet-18s across the 18 layers the network is named after, its 17 main-path convolutions and the final classifier. It runs over the 10,000 CIFAR-10 test images in batches of 1,000 on an NVIDIA RTX PRO 6000 Blackwell GPU, with fresh copies of the same weights for every run and the mean of five runs after a warm-up. pytorch-cka finished in 0.86 seconds, 6.8× faster than torch_cka and 30× faster than a plain NumPy version of the same estimator on the CPU. The work I removed grows cubically with batch size, so the gap should widen as batches grow.

elapsed: 0.00 s
pytorch-cka: 0.86 s
vs torch_cka: 6.8× faster

What I learned

Lessons I still use.

01
Derive it before optimizing it
Every implementation I read computed tr(K @ L) because that is how the formula is written. Working the estimator out term by term is what showed me the product was optional.
B³ → B²multiply-adds per layer pair
02
Don’t make the GPU wait on Python
An .item() inside a hot loop stalls the GPU until Python catches up. Keeping every value a tensor until the final division removed 1,998 syncs per batch.
1,998 → 0.item() syncs per batch
03
Memory has an owner
A forward hook keeps its object reachable, so del and gc can’t free what it points to. Scoping hooks to a with block and clearing features every batch cut retained memory from 4.26 GB to 0.18 GB.
4.26 → 0.18 GBretained after one pass
04
Measure against the real thing
My first version felt fast and was only 2.1× faster. Benchmarking against existing libraries on real models is what pointed at the loop, and later proved the rewrite.
5.87 s → 0.86 ssame GPU, same layers

Keep exploring.

ABC Fitness Orbit Explorer Machine Unlearning Comparator Multi-Agent Invoice Pipeline Get in touch

PyTorch-CKA

Results

It started as a bottleneckin my own tool.

Before the code,the linear algebra.

The trace never neededthe whole product.

Every layer pair,in one pass.

Keep the answer,drop the activations.

Same GPU, same layers,5.87 s down to 0.86 s.

Lessons I still use.

Derive it before optimizing it

Don’t make the GPU wait on Python

Memory has an owner

Measure against the real thing

Keep exploring.

It started as a bottleneck
in my own tool.

Before the code,
the linear algebra.

The trace never needed
the whole product.

Every layer pair,
in one pass.

Keep the answer,
drop the activations.

Same GPU, same layers,
5.87 s down to 0.86 s.