MoFO: Momentum-Filtered Optimizer for Mitigating Forgetting in LLM Fine-Tuning

## Abstract

Abstract Recently, large language models (LLMs) have demonstrated remarkable capabilities in a wide range of tasks. Typically, an LLM is pre-trained on large corpora and subsequently fine-tuned on task-specific datasets. However, during fine-tuning, LLMs may forget the knowledge acquired in the pre-training stage, leading to a decline in general capabilities. To address this issue, we propose a new fine-tuning algorithm termed Momentum-Filtered Optimizer (MoFO). The key idea of MoFO is to iteratively select and update the model parameters with the largest momentum magnitudes. Compared to full-parameter training, MoFO achieves similar fine-tuning performance while keeping parameters closer to the pre-trained model, thereby mitigating knowledge forgetting. Unlike most existing methods for forgetting mitigation, MoFO combines the following two advantages. First, MoFO does not require access to pre-training data. This makes MoFO particularly suitable for fine-tuning scenarios where pre-training data is unavailable, such as fine-tuning checkpoint-only open-source LLMs. Second, MoFO does not alter the original loss function. This could avoid impairing the model performance on the fine-tuning tasks. We validate MoFO through rigorous convergence analysis and extensive experiments, demonstrating its superiority over existing methods in mitigating forgetting and enhancing fine-tuning performance.

## 1 Introduction

The success of large language models (LLMs) lies in their strong capabilities in language understanding and generation.
Typically, LLMs are initially pre-trained on extensive corpora to acquire general capabilities, and subsequently, they are fine-tuned on smaller, task-specific datasets to adapt to particular tasks or domains (Dai and Le, ; Kenton and Toutanova, ; Radford et al., ). However, it has been observed that during the fine-tuning process, LLMs may forget the knowledge acquired in pre-training, leading to a decline in general capabilities (Lin et al., ; Chen et al., ; Dong et al., ; Korbak et al., ; Luo et al., ).
Therefore, addressing the issue of forgetting during fine-tuning has become an important research direction for LLMs.

In the literature, two classes of methods are commonly adopted to mitigate the forgetting: replay-based methods, and regularization-based methods.

- •
Replay-based methods leverage pre-training data during the fine-tuning process (Rolnick et al., ; Wang et al., ; Ouyang et al., ).
However, most open-source LLMs, such as the Llama series (Touvron et al., ), have not fully disclosed their pre-training datasets. Moreover, even with access to pre-training data, incorporating it into the fine-tuning process may significantly increase computational and memory costs.
- •
Regularization-based methods add penalty terms to the loss function,
encouraging the fine-tuned model to remain close to the pre-trained model, thereby reducing the risk of forgetting pre-training knowledge (Li et al., ; Kirkpatrick et al., ; Miceli-Barone et al., ; Panigrahi et al., ).
However, as we will present later, modifying the original loss function during fine-tuning may impair the model’s performance on the fine-tuning task.

In this paper, we design a replay-free and regularization-free method to mitigate forgetting during the fine-tuning process. We propose the Momentum-Filtered Optimizer (MoFO).
At each iteration, MoFO selects and updates only the parameters with the top $\alpha\%$ largest momentum magnitudes within each block, where $\alpha$ is the filtering hyperparameter. Here, the blocks refer to the parameters of different parts of the network (e.g., weight matrices and bias terms).

Our method is motivated by the following observation. The fine-tuning loss of LLMs has many minima(^1^11In this paper, we use the term ”minimum” (or ”minima” in the plural) to refer to a parameter configuration whose fine-tuning loss is near its lowest value in a small neighborhood, while acknowledging that this terminology may not strictly represent a local minimum of the fine-tuning loss function.) and these minima can vary significantly in their distances to the pretrained-model. We notice that the minima closer to the pre-trained model are less likely to experience forgetting. Inspired by this observation, MoFO manages to approach closer minima by dynamically updating a portion of parameters that are most effective for reducing fine-tuning loss at each iteration. Therefore, MoFO can reduce the risk of forgetting without sacrificing the fine-tuning performance.

Our contributions are as follows:

- •
We propose the MoFO algorithm, a new optimization method designed to mitigate the forgetting of pre-training knowledge during the fine-tuning process.
- •
We present an initial convergence analysis of the MoFO algorithm.
- •
We conduct experiments
on various tasks, demonstrating that
MoFO outperforms existing methods both in fine-tuning performance and mitigating forgetting.

## 2 Momentum Filtered Optimizer (MoFO)

### 2.1 Motivation

During fine-tuning, different training methods usually converge to different minima. We observe that these minima share similar fine-tuning loss but can vary significantly in their distances to the pre-trained model. Furthermore, minima closer are less likely to forget pre-training knowledge. Below, we provide an illustrative example of this phenomenon.

