Generalisation First, Memorisation Second? Memorisation Localisation for Natural Language Classification Tasks (2024)

Abstract: Memorisation is a natural part of learning from real-world data: neural models pick up on atypical input-output combinations and store those training examples in their parameter space. That this happens is well-known, but how and where are questions that remain largely unanswered. Given a multi-layered neural model, where does memorisation occur in the millions of parameters?Related work reports conflicting findings: a dominant hypothesis based on image classification is that lower layers learn generalisable features and that deeper layers specialise and memorise. Work from NLP suggests this does not apply to language models, but has been mainly focused on memorisation of facts.We expand the scope of the localisation question to 12 natural language classification tasks and apply 4 memorisation localisation techniques.Our results indicate that memorisation is a gradual process rather than a localised one, establish that memorisation is task-dependent, and give nuance to the generalisation first, memorisation second hypothesis.

## introduction memorisation in neural models is both concerning due to overfitting and privacy concerns, and desired because of information that needs to be stored, such as facts. with the recent surge in the number of models trained on very large, closedaccess training corpora, the nlp community has seen an increased interest in research questions that aim to improve our understanding of memorisation, such as: which data points are memorised, and can we extract those examples from models (e.g. @xcite @xcite ? how dependent is memorisation on model scale, architecture and training procedures (e.g. @xcite )? can we localise memorised information (i.e. pinpoint which weights, subcomponents or layers are most associated with storing that information) and edit models' memories (e.g. @xcite ? memorisation localisation is the central theme in this work, specifically localisation at the level of layers. there is a lack of consensus about which layers are particularly involved in memorisation in deep neural models, which may stem from the varying experimental setups and varying definitions of memorisation employed by different studies. work from computer vision (cv) mostly focused on memorisation of perfectly memorised mislabelled examples, positing that lower layers capture generalisable features while deeper layers memorise (e.g. @xcite (although that has been recently challenged by @xcite ). work from nlp often discusses memorisation of facts, for which lower @xcite , middle @xcite and final layers @xcite have all been mentioned as playing crucial roles in pre-trained language models (plms). and lately, with the availability of some open-access pre-training data, memorisation localisation for sequences memorised verbatim has gained traction, and initial results primarily point to lower layers @xcite . different articles investigating different types of memorisation arrive at different answers, and even work focused on a specific type (such as facts) does not always agree. we contribute an important piece of the puzzle in the memorisation localisation landscape, by employing a setup that is similar to @xcite we perform layer-wise localisation in fine-tuned models for 12 nlp classification tasks, enforcing memorisation by applying label perturbation to a data subset. we use four memorisation localisation methods ( §3), and first examine their accuracy in a control setup. afterwards ( §4), we address the main research question: in which layers does memorisation occur? the results do not always align -which underscores that we should not overly rely on one localisation methodbut do tell a coherent story. in §5, we introduce a visualisation technique (centroid analysis) to make this story more interpretable: memorisation is not sudden but gradual. together, layers gradually shift mislabelled examples to their newly assigned class, but when that happens is task-dependent: the better a model generalises to new data for a particular task, the more relevant deeper layers are for memorisation. figure 1 illustrates this. our findings beg a nuance of the generalisation first, memorisation second hypothesis, and we end with a discussion ( §6) of what our findings mean for localisation and model editing going forward. ## related work in this section, we discuss a broad range of related work on memorisation localisation from cv and nlp, focusing specifically on studies that discuss the role of different layers in deep neural networks. noise memorisation in cv multiple image classification studies concluded that deeper layers are more involved in memorisation than earlier layers, based on analyses of the memorisation of mislabelled or hard examples. this has been reported by work contrasting entire models (regular models and ones trained with randomised labels) @xcite @xcite , and work discussing how some examples are handled differently within one model: @xcite establish a positive correlation between prediction depth in image classification (the earliest layer that predicts the label) and example-level learning difficulty metrics. @xcite analyse hidden representations for image classification and report that memorisation of mislabelled examples occurs abruptly in late layers and late training epochs. rewinding models' final layers to earlier checkpoints reversed memorisation. contrary to previous work, @xcite report that for image classification datasets that were partially mislabelled, instead of memorisation being confined to specific layers, there are small sets of neurons dispersed over the full architecture involved in memorisation. memorisation of factual knowledge nlp memorisation localisation studies have primarily focused on factual knowledge, although only a subset of work in this direction discusses the roles of different layers. de @xcite first connected work from cv to fact memorisation in transformer, and train a hypernetwork to edit facts. their hypernetwork mostly edits the bottom layer of a six-layer transformer, and de cao et al. suggest this difference might be due to the change in modality. later work operated under the assumption of the knowledge neuron thesis, foot_0 assuming that facts are recalled from the training corpus through transformer's mlp weights, that act as a key-value memory @xcite , and that one may thus be able to identify mlp knowledge neurons @xcite : ( 1 although the articles above disagree in terms of whether higher or lower layers store factual information, work beyond model editing and knowledge neurons indicates factual information is already present in the lower layers @xcite show that facts and idioms are stored and retrieved in early layers in bert and gpt-2. deeper layers perform confidence boosting after retrieval. @xcite artificially block parts of the computation in fact retrieval for gpt-j and gpt-2 xl, and find that factual information is stored in the lower mlp sublayers, which is echoed by recent work from @xcite . ## methods to gain a good understanding of how memorisation is task-dependent, we combine binary classification tasks from (super)glue @xcite with tasks from more diverse domains and label set sizes. the tasks fall into four categories: generic natural language understanding (nlu), sentimentrelated tasks, hate speech detection and topic classification. we analyse four plms: bert-base (devlin et al., 2019), opt-125m @xcite , pythia-160m @xcite and gpt-neo-125m @xcite . we fine-tune each model separately for the 12 tasks. appendix d describes the models and our technical setup. 4 these 4 architectures are similar in size and have 12 layers each. in appendix b, we repeat a subset of the experiments with opt-1.3b. the plms (θ p ) are fine-tuned for 50 epochs, and checkpoints are stored when the training accuracy is near-ceiling (θ m 1 ), and at the end of training (θ m 2 ). we also train models using the original labels (θ o ), using the same random seeds as θ m 1 and θ m 2 . during fine-tuning, we freeze the input embeddings. results reported in §3.2 are based on one fine-tuning seed, and the remainder of the main paper computes results using three seeds. seeds control the data order and classification heads. ## results for memorisation localisation we now apply the localisation techniques to models in which all layers have been fine-tuned for one task at a time. the results indicate how important each layer is for memorisation, per dataset, per model. we cannot simply aggregate over all results (12 layers × 12 datasets × 4 localisation techniques × 4 models), because the absolute scores returned by different techniques are not directly comparable. we discuss the results per localisation technique. ## making memorisation interpretable via centroid analysis and probing this visualisation indicates that memorisation occurs through gradual changes from the first layers onward. this explains the results from the previous section, where we found that memorisation is not confined to individual layers and that lower layers were more successful in reverting memori-sation (in layer swapping/retraining) than deeper layers: memorisation starts early, and interventions are more successful when conducted before the hidden state has moved too far away from class y b . we can demonstrate this using the centroid analysis, applied while swapping six layers at the bottom or top with layers from θ o . for all tasks (see appendix a.3), we can prevent the noisy examples from moving to centroid a by replacing the bottom six, but for only a few tasks, replacing the top six has a similar effect. figure 11 shows this using rte and trec. swapping the bottom six prevents emitting the noisy class for rte and trec, but swapping the top six is more successful for trec than rte. task differences to summarise task differences, we compute two statistics: the crossing (the first layer in which the noisy mean is closer to a than to b) and classification initiation (the first layer without overlapping distributions for the two classes), shown in figure 12a . ## discussion we set out to perform memorisation localisation for natural language classification tasks by perturbing a subset of the labels and tracing those 'noisy' examples over layers. applying four localisation techniques to four models crystallised that memorisation is not local to specific layers but a cooperative process of weights from many layers. nonetheless, not all layers appear equally important. overall, early layers are more important than later ones: the model's manipulation of memorised examples starts in lower layers, and to prevent memorisation, early intervention is more successful than late intervention. we discussed results for 12-layer models in the main paper, but further experimentation with a 24-layer model (in appendix b) leads to similar findings of memorisation being gradual, with a similar ordering of tasks, but mid-range layers being more important than late layers. this is not in accordance with the generalisationfirst, memorisation-second hypothesis from cv (see §2), but does agree with more recent work on image classification by @xcite . it also aligns with related work on plms for fact memorisation and verbatim memorisation pointing to the lower layers @xcite , while also describing cooperative roles for earlier and deeper layers @xcite . the fact that memorisation is not local implies that editing model weights locally does not necessarily erase memorised information, even if a flipped label suggests this at the level of the output layer. this might be harmless when editing facts about named entities like cities (e.g. @xcite , yet is more worrisome when regarding memorisation of personal information (e.g. @xcite , and may be a reason why safety measures can easily be reversed in plms modified to reduce harmful outputs (e.g. @xcite . can we, due to the importance of early layers, conclude that our results falsify the generalisation first, memorisation second hypothesis? the results from §5 suggest that this question requires a nuanced answer due to the variation observed among tasks. the depth of memorisation is positively correlated with a model's generalisation capabilities, i.e. we do observe a generalisation first, memorisation second tendency but the better the model performs at a certain task, the stronger that tendency is. we consider better understanding what properties of a task direct memorisation to lower or higher layers an exciting avenue for future work.