Tokens for Learning, Tokens for Unlearning: Mitigating Membership Inference Attacks in Large Language Models via Dual-Purpose Training (2025)

Abstract: Large language models (LLMs) have become the backbone of modern natural language processing but pose privacy concerns about leaking sensitive training data. Membership inference attacks (MIAs), which aim to infer whether a sample is included in a model’s training dataset, can serve as a foundation for broader privacy threats. Existing defenses designed for traditional classification models do not account for the sequential nature of text data. As a result, they either require significant computational resources or fail to effectively mitigate privacy risks in LLMs. In this work, we propose DuoLearn, a lightweight yet effective empirical privacy defense for protecting training data of language models by leveraging token-specific characteristics. By analyzing token dynamics during training, we propose a token selection strategy that categorizes tokens into hard tokens for learning and memorized tokens for unlearning. Subsequently, our training-phase defense optimizes a novel dual-purpose token-level loss to achieve a Pareto-optimal balance between utility and privacy. Extensive experiments demonstrate that our approach not only provides strong protection against MIAs but also improves language modeling performance by around 10% across various LLM architectures and datasets compared to the baselines.

## introduction large language models (llms) have become the foundation of modern natural language processing with a wide range of applications in various domains @xcite . the rapidly increasing deployment of llms raises serious concerns about data privacy @xcite . llms have been shown to memorize the training data which can be later extracted by adversaries @xcite . membership inference attacks (mias) @xcite aim to infer whether a sample is included in a model's training data, serving as the foundation of broader privacy threats @xcite . due to the importance of understanding and mitigating mias, a significant amount of research has been conducted to design mia defenses @xcite . however, most defenses focus on general machine learning models for classification tasks and do not account for the sequential nature of text data, while advanced mias for llms have leveraged this property. for example, the series of min-k works @xcite uses the token-level loss on outlier tokens and significantly enhance mias for llms. thus, conventional data sanitization or regularization techniques have limited defense effectiveness @xcite . even though the classic differentially private (dp) training algorithm @xcite provides a strong defense, this approach comes at the inevitable cost of increased computation and reduced utility @xcite , which may not be desirable when the model trainer serves as the defender. in this paper, we propose a defense mechanism for membership inference attacks on llms -duolearn. a recent study @xcite reveals that using a carefully selected subset of tokens during training can match or even surpass the performance of using all tokens in language modeling. in the meantime, mias mainly exploit loss-based signals associated with a sample @xcite . we observe that during training, some certain tokens carry stronger mia signals than others, making the sample vulnerable to mias. thus, we leverage the token sequence nature of llms and propose a dynamic token selection strategy during training to proactively identify and categorize tokens into hard tokens (those the model struggles with) and memorized tokens (those with strong signals for mia risks). accordingly, we design a dual-objective loss function that performs learning via gradient descent on the hard tokens and unlearning via gradient ascent on the memorized tokens simultaneously in one backward pass, which makes the model learn useful information but not memorize specific training samples. our contributions can be summarized as follows: ## how do tokens contribute to membership inference risks? compared to conventional classification problems, membership inference attacks in language modeling have significant differences. in particular, each query in traditional classification models requires only one prediction. on the other hand, each query to language models involves multiple token predictions due to the sequential nature of text. this difference yields a question that how tokens contribute to overall sample-level membership inference risks. to answer this question, we perform a token-level analysis of membership inference risks. we calcu-late the mia signal for each token as its prediction loss calibrated by a reference model @xcite . a smaller signal value indicates that the model has a significantly higher confidence than the reference model on predicting the token. figure 1 (left) illustrates the histogram of mia signal values for the tokens of a member sample and a non-member sample (see figure 9 in appendix b for additional histograms). the nonmember sample distribution centers around zero, while the member sample skews to the negative side. consequently, the average aggregated mia signal is below zero for the member but around zero for the non-member, leading to membership inference risks. moreover, the mia signal values vary for different tokens, so some tokens contribute more to the membership inference risks than the others. figure 1 (right) illustrates the mia signal ranking of tokens of a member sample over training steps (see figure 10 in appendix b for additional samples). there is a complex changing dynamic in ranking between tokens before it becomes more stable at the end when the training converges. overall, the analysis suggests that the sample-level membership inference risk in language modeling stem from the cumulative effect of many tokens. this poses challenges for defense methods, as they need tokenlevel granularity to isolate and mitigate specific sources of leakage. additionally, it is non-trivial to develop a defense method that widely affects a large number of tokens without disrupting the complex token dependencies essential for model utility. ## proposed methodology -duolearn motivated by the analysis, we propose duolearna training framework that dynamically identifies hard tokens (those with higher calibrated losses) for learning and memorized tokens (those with lower backpropagation numbered ab89 has been shipped training @xmath0. this way, the model learns useful information without memorizing specific training samples. overview. we assume the model trainer is the defender and the goal is to mitigate the privacy risk of the training data in the trained model. we further assume the trainer can get access to an auxiliary dataset for better calibrating the mia signals, which can be a disjoint subset drawn from the same distribution of the training data. the general training process is illustrated in figure 2 . first, we train a reference model with the auxiliary dataset, which is feasible for the trainer based on previous works @xcite @xcite . then, during training of the target model, we use the token losses of the current training model calibrated by the reference model to dynamically identify hard tokens and memorized tokens in each training iteration. a dual-purpose loss function is used to keep the model simultaneously learning on hard and necessary tokens to enhance model utiilty while unlearning on memorized tokens to mitigate mia risks. reference modeling. reference model (θ ref ) shares an identical architecture with the training model (θ). we fine-tune a reference model on a small portion of the original dataset (denoted as t aux ) that can reflect the desired data distribution by standard causal language modeling (clm), i.e., implementing next-token-prediction cross entropy loss (l ce ): token selection. as our previous analysis, tokens contribute differently in membership inference risks. thus, considering all tokens equally is not optimal for privacy defense with respect to the utility and privacy trade-off. duolearn defines two sets of tokens: hard tokens (t h ) and memorized tokens (t m ). hard tokens are the tokens that the current training models (θ) have difficulty predicting, while memorized tokens are the tokens that the model has already memorized. to identify these two sets of tokens, we propose a token selection mechanism based on the prediction loss of each token calibrated by the reference model. we implement the score s(t i ) for each token t i which is the difference between the cross-entropy loss of the training model and the reference model, as used in previous works @xcite : the tokens with the highest scores are considered hard tokens t h (highest calibrated loss), while the tokens with the lowest scores are considered memorized tokens t m (lowest calibrated loss and strongest mia signals). let t be the set of all tokens in a batch. we select top k h hard tokens and bottom k m memorized tokens to form t h and t m , respectively. additionally, we introduce a threshold τ to filter out neutral tokens from t m which have scores close to zero or greater than zero, as these are not considered memorized. the token selection process is formulated as follows: we introduce a dual-purpose loss function designed to improve model performance on hard tokens while mitigating overfitting on memorized tokens. this loss function combines two components: the learning loss and the unlearning loss. the learning loss is the standard causal language modeling (clm) loss applied to the hard tokens t h . the unlearning loss, in contrast, is the negative clm loss applied to the memorized tokens t m , effectively performing gradient ascent. the dual-purpose loss is defined as follows, where α > 0 is a hyper-parameter that balances the learning and unlearning losses ## conclusion we introduced duolearn, an effective training framework defending against mias for llms. the extensive experiments demonstrate its robustness in protecting privacy while maintaining strong language modeling performance across various datasets and architectures. although our study fo4 following this setting and due to limited computing resources, we use a batch size of 8 which is much smaller than practical pretraining. all methods implement the same learning rate and are evaluated at their 25k-th iteration. cuses on fine-tuning and small scale pretraining due to computational constraints, duolearn can be seamlessly applied to large-scale pretraining, as in prior selective pretraining work @xcite . by categorizing tokens and treating them appropriately, duolearn opens a novel pathway for mia defense. future work can explore improved token selection strategies and multi-objective training approaches.