We conduct an experiment using the Pythia-160m model to illustrate this observation. We fine-tune this model on a subset of the FLAN dataset(^2^22The subset we use is ‘definite_pronoun_resolution_10templates,’ available at [https://huggingface.co/datasets/Muennighoff/flan](https://huggingface.co/datasets/Muennighoff/flan). This is a preprocessed version of the FLAN dataset that incorporates updates made to the FLAN datasets since the release of the original FLAN.) using two different optimizers: the Adam optimizer (Kingma and Ba, ) and the Lion optimizer (Chen et al., ). Figure LABEL:pythia_landscape_lion demonstrates that Adam and Lion converges to two different minima of the fine-tuning loss.
Figure LABEL:pythia_landscape_lion(a) reveals that the two minima share similar fine-tuning loss. The minimum reached by Adam is significantly closer to the pre-trained model. Figure LABEL:pythia_landscape_lion(b) indicates that the minimum reached by Adam (the closer minimum) has a lower pre-training loss than that reached by Lion (the farther minimum). Furthermore, we evaluate the forgetting of pre-training knowledge in some common sense reasoning tasks, which include HellaSwag (Zellers et al., ), ARC-Challenge, and ARC-Easy (Clark et al., ). Table 1 shows that Adam suffers from less accuracy degradation on average, indicating better preservation of the pre-training knowledge.

Therefore, we conjecture that the severity of forgetting during fine-tuning correlates with the distance between the fine-tuned model and the pre-trained model: in general, models that remain closer to the pre-trained model are less likely to experience forgetting. Consequently, we consider whether it is possible to find a method superior to the Adam optimizer that not only reaches the minimum during fine-tuning but also converges closer to the pre-trained model, thereby mitigating forgetting more effectively.

**Table 1: Pythia-160m’s accuracies on common sense tasks, after being fine-tuned with the Adam optimizer and Lion optimizer. Adam achieves less forgetting than Lion on average.**
|  | HellaSwag | ARC-easy | ARC-challenge | Average |
| --- | --- | --- | --- | --- |
| Pythia-160m | 30.1 | 39.6 | 23.8 | 31.2 |
| Adam | 28.3 | 37.4 | 22.1 | 29.3 |
| Lion | 26.5 | 29.0 | 24.1 | 26.5 |

We first discuss how to keep the model closer to the pre-trained model.
To achieve this goal, we recall the classical block coordinate descent (BCD) method (Tseng, ).
We believe that the BCD algorithm may converge to a minimum that is closer to the pre-trained model than the default full parameter fine-tuning with Adam.
This is because BCD only updates a subset of parameters at each iteration. Compared to full parameter updates, this method usually involves smaller adjustments to the parameters.

The remaining issue is how to design BCD that reaches a similar performance to the default methods.
An effective strategy is to prioritize updating parameters that have the greatest influence on reducing fine-tuning loss. A straightforward approach is to measure the parameter’s influence by the magnitude of its gradient. However, in the widely used Adam optimizer, momentum directly affects parameter updates, while gradients influence parameter updates indirectly by affecting the momentum.

Motivated by these discussions, to mitigate forgetting and achieve comparable performance in fine-tuning tasks, we will modify the Adam optimizer by updating the subset of parameters with the largest momentum magnitude. We will discuss more details in the next section.

Figure: Figure 2: Illustration of MoFO. Compared with Adam optimizer, MoFO updates only the parameters with the largest $\alpha\%$ momentum in each partition.
Refer to caption: extracted/5766953/figure/MOFO.jpg

### 2.2 Algorithm Formulation

We formally introduce the Momentum-Filtered Optimizer (MoFO) in Algorithm 1.
MoFO partitions all the parameters into $B$ fixed parts as shown in Line 4. At each iteration, MoFO selects the parameter entries with the largest $\alpha\%$ momentum magnitudes in each part as shown in Lines 10-13 of Algorithm 1, where the update fraction $\alpha\%$ is the pre-determined hyperparameter. The momentum filtering mechanism is illustrated in Figure 2.

We note that the NN parameters are naturally composed of different parts (e.g., weight matrices, bias terms) in the network architecture, and PyTorch’s backward propagation mechanism automatically returns the gradients of the loss with respect to each part of the parameters. To reduce computational complexity, we partition all parameters according to these fixed parts and select $\alpha\%$ of the parameter entries for BCD update within each part.

In Section 3.4, we will empirically demonstrate that MoFO’s momentum-based selection rule outperforms its gradient-based variant in fine-tuning tasks. It indicates that the momentum-based selection rule allows for better incorporation with Adam during the optimization process in fine-tuning.

MoFO efficiently selects and updates the most influential parameters, as dictated by the momentum’s magnitude, thus enhancing the fine-tuning process while alleviating the catastrophic forgetting of pre-training knowledge.

Figure: Algorithm 1 Momentum Filtered Optimizer (MoFO)

### 2.3 Initial Exploration

In this section, we empirically examine whether an LLM fine-tuned with MoFO converges to a minimum closer to the pre-trained model compared to the one fine-tuned with the Adam optimizer, and whether MoFO mitigates catastrophic forgetting in LLMs.

We fine-tune the Pythia-160m on the same dataset used in the experiment described at the beginning of Section 2, using both the Adam optimizer and MoFO. First, as shown in Figure LABEL:pythia_landscape(a), both MoFO and Adam optimizer achieve minimal fine-tuning loss, so switching Adam to MoFO does not lead to performance degeneracy. Second, as shown in Figure LABEL:pythia_landscape(b), the distance from the pre-trained model to the minimum reached by MoFO is approximately 20% of the distance to that reached by the Adam optimizer. This shows that MoFO significantly reduces the amount of parameter movement.

**Table 2: Pythia-160m’s performance on common sense tasks, after being fine-tuned with the Adam optimizer and MoFO. The results indicate that MoFO significantly mitigates catastrophic forgetting. Bold values denote the best results among these optimizers.**
|  | HellaSwag | ARC-easy | ARC-challenge | Average |
| --- | --- | --- | --- | --- |
| Pythia-160m | 30.1 | 39.6 | 23.8 | 31.2 |
| Adam | 28.3 | 37.4 | 22.1 | 29.3 |
| MoFO | 29.9 | 42.0 | 22.9 | 31.6 |

We further verify whether the reduced parameter movement (by MoFO) effectively mitigates the forgetting of general capabilities acquired by pre-training.
We evaluate the degree of forgetting from two perspectives: pre-training loss and evaluation benchmarks. First, in Figure LABEL:pythia_landscape(b), we find MoFO has a lower pre-training loss compared to the Adam optimizer, indicating that MoFO memorizes pre-training data better. Second, in Table 2, we evaluate the accuracies of Pythia-160m, fine-tuned using both the Adam optimizer and MoFO, on some widely-used common sense tasks, which measures the commonsense reasoning capabilities of
LLMs.
The results show that MoFO achieves less accuracy degradations on all three tasks than the Adam optimizer.
This indicates that MoFO not only helps maintain a smaller pre-training loss but also more effectively mitigates the forgetting of general capabilities.

In Section 4, we will provide further insights on why MoFO converges to a point closer to the pre-trained model.

### 2.4 Convergence Result

In this section, we conduct a convergence analysis of MoFO. We study a simplified version of our MoFO algorithm as a variant of gradient descent (GD), which is described by Algorithm 2.

Figure: Algorithm 2 GD version of MoFO

###### Theorem 1 (Convergence of MoFO) .

Suppose that the minimum value of the loss function is ${\mathcal{L}}^{*}$ and the gradient $\nabla{\mathcal{L}}$ is Lipschitz continuous with constant $L$. Then for the GD version of MoFO in Algorithm 2, if the learning rates satisfy $\eta_{t}=1/L$, it holds that

$$ $\min_{0\leq t<T}\|g_{t}\|_{\infty}=O(T^{-\frac{1}{2}}).$ $$

In summary, we demonstrate the convergence of a GD version of MoFO, providing theoretical support for the strong performance of MoFO in fine-tuning tasks. We note that it seems rather non-trivial to prove the convergence of the original version of MoFO: MoFO contains 1st and 2nd-order momentum in a fractional form, and such structure is known to be challenging to handle (Zhang et al., ). We leave the convergence of the original MoFO method as an interesting future direction.

###### Theorem 1 (Convergence of MoFO) .

## 3 Experiments

Now we verify the effectiveness of MoFO on instruction fine-tuning and continual fine-tuning.
We use Llama-2-7B (Touvron et al., ) and TinyLlama-1.1B (Zhang et al., ) as the base models for our experiments. We also provide studies on the choice of update fraction and update strategies of MoFO.

### 3.1 Experimental Settings

#### Datasets for instruction fine-tuning.

This group of datasets covers question-answer pairs from different domains like mathematical reasoning and code generation. Specifically, the datasets include:

- •
MetaMathQA (Yu et al., ). This dataset comprises 395K math question-answer pairs. Numerous studies indicate that LLMs significantly enhance performance metrics on mathematical benchmarks such as GSM8K after fine-tuning on this dataset. In this paper, we randomly select 10% of this dataset for training LLMs, which includes 33,000 question-answer pairs.
- •
Code-Alpaca (Chaudhary, ). This dataset contains 20,022 question-answer pairs. We fine-tune LLMs on this dataset to enhance their coding capabilities.

#### Datasets for continual fine-tuning.

We investigate the performance of MoFO in the continual fine-tuning scenario by implementing our approach on the TRACE benchmark dataset (Wang et al., ). TRACE benchmark is designed with a comprehensive set of 8 distinct tasks across various domains, including domain-specific knowledge, multilingual proficiency, code generation, and mathematical reasoning.

#### Metrics for instruction fine-tuning.

We introduce a set of widely used benchmarks to assess the performance and catastrophic forgetting effects on the general capabilities of LLMs after instruction fine-tuning. These benchmarks include:

- •
MMLU (Massive Multitask Language Understanding) (Hendrycks et al., ). It is a popular benchmark to evaluate factual knowledge of LLMs. This benchmark spans 57 diverse subjects, ranging from STEM fields and the humanities to social sciences. We report the 0-shot accuracy.
- •
Commonsense. We employ the widely recognized benchmarks ARC-Challenge, ARC-Easy (Clark et al., ), and HellaSwag (Zellers et al., ) to measure the commonsense reasoning capabilities of LLMs. In this paper, we refer to these benchmarks collectively as the Commonsense benchmark, and we use the average of these three metrics as the evaluation for this Commonsense benchmark.
- •
GSM8K (Cobbe et al., ). This benchmark consists of 8.5K high-quality grade school math problems. We evaluate the math capability of LLM on the test set of GSM8K through the LM Eval Harness framework (Gao et al., ). We follow the default implementation setting of LM Eval Harness and set the temperature hyperparameter as 0 and report 5-shot accuracy.
- •
HumanEval (Chen et al., ).
We adopt the widely used HumanEval to assess the coding capabilities of LLMs. It comprises 164 unique programming problems. We report the pass@10 performance.

#### Metrics for continual fine-tuning.

To evaluate the LLM’s performance in continual learning, we consider two key metrics in this scenario: Overall Performance (OP) (Chaudhry et al., ) and BackWard Transfer (BWT) (Lopez-Paz and Ranzato, ). Let $R_{t,i}$ be the test accuracy on task $T_{i}$ after we have gone through tasks $T_{1},\dots,T_{t}$ ($1\leq i\leq t\leq T$). The OP score measures the average accuracy on all the $T$ tasks
after the continual fine-tuning process, which is defined by

$$ $OP:=\frac{1}{T}\sum_{i=1}^{T}R_{T,i}.$ $$

The BWT score measures the average accuracy change of each task after learning new tasks, which is defined by

$$ $BWT:=\frac{1}{T-1}\sum_{i=1}^{T-1}(R_{T,i}-R_{i,i}).$ $$

These metrics provide a comprehensive assessment of the model’s ability to learn incrementally while retaining knowledge from past experiences. In this paper, we will utilize these scores to evaluate the effectiveness of our method in continual learning tasks.

### 3.2 Instruction Fine-Tuning

In this section, we aim to investigate the effectiveness of the MoFO approach in preserving general capabilities while learning fine-tuning tasks. We fine-tune Llama-2-7B on MetaMathQA and Code-Alpaca datasets, and compare the performance of MoFO against several baseline methods. The evaluation includes changes in performance on fine-tuning tasks and general capability metrics, with the pre-trained model’s performance serving as the reference point for comparison. The implementation details can be seen in B.

#### Baselines

We compare the proposed MoFO algorithm with some of the most widely used optimization techniques, which aim to alleviate forgetting, as our baselines. These methods include:

- •
Full FT refers to default full-parameter fine-tuning approach with the loss function $\mathcal{L}_{finetune}(\theta)$.
- •
$L_{2}$-regularization (Kirkpatrick et al., ) adds an $L_{2}$ regularization term to the original loss function $\mathcal{L}_{finetune}(\theta)$ to keep the fine-tuning model closer to the pre-trained model, thereby mitigating forgetting. The modified fine-tuning loss function is $\mathcal{L}_{finetune}(\theta)+\lambda_{2}\|\theta-\theta_{0}\|^{2}_{2}$, with the regularization hyperparameter $\lambda_{2}$ set to 1e-3.
- •
$L_{1}$-regularization (Panigrahi et al., ) adds an $L_{1}$ regularization term to the original loss function $\mathcal{L}_{finetune}(\theta)$ to promote sparsity in parameter movement. The modified fine-tuning loss function is $\mathcal{L}_{finetune}(\theta)+\lambda_{1}\|\theta-\theta_{0}\|_{1}$, with the regularization hyperparameter $\lambda_{1}$ set to 1e-6.
- •
Half Fine-tuning (HFT) (Hui et al., ) randomly updates half of the parameter blocks within each transformer layer at each iteration while the other half are frozen. HFT can be considered a specific case of the BCD algorithm.

**Table 3: The performance on the math task (GSM8K) and the score changes in general capabilities of Llama-2-7B after fine-tuning on the MetaMathQA dataset. The results show that MoFO achieves comparable performance in the math task, while significantly mitigating catastrophic forgetting of general capabilities. Bold values denote the best results among these methods.**
| Method | Performance |  | General Capability |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| GSM8K |  | Commensense | MMLU | HumanEval | Average |  |
| Llama-2-7B | 13.7 |  | 65.6 | 42.0 | 24.2 | 43.9 |
| Full FT | 49.4 |  | 62.3 | 36.6 | 16.1 | 38.3 |
|  | -3.3 | -5.4 | -8.1 | -5.6 |  |  |
| HFT | 47.5 |  | 65.5 | 42.3 | 23.6 | 43.8 |
|  | -0.1 | +0.3 | -0.6 | -0.1 |  |  |
| $L_{1}$-regularization | 39.0 |  | 65.1 | 38.3 | 27.4 | 43.6 |
|  | -0.5 | -3.7 | +3.2 | -0.3 |  |  |
| $L_{2}$-regularization | 44.5 |  | 65.5 | 39.2 | 25.9 | 43.5 |
|  | -0.1 | -2.8 | +1.7 | -0.4 |  |  |
| MoFO ($\alpha\%=15\%$) | 47.7 |  | 65.7 | 42.7 | 24.6 | 44.3 |
|  | +0.1 | +0.7 | +0.4 | +0.4 |  |  |

#### Results of fine-tuning on MetaMathQA.

We fine-tune Llama-2-7B on MetaMathQA using various baseline methods and present the experimental results on mathematical reasoning (GSM8K) and general capabilities in Table 3. These results demonstrate the effectiveness of our proposed MoFO algorithm in both optimization and mitigating catastrophic forgetting.

MoFO is compatible to the performance of Full FT and HFT on the math task, yet significantly outperforms these methods in preserving general capability. Specifically, Full FT shows a decline of $5\%$ in MMLU accuracy, while HFT experiences a drop of $0.6\%$ in HumanEval. In contrast, our MoFO not only maintains but slightly improves these general capability metrics by an average of $0.4\%$.

Compared to $L_{1}$ and $L_{2}$ regularization, our MoFO significantly outperforms these baselines on GSM8K by more than $3.2\%$ and $8.7\%$, respectively. Additionally, while $L_{1}$ and $L_{2}$ regularizations suffer from significant catastrophic forgetting on MMLU, MoFO shows no decline.

**Table 4: The performance on the coding generation task (HumanEval) and accuracy changes in general capabilities of Llama-2-7B after fine-tuning on the Code-Alpaca dataset. The results show that MoFO achieves comparable performance in the HumanEval while significantly mitigating catastrophic forgetting of general capabilities. Bold values denote the best results among these methods.**
| Method | Performance |  | General Capability |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| HumanEval |  | Commensense | MMLU | GSM8K | Average |  |
| Llama-2-7B | 24.2 |  | 65.6 | 42.0 | 13.7 | 40.4 |
| Full FT | 27.7 |  | 65.2 | 39.6 | 5.5 | 36.8 |
|  | -1.4 | -2.4 | -8.2 | -3.6 |  |  |
| HFT | 30.9 |  | 68.3 | 40.6 | 6.8 | 38.6 |
|  | +2.7 | -1.4 | -6.9 | -1.8 |  |  |
| $L_{1}$-regularization | 33.8 |  | 67.6 | 41.3 | 5.1 | 38.0 |
|  | +2.0 | -0.7 | -8.6 | -2.4 |  |  |
| $L_{2}$-regularization | 30.8 |  | 66.1 | 42.5 | 7.6 | 38.7 |
|  | +0.5 | +0.5 | -6.1 | -1.7 |  |  |
| MoFO ($\alpha\%=10\%$) | 33.7 |  | 67.3 | 42.7 | 8.0 | 39.3 |
|  | +1.7 | +0.7 | -5.7 | -1.1 |  |  |

#### Results of fine-tuning on Code-Alpaca.

We fine-tune Llama-2-7B on Code-Alpaca using various baseline methods and present the experimental results on coding generation and general capabilities in Table 4.

MoFO performs well on the fine-tuning task of code generation. It significantly outperforms Full FT, HFT, and $L_{2}$ regularization on HumanEval accuracy by more than $3\%$, and shows comparable results to $L_{1}$ regularization with only a $0.1\%$ difference.

In terms of general capability, MoFO demonstrates the least degradation compared to other baselines, with an average accuracy reduction of only $1.1\%$. Specifically, on MMLU, Full FT, HFT, and $L_{1}$ regularization experience degradations ranging from $0.7\%$ to $2.4\%$, while our method not only preserves but slightly improves the MMLU accuracy. On the math benchmark GSM8K, although all methods experience forgetting, our proposed MoFO algorithm exhibits a smaller decline of $5.7\%$, while other baseline methods exhibits drops ranging from $6.1\%$ to $8.2\%$.

In summary, our MoFO algorithm shows competitive performance in instruction fine-tuning while preserving the general capabilities of pre-trained LLMs, effectively alleviating catastrophic forgetting.

### 3.3 Continual Fine-Tuning

In this section, we explore the performance of our proposed MoFO in continual fine-tuning on the TRACE benchmark (Wang et al., ).
We sequentially train TinyLlama-1.1B on the TRACE dataset, which includes the eight tasks from different domains. The implementation details can be seen in B.

#### Performance upper bound.

Multi-task Learning (MTL) mixes samples from all eight distinct tasks together during training and typically achieves the highest OP score. We will use MTL as an upper bound for performance comparison.

#### Some orthogonal methods.

We consider several traditional methods from the field of continual learning. These methods can be orthogonal combined with MoFO to further enhance performance.

- •
Replay involves optimizing the model using current data along with a memory buffer containing old samples from previous tasks to mitigate catastrophic forgetting.
- •
Gradient of Episodic Memory (GEM) mitigates catastrophic forgetting by using gradients from old tasks to adjust the parameter updates during the training of new tasks.

**Table 5: The OP and BWT scores of TinyLlama-1.1B after fine-tuning on TRACE benchmark. The results show that MoFO outperforms Full FT and HFT in continual learning and can combine well with traditional continual learning methods. Bold values denote the best results among these methods in each group.**
|  | OP | BWT |
| --- | --- | --- |
| Full FT | 38.4 | -10.3 |
| HFT | 39.9 | -10.1 |
| MoFO $(\alpha\%=5\%)$ | 41.3 | -5.4 |
| GEM | 40.8 | -8.5 |
| GEM + MoFO $(\alpha\%=5\%)$ | 41.7 | -6.7 |
| Replay | 45.5 | 4.7 |
| Replay + MoFO $(\alpha\%=5\%)$ | 47.0 | 4.8 |
| MTL | 50.8 |  |

#### Results of continual fine-tuning.

We present the experimental results of sequentially fine-tuning TinyLlama-1.1B on the TRACE benchmark with various methods in Table 5. The results indicate that in continual fine-tuning, MoFO outperforms Full FT and HFT by at least $1.4\%$ on the OP score and by at least $4.7\%$ on the BWT score. Moreover, MoFO combines well with the Replay method, offering a $1.5\%$ performance gain on the OP metric compared to using Replay alone. Additionally, when combined with the GEM method, MoFO provides a $0.9\%$ improvement on the OP metric compared to using GEM alone.

In summary, these results underscore the superior performance of MoFO in continual fine-tuning and its effectiveness in alleviating catastrophic forgetting.

### 3.4 Further Analysis

In this section, we first investigate the impact of the update fraction of parameters in the MoFO algorithm at each iteration, and then explore the effects of different update strategies within MoFO.

#### Impact of update fraction of parameters in MoFO.

Following the setting in Section 4.2, we fine-tune Llama-2-7B on the MetaMathQA dataset using MoFO with varying update fractions of parameters at each iteration for 2 epochs. The experimental results of math reasoning (GSM8K) and average general capability performance changes are presented in Figure LABEL:ablation_fraction.

The parameter update fraction affects the fine-tuning performance. Figure LABEL:ablation_fraction(a) shows that larger update fractions can improve MoFO’s optimization effectiveness. Furthurmore, MoFO with a $5\%$ parameter update fraction is sufficient to achieve nearly $90\%$ of the performance of Full FT. When the update fraction exceeds $40\%$, MoFO’s performance matches that of Full FT.

The parameter update fraction also affects the preservation of general capabilities. Figure LABEL:ablation_fraction(b) indicates that MoFO avoids forgetting in general capabilities when the parameter update fraction is below $20\%$. Beyond the threshold of $20\%$, further increases in the parameter update fraction lead to a decline in general capabilities. Despite this, MoFO still forgets significantly less than Full FT.

In summary, MoFO can preserve pre-training knowledge and significantly enhance fine-tuning performance by choosing a moderate update fraction, avoiding the extremes of too small or too large fractions.

**Table 6: The performance on the math reasoning task (GSM8K) and accuracy changes of general capabilities of Llama-2-7B after fine-tuning on MetaMathQA using different updating strategies in MoFO.**
| Method | Performance |  | General Capability |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| GSM8K |  | Commensense | MMLU | HumanEval | Average |  |
| Pre-trained model and Full FT |  |  |  |  |  |  |
| Llama-2-7B | 13.7 |  | 65.6 | 42.0 | 24.2 | 43.9 |
| Full FT | 49.4 |  | 62.3 | 36.6 | 16.1 | 38.3 |
|  | -3.3 | -5.4 | -8.1 | -5.6 |  |  |
| BCD based Methods ($\alpha\%=10\%$) |  |  |  |  |  |  |
| Randomized BCD | 35.0 |  | 65.8 | 41.1 | 25.1 | 44.0 |
|  | +0.2 | -0.9 | +0.9 | +0.1 |  |  |
| Gradient-filtered BCD | 40.2 |  | 66.0 | 41.6 | 28.0 | 45.2 |
|  | +0.4 | -0.4 | +3.8 | +1.3 |  |  |
| MoFO | 45.4 |  | 65.7 | 43.5 | 27.4 | 45.5 |
|  | +0.1 | +1.5 | +3.2 | +1.6 |  |  |

#### Impact of update strategy in MoFO.

In addition to MoFO, we consider two other BCD methods, randomized BCD and gradient-filtered BCD. Randomized BCD updates a random subset of parameters at each iteration. Gradient-filtered BCD selects the filter based on gradient magnitudes rather than the momentum magnitudes used in MoFO. Specifically, line 11 in Algorithm 1 is replaced by:

$$ $(\texttt{FILTER}_{t}^{(k)})_{i}=1\text{ if }|(g^{(k)}_{t})_{i}|\text{ is among the top-}\alpha\%\text{ of all }|g^{(k)}_{t}|\text{'s entries else }0.$ $$

We fine-tune Llama-2-7B on MetaMathQA using these three methods with $10\%$ parameter update fraction and present the results in Table 6.

Experimental results show that all three BCD methods exhibit significantly less forgetting compared to Full FT, demonstrating the effectiveness of BCD algorithms in mitigating catastrophic forgetting.

In terms of GSM8K performance, our proposed MoFO method significantly surpasses both Gradient-filtered BCD and randomized BCD, indicating that updating parameters with the largest momentum leads to strong optimization power.

## 4 Why MoFO Converges to a Closer Point

In Section 3, we discuss how the distance of parameter movement during fine-tuning impacts the retention of pre-training knowledge. Drawing on these insights and the BCD method, we introduce MoFO. Here, we propose the following question as follows.

Question: Why does MoFO converge closer to the pre-trained LLMs than those of Adam?

We attempt to answer this question by the following toy example. We denote $\Theta=(\theta_{1},\theta_{2})\in\mathbb{R}^{2}$ to be the trainable parameters of our model and make the following assumptions:

- •
The pre-training loss is $\mathcal{L}_{pretrain}(\Theta)=\theta_{1}^{2}+\theta_{2}^{2}$ and the model has been trained to the global minimum $\Theta_{0}=(0,0)$ during the pre-trained phase.
- •
The fine-tuning loss is $\mathcal{L}_{finetune}(\Theta)=(\theta_{1}-1)^{2}(\theta_{2}-1)^{2}$. In this case, any global optimum of $\mathcal{L}_{finetune}$ lies in the set $\{(1,\theta_{2}):\theta_{2}\in\mathbb{R}\}\cup\{(\theta_{1},1):\theta_{1}\in
\mathbb{R}\}$, which is a union of two straight lines.

For full-parameter fine-tuning with Adam, starting from $\Theta_{0}=(0,0)$, the model converges to $(1,1)$ during the fine-tuning phase along the orange arrow in Figure LABEL:example_landscape, with a pre-training loss of 2. In contrast, when applying MoFO, the model converges to $(1,0)$ during the fine-tuning phase along the green arrow in Figure LABEL:example_landscape, resulting in a pre-training loss of 1. This demonstrates that MoFO can converge to a minimum that is closer to the pre-training model, thereby mitigating forgetting.

Insights:
In this example, we find that when a loss function has multiple distinct minima, they can be considered as different attractors. These attractors can influence the gradient direction of a pre-trained model, possibly drawing the model’s weights away from the nearest minimum. Specifically, full-parameter gradient descent based methods may converge to the balanced point of these attractors’ influences, which is the orange point in Figure LABEL:example_landscape(a). On the contrary, MoFO addresses this issue by updating only a subset of parameters during each iteration. This selective updating rule reduces interference among attractors, allowing the model to converge to a closer minimum.

## 5 Related Works

#### Catastrophic forgetting

Catastrophic forgetting, a significant issue where models forget previously learned information upon learning new data, has received considerable attention in machine learning (McCloskey and Cohen, ; Goodfellow et al., ; Kemker et al., ; Ramasesh et al., ; Liu et al., ). In the realm of LLMs, there has been a growing body of recent works (Luo et al., ; Kotha et al., ; Shi et al., ; Wu et al., ) focusing on the forgetting of models’ knowledge during the fine-tuning process.

Researchers have proposed numerous methods to alleviate forgetting in continual learning, which involves learning a sequence of tasks. These methods are not limited to learning in a sequential manner and can be applied to broader paradigms. Generally, three primary approaches are used: replay-based methods, regularization-based methods, and architecture-based methods.

- •
Replay-based methods leverage past experiences to facilitate the learning of new tasks. The most straightforward implementation of this approach is experience replay, which involves maintaining old samples in a buffer and replaying them during incremental training (Rolnick et al., ) with some variants enhancing performance or optimizing memory usage (Aljundi et al., ; Hayes et al., ; Cha et al., ; Chaudhry et al., ; Riemer et al., ). Several other variants utilize gradient information from past tasks (Lopez-Paz and Ranzato, ; Riemer et al., ; Chaudhry et al., ; Farajtabar et al., ; Aljundi et al., ; Chaudhry et al., ; Tiwari et al., ). In LLMs, several works (Yin et al., ; Wang et al., ; Ouyang et al., ) propose replay-based methods to mitigate forgetting. We emphasize that our MoFO method is orthogonal to replay-based methods. MoFO can be integrated into these replay strategies to further enhance their effectiveness.
- •
Regularization-based methods introduce constraints to preserve old knowledge.
Several studies add regularization terms to the loss functions to penalize parameter changes and mitigate forgetting (Kirkpatrick et al., ; Aljundi et al., ; Zenke et al., ; Li et al., ; Ritter et al., ; Kumar et al., ). Some other works apply regularization to the embedding or output changes (Li and Hoiem, ; Rannen et al., ; Buzzega et al., ; Huang et al., ). Unlike these approaches, MoFO does not modify the loss function and as a result, MoFO reached better fine-tuning performance (see Table 3 and 4).
- •
Architecture-based methods balance the goals of learning new knowledge and keeping the old one through model architectural modifications. LoRA (Hu et al., ), as the most popular parameter-efficient fine-tuning (PEFT) methods, modifies the model architecture by freezing the pre-training weights and introducing low-rank trainable matrices. Empirical works shows that LoRA forgets less but learns less during fine-tuning (Biderman et al., ). Some variants of LoRA find applications in continual learning of LLMs (Ren et al., ; Wang et al., ). In comparison, our MoFO still allows for high-rank updates to achieve better fine-tuning performance.
Another line of architecture-based methods focuses on model merging. The idea stem from the understanding that the task-specific knowledge is located at a small subspace of the weight space (Panigrahi et al., ; Gueta et al., ; Zhu et al., ). Consequently, various model merging methods have been proposed to simultaneously retain pre-training knowledge and improve fine-tuning task performance (Panigrahi et al., ; Yadav et al., ; Yu et al., ). However, these methods require an additional post-fine-tuning process before model merging. In contrast, our method only requires only one fine-tuning stage.

#### Block coordinate descent

Block Coordinate Descent (BCD) involves iteratively optimizing over a block of coordinates while holding the others constant. The foundational work of Tseng () provides a comprehensive analysis of the convergence properties of BCD under certain conditions. Subsequent research has explored various BCD variants (Hong et al., ), including randomized
BCD (Nesterov, ; Richtárik and Takáč, ; Lu and Xiao, ), cyclic BCD (Sun and Hong, ), and greedy BCD (Nutini et al., ). Among these, the greedy variant, also known as Gauss-Southwell BCD method, has drawn attention due to its ability to prioritize coordinates that yield the most substantial improvement in each iteration, thereby potentially accelerating convergence.

In the realm of machine learning, BCD has also found applications (Nutini et al., ). For example, Luo et al. () leverages BCD to perform memory-efficient fine-tuning of LLM and Xu and Zhang () use random masking to perform this. In federated learning, Rothchild et al. () adopts top-$k$ momentum value unsketch rather than our top-$k$ momentum filtering to tackle communication bottleneck and convergence issues.
A recent work, Hui et al. () addresses catastrophic forgetting during fine-tuning of LLMs by selectively freezing half of the parameters during training. Our approach is akin to a more efficient greedy BCD, achieving superior performance in fine-tuning tasks and alleviating forgetting better.

## 6 Conclusion

This paper presents the Momentum-Filtered Optimizer (MoFO), a new approach designed to mitigate the crucial issue of pre-training knowledge forgetting in LLMs during fine-tuning. By selectively updating the parameters with the largest momentum magnitudes in each parameter block, MoFO converges to a point closer to the pre-trained model
compared to full-parameter fine-tuning and effectively preserves pre-trained knowledge. Our experimental results demonstrate that MoFO not only significantly alleviates catastrophic forgetting but also surpasses the performance of traditional fine-tuning methods. Future work will explore further optimizations and potential applications of MoFO in multimodal LLMs.

## Appendix A Proof of Theorem 1

Before providing the proof, we will give some preliminary information.

Let’s assume the parameter space $\mathbb{R}^{d}$ is decomposed as $\mathbb{R}^{d}\cong\mathbb{R}^{d_{1}}\times\mathbb{R}^{d_{2}}\times\dots\times
\mathbb{R}^{d_{B}}$. Here, $B$ represents the number of parameter matrices in the neural network, and $d_{k}$ denotes the product of the dimensions (i.e., the number of rows multiplied by the number of columns) of the $k$-th parameter matrix. For any $z\in\mathbb{R}^{d}$, we denote

$$ $z={\rm Concat}(z^{(1)};z^{(2)};\dots;z^{(B)}),$ $$

where $z^{(k)}\in\mathbb{R}^{d_{k}}$ for any $1\leq k\leq B$.

###### Definition 1 .

For any $z\in\mathbb{R}^{d}$, we define

$$ $\|{z}\|_{\text{top-}\alpha\%}:=\|z\odot{\mathbf{e}}_{S}\|_{2},$ $$

where ${\mathbf{e}}_{S}={\rm Concat}({\mathbf{e}}^{(1)}_{S_{1}};{\mathbf{e}}^{(2)}_{S
_{2}};\dots;{\mathbf{e}}^{(B)}_{S_{B}})$. Here,

$$ $S_{k}=\{i\in[d_{k}]:|z^{(k)}_{i}|\text{ is among the top-}\alpha\%\text{ of all }|z^{(k)}|\text{'s entries }(|z^{(k)}_{1}|,|z^{(k)}_{2}|,\dots,|z^{(k)}_{d _{k}}|)\}$ $$

and ${\mathbf{e}}_{S_{k}}^{(k)}$ is a $d_{k}$-dimensional vector where the $i$-th entry is 1 if $i\in S_{k}$, and 0 otherwise.

###### Lemma 1 .

$\|{\cdot}\|_{\text{top-}\alpha\%}$ is a norm in ${\mathbb{R}}^{d}$.

###### Proof.

By Definition 1, we get

$$ $\|{z}\|_{\text{top-}\alpha\%}^{2}=\|z\odot{\mathbf{e}}_{S}\|^{2}_{2}=\sum_{k=1 }^{B}\|z^{(k)}\odot{\mathbf{e}}_{S_{k}}^{(k)}\|^{2}_{2}$ (1) $$

First, if $\|{z}\|_{\text{top-}\alpha\%}=0$, then by (1), $\|z^{(k)}\odot{\mathbf{e}}_{S_{k}}^{(k)}\|_{2}=0$ for any $1\leq k\leq B$. Thus,

$$ $\|z^{(k)}\|_{\infty}=\mathop{\rm argmax}_{1\leq i\leq d_{k}}|z^{(k)}_{i}|\leq \|z^{(k)}\odot{\mathbf{e}}_{S_{k}}^{(k)}\|_{2}=0.$ $$

So $z^{(k)}$ is a zero vector for any $1\leq k\leq B$ and then $z$ is a zero vector.

Second, for any given $c\in\mathbb{R}_{+}$, $\{|z^{(k)}_{i}|\}_{1\leq i\leq d_{k}}$ and $\{|cz^{(k)}_{i}|\}_{1\leq i\leq d_{k}}$ have the same order. So $z$ and $cz$ share the same filter ${\mathbf{e}}_{S}$ and

$$ $\|{cz}\|_{\text{top-}\alpha\%}=\|cz\odot{\mathbf{e}}_{S}\|_{2}=c\|z\odot{ \mathbf{e}}_{S}\|_{2}=c\|{z}\|_{\text{top-}\alpha\%}.$ $$

Third, for the vectors $x,y\in\mathbb{R}^{d}$, we denote the greedy filter of $x,y,x+y$ by $S,S^{\prime},S^{\prime\prime}$, respectively.

$$ $\displaystyle\|{x+y}\|_{\text{top-}\alpha\%}$ $\displaystyle=\|(x+y)\odot{\mathbf{e}}_{S^{\prime\prime}}\|_{2}$ $\displaystyle=\|x\odot{\mathbf{e}}_{S^{\prime\prime}}+y\odot{\mathbf{e}}_{S^{ \prime\prime}}\|_{2}$ $\displaystyle\leq\|x\odot{\mathbf{e}}_{S^{\prime\prime}}\|_{2}+\|y\odot{ \mathbf{e}}_{S^{\prime\prime}}\|_{2}$ $\displaystyle\leq\|x\odot{\mathbf{e}}_{S}\|_{2}+\|y\odot{\mathbf{e}}_{S^{ \prime}}\|_{2}$ $\displaystyle=\|{x}\|_{\text{top-}\alpha\%}+\|{y}\|_{\text{top-}\alpha\%}.$ $$

∎

###### Proof of Theorem 1 .

By the definition of the simple version of MOFO in Algorithm 2, we have

$$ $\Theta_{t+1}-\Theta_{t}=-\eta_{t}\hat{m}_{t}\odot\texttt{FILTER}_{t}=-\eta_{t} g_{t}\odot\texttt{FILTER}_{t},$ $$

where $\texttt{FILTER}_{t}$ is actually ${\mathbf{e}}_{S}$ in Definition 1.

By the Lipschitz condition of $\nabla\cal{L}$, we have

$$ $\displaystyle{\mathcal{L}}(\Theta_{t+1})$ $\displaystyle\leq{\mathcal{L}}(\Theta_{t})+\langle g_{t},\Theta_{t+1}-\Theta_{ t}\rangle+\frac{L}{2}\|\Theta_{t+1}-\Theta_{t}\|_{2}^{2}$ $\displaystyle={\mathcal{L}}(\Theta_{t})+\left(\frac{L}{2}\eta_{t}^{2}-\eta_{t} \right)\|g_{t}\odot\texttt{FILTER}_{t}\|^{2}_{2}$ $\displaystyle={\mathcal{L}}(\Theta_{t})+\left(\frac{L}{2}\eta_{t}^{2}-\eta_{t} \right)\|g_{t}\odot\texttt{FILTER}_{t}\|^{2}_{2}$ $\displaystyle\leq{\mathcal{L}}(\Theta_{t})-\frac{1}{2L}\|{g_{t}}\|_{\text{top- }\alpha\%}^{2},$ $$

where the second inequality becomes an equality if the learning rate $\alpha=1/L$.

Thus, we conclude that

$$ $\min_{0\leq t<T}\|{g_{t}}\|_{\text{top-}\alpha\%}\leq\left(\frac{1}{T}\sum_{t= 0}^{T-1}\|{g_{t}}\|_{\text{top-}\alpha\%}^{2}\right)^{\frac{1}{2}}\leq\left( \frac{2L}{T}({\mathcal{L}}(\Theta_{0})-{\mathcal{L}}(\Theta_{T}))\right)^{ \frac{1}{2}}\leq\frac{1}{\sqrt{T}}\left(2L({\mathcal{L}}(\Theta_{0})-{\mathcal {L}}^{*})\right)^{\frac{1}{2}}.$ $$

∎

###### Definition 1 .

###### Lemma 1 .

###### Proof.

###### Proof of Theorem 1 .

## Appendix B Implementation Details

Instruction fine-tuning. In our instruction fine-tuning experiments, we follow the implementation of Ivison et al. []. The learning rate is set to 2e-5 with a cosine decay scheduler. For fine-tuning on the MetaMathQA dataset, we set the maximum sequence length to 1024, the batch size to 128, and we train the Llama-2-7B for 2 epochs. The parameter update fraction for MoFO is set to $15\%$. For fine-tuning on the Code-Alpaca dataset, we set the batch size as 64 and the parameter update fraction as $10\%$, while keeping other settings the same.

Continual fine-tuning. In our continual fine-tuning experiments, we follow the default settings of the TRACE benchmark. We sequentially train TinyLlama-1.1B on the TRACE benchmark datasets: C-STANCE, FOMC, MeetingBank, Py150, ScienceQA, NumGLUE-cm, NumGLUE-ds, and 20Minuten for 5, 3, 7, 5, 3, 5, 5, and 7 epochs, respectively. We use a learning rate of 1e-5 with a cosine decay schedule and a batch size of 64. The parameter update fraction for MoFO is set to $5\%$.

All experiments are conducted on four A800 (80GB) GPUs.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. The fraction of MoFO was varied in the experiments. Appendix E.1 only discusses its impact on the Llama-2-7B model. How does this hyperparameter affect models of different sizes?

2. MoFO updates only a small subset of parameters at each iteration, and the learning rate decreases over time. Will this lead to reduced fine-tuning efficiency?
soundness: 3
strengths: 1. This paper provides a comprehensive and reasonable proof of the convergence.
2. The effectiveness of the method is verified with LLMs of different sizes.
3. The convergence direction of MoFO and baselines is visualized using loss landscape, proving that MoFO converges to a closer point.
summary: The paper presents a novel fine-tuning algorithm named Momentum-Filtered Optimizer (MoFO). The primary contribution of the paper is addressing the issue of catastrophic forgetting in LLMs during fine-tuning. MoFO selectively updates model parameters with the largest momentum magnitudes to maintain proximity to the pre-trained state, thereby mitigating knowledge forgetting. The authors claim that MoFO achieves similar performance to the default fine-tuning algorithm while effectively preserving pre-training knowledge, without requiring access to pre-training data. The paper is rigorous in its convergence analysis and experiments, demonstrating a good trade-off in maintaining general capabilities and downstream task performance.
weaknesses: In Section 4, the loss landscapes for HFT and LoRA were not reported.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Add discussions and reference for more works (see list in Weakness 1) on LLM CF mitigation, especially justify your advantage compared to them
2. Add baselines (see list in Weakness 1) and comparison studies (i.e. is the combination of MoFO and replay better than both MoFO only and replay only?) correspondingly in experimental studies. You might also need to better verify (possibly by adding references) why the experimental success of the combination of MoFO with your choice of HFT, GEM and replay would be representative enough to suggest the effectiveness of combining MoFO with othorgonal strategies.
soundness: 3
strengths: 1. The paper deals with an important problem in LLM application and gives a practical solution with fair theoretical support.
2. The illustation and analysis of MoFO strategy is clear and in detail.
3. Authors supported their method with experiments covering various tasks and settings.
summary: This paper presents a novel fine-tuning algorithm called the Momentum-Filtered Optimizer (MoFO), designed to address the issue of knowledge forgetting in large language models (LLMs) during fine-tuning. As a replay and regularization free strategy, MoFO iteratively selects and updates model parameters based on the largest momentum magnitudes, extending greedy block coordinate descent (BCD) methods. The algorithm achieves comparable fine-tuning performance to standard methods while effectively mitigating knowledge forgetting, without requiring access to pre-training data. The authors validate MoFO through rigorous convergence analysis and extensive experiments, demonstrating its superiority in reducing forgetting, making it particularly suitable for scenarios involving checkpoint-only open-source LLMs.
weaknesses: 1. The discussion of this paper did not cover other algorithms that are non replay and non regularization, and therefore doesn't give enough support for the motivation and novelty of their work. 
(including but not limited to: https://arxiv.org/pdf/2309.06256, https://arxiv.org/pdf/2404.10306, https://arxiv.org/abs/2302.03241 etc.)
A brief literature review section might be added specifically comparing MoFO to these other approaches, highlighting key differences and potential advantages.
2. The number and quality of baselines compared in experiments are too limited; the experimental validation of orthogonal methods is somehow unconvincing without proper comparisons.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. In Table 1, why LoRA perform well on HumanEval but worse on other tasks? Also, in both Table 1 and Figure 4, the performance of MoFO is not significantly better than HFT.

2. The study of the MoFO approach applied on other optimization methods beyond Adam is recommended.

3. I would suggest the authors include  an efficiency analysis for MoFO.
soundness: 3
strengths: 1. The paper provides a simple but elegant approach to mitigating forgetting in LLM fine-tuning.

2. Empirical results are convincing.

3. The paper is well-written with illustrative figures.
summary: This paper proposes a new method, named Momentum-Filtered Optimizer (MoFO) for finetuning LLMs while mitigating the forgetting of the knowledge captured by the pre-trained model. The proposed method leverages block coordinate descent (BCD), where the blocks of parameters to update are chosen based on the top-\alpha largest momentum values. The convergence result of MoFO is discussed. Empirical results are also provided to justify the advantages of MoFO.
weaknesses: 1. Using the default partitioning of model parameters as implemented in PyTorch is quite an ad-hoc choice. Is there any better or more principled way to partition the model parameters?

2. The convergence result is good, but it does not show the advantages of MoFO over the traditional Adam method, especially in mitigating forgetting.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 4
questions: N/A
soundness: 3
strengths: 1. The paper follows a clear, well-thought-out structure that facilitates easy reading and comprehension;

2. The experimental results appear encouraging.
summary: The authors focus on mitigating forgetting of pre-training knowledge in LLM fine-tuning. They propose update models with only a subset of parameters based on momentum during fine-tuning. Experiments on public datasets are conducted to demonstrate the effectiveness of the proposed method.
weaknesses: 1. In Section 2.1, the authors show a strong correlation between forgetting and the distance from the pre-trained model with experiments. However, no analysis or proof is provided to explain the reason.

2. Considering that MoFO only uses a subset of parameters during fine-tuning, it may require more time or iterations compared to other methods to reach a minimum loss. It could be helpful to provide a time cost analysis section.

3. Clarification on the selection process for $\alpha$ would enhance the understanding of this parameter's role.

4. Table 5 demonstrates that MoFO outperforms Gradient-filtered BCD and MV-filtered BCD. A deeper analysis of the factors contributing to MoFO's superior performance is needed to provide further insight into these results.

5. In Fig. 4, the differing numbers of points for various methods and the significance of lighter colors could be further clarified for better interpretability.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. In Section 2.1 and Figure 2, how do you exactly obtain the distance between the original pre-trained state and fine-tuning model's state? Also, could you explain your experiment setting here, such as did you fine-tune all parameters here? 
2. Based on your findings where **models that remain closer to their pre-trained state tend to preserve more pre-training knowledge**. A naive solution based on your findings is to adopt the PEFT fine-tuning methods, which only train a small part of model parameters, such as only tuning the parameters of the attention layer of an LLM. Why can we not directly use them to mitigate the forgetting problem? 
3. In Section 2.2, can you detail your parameter partition algorithm, especially under the scope of LLMs? 
4. In the experiment setting, is it fair to directly compare LoRA with the fully fine-tuned MoFO? It looks like a more fair comparison is to set them under the same trainable parameters set, like comparing LoRA with MoFO(PEFT version), and comparing Default FT with MoFO(FT version).
soundness: 3
strengths: 1. The paper works on an intriguing and important problem, which focuses on mitigating the forgetting problem in LLM fine-tuning.
2. The paper demonstrates its performance on multiple datasets and tasks.
summary: This paper focuses on the forgetting problem of LLM fine-tuning. It proposes a new selective-updating mechanism to mitigate forgetting in the original data domain.
weaknesses: 1. There is no related work section in the main paper, which makes it difficult for readers to quickly follow up on the relevant topic. I suggest authors add this part in their revised version (not in Appendix).
2. The presentation of some sections is not very clear, such as how the authors calculate the distance of different model states in Section 2.2 and how they make the parameter partition in Section 2.3. 
3. I am a little bit confused about some parts of the experiment setting, please refer to **Questions**.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: See weakness
soundness: 2
strengths: (a) The paper presents a convergence analysis of the proposed optimizer.

(b) It includes experiments on a diverse set of LLM models and a range of advanced NLP fine-tuning tasks.
summary: The motivation of this paper is to address the challenge of forgetting in large language models (LLMs) during fine-tuning. The main contribution is the proposed Momentum-Filtered Optimizer (MoFO), a fine-tuning algorithm based on block coordinate descent that selectively updates high-momentum parameters to prevent forgetting.
weaknesses: (a) Unclear Motivation: In the Introduction (Lines 54-55), the statement “However, regularization-based methods alter the original (fine-tuning) loss function, which can potentially affect the model’s performance on fine-tuning tasks” lacks citations or experimental support, weakening the rationale for excluding regularization-based, replay-free fine-tuning methods as baselines. Providing examples or references of such methods that negatively impact fine-tuning performance would help clarify this claim. Additionally, explaining how modifications to the “original” loss function might specifically affect performance would strengthen the motivation for a regularization-free approach. The approach of keeping fine-tuned model parameters close to pre-trained models to prevent forgetting is already well-established in continual learning [1], limiting the novelty of this finding.

(b) Insufficient Baselines: The exclusion of regularization-based fine-tuning methods as baselines lacks sufficient justification. Adding comparisons with well-known regularization-based approaches, such as proximal gradient and EWC-based methods [1], would provide a more thorough evaluation. Furthermore, the current baselines mainly focus on parameter-efficient fine-tuning techniques, omitting full fine-tuning methods relevant for a comprehensive comparison in the LLM context. Including full fine-tuning baselines, as shown in recent work on LLMs [2], would enhance the fairness and depth of the comparative analysis.

(c) Lack of Clear Novelty and Contribution: The specific contributions and novelty of the proposed method in addressing forgetting in fine-tuning LLMs are difficult to discern, as the forgetting issue is not unique to LLMs and has been explored in other settings. Clarifying how this approach differs from or builds upon existing forgetting mitigation techniques, especially in the LLM context, would make its contributions more apparent.

(d) Minor Issues: The use of non-integer epochs is unclear (Line 103); specifying iterations would be more appropriate. Additionally, please clarify the batch size and learning rate used in each experiment, as these details are essential for reproducibility.

[1] Kirkpatrick, J., et al. “Overcoming catastrophic forgetting in neural networks.” Proceedings of the National Academy of Sciences 114.13 (2017): 3521-3526.

[2] Lv, K., et al. “Full parameter fine-tuning for large language models with limited resources.” arXiv preprint arXiv:2306.09782 (2023).

Meta-review: I have read all the materials of this paper including the manuscript, appendix, comments, and response. Based on collected information from all reviewers and my personal judgment, I can make the recommendation on this paper, reject. No objection from reviewers who participated in the internal discussion was raised against the reject recommendation.

**Research Question**

This paper considers the forgetting problem of LLM fine-tuning, i.e., the fine-tuned model performs worse on the original general tasks.

**Challenge Analysis**

The authors argue that two types of existing solutions need to access the original pre-training data or gradients, which are impractical in some scenarios. 

**Philosophy**

The authors aim to solely employ the well-trained LLM model to tackle the above forgetting problem. To achieve this, the authors provide some empirical analysis, demonstrating the fine-tuned model performs well in the forgetting problem if the fine-tuned model is close to the original model. Thus, the authors aim to just update some partial parameters during the fine-tuning process. Here, the logic is not smooth. Moreover, the authors fail to verify their method can lead to a close-to-original model. (Major concern I)

**Technique**

The authors employ BCD and select the largest gradients to update. I am a little confused why updating largest gradients leads to a close model.  (Major concern II)

**Theoretical Analysis**

The authors demonstrate the convergence of the proposed new optimizer. 

**Experiments**

1.	The experimental results seem promising. 
2.	The Pareto analysis is a plus.
3.	However, the parameter $\alpha$ varies on different datasets. The authors need to provide how to set $\alpha$ in practice. (Major concern III; this is the key reason for my rejection recommendation.)
4. Again, the authors need to demonstrate the fine-tuned model by the proposed method is close to the original model. Appendix E2 is a supplementary evidence, but it is not enough, where all baseline methods should be included.


**Presentation**

1.	Table 3 can be improved in terms of visualization. 
2.	Section 5 can be moved forward.

In general, this paper is very close to the bar. I would like to see an improved version in the future conference.

Decision: Reject