PaperlayerView exact revision on arXiv

Source-linked AI summary

Reliability of CKA as a Similarity Measure in Deep Learning

MohammadReza Davari, Stefan Horoi, Amine Natik, Guillaume Lajoie, Guy Wolf, Eugene Belilovsky

arXiv:2210.16156v2cs.LGcs.AIcs.CV

TL;DR

Comparing neural representations is difficult, and linear CKA is widely used for that purpose. This paper characterizes how simple transformations affect CKA and tests ways to manipulate its values. It finds that CKA can change substantially while model behavior remains largely intact, motivating caution in its use.

  • Problem

    Comparing high-dimensional neural representations is difficult, yet CKA has become widely used to draw conclusions about their similarity.

    Source: §1 · pp. 1–2

  • Method

    The paper characterizes CKA’s sensitivity to simple transformations and empirically tests cases where CKA gives unexpected results or can be manipulated while preserving model behavior.

    Source: §1 · p. 2, §4 · p. 5

  • Results

    CKA can be manipulated to show strong similarity between early and later layers, or between all layers, with network performance barely hindered; the analysis also explains sensitivity to outliers and high-variance directions.

    Source: §3 · p. 5, §4.3 · p. 8, §5 · p. 9

  • Takeaways & Limitations

    The authors call for caution when using linear CKA, especially when model production is unknown, inconsistent, or uncontrolled.

    Source: §5 · p. 9

    The authors note that some problematic transformations may not occur in many applications, and their empirical subset translations were manually added for controlled experiments.

    Source: §4.2 · p. 7, §5 · p. 9

Abstract

from arXiv · show

Comparing learned neural representations in neural networks is a challenging but important problem, which has been approached in different ways. The Centered Kernel Alignment (CKA) similarity metric, particularly its linear variant, has recently become a popular approach and has been widely used to compare representations of a network's different layers, of architecturally similar networks trained differently, or of models with different architectures trained on the same data. A wide variety of conclusions about similarity and dissimilarity of these various representations have been made using CKA. In this work we present analysis that formally characterizes CKA sensitivity to a large class of simple transformations, which can naturally occur in the context of modern machine learning. This provides a concrete explanation of CKA sensitivity to outliers, which has been observed in past works, and to transformations that preserve the linear separability of the data, an important generalization attribute. We empirically investigate several weaknesses of the CKA similarity metric, demonstrating situations in which it gives unexpected or counter-intuitive results. Finally we study approaches for modifying representations to maintain functional behaviour while changing the CKA value. Our results illustrate that, in many cases, the CKA value can be easily manipulated without substantial changes to the functional behaviour of the models, and call for caution when leveraging activation alignment metrics.

1 INTRODUCTION

Comparing high-dimensional neural representations is difficult, and linear CKA has become a widely used way to compare them. The paper examines whether CKA remains reliable under simple transformations and shows that its values can change substantially without noticeable changes in model behavior.

1 INTRODUCTION

  • High-dimensional internal representations make comparing neural networks fundamentally difficult, motivating measures such as CKA.Linear CKA has become the preferred measure after prior work found no real benefit from nonlinear kernels in its considered experiments.Source: §1 · p. 1
  • The paper characterizes how simple, naturally occurring transformations can change CKA, including transformations that preserve function or linear separability.It also empirically demonstrates that CKA can be manipulated to high or low values without significant changes to model behavior.Source: §1 · p. 2
  • The authors revisit findings about CKA and explain its sensitivity to outliers and transformations that preserve linear separability, both relevant to representation comparisons.They also connect their theoretical analysis to empirical tests of CKA reliability.Source: §1 · p. 2

2 BACKGROUND ON CKA AND RELATED WORK

This section introduces CKA as an alignment of centered similarity structures and situates it among representation-comparison methods. It highlights prior applications and the paper’s focus on linear CKA and its sensitivity.

