Original: O N THE I NADEQUACY OF CKA AS A M EASURE OF S IMILARITY IN D EEP L EARNING
Revised: O N THE I NADEQUACY OF CKA AS A M EASURE OF S IMILARITY IN D EEP L EARNING

Original: A BSTRACT
Revised: A BSTRACT

Original: Comparing learned representations is a challenging problem which has been approached in different ways.
Revised: Comparing learned representations is a challenging problem which has been approached in different ways.

Original: The CKA similarity metric, particularly it’s linear variant, has recently become a popular approach and has been widely used to compare representations of a network’s different layers, of similar networks trained differently or of models with different architectures trained on the same data.
Revised: The CKA similarity metric, particularly it’s linear variant, has recently become a popular approach and has been widely used to compare representations of a network’s different layers, of similar networks trained differently or of models with different architectures trained on the same data.

Original: CKA results have been used to make a wide variety of claims about similarity and dissimilarity of these various representations.
Revised: CKA results have been used to make a wide variety of claims about similarity and dissimilarity of these various representations.

Original: In this work we investigate several weaknesses of the CKA similarity metric, demonstrating situations in which it gives unexpected or counterintuitive results.
Revised: In this work we investigate several weaknesses of the CKA similarity metric, demonstrating situations in which it gives unexpected or counterintuitive results.

Original: We then study approaches for modifying representations to maintain functional behaviour while changing the CKA value.
Revised: We then study approaches for modifying representations to maintain functional behaviour while changing the CKA value.

Original: Indeed we illustrate in some cases the CKA value can be heavily manipulated without substantial changes to the functional behaviour
Revised: Indeed we illustrate in some cases the CKA value can be heavily manipulated without substantial changes to the functional behaviour

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: In the last decade, increasingly complex deep learning models have taken over machine learning and have helped us solve, with remarkable accuracy, a multitude of tasks across a wide array of domains.
Revised: In the last decade, increasingly complex deep learning models have taken over machine learning and have helped us solve, with remarkable accuracy, a multitude of tasks across a wide array of domains.

Original: Due to the size and flexibility of these models it has been particularly hard to study and understand exactly how they solve the tasks we use them on.
Revised: Due to the size and flexibility of these models it has been particularly hard to study and understand exactly how they solve the tasks we use them on.

Original: A helpful framework for thinking about these models is that of representation learning , where we view artificial neural networks (ANNs) as learning internal representations of the input data that are progressively more expressive in terms of the input’s semantic information as we go deeper into the ANN layers.
Revised: A helpful framework for thinking about these models is that of representation learning , where we view artificial neural networks (ANNs) as learning internal representations of the input data that are progressively more expressive in terms of the input’s semantic information as we go deeper into the ANN layers.

Original: In practice, it is often of interest to analyze and compare the representations of multiple ANNs, however, the typically high dimensionality of ANN internal representation spaces makes this a fundamentally difficult task.
Revised: In practice, it is often of interest to analyze and compare the representations of multiple ANNs, however, the typically high dimensionality of ANN internal representation spaces makes this a fundamentally difficult task.

Original: To address this problem, the machine learning community has tried finding meaningful ways to compare ANN internal representations and various representation (dis)similarity measures have been proposed.
Revised: To address this problem, the machine learning community has tried finding meaningful ways to compare ANN internal representations and various representation (dis)similarity measures have been proposed.

Original: Recently, Centered Kernel Alignment (CKA) (Kornblith et al., 2019) was proposed and shown to perform significantly better on representation similarity “sanity checks” than past methods such as linear regression or CCA based methods (Raghu et al., 2017; Morcos et al., 2018).
Revised: Recently, Centered Kernel Alignment (CKA) (Kornblith et al., 2019) was proposed and shown to perform significantly better on representation similarity “sanity checks” than past methods such as linear regression or CCA based methods (Raghu et al., 2017; Morcos et al., 2018).

Original: While
Revised: While

