Understanding Memorization using Representation Similarity Analysis and Model Stitching

Authors: aishwarya  gupta
indranil  saha
piyush  rai
marco  fumero
clementine  domine
zorah  lähner
donato  crisostomi
luca  moschella
kimberly  stachenfeld
Venue: NeurIPS_cc_2024_Workshop_UniReps
Published: 2024-09-24 13:48:29

Abstract: It is well-known that deep neural networks can memorize even randomly labeled training data, raising questions on our understanding of their generalization behavior. However, despite several priors efforts, the mechanism and dynamics of how and where in the network memorization takes place is still not well understood, with contradictory findings in the literature. In this work, we use representation similarity analysis methods, in particular Centered Kernel Alignment (CKA) and model stitching, as a simple but effective way to guide the design of experiments that help shed light on the learning dynamics of different layers in a deep neural networks that are trained using random (i.e., noisy) labels. Our results corroborate some of the previous findings in the literature, provide new insights into the representations learned by the layers of the network when trained with various degrees of label noise, and also provide guidance into how techniques such as model stitching could be best leveraged to understand the functional similarity of a model that has memorized with another model that has good generalization.

Main paper: 
 introduction: 
 Deep Neural Nets (DNN) are widely used across many domains because of their strong generalization ability on unseen inputs. However, they are also prone to memorization, i.e., they can even fit randomly labeled training datasets perfectly [25] . Memorized inputs cannot be explained using generalizable features [24] , and usually represent atypical inputs in long-tailed distributions [9, 8] , or the underlying noise/outliers in the training dataset [25] . Though sometimes memorization can be benign [1] and may help in generalization on downstream tasks [20] , it poses privacy and security concerns [17, 4] and needs to be studied and analyzed carefully. However, the precise mechanism and dynamics of memorization are not well-understood and prior works have even reported contradictory findings, e.g., works that report that memorization happens only in deeper layers [19, 18, 2] , to also works that report that memorization is dispersed and localized across many layers of the network [21] .Representation similarity analysis [14] , in particular, Centered Kernel Alignment (CKA) [15] , has been used to analyze and compare representations learned by different layers of the same network or of different networks. These methods have also been used in comparing generalized and memorized networks [7] , and for comparing layer-wise representations of the model in problem settings such as transfer learning [10] . However, more thorough investigations of memorization in a DNN, e.g., the impact of varying levels of label randomization (i.e., degree of noise) over the learning dynamics of different layers is still lacking. In this work, we leverage CKA to compare layer-wise representations of a DNN with high generalization and negligible memorization to another DNN with weaker generalization and high memorization. We simulate this by training a model on the clean dataset (referred to as a clean/noise-free model) and a dataset comprising a fraction of randomly labeled inputs (referred to as a noisy model) respectively. This helps us gain insights into the (dis)similar learning dynamics between different layers of generalized and memorized models. Moreover, we further investigate the functional similarity [3, 13] of these layers by stitching early layers of the generalized model to the latter layers of the memorized model and vice versa [6] . This provides insights into the representational and functional aspects of these layers that can enhance the understanding of memorization and can help formulate methodologies to mitigate and fix it. 
background: 
 Some of the previous works on understanding memorization in DNNs have used methods such as prediction depth [23] of each input, and manifold analysis techniques [5] . In contrast, we leverage representation similarity measures which provide a simple way to compare different layers of the same network or of two different networks. These measures provide a simple method to track how the layers evolve and eventually converge. In particular, we use Centered Kernel Alignment (CKA) [15] which is computed using the normalized Hilbert Schmidt Independence Criterion (HSIC) [11] for comparing representations of two layers from the same network or from two different networks. For any two layers l and k of sizes d l and d k , respectively, with their corresponding layer activations denoted byX ∈ R N ×d l and Y ∈ R N ×d k on a dataset D of size N , CKA is defined as CKA(l, k) = HSIC(X,Y ) √ HSIC(X,X),HSIC(Y,Y ), where HSIC(X, Y ) = 1 (N -1) 2 trace(XHY H), andH = I N -1 N 11T denotes the centering matrix used to center the matrices X and Y . In addition to CKA, we also leverage model stitching [3] which can also be used to analyze representations learned by DNNs. Stitching connects the bottom layers of one network with the top layers of another network with a simple trainable layer in between, with the performance boost/drop indicating the degree to which the representation learned by the two models are similar to each other. 
