Open source · Machine learning
PyTorch-CKA
A one-call PyTorch toolkit for measuring and visualizing representation similarity across models, layers, and checkpoints.- Role
- Author
- Time frame
- Open source
- Python
- PyTorch
- CUDA
Why this project
CKA was standard. The workflow around it was not.
Centered Kernel Alignment is a common way to ask whether two networks represent a dataset similarly across layers, models, or training stages. The math is well known. The scaffolding is not.
Most research pipelines still re-implement feature extraction, batching, kernel computation, and plotting for every experiment. I started the library because I wanted reliable CKA measurements without rebuilding that stack each time.
Goals & constraints
Make the common experiment a single call.
- One function call for the common case: models, dataloaders, and target layers.
- A second path for people who already have activations.
- Stay practical on larger layer-wise comparisons: fast on GPU, careful with memory.
- Include visualizations so results can be inspected without a second toolkit.
Approach
Design the API around the experiment, not the kernel.
The library only works if a researcher can start an experiment without first becoming a CKA implementer. Speed and memory matter, but they come after a path that is short enough to reuse.
One compute call, two entry points
compute_cka handles models and dataloaders. cka_from_features accepts pre-extracted activations, including mixed layer shapes. The split keeps the happy path short without blocking advanced setups.
Vectorized PyTorch on GPU, not a Python loop over layers
The original pain was wall-clock time. I implemented the core path with vectorized ops and GPU execution so a full ResNet-18 layer-wise comparison could finish in a research loop, not overnight.
Accumulate HSIC in batches and free activations
Speed is useless if the run OOMs. Batch-wise HSIC accumulation and explicit deallocation cut GPU memory retention by about 96%, and the same path supports Hugging Face models, DataParallel, and DDP.
Implementation
What made the toolkit worth publishing.
A 44× faster core path
On ResNet-18 layer-wise comparisons over CIFAR-10 — 18 representational layers, NVIDIA H100 — the library is up to 44× faster than a common reimplementation baseline.
Visualizations that ship with the measurement
Heatmaps cover layer-to-layer similarity. Trend plots cover checkpoints and epochs. Researchers can inspect a matrix without bolting on a separate plotting stack.
Published where researchers already install tools
The package is on PyPI as pytorch-cka, so the workflow is pip install plus one function call.
Outcome
A reusable CKA workflow on PyPI.
The library is public on GitHub and PyPI, with about 400 monthly downloads. The useful outcome is not the number. It is that a representation-similarity experiment can start from a documented API instead of another notebook rewrite.
Reflection
Setup friction was the real product problem.
The kernel was the obvious problem. The product problem was setup friction. A one-call API and a feature-only escape hatch mattered as much as the 44× speedup.
Memory was the other surprise. Making CKA faster without making it smaller would have failed on the exact layer-wise comparisons the library exists to run.