2 BACKGROUND ON CKA AND RELATED WORK

  • CKA compares centered similarity structures through HSIC; linear CKA uses inner-product kernels, with K = XX⊤ and L = YY⊤.The section describes HSIC as the trace of the centered kernel matrices’ product, normalized by (n−1)^2.Source: §2 · p. 2
  • Prior studies used CKA to compare network depth and width, transformer and CNN representations, and layers implicated in forgetting.The section also reports that CKA has been used to study initialization effects and similarity among later layers in high-capacity models.Source: §2 · p. 3
  • Earlier work reported CKA’s insensitivity to removing low-variance components and sensitivity to dominant data points; this paper adds theoretical explanations and arbitrary-value manipulation.Ding et al. (2021) linked removal of low-variance components to reduced probing accuracy, while Nguyen et al. (2022) associated block structure with a few dominant points.Source: §2 · p. 3
  • Because linear CKA is the form most often used in practice, the paper focuses its theoretical analysis on that variant, while noting empirical problems with RBF-kernel CKA.The authors discuss a possible extension of the theory to nonlinear kernels.Source: §2 · p. 3

3 CKA SENSITIVITY TO SUBSET TRANSLATION

The paper analyzes subset translations of representations and shows that they can lower linear CKA, including when they preserve functionally relevant properties. The results explain sensitivity to outliers and show that CKA can understate similarity even when linear separability and margins are preserved.

3 CKA SENSITIVITY TO SUBSET TRANSLATION

  • Subset translations can significantly reduce CKA between an original representation set and its transformed copy, regardless of class and direction except in Corollary 4.The transformation shifts points outside a chosen subset S by c in direction v; the figure illustrates the general case, a single-point case, and a separability-preserving case.Source: §3 · pp. 3–4
  • The theorem’s expression depends on the unchanged-point proportion and the participation ratio, an effective dimensionality estimate for the representations.The section notes that neural representations’ effective dimensionality is often far below their actual dimension, and that weight decay biases relevant norms toward small values.Source: §3 · pp. 3–4
  • As an outlier’s displacement grows, CKA drops sharply even though the representation sets otherwise match.For a single outlier, the small unshifted fraction makes Γ(ρ) approximately 1 divided by the number of examples, which can dominate the expression when the participation ratio is small.Source: §3 · pp. 4–5
  • A subset can be translated while preserving linear separability by the same hyperplane and maintaining the same margins, yet CKA can decrease.The paper frames these properties as relevant to generalization and predictive behavior.Source: §3 · p. 5
  • The authors leave a theoretical analysis of nonlinear CKA outside the paper’s scope, but hypothesize that suitable transformations in an implicit feature space could extend the results.They note that finding such transformations would be harder than applying simple translations in representation space.Source: §3 · p. 5

4 EXPERIMENTS AND RESULTS

The experiments test whether CKA reflects feature similarity and usefulness, and whether its values can be manipulated without noticeably changing model behavior. They find counterintuitive comparisons and show that CKA maps can be substantially changed while performance is largely maintained.

4.1 CKA EARLY LAYER RESULTS

  • Early layers show relatively high CKA across generalized, memorized, and random networks, yet their filters differ and are not equally useful.The authors use linear probes on original-task training data to quantify representation usefulness.Source: §4.1 · pp. 5–6, Figure 2 · p. 6, Figure 3 · p. 6

4.2 PRACTICAL IMPLICATIONS OF THEORETICAL RESULTS

  • Artificial subset translations drop linear and RBF CKA below 0.2 despite preserved geometry; on CIFAR10, class shifts drive CKA to zero despite >90% separability, and one outlier can strongly alter it.RBF CKA with σ equal to 0.2 times the median distance is unperturbed by the artificial translation, but is described as less informative.Source: Figure 4 · p. 6, §4.2 · p. 6, Figure 6 · p. 7

4.3 EXPLICITLY OPTIMIZING THE CKA MAP

  • CKA maps can be reshaped to boost early-to-late similarity or mimic thin-network maps with barely hindered performance; optimization can displace a subset of one class.Vision transformers show similar results for mapping wider-network CKA maps to thin-network maps.Source: Figure 5 · p. 7, §4.3 · pp. 8–9, Figure 7 · p. 8

5 DISCUSSION AND CONCLUSION