Original: CKA can capture different notions of similarity between points in representation space by using different kernel functions, it was empirically shown in the original paper that there are no real benefits in using CKA with a non-linear kernel over it’s linear counterpart.
Revised: CKA can capture different notions of similarity between points in representation space by using different kernel functions, it was empirically shown in the original paper that there are no real benefits in using CKA with a non-linear kernel over it’s linear counterpart.

Original: As a result, linear CKA has been the preferred representation similarity measure of the machine learning community in the past three years and all other similarity measures have been, for all practical purposes, retired.
Revised: As a result, linear CKA has been the preferred representation similarity measure of the machine learning community in the past three years and all other similarity measures have been, for all practical purposes, retired.

Original: CKA has been utilized in a number of works to make conclusions regarding the similarity between different models and their behaviours such as wide versus deep ANNs (Nguyen et al., 2021) and transformer versus CNN based ANNs (Raghu et al., 2021).
Revised: CKA has been utilized in a number of works to make conclusions regarding the similarity between different models and their behaviours such as wide versus deep ANNs (Nguyen et al., 2021) and transformer versus CNN based ANNs (Raghu et al., 2021).

Original: They have also been used to draw conclusions about transfer learning (Neyshabur et al., 2020) and catastrophic forgetting (Ramasesh et al., 2021).
Revised: They have also been used to draw conclusions about transfer learning (Neyshabur et al., 2020) and catastrophic forgetting (Ramasesh et al., 2021).

Original: However, very recent work has questioned CKA’s supremacy as a representation similarity measure by highlighting some of its shortcomings.
Revised: However, very recent work has questioned CKA’s supremacy as a representation similarity measure by highlighting some of its shortcomings.

Original: In particular, Ding et al.
Revised: In particular, Ding et al.

Original: (2021) demonstrated that CKA lacks sensitivity to changes that affect functional behaviour being unable to detect removal of principal components from the analyzed representations even when this removal significantly decreases probing accuracy.
Revised: (2021) demonstrated that CKA lacks sensitivity to changes that affect functional behaviour being unable to detect removal of principal components from the analyzed representations even when this removal significantly decreases probing accuracy.

Original: Williams et al.
Revised: Williams et al.

Original: (2021) discussed how CKA does not respect the triangle inequality, which makes it problematic to use CKA values as a similarity measure in downstream analysis
Revised: (2021) discussed how CKA does not respect the triangle inequality, which makes it problematic to use CKA values as a similarity measure in downstream analysis

Original: [Figure]
Revised: [Figure]

Original: C K A
Revised: C K A

Original: CKA vs. Depth
Revised: CKA vs. Depth

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: tasks.
Revised: tasks.

Original: In this work, we further explore the inadequacies of linear CKA as a similarity measure in deep learning and showcase scenarios in which it behaves in undesirable ways, giving either high similarity measures in scenarios where the similarity should be low or the opposite.
Revised: In this work, we further explore the inadequacies of linear CKA as a similarity measure in deep learning and showcase scenarios in which it behaves in undesirable ways, giving either high similarity measures in scenarios where the similarity should be low or the opposite.

Original: We also present cases where the CKA value can be heavily manipulated to be either high or low without significant changes to the functional behaviour of the underlying ANNs.
Revised: We also present cases where the CKA value can be heavily manipulated to be either high or low without significant changes to the functional behaviour of the underlying ANNs.

Original:  P ROBLEM S TATEMENT AND B ACKGROUND ON CKA
Revised:  P ROBLEM S TATEMENT AND B ACKGROUND ON CKA

Original: Let X ∈ R n × d 1 denote a set of ANN internal representations, i.e. the neural activations of a specific layer with d 1 neurons in a network, in response to n ∈ N input examples.
Revised: Let X ∈ R n × d 1 denote a set of ANN internal representations, i.e. the neural activations of a specific layer with d 1 neurons in a network, in response to n ∈ N input examples.

Original: Let Y ∈ R n × d 2 be another set of such representations generated by the same input examples but possibly at a different layer of the same, or different, deep learning model.
Revised: Let Y ∈ R n × d 2 be another set of such representations generated by the same input examples but possibly at a different layer of the same, or different, deep learning model.

