Memorization Sinks: Isolating Memorization during LLM Training
Authors: Gaurav Rohit Ghosal
Pratyush Maini
Aditi Raghunathan
Venue: ICML 2025 poster (2025)
Keywords: Memorization
Localization
Unlearning
Primary area: deep_learning->large_language_models
TL;DR: We propose a method to disentangle sequence memorization and general language model capabilities during pretraining.

Abstract: Large language models are susceptible to memorizing repeated sequences, posing privacy and copyright concerns. A popular mitigation strategy is to remove memorized information from specific neurons post-hoc. However, such approaches have shown limited success so far. In a controlled setting, we show that the memorization of *natural* sequences (those that resemble linguistically plausible text) become *mechanistically entangled* with general language abilities, thereby becoming challenging to remove post-hoc. In this work, we put forward a new paradigm of MemSinks that promotes isolation of memorization by design. We leverage a sequence identifier to activate a unique set of memorization neurons for each sequence across repetitions. By analyzing the dynamics of learning and forgetting, we argue that MemSinks facilitates clean isolation of memorized content, making it easier to remove without compromising general language capabilities. We implement MemSinks at the billion-parameter and billion-token scale, and observe both effective isolation and strong generalization. To our knowledge, this is the first proof-of-concept on real data demonstrating that simultaneous generalization and isolation is achievable. We open-source our code at http://github.com/grghosal/MemSinks.

Memorization Sinks: Isolating Memorization during LLM Training
Gaurav R. Ghosal Pratyush Maini 12 Aditi Raghunathan
Abstract
mation from models post-hoe, including fine-tuning based approaches (Maini et al., 2022; Barbulescu & Triantafil- lou, 2024). However, presently all such methods present substantial tradeoff between removing memorization and preserving general model capability.
Large language models are susceptible to mem- orizing repeated sequences, posing privacy and copyright concerns. A popular mitigation strat- egy is remove memorized information from specific neurons post-hoc. However, such ap- proaches have shown limited success so far. In a controlled setting, we show that the memoriza- tion of natural sequences (those that resemble linguistically plausible text) become mechanisti- cally entangled with general language abilities, thereby becoming challenging to remove post- hoc. In this work, put forward new paradigm of MemS ink that promotes isolation of memo- rization by design. We leverage sequence iden- tifier to activate a unique set of memorization neurons for each sequence across repetitions. By analyzing the dynamies of learning and forget- ting, we argue that MemSinks facilitates clean isolation of memorized content, making it ier to remove without compromising general lan- guage capabilities. We implement MemS inks at the parammerr and billion-token scale, and observe both effective isolation and strong generalization. To our knowledge, this is the first proof-of-concept on real data demonstrating that simultaneous generalization and isolation is achievable. We open-source code at http: /MemS inks. //github. com/ hatpm/gg/ghaa/mmms
We identify a core reason for the tradeoff: standard training induces mechanistic entanglement, where the same compo- nents support both generalization and memorization. This occurs when memorization relies mechanisms also used for general language understanding. In controlled setting, we show that natural-looking repeated sequences are memo- rized with strong entanglement, making post-hoc removal difficult without harming general performance. We further prove that gradient descent has an implicit bias toward such entangled solutions, suggesting that mechanistic entangle- ment is inherent to current training methods.
Can new training approaches better disentangle memorization and general language capabilities?
We begin with natural disentanglement attempt (Section 4; see also (Cloud et al., 2024)): restricting gradient updates from repeated sequences to designated "memorized com- ponents", while the remaining "general components" learn only from non-repeated data. This approach has two major flaws. First, it weakens generalization by depriving general components of all training signal from potentially high- quality repeated sequences. Second, and more subtly, gen- cralization further degrades when memorized components are removed at inference: the memorization components are active the forward pass during training, causing the gen- eral components to implicitly rely on them. Their post-hoc removal breaks this dependence and harms performance.
1. Introduction
Large language models often memorize sequences seen fre- quently throughout pretraining (Carlini al., 2023; Nasr et al., 2023), posing concerns in privacy, copyright, and membership inference. There is significant research on the problem of "unlearning" or removing memorized infor-
Drawing lessons from the failures above, introduce Mem- orization Sinks (MemSing inks) inspired by the previous study of localization in Maini et al. (2023).
Memorization Sinks selectively activate sequence- specific memorization components alongside shared gen- cralization components. Dropping memorization com- ponents erases the memorization of the corresponding sequence.
'Department of Machine Learning, Carnegie Mellon Uni- versity, Pittsburgh. Pennsylvania, USA 'DataologyAl, Red- wood City, CA. USA. Correspondence to: Gaurav Ghosal <gghosal andrew.emu.edu.
Proceedings of the 4dnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Note that this approach addresses the implicit dependence

---

Memorization Sinks: Isolating Memorization during LLM Training
Standard Training
Active neuron
Inactive neuron
MemSinks
Generalizing gradient
Memorization
gradients
$1
S1
Large memorization gradient
Small memorization gradient
Memorization gradient affects all neurons
Generalizing gradient updates shared neurons
Memorization already fit
propoutprevents prevents forgetting memorization
Figure 1. Conceptual Intuition of MemSinks. We conceptually partition the learning sigmal from each example into "generalization" and "memorization" components, On the left, we show that standard training can store memorization signal in any neurons. In MemSinks,, we provide a set memorization sink neurons which are shielded from forgetting induced by other examples. As a result, (a) memorization accumulates in these neurons and (b) once the sink neurons fit memorization, is no longer reinforced elsewhere in the model.
issue identified earlier, and also allows the shared compo- nents to learn from repeated sequences. But what prevents the shared components from containing memorized informa- tion? Does it suffer from the same mechanistic entanglement issue of standard training?
that MemS inks can disentangle memorization from gener- alization in realistic settings (Section 5.2).
We perform various ablations on the impact of design choices for implementing Mems inks on small-scale TinyStories pre-training task in Sections 5.1 and 5.3. To as- sess the practicality of Mems' i inks, we investigate two key axes: (1) scalability with respect model and data size, and (2) robustness noise in sequence ID assignment. While our proof-of-concept assumes access to perfect sequence IDs, real-world deployments would require tolerance to imperfect or noisy estimates. We find that MemS inks is robust to small levels of noise in the sequence IDs (up to 10%) and works across a range model sizes. Importantly, the benefits of MemSi inkk scale with increasing model size, suggesting promise in larger scales.
The magic lies in the difference in the training dynamies of generalizing and memorization signal, which has been studied in disparate lines of work. Generalizing signal is consistently amplified across training sequences (Chatterjee, 2020). However, memorization signals of different training sequences interfere with each other, resulting in cyclical forgettigg dynamic (Toneva et al., 2018): learning when the repeated sequence is seen, but forgetting when training a different sequence due to interference. In standard training, there is no separation of model compo- nents this forgetting dynamic occurs throughout the model leading to mechanistic entanglement. However, in MemS nks, this cycle is broken by allocating separate components per repeated sequence that are protected from interference with other sequences. As a result, memo- rized information is not reinforced in the shared components across repetitions of sequences.
In summary, we argue that post-hoc unlearning on standard trained models is fundamentally limited due to inevitable mechanistic entanglement between memorization and gener- alization components. We propose a new training paradigm, MemS i inks, and provide a conceptual explanation for how it enables clean disentanglement. Empirically, we show that MemSinks supports post-hoe removal of memorized content without compromising performance real-world data. We also demonstrate its practicality: the approach scales well, with gains preserved, and even amplified, as we scale up models and datasets, and it remains robust under noisy sequence ID metadata. Ultimately, MemSi nks of- fers concrete path forward for the challenge of removing memorization without degrading model capabilities.
How well does MemSing disentangle memorization and generalization in practice? We train 360M and 1.7B SmolLM models (Allal et al., 2025) on the SlimPajama dataset (Shen et al., 2024), with repeated sequences drawn from TinyStories. Standard training leads to strong mem- orization where repeated sequences show much lower loss than held-out ones. With MemS inks, dropping the mem- orization components closes over 50% of this loss gap, mitigating memorization. Furthermore, MemS i nkks (with- out memorization components) matches the validation loss of standard training and significantly outperforms a de- duplication baseline, preserving the benefits of repeated data for generalization. This provides a proof-of-concept
2. Related Works
Forgetting Memorized Sequences. With the discovery of memorization of sequences in large language models (Car- lini al., 2023; Nasr et al., 2023), there has been interest
2

---

Memorization Sinks: Isolating Memorization during LLM Training
in techniques to remove memorization in LLMs post-hoc. One class of such methods involves further training of all language model parameters in order to reduce the likelihood of a memorized sequence. For example Thudi et al. (2022) presents the simple technique of simply training to increase the loss of memorized examples. Liu et al. (2022) further regularizes this by concurrently minimizing the loss on a retain set of validation set examples. Other works have examined using preference-basee training incentivize the replacement of memorized sequence with a safe alternative (Zhang et al., 2024; Maini et al., 2023). However, currently all such unlearning methods are prone to degrading general model capabilities, beyond the desired unlearning target (Maini et al., 2023). Moreover, they require the computa- tional expense of retraining all weights and in many cases access to additional data in the form of a retain set.
3. The Pitfalls of Post-Hoe Localization
A substantial body of prior work seeks to remove memo- rized content from LLMs through post-hoc updates. These approaches attempt preciscly target the mechanisms re- sponsible for memorization, while preserving general ca- pabilities. Their success, however, hinges on the assump- tion that memorization and generalization are supported by distinct, non-overlapping mechanisms. In this section, we conduct controlled experiments to examine when such mechanistic separation holds. We compare two settings: (1) highly atypical canary sequences, commonly used in prior memorization studies, and (2) sequences resembling natural text from the pretraining corpus. We find that memoriza- tion of these more natural sequences is significantly harder to remove post-hoc, suggesting there may be mechanistic entanglement between memorization and general capabili- ties. an analytical setting, we show that gradient descent actively prefers entangled solutions.
Localization of Memorization. A promising class of meth- ods for removing memorization seeks to localize specific model components, such as neurons, that are responsible for storing memorized sequences (Maini et al., 2023). Once such neurons are identified, they can simply be dropped out to ensure unlearning. These methods have the advantage of avoiding costly full model training and seek avoid model degradation by minimizing the number of parameters changed. Localization techniques have achieved success in removing personally identifying information (PII) (Chen et al., 2024), as well as when there exist distinctive mem- orization triggers (Stochr et al., 2024). However, all such methods present tradeoff between removing memorization and preserving model performance (Chang et al., 2024b).
3.1. Experimental Setting
We train models on two controlled settings designed to induce different types of memorization: repeated natural text sequences and highly atypical canaries. We then test whether the memorized sequences can be removed from the model without inducing model degradation.
Datasets. We conduct our experiments in a controlled setting using a subset of the TinyStories dataset (Eldan & Li, 2023). In our first setting, we randomly sample 100 stories from the TinyStories training set and repeat them 128 times TTS-Repetái t tin). Resultantly, memorized sequences are composed of natural text (coming from the same distri- bution as all other sequences). As a comparison, we study a second canary-bused setting where memorized sequences are atypical (TS-Canary). We concatenate random se- quences of tokens (Canaries) to 100 stories and repeat them 128 times in training. While TS-Canary more closely resembles memorization of label noise or atypical examples in supervised settings, it may be less representative of mem- orized documents in LLMs. In both cases, we additionally include 20,000 un-repeated SinyStories sequences (in total we train for 166 tokens).
Cftture-of-Expeer Models. Sparse mixture-of-expert models (S-MOE) have been studied in prior works as a means to expand model capacity without increasing com- pute costs (Shazeer et al., 2017). Like MemSi inks, MOE models selectively activate fully-connected neurons depend- ing on the input token. However, the selection of neurons in SMOE is typically learned subject to load-balancing con- straints (as opposed to explicitly enforced as in MemS sss.). Consequently, it is unclear what degree of localization is achieved by SMOE. Zoph et al. (2022) find that experts gen- erally specialize based off of low-level syntactic structures rather than semantic features. Dai et al. (2024) find that there is often redundancy across experts in MOEs, challeng- ing the idea that they reliably localize information within experts. Other works (Gururangan al., 2021; Park al., 2025) provide some evidence of localization in MOE mod- els: however they measure localization relative to coarse attributes such as language and general topics. In this work we study the far more fine-grained problem of localizing memorization of specific documents or sequences.
Evaluation Metrics. We measure sequence forgetting as the difference in loss on repeated sequences before and after localization and dropout (higher is better). We measure the model degradation as the difference between the validation loss before and after removal (higher is better). This reflects that we hope to avoid increases in validation loss when removing memorization.
3

---