The paper concludes that linear CKA can be sensitive to transformations that preserve important functional characteristics, and its values can be manipulated without changing model behavior. It calls for caution in uncontrolled comparisons and recommends using multiple similarity measures alongside interpretable methods and visualizations.

5 DISCUSSION AND CONCLUSION

  • CKA can assign low similarity to representations linked by affine transformations, high similarity to representations from unlike networks, and arbitrarily low or high values despite preserved functional behavior.The theoretical analysis also identifies sensitivity to outliers and high-variance directions, and cases where linear separability and shared margins do not prevent CKA from diminishing.Source: §5 · p. 9
  • The authors call for caution, especially when model-production procedures are unknown, inconsistent, or uncontrolled, while noting some problematic transformations may be uncommon.They emphasize that a single scalar in [0, 1] condenses rich geometric information and that more work is needed to relate similarity measures to transformations and model behavior.Source: §5 · p. 9
  • They recommend combining similarity measures with interpretable comparisons, including orthogonal Procrustes, activation sparsity, and data visualizations.The authors suggest relating results to the specific notion of similarity each measure quantifies.Source: §5 · p. 9

A EXPERIMENTAL DETAILS

The appendix estimates CKA by averaging minibatch HSIC values, using linear or RBF kernels. Its stated estimator uses unbiased HSIC and is described as independent of batch size.

A.1 MINIBATCH CKA

  • CKA is estimated by averaging HSIC values across minibatches, with linear and RBF kernels used in experiments.Each minibatch contains samples from two network layers’ representations.Source: §A.1 · p. 15
  • The compared minibatches contain m samples, while the two layers may have different neuron counts, n1 and n2.The representations are denoted Xb ∈ R^(m×n1) and Yb ∈ R^(m×n2).Source: §A.1 · p. 15
  • The minibatch estimator uses the unbiased HSIC estimator HSIC1 and is stated to be independent of batch size.The minibatch kernel matrices are formed from each layer’s activations.Source: §A.1 · p. 15

A.2 NETWORK ARCHITECTURE

The appendix describes a nine-layer classifier with eight convolutional layers and a final fully connected layer. It also specifies convolution settings and a wider ResNet-34 configuration for separate experiments.

A.2 NETWORK ARCHITECTURE

  • The Sec. 4.1 classifier has nine layers: eight convolutional layers followed by a fully connected classification layer.ReLU is used throughout, and each convolutional layer is followed by batch normalization.Source: §A.2 · p. 15
  • The Sec. 4.1 convolution kernels are 7 × 7 in the first two layers and 3 × 3 in the others, with zero padding and stride one.The listed kernel counts from lower to higher layers are [16, 16, 32, 32, 32, 64, 64].Source: §A.2 · p. 15
  • The wider-network experiments use ResNet-34 with channels scaled up; the network for Figures 5 and 6 instead uses 3 × 3 kernels in every layer.Source: §A.2 · p. 15

A.3 TRAINING DETAILS

The appendix reports training schedules for the experiments and a validation-accuracy rule for dynamically balancing distillation and CKA-map losses. The rule adjusts λ according to the accuracy gap from the original model.

A.3 TRAINING DETAILS

  • The generalized and memorized Sec. 4.1 models were trained for 100 epochs with AdamW, initial learning rate 1e-3, weight decay 5e-4, and cosine scheduling.Source: §A.3 · p. 15
  • The Sec. 4.3 base model was trained for 200 epochs at initial learning rate 0.01, while target-map models used 30 epochs and learning rate 1e-3.Target-map training used initial λ = 500, dynamically scaled by 0.8 per iteration, a 500-step warm-up, and distillation temperature 0.2.Source: §A.3 · p. 15
  • For the wider-network experiments, original models were trained for 100 epochs and target-map-optimized models for 15 epochs.Source: §A.3 · p. 15

A.4 CKA MAP LOSS BALANCE

  • λ is scaled down when validation accuracy falls more than the threshold below original accuracy, and scaled up otherwise to encourage closer agreement with the target CKA maps.Validation accuracy serves as a surrogate for preservation of the network’s representations.Source: §A.4 · pp. 15–16

