Note
Centered Kernel Alignment (CKA) [1] is a similarity index between representations of features in neural networks, based on the Hilbert-Schmidt Independence Criterion (HSIC) [2]. Given a set of examples, CKA compares the representations of examples passed through the layers that we want to compare.
Given two matrices
Then, the
where
Note
However, naive computation of linear CKA (i.e.: the previous equation) requires maintaining the activations across the entire dataset in memory, which is challenging for wide and deep networks [3].
Therefore, we need to define the unbiased estimator of HSIC so that the value of CKA is independent of the batch size
where
with
This project requires python >= 3.10. As first step, clone the repository
git clone https://github.com/RistoAle97/centered-kernel-alignment
Then, you can install the necessary packages with
Warning
This will install PyTorch not compiled with CUDA if you are on Windows. If you want to use your GPU during the computation, you should follow the official site
pip install -r requirements.txt
Take a look at the examples
directory to understand how to compute CKA in different scenarios.
Model compared with itself | Different models compared |
---|---|
![]() |
![]() |
[1] Kornblith, Simon, et al. "Similarity of neural network representations revisited." International Conference on Machine Learning. PMLR, 2019.
[2] Wang, Tinghua, Xiaolu Dai, and Yuze Liu. "Learning with Hilbert–Schmidt independence criterion: A review and new perspectives." Knowledge-based systems 234 (2021): 107567.
[3] Nguyen, Thao, Maithra Raghu, and Simon Kornblith. "Do wide and deep networks learn the same things? uncovering how neural network representations vary with width and depth." arXiv preprint arXiv:2010.15327 (2020).
This project is also based on the following repositories:
- representation_similarity (original implementation).
- PyTorch-Model-Compare (nice PyTorch implementation that employs hooks).
-
CKA.pytorch (minibatch version of CKA and useful batched implementation of
$HSIC_1$ ).
This project is MIT licensed.