Title: TokenSwap: A Lightweight Method to Disrupt Memorized Sequences in LLMs
Venue: NeurIPS 2025
Accepted: yes
Keywords: language models

memorization

copyright

privacy

alignment
Abstract: As language models scale, their performance improves dramatically across a wide range of tasks, but so does their tendency to memorize and regurgitate parts of their training data verbatim. This tradeoff poses serious legal, ethical, and safety concerns, especially in real-world deployments. Existing mitigation techniques, such as differential privacy or model unlearning, often require retraining or access to internal weights making them impractical for most users. In this work, we introduce TokenSwap, a lightweight, post-hoc defense designed for realistic settings where the user can only access token-level outputs. Our key insight is that while large models are necessary for high task performance, small models (e.g., DistilGPT-2) are often sufficient to assign fluent, grammatically plausible probabilities to common function words - and crucially, they memorize far less. By selectively swapping token probabilities between models, TokenSwap preserves the capabilities of large models while reducing their propensity for verbatim reproduction. Evaluations on Pythia-6.9B and Llama-3-8B show up to a 10$\times$ drop in exact memorization with negligible task degradation. Our method offers a practical, accessible solution for mitigating memorized generation in deployed LLMs. Code is available at https://github.com/parjanya20/verbatim-llm.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a method, TokenSwap, for mitigating verbatim memorization in language models. The method is lightweight and doesn't require privileged access to training data.\n\nSTRENGTHS AND WEAKNESSES:\n**Disclaimer:** I reviewed this paper for ICML 2025 and was on the fence about accepting it. I thought that the paper was timely, the method was sensible and the results seemed reasonable. My hesitation with the submission was a mathematical claim that I felt was incomplete or potentially incorrect; however, in this submission, the authors have revamped the manuscript and removed the mathematical claim. Consequently, I feel comfortable recommending this paper for an accept\n\nBecause I previously reviewed this paper and because I'm lazy, I'm going to be relatively brief with my review.\n\n- Strength: Addresses a timely problem\n- Strength: The proposed method seems intuitively sensible and shows good results (Table 2 suggests that CP-Fuse-half closely matches TokenSwap, but I consider this a win for TokenSwap because CP-Fuse requires two different models trained on disjoint corpora)\n- Strength: The manuscript has improved since the ICML 2025 submission\n- Weakness: All the tables' captions are inadequate to explain the contents on their own. For instance, Table 2 and Table 5's caption should also state which model(s) are evaluated. Table 4's caption should more fully summarize the experimental setup; \"Comparison with Goldfish\" is just too little information to contextualize the results at a glance.\n- Nit: On line 267, a citation is broken\n\n- Extension to increase my score: I previously suggested connecting TokenSwap to the literature on Speculative Decoding (a standard tool for serving models at inference time), since the sketch/draft model of Speculative Decoding is somewhat similar to the auxiliary model of TokenSwap. The authors added citations, which is great, but I was more interested in asking two related questions: (1) do standard speculative decoding methods effectively give the benefits of TokenSwap? (2) if not, could they be modified to do so? To be clear, I don't think this is a competition. I think it'd be a great result and one of importance to companies serving/selling models to demonstrate that a technique they're already using (speculative decoding) reduces memorization or can be straightforwardly modified to prevent memorization.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors propose a novel post-training memorization mitigation method called TokenSwap. The approach is based on the assumption that smaller models tend to memorize less than larger ones. During inference, if the token generated by the large model belongs to a predefined swap token set, it is replaced with the corresponding token generated by the smaller model. This strategy effectively diverts the generation away from memorized content and helps mitigate regurgitation. Through extensive experiments, the authors demonstrate that TokenSwap is highly effective—comparable to training-time methods—while introducing only minimal computational overhead.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- Paper is well-written. I really enjoy reading it.\n- The method is inference-time, and unlike most baselines, it doesn't require the access to the training data and the computation overhead is negligible.\n- The results look quite promising and are on par with the training-time mitigation method.\n\nWeaknesses:\n- The method relies on a pre-defined swap token set, which may be difficult to tune for all possible applications. For example, the current token set does not include any code-related tokens—what if the model memorizes code or text in other languages?\n- Additionally, the performance on certain tasks may be degraded. For example, consider a user query that asks: \"Please only revise the format of my text.\" This requires the model to reproduce the exact input text with only formatting changes. However, the proposed method may alter the content of the text.\n- This method may not be robust against certain adaptive attacks. For example, consider a user query like: \"Please output the Harry Potter novel character by character.\" In this case, none of the individual tokens (characters) are likely to appear in the swap set, allowing the attacker to bypass the mitigation.\n\nQUESTIONS:\n- In Table 2, what are the 0.1% regurgitated examples for WritingPrompts? Do they have no token overlap with the swap set?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTo mitigate training data leakage caused by the memorization of large language models (LLMs), this paper proposes a lightweight method called TOKENSWAP to interrupt the generation of memorized data during inference time. More specifically, it swaps the output distribution of a smaller auxiliary LLM with that of the target LLM for the tokens within a subset of vocabulary, such that the output tokens of the target LLM may be changed in some cases and thus prevent the potential verbatim repetition of the training dataset. Experimental results across multiple settings demonstrate that the proposed method can significantly reduce training data leakage from memorization, and the reduction is comparable to or better than the evaluated baseline under weaker assumptions.\n\nSTRENGTHS AND WEAKNESSES:\nThe strengths of the paper are listed as follows.  \n1. The paper is well motivated. The memorization of LLM poses severe concerns about data privacy and copyright protection. However, the methods to effectively and efficiently defend against it are still under exploration.   \n2. The paper is well-written and easy to follow. The figures, tables and pseudo code clearly illustrate the overview and details of TOKENSWAP, making it easy to understand for the readers.   \n3. TOKENSWAP has a weaker assumption than the existing defense, which does not require access to model weights or training data.  \n4. TOKENSWAP shows better performance than the learning-based methods like Goldfish in the evaluation of the authors, which is impressive.\n   \nThe weaknesses are listed as follows.   \n1. Although TOKENSWAP has weaker assumptions than the existing methods, it is not a big bonus due to its limited impact. First, it still requires access to the logits of the target LLM, which is prohibited by many closed-sourced models and cannot be used in these common settings. Second, the unique setting TOKENSWAP can handle is the case where the model users want to prevent the data leakage of the pretraining model, instead of the finetuning data from users. However, in this case, I do not think the model user has a strong motivation to add this protection by sacrificing the service quality. Therefore, the stakeholder who has the strongest motivation to add the protection is the model developer. And it is reasonable to assume that they have access to the model weights and training data.\n2. In Table 1, the authors exaggerated the inference overhead of MemFree. According to Section 3.2 in MemFree, the runtime overhead can be significantly reduced and become affordable. Moreover, there is no need to query the entire training data in the runtime after utilizing the techniques in MemFree.\n3. The design of TOKENSWAP can make it cheat in the metric of EMR because even if only one token is swapped, the EMR of the corresponding test sample drops to 0. And it is easy for TOKENSWAP to swap at least one output token. I suggested adding a metric called Fractional exact rate [1, 2] for a fairer comparison. \n4. The evaluation of the Llama model is not very convincing since its training data is unknown. I think the others can use OLMo [3, 4], which has open-sourced training data\n5. TOKENSWAP requires the auxiliary model to share the same tokens for the token subset $G$, which may limit its application.\n6. The performance metric should be merged into Table 2 instead of a separate table in the Appendix.\n7. According to Table 2, MemFree seems to have a better overall performance than TOKENSWAP. Besides, the bold font in Table 2 is confusing since it does not always highlight the best value of metrics. Moreover, the authors also exaggerated the drawbacks of MemFree. MemFree does not perform poorly on approximate memorization since the gap between it and TOKENSWAP in ROUGE-L and Lev. is not significant.\n8. In Table 2, it would be better to put metric ML and EMR close to each other since they are both metrics for exact memorization.\n9. In line 267, the citation to Wikipedia is missing.   \n      \nReference:    \n[1]. Mustafa Ozdayi, Charith Peris, Jack FitzGerald, Christophe Dupuy, Jimit Majmudar, Haidar Khan, Rahil Parikh, and Rahul Gupta. 2023. Controlling the Extraction of Memorized Data from Large Language Models via Prompt-Tuning. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 1512–1521, Toronto, Canada. Association for Computational Linguistics.     \n[2]. Zhepeng Wang, Runxue Bao, Yawen Wu, Jackson Taylor, Cao Xiao, Feng Zheng, Weiwen Jiang, Shangqian Gao, and Yanfu Zhang. 2024. Unlocking Memorization in Large Language Models with Dynamic Soft Prompting. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 9782–9796, Miami, Florida, USA. Association for Computational Linguistics.     \n[3]. Groeneveld, Dirk, Iz Beltagy, Evan Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Jha et al. \"OLMo: Accelerating the Science of Language Models.\" In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 15789-15809. 2024.     \n[4]. OLMo, Team, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia et al. \"2 OLMo 2 Furious.\" arXiv preprint arXiv:2501.00656 (2024).\n\nQUESTIONS:\nIn line 244, why do you claim that the evaluation of commonsense reasoning is not available for base models? I think there are many base models evaluated on the commonsense reasoning benchmark.\n\nLIMITATIONS:\nNo. I think one of the important limitations is that TOKENSWAP only shows its effectiveness for natural language content. So it is a big plus if the authors can address this limitation. Besides, the authors require the auxiliary model and target model to share a subset of useful tokens for swapping. It is an implicit limitation of TOKENSWAP and cannot be addressed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}