Original: It is standard practice to center these representations column-wise (feature or “neuron” wise) before analyzing them.
Revised: It is standard practice to center these representations column-wise (feature or “neuron” wise) before analyzing them.

Original: We are interested in representation similarity measures which try to capture a certain notion of similarity between X and Y .
Revised: We are interested in representation similarity measures which try to capture a certain notion of similarity between X and Y .

Original: CKA is one such similarity measure based on the Hilbert-Schmidt Independence Criterion (HSIC)
Revised: CKA is one such similarity measure based on the Hilbert-Schmidt Independence Criterion (HSIC)

Original: (Gretton et al., 2005) that was presented as a means to evaluate independence between random variables in a non-parametric way.
Revised: (Gretton et al., 2005) that was presented as a means to evaluate independence between random variables in a non-parametric way.

Original: In the linear case it takes the form:
Revised: In the linear case it takes the form:

Original: [Equation]
Revised: [Equation]

Original: Where subscript F denotes the Frobenius norm of a matrix.
Revised: Where subscript F denotes the Frobenius norm of a matrix.

Original: Linear CKA can then be computed as:
Revised: Linear CKA can then be computed as:

Original: [Equation]
Revised: [Equation]

Original: Intuitively, HSIC computes the similarity structures of X and Y through the Gram matrices of internal products ( XX ⊤ and Y Y ⊤ ) and then compares these similarity structures by computing their alignment through the trace of XX ⊤ Y Y ⊤ .
Revised: Intuitively, HSIC computes the similarity structures of X and Y through the Gram matrices of internal products ( XX ⊤ and Y Y ⊤ ) and then compares these similarity structures by computing their alignment through the trace of XX ⊤ Y Y ⊤ .

Original:  E XPERIMENTS AND R ESULTS 3. CKA E ARLY L AYER R ESULTS
Revised:  E XPERIMENTS AND R ESULTS 3. CKA E ARLY L AYER R ESULTS

Original: CKA values are often treated as a surrogate metric to measure the usefulness and similarity of a network’s learned features when compared to another network (Ramasesh et al., 2021).
Revised: CKA values are often treated as a surrogate metric to measure the usefulness and similarity of a network’s learned features when compared to another network (Ramasesh et al., 2021).

Original: In order to analyze this common assumption, we compare the features of: (1) a network trained to generalize on the CIFAR10 image classification task (Krizhevsky et al., 2009), (2) a network trained to “memorize” the CIFAR10 images (i.e.permuted target labels), and (3) an untrained randomly initialized network (for network architecture and training details see the Appendix).
Revised: In order to analyze this common assumption, we compare the features of: (1) a network trained to generalize on the CIFAR10 image classification task (Krizhevsky et al., 2009), (2) a network trained to “memorize” the CIFAR10 images (i.e. target labels are random), and (3) an untrained randomly initialized network (for network architecture and training details see the Appendix).

Original: As show in Fig.
Revised: As show in Fig.

Original: 1, early layers of these networks should have very similar representations given the high CKA values.
Revised: 1, early layers of these networks should have very similar representations given the high CKA values.

Original: Under the previously presented assumption, one should therefore conclude that the learned features at these layers are relatively similar and equally valuable.
Revised: Under the previously presented assumption, one should therefore conclude that the learned features at these layers are relatively similar and equally valuable.

Original: However this is not the case, we can see in Fig.
Revised: However this is not the case, we can see in Fig.

Original: 2 that the convolution filters are drastically different across the three networks.
Revised: 2 that the convolution filters are drastically different across the three networks.

Original: Moreover,
Revised: Moreover,

Original: Fig.
Revised: Fig.

Original: 2 elucidates that considerably high CKA similarity values for early layers, does not necessarily translate to more useful, or similar, captured features.
Revised: 2 elucidates that considerably high CKA similarity values for early layers, does not necessarily translate to more useful, or similar, captured features.