Memorization Sinks: Isolating Memorization during LLM Training
3.2. Empirical Observations
gled solutions can be preferred by the training process.
We show the results of our analysis in Figure 2. We ob- serve that both post-hoe methods achieve limited success and struggle particularly to remove typical memorized se- quences from TS-Repetit ition.
We consider a setting where the model features are sepa- rated into generalizing subspace Sem and a memorization subspace Smen- We assumee that memorization can be imple- mented either with an entangled solution, W Wen (reusing the features Smem) or a disentangled solution Wdis using the orthogonal space Smom- Our complete setting is described in Appendix D. Under some assumptions. we prove that gradient descent is biased towards reusing the features Syem:
Localization Methods Achieve Only Partial Success. In Figure 2(a) we show the trade-off in sequence forgetting and model degradation of pruning. We observe in both settings that dropping out the identified neurons leads to an increase in the memorized sequence's loss, suggesting some success in localization. There are similar trends in Figure 2(b) for integrated gradients, although we observe it generally produces less model degradation than pruning. Additionnlly, we see that integrated gradients is less effective in removing memorization in TS Repetep ition, while being highly effective in TS-Canary.
Theorem 3.1 (Informal: Natural Sequence Memorization is Entangled). Consider training I f(x) = W Woo W We on Drre as in the setting of Appendix D. Then gradient flow does not converge W dis-
We defer the full proof to Appendix D, but discuss the intuition of our result here. Our result follows from past analyses of the minimum- norm bias of gradient flow. As result of this bias, memorization perturbs the generalizing subspace and any post-hoc unlearning methods must also perturb S.e-- Had memorization been implemented with Snoise. unlcarning could be accomplished through updates orthogonal to the general capabilities (reducing the risk of degrading general capabilities).
Natural Sequence are Particularly Challenging to Re- move. Across both methods, we find that applying post hoc methods to rS-Repetition results in greater model degradation than TS- TS-Canary. This difference is particu- larly pronounced for integrated gradients. Recall that the memorized sequences in TS Te-Repet on are natural similar to the non-repeated training data and the valida- tion set. Our results suggest the memorization of typical sequences may not be cleanly isolated from the model mech- anisms responsible for general capabilities.
Ultimately, our findings in this section challenge the feasibil- ity of relying on post-hoc approaches to remove memorized information. We show empirically and theoretically that memorization can be implemented closely entangled with general capabilities. In the remainder of the paper, we study ways to induce models to isolate the memorization of natural text sequences during training.
No Clear Separation Between Natural Sequence Memo- rization and Generalization. In Figure 2(c), we plot the validation and memorization of a model trained on TS-Repet in.. We see that the loss on repeated se- quences and the validation set descend simultaneously. The simultaneous learning of memorization and generalization seen here underscores the tight interplay between natural sequence memorization and general capabilities, contrast- ing form settings such as label noise where memorization occupies a distinct phase of training (Li et al., 2020).
4. The Mirage of Forced Localization
In Section 3, we found that standard training techniques can result in memorization being difficult to decouple from the model's general capabilities. This suggests it may be necessary to specifically pre-train models for the ability to eliminate memorized information downstream.
Ultimately, our empirical results highlight that removing memorized natural text can be especially challenging for simple post-hoe methods. Next, we provide an explanation of this phenomena in a theoretical setup.
Ideally, we would like to train models where sequence- specific memorization is handled by a distinct set of com- ponents (the memorization component), while general capa- bilities are encoded in separate generalization component. This would enable removal of memorized information sim- ply by modifying the memorization component, without risking any damage to the generalizing mechanisms. The most straightforward way to enforce this structure is by constraining the training gradients: updates from repeated, likely-to-be-nen beeeoriiee sequences can be restricted to the memorization component, while updates from all other se- quences update the generalization component.
3.3. Theoretical Analysis of Mechanistic Entanglement
Our empirical results suggest that not all memorized se- quences are equally easy to remove. While atypical canaries can be removed with minimal model degradation, remov- ing natural memorized sequences tends to induce model degradation. Intuitively, this suggests that the memorization of natural text sequences is closely entangled with mech- anisms responsible for general language modeling (thus being difficult to fully extricate post-hoc). In this section, we theoretically illustrate that such mechanistically entan-
While the concept of forced localization is intuitively appeal- ing, we empirically find two major drawbacks. First, leads
4

---

Memorization Sinks: Isolating Memorization during LLM Training
1.75
Sequence Torgettin
1.50
Sequence Forgetting
2.5
TS-Canary
Standard: Mem Loss
TS-Repetition
Standard: Val Loss
1.25
2.0
Validation Loss
1.00
1.5
0.75
1.0
0.50
TS-Canary
0.5
0.25
TS-Repetition
0.0
0.00
-0.10-0.08- 1.00000000.00000
0.00
-0.04
-0.03
-0.02 -0.01 0.00
25
Model Degradation
Model Degradation
Training Step
(a) Pruning
(b) Integrated Gradients
(c) Learning Curve
Figure Study Localization (a) We plot the unlearning-mode degradation tradcoff of pruning by varying the number of dropped out neurons and demonstrate the method struggles unlearn sequences of both kinds. (h) We plot the performance of integrated gradients and demonstrate that although it mitigates model degradation in both cases, it struggles with removing typical sequences. (c) Loss curve when training TS- rS-Repetit ion. We observe that memorization decreases alongside the validation loss, indicating that the model gains capability even as memorizes sequences.
have their gradients routed to the generalization neurons. We provide complete experimental details in Appendix E.
8-
Gradient Mask Keep All
Gradient Wask-Dropout
Validation Loss
Standard: With Rep
6
4.2. Empirical Findings
5
Forced Localization Impairs Generalization. We ob- serve that the performance of gradient masking is inferior to standard model, even before memorization neurons are removed (Figure 3). This observation renders gradient masking impractical, as it significantly worsens the model's general capabilities. Our findings here highlight that some "shared" neurons those that are updated by sequences are essential to maximally aggregate generalizing features. However, this suggests potential tradeoff between gener- alization and disentangling memorization: shared neurons could entangle memorization with generalization.
4-
3-
5
10
25
Training Step
Figure 3. Impact of Gradient-Maskee Training Validation Performance. We compare the validation loss of gradient-masker training with (Gradient Mask-Dropout) and without (Gradient Mask-Keepall) memorization neurons removed to a standard train- ing run (Standard: With Rep). We observe that (a) gradient- masked training learns slowly relative to standard training, achieving a sig- nificantly worse validation loss and (b) dropping out memorization neurons degrades validation performance as training progresses.
Localization Alone is Not Sufficient for Removal. In Fig- ure 3, we further compare whether the localization induced by Gradient Masking translates to distortion-fre removal of memorization. Intuitively, if generalization and memo- rization truly separated, it should be possible to remove memorization neurons without harming validation perfor- mance. In Figure 3 we examine the effect of removing memorization (seen in the gap between Gradient Mask- Keep All and Gradient Mask-Dropout, We observe that the validation loss becomes sensitive to the removal of memo- rization neurons as training progresses dropping out the memorization component exhibits increasing loss. This reveals, somewhat surprisingly, that localization alone does not guarantee degradation-frer unlearning. In Section 4.3, we study this phenomena in an analytical setup.
to significantly worse generalization compared to standard training. This suggests that shared components- -thooe up- dated by gradients from all data points-ane are crucial for max- imally learning general capabilities. Second, we find that removing memorization components still degrades model performance, indicating failure to make memorization fully independent of the general capabilities.
4.1. Forcing Localization By Masking Gradients
We adopt similar methodology to Cloud et al. (2024). In each layer we partition the intermediate neurons in the ML.P into generalization and memorization neurons. During train- ing, the gradients in the MLP layer from repeated sequences are masked to only modify weights corresponding to mem- orization neurons and conversely, non-repeated sequences
4.3. Co-Adaptation of Model Components
Previously, we made an unexpected finding: segregating the memorization to specific neurons does not automatically
5

---

Memorization Sinks: Isolating Memorization during LLM Training
enable straightforward removal. In this section, we theoret- ically examine why this sensitivity arises by studying the training dynamics of simplified gradient masking setup (described in Appendix F). We show a phenomena of neu- ron co-adaptation in which separately trained neurons can nevertheless become sensitive to each other's removal.
We propose allow memorization to build up in pre- specified memorization sink neurons by shielding them from interfering updates. Intuitively, by setting aside a known and stable location for storing memorization, we reduce the need for it to be reinforced across all model parameters. After training, removal of memorization can be achieved by simply deleting the memorization sinks.
Theorem 4.1 (Informal: Co-Adaptation between Model Components). Consider training a two layer linear neural net f with the forced-localizane
Our design requires two properties of the memorization sink neurons: (a) they must experience less interference from unrelated sequences and (b) they must avoid co-adaptation with the rest of the model (as discussed in Section 4), We propose that both objectives can be achieved through dropout. Specifically, a different sub- set of sink neurons is activated for each input sequence. This selective activation reduces the frequency of interfer- ing gradient updates sink neurons, shielding them from forgetting. At the same time, their infrequent activation also regularizes the rest of the model against co-adaptation.
scheme in Appendix F. Let f denote a two layer linear neural net trained in a standard manner on Dpre' \ (s"mem, yymem) and fa denote f with memorization neuron dropped out, Then we have for a constant e >0 0
((xx))-((x) 2 NN
Intuitively, the proof of Theorem 4.1 shows that although the memorized example is isolated to a memorization neu- ron, dropping out this neuron is insufficient to recover a model that was trained without seeing (smem, m)m)) ). Our proof relies on the fact that the forward pass activation of the memorization neuron affects the gradient update to the generalization neurons. Our analysis also predicts that this co-adaptation increases with additional training, mirroring the empirical results in Figure 3.
Implementation. We implement MemS inks in the hidden layer of transformer MLPs, as prior works have found MLPs play a crucial role in storing memorization (Nanda et al., 2023; Geva et al., 2021). We set fraction g of the hidden layer neurons in each layer generalization neurons (which are activated across all sequences), while the remaining fraction serve as memorization sink neurons. We assign each sequence in the pretraining dataset an ID and mask the memorization sink neurons deterministically as function of the sequence ID. During all evaluations, we remove memorization sink neurons. Further implementation details are provided in Appendix G.
Taken together, our findings reveal that the intuitive idea of forcing localization of sequences falls short in erucial ways. Surprisingly, it is both too disruptng the model's ability to learn general features-ants too permissive, al- lowing entanglement between memorization and general capabilities to persist. Can we overcome these opposing failure modes? In the next section, we show that a more sub- tle and targeted intervention in model training dynamics can unlock finer control over memorization and generalization.
Experimental Details. We train GPT Medium model (same as all previous experiments), where 70% of MLP neurons are shared and the remaining 30% are allocated to the pool of memorization neurons. We emphasize that there are frr less memorization neurons than total sequences. Thus, we do not assume each sequence can be allocated its own memorization neurons. We set the memorization neuron dropout ratio p to 0.3, but explore other choices in Section 5.3. We train on the TS RSs-Repetite itn dataset from Section 3.
5. Channeling Training Dynamics with MemSinks
In Section 4, we found that rigidly partitioning gradients to achieve disentanglement is insufficient. Our result sug- gest the need enforce more precise separation between memorization and general capability. However, it is chal- lenging to exactly specify what consitutes memorization versus generalization during the training process.
5.1. Validation of MemSinks in TinyStories
MemSinks Enable Learning Across Sequences. In Figure 4(a), we compare the validation loss of Mems i nks with standard training with and without repeated documents. Firstly, note that standard training with repeated sequences outperforms filtering them out. This indicates that the model does learn general capabilities from observing documents repeated multiple times in our setting. Next, we compare the standard trained models with MemSinks. We observe that (evaluating without the memorization neurons), MemS inks achieves comparable validation loss to standard training.
To address this, we leverage the difference in training dy- namies of memorization and generalization. During training, generalizing signals are reinforced across distinct sequences and are broadly amplified (Chatterjee, 2020). Sequence- specific memorization, on the other hand, experiences in- terference from other examples and is gradually forgotten. As a result, repeated sequences are cyclically learned and forgotten throughout the training process, thereby becoming defused throughout the model.
6

---