experiments: 
 To understand the dynamics of memorization in a DNN, we experiment with two different architectures -a customized CNN and ResNet-18 [12] model trained on the SVHN [22] and CIFAR-10 [16] datasets, respectively. We create noisy versions of each dataset by randomly sampling r = {5%, 10%, 15%, 20%, 25%, 37.5%, 50%, 67.5%} percentage of total images from its training dataset and assigning them random labels. We then train a model for each value of r, including r = 0 (i.e., clean train dataset) as a noise-free baseline model. During each noisy model's training, checkpoints are saved at regular intervals and are used to study layer-wise memorization. In all our experiments, we use the test split of the dataset to compute the CKA metric. Further training details are mentioned in the supplementary material.In this experiment, we study the convergence of different layers of the network as the noise ratio r increases in the training data by evaluating the similarity of the intermediate checkpoints with the final trained model. For each value of r, the similarity between the i th layer of the saved checkpoints and the corresponding i th layer in the final trained model is computed using CKA. In Figure 1 , we show the convergence of ResNet-18 model trained on CIFAR-10 dataset for varied r. Note again that r = 0 represents the model trained on the clean dataset. The convergence of CNN models trained on noisy SVHN dataset is included in the supplementary.As shown in Figure 1 , the early layers converge quickly whereas the deeper layers converge slowly. This is consistent with what is observed in earlier works as well [5, 23] . Moreover, irrespective of the noise ratio r, the convergence of the early layers follows a similar trend to that of the noise-free model, suggesting a low influence of data noise on these layers. In contrast, the convergence rate of the deeper layers is directly proportional to the noise ratio; the models trained with large r exhibit the slowest convergence. Based on these observations, we hypothesize that irrespective of r, the early layers of the noisy models are mostly functionally similar to the early layers of the noise-free model, suggesting low memorization for the very early layers.To validate this hypothesis, we leverage model stitching [6] to evaluate functional similarity. For each value of r, we create a stitched model by stitching the early k layers of the noisy model to the latter layers of the noise-free model starting from the (k + 1) th layer using either a 1 × 1 convolution or Next, we perform another layer-wise comparison between the noisy and noise-free models to further analyze memorization and potential ways to mitigate it. For all values of r, we compute the CKA score between corresponding layers of the noisy model at varied checkpoints and layers of the finally trained noise-free model and report the findings in Figure 3 . The figure shows that, during the early training phase (e.g., epoch 50), all layers of all the models learn representations that are similar to those learned by the noise-free model. However, as training progresses, the representation similarity between corresponding layers of the noisy and noise-free models starts decreasing and the difference becomes prominent in deeper layers. The drop in similarity is proportional to the noise ratio present in their respective training datasets. This is consistent with the previous findings of early stopping [18] and the efficacy of many early learning regularization approaches [19, 2] in mitigating the impact of noise.To investigate further, we apply the stitching idea again but, this time, stitch the early layers of the noise-free (clean) model with the deeper layers of the noisy model. As shown in the second column of Figure 2 , stitching early layers/blocks of the clean, noise-free model to latter layers/blocks of the noisy models does not improve the stitched model's performance, corroborating the fact that early layers of noisy models do not have significant memorization. However, stitching the clean model up to half its depth to the noisy model shows performance improvement. For instance, for ResNet-18 model, stitching the noise-free model up to the second residual block improves the noisy model's performance substantially; with further stitching resulting in performance at par with that of the noise-free model, irrespective of the r value. The similar trend is observed for CNN model trained on noisy SVHN dataset. (We further validate our finding by computing the similarity between the layers of the stitched and clean model and discuss it in the supplementary material). However, the converse does not hold i.e., stitching a noisy model up to half of its depth to the clean model does not degrade its performance substantially as shown by the "blue" and "orange" colored lineplots in the left subplot of Figure 2(b) . This shows that, unlike conventional wisdom, fixing only the last layers may be insufficient to mitigate memorization. Instead, fixing a few of the early layers of the noisy model could also contribute in helping mitigate memorization, potentially because of a cascaded overall effect these early layers might also have (as opposed to the relatively less significant effect of fixing just the end layers), and help achieve high generalization performance similar to the clean model. This experiment indicates that the effect of memorization may actually also be present in the early layers and resonates with the recent insights that memorization may not be limited to the last layers but is localized in a subset of neurons across multiple layers (including the early layers) of the model [21] . Thus the difference in the representations learned in the last layers of a memorized network as compared to that of a generalizable network could possibly be due to a cascading of subtle but important differences in earlier layers' representations too, and not solely due to the last layers learning a different representation all by themselves. 