Original: 3. S ENSITIVITY TO AFFINE TRANSFORMATIONS
Revised: 3. S ENSITIVITY TO AFFINE TRANSFORMATIONS

Original: We consider artificially generated representations X ∈ R n × d to which we apply affine transformations to obtain Y ∈ R n × d .
Revised: We consider artificially generated representations X ∈ R n × d to which we apply affine transformations to obtain Y ∈ R n × d .

Original: We generate X by sampling 10K points  from the1K-dimensional unit
Revised: We generate X by sampling 10K points uniformly from the

Revised: 1K-dimensional unit centered at the and 10K points a similar cube centered at (1 . 1 , 0 , 0 , . . .

Revised: so the points from the two cubes are linearly separable along the first

Original: The first transformation we consider is translating the subset of representations sampled from the second cube by c ∈ R in direction v ∈ R d sampled from the the d -dimensional ball, we plot the CKA values between X and Y as a function of c in Fig. 3a).
Revised: The first transformation we consider is translating the subset of representations sampled from the second cube by c ∈ R in direction v ∈ R d sampled from the the d -dimensional ball, we plot the CKA values between X and Y as a function of c in Fig. 3a).

Original: This transformation entirely preserves the topological structure of the representations as well as their local geometry since the points sampled from each cube have not moved with respect to the other points sampled from the same cube and the two cubes are still separated, only the distance between them has been changed.
Revised: This transformation entirely preserves the topological structure of the representations as well as their local geometry since the points sampled from each cube have not moved with respect to the other points sampled from the same cube and the two cubes are still separated, only the distance between them has been changed.

Original: Despite these multiple notions of “similarity” between X and Y being preserved, the CKA values quickly drop below 0.2 even for relatively small translation distances.
Revised: Despite these multiple notions of “similarity” between X and Y being preserved, the CKA values quickly drop below 0.2 even for relatively small translation distances.

Original: The second transformation considered is a multiplication by a matrix M ∈ R d × d whose elements are sampled from a Gaussian with mean µ and standard deviation σ .
Revised: The second transformation considered is a multiplication by a matrix M ∈ R d × d whose elements are sampled from a Gaussian with mean µ and standard deviation σ .

Original: We verify the invertibility of M since it is not guaranteed and only keep invertible matrices.
Revised: We verify the invertibility of M since it is not guaranteed and only keep invertible matrices.

Original: We show the CKA values between X and the transformed Y in Fig.
Revised: We show the CKA values between X and the transformed Y in Fig.

Original: 3b).
Revised: 3b).

Original: Since this is a invertible linear transformation we would expect it to only modestly change the representations in X and the CKA value to be only slightly lower than 1.
Revised: Since this is a invertible linear transformation we would expect it to only modestly change the representations in X and the CKA value to be only slightly lower than 1.

Original: However, we observe that even for small values of µ and σ , CKA drops to 0, which would indicate that the two sets of representations are dissimilar and not linked by a simple transformation.
Revised: However, we observe that even for small values of µ and σ , CKA drops to 0, which would indicate that the two sets of representations are dissimilar and not linked by a simple transformation.

Original: 3. C ONSTRUCTING A LOW CKA SIMILARITY
Revised: 3. C ONSTRUCTING A LOW CKA SIMILARITY

Original: We further experiment with CKA’s sensitivity to translations but in a more realistic setting, while preserving important semantic information.
Revised: We further experiment with CKA’s sensitivity to translations but in a more realistic setting, while preserving important semantic information.

Original: We consider the 9 layers CNN presented in Section 6.1 of Kornblith et al.
Revised: We consider the 9 layers CNN presented in Section 6.1 of Kornblith et al.

Original: (2019) trained on CIFAR10.
Revised: (2019) trained on CIFAR10.

Original: When trained on classification tasks, ANNs tend to learn increasingly complex representations of the input data that can be more or less linearly separated into classes by the last layer of the network.
Revised: When trained on classification tasks, ANNs tend to learn increasingly complex representations of the input data that can be more or less linearly separated into classes by the last layer of the network.