Memorization Sinks: Isolating Memorization during LLM Training
3.0
Memorized Example Loss
7
MemSinks
MemSinks
Standard: No Rep
6-
Standard: No Rep
Sequence Forgetting
2.5
6
Standard: With Rep
5
Standard: With Rep
2.0
Val Loss
5-
4
1.5
3
1.0
Pruning
4
2-
0.5
MemSinks
3
Standard: No Rep
0.0
5
20
25
20
25
0.20
-0.15
-0.10 - 0.05 0.00
Eval Steps
Training Step
Model Degradation
(a) Validation Loss
(b) Loss on Memorized Examples
(c)
Unlearning-Modee Degradation Tradeoff'
Figure 4. Performance of MemSinks (a) We find that MemSinks achieves a comparable validation loss normally trained model on TS-Repet itiinn outperforming a model trained without repeated sequenees. (b) We show the loss of MemS mmmiink on the repeated sequences, showing that it memorizes significantly less than normally trained model. (c) We compare the sequence forgetting- modell degradation tradeoff' of MemSinks, relative to the post-hoc methods tested in Section 3, finding KemSinks outperforms both. We compute the model degradation for MemSinks and Standard: No Rep as difference in the validation loss relative a standard trained model on IS-Repet iinn.
Dropping Out MemSinks Removes Memorization. In Figure 4(b), we show the loss on the repeated TinyStories documents. A standard trained model memorizes these se- quences during training, achieving close to 0 loss on them. We observe that dropping out the memorization neurons significantly increases the loss on these sequences, increas- ing the loss to roughly 66% a standard trained model that does not memorize. Interestingly, we observe that in the latter part of training, the loss of sequence-tied dropout on memorized sequences begins to increase. This suggests that while shared neurons may initially implement some memorization, further training reverses this.
6(a) in the validation loss gap between the deduplication baseline (where the TinyStories are seen only once) and the full repetition baseline. We observe that the validation loss attained by MemSink outperforms the deduplication baseline (achieving validation loss comparable to the re- peated baseline). This provides evidence that MemS i inks can leverage the generalization benefits of repeated data.
MemSinks Mitigates Memorization. Although repeat- ing data has benefits for generalization, it also results in memorization as evidenced by the gap between the vali- dation loss and the loss on training examples. However, as shown in Figure 4(b) we find that training with MemS i nks achieves a significantly higher loss on the training data, sub- stantially closing the gap between validation and training loss (by at least 50%). a setting with more extensive repetition (20x), additionally find that the mitigation of memorization by MemSinks is more pronounced.
5.2. Larger Scale Experiments with MemSinks
In the previous section we studied MemS i nnks in small scale setting. Now, we provide empirical validation of MemSinks on a larger scale, involving open pretraining datasets. We concentrate on settings where repeating data beneficial for generalization (i.e. by upsampling rare domain) but we also wish avoid exactly memorizing the specific repeated documents.
Collectively, our findings in this section provide promising proof-of-concee evidence that MemS i inks can effectively separate memorization and generalization the presence of heterogeneous and noisy data that constitute the real-world LLM pretraining corpora.
Setting. We focus pretraining SmolLM 360 and 1.7 billion models on 1 billion and 2 billion tokens, respectively. The majority of the pretraining dataset consists of a sub- set the SlimPajama dataset, with small amount being sampled from TinyStories. Our goal is to simulataneously improve performance on validation TinyStories data while also avoiding memorization of the repeated TinyStories training sequences. Full details are given in Appendix I.
5.3. Practicality of MemSinks
In this section, we examine the robustness of MemS nks to various hyperparameters and variation in the training pro- cess. We focus the loss memorization examples as measure of
success in isolating memoriza- tion. We perform the experiments here in the small-scale Tiny Stories setting from Section 5.1.
MemSinks Preserve Benefits of Repeated Data. In the settings we study, validation performance on TinySto- ries benefits from repetition (by upweighting this relatively under-represented distribution). This can be seen Figure
Impact of MemSinks Activation Ratio. A key hyper- parameter for MemS i nks is the fraction of memorization
7

---

Memorization Sinks: Isolating Memorization during LLM Training
Memorized Example Loss
p=0.7
344M
Memorized Example Loss
7
0.01
1.9
p= 0.5
Forgettins
0.05
p= 0.3
1.8
170M
0.1
p= 0.1
0.5
4
Sequence
1.7
3
1.6
51M
2-
32M
1.5
-0.06
-0.04
-0.02
0.00
5
15 Training Step
25
Training Step
Model Degradation
(a) Impact of Activation Ratio
(h) Impact of Model Size
(c) Robustness Sequence ID Noise
Figure 5. Practicality of MemSinks (a) We study the impact of the fraction of memorization neurons activated (p) on any given sequence. We find that the model generally robust to this choicc, but activating too many can interfere the isolation of memorization. (b) We plot the tradeoff between removing memorization and model degradation across model sizes, finding that larger models achieve better tradeoff. (c) study the impact sequence noise d, where a fraction repeated documents have an inconsistent ID. We find that MemSinks withstands small amounts of noise (up to 10%).
MemSinks Val. Loss
Memsinns Val. Loss
4.5
MemSinks Loss
MemSinks
10x Val Loss Basaline
4.0
20x Val. Loss Baseline
10x Memn Loss Baseline
20x Mom. Loss Baseline
1x Val. Loss Baseline
3.5
Vol. Loss Baseline
4.0
3.0
3.5
2.5
2.0
13
2.0
100 Training Iteration
Training Step
(a) 360M Parameters
(b) 1.7B Parameters
Figure 6. Larger Scale Experiments on MemSinks We plot the validation and memorization losses of (a) 360M and (b) 1.7B SmolLM- style models mixture of SlimPajama and TinyStories data. We compare with standard training with and without repeated Tiny Stories data. Our results demonstrate that across both settings, MomS i nks able to mitigate memorization (relative to standard training) without harming validation loss (outperforming or comparable deduplicated training).
neurons activated on any given sequence (p). This con- trols across how many documents particular memorization neuron is activated, capacity is allocated to a given sequence. We see that MemSink inkks is generally robust to this parameter, with a relatively small value = 0.3 performing best. The per- formance of Kemmiinks does breaks down at higher levels (i.e. p 0.7). We hypothesize this breakdown arises as result of insufficient shielding of memorization sinks from forgetting and further investigate this in Section 5.4.
as indicated by the comparably high losses on repeated sequences (relative normally trained model which at- tains nearly 0 loss). On the other hand, we find that model degradation (the increase in validation loss compared to standard trained model of the same size) does grow as the model architecture becomes smaller. However, even on smaller models MemS inks outperforms post-hoc methods as shown Figure 4(c). Thus, while model size plays a role in the success of MemSinks, the method has benefits in small models well.
well as how much memorization
Another concern is that MemSinks could necessitate using significantly larger models. In Figure 5(b), we test the performance of MemSinks on a range of model sizes and find that it is capable of isolating memorization across model scales
Impact of Model Size.
Robustness to Masking Noise. Implementing MemS i nks requires that repetitions of the same sequence be presented with consistent mask over the memorization sink neu- rons. Now, we investigate the robustness of MemSing inks when this masking occurs inconsistently. Whenever re-
8

---

Memorization Sinks: Isolating Memorization during LLM Training
Memorized Example Loss
1.4
that sequence-specific memorization accumulates quickly in memorization sinks (relative to standard training). In Theorem H.2, we formally demonstrate that the amount of forgetting experienced on given example depends on the number of steps taken on other examples, as well as the sim- ilarity of the other examples with the memorized example (measured in the cosine similarity). Based on this result, we show that the amount of memorization in the shared neurons can be upper bounded by quantity dependent on the mask- ing ratio, while it can be lower bounded in the memorization neurons. Our results theoretically mirror the trends seen in Figure 5(a).
Standard: With Rep
1.2
MemSinks
1.0
0.8
0.6
0.4
0.2
0.0
650
700
750 800 850 900 9501000 Traininq Step
Figure 7. Learning and Forgetting Dynamics Standard Training and MemSinks. We track the training loss single memorized example across the training trajectory for both standard training and McmSinks. We find that McmS 1 nks experiences lower amplitude learning- forgetting cycles than standard training.
6. Discussion
In this work, we have studied a crucial challenge in the responsible deployment of large-languaee models: remov- ing memorized information without harming general model capabilities. Although extensive prior work has focused on post-hoc removal strategies, we highlight an important short- coming of this approach: memorization can often leverage general capabilities making degradation- free post-hoe re- moval impractical. Our findings argue for the need to more explicitly structure models allow post-hoc unlearning to be performed reliably. Towards this paradigm, our work sheds light on important considerations for promoting this structure - highlighting why naive but intuitive approaches can be impractical. On the other hand, we uncover a frame- work for leveraging training dynamics to isolate memoriza- tion without tradeoff's in validation performance.
peated sequence is encountered, we randomly perturb its ID with probability d. Our results in Figure 5(c) show that MemSinka is robust to relatively small values of d, up to 10%. This suggests that having some amount of noise in the sequence IDs is permissible. On the other hand, when sequence IDs are highly inconsistent across repetitions (50% noise), MemS Sinks fails to isolate memorization as effec- tively. This verifies that consistency of sequence IDs across repetitions is an important factor behind its success.
5.4. Exploring the Mechanism of MemSinks
Recall that the motivation for MemS ins: to redirect sequence-specific memorization towards sink neurons. Ear- lier, we made an intuitive hypothesis that this could be accomplished by ensuring that these memorization sink neu- rons are activated less frequently. Here, we formalize this argument via empirical and theoretical investigations.
Limitations. As the focus of this paper has primarily been conceptual, much of the experiments are performed in small-scale and controlled settings, As such, further inves- tigation is necessary to examine the impact of MemS i nks on model capabilities that arise at larger scales. addition, the current implementation of MemSi nks crucially relies on meta-data which groups duplicated examples (so they can be given appropriate dropout masks). Although we pro- vide evidence that MemS inks is robust to some level of inconsistency in these annotations, it is important to further investigate efficient techniques to generate this meta-data and how robust MemSi i nkk remains in larger-scale settings. Finally, an important avenue for future work is verifying that MemSi Binssnns is robust to adversarial extraction techniques.
Dynamics of Memorization Sink Neurons. We first study the learning dynamics of the memorization sinks. We rerun our TinyStories pretraining setup with a single repeated sequence that is observed every 40 gradient steps. We track the training loss of this sequence over the course of training for both standard training and MemSi ink.. Recall that in MemSi ink, the training loss on sequence uses forward pass with the shared neurons and the sequence's 's assigned memorization neurons activated. We track this training loss as it determines the size of the update experienced by the model from observing the sequence. We see that later in training, standard training continues to experience high- amplitude learnig/forgeetinng cycles. MemSi ink,, on the other hand experiences less such fluctuations, maintaining a lower train loss on the repeated sequence. Our findings suggest that memorization sink neurons fit memorization and experience protection from forgetting.
Future Directions& & Beyond Memorization. Our work has primarily focused on verbatim memorization. However, changing the annotations used by MemS inks could allow for localization at various levels of granularity. For example, using domain or topic annotations could enable unlearn- ing different data sources or topics. Future research can also examine the benefits of localization beyond unlearning. For example, better localization of facts could enable more reliable editing of knowledge.
How Memorization Accumulates. Previously, we saw

---