conclusion: 
 Using representation similarity analysis (in particular, CKA) and model stitching, we have analyzed the learning dynamics of various layers of DNNs that memorize because of being trained on randomly labeled data. Our empirical analysis both corroborates earlier findings in the literature, as well as provides new insights. In particular, it sheds light on how the early and the deeper layers of the network evolve during training, the effect of the amount of noise in the labels on their convergence, and also provided insights about which parts of the networks should be fixed to mitigate memorization. 
 
 Appendix: Proceedings of the II edition of the Workshop on Unifying Representations in Neural Models (UniReps 2024). 


Decision: In light of the positive reviewers' feedback and relevancy of the submission, we are pleased to accept this paper for presentation at UniReps 2024. We kindly ask the authors to incorporate the reviewers' suggestions and feedback in the final camera-ready version of the manuscript.

Reviews: review: title: title: Late layers memorize more compared to early ones
main_review: main_review: This paper explores the dynamics of memorization in deep neural networks (DNNs), focusing on how layers in these networks behave when trained with noisy data. Authors use Centered Kernel Alignment (CKA) and model stitching as tools to analyze representations learned by neural networks and the learning dynamics at various layers. They compare models trained on clean data and models trained on varying levels of label noise. They conducted experiments on ResNet-18 using the CIFAR-10 dataset, varying the noise in the training labels to study how different layers of the network were affected. The findings suggest that early layers are less affected by noise and retain functional similarity to clean models, while deeper layers are more prone to memorization. They also propose that fixing early layers can help mitigate memorization.

**Strengths:**

Novel Approach to Memorization: The paper brings in CKA and model stitching to analyze memorization in neural networks, which is a fresh approach.

Important Findings on Early and Deep Layers: Authors observed that early layers are less prone to memorization and remain functionally similar to clean models. 

Practical Implications: The suggestion that fixing early layers can help mitigate memorization is a valuable insight for researchers and practitioners dealing with noisy datasets. It opens up new possibilities for improving model generalization in the presence of noise.


**Weaknesses:**

Limited Dataset and Architecture Choices: While the use of CIFAR-10 and ResNet-18 is standard for many experiments, these choices might limit the generalizability of the findings. Memorization behaviors might differ with larger or more complex datasets, or different architectures (e.g., transformers or more recent models). Expanding the experiments to other datasets and models would strengthen the conclusions.
title: title: Paper provides confirmatory understanding of older findings using CKA; model-stitching gives better insight
main_review: main_review: **Summary**
- This paper investigates the dynamics of memorization in deep neural networks (DNNs) using Centered Kernel Alignment (CKA) and model stitching techniques. The authors train ResNet18 on CIFAR-10 with varying degrees of label noise and analyze the convergence and similarity of layer representations. Findings reveal early layers extract generalizable features, whereas memorization appears localized to deeper layers.

**Strengths**
1. Strong clarity of method. Experimental procedure is elegant, simple and clearly explained.
2. Valid use of representation similarity (CKA) to gain understanding of where DNN memorization happens
3. The model stitching experiment provides improved insight in a causal / intervention-based way

**Weaknesses**
1. While the paper's findings regarding CKA and memorization in deeper layers are insightful, they are not entirely novel, as similar procedures have been used in previous works (e.g. ref 22). However, the confirmation of these results using different methods does add value to the field. The paper's main interesting and original contribution lies in the model stitching experiment, which challenges prevailing notions by showing that mitigations against memorization should occur in early layers too
2. It would have been nice to offer more mechanistic proposals on why the CKA method and model-stitching methods arrived at seemingly different conclusions, but since this is an extended abstract the brevity is understandable
title: title: A Comprehensive Analysis of model memorization
main_review: main_review: This paper provides a novel investigation into the dynamics of memorization in deep neural networks (DNNs), specifically using two core techniques: Centered Kernel Alignment (CKA) and model stitching. The paper successfully addresses gaps in the literature about whether memorization is confined to certain layers or dispersed throughout the network, offering new insights into how DNNs behave under noisy training conditions.

**Strengths:**
1. **Innovative Use of CKA and Model Stitching:**

The paper leverages CKA and model stitching technique, providing experimental support for conclusions about functional similarities between noisy and noise-free networks.

2. **Thorough Experiments:**

The experiments use varied noise levels on the CIFAR-10 dataset, and the findings are backed by multiple checkpoints and extensive layer-wise analysis.

3. **Addressing Prior Contradictions:**

The paper successfully reconciles conflicting findings in the literature. It shows that memorization can indeed be dispersed, but is more significant in deeper layers, providing a clearer understanding of where and how memorization occurs.

**Weaknesses:**
1. **Limited Applicability Beyond CIFAR-10:**

While the findings are compelling, the scope of the dataset (CIFAR-10) is somewhat narrow. Extending the experiments to more complex datasets like ImageNet or tasks beyond image classification would provide better generalization of the results.

