Open source · Machine learning
PyTorch-CKA
A one-call PyTorch toolkit for measuring and visualizing representation similarity across models, layers, and checkpoints.- Python
- PyTorch
- CUDA
Overview
Turning representation similarity analysis into a reusable research workflow.
Centered Kernel Alignment (CKA) is a standard way to measure how similar neural network representations are across layers, models, or training stages. Despite that, many research pipelines still re-implement feature extraction, batching, kernel computation, and plotting for every experiment. That friction slows iteration and makes it harder to compare runs under consistent conditions.
PyTorch-CKA packages the full workflow into a practical library: pass models, dataloaders, and target layers to a single compute call, or feed pre-extracted features directly when activations already exist. The core path is built on vectorized PyTorch operations with GPU acceleration and explicit memory deallocation so larger layer-wise comparisons stay practical. It also supports common research setups such as Hugging Face models, DataParallel, and Distributed Data Parallel.
Beyond the matrix itself, the library includes reusable visualizations—heatmaps for layer-to-layer similarity and trend plots for checkpoint or epoch comparisons—so results can be inspected without bolting on a separate analysis stack. The package is published on PyPI as pytorch-cka and is designed for researchers who want reliable CKA measurements without rebuilding the scaffolding around them.
Contribution
What I worked on
- Designed a one-call API that covers model-and-dataloader workflows and feature-only comparison paths.
- Accelerated core CKA computation by up to 44× on ResNet-18 layer-wise comparisons through vectorized PyTorch ops and GPU execution.
- Added memory-conscious activation handling and support for Hugging Face models, DataParallel, and DDP.
- Shipped reusable heatmap and trend visualizations for layer-wise and checkpoint-level analysis, and released the toolkit on PyPI.