B ADDITIONAL RESULTS

The invertible-matrix experiment shows that linear CKA can drop sharply under transformations expected to preserve representation relationships. Whether invariance to such transformations is desirable depends on the data-size and layer-width setting.

B.1 CKA SENSITIVITY TO INVERTIBLE LINEAR TRANSFORMATIONS

  • For small Gaussian-matrix parameters, linear CKA between X and Y = XM can drop to 0 even though M is invertible.The figure plots CKA against the matrix-element mean and standard deviation across 10 random instantiations.Source: §B.1 · p. 16, Figure 9 · p. 16
  • Invariance to invertible linear transformations may be desirable when dataset size exceeds layer width, as in smaller networks or final layers, but not in every context.Source: §B.1 · p. 16

B.2 WIDER NETWORKS

Across these additional experiments, CKA maps and values can be manipulated in several settings, including under distribution shift and with a Vision Transformer. Some manipulations preserve performance, while high CKA alone does not guarantee similarly useful features.

B.2 WIDER NETWORKS

  • Optimized ResNet-34 CKA maps continued to resemble the thin-network target under the CIFAR10-to-Patch Camelyon shift, with similar findings for VGG-style networks.The maps were optimized using CIFAR10 and evaluated on Patch Camelyon; the VGG-style experiment broadened the architecture scope.Source: §B.2 · pp. 16–17, Figure 10 · p. 17, Figure 11 · p. 17

B.3 RBF KERNEL

  • RBF kernel CKA values can also be manipulated across additional bandwidths when the values are meaningful.Source: §B.3 · p. 18

B.4 CKA MAP OPTIMIZATION VIA LOGISTIC LOSS

  • Replacing distillation loss with cross-entropy loss produced similar CKA-map manipulation results, including strong early-to-late layer similarity with only slight performance loss.The experiments maximize similarity between the first and last layer, maximize similarity across all layers, or minimize it across all layers.Source: §B.4 · p. 18, Figure 12 · p. 18, Figure 13 · p. 19

B.5 CKA OPTIMIZATION OF VIT

  • An eight-block CIFAR10 Vision Transformer achieved the specified CKA-map manipulations with minimal performance loss, supporting the approach's model-agnostic scope.The manipulations target first-to-last block similarity, all-block similarity maximization, or all-block similarity minimization.Source: §B.5 · p. 19, Figure 14 · p. 19, Table 1 · p. 19

B.6 CLOSER LOOK AT THE EARLY LAYERS

  • Generalized, memorized, and randomly populated networks showed relatively high early-layer CKA across same- and different-seed comparisons, but high CKA did not imply similarly useful features.Feature usefulness was assessed with linear-probe accuracy on the CIFAR10 test set.Source: §B.6 · p. 20, Figure 15 · p. 20, Figure 16 · p. 20

B.7 OUT OF DISTRIBUTION EVALUATION

Models whose CKA maps were modified while maintaining test accuracy also showed effectively unchanged average performance on CIFAR-10-C corruptions, despite being trained only on uncorrupted data.

B.7 OUT OF DISTRIBUTION EVALUATION

  • Average out-of-distribution performance on CIFAR-10-C remained effectively unchanged after CKA-map manipulation, although the models were trained only on uncorrupted data.Source: §B.7 · p. 21

C PROOFS

The proofs analyze a representation set in which a subset of points is translated, then center the modified set and express linear CKA through HSIC terms. The derivation uses variance and participation ratio to simplify the limiting expression.

C PROOFS

  • The proof partitions X into unchanged C1 and translated C2 points, forming the modified set by shifting C2 by c in direction v.Source: C PROOFS · p. 22
  • Because the modified representation set is not centered, the proof computes its mean and subtracts it before evaluating linear CKA through three HSIC terms.Source: C PROOFS · p. 22
  • The derivation relates total representation variance to eigenvalues and uses participation ratio as an effective-dimensionality estimate in the limiting expression.Source: C PROOFS · pp. 26–27
Loading 2210.16156v2…