2. **Stitching Not Fully Explored:**

As shown in Figure 2, for a model trained under the same noisy percentage, respectively replace the first one/two layers and the last one/two layers with the corresponding layer of the noise-free trained model, and it can be found that the test accuracy rate of the latter is always higher than that of the former, which seems to be in contradiction with the conclusion in the article (L129-137).

3. **Lack of Conceptual Explanations:**

The paper lacks clear definitions for key concepts such as memorization and generalization, which are crucial for the reader to fully grasp the contributions. Additionally, CKA and other methods, while used effectively, are not explained in sufficient detail for readers unfamiliar with these techniques.

4. **Minor Typos:**

There are several minor errors in the paper, such as a typo in line 93 where it should refer to "Figure 2(a)" instead of "Figure 3(a)." Additionally, the legends in Figure 2 are inconsistently labeled with different capitalizations, which could confuse readers.

**Originality and Significance:**

The paper makes a significant contribution by improving our understanding of memorization in DNNs through representation similarity and model stitching. These tools offer a fresh perspective on comparing noisy and noise-free models, leading to actionable insights into the learning dynamics of DNN layers. Additionally, the insights into functional similarity between layers from different networks could be useful in designing more robust and generalizable architectures.

**Clarity:**

The paper is well-written and clear, with logical flow and detailed descriptions of the methods. The figures and plots effectively illustrate key points. However, certain technical aspects of CKA and model stitching may benefit from additional explanation for readers less familiar with these techniques.
title: title: Review
main_review: main_review: This paper investigates the mechanisms of memorization in deep neural networks (DNNs) trained with random labels. The experiments use CKA and to compares layer-wise representations between DNNs with high generalization and low memorization against those with weaker generalization and high memorization. The results show that early layers converge quickly and are less influenced by label noise, while deeper layers exhibit slower convergence and greater sensitivity to noise levels. Interestingly, the research demonstrates that fixing early residual blocks can effectively mitigate memorization without significantly compromising model accuracy, challenging the traditional focus on only the final layers.

Strengths:
- Interesting findings. Shows that early layers are less impacted by memorization than deeper layers from a representation similarity perspective. Also shows how representation similarity changes over the course of training.
- Challenges the conventional wisdom re: fixing only the last layers suffices to mitigate memorization.
- Use of model stitching to supplement CKA findings. I like the use of model stitching for validating trends/findings identified via the CKA representation similarity analysis, which by itself can be unreliable [4].

Weaknesses:
- Novelty. As mentioned in the paper, there are quite a few papers that study memorization from a representation similarity perspective. So, a discussion contrasting the findings in this paper more clearly to specific findings in previous work would be useful.

- Analysis of memorization specific to random labels setting. To me, it is unclear whether these findings are specific to memorization induced via random labels or whether they are more general and can be observed in other more realistic settings as well [2,3].

Suggestions for improvement:
- Would be interesting to compare models trained with different levels of random label noise at a more fine-grained example level. One way to do this is to use similarity metrics based on data attribution (e.g., ModelDiff [5] or Influence embeddings [6]) to study how model predictions of generalizable DNN and memorization DNN functionally differ at the example level. Does memorization completely change how a model relies on individual training examples?


[1] MohammadReza Davari, Stefan Horoi, Amine Natik, Guillaume Lajoie, Guy Wolf, and Eugene Belilovsky. Reliability of cka as a similarity measure in deep learning. In The Eleventh International Conference on Learning Representations

[2] Carlini, N., Ippolito, D., Jagielski, M., Lee, K., Tramer, F. and Zhang, C., 2022. Quantifying memorization across neural language models. arXiv preprint arXiv:2202.07646.

[3] Somepalli, G., Singla, V., Goldblum, M., Geiping, J. and Goldstein, T., 2023. Understanding and mitigating copying in diffusion models. Advances in Neural Information Processing Systems, 36, pp.47783-47803.

[4] Davari, M., Horoi, S., Natik, A., Lajoie, G., Wolf, G. and Belilovsky, E., 2022. Reliability of cka as a similarity measure in deep learning. arXiv preprint arXiv:2210.16156.

[5] Shah, H., Park, S.M., Ilyas, A. and Madry, A., 2023, July. Modeldiff: A framework for comparing learning algorithms. In International Conference on Machine Learning (pp. 30646-30688). PMLR.

[6] Wang, F., Adebayo, J., Tan, S., Garcia-Olano, D. and Kokhlikyan, N., 2024. Error discovery by clustering influence embeddings. Advances in Neural Information Processing Systems, 36.