Manifold Approximation leads to Robust Kernel Alignment
Organizations: MIT
Abstract
Centered kernel alignment (CKA) is a popular metric for comparing representations, determining equivalence of networks, and conducting neuroscience research. However, CKA does not account for the underlying manifold and relies on numerous heuristics that cause it to behave differently at different scales of data. In this work, we propose Manifold-approximated Kernel Alignment (MKA), which incorporates manifold geometry (i.e., local neighborhood) into the alignment task. We derive a theoretical framework for MKA. Our construction removes representation-dependent degree weighting, yields a closed-form alignment score, and gives each row a probabilistic interpretation as a local Markov transition distribution. We further establish kernel- and score-level stability with respect to the neighborhood size. We perform empirical evaluations on synthetic datasets and real-world examples to characterize and compare MKA to its contemporaries. In controlled synthetic experiments, MKA is consistently less sensitive to hyperparameters and captures the selected topology-preserving relationships. On the large-scale benchmark, MKA achieves its strongest aggregate performance in vision, is competitive on graphs, and is weaker than CKA in several NLP evaluations. Our findings suggest that manifold-aware kernel alignment provides a robust foundation for measuring representations and is potentially useful when local neighborhood geometry is important, with possible applications in representation learning.
Figures & tables
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
| Method | Distance Matrix | kNN | Kernel/Weight | Score Computation | Residual | |
|---|---|---|---|---|---|---|
| CKA | 0.0120 | – | 0.0214 | 0.0036 | 0.0007 | |
| kCKA | – | 0.0592 | 0.0081 | 0.0037 | 0.0007 | |
| MKA | – | 0.0412 | 0.3950 | 0.0016 | 0.0008 | |
| CKA | 0.6522 | – | 1.9994 | 0.5585 | 0.0116 | |
| kCKA | – | 0.6778 | 0.6142 | 0.5657 | 0.0069 | |
| MKA | – | 0.6275 | 3.9232 | 0.2890 | 0.0070 |
| Test | Accuracy Correlation | ||||||
|---|---|---|---|---|---|---|---|
| Dataset | IN100 | ||||||
| Architecture | RNet18 | RNet34 | RNet101 | VGG11 | VGG19 | ViT B32 | ViT L32 |
| CKA | 0.33 | -0.09 | 0.08 | 0.02 | -0.22 | -0.23 | 0.02 |
| CKA ( ) | 0.29 | -0.06 | 0.01 | 0.00 | -0.10 | -0.28 | 0.09 |
| CKA ( ) | 0.02 | 0.19 | -0.08 | -0.14 | -0.23 | -0.15 | 0.02 |
| kCKA ( ) | 0.00 | -0.01 | 0.15 | 0.09 | 0.58 | -0.16 | 0.10 |
| Test | JSD Correlation | Disagreement Correlation | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Dataset | IN100 | IN100 | ||||||||||||
| Architecture | RNet18 | RNet34 | RNet101 | VGG11 | VGG19 | ViT B32 | ViT L32 | RNet18 | RNet34 | RNet101 | VGG11 | VGG19 | ViT B32 | ViT L32 |
| CKA | 0.26 | 0.02 | 0.30 | -0.09 | 0.04 | 0.02 | -0.12 | 0.36 | 0.00 | 0.29 | -0.01 | -0.25 | 0.00 | -0.05 |
| CKA ( ) | 0.20 | -0.13 | 0.30 | -0.08 | 0.19 | 0.03 | -0.21 | 0.33 | -0.14 | 0.35 | -0.00 | -0.21 | 0.04 | -0.16 |
| CKA ( ) | 0.07 | -0.03 | -0.26 | 0.02 | -0.13 | 0.32 | -0.61 | 0.35 | -0.18 | 0.08 | 0.39 | 0.13 | 0.35 | -0.48 |
| kCKA ( ) | -0.34 | 0.35 | 0.43 | -0.02 | 0.24 | 0.32 | -0.32 | -0.56 | 0.45 | 0.18 | -0.18 | -0.17 | 0.22 | -0.17 |
| Evaluation | AUPRC | Conformity Rate | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Dataset | IN100 | IN100 | ||||||||||||
| Architecture | RNet18 | RNet34 | RNet101 | VGG11 | VGG19 | ViT B32 | ViT L32 | RNet18 | RNet34 | RNet101 | VGG11 | VGG19 | ViT B32 | ViT L32 |
| CKA | 1.00 | 1.00 | 1.00 | 1.00 | 0.77 | 0.52 | 0.81 | 1.00 | 1.00 | 1.00 | 1.00 | 0.77 | 0.65 | 0.83 |
| CKA ( ) | 1.00 | 1.00 | 1.00 | 1.00 | 0.77 | 0.78 | 0.86 | 1.00 | 1.00 | 1.00 | 1.00 | 0.75 | 0.81 | 0.87 |
| CKA ( ) | 0.65 | 0.70 | 0.82 | 0.80 | 0.76 | 0.77 | 0.97 | 0.71 | 0.74 | 0.87 | 0.84 | 0.73 | 0.75 | 0.99 |
| kCKA ( ) | 1.00 | 1.00 | 1.00 | 0.88 | 0.73 | 0.57 | 0.90 | 1.00 | 1.00 | 1.00 | 0.95 | 0.70 | 0.73 | 0.97 |
| Evaluation | AUPRC | Conformity Rate | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Dataset | IN100 | IN100 | ||||||||||||
| Architecture | RNet18 | RNet34 | RNet101 | VGG11 | VGG19 | ViT B32 | ViT L32 | RNet18 | RNet34 | RNet101 | VGG11 | VGG19 | ViT B32 | ViT L32 |
| CKA | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 0.86 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 0.96 |
| CKA ( ) | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 0.90 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 0.97 |
| CKA ( ) | 0.77 | 0.62 | 0.66 | 0.80 | 0.74 | 1.00 | 0.79 | 0.94 | 0.85 | 0.88 | 0.93 | 0.94 | 1.00 | 0.93 |
| kCKA ( ) | 0.94 | 0.92 | 0.98 | 1.00 | 1.00 | 1.00 | 0.91 | 0.99 | 0.98 | 1.00 | 1.00 | 1.00 | 1.00 | 0.97 |
| Evaluation | Spearman | Conformity Rate | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Dataset | IN100 | IN100 | ||||||||||||
| Architecture | RNet18 | RNet34 | RNet101 | VGG11 | VGG19 | ViT B32 | ViT L32 | RNet18 | RNet34 | RNet101 | VGG11 | VGG19 | ViT B32 | ViT L32 |
| CKA | 0.97 | 0.79 | 0.97 | 0.88 | 0.93 | 1.00 | 1.00 | 0.95 | 0.98 | 0.99 | 0.90 | 0.92 | 1.00 | 1.00 |
| CKA ( ) | 0.90 | 0.73 | 0.91 | 0.97 | 0.94 | 1.00 | 1.00 | 0.94 | 0.95 | 0.97 | 0.95 | 0.93 | 1.00 | 1.00 |
| CKA ( ) | 0.97 | 0.89 | 0.84 | 1.00 | 0.84 | 0.95 | 1.00 | 0.96 | 0.96 | 0.96 | 1.00 | 0.94 | 0.93 | 1.00 |
| kCKA ( ) | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 |
| Representation | CLS Token | |
|---|---|---|
| Test | Accuracy Correlation | |
| Architecture | BERT | ALBERT |
| CKA | 0.11 | 0.26 |
| CKA ( ) | 0.08 | 0.09 |
| CKA ( ) | -0.17 | -0.24 |
| kCKA ( | 0.03 | -0.24 |
| Representation | CLS Token | Mean-pooled Token | ||||||
|---|---|---|---|---|---|---|---|---|
| Test | JSD Correlation | Disagreement Correlation | JSD Correlation | Disagreement Correlation | ||||
| Architecture | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT |
| CKA | 0.36 | 0.29 | 0.15 | 0.57 | 0.51 | 0.06 | 0.38 | 0.18 |
| CKA ( ) | 0.31 | 0.05 | 0.50 | 0.22 | 0.53 | 0.02 | 0.49 | 0.16 |
| CKA ( ) | -0.25 | -0.11 | 0.27 | -0.04 | 0.13 | -0.06 | -0.25 | 0.05 |
| kCKA ( ) | 0.27 | 0.03 | 0.33 | -0.00 | 0.45 | 0.05 | 0.41 | 0.10 |
| Representation | CLS Token | Mean-pooled Token | ||||||
|---|---|---|---|---|---|---|---|---|
| Evaluation | AUPRC | Conformity Rate | AUPRC | Conformity Rate | ||||
| Architecture | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT |
| CKA | 0.75 | 0.80 | 0.89 | 0.93 | 0.66 | 0.45 | 0.86 | 0.71 |
| CKA ( ) | 0.74 | 0.79 | 0.89 | 0.93 | 0.58 | 0.45 | 0.81 | 0.68 |
| CKA ( ) | 0.56 | 0.54 | 0.82 | 0.75 | 0.45 | 0.31 | 0.66 | 0.53 |
| kCKA ( ) | 0.69 | 0.46 | 0.83 | 0.68 | 0.50 | 0.43 | 0.73 | 0.59 |
| Representation | CLS Token | Mean-pooled Token | ||||||
|---|---|---|---|---|---|---|---|---|
| Evaluation | AUPRC | Conformity Rate | AUPRC | Conformity Rate | ||||
| Architecture | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT |
| CKA | 0.59 | 0.63 | 0.88 | 0.67 | 0.55 | 0.56 | 0.85 | 0.64 |
| CKA ( ) | 0.58 | 0.67 | 0.87 | 0.81 | 0.55 | 0.56 | 0.84 | 0.66 |
| CKA ( ) | 0.50 | 0.43 | 0.79 | 0.58 | 0.55 | 0.32 | 0.74 | 0.52 |
| kCKA ( ) | 0.57 | 0.61 | 0.83 | 0.65 | 0.56 | 0.58 | 0.82 | 0.66 |
| Representation | CLS Token | Mean-pooled Token | ||||||
|---|---|---|---|---|---|---|---|---|
| Evaluation | AUPRC | Conformity Rate | AUPRC | Conformity Rate | ||||
| Architecture | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT |
| CKA | 0.44 | 0.85 | 0.84 | 0.84 | 0.36 | 0.74 | 0.79 | 0.78 |
| CKA ( ) | 0.45 | 0.85 | 0.80 | 0.86 | 0.34 | 0.72 | 0.77 | 0.78 |
| CKA ( ) | 0.43 | 0.69 | 0.71 | 0.74 | 0.35 | 0.42 | 0.56 | 0.55 |
| kCKA ( ) | 0.37 | 0.76 | 0.71 | 0.78 | 0.35 | 0.63 | 0.73 | 0.73 |
| Representation | CLS Token | Mean-pooled Token | ||||||
|---|---|---|---|---|---|---|---|---|
| Evaluation | Spearman | Conformity Rate | Spearman | Conformity Rate | ||||
| Architecture | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT | BERT | ALBERT |
| CKA | 0.99 | 0.98 | 0.98 | 0.97 | 1.00 | 0.99 | 0.99 | 0.99 |
| CKA ( ) | 0.99 | 0.98 | 0.98 | 0.97 | 1.00 | 1.00 | 0.99 | 1.00 |
| CKA ( ) | 0.92 | 0.95 | 0.90 | 0.93 | 0.95 | 1.00 | 0.96 | 1.00 |
| kCKA ( ) | 0.91 | 0.90 | 0.90 | 0.89 | 1.00 | 1.00 | 1.00 | 1.00 |
| Evaluation | Spearman | |||||||||
|---|---|---|---|---|---|---|---|---|---|---|
| Dataset | Cora | Flickr | OGBN-Arxiv | |||||||
| Architecture | GCN | SAGE | GAT | PGNN | GCN | SAGE | GAT | GCN | SAGE | GAT |
| CKA | 0.04 | -0.15 | 0.03 | -0.01 | 0.50 | 0.41 | -0.14 | -0.07 | 0.14 | -0.03 |
| CKA ( ) | 0.09 | -0.22 | 0.00 | -0.06 | 0.52 | 0.17 | -0.16 | -0.12 | 0.18 | -0.13 |
| CKA ( ) | 0.01 | -0.29 | 0.07 | 0.03 | 0.43 | 0.14 | 0.11 | -0.22 | 0.17 | -0.28 |
| kCKA ( ) | 0.07 | -0.25 | -0.01 | 0.11 | 0.42 | -0.22 | -0.28 | -0.06 | 0.23 | -0.16 |
| Type | Grounding by Prediction | |||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Test | JSD Correlation | Disagreement Correlation | ||||||||||||||||||
| Evaluation | Spearman | Spearman | ||||||||||||||||||
| Dataset | Cora | Flickr | OGBN-Arxiv | Cora | Flickr | OGBN-Arxiv | ||||||||||||||
| Architecture | GCN | SAGE | GAT | PGNN | GCN | SAGE | GAT | GCN | SAGE | GAT | GCN | SAGE | GAT | PGNN | GCN | SAGE | GAT | GCN | SAGE | GAT |
| CKA | 0.73 | 0.12 | 0.59 | -0.23 | 0.54 | 0.35 | 0.12 | 0.18 | 0.19 | 0.38 | 0.65 | -0.05 | 0.44 | 0.01 | 0.26 | 0.32 | 0.02 | 0.13 | 0.14 | 0.21 |
| CKA ( ) | 0.73 | 0.30 | 0.52 | -0.08 | 0.56 | 0.18 | 0.15 | 0.22 | 0.22 | 0.25 | 0.64 | -0.06 | 0.34 | 0.14 | 0.29 | 0.15 | 0.01 | 0.21 | 0.18 | 0.09 |
| Evaluation | AUPRC | Conformity Rate | ||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Dataset | Cora | Flickr | OGBN-Arxiv | Cora | Flickr | OGBN-Arxiv | ||||||||||||||
| Architecture | GCN | SAGE | GAT | PGNN | GCN | SAGE | GAT | GCN | SAGE | GAT | GCN | SAGE | GAT | PGNN | GCN | SAGE | GAT | GCN | SAGE | GAT |
| CKA | 0.48 | 0.56 | 0.43 | 0.25 | 0.88 | 0.42 | 0.31 | 1.00 | 1.00 | 1.00 | 0.72 | 0.78 | 0.54 | 0.64 | 0.96 | 0.50 | 0.61 | 1.00 | 1.00 | 1.00 |
| CKA ( ) | 0.50 | 0.50 | 0.52 | 0.26 | 0.92 | 0.45 | 0.38 | 1.00 | 0.96 | 1.00 | 0.75 | 0.77 | 0.76 | 0.65 | 0.98 | 0.66 | 0.69 | 1.00 | 0.98 | 1.00 |
| CKA ( ) | 0.74 | 0.72 | 0.73 | 0.21 | 0.78 | 0.42 | 0.23 | 1.00 | 0.90 | 0.99 | 0.93 | 0.88 | 0.86 | 0.52 | 0.93 | 0.50 | 0.57 | 1.00 | 0.98 | 1.00 |
| kCKA ( ) | 0.43 | 0.42 | 0.42 | 0.24 | 0.74 | 0.43 | 0.37 | 1.00 | 0.74 | 0.92 | 0.59 | 0.52 | 0.52 | 0.58 | 0.88 | 0.61 | 0.66 | 1.00 | 0.88 | 0.98 |
| Evaluation | AUPRC | Conformity Rate | ||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Dataset | Cora | Flickr | OGBN-Arxiv | Cora | Flickr | OGBN-Arxiv | ||||||||||||||
| Architecture | GCN | SAGE | GAT | PGNN | GCN | SAGE | GAT | GCN | SAGE | GAT | GCN | SAGE | GAT | PGNN | GCN | SAGE | GAT | GCN | SAGE | GAT |
| CKA | 0.67 | 0.82 | 0.80 | 0.42 | 0.41 | 1.00 | 0.37 | 1.00 | 1.00 | 1.00 | 0.85 | 0.92 | 0.95 | 0.62 | 0.78 | 1.00 | 0.75 | 1.00 | 1.00 | 1.00 |
| CKA ( ) | 0.73 | 0.81 | 0.82 | 0.42 | 0.75 | 1.00 | 0.53 | 1.00 | 1.00 | 1.00 | 0.91 | 0.92 | 0.96 | 0.55 | 0.94 | 1.00 | 0.80 | 1.00 | 1.00 | 1.00 |
| CKA ( ) | 0.79 | 0.85 | 0.78 | 0.42 | 0.97 | 1.00 | 0.29 | 1.00 | 1.00 | 1.00 | 0.95 | 0.96 | 0.93 | 0.55 | 0.99 | 1.00 | 0.60 | 1.00 | 1.00 | 1.00 |
| kCKA ( ) | 0.77 | 0.86 | 0.84 | 0.42 | 1.00 | 1.00 | 0.55 | 1.00 | 1.00 | 1.00 | 0.94 | 0.95 | 0.97 | 0.56 | 1.00 | 1.00 | 0.84 | 1.00 | 1.00 | 1.00 |
| Evaluation | AUPRC | Conformity Rate | ||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Dataset | Cora | Flickr | OGBN-Arxiv | Cora | Flickr | OGBN-Arxiv | ||||||||||
| Architecture | GCN | SAGE | GAT | GCN | SAGE | GAT | GCN | SAGE | GCN | SAGE | GAT | GCN | SAGE | GAT | GCN | SAGE |
| CKA | 0.13 | 0.76 | 0.91 | 0.64 | 0.80 | 0.57 | 1.00 | 1.00 | 0.37 | 0.92 | 0.97 | 0.85 | 0.94 | 0.77 | 1.00 | 1.00 |
| CKA ( ) | 0.13 | 0.80 | 0.98 | 0.69 | 0.94 | 0.58 | 1.00 | 1.00 | 0.38 | 0.93 | 0.99 | 0.86 | 0.99 | 0.80 | 1.00 | 1.00 |
| CKA ( ) | 0.13 | 0.93 | 0.93 | 0.75 | 0.42 | 0.55 | 1.00 | 0.99 | 0.38 | 0.99 | 0.97 | 0.90 | 0.52 | 0.84 | 1.00 | 1.00 |
| kCKA ( ) | 0.13 | 0.76 | 0.96 | 0.74 | 0.82 | 0.60 | 1.00 | 0.96 | 0.37 | 0.92 | 0.99 | 0.89 | 0.94 | 0.83 | 1.00 | 0.99 |
| Evaluation | Conformity Rate | Spearman | ||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Dataset | Cora | Flickr | OGBN-Arxiv | Cora | Flickr | OGBN-Arxiv | ||||||||||||||
| Architecture | GCN | SAGE | GAT | PGNN | GCN | SAGE | GAT | GCN | SAGE | GAT | GCN | SAGE | GAT | PGNN | GCN | SAGE | GAT | GCN | SAGE | GAT |
| CKA | 0.99 | 1.00 | 1.00 | 0.99 | 0.81 | 0.99 | 0.63 | 0.88 | 0.70 | 0.90 | 0.99 | 1.00 | 1.00 | 0.99 | 0.73 | 0.99 | 0.48 | 0.86 | 0.58 | 0.91 |
| CKA ( ) | 1.00 | 1.00 | 1.00 | 1.00 | 0.86 | 0.99 | 0.80 | 0.84 | 0.80 | 0.93 | 1.00 | 1.00 | 1.00 | 1.00 | 0.79 | 0.99 | 0.66 | 0.75 | 0.73 | 0.93 |
| CKA ( ) | 1.00 | 1.00 | 1.00 | 1.00 | 0.97 | 0.99 | 0.85 | 0.91 | 0.96 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 0.98 | 0.99 | 0.64 | 0.84 | 0.91 | 1.00 |
| kCKA ( ) | 1.00 | 1.00 | 1.00 | 0.91 | 0.95 | 0.95 | 0.96 | 0.88 | 0.98 | 1.00 | 1.00 | 1.00 | 1.00 | 0.91 | 0.90 | 0.97 | 0.95 | 0.85 | 0.99 | 1.00 |
Kernel Methods
Kernel methods use similarity functions and reproducing kernel spaces to model nonlinear relationships, quantify uncertainty, and compare data.
See all 230 papers