Mitigating Forgetting in LLM Fine-Tuning via Low-Perplexity Token Learning

We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.

weakness 1. simplified view of STM cause training instable due to unnatural gap.

Thank you for the insightful feedback. In our experiments with Self-Output and Rephrase, we observe that high-perplexity tokens are often rewritten into low-perplexity alternatives, leading to improved performance, lower training loss, and faster convergence, resulting in a smoother training curve compared to using ground truth data. Similarly, STM exhibits consistent trends in both training stability and performance. We believe that STM primarily filters out high-perplexity tokens that are unhelpful for training. Even when such tokens carry important information, the model is often able to recover it through contextual learning from the surrounding unmasked tokens. Our results on the MATH dataset support this: as shown in Appendix D, Figures 8, 9, and 10, STM removes a higher proportion of symbols and special characters compared to Self-Output, yet still improves both BWT and TI scores, demonstrating enhanced target-task performance and reduced forgetting. We agree that future extensions of STM could benefit from more nuanced selection mechanisms, e.g., conditioning on attention patterns or incorporating token roles (e.g., logical operators vs. distractors), to better distinguish between helpful and harmful high-perplexity tokens. We have added this direction to our discussion section as a promising avenue for future work.

Weakness 2. Limited theoretical justification

Thank you for raising this important point. As noted in our limitations checklist, this work does not aim to provide a theoretical formulation and instead focuses on a comprehensive empirical investigation. To that end, we conducted extensive experiments across multiple model families, scales, datasets, and fine-tuning strategies (see Tables 2, 3, 4; Appendix C), all of which consistently demonstrate that low-perplexity token training, via STM or Self-Output, yields better preservation of general capabilities and improved target-task performance. Furthermore, in Appendix J, we analyze LoRA rank and weight updates and observe that STM and Self-Output generally result in lower-rank adaptations and smaller update magnitudes. This suggests that low-perplexity training tends to induce more localized updates, which may help retain pre-trained capabilities during task adaptation. While we do not provide a theoretical justification, we believe the consistent empirical trends observed across varied setups offer practical insights into the behavior of low-perplexity training. We agree that building a formal theoretical framework to explain these effects is an important direction for future work, and our findings may serve as a foundation for such efforts.

Weakness 3. Selection of STM make less steadily applicable in real-world application.

We thank the reviewer for the insightful question. While selecting an appropriate threshold for STM can be challenging when deploying STM to entirely new domains, our empirical findings suggest that the current threshold setting (e.g., τ=2.5) is remarkably stable across a diverse range of real-world domains in our work: programming, mathematical reasoning, and factual knowledge retrieval. These domains are representative of practical applications where fine-tuning stability and generalization are critical. To further illustrate real-world applicability, we also apply STM to instruction-following tasks, IFEVAL. As shown below, STM consistently mitigates forgetting and improves performance relative to baseline fine-tuning across different learning rates and models:

Instruction following (IFEVAL)

While task-specific variability in perplexity distribution could affect optimal masking behavior in principle, our current result suggests that STM is already broadly applicable in practical settings. While we view more adaptive or data-aware thresholding mechanisms as a valuable future work, our results indicate that STM, as a simple and practical first step, already delivers robust performance across real-world fine-tuning tasks.

Q1. removal of high ppl token is a specific case for STM.

