PaperlayerView exact revision on arXiv

Source-linked AI summary

Similarity of Neural Networks with Gradients

Shuai Tang, Wesley J. Maddox, Charlie Dickens, Tom Diethe, Andreas Damianou

arXiv:2003.11498v1cs.LGstat.ML

TL;DR

Comparing neural networks requires a suitable similarity measure and a representation that captures the learned model. This paper combines feature and gradient information in kernel representations and uses sketching with similarity indices to support comparisons across datasets of different sizes. The authors report meaningful similarity scores reflecting relationships among models through their training datasets and learning tasks.

  • Problem

    Comparing trained networks requires an informative representation and similarity index, while existing indices such as CKA or NBS cannot directly handle datasets with different sample counts.

    Source: §1 · p. 1

  • Method

    The method combines learned features and their loss gradients in kernel representations, then compares them with CKA or NBS augmented by sketching.

    Source: §1 · p. 1

  • Results

    The method provides meaningful similarity scores illustrating model relationships based on both training datasets and learning tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Analysis of the kernel mean embedding norm provides insight into the representation and favors CKA over NBS as a similarity index.

    Source: §6 · p. 8

    A scalar Fisher-Rao norm alone is insufficient for comparing models, since equal norms do not necessarily imply that models implement the same function.

    Source: §2 · p. 2

Abstract

from arXiv · show

A suitable similarity index for comparing learnt neural networks plays an important role in understanding the behaviour of the highly-nonlinear functions, and can provide insights on further theoretical analysis and empirical studies. We define two key steps when comparing models: firstly, the representation abstracted from the learnt model, where we propose to leverage both feature vectors and gradient ones (which are largely ignored in prior work) into designing the representation of a neural network. Secondly, we define the employed similarity index which gives desired invariance properties, and we facilitate the chosen ones with sketching techniques for comparing various datasets efficiently. Empirically, we show that the proposed approach provides a state-of-the-art method for computing similarity of neural networks that are trained independently on different datasets and the tasks defined by the datasets.

1. Introduction

The paper targets comparison of neural networks across models and datasets, requiring both an informative model representation and a suitable similarity index. It combines feature and gradient information in kernel representations and augments similarity indices with sketching to support comparisons across datasets of different sizes.

  • The comparison problem involves choosing both a suitable similarity index and an informative representation of each trained network.The motivating applications include transfer learning, meta learning, and imitation learning.Source: §1 · p. 1
  • The method represents a network using both learned feature vectors and loss gradients with respect to those features, combined through a Hadamard product of separate kernel matrices.Feature vectors reflect training data, while gradients reflect the learning task defined by the dataset.Source: §1 · p. 1
  • Sketching augments CKA or NBS to compare representations from datasets of different sizes while making expensive matrix operations more tractable.Sketches project representations to a chosen dimension and approximately preserve their spectral structure.Source: §1 · pp. 1–2
  • With a linear kernel, CKA and NBS provide isotropic-scaling invariance; the proposed approach also seeks invariance to hidden-unit rotations and permutations.Sketching addresses invariance to varying numbers of samples across datasets.Source: §1 · pp. 1–2
  • Empirically, the method provides meaningful similarity scores that illustrate relationships among models through both their training datasets and learning tasks.Source: §1 · p. 2

2. Related Work

Prior work compares networks through feature-based similarity, gradient-derived representations, and scalar measures such as the Fisher-Rao norm. The paper positions its feature-and-gradient kernel representation as a way to compare models while addressing limits in those approaches.

  • Feature-based studies use correlation scores or kernel projections into an RKHS to compare neural-network representations.Raghu et al. (2017) and Morcos et al. (2018) study correlations; Gretton et al. (2005) is cited for kernel-based projection.Source: §2 · p. 2
  • Gradient-derived kernels include the Fisher kernel and Neural Tangent Kernel, but gradient information has seen limited use in network comparison because of computational complexity.The passage also notes applications of gradient information in large-scale learning tasks.Source: §2 · p. 2
  • Task2Vec represents tasks using diagonal Fisher Information matrix terms from a reference model and compares representations with cosine similarity.It requires a reference model, can give unreasonable scores when multiple models are supplied, and is not easily parallelised because gradients are with respect to parameters.Source: §2 · p. 2
  • The Fisher-Rao norm combines feature and gradient information and is described as an indicator of model complexity and generalisation ability, but a scalar value is insufficient for model comparison.Deep linear networks can have the same FR norm across layers while implementing varying functions.Source: §2 · p. 2
  • The proposed representation constructs a kernel matrix from both training data features and gradients that reflect the learning task.The authors report it is considerably more appropriate for comparing models than considering feature vectors alone.Source: §2 · p. 2
  • CKA was reported to convey feature-map relations better than CCA and Kernel CCA, but its standard use is limited to same-dataset comparisons.The paper proposes sketching datasets to a common size before applying CKA.Source: §2 · p. 3

3. Method

The method represents a network on a dataset with a kernel built from both layer features and loss gradients, then compares kernels with CKA or NBS and sketching. This design targets comparisons across models trained on datasets of different sizes.