Original: With this in mind, we argue that a relevant way in which two sets of representations can be “similar” in practice is if they are linearly separable by the same hyperplanes in parameter space, the same linear classifier can then correctly classify both sets of representations.
Revised: With this in mind, we argue that a relevant way in which two sets of representations can be “similar” in practice is if they are linearly separable by the same hyperplanes in parameter space, the same linear classifier can then correctly classify both sets of representations.

Original: Given X , the network’s internal representations of 10k training images at the last layer before the output, we can create Y by translating examples corresponding to a certain class in a way that does not affect the linear separability of the representations.
Revised: Given X , the network’s internal representations of 10k training images at the last layer before the output, we can create Y by translating examples corresponding to a certain class in a way that does not affect the linear separability of the representations.

Original: We use an SVM classifier to extract the hyperplanes in parameter space, which best separate the data (with approx.
Revised: We use an SVM classifier to extract the hyperplanes in parameter space, which best separate the data (with approx.

Original: 91% success rate).
Revised: 91% success rate).

Original: We then translate a subset of the representations in a direction which will never cross these hyperplanes.
Revised: We then translate a subset of the representations in a direction which will never cross these hyperplanes.

Original: We plot the CKA values between X and Y according to the translation distance in Fig. 3c).
Revised: We plot the CKA values between X and Y according to the translation distance in Fig. 3c).

Original: The CKA values quickly drop to 0 , despite the existence of a linear classifier that can classify both sets of representations into the correct classes with > 90% accuracy.
Revised: The CKA values quickly drop to 0 , despite the existence of a linear classifier that can classify both sets of representations into the correct classes with > 90% accuracy.

Original: In Fig.
Revised: In Fig.

Original: 3c) we also examine linear CKA’s sensitivity to outliers.
Revised: 3c) we also examine linear CKA’s sensitivity to outliers.

Original: Plotted are the CKA values between the set of training set representations and the same representations but with a single point being translated far from its original location.
Revised: Plotted are the CKA values between the set of training set representations and the same representations but with a single point being translated far from its original location.

Original: While the translation distance needed to achieve low CKA values is relatively high, the fact that the position of a single point out of tens of thousands can so drastically influence the CKA value raises doubts about CKA’s effectiveness as a similarity metric.
Revised: While the translation distance needed to achieve low CKA values is relatively high, the fact that the position of a single point out of tens of thousands can so drastically influence the CKA value raises doubts about CKA’s effectiveness as a similarity metric.

Original: 3. O PTIMIZING CKA M AP
Revised: 3. O PTIMIZING CKA M AP

Original: The CKA map, commonly used to analyze network architectures (Kornblith et al., 2019) and their inter-layers similarities, is a matrix M , where M [ i, j ] is the CKA value between the activations of layers i , and j of a network.
Revised: The CKA map, commonly used to analyze network architectures (Kornblith et al., 2019) and their inter-layers similarities, is a matrix M , where M [ i, j ] is the CKA value between the activations of layers i , and j of a network.

Original: We set to directly manipulate the CKA map of a trained network f θ ∗ , 
Revised: We set to directly manipulate the CKA map of a trained network f θ ∗ , by

Original: [Figure]
Revised: [Figure]

Original: Correlating 1st & Last Layer Acc.
Revised: Correlating 1st & Last Layer Acc.

Original: Correlating all Layer Acc. 82.
Revised: Correlating all Layer Acc. 82.

Original: Differentiating All Layers Acc.
Revised: Differentiating All Layers Acc.

Original: [Figure]
Revised: [Figure]

Original: Acc. 85.7% 0.
Revised: Acc. 85.7% 0.

Original: by adding the desired CKA map to its optimization objective, while maintaining its original outputs via distillation loss (Hinton et al., 2015).
Revised: adding the desired CKA map, M target , to its optimization objective, while maintaining its original outputs via distillation loss (Hinton et al., 2015).

Original: We further optimize f θ ∗ over the training set ( X, Y ) following the equation:[Equation]
Revised: We further optimize f θ ∗ over the training set ( X, Y ) following the equation:

Original: [Equation]
Revised: [Equation]

Original: Where L map =
Revised: Where L map =

Revised: (cid:80) i,j ln cosh (cid:0) M

Original: ln cosh (cid:0) M f θ ( X ) − M  target (cid:1) .
Revised: [ i, j ] f θ ( X ) − M [ i, j ] target (cid:1) .

Original: Fig.
Revised: Fig.

Original: 4 shows the CKA map of f θ ∗ along with the CKA map of three scenarios we investigated: (1) maximizing the CKA similarity between the 1 st and last layer, (2) maximizing the CKA similarity between all layers, and (3) minimizing the CKA similarity between all layers (for network architecture and training details see Appendix).
Revised: 4 shows the CKA map of f θ ∗ along with the CKA map of three scenarios we investigated: (1) maximizing the CKA similarity between the 1 st and last layer, (2) maximizing the CKA similarity between all layers, and (3) minimizing the CKA similarity between all layers (for network architecture and training details see Appendix).

Original: In cases (1) and (2), the network performance is barely hindered by the manipulations of its CKA map.
Revised: In cases (1) and (2), the network performance is barely hindered by the manipulations of its CKA map.

Original: This is surprising and contradictory to the previous findings (Kornblith et al., 2019; Raghu et al., 2021) as it suggests that it is possible to achieve a strong CKA similarity between early and later layers of a well-trained network.
Revised: This is surprising and contradictory to the previous findings (Kornblith et al., 2019; Raghu et al., 2021) as it suggests that it is possible to achieve a strong CKA similarity between early and later layers of a well-trained network.

Original: We further manipulated the CKA map of f θ ∗ to produce a series of comical CKA maps (Fig. 5).
Revised: We further manipulated the CKA map of f θ ∗ to produce a series of comical CKA maps (Fig. 5).

Original: Although the network CKA maps seen in Fig.
Revised: Although the network CKA maps seen in Fig.

Original: 5 closely resemble their respective targets, it should be noted that we prioritized maintaining the network outputs, and ultimately its accuracy by choosing small λ values in Eq.
Revised: 5 closely resemble their respective targets, it should be noted that we prioritized maintaining the network outputs, and ultimately its accuracy by choosing small λ values in Eq.

Original: Higher values of λ results in stronger agreements between the target and network CKA maps at the cost of performance.
Revised: Higher values of λ results in stronger agreements between the target and network CKA maps at the cost of performance.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We continue a very recent line of research by showing situations in which linear CKA is inadequate as a representation similarity measure.
Revised: We continue a very recent line of research by showing situations in which linear CKA is inadequate as a representation similarity measure.

Original: Linear CKA attributes low similarity to sets of representations that are directly linked by simple affine transformations preserving important functional characteristics and it attributes high similarity to representations from very different networks.
Revised: Linear CKA attributes low similarity to sets of representations that are directly linked by simple affine transformations preserving important functional characteristics and it attributes high similarity to representations from very different networks.

Original: Furthermore, we can manipulate CKA to give arbitrarily low/high values while preserving functional behaviour.
Revised: Furthermore, we can manipulate CKA to give arbitrarily low/high values while preserving functional behaviour.

Original: An important direction of improvement for our work is to theoretically study CKA’s sensitivity to some of the used transformations.
Revised: An important direction of improvement for our work is to theoretically study CKA’s sensitivity to some of the used transformations.

Original: We hope our work inspires ML researchers to also consider other similarity measures and potentially different data analysis and visualization algorithms when studying neural network representations as opposed to relying solely on CKA to draw conclusions.
Revised: We hope our work inspires ML researchers to also consider other similarity measures and potentially different data analysis and visualization algorithms when studying neural network representations as opposed to relying solely on CKA to draw conclusions.

Original: It would also be interesting to see whether or not results obtained using other methods of analysis would support scientific conclusions already established using linear CKA.
Revised: It would also be interesting to see whether or not results obtained using other methods of analysis would support scientific conclusions already established using linear CKA.