To ensure STM is generalized enough, we take 5 different tasks, 5 different models ranging from 2B to 9B, originating from 4 different model families. And our conclusions still hold for such similar threshold selection based on the token filtering ratio. In addition to our scenarios, concurrent work of VLM (Liao et al, https://arxiv.org/abs/2504.15362) also has similar findings with generated data training in a multi-modal application, which shows STM’s possibility to apply in broader applications.

Q2. any challenge of threshold selection of STM.

To address the challenge for generalization of threshold selection, we further apply our threshold selection on other new models we did not mainly discuss in this paper as well: OLmo 7B and Gemma 9B results in Appendix C Table 7, holding the same conclusion that STM benefits both target improvement and reduction in forgetting. It is still unknown how STM affects the learning when input contains a higher proportion of difficult, higher ppl context, e.g. specific domain data adaptation or hard-to-verify data, e.g. reasoning process, it makes us believe that considering not only tokenwise, but also attentionwise masking could benefit the model training as well, which is a potential future work to investigate.

Thank you for your thoughtful comment. We agree that theoretical justifications are important and appreciate your perspective. While we believe that our empirical findings offer meaningful contributions—particularly in the context of large language models—we would greatly appreciate it if you could clarify the type of theoretical justification you are referring to (e.g., convergence guarantees, optimality results, or another aspect).

Additionally, we would be grateful if you could elaborate on what is meant by the method being “unintuitive.” A more detailed explanation would help us better understand and address your concern.

We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.

Lower baseline finetuning results than original models.

We appreciate the insightful concern. In related work [9], it discusses how instruction tuning’s copy pattern could possibly harm the performance, and could also affect the response quality of LLM. Such possibly “overfitting” problems could be solved either by carefully tuning the hyperparameters or other regularization methods. Here we provide carefully tuned baseline fine-tuning results with very small and not practically used learning rates in the previous rebuttal comment, which can improve the performance from the original model in terms of TI. While it still fails to surpass the performance of STM with a larger and commonly used lr. And STM sustains overall improvement of TI and BWT across learning rate, also showing a low sensitivity for learning rate with STM. In addition, in our work, we try to address the forgetting problem from fine-tuning, which is also mentioned by related work [5,35,36,40]. Thus, the positive TI metrics are not the main focus, but the BWT metrics. Yet we think it would be better to mention details about how fine-tuning results could suffer from easily overfitting/pattern-copy failure and catastrophic forgetting in our future version.

Only fintuning and evaluating on the similar domain at the same time

Thank you for your concern. In section 4.1, we simply brief the target and non-target tasks and datasets used for experiments. We actually target to train on two datasets from two domains respectively, and test the rest of the four non-target datasets from three different domains altogether, as shown in Table 8. For example, we fine-tune on MBPP for the coding domain, but we evaluate not only against BIRD, but also the MATH-test, GSM8k, and ARC-challenge dataset together, and calculate the average BWT score for forgetting metrics calculation, and evaluate on MBPP-test for TI score for target task improvement calculation. We understand your concern about similar domains between training and evaluation, so we consider all three domains (programming, math, and knowledge-based QA) for evaluation at once under the suitable limitations ( which is the need for label-verifiable tasks for generated data baseline) that fit our baselines. Furthermore, we believe that training on programming tasks with robust training like STM should ensure improvements on similar tasks as well. In many applications, this does not always hold for SFT; if in similar tasks SFT already forgets, then it will forget even more in truly disparate tasks than STM.

Masking prevent model learning from challenging and valuable parts.

We understand the concern for STM to “underfit” some tasks with important key tokens being masked while training. Thus, we do have a simple trade-off study about the filtered ratios of STM threshold vs performance in Table 3 and the filtered token distribution study in Appendix D, Figures 8-10: STM masks more high ppl tokens in math, including symbols and special characters for mathematical calculation, than self-generated data. But the performance of TI is still positive and better than baseline fine-tuning, showing that the model still improves without knowing those important key symbols. We believe it is very likely that models can still learn novel concepts or logics based on the correct context we left in the training data, which, in our experiment, shows that STM has comparable performance with the self-output baseline. To address the difficulty of the problem, we could also explore curriculum learning (learn easy first, then hard later) to possibly avoid eliminating all higher perplexity tokens of challenging tasks. And we would like to emphasize that STM is built on top of SFT, and tasks that SFT could train on are also applicable for STM to address forgetting problems. Some open-ended tasks would not be suitable, such as RLHF or preference learning related tasks, or tasks with dynamic or unverifiable labels.

Unsuitable for the task of stylistic variety, nuance, and creativity:

We agree there may be a possible limitation to consider forgetting training on only two domains. Thus, we also include a knowledge QA task like ARC-Challenge to evaluate the forgetting metric of models. Likewise, in the previous rebuttal, we believe most tasks SFT could train on should be covered by STM to address forgetting well. While in open-domain tasks including creative generation tasks, pure SFT may still underperform other methods due to there are few clear or gold standards of data that can be learned, but more preference learning strategies can work like RLHF, prompting, or LLM discussion (Lu et al) solutions. We think it’s very important to consider how different traits/dimensions of high ppl training data differ when applying STM, and develop different masking strategies like conditional masking (a more finetuned masking way) or curriculum learning (learning order of data by difficulty) as our future work to explore how to address learning with different traits of data. In this paper, we would like to start by figuring out that the root cause of forgetting is not purely because of generated data, but the high perplexity tokens in the ground truth training data.

We acknowledge your concern about the performance drop on in-domain datasets, in our other rebuttal poeA, we have conducted additional sweep of learning rate, which reveals that the initial performance drop was indeed partially due to suboptimal hyperparameter selection. With appropriate learning rate baseline finetuning can achieve positive gains on in-domain tasks (+1.33% on MBPP).

However when compared with STM methods, STM consistently outperforms baseline fine tuning on in-domain tasks. This demonstrates that STM is not only more effective but also more robust to hyper-parameter choices.

MBPP results :

We deeply appreciate the reviewer’s insightful and detailed feedback. Below, we address each of your questions in turn.

Weakness 1. training over time and lower diversity with STM

Thank you for the insightful concern. Although our method manages to address forgetting after supervised fine-tuning instead of diversity, we still understand the reviewer’s concerns about the difficulty of problems and diversity over multiple training. We think the difficulty distribution of problems is an unexplored topic of STM training and a great future work to do. For multiple times of training sessions, in Appendix G, Table 15, we have investigated a simple scenario of STM applied to continual training, which is training on a model with different datasets sequentially. The results showed that forgetting is also mitigated during each turn of training and enhances the performance of related domains as well. Based on the reviewer's suggestion, we also observe how diversity trades off with forgetting after training the model on the MBPP coding dataset. The experiment is tested on 100 MATH questions. Each instance’s reasoning response is generated 4 times, and average self-bleu scores are calculated. The results showed that even with a similar diversity (self-bleu), SFT results in more forgetting, which also emphasizes STM’s capability to mitigate forgetting. We understand the reviewer’s concern that training could reduce diversity, but we think it is not caused by STM, as the diversity is almost the same with pure SFT. However, we believe that STM should focus on non-open domains that SFT could apply to, which have verifiable answers for models to train on, so that STM could enhance SFT with less forgetting.

Proposed experiment on math reasoning x100 with MBPP trained Llama3 8B-it models

Limiation: exploration on reasoning dataset, including instruction following, safety and MMLU.

We thank the reviewer for the critical insight. We think it’s very crucial to explore more tasks that are not covered in our work, as our limitation section. We also experimented with a non-target knowledge task (arc) other than just programming and mathematical datasets for forgetting dataset metric comparison (BWT). And reasoning dataset is also a great direction to investigate, while in this paper, to compare with baselines like self-output and rephrase, reasoning dataset is not easy to construct because we do not have a correct evaluation for the generated reasoning data, which could bring noise to the training as well. Because our core concept is to figure out that forgetting stems from high perplexity tokens instead of purely model-generated data, we need to compare generated baselines with applicable datasets. We believe tasks like knowledge QA, reasoning, safety, or instruction following tasks could be great future works to explore STM’s training effectiveness as well. Here we prepare a simple test for Instruction following on Gemma2-2B-it and Llama3-8B-it. We can still see that even though small lr prevents instruction following capability from decreasing, STM manages to maintain and even improve such capability better. Instruction following (IFEVAL)

As for exploration in safety, we have tested the results on the safety benchmark (advbench). Due to the limited training domains of programming and coding only, we found that neither baseline fine-tuning nor STM would cause degradation in the safety benchmark (advbench), which likely means that the difficulty of training tasks or related capability after training does not affect the performance on safety response.

Q1. lower improvement of STM than self-output method in math (reasoning)

We thank the reviewer for raising the insightful concern. First of all, we believe the lower perplexity token trait of self-output data benefits the training to reduce or even improve forgetting. Second, we have initially observed some filtered tokens case study for the MATH task in Appendix D, figures 8,9, and 10. It does show that compared to self-output data, STM’s MATH ground truth data intrinsically filters out higher proportions of high perplexity symbolic and special character tokens, which could be important logics or reasoning processes, resulting in lower TI. While in BWT the averaged scores differ in less than 1% of the original scores, which could result from randomness of inference results of LLM and indicate a good preservation of generalization for both STM and Self-output.. To address the lower TI issue, we believe STM with conditional tokens masking could be a potential direction to explore when models could learn which high perplexity and “non-functional” tokens to filter.

Q2. STM effectiveness on more general tasks like instruction following, safety or MMLU

We appreciate the concern raised by the reviewer. As we mentioned in the Limitation section, to indicate the forgetting stemming from low perplexity, we managed to compare the approaches that also bring low perplexity training data: Self-output and rephrase. However, these approaches are limited when labels of generated results are hard to evaluate correctly. Besides, we tested ARC for a knowledge QA based dataset; we do believe STM-trained models also mitigate forgetting in general tasks like MMLU, Safety, and Instruction following, since these tasks intrinsically contain high perplexity tokens to filter. Therefore, there is also a great potential to investigate STM effectiveness in those datasets. For further study, please see the previous comment about the results in the instruction following task. The trend of instruction following meets our expectations. And the comments we left about the safety benchmark in the previous comment as well.

Q3. further investiagtion on the trheshold selection of STM cause on diversity, safety

Thank you for the insightful feedback. We have an initial study about high ppl tokens distribution analysis in the MATH dataset and the generated dataset in Appendix D, Figures 8~10. In the analysis, we find a different proportion of high ppl in the generated dataset and ground truth data. But removing the high ppl tokens in ground truth data for STM training does not harm the target task and non-target task performance. To address more aspects like diversity, please see the previous initial study on self-bleu diversity on math reasoning. And for the safety issue, as we mentioned in weakness 2. Safety may not be a forgetting issue for model training with coding and math domains. STM and baseline fine-tuning rarely forget safety in advbench in our case.

I would like to thank the authors for their detailed response. These additional experiments and proposed methods would strengthen the paper further. I will be maintaining my score.

We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.

lr sweeping

We thank the reviewer for highlighting the importance of learning rate tuning in SFT performance and its impact on forgetting. We apologize for the lack of clarity in our original submission and will include additional details in Appendix E in the camera-ready version. Following prior works [9] and the Llama 2 finetuning recipe (Touvron et al.), our initial experiments used learning rates ranging from 1e-4 to 5e-5, with batch size from 4 to 16, the AdamW optimizer (bnb_8bit), a weight decay of 0.05, and LoRA rank 16. We also provide a weight decay analysis in Appendix I (Fig. 11). In response to the reviewer’s suggestion, we conducted additional learning rate sweeps down to 1e-7 for both Llama 3 8B and Gemma2 2B on the MBPP dataset. We found that while baseline SFT improves at smaller learning rates (e.g., TI: -10.6% → 1.33%), the forgetting problem (BWT) persists (e.g., -1.6%). In contrast, STM consistently outperforms baseline across all learning rates with stronger target task performance and substantially lower forgetting. Notably, STM achieves better BWT and stable performance even at small learning rates (e.g., TI: 2.23%, BWT: +1.39% at 1e-7). These results demonstrate that STM is more robust to learning rate choices, requires less learning rate tuning, and achieves more stable training dynamics, highlighting its practical advantage in real-world fine-tuning scenarios.

Baseline Training vs STM:

checking Rephrase with ground truth data

We appreciate the concern raised by the reviewer.
We have a sanity check on the output of Rephrase, and remove the incorrect label output as the final training set for Rephrase to ensure that all the final outputs of Rephrase are correct. But we did not evaluate the intermediate output of Rephrase, e.g., the reasoning or CoT process it could generate. This process is partially reflected in Appendix Table 13, which shows the number of training samples retained after filtering. We agree that this could have been more clearly stated and will improve the wording and methodology description in the future version to better reflect this filtering step. We appreciate the opportunity to clarify this point.

perplexity of the ground truth reponse to larger model

We appreciate the reviewer’s insightful comment. As noted in our limitations, we did not explore models >10B due to resource constraints. This reflects a common real-world scenario, where most research teams work with models under 10B due to limited compute access. In practice, small LLMs (<10B) are commonly used for both training and deployment, and we believe STM remains highly relevant in such settings. In addition, in Appendix C, we have tested another newer 7B (OLMo) and larger Gemma 2 9B model results, and the results are consistent with our hypothesis. Due to time limitations, we haven’t finished the whole evaluation on the 70B model, but we will include the analysis in the future version.

focus on math and code, forsee any challenge

We thank the reviewer for thoughtful feedback. We included a knowledge-based dataset as a non-target evaluation set for BWT calculation in Section 4.1. But to brief more about such limitation, the reason we focus on training code and math is that the labels are easily verifiable, and our self-output baselines are limited in such a setting for comparison. We do believe STM applies to other tasks like Knowledge QA (as ARC we mentioned in 4.1), or safety and instruction following tasks mentioned in the Limitation section. A particular challenge we foresee could be the different proportions of low perplexity tokens that exist in those tasks. And specific function words may be exceptions for filter, a much smarter masking is potentially for future work, e.g., not only on token masking but also attention to specific neurons. And based on such a setting, how do we automatically learn a good threshold for STM is also critical. To explore other domains, we also compare STM with baseline fine-tuning on the instruction following domain:

Instruction following (IFEVAL). We found a similar trend for STM that sustains or even improves the non-target domain very well.