Memorization Sinks: Isolating Memorization during LLM Training
Impact Statement
ing in large language models, 2024. URL https:
/ //aaa/i..rrgaa////////// 0009977770077.7.
This work aims to advance the field of Machine Learning by proposing a method to localize and manage memorized sequences in language models. Our approach, MemSinks. directly tackles critical privacy concern: once sequence is memorized by model, unlearning it is difficult without extensive fine-tuning. By systematically isolating memo- rized information, our method enables safer, more targeted removal of private or sensitive text.
Bourtoule, L., Chandrasckaran, V., Choquette- Choue, C. A., Jia, H., Travers, A., Zhang. B., Lie, D., and Papernot, N. Machine unlearning. In 2021 IEEE Symposium on Security and Privacy (SP), pp. 141-159. TEEE, 2021.
Carlini, N., Ippolito, D.. Jagielski, M., Lee, K., Tramer, F., and Zhang, C. Quantifying memorization across neural language models, 2023. URL abs / 2202.2222202 646
From an ethical perspective, the ability to better control and remove specific memorized sequences has potential societal benefits, particularly regarding user privacy and data protec- tion. On the other hand, any mechanism that manipulates model internals could be misused for censorship or selective information removal if deployed irresponsibly. Mitigating such risks requires robust governance and transparent poli- cies on what content may be unlearned.
Chang, H., Park, J., Ye, S., Yang, S., Seo, Y., Chang, D.-S., and Seo, M. How do large language models acquire factual knowledge during pretraining?, 2024a. URL https:/ /ttp/aaxiiv////////////////t http://att.rrgabs////////// wtt.:////..b.........o/ 11813.
Chang, T.-Y., Thomason, J., and Jia, R. Do localization methods actually localize memorized data in llms? a tale of two benchmarks, 2024b. URL aattps://arae org/ abs aww///s11111 11.22311990000
Our proposal also contributes the broader conversation about data govemnacee in large-scale AI systems. As lan- guage models continue to grow in size and capability, bal- ancing powerful generalization with privacy safeguards will become increasingly essential. We believe this work offers a practical step forward, but emphasize the need for interdisci- plinary collaboration spanning technical, ethical, and legal domains- to ensure that memory localization technologies support socially responsible AI systems.
Chatterjee, S. Coherent gradients: An approach to under- standing generalization in gradient descent-based opti- mization. arXiv preprint arXiv:2002. 10657, 2020.
Chen, R., Hu, T., Feng, Y., and Liu, Z. Learnable privacy neurons localization in language models, 2024. URL 10989.
https://arxiva http//aaiivvr.rr.ae...//// abs/ 2405. 1498.
Acknowledgements
Cloud, A., Boldman-Wetzleds J., Wybitul, E., Miller, J., and
Turner, A. M. Gradient routing: Masking gradients to localize computation in neural networks, 2024. URL wwwpg/aa///////eeeae//eeeeeeeeee abs///////// 2410 04332.
We gratefully acknowledge support from Apple, NSF and the A12050 program at Schmidt Sciences. The authors would like thank Jacob Springer, Chen Wu, Neil Kale, and Ziqian Zhong for their very helpful feedback and dis- cussions during the onset of the project. Additionally, the authors are extremely grateful to Fahim Tajwar, Abitha Thankaraj, Tori Qiu, and Naveen Raman for their very help- ful review and feedback on the finalized version the paper. The authors are also very grateful to Sachin Goyal for his help and advice on using the LitGPT Library.
https: /ttps://arxvvvooo www.xiiioorre/eeeeee
Dai, D.,, Deng, C., Zhao, C., Xu, R. X., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., Xie, Z., Li, Y. K., Huang, P., Luo, F., Ruan, C., Sui, Z., and Liang, W. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, 2024. URL https: //arxiv. org/ abs / 22002 06066.
Eldan, R. and Li, Y. Tinystories: How small can language models be and still speak coherent english?, 2023. URL https: //aa/aavvvv./////////aa/////aaiv httg:/bb/223000075555
References
AI, Lightning-AI//I
Litgpt.
https::/gitee com/ 2023.
Feldman, V. and Zhang, C. What neural networks mem- orize and why: Discovering the long tail via influence estimation, 2020. URL https: /ttp/aaxvvv.../ org abs/
Allal, I.. B., elie, Marafioti, A., Lozhkov, A., Lhoest, Q.
Noyan, M., Farré, M., Blázquez, G. M., Tazi, N., yousan, Bartolome, A., Ferdman, E., Isbarov, J., vb, and Cheng, L. huggingfacce/smolln. 3 2025. URL ht / com/hugg www.hhuggggggnnn..eoo smollm.
Geva, M., Schuster, R., Berant, J., and Levy, O. Trans- former feed-forward layers are key-value memories. In Moens, M.-F., Huang, X., Specia, L., and Yih, S. W.-t. (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. pp. 5484-5495, Online and Punta Cana, Dominican
Barbulescu, G.-O. and Triantafillou, P. To each (textual sequence) its own: Improving memorized-data unlearn-
10

---

Memorization Sinks: Isolating Memorization during LLM Training
Republic, November 2021. Association for Computa- tional Linguisties. doi: 10. 00.18633/11221..mmmppp
Park, J., Ahn, Y. J., Kim, K.-E., and Kang. Monet:
Mixture of monosemantic experts for transformers, 2025.
446. URL https: emnlp-main.442
org/2011.
URL https: /ttp:s://aa/vooo// wwt./ooggg///aaooooe..ooooeeaa/t 04139.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q.,
Gururangan, S., Lewis, M., Holtzman, A.. Smith, N. A., and Zettlemoyer, I.. Demix layers: Disentangling domains for modular language modeling, 2021. URL https: //arxiv. org/ abs 2abs/2100 2108.00000 05036.
Hinton, G., and Dean, J. Outrageously large neural net- works: The sparsely- mixture-of-exxpets layer, 2017. www.ariiiiiv...o wwwpw/wwww.w.////////../
URL
Shen, Z., Tao, T., Ma, L.., Neiswanger, W., Liu, Z., Wang, H., Tan, B., Hestness, J., Vassilieva, N., Soboleva, D., and Xing, E. Slimpajama-dc: Understanding data combi- nations for Ilm training. 2024. htLp:://://art
Harutyumyan, H., Reing, K., Steeg, V., and Galstyan, A. Improving generalization by controlling label-noise information in neural network weights, 2020. URL https: /ttps://arzee /ttp///aavvvv/v. org/ abs/ 2002 07933.
org / abs / 22399 10818.
Li, M., Soltanolkotabi, M., and Oymak, S. Gradient de- scent with early stopping is provably robust to label noise for overparameterized neural networks. In Chiappa, S. and Calandra, R. (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 4313-4324, PMLR, 26-28 Aug 2020. URL https: pttctts:/prooeeeeeeceee /ptpceeeiigggssseeeeeeee mlr press// v1 v108/1i20j html.
Stochr, N., Gordon, M., Zhang, C., and Lewis, O. Localizing paragraph memorization in language models, 2024. URL https: www.aaaivvvv...... abs 19851.
Thudi, A., Deza, G., Chandrasekaran, V., and Papernot, N. Unrolling sgd: Understanding factors influencing ma- chine unlearning, 2022. httpst ptattps:/arxxi..orrt abs 13398.
Toneva, M., Sordoni, A., Combes, R. T. d., Trischler, A., Bengio, Y., and Gordon, G. An empirical study of example forgetting during deep neural network learning. arXiv preprint /rXiv:/8/2.05// 2018.
Liu, B., Liu, Q., and Stone, P. Continual learning and private unlearning, 2022. URL 2203. 12817.
org/ abs/
Triantafillou, E., Pedregosa, F., Hayes, J., Kairouz, P., Guyon, I., Kurmanji, M., Dziugaite, G. K., Triantafillou, P., Zhao, K., Hosoya, S., Junior, C. S. Dumoulin, V., Mitliagkas, I., Escalera, S., Wan, J., Dane, S., Demkin, M., and Reade, W. Neurips 2023 machine unleamn ing. https: com/ competitions earning,, 2023.
Maini, P., Garg, S., Lipton, 7., and Kolter, J. Z. Char- acterizing datapoints via second-split forgetting. In Koyejo, S., Mohamed, S., Agarwal, A., Bel- grave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, vol- ume 35, pp. 30044-30057. Curran Associates, Inc., 2022. URL ftt.ss//ppper//2222///tt oc ccc/paper_fillseeeeeeeeeeessssee
Kaggle.
A. Vladarean, M.-L., Pillaud-Vivier L.. and Flammarion, N. On the spectral bias of two-layer lin- ear networks. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https:
pdf.
Maini, P., Mozer, M. C., Sedghi, H., Lipton, Z. C., Kolter, J. 7.,, and Zhang, C. Can neural network memorization be localized?, 2023. URL aattps://arxtt org/ abs/ 2307. 09542.
/ //openreview.o
Zhang, R., Lin, L., Bai, Y., and Mei, S. Negative preference optimization: From catastrophic collapse to effective un- learning, 2024. URL https: http://wwwwww////eeeee///o /abs/ 2404.05868.
Nanda, N., Rajamanoharan, S., Kramar, J., and Shah, R. Fact finding: Attempting to reverse- engineer factual recall on the neuron level, Dec 2023. URL
Zoph, B., Bello, I., Kumar, S., Du, N., Huang, Y., Dean, J., Shazeer, N., and Fedus, W. St-moe: Designing stable and fact fact-finning--tttemptie to- to-revesse-eng www.ennnnnneenneeeeee tannfffrrrbll &xper models, 2022. URL https: / waaxiiv.org/aat
forum.
Nasr, M., Carlini, N., Hayase, J., Jagielski, M., Cooper, A. F., Ippolito, D., Choquette-Choo, A., Wallace, E., Tramèr, F., and Lee, K. Scalable extraction of training data from (production) language models, 2023. URL https: //apsi///rxiv wwwpxiiioorr////oooo org/abs/
08906.
abs /2111 17035.

---

Memorization Sinks: Isolating Memorization during LLM Training
A. Extended Related Works
Memorization in Classification Settings. A large body work has focused on unlearning or forgetting memorized information from neural models, especially in the classification domain. This includes methods such as SISA (Bourtoule al., 2021) that exactly unlearn information by maintaining multiple model copies, and recent influx of approximate unlearning approaches (Triantafillou et al., 2023) that aim to perform post-hoc procedures a model in order to remove information in question. The study of memorization in classification settings is connected to label noise and atypical examples (Harutyunyan et al., 2020; Feldman & Zhang, 2020). On the other hand, in this work, we concentrate on more natural or typical memorized points which often arise in the large language model setting.
B. Implementation Details of TinyStories Training
Implementation and Architecture. We use the nanoGPT library perform standard pretraining of the models. We train GPT-2-Medium like architecture with embedding dimension 1024 and times expansion in the MLP layer. We used 24 layers, the resulting model had approximately 344 M parameters.
Table I. Hyperparameter Tuning for Standard Training
Parameter
Values
Max Learning Rate
(6e-5,6e-4.6e-3)
Weight Decay
[Ie-5,ee3,,ee- Max Learning Rate
Min Learning Rate
LR Decay Steps
Total Training Steps
Hyperparameter Tuning, We set the hyperparameters for our training as shown in Table For parameters denoted in sets, we tuned over choices of these parameters relative to the validation loss. We also performed early stopping on the validation loss, but generally found that overfitting did not occur.
C. Implementation Details of Post-hoc Localization Techniques
We generally follow the methodology used in (Chang et al., 2024b) and directly used their code as released online. We restrict our attention to their Hard-Concrete and Integrated Gradients methods presented in the papers.
Hyperparameters: Hard Concrete. We tuned x, the & loss coefficient used in training the pruning mask M over the values (100, 500, 1000) on a tuning set of sequences. Additionally, we tuned the number of pruning iterations in the range {1000, 2000, 4000). The remainder of hyperparameters were set to the optimal values reported by (Chang al., 2024a). We tuned relative to the lowest validation loss achieved after dropping out the identified neurons.
Hy Hyperparameters: Integrated Gradients. For Integrated Gradients, the only hyperparameter was the number of IG steps. As result, we set this to the value reported in the paper, 16.
Dropout Procedure. Following the computation mask scores by either Hard Concrete or attribution scores by Integrated Gradients, we sorted the neurons in each layer by these scores. Given a dropout parameter r, we dropped out an proportion of the neurons in each layer, as was performed in (Chang et al., 2024a).
D. Disentanglement in Standard Training
Setup. We will consider a dataset with n examples seen once Donce = { b: and a memorization sequence memem _ mem) that is repeated l times. Thus, explicitly, our training dataset consists of D Drre = Donce U (U&(s , mem)) )). For the purposes of this theory, we will assume that the Rd can be partitioned into a semantic subspace Stemantic
and a memorization subspace Smuse:; explicitly we will write that ((s)) (
(i)
(ts) m
We will make the
mem
assumption that Vi E [1, N] (0) TT (s(s (mem) = 0. This can be interpreted as enabling us to uniquely identify a
(1)
particular example in the noise subspace which would facilitate disentanglement of memorization. For convenience, we will consider the matrix implemented by W Wpwy Wte i and we will denote W = argmin pegmin,(w,,p) L(W [W]| (i.e. the
12

---

Memorization Sinks: Isolating Memorization during LLM Training
minimum norm solution when W when the memorization example is not seen). We will assume that rank (W*) > 2. We will assume that we can achieve 0 loss using only the subspace Ssemantc- In particular, denoting the singular value decomposition of W = U**,"*" this translates to the condition that (V*)(VV) C S Ssemantic- We will denote the term Amem = y f()(s)))) (the change in prediction necessary to memorize s sm and assume || || > 1. We will Vi E [1, rank [] (i.e. the generalizing solution has low norm) and mem
assume that all the singular values of oi (W*) < 2F that the memorization activations are unit norm ([(((s)
= 1). Finally, we will assume that the A 1 col(W*)
mem
Desiderata of an Unlearnable Model Inutitively, the problem of unlearning of efficiently finding a model which behaves as though it never saw the repeated example (i.e. recover model such that W WW") However, we also wish to accomplish in ways that are efficient computationally and do not require significant access to pre (as the pretraining dataset can often be inaccessible in downstream updates). This precludes simply retraining the model on the set Pre (s y. Intuitively, this is straightforward if the memorization of (smem .y mymem) makes use of the noise features in order to implement the memorization. Concretely, consider the solution W = =W" + Amem (cb(s )
) T. This solution implements memorization in a disentangled way by isolating it in a direction (in input-space) that is completely orthogonal to the semantically meaningful, generalizing subspace. As a result, (s m)m)) can be unlearned using updates that have no effect on the predictions of any (mem)
other sequence. This is becase we assume that Vi E [1 (Ø(s) (i) ) T((s) ( mem = 0 and thus for any unlearning update mom
uv where v E Smem: we have that Vi (W + uv ())(()) = ())(s())
D.1. Analysis of Natural Sequence Memorization
We will begin by describing the two possible solutions for memorizing
Definition D.1 (Disentangled Memorizing Solution). A learned parameter W implements the Disentangled Memorizing Solution if Wai can be written as W* + cuv where E Ssem and Was achieves 0 training loss on Ppr
This definition follows along from our previous discussion of the desiderata of an unlcarnable model. In particular, we constrain memorization to be stored using the orthogonal memorization component. This implies that the memorized sequence can be unlcarned simply by taking gradient steps which do not impact the predictions unrelated sequence. Hence, we are guaranteed not to distort the mechanisms responsible for the model's general capabilities.
Finally, we will discuss entangled memorizing solutions-those that require modifying the general capability subspace of the model and hence run the risk of incurring model degradation in subsequent unlearning.
Definition D.2 (Entangled Memorizing Solution). A learned parameter Wen
eni implements the entangled memorizing solution
if W Wen achieves 0 training loss on pre and W Went can be written as Σ
o. ofAa) A +
"iv
1 where the Vi,...V;
are the original right singular vectors of W and ü; are a potentially shifted set of singular vectors and ü,+1, Vi+1 are a new set of vectors and Vi E k] we have that of S <WW)): + A and ok+1 *+I
||اسل|| where k: rank(W**)
k+1
Intuitively, entangled memorizing solution reuses and shifts vectors that are used by the generalizing solution. We impose the condition on the singular values to model that the implementation of memorization is "diffused" across multiple of the generalizing semantically meaningful features.
Theorem D.3 (Memorization of Natural Sequences is Entangled).. Consider training a two layer linear network f(z) W [(z)=WpmjjWW ponjWWse on Dpre with the squared loss. Suppose that (smemm yymemm is natural. Let W be the final function learned by gradient flow and U* Emmon (V (ine!))" T be its singular value decomposition. Then, W Wine + W dis-
We will first provide some intuition for this statement. Intuitively, f could either memorize (s mem) by (a) making use of the noise component of (smm - thereby isolating the change in the model induced by memorization orthogonal to the generalizing features or (b) by implementing the memorization by shifting the generalizing features. We will consider the setting in which it is possible to implement memorization in both ways. However, unlearming memorization implemented with the generalizing features carries the risk of distorting general capabilities whereas memorization isolated to the orthogonal space can be easily and robustly removed.
As result, the primary intuition behind Theorem D..3 is to demonstrate that the training dynamics will prefer more harder to unlearn solutions when memorizing natural sequences.
Before providing the proof of Theorem D.3, we will restate an important result on the implicit bias of gradient flow in two-layer neural networks from Varre et al. (2023). We will first specify a set initializations for which the result holds.
13

---

Memorization Sinks: Isolating Memorization during LLM Training
Table 2. Hyperparameter Tuning for Sequence-Tice Tied Dropout
Parameter
Values
Max Learning Rate
(6e-5,6e-4,6e3)
Weight Decay
1e-5,,e-3,1e-1 Max Learning Rate
Min Learning Rate
10 Total Training Steps {0.7.0.9.0.95)
LR Decay Steps
9
Definition D.4 (Orthogonal Feature Initialization). A two layer linear network W proj Wie has orthogonal feature initialization if Wi / where P E [RAX|[PP I} for some
and W Wpay proj =
Theorem D.5. Consider training Iwo layer neural network parameterized by (W proj. Wie with gradient flow and orthogonal feature initialization. Let 3 W prooy We.- We will denote the set of optimal parameters for given pretraining dataset Drre as pDeee)
projs
1. The parameters converge to a global optima (i.e. lime-+oo (W Wmng(t) Wse(t)) E (DDpe)
2. The effective linear predictor converges to a min-norm solution, formally lime-+00 B(t) = 5** where 3 (DD,
Using the result from Theorem D.5, the proof of Theorem D.3 is straightforward and simply requires comparing the Frobenius norms of the disentangled and entangled memorizing solutions. In particular, because [|o(s()) ( )||2 = 1, we must have that
|[W dis 1 = | || Amem|) ?
(the remaining singular value comes from the orthogonal rank-1 term uv On the other
hand, consider the entangled memorizing solution. We have that Weaull !|: <
| +
or (W*);
||Ammall I+1 + |m S
k+1
|| 2 ااسسو <
W || + ||
Hence Wdis cannot be the minimum norm solution and gradient flow will not converge to it.
E. Implementation of Gradient Masking
We generally follow the implementation outlined in (Cloud et al., 2024). We partition each MLP layer into memorization and generalization neurons. We tune this delineation of memorization and generalization neurons by the proportion of generalization neurons We additionally partition our dataset into examples seen once and the repeated examples. During training, we mask the gradients in each MLP layer such that the gradients from the repeated examples update only a the memorization block, whereas gradients of all other examples are routed to the generalization block.
Hyperparameter Tuning. We show the hyperparameters tuned for this method in Table 2. Hyperparameter denoted in sets are tuned relative to the validation loss before dropping out memorization neurons.
F. Analysis of Gradient Masking
Data Distribution. Like the previous analysis, we will consider a dataset of N (s;, (s) pairs which are only seen once, and a special (Smem: Ymem) that is repeated training and we wish to isolate. Thus, the total training set D ((s,, N U (Smem, Vmnem).
Model Structure. We will use the same two layer neural network structure f(z W proj Wt However, for simplicity, we will consider only training the second layer (i.e. W Wroj- Further more, to implement gradient masking, we will partition the hidden space of the MLP into two components
memorization component W mem and generalizing component W mem Thus the prediction
W gen- In particular, we will assume the structure W proj [ gen
f'xx) could be written as W gen (W + W mem ( Wem))mem- Here, we use the shorthand (Wr.r)gen to denote the entries of the hidden layer activations that are input to the generalization neurons, while (Wicz) denotes those that are input to the memorization neurons. We will denote the dropped out model fa as the model with the generalization neurons dropped (zero-ed) out. Concretely fa = W gen
0] X
14

---

Memorization Sinks: Isolating Memorization during LLM Training
Training Procedure. We introduce the following simplified version of the gradient masking scheme examine empirically,
Concretely, prescribe the following SGD--lke update rules for W gen and W mem- given learning rate ry To keep track of the parameter values at different time-steps, will use superscript.
WT + gen
7 wg. [(f(x)))) if # 8mem
WT+1
gen
W
T
otherwise
gen
Similarly, we will consider the following for W
mem:
WT mem
+ VV wm.. [ffo((),
if Smem
WT+1
mem
W mem
otherwise
In both cases, consider C be the squared loss function. Intuitively, W mem receives gradient only from the mem and no other examples whereas Wgen gen receives gradient from all but Smem- Thus, by training construction, we could say that the mapping (5mem: Umem) is stored only in W mem as its gradient cannot change any of the other parameters.
Ground-Truth Unlearned Model. When performing unlearning, we wish to obtain a model that behaves as if it had never observed (Smem, ymem). We denote such a model (trained under standard training with parameters 7 as f).
Initialization. We assume that W gen (0) = 0 and that W mem (0) O. We consider Wie is initialized at a non-zero value and does not change throughout training.
We will assume that both f and f observe the data points in the same order (i.e. (81, y1),, ..., (N Bloaa) We will assume that I observes (Smem: Umem) once timestep T. With the setup in place, model fa and the
Order of Observations.
(81,91), will now introduce a lower bound in the difference between the predictions of dropped ground-truth truth unlearned model f. Concretely, we show:
Theorem F.1 (Removing Memorization After Gradient Routing). Suppose that we train f on DLM U (Smem- Umem). Suppose that there have been N gradient steps taken since the last time W mem updated. Additionally, assume that is trained using the same hyperparameters but without observing (smem, yyme)))
[a(x) - f(x)) 2 (N)(1->))* ||x||2
where c mins.j 8, sj > O.
We now proceed with the proof:
Proof We will use the notation W Wmem; W gen to denote the parameters of f and W mem: w (which does not observe the (5mem: Mmem) but sees all other points in the same order.
gen to denote the parameters of f
Observe that, in the case of f. W mem remains constant throughout the training process does W mmo 0 by Equation F. Next, recall that by Equation F, we have that WT gen
WT-1 We then have the following.
N-1 (T)
N-1
N-1
W gen
(T+N-1)
W gen
II [I 18,8 )) + 7
YESk
II (I (s,s))))
j=k+1
W
mem S
II (I 7y,8)) j=k+1
k=1
and
N-1 (T--1)
N-1
N-1
(T
(1-
+1
Σ
Y
II (I - 71
gen
gen
j j=*+1
Note that since WT WT- W a and Wan an is calculated as
and the order of observing data points is the same the parameter-space difference between
all
N
N-1
W
W mem S
II (I 7s,5,)
gen
gen
j=k+1
15

---

Memorization Sinks: Isolating Memorization during LLM Training
Next observe that
1 ( f (x) (W (Wi++n1) (T+ + * W x
gen
mem
(W ( W W x
gen
(W (W[T+n11
w (TT-n-2)) X
gen
gen
N-1
W mem S/5k)
II j=k++
(I 7s,8, 1)x
Thus, we wish to lower bound
N
N
W mem StSS
T)
II (I - 78,8] ))x j=k+1
First we note that of eigenvalues use this to rewrite the above term as:
N N-1 j=b+1 (I - 7 }}x E span( [ and by the repeated application of variational characterization
[| j= (15--,8,) ))x{11 2 (1 ~1 |xx|2- Now, we will recall that W mom oty memss T
mem- meme
We will
N-1
N-1
1 [|||skll>
Σ
(S
II jmk++
(I (Dsss,,))
Upon repeated application of the prior result and Lemma 2, we then have that
N-1
N-1 (1 (1-7) 1
(s S T
II
(I -
78)5])x)11 1
j=++1
N-1
(1 - 7) N-1 ||x||| c²
k=1 2 (N 1)(1 -7) |x|1| 62
by positivity of the summands we can eliminate the absolute value symbol in step 2 and we have used that for all i N - (1- 2 (1 >- which follows from y < 1. This yields the desired lower bound.
Lemma F.2 (Lower Bound of Outer Product in span(fki,... (k1,... k₂))), Suppose ming, k. k, span( (k1, ...; k,)), and consider i, E [1, Then
>
0. x
kj x|12 2 c ||||||
Proof k3 x is a scalar so we have that kkk x||2 = ||kl|| xl. x).
Since x is in the span of (ki}, we have that [k] x 2 2c|1x
G. Implementation of MemSinks
Model Architecture and Implementation. We used the same model architecture as reported in Appendix B. We set the first g fraction of neurons in each MLP as the "shared neurons" and left the remaining g fraction the memorization neuron pool. We applied the dropout layer after the GeLU activation function, prior to the downprojection layer.
Assignment of Sequence IDs. We sequentially numbered the sequences in the TinyStories training set and use these indices as the sequence IDs.
Hy Hyperparameter Tuning. Table 3, we show the hyperparameter ranges tuned over for MemS inks. Hyperparameters denoted in sets were tuned over using the validation loss when the memorization is dropped out.
16

---

Memorization Sinks: Isolating Memorization during LLM Training
Table 3. Hyperparameter Tuning for Sequence-Tice Tineeneee-eeed Dropout
Parameter
Values
Max Learning Rate
(6e-5,6e-4,6e3)
Weight Decay
1e--,,ee3,.1e-11 Max Learning Rate
Min Learning Rate
10 Total Training Steps (0.1.0.3.0.500.) (0.1,0.3.0.500.7)
LR Decay Steps
g
p
H. Analysis of MemSinks
H.1. Formalization of Training Process
Architecture. For simplicity, we study the training dynamics of an MLP layer f(x) = W proj Wax,, where W prooj Rddxduus Wi: E Rdd Here, demb refers to the embedding size of the model and d, refers to the number hidden neurons in the MLP. Given sequence s, we consider that f takes in the final position embedding of s, which we denote o(s) and directly outputs the logits of the next token (i.e. softmax(f(o(s)) is probability distribution over the next token in sequence
Wryy
For convenience, we will denote the hidden activations of sequence as (s). In our analysis, we will assume that the activation space of z(s) can be split into two subspaces z(s) = (z($) ( stared z(s) m These components will correspond to our choice of shared and memorization neurons. We will additionally consider Wte frozen throughout training and mainly study the training dynamics of W proj. Thus for convenience, we will also decompose Wie into two column-blocks (corresponding to the shared and memorization neurons, respectively):
W pmj
W
Wmee
proj
proj
Data Setup. We will treat our data (embedding, next token) pairs. We consider we have repeated sequence memn with corresponding next token mem, Next, we will assume we have large dataset of sequences seen only once during training Donce = (('))) (s 3)3. For simplicity, we will consider the case where Vi ( o e Since we treat W proj as frozen, we will also define Eshared = min z(s mem) (atame)) 2s shured and likewise that mem = min; z(s"em) ) (s)( mem- Intuitively, these quantities lower bound how similar the activations in the shared and memorization neurons are between the repeated example and any other example. For simplicity we will assume that the ||z('||
for all ( and that the parameter [[ pmujl/2 < Cyoe remains bounded throughout training. Finally we assume that the ouput embeddings e are mutually orthogonal. 2
standard training, we study the training trajectory (with learning rate 7) of minimizing the cross entropy loss with respect the parameter W Wroj-- We consider training with batch size 1.
Training Process.
H.2. Forgetting Under Normal Training Dynamics
To begin, we introduce result on the softmax with bounded inputs
Theorem H.1 (Sofumax on & bounded vectors). Consider E IR and suppose || S C. Then max, (o(x)) S d-I and (o))()
Proof. (x);
exp(ars) exp(ay)
exp (d)))((dd))xppcc
exp (C))
(2C
Likewise a(x); >
2C +d-+
exp(-C) exp(-C)+(dD) exp(C)
exp(-20) exp(-20)+(d-1)
exp(-20) (C))
Given our assumption that 1 1012 and the bounded parameter norm assumption W
projlla < C.ss it follows that 2
||W poojz( < Sy. By Theorem H.I, have that the entries of exp(-Cma) < (f((z()) the remainder the theory, we denote Cynn
exp(Crm) element wise. In donb
dant
expí Cpm)) and Cyna deate
exp(C)ea))
deub
We will first show that the memorization of the repeated sequence Smem is forgotten when we take intervening steps on non-repeated sequences (9) ,..., ((+t)) Formally, we have the following proposition. Formally, suppose that after step i,
17

---

Memorization Sinks: Isolating Memorization during LLM Training
have just seen sfmem), Then we will show that the logit (mmem) decreases during subsequent training steps i through i + nt. For this analysis, we will focus the dynamics the shared neurons.
Theorem H.2 (Forgetting in Standard Training). Suppose take a gradient step s m/me)) at gradient step i and subsequently make gradient updates on non-repeated sequences s ...,..(ttm,, After the m gradient steps, we have that (emee)) [ f(t+m) (ymem) S (emen)) (10)) 71M66min
Proof. Only the parameter W Wmjj changes throughout training, so we can restrict our attention to its dynamics. We have that the gradient of W proj on the sequence-next token pair (z, e)
aL OW proj
(e - ((f(())
Now let W proj denote the parameter value of W proj after the i-th observation. We have that
W proj
(++m)
W
(1)
(e - (2(()))2))))2
(1)
proj
where we will denote f(i+t) as the model with parameter W proj- Then, have that the logit the correct next token for memorized example mem is
m
mem) T (2
(e (emem) T (x)(mem) + memem) T
Σ (el)) - (z mem oo (z()
mem)
j=1
Now, since we have that the token embeddings are orthogonal, we can rewrite this as
(e mmm)) T ((++m) (2 mem) = (e T (z) (x mem
(z (z (2 mmm)
Note that by the assumption of bounded norm for W proj- we have that (eme T o off(iti) 2 Cmin (defined earlier).
Note also the assumption that z (i)
(z mam)
2 e Vi. This implies that
m Σ
(e mem T (2+tm))mmmmm S (e mem) T mom)
cCmin
(2)
This immediately yields our desired claim.
Next, we will show that the seqTD accumulates memorization in the memorization neurons, as formalized in the following theorem. This theorem also crystalizes some key quantities relating to gradient interference. First of all, we see that the forgetting depends on the number of further gradient steps taken after seeing me Secondly, we observe that the impact of forgetting dynamics is influcned by how similar the activation of neurons are amongst different examples: controlled by e. The first observation immediately suggests that if some neurons were activated less often, then those neurons would be effectively "store" more memorization.
H.3. Analysis of MemSinks
Theorem H.3 (MemSinks Accumulates Memorization in Memorization Neurons). Consider training MemSinks, where the memorization neurons are activated p fraction of non-repeated examples. We will assume that the model is trained from 0 initialization. Denote the MLP fmem-dropped as the model with memorization neurons dropped out and gen-dropped as the model with the generalization neurons dropped out, Suppose that the model is trained for N total steps and the repeated sequence mom is observed k times. Then we have at the end of training
18

---

Memorization Sinks: Isolating Memorization during LLM Training
1. (emem) T ony ( (d(s)))))) S 7k(1 Camin) - 7(N siired
2. (emom) mem-oonl (oleently((((() ( 2 7k(1 Cmax) Y(N k) pp Cmax
where Cmin and Cmaa constants depending on an upper bound of the parameter norm of W pmj--
Proof. Our argument resembles the proof of Theorem H.2, and we will rely on the intuition therein. For reference, we will write the gradients for the components of W proj below.
aL
= (e ozf(z)) Zhared
UW shared proj
and likewise
aL
(e - of off())
W mem
-ontynnyyyy genn-onde (2 mmmm At any point in training, recall that upper and lower bound the value Cmin S (T (f (2mmm) I Cmaz- As such, observe that (cl)) (f(mmm) received updates upper bounded by 7(1 Criin) (from the i obervations of mem and (N - k) updates upper bounded by mman (from the remaining
We will first examine (emmm) T p(nd
(N k)) observations of the ( This yields the desired claim for (1).
Now, for claim (2) observe that the component (em)m)) T f((n) (again, from the kt observations of m but only p(N - k) updates from other observations, which can likewise be lower bounded by mmmm This immediately implies the desired claim (2)
mmtm-onty (zmem) receives A updates lower bounded by (1 - Cmaz)
This theorem formalizes the notion that memorization "accumulates" in the memorization neurons when they are shielded from the interference of other sequences sufficiently. In our theory, the extent to which this occurs is dependent on two quantities (1) the fraction of non-repeated sequences for which the memorization neurons active and (2) the similarity of activations of the repeated example and non-repeated example in the memorization neurons. Relative to algorithm design, however, we will generally only have control over and so we will consider ésared Emem out of convenience. Our analysis demonstrates that when p is set appropriately low. Some calculation demonstrates that when
Curin. [N-k) (Cmax Cmin), then we will have a seperation in the logits of mem where the memorization neurons primarily contain the memorized example.
I. Large Scale Implementation Details
Here, we report all experimental and implementation details for the large-scale experiments conducted in this paper
Pretraining Data Mixture. We consider pretraining datasets of 1 and 2 billion parameters (respectively for the 360M model and 1.7B model). The majority of pretraining tokens are sourced from a subset of the SlimPajama dataset (Shen et al., 2024). However, we consider setting in which we have an under-sampled domain of 5000 TinyStories examples. We ensure that these 5000 TinyStories (Eldan & Li, 2023) documents are included in the pretraining data before filling the remaining tokens from SlimPajamas.
Table Hyperparameters for Standard Training (360M)
Parameter
Values
Weight Decay
0.1
Max Learning Rate
5e-4
Min Learning Rate
Max Learning Rate
LR Decay Steps
Total Training Steps 1024
Batch Size
19

---

Memorization Sinks: Isolating Memorization during LLM Training
Table 5. Hyperparameters for Standard Training (1.7B)
Parameter
Values
Weight Decay
0.1
Max Learning Rate
5e-5
Min Learning Rate
Max Learning
LR Decay Steps
Total Training Steps 1024
Batch Size
Model Architectures. We pretrain using SmolLM family models size 360M and 1.7B. For standard training runs, we do not modify the architecture any way. We describe the modifications necessary for MemS inks below.
Training Hyperparameters. We report the hyperparameters used for training the models (at size 360M and 1.7B, respectively) in Tables 4 and 5, For fairness, we use the same hyperparameters train MemSinks.
L.. Implementation of MemSinks
We will release a full implementation of MemSing s at mentation builds off of the LitGPT library (AI, 2023).
Our imple-
Assignment of Sequence IDs. We assign each document (example) in the pretraining corpus with sequence ID that is hash of the tokens within that document. We implemented a tensorized hash function in pytorch which will be released with our code.
Tokenization and Packing. When implementing MemSinks, is necessary to imput the sequence identification meta-data into the forward pass of the transformer model. In order to accomplish this while still leveraging the optimized streaming data loaders typically used in pretraining, we interleaved the sequence IDs into the stream of tokens saved during the tokenization step. That is, the pretraining data stream contained tokens every even position and the previous tokens ID code at every odd position. This enables seamless loading of the sequence ID data alongside the tokens during training. During training, the sequence II) data and tokens are separated using simple indexing operation. As opposed to the smaller-scale TinyStories setting, we pass sequence ID with each token to enable training across document boundaries and prevent the necessity padding tokens.
Online Computation of MemSinks Masks. Another crucial challenge faced when scaling MemS nkks is the need to deterministically generate neuron mask for each token in the sequence. While we used the standard Pytorch seeded random number generator during our smaller-scale experiments, this becomes impractical when different tokens in batch have separate masks. As Pytorch does not support hatched-seeding in its random number generator, we implemented a tensorized linear congruential random number generator to generate neuron masks online during training. This enabled us to avoid needing to precompute masks across a billion token training set.
Placement MemSinks. We implemented MemS inks at every hidden MLP layer by partitioning the neurons into generalizing and memorization sink groups. We then apply the mask over the memorization neurons. For efficiency, used the same masking across all MLP layers. As the SmolLM (Allal et al., 2025) family models uses GatedMLP implementation, apply MemSinks in the output-space after the gating.
MemSinks Hyperparameters. For MemSinks, for both model sizes, we tune across hyperparameter choices given in
Table 6. Hyperparameters for MemSing inks (both sizes)
Parameter
Values
9
{0.05, 0.1, 0.3, 0.5}
{0.1,0.3.0.5)
Table
20


Reviews:
[{"review_id": "vfAWpbz4ub", "reviewer": "Reviewer_Q8Dm", "summary": "The paper introduces a training strategy called Sequence-Tied Dropout (SeqTD) for large language models that aims to isolate memorized sequences into a specific subset of neurons while still allowing the model to learn general language patterns. The authors argue that standard training causes memorization to be entangled with general knowledge, making post-hoc removal of sensitive or copyrighted text problematic. The proposed method enforces a consistent dropout mask for repeated sequences based on sequence IDs, which channels memorization signals into designated neurons. Experiments on a modified TinyStories dataset demonstrate that SeqTD can “unlearn” repeated sequences effectively without degrading overall model performance, and theoretical analyses are provided to support the observed dynamics.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful and constructive review. We're glad that:\n- (i) you found our **theoretical analysis sound**,  \n- (ii) appreciated the **clarity and motivation of our experiments**, and  \n- (iii) recognized the novelty of using dropout for disentangling memorization from generalization.\n\nWe've added new experiments and clarifications that we believe directly address your suggestions.\n\n---\n### **Theme 1: Experimental Scope & Generalization**\n> *\"Experiments are limited to a small-scale dataset and do not explore more challenging benchmarks.\"*\n\nTo expand the scope and test robustness, we've added:\n\n🔗 *Full results available here:* [Model Scaling Experiment Results](https://shorturl.at/md9cQ), [Token Scaling Experiment Results](https://shorturl.at/Aa7VZ),  [Metrics Besides Loss](https://shorturl.at/M7UfM)\n\n- **Model scaling experiments** across four model sizes: 32, 51, 170, 344 million parameters. Our results show that SeqTD mitigates memorization across all model sizes but the benefits of SeqTD increase with scale (no model degradation at 344M parameters).\n- **Token Scaling Experiments**  We train a 350M parameter model on a larger-scale corpus of 1 billion tokens containing a mix of TinyStories and SlimPajama. Our results demonstrate that SeqTD mitigates memorization (reducing the memorization-validation loss gap by 2.5x) while outperforming a data-deduplication baseline on validation loss.\n- **Metrics Besides Loss** We validate SeqTD using generation-based extraction metrics of memorized sequence *token accuracy* and *memorized sequence perplexity rank*. We demonstrate SeqTD drops token accuracy on memorized sequences from >90%  to <10%, suggesting robust mitigation of memorization.\n\n\n---\n### **Theme 2: Clarifying Mechanism via Theory**\n> *\"More practical insights into the theoretical assumptions.\"*\n\n We've surfaced the core intuition behind our theoretical results through the following simplified theorem (based on the theory introduced in Appendix E, page 14)\n\n> **Simplified Theorem (E.3):**  \n> *When a set of memorization neurons is activated consistently across repetitions of a given sequence and less frequently on other sequences, sequence memorization accumulates in these neurons and away from neurons activated on all sequences (shared neurons).*\n\nThe key idea is that memorization neurons (by virtue of being less frequently activated) are shielded from interference from other gradient updates. As a result sequence memorization accumulates and is preserved in these neurons. On the other hand, memorization stored in the generalization neurons experiences forgetting dynamics and is eliminated as a result. **We empirically verify this proposed mechanism in Figure 6 of our paper by showing that SeqTD experiences smaller amplitude learning-forgetting cycles.*\n\n---\n### **Theme 3: Clarifying Evaluation Objectives (CotaEval)**\n> *\"Evaluation would benefit from using standardized copyright benchmarks like cotaeval.\"*\n\nThank you for the suggestion. While we agree with the **long-term motivation around copyright and privacy**, we want to clarify that this paper is a **foundational investigation into training dynamics**, not a compliance benchmark. The datasets used (e.g., TinyStories) are synthetic or constructed for studying memorization patterns—not real-world content where cotaeval would apply.\n\n\n---\n### **Theme 4: Literature Review and Broader Context**\n> *\"Expand discussion to include recent works on copyright compliance.\"*\n\nWe agree and have expanded our Related Work section in our internal draft to include the copyright compliance literature you suggested. Thanks for the suggestion!\n\n\n---\n**Table 1: Summary of Actions for Reviewer Q8Dm**  \n| Concern                      | Action Taken                                             | Expected Outcome                              |\n|------------------------------|----------------------------------------------------------|------------------------------------------------|\n| Small-scale experiment scope | Added model-size sweep + 1B-token training               | Demonstrate scaling properties and robustness         |\n| Evaluation depth             | Added perplexity rank + token accuracy metrics           | Provide stronger insights on model behavior    |\n| Mechanism clarity            | Moved and Simplified Theorem E.3             | Integrate theoretical foundation with empirical observations     |\n| CotaEval request             | Respectfully deferred; explained paper focus             | Clarified evaluation scope and contribution    |\n| Literature review gaps       | Added broader copyright/unlearning papers                | Situate paper in practical context             |\n\nWe believe these additions substantially strengthen the paper while maintaining its focus on the core theoretical and empirical contributions. We appreciate your thoughtful suggestions that helped us improve the work's clarity, scope, and connection to the broader literature.", "claims_and_evidence": "Claims: The paper claims that (1) traditional training leads to entangled memorization and generalization, (2) post-hoc localization methods such as pruning and gradient attribution are insufficient for typical repeated sequences, and (3) SeqTD can isolate memorization into dedicated neurons with minimal impact on overall performance.\n\nEvidence: The authors support their claims with experiments comparing validation loss and sequence loss under different settings and provide theoretical analysis to explain the learning/forgetting dynamics. However, while the evidence is convincing for the controlled TinyStories setup, the experiments are limited to a single, small-scale dataset and do not explore more challenging benchmarks.", "methods_and_evaluation_criteria": "The methodology of partitioning MLP neurons into shared and memorization pools and enforcing a consistent dropout mask per sequence is clearly described.\n\nThe evaluation focuses on two key criteria: the loss increase on repeated sequences (indicating successful unlearning) and the validation loss (indicating generalization).\n\nAlthough these metrics are reasonable, the evaluation would benefit from incorporating additional standardized benchmarks—such as those similar to cotaeval—to more comprehensively assess the impact on copyright compliance and overall performance.", "theoretical_claims": "The paper provides a series of theoretical results (e.g., Theorems E.1–E.3) to formalize the dynamics of memorization and forgetting under standard training and SeqTD.\nWhile the proofs appear sound under the stated assumptions, a more detailed review of the derivations is necessary to fully validate the claims. No major issues were detected, but additional discussion on the assumptions’ practicality in real-world scenarios would strengthen the work.", "experimental_designs_or_analyses": "The experiments on TinyStories illustrate the core ideas of SeqTD. However, the experimental design is limited in scope; testing on only a small-scale dataset does not fully demonstrate the method’s scalability or robustness in real-world applications.\n\nAdditionally, more thorough comparisons with alternative methods (beyond the baseline post-hoc localization techniques) would help in assessing the relative strengths and weaknesses of SeqTD.", "relation_to_broader_scientific_literature": "The paper is well-situated within the literature on memorization in large language models and unlearning methods. It addresses issues raised in works on model memorization and the challenges of removing sensitive content from pre-trained models.", "essential_references_not_discussed": "the literature review could be expanded to incorporate recent studies in copyright compliance and evaluation frameworks, such as:\n\n1. “Copyright Violations and Large Language Models”\n1. “Foundation Models and Fair Use”\n1.  “Evaluating Copyright Takedown Methods for Language Models”\n1. “LLMs and Memorization: On Quality and Specificity of Copyright Compliance”\n1. “SHIELD: Evaluation and Defense Strategies for Copyright Compliance in LLM Text Generation”\n1. “Digger: Detecting Copyright Content Misusage in Large Language Model Training”\n1. “Speak, Memory: An Archaeology of Books Known to ChatGPT/GPT-4”\n1. “Avoiding Copyright Infringement via Large Language Model Unlearning”\n1. “CopyBench: Measuring Literal and Non-Literal Reproduction of Copyright-Protected Text in Language Model Generation”\n1. “Preventing Verbatim Memorization in Language Models Gives a False Sense of Privacy”\n\nIncorporating these references would provide a broader context and strengthen the paper’s discussion regarding practical applications and limitations in copyright-sensitive domains.", "other_strengths_and_weaknesses": "Strengths:\n\n- The paper proposes a novel and conceptually clear method for disentangling memorization from general model capabilities.\n- The theoretical analysis supports the empirical findings, and the idea of sequence-specific dropout is interesting and innovative in this context.\n\nWeaknesses:\n\n- The experimental evaluation is limited to a small, controlled dataset (TinyStories) and does not test scalability or robustness in more realistic settings.\n- The literature review could be expanded to better position the work within recent developments in copyright compliance and related evaluation methods.", "other_comments_or_suggestions": "It would be beneficial to include experiments on larger datasets or more realistic benchmarks to validate the general applicability of SeqTD.\n\nConsider incorporating a standardized evaluation package (e.g., cotaeval or similar ones) to strengthen the experimental analysis.\n\nExpanding the literature review to discuss related works in copyright evaluation and unlearning methods would improve the context and relevance of the paper.", "questions_for_authors": "See above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "And96L6WwB", "reviewer": "Reviewer_hZ7x", "summary": "This work tackles privacy risks in LLMs from memorizing repeated sequences. Current post-hoc neuron isolation methods fail for data entangled with general capabilities. The authors propose SeqTD, a training method splitting neurons into shared (generalization) and memorization groups. By activating fixed memorization neurons for repeated sequences and shielding shared ones, SeqTD enables precise removal of memorized content without performance loss.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive and balanced review. We're glad you found:\n- (i) the **formulation of SeqTD clear and valuable**,  \n- (ii) our **experiments well-designed within the scope**, and  \n- (iii) the **motivation around privacy and unlearning promising**.\n\nWe've carefully addressed your main concerns through new experiments and clarifications.\n\n---\n### **Theme 1: Scalability and Real-World Applicability**\n> *\"The absence of benchmarks on larger models or real-world data limits practical insights.\"*\n\nThis is an excellent point. To address this concern, we've conducted **additional scaling experiments** across multiple model sizes and more realistic training regimes:\n\n🔗 *Full results available here:* [Model Scaling Experiment Results](https://shorturl.at/md9cQ), [Token Scaling Experiment Results](https://shorturl.at/Aa7VZ)\n\n- **Model scaling experiments** across four model sizes: 32M, 51M, 170M, 344M parameters. Our results show that SeqTD mitigates memorization across all model sizes, but the benefits of SeqTD increase with scale (no model degradation at 344M parameters)\n- **Token Scaling Experiments**  We train a 350M parameter model on a larger-scale corpus of 1 billion tokens containing a mix of TinyStories and SlimPajama. Our results demonstrate that SeqTD mitigates memorization (reducing the memorization-validation loss gap by 2.5x) while outperforming deduplication.\n- **Cross-architecture validation** To ensure the approach generalizes beyond a single model family, we perform experiments on the SmolLM-2 model family which uses a Gated MLP in contrast to the GPT-2 models trained in our original submission.\n---\n### **Theme 2: Intuition for Disentanglement Mechanism**\n> *\"The paper lacks sufficient intuition about how SeqTD disentangles memorization.\"*\n> \nThank you for highlighting this critical gap in our explanation. We've significantly enhanced the paper with clearer intuition about how and why SeqTD works by elevating the core insight from our theory in Appendix E (page 14):\n> **Simplified Theorem (E.2):**  \n> *When repeated sequences consistently activate a fixed subset of neurons (via sequence-tied dropout), shared neurons gradually forget these patterns due to interference from other examples, while memorization neurons retain them.*\n> \nThis theorem (pg. 13) formalizes the key mechanism: memorization neurons (activated less frequently) are shielded from forgetting during pretraining and accumulate sequence memorization. Thus, SeqTD can control where memorization accumulates by controlling how often neurons are active in pretraining.\n\n---\n### **Theme 3: Comparison to Gradient Routing**\n> *\"The authors may want to include more discussion on how SeqTD compares with Gradient Routing [1].\"*\n\nWe appreciate this suggestion and have a dedicated comparison section in Section 4. In particular we discuss:\n- **Empirical Comparison**: As shown in Section 4.1 (pg. 4-5), our experiments with gradient routing demonstrated that such approaches can hinder cross-sequence learning and don't fully isolate memorization, whereas SeqTD preserves general capabilities while enabling more effective post-hoc removal\n- **Mechanistic Differences**: As discussed in Section 4.1, there are two crucial differences between SeqTD and Gradient Routing. Firstly, SeqTD allows for a pool of shared neurons (that are updated by all sequences) and our empirical results show this is crucial for generalization across sequences. Secondly, the dropout performed by SeqTD disincentivizes co-adaptations between general and memorization neurons. In gradient routing, we empirically find such co-adaptations develop which leads to an additional drop in performance after removing neurons.\n\n---\n### **Theme 4: Miscellaneous Improvements**\nThank you for your helpful suggestions about presentation consistency. We duly take note of all of them and they have been updated in our internal overleaf draft.\n\n---\n**Table 2: Summary of Actions for Reviewer hZ7x**  \n| Concern                        | Action Taken                                          | Expected Outcome                         |\n|--------------------------------|-------------------------------------------------------|------------------------------------------|\n| Lack of scaling results        | Added experiments on model and token scaling | Demonstrated improved effectiveness at scale |\n| Intuition for disentanglement | Added simplified theorem and integrated with empirical observation| Clearer mechanistic understanding |\n| Gradient Routing comparison    | Added detailed empirical comparison and discussion of mechanistic differences | Clarified relationship to prior work |\n| Presentation consistency | Standardized formatting, improved figures, fixed typos | Enhanced readability |\n\nPlease let us know if these revisions strengthen the paper and address your concerns about scalability, theoretical foundations, and comparisons to prior work. We're grateful for your thoughtful feedback!", "claims_and_evidence": "The main problem that SeqTD addresses is the challenge of isolating and removing memorized sequences from language models while preserving their general capabilities (localization). Existing post-hoc methods prove ineffective or significantly degrade model performance when memorized sequences are statistically similar to the broader training distribution. SeqTD offers a training-time approach that successfully identifies and removes such problematic memorization while maintaining the model's overall generalization abilities, demonstrating superior performance compared to previous methods.", "methods_and_evaluation_criteria": "The evaluation approach is comprehensive and effective, employing appropriate metrics to assess sequence forgetting and model preservation capabilitie. For the method part, the authors may need to provide more insights why pretraining is better than post-hoc. Besides, the authors may also want to have more analysis about the difference of their method and [1].\n\n\n[1] Gradient Routing: Masking Gradients to Localize Computation in Neural Networks, arxiv 2024.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Experiments are rigorous in controlled settings: repeated vs. atypical sequences, ablation studies on dropout ratios, and noise tolerance . However, the absence of benchmarks on larger modelsor real-world data limits practical insights. The comparison to post-hoc methods is a strength, but including more recent baselines would enhance relevance.", "relation_to_broader_scientific_literature": "This works seem to be a new approach for localization of memorization. It may help privacy research for achieving a better trade-off between unlearning and performance.", "essential_references_not_discussed": "[1] is cited but more discussion is needed.\n\n\n[1] Gradient Routing: Masking Gradients to Localize Computation in Neural Networks, arxiv 2024.", "other_strengths_and_weaknesses": "This is generally a new and valuable method with comprehensive evaluation for the presented model. However, several critical aspects are underdeveloped: \n1. deeper insights into why the training-time approach outperforms post-hoc methods would strengthen the theoretical foundation; \n2. scalability questions remain regarding larger models and diverse datasets beyond TinyStories, raising concerns about the method's applicability to production-scale LLMs; \n3. the paper lacks sufficient intuition about how SeqTD effectively disentangles memorization from generalization capabilities, which are often deeply intertwined in LLMs. While the implementation details and empirical results are thoroughly presented, the underlying mechanisms and theoretical guarantees for successful disentanglement deserve more attention.", "other_comments_or_suggestions": "I suggest addressing several minor inconsistencies in the manuscript: standardize the formatting of quotation marks and italics, currently some of them are confusing; maintain consistency with hyphenated terms like \"trade-off\"/\"tradeoff\" and \"pretraining\"/\"pre-training\"; and convert rasterized figures to vector graphics to improve visual quality and readability.", "questions_for_authors": "What mechanisms explain how your training-time approach successfully disentangles memorization from generalization capabilities in language models, and how might these insights scale to larger models trained on diverse datasets?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wIXv4eeMTW", "reviewer": "Reviewer_hkbC", "summary": "The paper presents an investigation into sequence memorization in large language models (LLMs) and introduces Sequence-Tied Dropout (SeqTD) as a novel method to isolate memorization while maintaining generalization capabilities. The key argument is that typical memorization is not confined to specific neurons in standard training, making it difficult to remove without affecting overall model performance. The proposed SeqTD approach mitigates this issue by partitioning neurons into shared and memorization-specific groups, ensuring that memorization accumulates in a fixed subset of neurons while allowing general knowledge to remain broadly distributed.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed and insightful feedback. We really appreciate that you found:\n- (i) SeqTD to be a **simple yet novel idea**,  \n- (ii) the **methodology rigorous**, and  \n- (iii) our **ablations helpful and revealing**.\n\nWe've made specific additions to respond to your concerns.\n\n---\n### **Theme 1: Figure 4 (a) Validation Loss Dynamics**\n> *\"Figure 4(a) shows that validation loss without repetition is lower than with repetition. Can you clarify why this happens?\"*\n\n🔗 *Full results available here:* [Token Scaling Experiment Results](https://shorturl.at/Aa7VZ)\n\nYou've identified a crucial insight. In our initial setup, we treated data intervention (deduplication) as separate from model intervention (SeqTD), assuming only model interventions were appropriate for our analysis of disentangling memorization and generalization in LLMs. In Figure 4(a), the \"without repetition\" line represents an oracle baseline that involves direct data manipulation (deduplication). We initially considered this an \"unfair\" comparison because:\n\n- Our research focuses specifically on model-based interventions that can be applied without modifying training data, in order to understand the dynamics of generalization and memorization better.\n- Data deduplication was hence considered an oracle that represents an ideal but unreachable solution for the minimum validation loss on the trainig distribution.\n\nTo address this ambiguity, we've conducted new experiments with a more realistic setting where repetition is actually beneficial to generalization:\n\n- We constructed a mixed dataset where 99% of tokens are from SlimPajama and 1% from TinyStories\n- The objective is to achieve optimal validation loss on TinyStories while minimizing memorization\n- In this realistic setting, some repetition is actually helpful (unlike our previous setup)\n- The new \"validation optimal\" occurs when TinyStories data is upsampled 10x (11M tokens)\n\nIn the rebuttal [slides](https://shorturl.at/DqABg), we show in this setting, SeqTD significantly exceeds the deduplication baseline (reducing performance gap by 4x), while also reducing the loss gap between memorized and validation examples by 2.5x. This demonstrates SeqTD's capability to enable learning from repeated examples while mitigating memorization.\n\n### **Theme 2: Figure 4 (b) \"dip-and-rise\" pattern in memorization loss**\n> Regarding the \"dip-and-rise\" pattern in memorization loss with SeqTD:\n\n- This can be explained by the conceptual intuition for SeqTD (Section 5.3) There is initially some memorization learned in the shared neurons early in training. Once the memorization neurons fit the repeated sequences, further gradient steps no longer reinforce memorization in the generalizing neurons. At this stage, the interference from other sequences forgets the memorization learned in the shared neurons -- resulting in the rising memorization loss. \nImportantly, our larger token scale training [results](https://shorturl.at/B51HI) do not exhibit this behavior, suggesting it may be eliminated when generalization neurons don't have the capacity to memorize sequences on their own.\n\n---\n### **Theme 3: Evaluation Clarity and Metric Expansion**\n> *\"The forgetting evaluation could be more detailed beyond just loss increase.\"*\n\n🔗 *Full results available here:*  [Google Slides – Memorization Metrics](https://shorturl.at/M7UfM)\n\nWe've enhanced our evaluation with new results:\n\n- Perplexity-based ranks for memorized examples with and without SeqTD.\n- Memorization token accuracy metrics to complement loss measurements\n- Clear comparisons between SeqTD, standard training, and deduplication baselines\n\nThese additions provide multiple perspectives on the effectiveness of SeqTD beyond simple loss measurements. The observed patterns in general follow the same trends as seen with loss based metrics.\n\n### **Theme 4: Scaling experiments**\nWhile not a key concern in your review, we have developed new scaling trends and experiments based on other reviewer comments. If this is of interest, you can read our respons to Reviewer hZ7x.\n\n---\n**Table 3: Summary of Actions for Reviewer hkbC**  \n| Concern                         | Action Taken                                          | Expected Outcome                          |\n|----------------------------------|-------------------------------------------------------|-------------------------------------------|\n| Figure 4 (a) ambiguity               | Added larger scale experiments where repetition is beneficial | Clarify observed loss behavior and utility of SeqTD |\n| Loss dynamics explanation        | Added conceptual mechanism behind rising mem loss | Explain \"dip-and-rise\" pattern in memorization loss |\n| Evaluation diversity             | Added perplexity ranks + token accuracy              | Strengthen metrics beyond loss            |\n\nWe would appreciate any suggestions for additional metrics or visualizations you believe would further strengthen our paper.", "claims_and_evidence": "The claims made in the paper are generally well-supported:\n\n*  The claim that memorization is entangled with generalization in standard training is backed by experiments showing that removing neurons responsible for memorization also degrades model performance.\n*   The assertion that SeqTD enables controlled memorization isolation is demonstrated through experiments where SeqTD effectively removes memorization with minimal performance loss.\n*   The paper provides empirical evidence to support the claim that memorization neurons accumulate sequence-specific information while shared neurons maintain generalization.\n*   However, Figure 4(a) presents some inconsistencies where validation loss without repetition is lower than with repetition. This result seems counterintuitive. The authors should clarify this discrepancy.", "methods_and_evaluation_criteria": "The methods used are well-aligned with the problem:\n\n* The experiments use TinyStories, TS-Repetition dataset, a reasonable choice for studying memorization in small-scale settings.\n*    The evaluation metrics, including sequence forgetting and validation loss before and after neuron dropout, appropriately measure the effectiveness of SeqTD.\n* The comparison with existing methods such as gradient attribution and pruning strengthens the validity of the results.\n*  However, the experimental setup may not fully capture how SeqTD would behave in larger-scale LLMs, but the for the stated goal of the paper the choice of model and training suffices.", "theoretical_claims": "There are no theory claims in the main paper.", "experimental_designs_or_analyses": "The experimental design is generally solid, but some aspects could be improved:\n* The results in Figure 4(a) are somewhat inconsistent with expected behavior; typically, repeated sequences should reinforce learning, but the validation loss is lower without repetition. This raises questions about whether the small-scale training setup accurately reflects real-world LLM behavior.\n* Certain aspects need more discussion such as with SeqTD, the mem loss goes down and then up, this is not the case with standard training, what is the mechanism that drives this. Since early stopping with SeqTD could be detrimental to the privacy.", "relation_to_broader_scientific_literature": "Most of the relevant memorization localization papers are included", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths\n\n* The proposed SeqTD method is simple and novel (always a great combination) and well-supported empirical evidence.\n* The methodology is rigorous and the ablation studies add depth to the evaluation.\n\nWeaknesses\n* The results in Figure 4(a) raise concerns about whether the small-scale training setup accurately represents LLM behavior.\n* The forgetting evaluation could be more detailed beyond just loss increase. For example perplexity based ranks could be provided with and without the mitigation strategy.\n* Results also suggest that training without repetition has better performance for memorized examples compared to SeqTD, this there is a gap to potential upper limit\n* Certain aspects need more discussion such as with SeqTD, the mem loss goes down and then up, this is not the case with standard training, what is the mechanism that drives this. Since early stopping with SeqTD could be detrimental to the privacy.", "other_comments_or_suggestions": "Typo\n-  line 235 should use ` in latex rather than ' (i.e. open quote)", "questions_for_authors": "1. Figure 4(a) shows that validation loss without repetition is lower than with repetition. Can you clarify why this happens?\n\nThis seems inconsistent. If the small-scale setup does not reflect behavior in large-scale LLMs, that should be discussed explicitly. Please clarify this result.\n\n2. Can the authors discuss why with SeqTD, the mem loss goes down and then up, this is not the case with standard training, what is the mechanism that drives this?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eGb50RNBT8", "reviewer": "Reviewer_iFmB", "summary": "The paper studies memorization vs general capabilities in models by introducing their Sequence-Tied Dropout (SeqTD). This method a pool of shared neurons and a set of memorization neurons. They use the sequence ID to determine which memorization neurons to use. Note that there is some overlap between the sequences and an neuron.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the sharp and honest feedback. We're happy that you found:\n- (i) our ablations **insightful and well-done**,  \n- (ii) the method **interesting for isolating memorization**, and  \n- (iii) the paper a **valuable starting point** for deeper investigation.\n---\n### **Theme 1: Scale & Training Practicality**\n> *\"This method is not convincingly scalable—token counts and compute are low.\"*\n\nYou raise a crucial point about scalability. To address this, we've expanded our experiments significantly:\n\n🔗 *Full results available here:* [Model Scaling Experiment Results](https://shorturl.at/md9cQ), [Token Scaling Experiment Results](https://shorturl.at/Aa7VZ)\n- **Model scaling experiments** across four model sizes: 32M, 51M, 170M, 344 million parameters. Our results show that SeqTD mitigates memorization across all model sizes, but that the benefits of SeqTD increase with scale (no model degradation at 344M parameters). This provides evidence that *the benefits of SeqTD improve at scale*.\n- **Extended Token Scaling Experiments**  We train a 350M parameter model on a larger-scale corpus of 1 billion tokens containing a mix of TinyStories and SlimPajama. Our results demonstrate that SeqTD mitigates memorization (reducing the memorization-validation loss gap by 2.5x) while outperforming a deduplication baseline in validation loss.\n- **Explicitly stated token counts** in Section 5.1: ~16M tokens in initial experiments ((20K examples × 1 repetition + 100 examples × 128 repeats) × ~500 tokens per example)\n\n**Sequence hashing concern** \nWe address this with our sequence ID noise experiments (Fig. 5a). SeqTD remains effective even when 10% of sequence repetitions receive inconsistent IDs, demonstrating robustness to the n-gram overlap issue you raised. This suggests that perfect ID consistency is not required for SeqTD.\n\n---\n### **Theme 2: Evaluation Beyond Loss**\n> *\"Would prefer tests that evaluate memorization via generation.\"*\n\nTo demonstrate real-world utility, we've added prompt-based generation tests that provide much stronger evidence of SeqTD's effectiveness:\n\n🔗 *Full results available here:* [Google Slides – Memorization Metrics](https://shorturl.at/M7UfM)\n\n- **Prompt continuation tests** show standard models regenerate memorized text nearly verbatim (>95% token match), while SeqTD models fail to reproduce memorized content (<20% token match)\n- **Memorization rank metrics** reveal that SeqTD causes an 8× rank degradation for memorized content.\n\nThese new evaluations directly address your concern by showing the loss function results presented in the paper mirror two generation-based metrics.\n\n---\n### **Theme 3: Theoretical Foundation**\n> *\"They included some analysis on MLPs in the appendix. I do not think that these particularly contribute to the paper...\"*\n\nYou're right that the theoretical components need stronger connections to our practical claims. We've integrated key theoretical insights into Section 5.2 with a simplified version of Theorem E.3:\n\n> **Simplified Theorem:**  \n> *When p (memorization neuron activation rate) is sufficiently low, memorization accumulates in dedicated neurons rather than shared ones. This accumulation increases as memorization neurons are activated less frequently*\n\nThis theorem sheds important light on the scaling behaviors of SeqTD. As model scale increases and larger memorization pools are feasible (with each neuron active on less sequences), the localization achieved by SeqTD should improve. We validate this in our model scaling experiments. Our theory additionally reveals the crucial role of the p_mem parameter (which controls memorization neuron activation) as we discuss in Section 5.2 (pg 6).\n\n---\n**Table 4: Summary of Actions for Reviewer iFmB**  \n| Concern                        | Action Taken                                           | Expected Outcome                                 |\n|--------------------------------|--------------------------------------------------------|--------------------------------------------------|\n| Scalability and token scale    | Added 1B-token runs at 355M parameters | Validate claims in realistic training regimes    |\n| Memorization evaluation        | Added generation-based memorization metrics     | Demonstrate practical benefits beyond loss metrics |\n| Theoretical foundation         | Integrated simplified theorem with empirical validation | Connect theory to observed scaling behavior      |\n| Training details transparency  | Added explicit token counts and training setup tables  | Clarify experimental conditions and reproducibility |\n\nWe acknowledge that more work remains to fully validate SeqTD across the full spectrum of model scales and datasets. However, our new experiments provide compelling evidence that the method's benefits increase with scale, and the theoretical foundations help explain why. We look forward to your suggestions for additional experiments or analyses that would further address your concerns.", "claims_and_evidence": "They claim they propose a scalable and practical method that enables post-hoc isolation while preserving general model capabilities. I do not think this method is scalable nor practical. However, I do agree that it enables post-hoc isolation.", "methods_and_evaluation_criteria": "They pretrain on only one model size on one dataset using SeqTD. They compare to other methods to show that they are better at evaluating sequence forgetting (loss before and after on repeated sequences) and model degradation (loss before and after on val set). I do think only using loss limits this study, and that using generated sequences for sequence forgetting might have been better. Additionally, it would be good to study additional model sizes maybe smaller and larger (if compute permits).", "theoretical_claims": "They included some analysis on MLPs in the appendix. I do not think that these particularly contribute to the paper, but I did NOT carefully review the theorems.", "experimental_designs_or_analyses": "See Methods/Evaluations", "relation_to_broader_scientific_literature": "The paper situates this work mainly with Chang et al. 2024b. (https://aclanthology.org/2024.naacl-long.176/).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths\n- In-depth study of the methods with interesting ablations\n- Seem to be able to disentangle sequence memorization and general capabilities.\n\nWeaknesses\n- I am not convinced by the authors' arguments that this method is scalable and practical. Like n-gram overlap could be very high between different sequences but even with the proposed hashing technique, this method would map the sequence onto different memorization neurons. Additionally, I think the relationship of p-shared parameters and model size is not well understood. There are a lot of experiments to be made before I would be convinced of such a claim.\n- Limited Experiments -- although the authors do a good job with running interesting ablations, I do not think enough datasets and models were considered. Additionally, the models from my calculations were trained for 100-200 of million tokens (2M examples x ~50-100 tokens per example) instead of billions. (Why is it not clear how many tokens were considered in the paper?) I think it makes it hard to be convinced of some of the claims in the paper when the training runs are very small with no breath in terms of dataset or model size.", "other_comments_or_suggestions": "Add details in the paper about the number of tokens and training steps in the experimental setting.", "questions_for_authors": "See Weaknesses and suggestions please.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper studies the question of identifying neurons responsible for memorization in LLMs.
The paper shows that standard training typically does not allow for memorization to be disentangled from model capabilities, and so removing neurons responsible for memorization also hurts utility.
The paper thus suggests a training recipe to force a more clear separation between neurons responsible for memorization and neurons responsible for generalization.

Reviewers found the idea interesting, but raised some concerns about the limited experimental scale, as well as a lack of justification for how the proposed method actually promotes disentanglement of neurons between memorization and generalization.
Despite these shortcomings, I think the paper presents an interesting new idea, which can lead to follow-up work on better understanding the learning dynamics of memorization.
Decision: Accept (poster)