3.1. Representation

  • At layer l, the representation combines each sample’s feature vector with the loss gradient with respect to that layer’s output in a kernel matrix.The feature vectors encode sample locations, while the gradients indicate how those projected samples should move to minimize loss.Source: §3.1 · p. 3
  • With log-likelihood loss, the expected gradient tends to zero, so the gradient component becomes uninformative; experiments use a smoothed predictive distribution with β = 0.5.The smoothing parameter is set to β = 0.5 in all experiments.Source: §3.1 · p. 3

3.2. Similarity Index

  • CKA and NBS compare the constructed kernel representations; with a linear kernel, both are invariant to isotropic scaling and embedding-space rotation.NBS is described as a normalized Bures metric derived from optimal transport between distributions.Source: §3.2 · pp. 3–4

3.3. Sketching the Kernel Matrix

  • Sketching projects kernel representations to fixed-size summaries, enabling comparisons of different shapes and reducing similarity computation from O(N^3) to O(M^3).Sketching preserves the top principal directions with high probability, but the authors note that theoretical guarantees on similarity or sketch accuracy are difficult to give.Source: §3.3 · pp. 4–5

3.4. Our Method in a Nutshell

  • The method uses sparse CountSketch to hash signed input rows into buckets, then compares the resulting kernel summaries; representation construction needs one forward and backward call per sample.The similarity index takes two sketched representations as input and returns a similarity score.Source: §3.3 · p. 4, §3.4 · p. 5

4. Other Similarity Measures

The section keeps CKA and NBS as similarity indices while considering alternative ways to combine feature and gradient information, and includes Task2Vec and Fisher-Rao norm as comparators.

4. Other Similarity Measures

  • Task2Vec and Fisher-Rao norm are included as closely related comparison methods; the Fisher-Rao norm supplies a scalar representation for each network.The section introduces a distance based on the difference between the two scalar representations.Source: §4 · p. 5
  • The proposed alternatives change how feature and gradient information is merged while retaining CKA or NBS; the resulting three combinations remain positive semidefinite kernel matrices.One listed combination is K = (1/2 I_N + 1/2 T_fg)K_f(1/2 I_N + 1/2 T_fg).Source: §4 · p. 5

5. Experimental settings and Analysis

Across independently trained models, combining features with gradients makes CKA and NBS produce similarity scores matching the stated task-based expectations. Residual-block comparisons and kernel-mean-embedding analyses further distinguish dataset relationships, while the embedding norm also supports reference-model comparisons.

5. Experimental settings and Analysis

  • Experiments use 33 ResNets trained to convergence across CIFAR-10, CIFAR-100, and SVHN, with comparisons computed at residual-block outputs.The set comprises ten independently initialized ResNet-18s and one ResNet-50 per dataset.Source: §5 · p. 5

5.1. Similarity of Independently Trained Models

  • Combining features and gradients lets both CKA and NBS match the three stated similarity expectations for independently trained models.Feature-only comparisons give high scores for models trained on the same dataset but low and uninterpretable scores across datasets.Source: §5.1 · p. 6
  • Task2Vec gives uninformative scores for independently trained models, while Fisher-Rao Distance fails the stated desiderata and encodes little dataset information.Task2Vec requires a reference model to compute each task’s Fisher information matrix.Source: §5.1 · p. 6

5.2. Similarity between Individual Residual Blocks

  • The norm-based comparison readily separates SVHN when the base model was trained on natural images, unlike the other two quantities.Source: §5.2 · p. 7
  • CIFAR-10 and CIFAR-100 show similar layerwise patterns, while CIFAR-100 layers remain highly similar to layers before ResBlock8 in CIFAR-10 and SVHN models.ResBlocks 9–10 of SVHN models are dissimilar to those blocks in CIFAR models; the authors suggest harder-task functions may transfer to easier tasks and caution against over-interpretation.Source: §5.2 · p. 7

5.3. Norm of Kernel Mean Embedding

  • With a CIFAR-10-trained model, the embedding norm rates SVHN higher than CIFAR-100, whereas the CKA normalizer struggles and the NBS normalizer fails.This matches the authors’ observation that CIFAR-100 is an easier fine-tuning target than SVHN for a CIFAR-10 model.Source: §5.3 · p. 8
  • The CKA-normalized embedding norm gives higher values when source and target datasets match, while the NBS-normalized norm consistently assigns high values to SVHN.In these experiments, the authors find CKA gives a more reasonable similarity illustration than NBS.Source: §5.3 · p. 8

5.4. A Comparison Method Given a Reference Model

  • The embedding norm can serve as a similarity score between a reference model and a new dataset, with lower time and space complexity than Task2Vec.Its sketch-defined representation size is fixed, and its gradient computation can be parallelized in batch mode.Source: §5.4 · p. 8

6. Conclusion

The conclusion presents a representation combining feature and gradient vectors, compared through sketched similarity indices, and reports that embedding-norm analysis favors CKA over NBS. The sketch appendix describes CountSketch and the theoretical and practical properties of the transforms used.