Original: learning
Revised: learning

Original: net/
Revised: net/

Original:  A PPENDIX 5. N ETWORK A RCHITECTURE
Revised:  A PPENDIX 5. N ETWORK A RCHITECTURE

Original: In Sec.
Revised: In Sec.

Original: 3.1 we use a 9 layer neural network; the first 8 of these layers are convolution layers and the last layer is a fully connected layer used for classification.
Revised: 3.1 we use a 9 layer neural network; the first 8 of these layers are convolution layers and the last layer is a fully connected layer used for classification.

Original: We use ReLU (Nair & Hinton, 2010) throughout the network.
Revised: We use ReLU (Nair & Hinton, 2010) throughout the network.

Original: The kernel size of every convolution layer is set to (3 , 3) except the first two convolution layers, which have (7 , 7) kernels.
Revised: The kernel size of every convolution layer is set to (3 , 3) except the first two convolution layers, which have (7 , 7) kernels.

Original: All convolution layers follow a padding ofand a stride of 1.
Revised: All convolution layers follow a padding ofand a stride of 1.

Original: Number of kernels in each layer of the network, from the lower layers onward follows: [16 , 16 , 32 , 32 , 32 , 64 , 64] .
Revised: Number of kernels in each layer of the network, from the lower layers onward follows: [16 , 16 , 32 , 32 , 32 , 64 , 64] .

Original: In this network, every convolution layer is followed by batch normalization (Ioffe & Szegedy, 2015).
Revised: In this network, every convolution layer is followed by batch normalization (Ioffe & Szegedy, 2015).

Original: The network we used in Sec.
Revised: The network we used in Sec.

Original: 3.4 is similar to the network we just described, except the kernel size for all layers are set to (3 , 3) .
Revised: 3.4 is similar to the network we just described, except the kernel size for all layers are set to (3 , 3) .

Original: 5. T RAINING D ETAILS
Revised: 5. T RAINING D ETAILS

Original: The models in Sec. 3.1, both the generalized and memorized network, were trained for 100 epochs using AdamW (Loshchilov & Hutter, 2017) optimizer with a learning rate (LR) of 1 e - 3 and a weight decay of 5 e - 4 .
Revised: The models in Sec. 3.1, both the generalized and memorized network, were trained for 100 epochs using AdamW (Loshchilov & Hutter, 2017) optimizer with a learning rate (LR) of 1 e - 3 and a weight decay of 5 e - 4 .

Original: The LR is follows cosine LR scheduler (Loshchilov & Hutter, 2016) with an initial
Revised: The LR is follows cosine LR scheduler (Loshchilov & Hutter, 2016) with an initial

Original: LR stated earlier.
Revised: LR stated earlier.

Original: The training of the base model (original) model in Sec.
Revised: The training of the base model (original) model in Sec.

Original: 3.4 follows the same training procedure as of the models from Sec. 3.1, except in this setting we train the model for 200 epochs, with an initial
Revised: 3.4 follows the same training procedure as of the models from Sec. 3.1, except in this setting we train the model for 200 epochs, with an initial

Original: LR of 0.01.
Revised: LR of 0.01.

Original: All other models in Sec.
Revised: All other models in Sec.

Original: 3.4 (with a target CKA map to optimize) are also trained with similar training hyperparameters to that of the base model, except the followings: (1) these models are only trained for 50 epochs.
Revised: 3.4 (with a target CKA map to optimize) are also trained with similar training hyperparameters to that of the base model, except the followings: (1) these models are only trained for 50 epochs.

Original: (2) the objective function includes a hyperparameter λ (see Eq. 3), which we set to 100 for all models.
Revised: (2) the objective function includes a hyperparameter λ (see Eq. 3), which we set to 100 for all models.

Original: The cosine LR scheduler includes a warm-up step ofoptimization steps.
Revised: The cosine LR scheduler includes a warm-up step ofoptimization steps.

Original: (4) the LR is set to 1 e - 3 .
Revised: (4) the LR is set to 1 e - 3 .