6. Conclusion

  • The proposed kernel representation combines feature and gradient vectors, while sketching enables comparisons between tasks with different sample counts.The conclusion says the representation informs similarity with both the input-data distribution and model construction.Source: §6 · p. 8
  • Analysis of the predictive-distribution kernel mean embedding norm provides insight into the representation and favors CKA over NBS.Source: §6 · p. 8
  • For future work, the authors propose exploring the method as a recommender system for large collections of pretrained models and query datasets.Source: §6 · p. 9

A. Properties of Sketches

  • CountSketch assigns each input column uniformly at random to a row and gives it a random sign; despite weak dependence of its failure probability on δ, streaming and hashing can be faster than other sketches for dense inputs.Source: A. Properties of Sketches · p. 11
  • The authors use CWT in practice for speed although it lacks a worst-case JLT guarantee; their stated reason is that it preserves directions spanning the input matrix's column space.The stated general bounds apply to JLT matrices, while CWT preserves norms for points in a subspace rather than arbitrary point sets.Source: A. Properties of Sketches · p. 11

B. Training details of Models

The appendix records the model-training setup and sketching hyperparameter choices, including the tested sketch sizes and the setting used for reported results.

B. Training details of Models

C. Hyperparameters in Sketching

  • The authors tested sketch sizes M = 512 and M = 128, found no significant difference, and report main results with M = 512.They describe the feature-map matrix for neural-network classification as low-rank and say sketching batch size did not affect results.Source: C. Hyperparameters in Sketching · pp. 11–12

D. Proofs

This appendix section introduces proofs of inequalities and equalities from the main paper, including a proof of the inequality in Eq. 14.

D. Proofs

  • The appendix provides short proofs of inequalities and equalities stated in the main paper.Source: D. Proofs · p. 12

D.1. Proof of the inequality in Eq. 14

  • The Eq. 14 proof invokes the Araki-Lieb-Thirring inequality for 0 ≤ r ≤ 1, q ≥ 0, and positive semidefinite A and B.Source: §D.1 · p. 12
  • For Eq. 14, the proof sets B = K1, r = 1/2, and q = 1, then derives the first inequality.Source: §D.1 · p. 12

D.2. Proof of the inequality in Eq. 18

The appendix derives a bound on the kernel mean embedding norm, shows how to compute that norm from the kernel matrix, and proves a sketching bound for transformed representations.

D.2. Proof of the inequality in Eq. 18

  • Jensen's inequality yields the bound ||µp|| ≤ ||θ||fr/(L + 1).Source: §D.2 · p. 12

D.3. Proof of the equality in Eq. 5.3

  • Once the kernel matrix is constructed, the kernel mean embedding norm can be computed efficiently.Source: §D.3 · p. 12

D.4. Proof of the bound in Eq. 15

  • For JLT S and matrices Ψ1 and Ψ2, the proof bounds inner products between sketched columns by the original inner products with ϵ-scaled norm terms.The derivation sums over the columns of both matrices.Source: §D.4 · pp. 12–13

D.5. Proof of the bound in Eq. 16

The Sketched CKA deviation is bounded relative to the original CKA score. The accompanying figures compare representations using CKA and NBS, including feature-only, gradient-only, and combined feature-and-gradient representations.

D.5. Proof of the bound in Eq. 16

  • The absolute difference between Sketched CKA and original CKA is bounded by max{4ε/(1+ε)^2, 4ε/(1−ε)^2} times original CKA, which simplifies to 4ε/(1−ε)^2 times original CKA (46).Source: §D.5 · p. 13

E.1. Model Comparison with different representations

  • Under CKA, the proposed representation gives an illustration matching the stated desiderata; the NBS figure compares models using different representations.Source: Figure 7 · p. 14, Figure 8 · p. 15

E.2.1. FEATURE-ONLY REPRESENTATION

  • The feature-only CKA figure compares every pair of ResNet-18 residual blocks across three chosen tasks.Source: Figure 9 · p. 16

E.2.2. GRADIENT-ONLY REPRESENTATION

  • The gradient-only CKA figure makes the same residual-block comparison for ResNet-18 across three chosen tasks.Source: Figure 10 · p. 17

E.2.3. OUR REPRESENTATION WITH FEATURES AND GRADIENTS

F. Pretrained ImageNet models as the base model

The appendix examines a scalar reference for how much a pretrained base model may need to adapt to a dataset. Across pretrained ImageNet models and layers, the scalar distinguishes task values, with SVHN higher than CIFAR10 and CIFAR100.

F. Pretrained ImageNet models as the base model

G. Kernel Ridge Regression as Classification

The section explores classification without fine-tuning by applying kernel ridge regression to sketched data and labels from a pretrained model. Performance improves for layers closer to the output, and the method performs better overall than one-epoch last-layer fine-tuning, though it remains below fine-tuning to convergence.

G. Kernel Ridge Regression as Classification

Loading 2003.11498v1…