Excited to share our paper on continual learning & reducing forgetting at #ICML2026 tomorrow! Poster Session 3: HALL A #2010 Wed, Jul 8, 2026 • 9:30–11:15 AM CST Unfortunately, I’m not at ICML; but my students @SenmiaoWang will be happy to chat.

Forgetting is a challenge for continual learning, including sft.

1. Finding: Correlation of distance and forgetting. We found that optimizers can converge to minima with similar SFT loss with different distances from the pre-trained model. The farther the model moves, the more pre-trained knowledge it tends to lose.

2. Algorithm: MoFO, a Momentum-Filtered Optimizer for LLM fine-tuning. At each step, MoFO updates only the parameters with the largest momentum magnitudes within each block, leaving the rest unchanged. This biases toward closer solutions, preserving more pre-trained knowledge while maintaining strong task performance.

No pre-training data required. No extra model-weight storage.

Paper: https://arxiv.org/pdf/2407.20999 [TMLR, J2C Certification]

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

Source: https://arxiv.org/html/2407.20999 Selective Rule Based on Momentum Magnitudes

Motivated by the correlation between forgetting and the distance from the pre-trained model, we seek to design an optimizer that encourages the fine-tuned model to keep closer to the pre-trained model. To achieve this, we draw inspiration from the classical block coordinate descent (BCD) method(Tseng,2001), which updates only a subset of parameters during each iteration. We anticipate that restricting updates to a subset of parameters—similar to the BCD approach—will result in smaller overall deviations from the pre-trained model compared to full-parameter updates across all iterations, thereby mitigating the forgetting of pre-training knowledge.

To further accelerate convergence under the BCD framework, we adopt Gauss-Southwell rule, i.e., the greedy rule(Nutini et al.,2015). Gauss-Southwell rule selects the parameters with the largest gradients at each iteration, as those are expected to yield the greatest immediate reduction in the loss. It has also been shown inNutini et al. (2015)that BCD using the Gauss-Southwell rule—also referred to as greedy BCD—can converge faster than the traditional random BCD. However, BCD algorithms, including greedy BCD, are mostly developed based on the GD or SGD framework, but in LLM training, Adam has replaced SGD as the default optimizer(Zhang et al.,2024c). Our experiments in Section4.4show that directly following the Gauss–Southwell rule in Adam—that is, always updating parameters with large gradients—does not lead to satisfactory performance on fine-tuning tasks.

Adam inherently incorporates momentum term in parameter updates. Therefore, we propose to modify the Adam optimizer to update only the parameters with thelargest momentum magnitudes. By focusing on partial yet significant updates, our method, named MoFO, aims to effectively fine-tune models while maintaining closer to their pre-trained state. We first introduce MoFO in the next subsection. Further theoretical analysis and empirical exploration in the selection rule will be provided in Section3.2and4.4, respectively.

2.2Formulation of MoFO

Algorithm 1Momentum Filtered Optimizer (MoFO)1:Input: Filtering threshold

α\alpha, number of partitions

BBwith the

kk-th partition of size

dkd_{k}, hyperparameters

β1,β2\beta_{1},\beta_{2}of Adam optimizer, learning rate schedule

{ηt}\{\eta_{t}\}.

2:Initialize

m0,v0m_{0},v_{0}as zero tensors.

3:foriteration

ttfrom

1,2,…1,2,\dotsuntil convergedo

4:forpartition

kkfrom

11to

BBdo

5:

gt(k)=∇(k)ℒf​i​n​e​t​u​n​e​(θt−1)g^{(k)}_{t}=\nabla_{(k)}\mathcal{L}_{finetune}(\theta_{t-1}) 6:

mt(k)=β1​mt−1(k)+(1−β1)​gt(k)m^{(k)}_{t}=\beta_{1}m^{(k)}_{t-1}+(1-\beta_{1})g^{(k)}_{t} 7:

vt(k)=β2​vt−1(k)+(1−β2)​gt(k)∘gt(k)v^{(k)}_{t}=\beta_{2}v^{(k)}_{t-1}+(1-\beta_{2})g^{(k)}_{t}\circ g^{(k)}_{t} 8:

m^t(k)=mt(k)/(1−β1t)\hat{m}^{(k)}_{t}=m^{(k)}_{t}/(1-\beta_{1}^{t}) 9:

v^t(k)=vt(k)/(1−β2t)\hat{v}^{(k)}_{t}=v^{(k)}_{t}/(1-\beta_{2}^{t}) 10:forentry index

iifrom

11to

dkd_{k}do

11:

[FLTα(k)​(mt)]i=1[\texttt{FLT}^{(k)}_{\alpha}(m_{t})]_{i}=1if

|(mt(k))i||(m^{(k)}_{t})_{i}|is within the top-

α\alphaof

|mt(k)||m^{(k)}_{t}|’s valueselse0

12:endfor

13:

θt(k)=θt−1(k)−ηt⋅(m^t(k)⊙FLTα(k)​(mt))/v^t(k)\theta_{t}^{(k)}=\theta_{t-1}^{(k)}-\eta_{t}\cdot(\hat{m}^{(k)}_{t}\odot\texttt{FLT}^{(k)}_{\alpha}(m_{t}))/\sqrt{\hat{v}^{(k)}_{t}}# Momentum Filtering

14:endfor

15:

θt=Concat​(θt(1),…,θt(B))\theta_{t}=\texttt{Concat}(\theta_{t}^{(1)},\dots,\theta_{t}^{(B)}) 16:endfor

Figure 2:Illustration of MoFO.We formally introduce the Momentum-Filtered Optimizer (MoFO) in Algorithm1. First, all model parameters are partitioned intoBBblocks. At each iteration, MoFO first computes the gradient and momentum terms for parameters in each block following the standard rule of Adam, as shown in Lines 5-9. Then, MoFO selects and updates the parameter entries with the largestα\alphamomentum magnitudes in each parameter block, as shown in Lines 10-13, where the update fractionα\alphais a pre-determined hyperparameter. This momentum filtering mechanism is illustrated in Figure2.

Mathematically, the filter can be represented as follows. Consider a momentum vectorm=(m(1),…,m(B))m=(m^{(1)},\dots,m^{(B)}), where eachm(k)∈ℝdkm^{(k)}\in\mathbb{R}^{d_{k}}corresponds to thekk-th block of parameters with dimensionalitydkd_{k}. The top-α\alphafilter, denoted asFLTα​(m)\texttt{FLT}_{\alpha}(m), is defined asFLTα​(m)=(FLTα(1)​(m),…,FLTα(B)​(m))\texttt{FLT}_{\alpha}(m)=(\texttt{FLT}_{\alpha}^{(1)}(m),\dots,\texttt{FLT}_{\alpha}^{(B)}(m)), where theii-th entry ofFLTα(k)​(m)\texttt{FLT}_{\alpha}^{(k)}(m)is given by

[FLTα(k)​(m)]i={1if​|mi(k)|​is within the top-​α​of​|m(k)|​values,0otherwise,\displaystyle\left[\texttt{FLT}_{\alpha}^{(k)}(m)\right]_{i}= fori=1,2,⋯,dki=1,2,\cdots,d_{k},k=1,2,⋯,Bk=1,2,\cdots,B. In our Momentum-Filtered Optimizer (MoFO), this filterFLTα\texttt{FLT}_{\alpha}is applied to the momentummtm_{t}, selecting the entries with the largest magnitudes for updating.

For the parameter partitioning, we note that the network architecture is naturally composed of different modules (e.g., weight matrices, and bias terms). In the PyTorch implementation, the parameters of different modules (along with their gradients and momenta) are naturally stored in separate data tensors. Therefore, we adopt the default partitioning of model parameters as implemented in PyTorch. For Transformers, this means that parameters such as query (Q), key (K), value (V) weights in the attention layers, as well as feed-forward network (FFN) weights, are grouped into distinct partitions following PyTorch’s default scheme. This allows us to select and update the top-α\alphaparameters in each block without introducing much implementation overhead. See AppendixC.4for further explanation of the partitioning.

At each iteration, MoFO efficiently selects and updates the most “influential” parameters, as dictated by the momentum’s magnitude, while keeping other parameters fixed. We argue that filtering the momentum is more effective than filtering the gradient. In Section4.4, we will empirically demonstrate that MoFO’s momentum-based filtering rule outperforms other filtering rules in fine-tuning tasks.

3Theoretical Analysis

3.1Convergence Result

In this section, we present the convergence result of MoFO for non-convex loss functions. For the simplicity of analysis, we consider the full-batch version of MoFO, with hyperparameters satisfying the following assumption.

Assumption 1.

Loss functionℒ\mathcal{L}is lower bounded byℒ∗\mathcal{L}^{*}. The gradient∇ℒ\nabla\mathcal{L}is Lipschitz continuous with constantLL.

Theorem 1(Convergence of MoFO).

Suppose that the first- and second-order momentum hyperparametersβ1\beta_{1}andβ2\beta_{2}satisfy0<β1<β2<10<\beta_{1}<\sqrt{\beta_{2}}<1. The learning rate schedule at stepttisηt=η/t\eta_{t}=\eta/\sqrt{t}for someη>0\eta>0. Then, under Assumption1, MoFO satisfies

min0≤t≤T−1⁡‖∇ℒ​(θt)⊙FLTα​(∇ℒ​(θt))‖1=𝒪​(log⁡TT)as​T→∞.\min_{0\leq t\leq T-1}\|\nabla\mathcal{L}(\theta_{t})\odot\texttt{FLT}_{\alpha}\big(\nabla\mathcal{L}(\theta_{t})\big)\|_{1}=\mathcal{O}\left(\frac{\log T}{\sqrt{T}}\right)\quad\text{as }T\to\infty.Moreover, this bound directly implies

min0≤t≤T−1⁡‖∇ℒ​(θt)‖p=𝒪​(log⁡TT)as​T→∞,\min_{0\leq t\leq T-1}\|\nabla\mathcal{L}(\theta_{t})\|_{p}=\mathcal{O}\left(\frac{\log T}{\sqrt{T}}\right)\quad\text{as }T\to\infty,for anyp∈[1,∞]p\in[1,\infty].

Although MoFO is designed to mitigate forgetting by updating only a small subset of parameters at each step, it is guaranteed to converge to a critical point of the fine-tuning loss function under the Lipschitz smoothness assumption. This result provides theoretical evidence that MoFO can achieve competitive performance in fine-tuning tasks.

Proof Sketch of Theorem1:.

Our proof is inspired by the convergence analysis for full-batch Adam inShi et al. (2021), but we introduce additional techniques tailored to MoFO’s filtering mechanism. We will highlight these additions precisely at the points where they arise below.

Letgt=∇ℒ​(θt−1)g_{t}=\nabla\mathcal{L}(\theta_{t-1}). A central step is to establish, for suitable constantsC1,C2>0C_{1},C_{2}>0,

C1t​‖gt‖≤ℒ​(θt−1)−ℒ​(θt)+C2t,\frac{C_{1}}{\sqrt{t}}\|g_{t}\|\leq\mathcal{L}(\theta_{t-1})-\mathcal{L}(\theta_{t})+\frac{C_{2}}{t},(1)Summing this inequality fromt=1t=1toTTand using∑t=1Tt−1=log⁡T+𝒪​(1)\sum_{t=1}^{T}t^{-1}=\log T+\mathcal{O}(1)yields the convergence result for Adam in terms of a diminishing norm of gradient, given by

min1≤t≤T⁡‖gt‖=𝒪​(log⁡TT).\min_{1\leq t\leq T}\|g_{t}\|=\mathcal{O}\left(\frac{\log T}{\sqrt{T}}\right).(2)

Choice of norm and two subgoals.

In finite-dimensional spaces, all norms are equivalent, so convergence statements like (1)–(2) can be expressed in any fixed norm up to norm-equivalence constants. In practice, however, specific analyses instantiate (1) with a particular norm:Shi et al. (2021)work with theL1L_{1}norm‖gt‖1\|g_{t}\|_{1}for full-batch Adam, whereas our full-batch MoFO analysis will use theL1,top-​αL_{1,\text{top-}\alpha}norm‖gt⊙FLTα​(gt)‖1\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\|_{1}, which will be defined in AppendixA.1.

To keep the logic precise, we separate our argument into two subgoals:

Step I (Key inequality for MoFO).Show that there exist constantsC1,C2>0C_{1},C_{2}>0(independent oftt) such that

C1t​‖gt⊙FLTα​(gt)‖1≤ℒ​(θt−1)−ℒ​(θt)+C2t.\frac{C_{1}}{\sqrt{t}}\,\big\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\big\|_{1}\;\leq\;\mathcal{L}(\theta_{t-1})-\mathcal{L}(\theta_{t})+\frac{C_{2}}{t}.(3)Here, we recall thatFLTα​(⋅)\texttt{FLT}_{\alpha}(\cdot)preserves theα\alpha-fraction of largest-magnitude coordinates in each partition and zeros out the rest.

Step II (Norm property of the left-hand side of (3)).*This component is one of our new technical ingredients (absent fromShi et al. (2021)).*We first verify that the mappingx↦‖x⊙FLTα​(x)‖1x\mapsto\|x\odot\texttt{FLT}_{\alpha}(x)\|_{1}defines a norm onℝd\mathbb{R}^{d}, which is referred to asL1,top-​αL_{1,\text{top-}\alpha}norm. This is proved in Proposition1(AppendixA.1) by checking nonnegativity and definiteness, positive homogeneity, and the triangle inequality. Moreover, for theLpL_{p}upper bound in Theorem1, we use the norm equivalence between theL1,top-​αL_{1,\text{top-}\alpha}norm and theLpL_{p}norm, as shown in Lemma2.

With Step II established, it remains to prove the key inequality in Step I.A direct adaptation ofShi et al. (2021)to MoFO is infeasible due to structural differences. We proceed as follows:

1. (i)Recap key elements fromShi et al. (2021);
2. (ii)Identify challenges in extending them to MoFO;
3. (iii)Resolve these challenges by carefully handling the momentum filter.

*Part (i): Key elements ofShi et al. (2021).*For full-batch Adam with bias-corrected momentsm^t,v^t\hat{m}_{t},\hat{v}_{t}and the learning rate scheduleηt=η/t\eta_{t}=\eta/\sqrt{t}, the parameter update is

θt−θt−1=−ηt⋅m^tv^t.\theta_{t}-\theta_{t-1}=-\frac{\eta}{\sqrt{t}}\cdot\frac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}}}. By theLL-smoothness of the loss and the descent lemma,

ηt​∑i=1dgi,t​m^i,tv^i,t≤ℒ​(θt−1)−ℒ​(θt)+L2​‖θt−θt−1‖22.\frac{\eta}{\sqrt{t}}\sum_{i=1}^{d}g_{i,t}\frac{\hat{m}_{i,t}}{\sqrt{\hat{v}_{i,t}}}\leq\mathcal{L}(\theta_{t-1})-\mathcal{L}(\theta_{t})+\frac{L}{2}\|\theta_{t}-\theta_{t-1}\|_{2}^{2}.(4) Lemma5in AppendixA.2lower-bounds the per-coordinate contribution as

gi,t​m^i,tv^i,t≥A​|gi,t|−Bt,g_{i,t}\frac{\hat{m}_{i,t}}{\sqrt{\hat{v}_{i,t}}}\geq A|g_{i,t}|-\frac{B}{\sqrt{t}},for some constantsA,B>0A,B>0. Substituting this into (4), summing over coordinates, and controlling the quadratic term in (4) yields the fundamental inequality (1) with theL1L_{1}-norm:

C1t​‖gt‖1≤ℒ​(θt−1)−ℒ​(θt)+C2t,\frac{C_{1}}{\sqrt{t}}\|g_{t}\|_{1}\leq\mathcal{L}(\theta_{t-1})-\mathcal{L}(\theta_{t})+\frac{C_{2}}{t}, *Part (ii): Challenges in extending to MoFO.*For the full-batch version of MoFO, the parameter update becomes

θt−θt−1=−ηt​m^tv^t⊙FLTα​(mt),\theta_{t}-\theta_{t-1}=-\frac{\eta}{\sqrt{t}}\frac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}}}\odot\texttt{FLT}_{\alpha}(m_{t}),i.e., the step is filtered by momentum magnitudes (NOT gradient magnitudes). Building upon the convergence analysis ofShi et al. (2021)in Part (i) leads to

C1t​‖gt⊙FLTα​(mt)‖1≤ℒ​(θt−1)−ℒ​(θt)+C2t,\frac{C_{1}}{\sqrt{t}}\|g_{t}\odot\texttt{FLT}_{\alpha}(m_{t})\|_{1}\leq\mathcal{L}(\theta_{t-1})-\mathcal{L}(\theta_{t})+\frac{C_{2}}{t},(5)which brings a notable difference to the target key inequality (3): the inequality (5) applies the momentum filterFLTα​(mt)\texttt{FLT}_{\alpha}(m_{t}), whereas the desired bound applies the gradient filterFLTα​(gt)\texttt{FLT}_{\alpha}(g_{t}). This introduces a non-trivial challenge because the‖gt⊙FLTα​(mt)‖1\|g_{t}\odot\texttt{FLT}_{\alpha}(m_{t})\|_{1}being small does not naturally imply that‖gt⊙FLTα​(gt)‖1\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\|_{1}is small: large entries ofgtg_{t}might be excluded if their momenta lags. Consequently, additional analysis is required to bound the discrepancy between them.

Part (iii): How we overcome the challenge.**Here we introduce another new technique to address the discrepancy between‖gt⊙FLTα​(mt)‖1\|g_{t}\odot\texttt{FLT}_{\alpha}(m_{t})\|_{1}and‖gt⊙FLTα​(gt)‖1\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\|_{1}.We deal with the challenge via Lemma1in AppendixA.1, which establishes that theL1L_{1}-deviation in filtered outputs isLipschitz-stableunder input perturbations. Specifically, for anyx,y∈ℝdx,y\in\mathbb{R}^{d}:

‖x⊙FLTα​(x)‖1−‖x⊙FLTα​(y)‖1⏟top-​α​filtered error≤2​‖x−y‖1⏟input error.\underbrace{\left\|x\odot\texttt{FLT}_{\alpha}(x)\right\|_{1}-\left\|x\odot\texttt{FLT}_{\alpha}(y)\right\|_{1}}_{\text{top-}\alpha\text{ filtered error}}\leq 2\underbrace{\|x-y\|_{1}}_{\text{input error}}.This result points out that while the filterFLTα​(⋅)\texttt{FLT}_{\alpha}(\cdot)itself is unstable under perturbation, the filtered output remains controllable through. It effectively “smooths” the discontinuity that would plausibly prevent convergence analysis.

With the Lipschitz stability established, we now deal with the challenge. We first control theL1L_{1}distance between the bias-corrected momentumm^t=mt1−β1t\hat{m}_{t}=\frac{m_{t}}{1-\beta_{1}^{t}}and the gradientgtg_{t}by showing that‖m^t−gt‖1=𝒪​(1/t)\|\hat{m}_{t}-g_{t}\|_{1}=\mathcal{O}(1/\sqrt{t}). Second, we apply Lemma1withx=gtx=g_{t}andy=m^ty=\hat{m}_{t}and yield

‖gt⊙FLTα​(gt)‖1−‖gt⊙FLTα​(m^t)‖1≤2​‖m^t−gt‖1=𝒪​(1/t).\left\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\right\|_{1}-\left\|g_{t}\odot\texttt{FLT}_{\alpha}(\hat{m}_{t})\right\|_{1}\leq 2\|\hat{m}_{t}-g_{t}\|_{1}=\mathcal{O}(1/\sqrt{t}).Since the filtering functionFLTα​(⋅)\texttt{FLT}_{\alpha}(\cdot)is invariant under positive scaling, we can definitely replacem^t\hat{m}_{t}with the original momentummtm_{t}:

‖gt⊙FLTα​(gt)‖1−‖gt⊙FLTα​(mt)‖1≤2​‖m^t−gt‖1=𝒪​(1/t).\left\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\right\|_{1}-\left\|g_{t}\odot\texttt{FLT}_{\alpha}(m_{t})\right\|_{1}\leq 2\|\hat{m}_{t}-g_{t}\|_{1}=\mathcal{O}(1/\sqrt{t}). Combining it with (5) and subsuming residual𝒪​(1/t)\mathcal{O}(1/t)terms, we obtain our target key inequality (3):

C1t​‖gt⊙FLTα​(gt)‖1≤ℒ​(θt−1)−ℒ​(θt)+C2t.\frac{C_{1}}{\sqrt{t}}\left\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\right\|_{1}\leq\mathcal{L}(\theta_{t-1})-\mathcal{L}(\theta_{t})+\frac{C_{2}}{t}. In conclusion, our statement ofStep Iprimarily comprises:

• •Step I.1 (Part (i) and (ii)).Extending the Adam convergence framework(Shi et al.,2021)to incorporate MoFO’s momentum filtering mechanism, yielding (5);
• •Step I.2 (Part (iii)).Resolving the momentum-gradient filter discrepancy via our Lipschitz stability analysis (Lemma1), which is a new technique introduced in this work.

Combined with the norm property inStep II—another new ingredient not present inShi et al. (2021)—, these results enable us to establish Theorem1.

∎

We remark that the choice ofβ1\beta_{1}andβ2\beta_{2}in Assumption1aligns with that used in analyzing full-batch Adam(Shi et al.,2021). Furthermore, the use of a diminishing learning rate in Theorem1is crucial for ensuring the stability of updates and avoiding divergence in the optimization process.

In summary, Theorem1demonstrates that despite updating only a subset of parameters, MoFO maintains the same convergence rate as Adam. This highlights the theoretical robustness of the momentum filter design in MoFO. We believe this result could provide valuable insights into adaptive optimization methods with filtering mechanisms.

3.2Initial Analysis on Forgetting Mitigation

Does MoFO converge to a model that is closer to the pre-trained LLM than Adam, thereby reducing forgetting? In this subsection, we attempt to address this question by providing an initial theoretical analysis on an illustrative example.

Example 1.

Suppose the parameter space isℝd\mathbb{R}^{d}, and the updating ratio isα=1/d\alpha=1/d, i.e., only one coordinate is updated in each iteration. We assume thepre-training lossisℒpretrain​(θ)=12​‖θ‖22\mathcal{L}_{\rm pretrain}(\theta)=\frac{1}{2}\|\theta\|_{2}^{2}and that the model has been trained to the global minimumθpretrain=(0,0,…,0)\theta_{\rm pretrain}=(0,0,\dots,0)during the pre-training phase. Thefine-tuning lossis given byℒ​(θ)=∏i=1d(ai​θi−bi)2\mathcal{L}(\theta)=\prod_{i=1}^{d}(a_{i}\theta_{i}-b_{i})^{2}, whereai,bi>0a_{i},b_{i}>0for any1≤i≤d1\leq i\leq d. In this example, the set of global minima ofℒ​(θ)\mathcal{L}(\theta)is a union of hyperplanes:

S=⋃i=1dSi,where​Si:={θ∈ℝd:θi=bi/ai}.\displaystyle S=\bigcup_{i=1}^{d}S_{i},\quad\text{where }S_{i}=\{\theta\in\mathbb{R}^{d}:\theta_{i}=b_{i}/a_{i}\}.Here, eachSiS_{i}represents a hyperplane inℝd\mathbb{R}^{d}.

Remark 1.

We note that the loss landscapes of neural networks are generally non-convex(Liu et al.,2022). Here, we also adopt a non-convex fine-tuning lossℒ​(θ)\mathcal{L}(\theta). Further, we note that the set of global-minimaSSconsists of infinitely many minima spread across multiple hyperplanes, aligning with the observation on the degenerate structure of minima in neural networks(Lin et al.,2024b).

Theorem 2.

In Example1, if the learning rates are chosen appropriately, then MoFO converges to a minimumθMoFO∗\theta^{*}_{\rm MoFO}that is closer to the pre-training state than the minimumθAdam∗\theta^{*}_{\rm Adam}obtained by Adam, i.e.,

‖θMoFO∗−θpretrain‖2<‖θAdam∗−θpretrain‖2.\|\theta^{*}_{\rm MoFO}-\theta_{\rm pretrain}\|_{2}<\|\theta^{*}_{\rm Adam}-\theta_{\rm pretrain}\|_{2}.Moreover, MoFO attains a strictly lower pre-training loss, i.e.,

ℒpretrain​(θMoFO∗)<ℒpretrain​(θAdam∗).\mathcal{L}_{\rm pretrain}(\theta^{*}_{\rm MoFO})<\mathcal{L}_{\rm pretrain}(\theta^{*}_{\rm Adam}).

Next, we visualize the landscape of this example to provide additional intuition. Consider a simplified two-dimensional case. Specifically, letθ∈ℝ2\theta\in\mathbb{R}^{2}andℒ​(θ)=(θ1−1)2​(θ2−1)2\mathcal{L}(\theta)=(\theta_{1}-1)^{2}(\theta_{2}-1)^{2}.

Figure 3:The fine-tuning loss landscape and the training paths of different optimization methods. The color bar indicates the fine-tuning loss value.A logarithmic scale is applied to the loss values for better visualization. MoFO converges to a minimum closest to the pre-trained model.

In Example1, each hyperplaneSiS_{i}, which forms part of the global minima ofℒ\mathcal{L}, can be viewed as an attractor for the optimization algorithms. These attractors (SiS_{i}’s) influence the model’s update direction during training. As illustrated in Figure3, the attractors in this case are two straight lines:θ1=1\theta_{1}=1andθ2=1\theta_{2}=1. When using Adam, the model is simultaneously pulled by both attractors, causing it to move diagonally along the orange line and converge at(1,1)(1,1), with a resulting pre-training loss of 1.

In contrast, MoFO is influenced by only one attractor (θ1=1\theta_{1}=1), leading it to follow the green line and converge to(1,0)(1,0), which achieves a lower pre-training loss of 0.5444By symmetry, MoFO may also be influenced by another attractorθ2=1\theta_{2}=1and converges vertically to(0,1)(0,1).. We hypothesize that, for full-parameter fine-tuning (Adam), interference among multiple attractors drives convergence toward a “balanced” solution, whereas MoFO mitigates such interference by selectively updating only parameters with large momentum magnitudes. Such an updating rule helps MoFO move toward a single attractor, thereby converging closer and forgetting less. In addition to Adam and MoFO, we provide convergence paths of some other baseline methods in AppendixB.4.

We believe that the above analysis provides an initial insight into the effectiveness of MoFO. More in-depth analysis of MoFO in forgetting mitigation is left for future work.

4Experiments

4.1Experimental Settings

We verify the effectiveness of MoFO oninstruction fine-tuningandcontinual fine-tuning. We use Llama-2-7B(Touvron et al.,2023), Gemma-2B-IT(Team et al.,2024), and TinyLlama-1.1B(Zhang et al.,2024b)as our base models. The instruction fine-tuning datasets cover question-answer pairs from different domains like mathematical reasoning and medical knowledge. Specifically, the datasets include: MetaMathQA(Yu et al.,2024b), PMC-LLaMA-Instructions(Wu et al.,2024), Magicoder-Evol-Instruct(Wei et al.,2023). We randomly sample 39.5K and 51K instances from these datasets, respectively, for training the LLMs. Additionally, We investigate the performance of MoFO in the continual fine-tuning scenario by implementing our approach on the TRACE benchmark dataset(Wang et al.,2023b).

Evaluation metrics for instruction fine-tuning.We employ widely used benchmarks to assess the performance and potential forgetting effects on the general capabilities of LLMs after instruction fine-tuning. These benchmarks include MMLU(Hendrycks et al.,2021)(0-shot) for factual knowledge; ARC-Challenge, ARC-Easy(Clark et al.,2018), and HellaSwag(Zellers et al.,2019)(0-shot) for commonsense reasoning (CR); GSM8K(Cobbe et al.,2021)(5-shot) for mathematical reasoning; HumanEval (HEval)(Chen et al.,2021)(pass@10) for code generation; PubMedQA(Jin et al.,2019), MedMCQA(Pal et al.,2022), and MedQA(Jin et al.,2021)(0-shot) for medical question answering (MedQ)555For CR and MedQ, we report the average of the benchmarks they comprise.; IFEval (0-shot) for instruction following.

Evaluation 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.,2018)and BackWard Transfer (BWT)(Lopez-Paz & Ranzato,2017).

For more descriptions and implementation details of these metrics and datasets, see AppendixC.

4.2Instruction Fine-Tuning

In this section, we investigate the effectiveness of the MoFO algorithm in both preserving general capabilities and learning fine-tuning tasks. The implementation details are provided in AppendixC. The specific hyperparameter settings in each experiment are provided in AppendixC.3.

LLM Fine-tuning strategy baselines.We compare MoFO with the default fine-tuning approach and other methods designed to mitigate forgetting. These baselines include:Default fine-tuning (Default FT)refers to the full-parameter fine-tuning approach using the Adam optimizer.Half Fine-tuning (HFT)(Hui et al.,2024)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.LoRA(Hu et al.,2022)is a widely-used, parameter-efficient fine-tuning method. LoRA trains low-rank matrix adaptations on the base model’s weights. Recent work(Biderman et al.,2024)demonstrates that LoRA can mitigate forgetting.

Table 1:The performance of the fine-tuning task (math), measured by GSM8K, and the general capability scores of Llama-2-7B after fine-tuning on the MetaMathQA dataset. The results show that MoFO achieves comparable performance in the fine-tuning task, while significantly mitigating forgetting of general capabilities. Bold values denote the best results among these methods.MethodGSM8KGeneral CapabilityCRMMLUHEvalAvg.Llama-2-7B13.765.642.024.243.9Default FT49.462.336.616.138.3HFT47.565.542.323.643.8LoRA43.365.137.726.443.1MoFO47.765.742.724.644.3Results 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 Table1. We report the experimental results of LoRA under the best-performing hyperparameter configuration on the fine-tuning task. These results demonstrate the effectiveness of our proposed MoFO algorithm in both optimization and mitigating forgetting.

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

Comparison from a Pareto perspective.Generally, improving performance on the fine-tuning task and reducing forgetting are often a pair of competing objectives. It is intriguing to study how different fine-tuning methods balance this tradeoff. By adjusting the hyperparameters of different methods, we can observe a set of fine-tuned models, each representing a different tradeoff between fine-tuning performance and forgetting. The Pareto frontier formed by these models helps visualize the tradeoffs, and we can identify which method offers the best balance between fine-tuning and forgetting.

In this comparison, we also include traditional regularization methods such asL2L_{2}-regularization(Li et al.,2018)(denoted asL2L_{2}reg) andL1L_{1}-regularization(Panigrahi et al.,2023)(denoted asL1L_{1}reg), which are not specifically designed for large models. These methods modify the original fine-tuning lossℒf​i​n​e​t​u​n​e​(θ)\mathcal{L}_{finetune}(\theta)by adding a regularization term. ForL2L_{2}-regularization, the modified loss isℒf​i​n​e​t​u​n​e​(θ)+λ2​‖θ−θ0‖22\mathcal{L}_{finetune}(\theta)+\lambda_{2}\|\theta-\theta_{0}\|^{2}_{2}, and forL1L_{1}-regularization, it isℒf​i​n​e​t​u​n​e​(θ)+λ1​‖θ−θ0‖1\mathcal{L}_{finetune}(\theta)+\lambda_{1}\|\theta-\theta_{0}\|_{1}, whereλ2\lambda_{2}andλ1\lambda_{1}are the respective regularization hyperparameters.

We fine-tune the Llama-2-7B model on the MetaMathQA dataset usingL1L_{1}andL2L_{2}regularization, as well as LoRA, and compare their performance with MoFO. We present the results in Figure5(a)and plot Pareto optimal fronts666Since it is impractical to exhaust all hyperparameter configurations in real experiments, we present linear interpolation approximations of the Pareto fronts in Figure5(a).for these methods. Details of the hyperparameter configurations for this experiment are provided in AppendixC.3. These results show the effectiveness of the MoFO algorithm in both optimization and mitigating forgetting.

(a)The performance on the math task (GSM8K) and the scores in general capabilities of Llama-2-7B after fine-tuning on the MetaMathQA dataset. Only points on the Pareto front are shown as solid points, while the remaining points are presented as semi-transparent. The results show that compared withL1L_{1},L2L_{2}regularization, and LoRA across various hyperparameter configurations, the MoFO algorithm achieves a better Pareto front.The result reveals that MoFO consistently achieves a better Pareto front in comparison to baseline methods. When compared to regularization methods and LoRA, MoFO exhibits less forgetting and can even maintain general capabilities with comparable GSM8K accuracies. Additionally, MoFO outperforms regularization methods in math tasks when the magnitudes of forgetting are similar. We also note thatL1L_{1}andL2L_{2}regularization(Panigrahi et al.,2023; Li et al.,2018)require storing the pre-training weights throughout the entire fine-tuning process for regularization computation, which incurs additional memory overhead. For preliminary analysis on why MoFO might compare favorably toL1/L2L_{1}/L_{2}regularization, see AppendixF.5.

AppendixEreports additional instruction fine-tuning results. Specifically, we evaluate

• •LLM variants: Gemma-2B-IT and Llama-2-7B-Chat;
• •Domain-specific datasets: medical dataset (PMC-LLaMA-Instruct(Wu et al.,2024)), coding dataset (Magicoder-Evol-Instruct(Wei et al.,2023));
• •Different baselines: HMA(Lin et al.,2024a), CoFiTune(Zhang et al.,2024a), Soft-masking(Ke et al.,2023a;b).

MoFO Converges Closer to the Pre-trained Model.

In this part, we empirically investigate whether MoFO converges closer to the pre-trained model. Building on the fine-tuned models in Table1, we compare their distances to the pre-trained model. In addition, for theL1L_{1}andL2L_{2}regularization baselines in our Pareto analysis above, we select the models that achieve the best performance on the GSM8K benchmark (corresponding to the fine-tuning task).

Figure5(b)shows that models fine-tuned with MoFO are closer to the pre-trained model compared to other baseline methods.

(b)The distances for the fine-tuned Llama2-7B on MetaMathQA. The results show MoFO achieves minima closer to the pre-trained model.### 4.3Continual Fine-Tuning

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

Continual learning baselines.We consider several traditional methods from the field of continual learning to compare with MoFO. These methods can also be orthogonal combined with MoFO to further enhance performance.Replayinvolves optimizing the model using current data along with a memory buffer containing samples from previous tasks to mitigate forgetting, and we follow the implementation inWang et al. (2023b).Gradient of Episodic Memory (GEM)(Lopez-Paz & Ranzato,2017)mitigates forgetting by using gradients from old tasks to adjust the parameter updates during the training of new tasks.Elastic weight consolidation (EWC)(Kirkpatrick et al.,2017)uses a diagonal approximation of Fisher information matrix, which can be calculated by gradients from previous tasks, to regularize parameter updates.

Table 2:The OP and BWT scores of TinyLlama-1.1B after fine-tuning on TRACE benchmark. The results show that MoFO outperforms Default FT, HFT, GEM, and EWC in continual learning and can combine well with continual learning methods. Bold values denote the best results among these methods in each group.OPBWTDefault FT38.4-10.3HFT39.9-10.1MoFO41.3-5.4GEM40.8-8.5GEM + MoFO41.7-6.7EWC41.1-8.3EWC + MoFO43.2-4.4Replay45.54.7Replay + MoFO47.04.8Results of continual fine-tuning.We present the experimental results of sequentially fine-tuning TinyLlama-1.1B on the TRACE benchmark with various methods in Table2. The results indicate that in continual fine-tuning, MoFO not only outperforms other fine-tuning baselines but also surpasses GEM and EWC. Moreover, MoFO combines well with the Replay method, offering a1.5%1.5\%performance gain on the OP metric compared to using Replay alone. Moreover, MoFO works well in combination with EWC, yielding at least a 2.1% improvement in the OP metric over using EWC alone. Additionally, when combined with the GEM method, MoFO provides a0.9%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 forgetting.

4.4Impact of Update Strategy in MoFO

Table 3:The performance on the math reasoning task (GSM8K) and general capability scores of Llama-2-7B after fine-tuning on MetaMathQA using different updating strategies in MoFO. Bold values denote the best results among the BCD methods.MethodGSM8KGeneral CapabilityCRMMLUHEvalAvg.Llama-2-7B13.765.642.024.243.9Default FT49.462.336.616.138.3Random BCD35.065.841.125.144.0Grad BCD40.266.041.628.045.2MV BCD42.266.040.027.644.5MoFO45.465.743.527.445.5In addition to MoFO, we consider three other BCD methods with different filtering strategies:random BCD,gradient-filtered BCD, andMV-filtered BCD.Random BCDupdates a random subset of parameters at each iteration.Gradient-filtered BCDreplaces MoFO’s filterFLTα​(mt)\texttt{FLT}_{\alpha}(m_{t})withFLTα​(gt)\texttt{FLT}_{\alpha}(g_{t}), whileMV-filtered BCDusesFLTα​(mt/vt)\texttt{FLT}_{\alpha}(m_{t}/\sqrt{v_{t}}).

We fine-tune Llama-2-7B on MetaMathQA using these four methods with10%10\%parameter update fraction and present the results in Table3. Experimental results show that all four BCD methods exhibit significantly less forgetting compared to Default FT, demonstrating the effectiveness of BCD algorithms in mitigating forgetting.

In terms of GSM8K performance, our proposed MoFO method significantly surpasses Random BCD, Gradient-filtered BCD, and MV-filtered BCD, indicating that updating parameters with the largest momentum leads to strong optimization power. Additional comparative experiments on BCD filtering strategies are presented in AppendixF.3. More insights towards this result are provided in AppendixF.4.

4.5Furthur Analysis

Guidelines for settingα\alpha.Experiments show that setting the updating fractionα=15%\alpha=15\%works the best for most of our experiments; and5%−15%5\%-15\%all work quite well. We provide a more detailed guideline for determiningα\alphain AppendixD. The guideline involves randomly sampling a small proxy subset and performing a grid search over possibleα\alphavalues.

Efficiency Analysis.We provide an efficiency analysis on MoFO in AppendixF.2. The results show that MoFO requires only around4%−5%4\%-5\%additional training time compared with Default FT throughout the entire training process.

5Related Works

5.1Forgetting in Continual Learning

Catastrophic forgetting, a significant issue where models forget previously learned information upon learning new data, has received considerable attention in machine learning(McCloskey & Cohen,1989; Goodfellow et al.,2013; Kemker et al.,2018; Ramasesh et al.,2021; Verwimp et al.,2023; Liu et al.,2024). Traditional continual learning primarily focuses on addressing catastrophic forgetting insequential-task learningscenarios. In addition to investigating the forgetting of pre-training knowledge during fine-tuning, Section4.3conducts experimental studies on catastrophic forgetting in sequential-task fine-tuning processes, which aligns more closely with conventional continual learning paradigms.

Replay-based methods. In sequential-task learning, these methods leverage past experiences to facilitate the learning of new tasks. The most classical scheme is experience replay, which involves replaying data of past tasks during incremental training(Rolnick et al.,2019)(Aljundi et al.,2019a; Hayes et al.,2019; Cha et al.,2021; Chaudhry et al.,2019b; Riemer et al.,2019b). Other variants utilize gradient information from old tasks(Lopez-Paz & Ranzato,2017; Riemer et al.,2019a; Chaudhry et al.,2019a; Farajtabar et al.,2020; Aljundi et al.,2019b; Chaudhry et al.,2021; Tiwari et al.,2022). In LLMs,Yin et al. (2023); Wang et al. (2024b); Ouyang et al. (2022)propose replay-based methods to mitigate forgetting. While MoFO is a replay-free method, MoFO can be combined with replay strategies.

Regularization-based methods. These methods introduce constraints to the training process to preserve past knowledge, such as adding regularization to the loss functions(Kirkpatrick et al.,2017; Aljundi et al.,2018; Zenke et al.,2017; Li et al.,2018; Ritter et al.,2018; Kumar et al.,2023)or the embedding/output changes(Li & Hoiem,2017; Rannen et al.,2017; Buzzega et al.,2020; Huang et al.,2021; Cha et al.,2020). Some regularization-based approaches still rely on partial information from previous models(Kirkpatrick et al.,2017). In contrast, MoFO does not require past information and does not alter the original loss function, making it inherently orthogonal to regularization-based methods. To improve generalization and robustness to noise after instruction-tuning, several studies introduce explicit regularization(Li & Zhang,2021; Zhang et al.,2023). In particular,Zhang et al. (2023)proposes a Hessian-based penalty that encourages convergence to flatter minima. It is an interesting direction for future research to evaluate whether MoFO’s momentum filtering mechanism implicitly favors more stable minima, thereby further enhancing generalization and noise robustness.

Optimization-based methods. These methods focus on modifying the training algorithm to mitigate forgetting. In traditional continual learning, optimization-based methods commonly include, but are not limited to, gradient projection techniques(Wang et al.,2023a; Lopez-Paz & Ranzato,2017), meta-learning approaches(Beaulieu et al.,2020; Javed & White,2019), and strategies leveraging the structure of the loss landscapes(Mirzadeh et al.,2020a;b). When it comes to forgetting-mitigation in LLM training, recent studies have explored optimization strategies that update only a subset of parameters at each iteration. For instance,Hui et al. (2024)randomly freezes half of the model’s parameter modules and updatets the rest at each iteration.Ke et al. (2023b;a)introduce a soft-masking mechanism that selects parameters for update based on their importance values. Further,Zhang et al. (2024a)combines selective module updating with soft-masking. MoFO, which also falls into this category, updates parameters with largest momentum magnitudes at each iteration. Compared to these works, our study provides a theoretical convergence guarantee of our proposed method, thereby establishing its effectiveness in LLM fine-tuning.

Model merging methods. These methods balance learning new knowledge and retaining old knowledge by merging the new and past models. One line of research focuses on model averaging, which interpolates between the weights of different LLMs(Wortsman et al.,2022a;b; Eeckt et al.,2022; Yadav et al.,2024; Lin et al.,2023;2024a). Another line of research relies on the observation that task-specific knowledge largely resides in a subspace of the weight space(Ilharco et al.,2023; Panigrahi et al.,2023; Gueta et al.,2023; Zhu et al.,2024; He et al.,2024), and leverage task vectors or task localization to preserve pre-training knowledge in the fine-tuned models(Panigrahi et al.,2023; Yadav et al.,2024; Yu et al.,2024a).

Architecture-based methods. These methods modify the model’s architecture in training. LoRA(Hu et al.,2022), as the most popular parameter-efficient fine-tuning (PEFT) method, freezes the pre-training weights and introduces low-rank trainable matrices. Variants of LoRA are applied in continual learning for LLMs(Ren et al.,2024; Wang et al.,2023a). However, LoRA is observed to forget less but also learn less than default fine-tuning(Biderman et al.,2024). Apart from LoRA, Adapters(Houlsby et al.,2019)and BitFit(Zaken et al.,2021)are also well-known PEFT methods. In AppendixE.4, we include them as baselines for comparison and find that, while they are slightly less effective than MoFO in mitigating forgetting, their performance on the fine-tuning task is substantially worse than that of MoFO.

Other approaches adaptively expand model capacity or isolate partial weights to mitigate interference between new and old tasks(Wang et al.,2023a; Razdaibiedina et al.,2023). In contrast, MoFO updates a subset of parameters at each iteration, but does not alter the total trainable parameters.

5.2Block Coordinate Descent

Block Coordinate Descent (BCD) involves iteratively optimizing over a block of coordinates while holding the others constant. The foundational work ofTseng (2001)provides a comprehensive analysis of the convergence properties of BCD under certain conditions. Subsequent research has explored various BCD variants(Hong et al.,2017), including random BCD(Nesterov,2012; Richtárik & Takáč,2014; Lu & Xiao,2015), cyclic BCD(Sun & Hong,2015; Razaviyayn et al.,2013), and greedy BCD(Nutini et al.,2015). 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.,2022). For example,Luo et al. (2024)leverages BCD to perform memory-efficient fine-tuning of LLM andXu & Zhang (2024)uses random masking to perform this. In federated learning,Rothchild et al. (2020)adopts top-kkmomentum value unsketch rather than our top-kkmomentum filtering to tackle communication bottleneck and convergence issues. In LLMs, some concurrent works propose BCD-based algorithms leveraging task vectors to enhance fine-tuning performance(Li et al.,2024)and mitigate catastrophic forgetting in multi-task learning(Panda et al.,2024). Our approach can be regarded as a type of greedy BCD adapted to Adam, achieving good performance in fine-tuning tasks and alleviating forgetting.

6Conclusion and Limitations

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 achieves comparable performance to default fine-tuning but also effectively alleviates forgetting.

While this work provides a preliminary exploration of applying traditional block coordinate descent methods to mitigate forgetting in LLM training, several avenues remain open for further investigation. First, the current framework uses a uniform and consistent update fraction across all parameter blocks throughout training, whereas future work may explore adaptive update fractions and block-wise dynamic adjustments. Second, although our focus centers on forgetting mitigation during supervised fine-tuning, extending this methodology to downstream phases such as RLHF represents a promising direction for improving LLM development pipelines.

Acknowledgement

The authors would like to express our sincere gratitude to the reviewers for their insightful feedback during the discussion phase. The authors also thank Congliang Chen and Ziniu Li for their helpful suggestions. This paper is supported by NSFC (No. 12326608 and No. 12401409); Hetao Shenzhen-Hong Kong Science and Technology Innovation Cooperation Zone Project (No.HZQSWS-KCCYB-2024016); University Development Fund UDF01001491, the Chinese University of Hong Kong, Shenzhen; Guangdong Provincial Key Laboratory of Mathematical Foundations for Artificial Intelligence (2023B1212010001).

References

• Aljundi et al. (2018)Rahaf Aljundi, Francesca Babiloni, Mohamed Elhoseiny, Marcus Rohrbach, and Tinne Tuytelaars.Memory aware synapses: Learning what (not) to forget.InProceedings of the European conference on computer vision (ECCV), pp. 139–154, 2018.
• Aljundi et al. (2019a)Rahaf Aljundi, Eugene Belilovsky, Tinne Tuytelaars, Laurent Charlin, Massimo Caccia, Min Lin, and Lucas Page-Caccia.Online continual learning with maximal interfered retrieval.Advances in neural information processing systems, 32, 2019a.
• Aljundi et al. (2019b)Rahaf Aljundi, Min Lin, Baptiste Goujaud, and Yoshua Bengio.Gradient based sample selection for online continual learning.Advances in neural information processing systems, 32, 2019b.
• Beaulieu et al. (2020)Shawn Beaulieu, Lapo Frati, Thomas Miconi, Joel Lehman, Kenneth O Stanley, Jeff Clune, and Nick Cheney.Learning to continually learn.InECAI 2020, pp. 992–1001. IOS Press, 2020.
• Biderman et al. (2024)Dan Biderman, Jose Gonzalez Ortiz, Jacob Portes, Mansheej Paul, Philip Greengard, Connor Jennings, Daniel King, Sam Havens, Vitaliy Chiley, Jonathan Frankle, et al.LoRA learns less and forgets less.arXiv preprint arXiv:2405.09673, 2024.
• Buzzega et al. (2020)Pietro Buzzega, Matteo Boschini, Angelo Porrello, Davide Abati, and Simone Calderara.Dark experience for general continual learning: a strong, simple baseline.Advances in neural information processing systems, 33:15920–15930, 2020.
• Cha et al. (2021)Hyuntak Cha, Jaeho Lee, and Jinwoo Shin.Co2l: Contrastive continual learning.InProceedings of the IEEE/CVF International conference on computer vision, pp. 9516–9525, 2021.
• Cha et al. (2020)Sungmin Cha, Hsiang Hsu, Taebaek Hwang, Flavio P Calmon, and Taesup Moon.Cpr: classifier-projection regularization for continual learning.arXiv preprint arXiv:2006.07326, 2020.
• Chaudhry et al. (2018)Arslan Chaudhry, Puneet K Dokania, Thalaiyasingam Ajanthan, and Philip HS Torr.Riemannian walk for incremental learning: Understanding forgetting and intransigence.InEuropean Conference on Computer Vision, pp. 556–572, 2018.
• Chaudhry et al. (2019a)Arslan Chaudhry, Marc’Aurelio Ranzato, Marcus Rohrbach, and Mohamed Elhoseiny.Efficient lifelong learning with A-GEM.InInternational Conference on Learning Representations, 2019a.
• Chaudhry et al. (2019b)Arslan Chaudhry, Marcus Rohrbach, Mohamed Elhoseiny, Thalaiyasingam Ajanthan, Puneet K Dokania, Philip HS Torr, and Marc’Aurelio Ranzato.On tiny episodic memories in continual learning.arXiv preprint arXiv:1902.10486, 2019b.
• Chaudhry et al. (2021)Arslan Chaudhry, Albert Gordo, Puneet Dokania, Philip Torr, and David Lopez-Paz.Using hindsight to anchor past knowledge in continual learning.InProceedings of the AAAI conference on artificial intelligence, volume 35, pp. 6993–7001, 2021.
• Chen et al. (2021)Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al.Evaluating large language models trained on code.arXiv preprint arXiv:2107.03374, 2021.
• Chen et al. (2020)Sanyuan Chen, Yutai Hou, Yiming Cui, Wanxiang Che, Ting Liu, and Xiangzhan Yu.Recall and learn: Fine-tuning deep pretrained language models with less forgetting.InProceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 7870–7881, 2020.
• Chen et al. (2024)Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al.Symbolic discovery of optimization algorithms.Advances in neural information processing systems, 36, 2024.
• Clark et al. (2018)Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord.Think you have solved question answering? Try ARC, the AI2 reasoning challenge.arXiv preprint arXiv:1803.05457, 2018.
• Cobbe et al. (2021)Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al.Training verifiers to solve math word problems.arXiv preprint arXiv:2110.14168, 2021.
• Cui et al. (2024)Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Bingxiang He, Wei Zhu, Yuan Ni, Guotong Xie, Ruobing Xie, Yankai Lin, et al.Ultrafeedback: Boosting language models with scaled ai feedback.InForty-first International Conference on Machine Learning, 2024.
• Dai & Le (2015)Andrew M Dai and Quoc V Le.Semi-supervised sequence learning.Advances in neural information processing systems, 28, 2015.
• Dong et al. (2021)Xinshuai Dong, Anh Tuan Luu, Min Lin, Shuicheng Yan, and Hanwang Zhang.How should pre-trained language models be fine-tuned towards adversarial robustness?Advances in Neural Information Processing Systems, 34:4356–4369, 2021.
• Dozat (2016)Timothy Dozat.Incorporating nesterov momentum into adam.2016.
• Eeckt et al. (2022)Steven Vander Eeckt et al.Weight averaging: A simple yet effective method to overcome catastrophic forgetting in automatic speech recognition.arXiv preprint arXiv:2210.15282, 2022.
• Farajtabar et al. (2020)Mehrdad Farajtabar, Navid Azizan, Alex Mott, and Ang Li.Orthogonal gradient descent for continual learning.InInternational Conference on Artificial Intelligence and Statistics, pp. 3762–3773. PMLR, 2020.
• Gao et al. (2020)Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al.The Pile: An 800GB dataset of diverse text for language modeling.arXiv preprint arXiv:2101.00027, 2020.
• Gao et al. (2023)Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou.A framework for few-shot language model evaluation, 12 2023.URLhttps://zenodo.org/records/10256836.
• Goodfellow et al. (2013)Ian J Goodfellow, Mehdi Mirza, Da Xiao, Aaron Courville, and Yoshua Bengio.An empirical investigation of catastrophic forgetting in gradient-based neural networks.arXiv preprint arXiv:1312.6211, 2013.
• Gueta et al. (2023)Almog Gueta, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen.Knowledge is a region in weight space for fine-tuned language models.InFindings of the Association for Computational Linguistics: EMNLP 2023, pp. 1350–1370, 2023.
• Hayes et al. (2019)Tyler L Hayes, Nathan D Cahill, and Christopher Kanan.Memory efficient experience replay for streaming learning.In2019 International Conference on Robotics and Automation (ICRA), pp. 9769–9776. IEEE, 2019.
• He et al. (2024)Yifei He, Yuzheng Hu, Yong Lin, Tong Zhang, and Han Zhao.Localize-and-stitch: Efficient model merging via sparse task arithmetic.arXiv preprint arXiv:2408.13656, 2024.
• Hendrycks et al. (2021)Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt.Measuring massive multitask language understanding.InInternational Conference on Learning Representations, 2021.
• Hong et al. (2017)Mingyi Hong, Xiangfeng Wang, Meisam Razaviyayn, and Zhi-Quan Luo.Iteration complexity analysis of block coordinate descent methods.Mathematical Programming, 163:85–114, 2017.
• Houlsby et al. (2019)Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly.Parameter-efficient transfer learning for nlp.InInternational conference on machine learning, pp. 2790–2799. PMLR, 2019.
• Hu et al. (2022)Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al.LoRA: Low-rank adaptation of large language models.InInternational Conference on Learning Representations, 2022.
• Huang et al. (2024)Jianheng Huang, Leyang Cui, Ante Wang, Chengyi Yang, Xinting Liao, Linfeng Song, Junfeng Yao, and Jinsong Su.Mitigating catastrophic forgetting in large language models with self-synthesized rehearsal.arXiv preprint arXiv:2403.01244, 2024.
• Huang et al. (2021)Yufan Huang, Yanzhe Zhang, Jiaao Chen, Xuezhi Wang, and Diyi Yang.Continual learning for text classification with information disentanglement based regularization.InProceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2736–2746, 2021.
• Hui et al. (2024)Tingfeng Hui, Zhenyu Zhang, Shuohuan Wang, Weiran Xu, Yu Sun, and Hua Wu.Hft: Half fine-tuning for large language models.arXiv preprint arXiv:2404.18466, 2024.
• Ilharco et al. (2023)Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi.Editing models with task arithmetic.InInternational Conference on Learning Representations (ICLR). International Conference on Learning Representations, 2023.
• Ivison et al. (2023)Hamish Ivison, Yizhong Wang, Valentina Pyatkin, Nathan Lambert, Matthew Peters, Pradeep Dasigi, Joel Jang, David Wadden, Noah A Smith, Iz Beltagy, et al.Camels in a changing climate: Enhancing lm adaptation with tulu 2.arXiv preprint arXiv:2311.10702, 2023.
• Javed & White (2019)Khurram Javed and Martha White.Meta-learning representations for continual learning.Advances in neural information processing systems, 32, 2019.
• Jin et al. (2021)Di Jin, Eileen Pan, Nassim Oufattole, Wei-Hung Weng, Hanyi Fang, and Peter Szolovits.What disease does this patient have? a large-scale open domain question answering dataset from medical exams.Applied Sciences, 11(14):6421, 2021.
• Jin et al. (2019)Qiao Jin, Bhuwan Dhingra, Zhengping Liu, William Cohen, and Xinghua Lu.Pubmedqa: A dataset for biomedical research question answering.InProceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 2567–2577, 2019.
• Ke et al. (2023a)Zixuan Ke, Bing Liu, Wenhan Xiong, Asli Celikyilmaz, and Haoran Li.Sub-network discovery and soft-masking for continual learning of mixed tasks.InFindings of the Association for Computational Linguistics: EMNLP 2023, pp. 15090–15107, 2023a.
• Ke et al. (2023b)Zixuan Ke, Yijia Shao, Haowei Lin, Tatsuya Konishi, Gyuhak Kim, and Bing Liu.Continual pre-training of language models.InInternational Conference on Learning Representations (ICLR). International Conference on Learning Representations, 2023b.
• Kemker et al. (2018)Ronald Kemker, Marc McClure, Angelina Abitino, Tyler Hayes, and Christopher Kanan.Measuring catastrophic forgetting in neural networks.InProceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
• Kenton & Toutanova (2019)Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova.Bert: Pre-training of deep bidirectional transformers for language understanding.InProceedings of NAACL-HLT, pp. 4171–4186, 2019.
• Kingma & Ba (2014)Diederik P Kingma and Jimmy Ba.Adam: A method for stochastic optimization.arXiv preprint arXiv:1412.6980, 2014.
• Kirkpatrick et al. (2017)James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al.Overcoming catastrophic forgetting in neural networks.Proceedings of the national academy of sciences, 114(13):3521–3526, 2017.
• Korbak et al. (2022)Tomasz Korbak, Hady Elsahar, German Kruszewski, and Marc Dymetman.Controlling conditional language models without catastrophic forgetting.InInternational Conference on Machine Learning, pp. 11499–11528. PMLR, 2022.
• Kumar et al. (2023)Saurabh Kumar, Henrik Marklund, and Benjamin Van Roy.Maintaining plasticity via regenerative regularization.arXiv preprint arXiv:2308.11958, 2023.
• Li & Zhang (2021)Dongyue Li and Hongyang Zhang.Improved regularization and robustness for fine-tuning in neural networks.Advances in Neural Information Processing Systems, 34:27249–27262, 2021.
• Li et al. (2024)Haoling Li, Xin Zhang, Xiao Liu, Yeyun Gong, Yifan Wang, Yujiu Yang, Qi Chen, and Peng Cheng.Gradient-mask tuning elevates the upper limits of llm performance.arXiv preprint arXiv:2406.15330, 2024.
• Li et al. (2018)Xuhong Li, Yves Grandvalet, and Franck Davoine.Explicit inductive bias for transfer learning with convolutional networks.InInternational Conference on Machine Learning, pp. 2825–2834. PMLR, 2018.
• Li & Hoiem (2017)Zhizhong Li and Derek Hoiem.Learning without forgetting.IEEE transactions on pattern analysis and machine intelligence, 40(12):2935–2947, 2017.
• Lin et al. (2023)Yong Lin, Lu Tan, Hangyu Lin, Zeming Zheng, Renjie Pi, Jipeng Zhang, Shizhe Diao, Haoxiang Wang, Han Zhao, Yuan Yao, et al.Speciality vs generality: An empirical study on catastrophic forgetting in fine-tuning foundation models.arXiv preprint arXiv:2309.06256, 2023.
• Lin et al. (2024a)Yong Lin, Hangyu Lin, Wei Xiong, Shizhe Diao, Jianmeng Liu, Jipeng Zhang, Rui Pan, Haoxiang Wang, Wenbin Hu, Hanning Zhang, et al.Mitigating the alignment tax of rlhf.InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 580–606, 2024a.
• Lin et al. (2024b)Zhanran Lin, Puheng Li, and Lei Wu.Exploring neural network landscapes: Star-shaped and geodesic connectivity.arXiv preprint arXiv:2404.06391, 2024b.
• Liu et al. (2022)Chaoyue Liu, Libin Zhu, and Mikhail Belkin.Loss landscapes and optimization in over-parameterized non-linear systems and neural networks.Applied and Computational Harmonic Analysis, 59:85–116, 2022.
• Liu et al. (2024)Chengyuan Liu, Shihang Wang, Yangyang Kang, Lizhi Qing, Fubang Zhao, Changlong Sun, Kun Kuang, and Fei Wu.More than catastrophic forgetting: Integrating general capabilities for domain-specific llms.arXiv preprint arXiv:2405.17830, 2024.
• Lopez-Paz & Ranzato (2017)David Lopez-Paz and Marc’Aurelio Ranzato.Gradient episodic memory for continual learning.Advances in neural information processing systems, 30, 2017.
• Lu & Xiao (2015)Zhaosong Lu and Lin Xiao.On the complexity analysis of randomized block-coordinate descent methods.Mathematical Programming, 152:615–642, 2015.
• Luo et al. (2024)Qijun Luo, Hengxu Yu, and Xiao Li.Badam: A memory efficient full parameter training method for large language models.arXiv preprint arXiv:2404.02827, 2024.
• Luo et al. (2023a)Yun Luo, Zhen Yang, Fandong Meng, Yafu Li, Jie Zhou, and Yue Zhang.An empirical study of catastrophic forgetting in large language models during continual fine-tuning.arXiv preprint arXiv:2308.08747, 2023a.
• Luo et al. (2023b)Ziyang Luo, Can Xu, Pu Zhao, Qingfeng Sun, Xiubo Geng, Wenxiang Hu, Chongyang Tao, Jing Ma, Qingwei Lin, and Daxin Jiang.Wizardcoder: Empowering code large language models with evol-instruct, 2023b.
• McCloskey & Cohen (1989)Michael McCloskey and Neal J Cohen.Catastrophic interference in connectionist networks: The sequential learning problem.InPsychology of learning and motivation, volume 24, pp. 109–165. Elsevier, 1989.
• Mirzadeh et al. (2020a)Seyed Iman Mirzadeh, Mehrdad Farajtabar, Dilan Gorur, Razvan Pascanu, and Hassan Ghasemzadeh.Linear mode connectivity in multitask and continual learning.arXiv preprint arXiv:2010.04495, 2020a.
• Mirzadeh et al. (2020b)Seyed Iman Mirzadeh, Mehrdad Farajtabar, Razvan Pascanu, and Hassan Ghasemzadeh.Understanding the role of training regimes in continual learning.Advances in Neural Information Processing Systems, 33:7308–7320, 2020b.
• Nesterov (2012)Yu Nesterov.Efficiency of coordinate descent methods on huge-scale optimization problems.SIAM Journal on Optimization, 22(2):341–362, 2012.
• Nutini et al. (2015)Julie Nutini, Mark Schmidt, Issam Laradji, Michael Friedlander, and Hoyt Koepke.Coordinate descent converges faster with the gauss-southwell rule than random selection.InInternational Conference on Machine Learning, pp. 1632–1641. PMLR, 2015.
• Nutini et al. (2022)Julie Nutini, Issam Laradji, and Mark Schmidt.Let’s make block coordinate descent converge faster: faster greedy rules, message-passing, active-set complexity, and superlinear convergence.Journal of Machine Learning Research, 23(131):1–74, 2022.
• Ouyang et al. (2022)Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al.Training language models to follow instructions with human feedback.Advances in neural information processing systems, 35:27730–27744, 2022.
• Pal et al. (2022)Ankit Pal, Logesh Kumar Umapathi, and Malaikannan Sankarasubbu.Medmcqa: A large-scale multi-subject multi-choice dataset for medical domain question answering.InConference on health, inference, and learning, pp. 248–260. PMLR, 2022.
• Panda et al. (2024)Ashwinee Panda, Berivan Isik, Xiangyu Qi, Sanmi Koyejo, Tsachy Weissman, and Prateek Mittal.Lottery ticket adaptation: Mitigating destructive interference in llms.arXiv preprint arXiv:2406.16797, 2024.
• Panigrahi et al. (2023)Abhishek Panigrahi, Nikunj Saunshi, Haoyu Zhao, and Sanjeev Arora.Task-specific skill localization in fine-tuned language models.InInternational Conference on Machine Learning, pp. 27011–27033. PMLR, 2023.
• Radford et al. (2018)Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.Improving language understanding with unsupervised learning.2018.
• Ramasesh et al. (2021)Vinay Venkatesh Ramasesh, Aitor Lewkowycz, and Ethan Dyer.Effect of scale on catastrophic forgetting in neural networks.InInternational Conference on Learning Representations, 2021.
• Rannen et al. (2017)Amal Rannen, Rahaf Aljundi, Matthew B Blaschko, and Tinne Tuytelaars.Encoder based lifelong learning.InProceedings of the IEEE international conference on computer vision, pp. 1320–1328, 2017.
• Razaviyayn et al. (2013)Meisam Razaviyayn, Mingyi Hong, and Zhi-Quan Luo.A unified convergence analysis of block successive minimization methods for nonsmooth optimization.SIAM Journal on Optimization, 23(2):1126–1153, 2013.
• Razdaibiedina et al. (2023)Anastasia Razdaibiedina, Yuning Mao, Rui Hou, Madian Khabsa, Mike Lewis, and Amjad Almahairi.Progressive prompts: Continual learning for language models.arXiv preprint arXiv:2301.12314, 2023.
• Ren et al. (2024)Weijieying Ren, Xinlong Li, Lei Wang, Tianxiang Zhao, and Wei Qin.Analyzing and reducing catastrophic forgetting in parameter efficient tuning.arXiv preprint arXiv:2402.18865, 2024.
• Richtárik & Takáč (2014)Peter Richtárik and Martin Takáč.Iteration complexity of randomized block-coordinate descent methods for minimizing a composite function.Mathematical Programming, 144(1):1–38, 2014.
• Riemer et al. (2019a)Matthew Riemer, Ignacio Cases, Robert Ajemian, Miao Liu, Irina Rish, Yuhai Tu, and Gerald Tesauro.Learning to learn without forgetting by maximizing transfer and minimizing interference.InInternational Conference on Learning Representations, 2019a.
• Riemer et al. (2019b)Matthew Riemer, Tim Klinger, Djallel Bouneffouf, and Michele Franceschini.Scalable recollections for continual lifelong learning.InProceedings of the AAAI conference on artificial intelligence, volume 33, pp. 1352–1359, 2019b.
• Ritter et al. (2018)Hippolyt Ritter, Aleksandar Botev, and David Barber.Online structured laplace approximations for overcoming catastrophic forgetting.Advances in Neural Information Processing Systems, 31, 2018.
• Rolnick et al. (2019)David Rolnick, Arun Ahuja, Jonathan Schwarz, Timothy Lillicrap, and Gregory Wayne.Experience replay for continual learning.Advances in neural information processing systems, 32, 2019.
• Rothchild et al. (2020)Daniel Rothchild, Ashwinee Panda, Enayat Ullah, Nikita Ivkin, Ion Stoica, Vladimir Braverman, Joseph Gonzalez, and Raman Arora.Fetchsgd: Communication-efficient federated learning with sketching.InInternational Conference on Machine Learning, pp. 8253–8265. PMLR, 2020.
• Roziere et al. (2023)Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, et al.Code llama: Open foundation models for code.arXiv preprint arXiv:2308.12950, 2023.
• Shi et al. (2024)Haizhou Shi, Zihao Xu, Hengyi Wang, Weiyi Qin, Wenyuan Wang, Yibin Wang, and Hao Wang.Continual learning of large language models: A comprehensive survey.arXiv preprint arXiv:2404.16789, 2024.
• Shi et al. (2021)Naichen Shi, Dawei Li, Mingyi Hong, and Ruoyu Sun.Rmsprop converges with proper hyper-parameter.In9th International Conference on Learning Representations, ICLR 2021, 2021.
• Sun & Hong (2015)Ruoyu Sun and Mingyi Hong.Improved iteration complexity bounds of cyclic block coordinate descent for convex problems.Advances in Neural Information Processing Systems, 28, 2015.
• Suzgun et al. (2022)Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, et al.Challenging big-bench tasks and whether chain-of-thought can solve them.arXiv preprint arXiv:2210.09261, 2022.
• Team et al. (2024)Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al.Gemma: Open models based on gemini research and technology.arXiv preprint arXiv:2403.08295, 2024.
• Tieleman (2012)Tijmen Tieleman.Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude.COURSERA: Neural networks for machine learning, 4(2):26, 2012.
• Tiwari et al. (2022)Rishabh Tiwari, Krishnateja Killamsetty, Rishabh Iyer, and Pradeep Shenoy.Gcr: Gradient coreset based replay buffer selection for continual learning.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 99–108, 2022.
• Together (2023)Together.Redpajama, a project to create leading open-source models, starts by reproducing llama training dataset of over 1.2 trillion tokens.https://www.together.ai/blog/redpajama, 2023.Accessed: 2025-08-14.
• Touvron et al. (2023)Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al.Llama 2: Open foundation and fine-tuned chat models.arXiv preprint arXiv:2307.09288, 2023.
• Tseng (2001)Paul Tseng.Convergence of a block coordinate descent method for nondifferentiable minimization.Journal of optimization theory and applications, 109:475–494, 2001.
• Verwimp et al. (2023)Eli Verwimp, Rahaf Aljundi, Shai Ben-David, Matthias Bethge, Andrea Cossu, Alexander Gepperth, Tyler L Hayes, Eyke Hüllermeier, Christopher Kanan, Dhireesha Kudithipudi, et al.Continual learning: Applications and the road forward.arXiv preprint arXiv:2311.11908, 2023.
• Wang et al. (2024a)Liyuan Wang, Xingxing Zhang, Hang Su, and Jun Zhu.A comprehensive survey of continual learning: theory, method and application.IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024a.
• Wang et al. (2023a)Xiao Wang, Tianze Chen, Qiming Ge, Han Xia, Rong Bao, Rui Zheng, Qi Zhang, Tao Gui, and Xuanjing Huang.Orthogonal subspace learning for language model continual learning.InThe 2023 Conference on Empirical Methods in Natural Language Processing, 2023a.
• Wang et al. (2023b)Xiao Wang, Yuansen Zhang, Tianze Chen, Songyang Gao, Senjie Jin, Xianjun Yang, Zhiheng Xi, Rui Zheng, Yicheng Zou, Tao Gui, et al.Trace: A comprehensive benchmark for continual learning in large language models.arXiv preprint arXiv:2310.06762, 2023b.
• Wang et al. (2024b)Yifan Wang, Yafei Liu, Chufan Shi, Haoling Li, Chen Chen, Haonan Lu, and Yujiu Yang.Inscl: A data-efficient continual learning paradigm for fine-tuning large language models with instructions.InProceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 663–677, 2024b.
• Wang et al. (2020)Zirui Wang, Sanket Vaibhav Mehta, Barnabás Poczós, and Jaime G Carbonell.Efficient meta lifelong-learning with limited memory.InProceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 535–548, 2020.
• Weber et al. (2024)Maurice Weber, Daniel Y. Fu, Quentin Anthony, Yonatan Oren, Shane Adams, Anton Alexandrov, Xiaozhong Lyu, Huu Nguyen, Xiaozhe Yao, Virginia Adams, Ben Athiwaratkun, Rahul Chalamala, Kezhen Chen, Max Ryabinin, Tri Dao, Percy Liang, Christopher Ré, Irina Rish, and Ce Zhang.Redpajama: an open dataset for training large language models.NeurIPS Datasets and Benchmarks Track, 2024.
• Wei et al. (2023)Yuxiang Wei, Zhe Wang, Jiawei Liu, Yifeng Ding, and Lingming Zhang.Magicoder: Empowering code generation with oss-instruct.arXiv preprint arXiv:2312.02120, 2023.
• Wortsman et al. (2022a)Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al.Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time.InInternational conference on machine learning, pp. 23965–23998. PMLR, 2022a.
• Wortsman et al. (2022b)Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, et al.Robust fine-tuning of zero-shot models.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 7959–7971, 2022b.
• Wu et al. (2024)Chaoyi Wu, Weixiong Lin, Xiaoman Zhang, Ya Zhang, Weidi Xie, and Yanfeng Wang.Pmc-llama: toward building open-source language models for medicine.Journal of the American Medical Informatics Association, pp. ocae045, 2024.
• Xiao et al. (2024)Nachuan Xiao, Xiaoyin Hu, Xin Liu, and Kim-Chuan Toh.Adam-family methods for nonsmooth optimization with convergence guarantees.Journal of Machine Learning Research, 25(48):1–53, 2024.
• Xu & Zhang (2024)Jing Xu and Jingzhao Zhang.Random masking finds winning tickets for parameter efficient fine-tuning.arXiv preprint arXiv:2405.02596, 2024.
• Yadav et al. (2024)Prateek Yadav, Derek Tam, Leshem Choshen, Colin A Raffel, and Mohit Bansal.Ties-merging: Resolving interference when merging models.Advances in Neural Information Processing Systems, 36, 2024.
• Yin et al. (2023)Da Yin, Xiao Liu, Fan Yin, Ming Zhong, Hritik Bansal, Jiawei Han, and Kai-Wei Chang.Dynosaur: A dynamic growth paradigm for instruction-tuning data curation.InProceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 4031–4047, 2023.
• Yu et al. (2024a)Le Yu, Bowen Yu, Haiyang Yu, Fei Huang, and Yongbin Li.Language models are super mario: Absorbing abilities from homologous models as a free lunch.InForty-first International Conference on Machine Learning, 2024a.
• Yu et al. (2024b)Longhui Yu, Weisen Jiang, Han Shi, YU Jincheng, Zhengying Liu, Yu Zhang, James Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu.Metamath: Bootstrap your own mathematical questions for large language models.InThe Twelfth International Conference on Learning Representations, 2024b.
• Zaken et al. (2021)Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg.Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models.arXiv preprint arXiv:2106.10199, 2021.
• Zellers et al. (2019)Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.Hellaswag: Can a machine really finish your sentence?InProceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4791–4800, 2019.
• Zenke et al. (2017)Friedemann Zenke, Ben Poole, and Surya Ganguli.Continual learning through synaptic intelligence.InInternational conference on machine learning, pp. 3987–3995. PMLR, 2017.
• Zhang et al. (2024a)Hengyuan Zhang, Yanru Wu, Dawei Li, Sak Yang, Rui Zhao, Yong Jiang, and Fei Tan.Balancing speciality and versatility: a coarse to fine framework for supervised fine-tuning large language model.arXiv preprint arXiv:2404.10306, 2024a.
• Zhang et al. (2023)Hongyang R Zhang, Dongyue Li, and Haotian Ju.Noise stability optimization for finding flat minima: A hessian-based regularization approach.arXiv preprint arXiv:2306.08553, 2023.
• Zhang et al. (2024b)Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu.Tinyllama: An open-source small language model.arXiv preprint arXiv:2401.02385, 2024b.
• Zhang et al. (2024c)Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo.Why transformers need adam: A hessian perspective.arXiv preprint arXiv:2402.16788, 2024c.
• Zhang et al. (2024d)Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun.Adam-mini: Use fewer learning rates to gain more.arXiv preprint arXiv:2406.16793, 2024d.
• Zhao et al. (2014)Tuo Zhao, Mo Yu, Yiming Wang, Raman Arora, and Han Liu.Accelerated mini-batch randomized block coordinate descent method.Advances in neural information processing systems, 27, 2014.
• Zhu et al. (2024)Didi Zhu, Zhongyi Sun, Zexi Li, Tao Shen, Ke Yan, Shouhong Ding, Kun Kuang, and Chao Wu.Model tailor: Mitigating catastrophic forgetting in multi-modal large language models.arXiv preprint arXiv:2402.12048, 2024.

Appendix ATheoretical Analysis

AppendixAis organized into three self-contained parts. A roadmap is provided to guide the reader in locating and understanding each theoretical result efficiently.

Roadmap to Appendix A (Quick Reference).

AppendixA.1: Supplementary Analysis on the Top-α\alphaFilter

This section formalizes the top-α\alphafilterFLTα​(⋅)\mathrm{FLT}_{\alpha}(\cdot)and the induced quantity‖z⊙FLTα​(z)‖1\|z\odot\mathrm{FLT}_{\alpha}(z)\|_{1}that will serve as our working norm in the analysis. Proposition1verifies that this is indeed a norm. Lemma1establishes aLipschitz stabilityproperty for filtered outputs, which later lets us pass from the momentum-defined filtering to the gradient-defined filtering. Lemma2relates theL1,top-​αL_{1,\text{top-}\alpha}norm to standardLpL_{p}norms, which is used at the very end to translate the convergence bound to anyp∈[1,∞]p\in[1,\infty].

AppendixA.2: Proof of Theorem 1 (Convergence of MoFO)

For a high-level narrative, see the proof sketch accompanying Theorem1in the main text; here we provide a quick-reference map to the technical steps.

Lemma3bounds each step size and theℓ2\ell_{2}-movement of parameters using the number of active coordinates. Lemma4controls the drift of gradients across iterations viaLL-smoothness. Lemma5lower-bounds the per-coordinate alignment termgi,t​m^i,t/v^i,tg_{i,t}\hat{m}_{i,t}/\sqrt{\hat{v}_{i,t}}, which is the driver in the descent inequality. Lemma6shows that the bias-corrected momentumm^t\hat{m}_{t}tracksgtg_{t}inℓ1\ell_{1}at rateO​(t−1/2)O(t^{-1/2}). Combining these with the descent lemma yields a key inequality in which the filter isFLTα​(mt)\texttt{FLT}_{\alpha}(m_{t}):

C1t​‖gt⊙FLTα​(mt)‖1≤ℒ​(θt−1)−ℒ​(θt)+C2t.\frac{C_{1}}{\sqrt{t}}\|g_{t}\odot\texttt{FLT}_{\alpha}(m_{t})\|_{1}\leq\mathcal{L}(\theta_{t-1})-\mathcal{L}(\theta_{t})+\frac{C_{2}}{t}. Finally, Lemma1help convert it to the desired inequality with the gradient filterFLTα​(gt)\texttt{FLT}_{\alpha}(g_{t}); summing overttgivesmin0≤t≤T−1⁡‖∇ℒ​(θt)⊙FLTα​(∇ℒ​(θt))‖1=O​(log⁡T/T)\min_{0\leq t\leq T-1}\|\nabla\mathcal{L}(\theta_{t})\odot\texttt{FLT}_{\alpha}(\nabla\mathcal{L}(\theta_{t}))\|_{1}=O(\log T/\sqrt{T}). By Lemma2, this implies the sameO​(log⁡T/T)O(\log T/\sqrt{T})rate for anyLpL_{p}normmin0≤t≤T−1⁡‖∇ℒ​(θt)‖p\min_{0\leq t\leq T-1}\|\nabla\mathcal{L}(\theta_{t})\|_{p}.

AppendixA.3: Proof of Theorem2(Illustrative example: forgetting mitigation of MoFO)

In an illustrative example with updating ratioα=1/d\alpha=1/d, MoFO converges to a single-attractor minimum and is strictly closer to the pre-training state than Adam, thus attaining a lower pre-training loss. We prove this by mathematical induction.

AppendixA.4: Challenges and Potential Extensions of Theorem1to Nonsmooth Objectives

This subsection explains why extending the convergence analysis of Theorem1to nonsmooth objectives is nontrivial and outlines several possible future directions.

A.1Supplementary Analysis on the Top-α\alphaFilter

In this section, we provide supplementary analysis on our top-α\alphafilter, which serves as a preliminary for proving Theorem1in AppendixA.2.

As introduced in Section2.2, the entire parameter space is divided intoBBparts, with thekk-th part having a dimension ofdkd_{k}. We assume the parameter space isℝd\mathbb{R}^{d}, which can be expressed as the productℝd≅ℝd1×ℝd2×⋯×ℝdB\mathbb{R}^{d}\cong\mathbb{R}^{d_{1}}\times\mathbb{R}^{d_{2}}\times\dots\times\mathbb{R}^{d_{B}}. For anyz∈ℝdz\in\mathbb{R}^{d}, we represent it as:

z=Concat​(z(1),z(2),…,z(B)),z=\text{Concat}(z^{(1)},z^{(2)},\dots,z^{(B)}),wherez(k)∈ℝdkz^{(k)}\in\mathbb{R}^{d_{k}}for each1≤k≤B1\leq k\leq B.

Definition 1.

For anyz∈ℝdz\in\mathbb{R}^{d}, we define the top-α\alphafilter ofzzas

FLTα​(z):=Concat​(𝐞S1(1);𝐞S2(2);…;𝐞SB(B))∈ℝd,\texttt{FLT}_{\alpha}(z):={\rm Concat}({\mathbf{e}}^{(1)}_{S_{1}};{\mathbf{e}}^{(2)}_{S_{2}};\dots;{\mathbf{e}}^{(B)}_{S_{B}})\in\mathbb{R}^{d},where

Sk={i∈[dk]:|zi(k)|​ranks within the top-​α​of all​|z(k)|​’s entries​(|z1(k)|,|z2(k)|,…,|zdk(k)|)}S_{k}=\{i\in[d_{k}]:|z^{(k)}_{i}|\text{ ranks within the top-}\alpha\text{ of all }|z^{(k)}|\text{’s entries }(|z^{(k)}_{1}|,|z^{(k)}_{2}|,\dots,|z^{(k)}_{d_{k}}|)\}and𝐞Sk(k){\mathbf{e}}_{S_{k}}^{(k)}is adkd_{k}-dimensional vector where theii-th entry is 1 ifi∈Ski\in S_{k}, and 0 otherwise.

Remark 2.

To ensure that the top-α\alphafilterFLTα​(z)\texttt{FLT}_{\alpha}(z)is well-defined, when multiple entries share identical absolute values and including all of them in the setSkS_{k}would result in exceeding theα\alphathreshold of set size, the construction ofSkS_{k}prioritizes the entries with the smallest indices among those with the same absolute values.

Definition 2.

For anyz∈ℝdz\in\mathbb{R}^{d}, we define theL1,top-​αL_{1,\text{top-}\alpha}norm ofzzas

‖z‖1,top-​α:=‖z⊙FLTα​(z)‖1.\|{z}\|_{1,\text{top-}\alpha}:=\|z\odot\texttt{FLT}_{\alpha}(z)\|_{1}.

Proposition 1.

∥⋅∥1,top-​α\|{\cdot}\|_{1,\text{top-}\alpha}is indeed a norm inℝd{\mathbb{R}}^{d}.

Proof.

By Definition1, we get

‖z‖1,top-​α=‖z⊙FLTα​(z)‖1=∑k=1B‖z(k)⊙𝐞Sk(k)‖1.\|{z}\|_{1,\text{top-}\alpha}=\|z\odot\texttt{FLT}_{\alpha}(z)\|_{1}=\sum_{k=1}^{B}\|z^{(k)}\odot{\mathbf{e}}_{S_{k}}^{(k)}\|_{1}.(6) First, if‖z‖1,top-​α=0\|{z}\|_{1,\text{top-}\alpha}=0, then by (6),‖z(k)⊙𝐞Sk(k)‖1=0\|z^{(k)}\odot{\mathbf{e}}_{S_{k}}^{(k)}\|_{1}=0for any1≤k≤B1\leq k\leq B. Thus,

‖z(k)‖∞=arg​max1≤i≤dk⁡|zi(k)|≤‖z(k)⊙𝐞Sk(k)‖1=0.\|z^{(k)}\|_{\infty}=\operatorname*{arg\,max}_{1\leq i\leq d_{k}}|z^{(k)}_{i}|\leq\|z^{(k)}\odot{\mathbf{e}}_{S_{k}}^{(k)}\|_{1}=0.Soz(k)z^{(k)}is a zero vector for any1≤k≤B1\leq k\leq Band thenzzis a zero vector.

Second, for any givenc∈ℝ+c\in\mathbb{R}_{+},{|zi(k)|}1≤i≤dk\{|z^{(k)}_{i}|\}_{1\leq i\leq d_{k}}and{|c​zi(k)|}1≤i≤dk\{|cz^{(k)}_{i}|\}_{1\leq i\leq d_{k}}have the same order. Sozzandc​zczshare the same filterFLTα​(z)\texttt{FLT}_{\alpha}(z)and

‖c​z‖1,top-​α=‖c​z⊙FLTα​(c​z)‖1=c​‖z⊙FLTα​(z)‖1=c​‖z‖1,top-​α.\|{cz}\|_{1,\text{top-}\alpha}=\|cz\odot\texttt{FLT}_{\alpha}(cz)\|_{1}=c\|z\odot\texttt{FLT}_{\alpha}(z)\|_{1}=c\|{z}\|_{1,\text{top-}\alpha}. Third, for anyx,y∈ℝdx,y\in\mathbb{R}^{d}, we let

Sk′\displaystyle S_{k}^{\prime}={i∈[dk]:|xi(k)|​ranks within the top-​α​of all​|x(k)|​’s entries​(|x1(k)|,|x2(k)|,…,|xdk(k)|)},\displaystyle=\{i\in[d_{k}]:|x^{(k)}_{i}|\text{ ranks within the top-}\alpha\text{ of all }|x^{(k)}|\text{’s entries }(|x^{(k)}_{1}|,|x^{(k)}_{2}|,\dots,|x^{(k)}_{d_{k}}|)\},Sk′′\displaystyle S_{k}^{\prime\prime}={i∈[dk]:|xi(k)+yi(k)|ranks within the top-αof all\displaystyle=\{i\in[d_{k}]:|x^{(k)}_{i}+y^{(k)}_{i}|\text{ ranks within the top-}\alpha\text{ of all }|x(k)+y(k)|’s entries(|x1(k)+y1(k)|,|x2(k)+y2(k)|,…,|xdk(k)+ydk(k)|)}.\displaystyle\qquad\qquad\qquad|x^{(k)}+y^{(k)}|\text{’s entries }(|x^{(k)}_{1}+y^{(k)}_{1}|,|x^{(k)}_{2}+y^{(k)}_{2}|,\dots,|x^{(k)}_{d_{k}}+y^{(k)}_{d_{k}}|)\}.Then we have

FLTα​(x)=Concat​(𝐞S1′(1);𝐞S2′(2);…;𝐞SB′(B))andFLTα​(x+y)=Concat​(𝐞S1′′(1);𝐞S2′′(2);…;𝐞SB′′(B)).\displaystyle\texttt{FLT}_{\alpha}(x)={\rm Concat}({\mathbf{e}}^{(1)}_{S^{\prime}_{1}};{\mathbf{e}}^{(2)}_{S^{\prime}_{2}};\dots;{\mathbf{e}}^{(B)}_{S^{\prime}_{B}})\quad\text{and}\quad\texttt{FLT}_{\alpha}(x+y)={\rm Concat}({\mathbf{e}}^{(1)}_{S^{\prime\prime}_{1}};{\mathbf{e}}^{(2)}_{S^{\prime\prime}_{2}};\dots;{\mathbf{e}}^{(B)}_{S^{\prime\prime}_{B}}).By the construction ofSk′S^{\prime}_{k}, for any1≤k≤B1\leq k\leq B, we have

‖x(k)⊙𝐞Sk′′(k)‖1≤‖x(k)⊙𝐞Sk′(k)‖1.\|x^{(k)}\odot{\mathbf{e}}^{(k)}_{S^{\prime\prime}_{k}}\|_{1}\leq\|x^{(k)}\odot{\mathbf{e}}^{(k)}_{S^{\prime}_{k}}\|_{1}.So

‖x⊙FLTα​(x+y)‖1=∑k=1B‖x(k)⊙𝐞Sk′′(k)‖1≤∑k=1B‖x(k)⊙𝐞Sk′(k)‖1=‖x⊙FLTα​(x)‖1.\|x\odot\texttt{FLT}_{\alpha}(x+y)\|_{1}=\sum_{k=1}^{B}\|x^{(k)}\odot{\mathbf{e}}^{(k)}_{S^{\prime\prime}_{k}}\|_{1}\leq\sum_{k=1}^{B}\|x^{(k)}\odot{\mathbf{e}}^{(k)}_{S^{\prime}_{k}}\|_{1}=\|x\odot\texttt{FLT}_{\alpha}(x)\|_{1}.Similarly, it holds that

‖y⊙FLTα​(x+y)‖1≤‖y⊙FLTα​(y)‖1.\|y\odot\texttt{FLT}_{\alpha}(x+y)\|_{1}\leq\|y\odot\texttt{FLT}_{\alpha}(y)\|_{1}.Thus, we have

‖x+y‖1,top-​α\displaystyle\|{x+y}\|_{1,\text{top-}\alpha}=‖(x+y)⊙FLTα​(x+y)‖1\displaystyle=\|(x+y)\odot\texttt{FLT}_{\alpha}(x+y)\|_{1}=‖x⊙FLTα​(x+y)+y⊙FLTα​(x+y)‖1\displaystyle=\|x\odot\texttt{FLT}_{\alpha}(x+y)+y\odot\texttt{FLT}_{\alpha}(x+y)\|_{1}≤‖x⊙FLTα​(x+y)‖1+‖y⊙FLTα​(x+y)‖1\displaystyle\leq\|x\odot\texttt{FLT}_{\alpha}(x+y)\|_{1}+\|y\odot\texttt{FLT}_{\alpha}(x+y)\|_{1}≤‖x⊙FLTα​(x)‖1+‖y⊙FLTα​(y)‖1\displaystyle\leq\|x\odot\texttt{FLT}_{\alpha}(x)\|_{1}+\|y\odot\texttt{FLT}_{\alpha}(y)\|_{1}=‖x‖1,top-​α+‖y‖1,top-​α.\displaystyle=\|{x}\|_{1,\text{top-}\alpha}+\|{y}\|_{1,\text{top-}\alpha}.∎

We propose a lemma which is useful for the proof of Theorem1.

Lemma 1.

For anyx,y∈ℝdx,y\in{\mathbb{R}}^{d}, it holds that

‖x⊙FLTα​(x)‖1−‖x⊙FLTα​(y)‖1≤2​‖x−y‖1.\|x\odot\texttt{FLT}_{\alpha}(x)\|_{1}-\|x\odot\texttt{FLT}_{\alpha}(y)\|_{1}\leq 2\|x-y\|_{1}.

Proof.

By Proposition1,∥⋅∥1,top-​α\|{\cdot}\|_{1,\text{top-}\alpha}is a norm inℝd{\mathbb{R}}^{d}, so we have

‖x⊙FLTα​(x)‖1−‖x⊙FLTα​(y)‖1\displaystyle\quad\ \|x\odot\texttt{FLT}_{\alpha}(x)\|_{1}-\|x\odot\texttt{FLT}_{\alpha}(y)\|_{1}=‖x⊙FLTα​(x)‖1−‖y⊙FLTα​(y)‖1+‖y⊙FLTα​(y)‖1−‖x⊙FLTα​(y)‖1\displaystyle=\|x\odot\texttt{FLT}_{\alpha}(x)\|_{1}-\|y\odot\texttt{FLT}_{\alpha}(y)\|_{1}+\|y\odot\texttt{FLT}_{\alpha}(y)\|_{1}-\|x\odot\texttt{FLT}_{\alpha}(y)\|_{1}=‖x‖1,top-​α−‖y‖1,top-​α+‖y⊙FLTα​(y)‖1−‖x⊙FLTα​(y)‖1\displaystyle=\|{x}\|_{1,\text{top-}\alpha}-\|{y}\|_{1,\text{top-}\alpha}+\|y\odot\texttt{FLT}_{\alpha}(y)\|_{1}-\|x\odot\texttt{FLT}_{\alpha}(y)\|_{1}≤‖x−y‖1,top-​α+‖(y−x)⊙FLTα​(y)‖1\displaystyle\leq\|{x-y}\|_{1,\text{top-}\alpha}+\|(y-x)\odot\texttt{FLT}_{\alpha}(y)\|_{1}≤‖x−y‖1+‖y−x‖1\displaystyle\leq\|x-y\|_{1}+\|y-x\|_{1}=2​‖x−y‖1.\displaystyle=2\|x-y\|_{1}.∎

The lemma below quantifies the relationship betweenL1,top-​αL_{1,\text{top-}\alpha}andLpL_{p}norms forp∈[1,+∞]p\in[1,+\infty].

Lemma 2.

Assume the parameter spaceℝd{\mathbb{R}}^{d}is decomposed intoBBblocksℝd1×⋯×ℝdB{\mathbb{R}}^{d_{1}}\times\cdots\times{\mathbb{R}}^{d_{B}}andz=Concat​(z(1),…,z(B))z=\mathrm{Concat}(z^{(1)},\dots,z^{(B)})withz(k)∈ℝdkz^{(k)}\in{\mathbb{R}}^{d_{k}}. Then for anyz∈ℝdz\in{\mathbb{R}}^{d}and anyp∈[1,∞]p\in[1,\infty], it holds that

α​‖z‖p≤‖z‖1,top-​α≤(d​α+B)1−1p​‖z‖p.\alpha\|z\|_{p}\leq\|{z}\|_{1,\text{top-}\alpha}\leq(d\alpha+B)^{1-\frac{1}{p}}\|z\|_{p}.

Proof.

We first prove thep=1p=1case. Fix any blockkk. Write the absolute values in nonincreasing ordera1(k)≥⋯≥adk(k)≥0a_{1}^{(k)}\geq\cdots\geq a_{d_{k}}^{(k)}\geq 0, whereai(k):=|zi(k)|a_{i}^{(k)}:=|z_{i}^{(k)}|. Letmk:=|Sk​(z)|m_{k}:=|S_{k}(z)|; by the definition of the top-α\alphafilter (Definition1) and the tie-breaking rule in Remark2,mk=⌈α​dk⌉m_{k}=\lceil\alpha d_{k}\rceil. Since the average of the topmkm_{k}numbers is at least the overall average, we have

∑i∈Sk​(z)|zi(k)|≥mkdk​∑i=1dk|zi(k)|≥α​‖z(k)‖1.\sum_{i\in S_{k}(z)}|z_{i}^{(k)}|\geq\frac{m_{k}}{d_{k}}\sum_{i=1}^{d_{k}}|z_{i}^{(k)}|\geq\alpha\|z^{(k)}\|_{1}.Summing overkkyields‖z‖1,top-​α≥α​∑k‖z(k)‖1=α​‖z‖1\|{z}\|_{1,\text{top-}\alpha}\geq\alpha\sum_{k}\|z^{(k)}\|_{1}=\alpha\|z\|_{1}. The upper bound‖z‖1,top-​α≤‖z‖1\|{z}\|_{1,\text{top-}\alpha}\leq\|z\|_{1}is immediate sinceFLTα​(z)\texttt{FLT}_{\alpha}(z)is a{0,1}\{0,1\}filter.

Next, letp∈(1,+∞]p\in(1,+\infty]. The lower bound follows from

‖z‖1,top-​α≥α​‖z‖1≥α​‖z‖p,\|{z}\|_{1,\text{top-}\alpha}\geq\alpha\|z\|_{1}\geq\alpha\|z\|_{p},since‖z‖1≥‖z‖p\|z\|_{1}\geq\|z\|_{p}. For the upper bound, fixkkand apply Hölder’s inequality on the subsetSk​(z)S_{k}(z):

∑i∈Sk​(z)|zi(k)|≤(∑i∈Sk​(z)1q)1/q​(∑i=1dk|zi(k)|p)1/p=mk1/q​‖z(k)‖p=mk1−1p​‖z(k)‖p,\sum_{i\in S_{k}(z)}|z_{i}^{(k)}|\leq\Big(\sum_{i\in S_{k}(z)}1^{q}\Big)^{1/q}\Big(\sum_{i=1}^{d_{k}}|z_{i}^{(k)}|^{p}\Big)^{1/p}=m_{k}^{1/q}\|z^{(k)}\|_{p}=m_{k}^{1-\frac{1}{p}}\|z^{(k)}\|_{p},whereq=pp−1q=\frac{p}{p-1}and we used1/q=1−1/p1/q=1-1/p. Summing overkkand applying Hölder again to the finite sum∑k=1Bmk1−1/p​‖z(k)‖p\sum_{k=1}^{B}m_{k}^{1-1/p}\|z^{(k)}\|_{p}gives

‖z‖1,top-​α=∑k=1B∑i∈Sk​(z)|zi(k)|≤∑k=1Bmk1−1p​‖z(k)‖p≤(∑k=1Bmk)1−1p​(∑k=1B‖z(k)‖pp)1/p.\|{z}\|_{1,\text{top-}\alpha}=\sum_{k=1}^{B}\sum_{i\in S_{k}(z)}|z_{i}^{(k)}|\leq\sum_{k=1}^{B}m_{k}^{1-\frac{1}{p}}\|z^{(k)}\|_{p}\leq\Big(\sum_{k=1}^{B}m_{k}\Big)^{1-\frac{1}{p}}\Big(\sum_{k=1}^{B}\|z^{(k)}\|_{p}^{p}\Big)^{1/p}.We setm:=∑k=1Bmkm:=\sum_{k=1}^{B}m_{k}. Since∑k=1B‖z(k)‖pp=‖z‖pp\sum_{k=1}^{B}\|z^{(k)}\|_{p}^{p}=\|z\|_{p}^{p}(blocks are disjoint) andm=∑kmk=∑k⌈α​dk⌉≤α​d+Bm=\sum_{k}m_{k}=\sum_{k}\lceil\alpha d_{k}\rceil\leq\alpha d+B, we obtain

‖z‖1,top-​α≤m1−1p​‖z‖p≤(α​d+B)1−1p​‖z‖p.\|{z}\|_{1,\text{top-}\alpha}\leq m^{1-\frac{1}{p}}\|z\|_{p}\leq(\alpha d+B)^{1-\frac{1}{p}}\|z\|_{p}.Forp=+∞p=+\inftythe same argument applies withq=1q=1and yields∑i∈Sk​(z)|zi(k)|≤mk​‖z(k)‖∞\sum_{i\in S_{k}(z)}|z_{i}^{(k)}|\leq m_{k}\|z^{(k)}\|_{\infty}, hence‖z‖1,top-​α≤m​‖z‖∞\|{z}\|_{1,\text{top-}\alpha}\leq m\|z\|_{\infty}; the lower bound‖z‖1,top-​α≥α​‖z‖∞\|{z}\|_{1,\text{top-}\alpha}\geq\alpha\|z\|_{\infty}follows from‖z‖1≥‖z‖∞\|z\|_{1}\geq\|z\|_{\infty}. This completes the proof. ∎

A.2Proof of Theorem1(Convergence of MoFO)

Notation recap (AppendixA.2).

• •gt=∇ℒ​(θt−1)∈ℝdg_{t}=\nabla\mathcal{L}(\theta_{t-1})\in\mathbb{R}^{d}is the full-batch gradient at steptt.
• •Step size:ηt=η/t\eta_{t}=\eta/\sqrt{t}; hyperparameters satisfyβ1<β2<1\beta_{1}<\sqrt{\beta_{2}}<1.
• •First/second moments:mt,vt∈ℝdm_{t},v_{t}\in\mathbb{R}^{d}with updates (mi,t,vi,t)=((1−β1)​gi,t+β1​mi,t−1,(1−β2)​gi,t2+β2​vi,t−1);(m_{i,t},v_{i,t})=\big((1-\beta_{1})g_{i,t}+\beta_{1}m_{i,t-1},\,(1-\beta_{2})g_{i,t}^{2}+\beta_{2}v_{i,t-1}\big);bias-corrected first/second moments:m^t=mt/(1−β1t)\hat{m}_{t}=m_{t}/(1-\beta_{1}^{t}),v^t=vt/(1−β2t)\hat{v}_{t}=v_{t}/(1-\beta_{2}^{t}).
• •MoFO’s filterFLTα​(⋅)∈{0,1}d\texttt{FLT}_{\alpha}(\cdot)\in\{0,1\}^{d}: in each partitionk∈[B]k\in[B]of sizedkd_{k}(with∑kdk=d\sum_{k}d_{k}=d), we keep⌈dk​α⌉\lceil d_{k}\alpha\rceilentries and zero out the others. We write‖z⊙FLTα​(z)‖1\|z\odot\texttt{FLT}_{\alpha}(z)\|_{1}for theℓ1\ell_{1}-norm of the kept coordinates and‖x‖1,top-​α≜‖x⊙FLTα​(x)‖1\|{x}\|_{1,\text{top-}\alpha}\triangleq\|x\odot\texttt{FLT}_{\alpha}(x)\|_{1}.
• •MoFO’s update:θi,t−θi,t−1=−ηt​m^i,t/v^i,t\theta_{i,t}-\theta_{i,t-1}=-\,\eta_{t}\,\hat{m}_{i,t}/\sqrt{\hat{v}_{i,t}}ifFLTα​(mt)i=1\texttt{FLT}_{\alpha}(m_{t})_{i}=1, and0otherwise.

Our proof of Theorem1follows the convergence analysis of the full-batch Adam optimizer inShi et al. (2021), with novel adaptations to address the unique aspects of MoFO.

To maintain consistency with the notation used in MoFO (Algorithm1in Section2.2), we denote

zt=Concat​(zt(1),…,zt(B)),z_{t}=\texttt{Concat}(z_{t}^{(1)},\dots,z_{t}^{(B)}),wherezzrepresents the model parameterθ\theta, the gradientgg, the first moment estimatemm, or the second moment estimatevv. Notably, each of these variables belongs toℝd\mathbb{R}^{d}. Thus, for any1≤i≤d1\leq i\leq d, we can denotezi,tz_{i,t}as theii-th entry ofztz_{t}whenzzrepresentsθ\theta,gg,mm, orvv.

By the update rules of the first and second moment estimates, we have

mi,t=(1−β1)​gi,t+β1​mi,t−1,\displaystyle m_{i,t}=(1-\beta_{1})g_{i,t}+\beta_{1}m_{i,t-1},mi,0=0,\displaystyle\quad m_{i,0}=0,vi,t=(1−β2)​gi,t2+β2​vi,t−1,\displaystyle v_{i,t}=(1-\beta_{2})g^{2}_{i,t}+\beta_{2}v_{i,t-1},vi,0=0.\displaystyle\quad v_{i,0}=0.So by mathematical induction, for any1≤i≤d1\leq i\leq d, we have

mi,t=(1−β1)​∑s=1tβ1t−s​gi,sm_{i,t}=(1-\beta_{1})\sum_{s=1}^{t}\beta_{1}^{t-s}g_{i,s}(7)and

vi,t=(1−β2)​∑s=1tβ2t−s​gi,s2.v_{i,t}=(1-\beta_{2})\sum_{s=1}^{t}\beta_{2}^{t-s}g^{2}_{i,s}.(8)We will frequently use Equation (7) and (8) in the proofs of the subsequent lemmas and theorems.

Lemma 3.

For the full-batch version of MoFO with hyperparameters satisfyingβ1<β2<1\beta_{1}<\sqrt{\beta_{2}}<1,ϵ=0\epsilon=0, it holds that

|θi,t−θi,t−1|≤11−β2​(1−β1/β2)⋅ηt⋅FLTα​(mt)i,for any coordinate​1≤i≤d.|\theta_{i,t}-\theta_{i,t-1}|\leq\frac{1}{\sqrt{1-\beta_{2}}(1-\beta_{1}/\sqrt{\beta_{2}})}\cdot\eta_{t}\cdot\texttt{FLT}_{\alpha}(m_{t})_{i},\quad\text{for any coordinate }1\leq i\leq d.Moreover, it holds that

‖θt−θt−1‖2≤C​ηt,\|\theta_{t}-\theta_{t-1}\|_{2}\leq C\eta_{t},whereC=d​α+B1−β2​(1−β1/β2)C=\frac{\sqrt{d\alpha+B}}{\sqrt{1-\beta_{2}}(1-\beta_{1}/\sqrt{\beta_{2}})}.

Proof.

When theii-th entry is not in our filter at iterationtt, i.e.FLTα​(mt)i=0\texttt{FLT}_{\alpha}(m_{t})_{i}=0, we haveθi,t=θi,t−1\theta_{i,t}=\theta_{i,t-1}. Then

|θi,t−θi,t−1|=0=11−β2​(1−β1/β2)⋅ηt⋅FLTα​(mt)i.|\theta_{i,t}-\theta_{i,t-1}|=0=\frac{1}{\sqrt{1-\beta_{2}}(1-\beta_{1}/\sqrt{\beta_{2}})}\cdot\eta_{t}\cdot\texttt{FLT}_{\alpha}(m_{t})_{i}. When theii-th entry is in our filter, i.e.FLTα​(mt)i=1\texttt{FLT}_{\alpha}(m_{t})_{i}=1, by the weight updating rule of MoFO, we haveθi,t−θi,t−1=−ηt​m^i,t/v^i,t\theta_{i,t}-\theta_{i,t-1}=-\eta_{t}\hat{m}_{i,t}/\sqrt{\hat{v}_{i,t}}. We first analyzemi,tm_{i,t}andvi,tv_{i,t}.

By Equation (7) and (8), the first/second momentsmi,t,vi,tm_{i,t},v_{i,t}satisfy

|mi,t|\displaystyle|m_{i,t}|≤(1−β1)​∑s=1tβ1t−s​|gi,s|,\displaystyle\leq(1-\beta_{1})\sum_{s=1}^{t}\beta_{1}^{t-s}|g_{i,s}|,vi,t\displaystyle v_{i,t}=(1−β2)​∑s=1tβ2t−s​gi,s2≥(1−β2)​β2t−s​gi,s2,for any​1≤s≤t.\displaystyle=(1-\beta_{2})\sum_{s=1}^{t}\beta_{2}^{t-s}g^{2}_{i,s}\geq(1-\beta_{2})\beta_{2}^{t-s}g^{2}_{i,s},\quad\text{for any }1\leq s\leq t.So we get

|θi,t−θi,t−1|\displaystyle|\theta_{i,t}-\theta_{i,t-1}|=|−ηt​m^i,tv^i,t|=ηt​1−β2t1−β1t​|mi,t|/vi,t\displaystyle=\left|-\eta_{t}\frac{\hat{m}_{i,t}}{\sqrt{\hat{v}_{i,t}}}\right|=\eta_{t}\frac{\sqrt{1-\beta_{2}^{t}}}{1-\beta_{1}^{t}}|m_{i,t}|/\sqrt{v_{i,t}}≤ηt​1−β2t1−β1t​∑s=1t(1−β1)​β1t−s​|gi,s|(1−β2)​β2t−s​|gi,s|=ηt​1−β11−β1t​1−β2t1−β2​∑s=1t(β1/β2)t−s\displaystyle\leq\eta_{t}\frac{\sqrt{1-\beta_{2}^{t}}}{1-\beta_{1}^{t}}\sum_{s=1}^{t}\frac{(1-\beta_{1})\beta_{1}^{t-s}|g_{i,s}|}{\sqrt{(1-\beta_{2})\beta_{2}^{t-s}}|g_{i,s}|}=\eta_{t}\frac{1-\beta_{1}}{1-\beta_{1}^{t}}\sqrt{\frac{1-\beta_{2}^{t}}{1-\beta_{2}}}\sum_{s=1}^{t}(\beta_{1}/\sqrt{\beta_{2}})^{t-s}≤ηt1−β2​∑s=0t−1(β1/β2)s\displaystyle\leq\frac{\eta_{t}}{\sqrt{1-\beta_{2}}}\sum_{s=0}^{t-1}(\beta_{1}/\sqrt{\beta_{2}})^{s}≤ηt1−β2​(1−β1/β2).\displaystyle\leq\frac{\eta_{t}}{\sqrt{1-\beta_{2}}(1-\beta_{1}/\sqrt{\beta_{2}})}.Here, the last inequality holds because of the assumptionβ1<β2<1\beta_{1}<\sqrt{\beta_{2}}<1.

The parameter vector is partitioned intoBBblocks with sizes{dk}k=1B\{d_{k}\}_{k=1}^{B}and MoFO actually choose⌈dk​α⌉\lceil d_{k}\alpha\rceilentries to update in each partkkof parameters. Then for anyz∈ℝdz\in{\mathbb{R}}^{d}, we have

#​{1≤i≤d:FLTα​(z)i=1}=∑k=1B⌈dk​α⌉≤∑k=1B(dk​α+1)=d​α+B.\#\{1\leq i\leq d:\texttt{FLT}_{\alpha}(z)_{i}=1\}=\sum_{k=1}^{B}\lceil d_{k}\alpha\rceil\leq\sum_{k=1}^{B}(d_{k}\alpha+1)=d\alpha+B.Then for theL2L_{2}-norm of the parameter update, we have

‖θt−θt−1‖2\displaystyle\|\theta_{t}-\theta_{t-1}\|_{2}=(∑k=1d|θi,t−θi,t−1|2⋅FLTα​(mt)i)12\displaystyle=\left(\sum_{k=1}^{d}|\theta_{i,t}-\theta_{i,t-1}|^{2}\cdot\texttt{FLT}_{\alpha}(m_{t})_{i}\right)^{\frac{1}{2}}≤(ηt2(1−β2​(1−β1/β2))2⋅#​{1≤i≤d:FLTα​(z)i=1})12\displaystyle\leq\left(\frac{\eta_{t}^{2}}{(\sqrt{1-\beta_{2}}(1-\beta_{1}/\sqrt{\beta_{2}}))^{2}}\cdot\#\{1\leq i\leq d:\texttt{FLT}_{\alpha}(z)_{i}=1\}\right)^{\frac{1}{2}}≤d​α+B1−β2​(1−β1/β2)⋅ηt\displaystyle\leq\frac{\sqrt{d\alpha+B}}{\sqrt{1-\beta_{2}}(1-\beta_{1}/\sqrt{\beta_{2}})}\cdot\eta_{t}=C​ηt.\displaystyle=C\eta_{t}. ∎

Lemma 4.

Suppose that the gradient∇ℒ\nabla\mathcal{L}is Lipschitz continuous with constantLL. Suppose that the full-batch version of MoFO has the hyperparameters satisfyingβ1<β2<1\beta_{1}<\sqrt{\beta_{2}}<1,ϵ=0\epsilon=0and the learning rate scheduleηt=η/t\eta_{t}=\eta/\sqrt{t}. For any iteration stepst≥s≥1t\geq s\geq 1and any coordinateii, it holds that

|gi,t−gi,s|≤‖gt−gs‖2≤2​2​L​C​η​(t−s)t,|g_{i,t}-g_{i,s}|\leq\|g_{t}-g_{s}\|_{2}\leq\frac{2\sqrt{2}LC\eta(t-s)}{\sqrt{t}},whereC=d​α+B1−β2​(1−β1/β2)C=\frac{\sqrt{d\alpha+B}}{\sqrt{1-\beta_{2}}(1-\beta_{1}/\sqrt{\beta_{2}})}.

Proof.

Fix the iterations stepsttandsswitht≥s≥1t\geq s\geq 1. Since∇ℒ\nabla\mathcal{L}has Lipschitz constantLL, the gradient difference between the stepttandsssatisfies

|gi,t−gi,s|≤‖gt−gs‖2=‖∇ℒ​(θt−1)−∇ℒ​(θs−1)‖2≤L​‖θt−1−θs−1‖2.|g_{i,t}-g_{i,s}|\leq\|g_{t}-g_{s}\|_{2}=\|\nabla\mathcal{L}(\theta_{t-1})-\nabla\mathcal{L}(\theta_{s-1})\|_{2}\leq L\|\theta_{t-1}-\theta_{s-1}\|_{2}.(9) By Lemma3, for anyt>s≥1t>s\geq 1, the parameter difference satisfies

‖θt−1−θs−1‖2\displaystyle\|\theta_{t-1}-\theta_{s-1}\|_{2}≤∑u=st−1‖θu−θu−1‖2≤C​∑u=st−1ηu\displaystyle\leq\sum_{u=s}^{t-1}\|\theta_{u}-\theta_{u-1}\|_{2}\leq C\sum_{u=s}^{t-1}\eta_{u}≤C​η​∑u=st−11u≤C​η​∑u=st−12u−1+u≤2​C​η​∑u=st−1(u−u−1)\displaystyle\leq C\eta\sum_{u=s}^{t-1}\frac{1}{\sqrt{u}}\leq C\eta\sum_{u=s}^{t-1}\frac{2}{\sqrt{u-1}+\sqrt{u}}\leq 2C\eta\sum_{u=s}^{t-1}(\sqrt{u}-\sqrt{u-1})=2​C​η​(t−1−s−1)=2​C​η​(t−s)t−1+s−1\displaystyle=2C\eta(\sqrt{t-1}-\sqrt{s-1})=\frac{2C\eta(t-s)}{\sqrt{t-1}+\sqrt{s-1}}≤2​C​η​(t−s)t−1≤2​C​η​(t−s)t/2\displaystyle\leq\frac{2C\eta(t-s)}{\sqrt{t-1}}\leq\frac{2C\eta(t-s)}{\sqrt{t/2}}=2​2​C​η​(t−s)t.\displaystyle=\frac{2\sqrt{2}C\eta(t-s)}{\sqrt{t}}.Whent=s>1t=s>1, it is obvious that

‖θt−1−θs−1‖2=0≤2​2​C​η​(t−s)t.\|\theta_{t-1}-\theta_{s-1}\|_{2}=0\leq\frac{2\sqrt{2}C\eta(t-s)}{\sqrt{t}}.Combining it with (9), for anyt≥s≥1t\geq s\geq 1, we have

|gi,t−gi,s|≤‖gt−gs‖2≤2​2​L​C​η​(t−s)t.|g_{i,t}-g_{i,s}|\leq\|g_{t}-g_{s}\|_{2}\leq\frac{2\sqrt{2}LC\eta(t-s)}{\sqrt{t}}. ∎

Lemma 5.

Under the assumptions in Lemma4, for any iteration stept≥1t\geq 1and any coordinateii, it holds that

gi,t​m^i,tv^i,t≥1−β2​(|gi,t|−[2​2​β1(1−β1)2+41−β2]​L​C​ηt).g_{i,t}\frac{\hat{m}_{i,t}}{\sqrt{\hat{v}_{i,t}}}\geq\sqrt{1-\beta_{2}}\left(|g_{i,t}|-\left[\frac{2\sqrt{2}\beta_{1}}{(1-\beta_{1})^{2}}+\frac{4}{1-\beta_{2}}\right]\frac{LC\eta}{\sqrt{t}}\right).

Proof.

By Lemma4, we get

gi,t​gi,s=gi,t2−gi,t​(gi,t−gi,s)≥gi,t2−|gi,t|⋅|gi,t−gi,s|≥gi,t2−2​2​L​C​η​(t−s)t​|gi,t|.\displaystyle g_{i,t}g_{i,s}=g_{i,t}^{2}-g_{i,t}(g_{i,t}-g_{i,s})\geq g_{i,t}^{2}-|g_{i,t}|\cdot|g_{i,t}-g_{i,s}|\geq g_{i,t}^{2}-\frac{2\sqrt{2}LC\eta(t-s)}{\sqrt{t}}|g_{i,t}|.Then for the product of gradient and momentum, we have

gi,t​mi,t\displaystyle g_{i,t}m_{i,t}=(1−β1)​∑s=1tβ1t−s​gi,t​gi,s\displaystyle=(1-\beta_{1})\sum_{s=1}^{t}\beta_{1}^{t-s}g_{i,t}g_{i,s}(10)≥gi,t2⋅(1−β1)​∑s=1tβ1t−s−2​2​L​C​ηt​|gi,t|⋅(1−β1)​∑s=1tβ1t−s⋅(t−s)\displaystyle\geq g_{i,t}^{2}\cdot(1-\beta_{1})\sum_{s=1}^{t}\beta_{1}^{t-s}-\frac{2\sqrt{2}LC\eta}{\sqrt{t}}|g_{i,t}|\cdot(1-\beta_{1})\sum_{s=1}^{t}\beta_{1}^{t-s}\cdot(t-s)≥gi,t2⋅(1−β1)​∑s=0t−1β1s−2​2​L​C​ηt​|gi,t|⋅(1−β1)​∑s=1t−1s​β1s.\displaystyle\geq g_{i,t}^{2}\cdot(1-\beta_{1})\sum_{s=0}^{t-1}\beta_{1}^{s}-\frac{2\sqrt{2}LC\eta}{\sqrt{t}}|g_{i,t}|\cdot(1-\beta_{1})\sum_{s=1}^{t-1}s\beta_{1}^{s}.Since we have

∑s=0t−1β1s=1−β1t1−β1,∑s=1t−1s​β1s−1≤∑s=1∞s​β1s−1=dd​β1​(∑s=1∞β1s)=dd​β1​(β11−β1)=1(1−β1)2,\sum_{s=0}^{t-1}\beta_{1}^{s}=\frac{1-\beta_{1}^{t}}{1-\beta_{1}},\quad\sum_{s=1}^{t-1}s\beta_{1}^{s-1}\leq\sum_{s=1}^{\infty}s\beta_{1}^{s-1}=\frac{d}{d\beta_{1}}\left(\sum_{s=1}^{\infty}\beta_{1}^{s}\right)=\frac{d}{d\beta_{1}}\left(\frac{\beta_{1}}{1-\beta_{1}}\right)=\frac{1}{(1-\beta_{1})^{2}},(11)it holds that

gi,t​mi,t≥RHS of​(10)≥(1−β1t)​gi,t2−2​2​β1​L​C​η(1−β1)​t​|gi,t|.g_{i,t}m_{i,t}\geq\text{RHS of }(\ref{09031709})\geq(1-\beta_{1}^{t})g_{i,t}^{2}-\frac{2\sqrt{2}\beta_{1}LC\eta}{(1-\beta_{1})\sqrt{t}}|g_{i,t}|.(12) For the second momentvi,tv_{i,t}, we have

vi,t\displaystyle v_{i,t}=(1−β2)​∑s=1tβ2t−s​gi,s2≤(1−β2)​∑s=1tβ2t−s​(|gi,t|+|gi,s−gi,t|)2\displaystyle=(1-\beta_{2})\sum_{s=1}^{t}\beta_{2}^{t-s}g_{i,s}^{2}\leq(1-\beta_{2})\sum_{s=1}^{t}\beta_{2}^{t-s}(|g_{i,t}|+|g_{i,s}-g_{i,t}|)^{2}(13)≤(1−β2)​∑s=1tβ2t−s​(|gi,t|+2​2​L​C​η​(t−s)t)2=(1−β2)​∑s=0t−1β2s​(|gi,t|+2​2​L​C​η​st)2\displaystyle\leq(1-\beta_{2})\sum_{s=1}^{t}\beta_{2}^{t-s}\left(|g_{i,t}|+\frac{2\sqrt{2}LC\eta(t-s)}{\sqrt{t}}\right)^{2}=(1-\beta_{2})\sum_{s=0}^{t-1}\beta_{2}^{s}\left(|g_{i,t}|+\frac{2\sqrt{2}LC\eta s}{\sqrt{t}}\right)^{2}=|gi,t|2⋅(1−β2)​(∑s=0t−1β2s)+|gi,t|⋅4​2​L​C​ηt​(1−β2)​(∑s=1t−1s​β2s)\displaystyle=|g_{i,t}|^{2}\cdot(1-\beta_{2})\left(\sum_{s=0}^{t-1}\beta_{2}^{s}\right)+|g_{i,t}|\cdot\frac{4\sqrt{2}LC\eta}{\sqrt{t}}(1-\beta_{2})\left(\sum_{s=1}^{t-1}s\beta_{2}^{s}\right)+8​L2​C2​η2t​(1−β2)​(∑s=1t−1s2​β2s).\displaystyle\qquad+\frac{8L^{2}C^{2}\eta^{2}}{t}(1-\beta_{2})\left(\sum_{s=1}^{t-1}s^{2}\beta_{2}^{s}\right).Since we have

∑s=0t−1β2s\displaystyle\sum_{s=0}^{t-1}\beta_{2}^{s}=1−β2t1−β2≤11−β2,\displaystyle=\frac{1-\beta_{2}^{t}}{1-\beta_{2}}\leq\frac{1}{1-\beta_{2}},∑s=0t−1s​β2s−1\displaystyle\sum_{s=0}^{t-1}s\beta_{2}^{s-1}≤∑s=0∞s​β2s−1=dd​β2​(∑s=0∞β2s)=dd​β2​(11−β2)=1(1−β2)2,\displaystyle\leq\sum_{s=0}^{\infty}s\beta_{2}^{s-1}=\frac{d}{d\beta_{2}}\left(\sum_{s=0}^{\infty}\beta_{2}^{s}\right)=\frac{d}{d\beta_{2}}\left(\frac{1}{1-\beta_{2}}\right)=\frac{1}{(1-\beta_{2})^{2}},∑s=0t−1s2​β2s−1\displaystyle\sum_{s=0}^{t-1}s^{2}\beta_{2}^{s-1}≤∑s=0∞s2​β2s−1=β2​(∑s=0∞s​(s−1)​β2s−2)+∑s=0∞s​β2s−1\displaystyle\leq\sum_{s=0}^{\infty}s^{2}\beta_{2}^{s-1}=\beta_{2}\left(\sum_{s=0}^{\infty}s(s-1)\beta_{2}^{s-2}\right)+\sum_{s=0}^{\infty}s\beta_{2}^{s-1}=β2⋅d2d​β22​(∑s=0∞β2s)+1(1−β2)2=β2⋅d2d​β22​(11−β2)+1(1−β2)2\displaystyle=\beta_{2}\cdot\frac{d^{2}}{d\beta_{2}^{2}}\left(\sum_{s=0}^{\infty}\beta_{2}^{s}\right)+\frac{1}{(1-\beta_{2})^{2}}=\beta_{2}\cdot\frac{d^{2}}{d\beta_{2}^{2}}\left(\frac{1}{1-\beta_{2}}\right)+\frac{1}{(1-\beta_{2})^{2}}=2​β2(1−β2)3+1(1−β2)2\displaystyle=\frac{2\beta_{2}}{(1-\beta_{2})^{3}}+\frac{1}{(1-\beta_{2})^{2}}=1+β2(1−β2)3,\displaystyle=\frac{1+\beta_{2}}{(1-\beta_{2})^{3}},it holds that

vi,t≤RHS of​(13)\displaystyle v_{i,t}\leq\text{RHS of }(\ref{09031825})≤|gi,t|2+|gi,t|⋅4​2​β2​L​C​η(1−β2)​t+8​(1+β2)​β2​L2​C2​η2(1−β2)2​t\displaystyle\leq|g_{i,t}|^{2}+|g_{i,t}|\cdot\frac{4\sqrt{2}\beta_{2}LC\eta}{(1-\beta_{2})\sqrt{t}}+\frac{8(1+\beta_{2})\beta_{2}L^{2}C^{2}\eta^{2}}{(1-\beta_{2})^{2}t}≤|gi,t|2+|gi,t|⋅8​L​C​η(1−β2)​t+16​L2​C2​η2(1−β2)2​t\displaystyle\leq|g_{i,t}|^{2}+|g_{i,t}|\cdot\frac{8LC\eta}{(1-\beta_{2})\sqrt{t}}+\frac{16L^{2}C^{2}\eta^{2}}{(1-\beta_{2})^{2}t}=(|gi,t|+4​L​C​η(1−β2)​t)2.\displaystyle=\left(|g_{i,t}|+\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}\right)^{2}.Thus, we get

vi,t≤|gi,t|+4​L​C​η(1−β2)​t.\sqrt{v_{i,t}}\leq|g_{i,t}|+\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}. Recalling (12), we have

gi,t​mi,t\displaystyle g_{i,t}m_{i,t}≥(1−β1t)​(|gi,t|+4​L​C​η(1−β2)​t)​(|gi,t|−2​2​β1​L​C​η(1−β1t)​(1−β1)​t−4​L​C​η(1−β2)​t)\displaystyle\geq(1-\beta_{1}^{t})\left(|g_{i,t}|+\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}\right)\left(|g_{i,t}|-\frac{2\sqrt{2}\beta_{1}LC\eta}{(1-\beta_{1}^{t})(1-\beta_{1})\sqrt{t}}-\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}\right)+(1−β1t)⋅4​L​C​η(1−β2)​t​(2​2​β1​L​C​η(1−β1t)​(1−β1)​t+4​L​C​η(1−β2)​t)\displaystyle\quad+(1-\beta_{1}^{t})\cdot\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}\left(\frac{2\sqrt{2}\beta_{1}LC\eta}{(1-\beta_{1}^{t})(1-\beta_{1})\sqrt{t}}+\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}\right)≥(1−β1t)​(|gi,t|+4​L​C​η(1−β2)​t)​(|gi,t|−2​2​β1​L​C​η(1−β1t)​(1−β1)​t−4​L​C​η(1−β2)​t)\displaystyle\geq(1-\beta_{1}^{t})\left(|g_{i,t}|+\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}\right)\left(|g_{i,t}|-\frac{2\sqrt{2}\beta_{1}LC\eta}{(1-\beta_{1}^{t})(1-\beta_{1})\sqrt{t}}-\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}\right)≥(1−β1t)​vi,t​(|gi,t|−2​2​β1​L​C​η(1−β1t)​(1−β1)​t−4​L​C​η(1−β2)​t).\displaystyle\geq(1-\beta_{1}^{t})\sqrt{v_{i,t}}\left(|g_{i,t}|-\frac{2\sqrt{2}\beta_{1}LC\eta}{(1-\beta_{1}^{t})(1-\beta_{1})\sqrt{t}}-\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}\right). Therefore, for the bias-corrected first/second momentsm^i,t\hat{m}_{i,t}andv^i,t\hat{v}_{i,t}, it holds that

gi,t​m^i,tv^i,t=1−β2t1−β1t​gi,t​mi,tvi,t\displaystyle g_{i,t}\frac{\hat{m}_{i,t}}{\sqrt{\hat{v}_{i,t}}}=\frac{\sqrt{1-\beta_{2}^{t}}}{1-\beta_{1}^{t}}g_{i,t}\frac{m_{i,t}}{\sqrt{v_{i,t}}}≥1−β2t​(|gi,t|−2​2​β1​L​C​η(1−β1t)​(1−β1)​t−4​L​C​η(1−β2)​t)\displaystyle\geq\sqrt{1-\beta_{2}^{t}}\left(|g_{i,t}|-\frac{2\sqrt{2}\beta_{1}LC\eta}{(1-\beta_{1}^{t})(1-\beta_{1})\sqrt{t}}-\frac{4LC\eta}{(1-\beta_{2})\sqrt{t}}\right)≥1−β2​(|gi,t|−[2​2​β1(1−β1)2+41−β2]​L​C​ηt).\displaystyle\geq\sqrt{1-\beta_{2}}\left(|g_{i,t}|-\left[\frac{2\sqrt{2}\beta_{1}}{(1-\beta_{1})^{2}}+\frac{4}{1-\beta_{2}}\right]\frac{LC\eta}{\sqrt{t}}\right). ∎

Lemma 6.

Under the assumptions in Lemma4, for any iteration stept≥1t\geq 1and any coordinateii, it holds that

‖m^t−gt‖1≤2​2​β1​d​L​C​η(1−β1)2​t.\left\|\hat{m}_{t}-g_{t}\right\|_{1}\leq\frac{2\sqrt{2}\beta_{1}\sqrt{d}LC\eta}{(1-\beta_{1})^{2}\sqrt{t}}.

Proof.

Recalling the calculation of the momentummtm_{t}in (7), we get

mt=(1−β1)​∑s=1tβ1t−s​gs,\displaystyle m_{t}=(1-\beta_{1})\sum_{s=1}^{t}\beta_{1}^{t-s}g_{s},and

mt−(1−β1t)​gt=(1−β1)​∑s=1tβ1t−s​(gt−gs).\displaystyle m_{t}-(1-\beta_{1}^{t})g_{t}=(1-\beta_{1})\sum_{s=1}^{t}\beta_{1}^{t-s}(g_{t}-g_{s}). By Lemma4and Equation (11) in the proof of Lemma5, we get

‖m^t−gt‖2=‖mt1−β1t−gt‖2\displaystyle\|\hat{m}_{t}-g_{t}\|_{2}=\left\|\frac{m_{t}}{1-\beta_{1}^{t}}-g_{t}\right\|_{2}≤1−β11−β1t​∑s=1tβ1t−s​‖gt−gs‖2≤∑s=1tβ1t−s​‖gt−gs‖2\displaystyle\leq\frac{1-\beta_{1}}{1-\beta_{1}^{t}}\sum_{s=1}^{t}\beta_{1}^{t-s}\|g_{t}-g_{s}\|_{2}\leq\sum_{s=1}^{t}\beta_{1}^{t-s}\|g_{t}-g_{s}\|_{2}≤2​2​L​C​ηt​∑s=1tβ1t−s​(t−s)=2​2​L​C​ηt​∑s=0t−1s​β1s\displaystyle\leq\frac{2\sqrt{2}LC\eta}{\sqrt{t}}\sum_{s=1}^{t}\beta_{1}^{t-s}(t-s)=\frac{2\sqrt{2}LC\eta}{\sqrt{t}}\sum_{s=0}^{t-1}s\beta_{1}^{s}≤2​2​β1​L​C​η(1−β1)2​t.\displaystyle\leq\frac{2\sqrt{2}\beta_{1}LC\eta}{(1-\beta_{1})^{2}\sqrt{t}}. By Cauchy-Schwarz’s inequality, we have

‖m^t−gt‖1≤d​‖m^t−gt‖2≤2​2​β1​d​L​C​η(1−β1)2​t.\|\hat{m}_{t}-g_{t}\|_{1}\leq\sqrt{d}\|\hat{m}_{t}-g_{t}\|_{2}\leq\frac{2\sqrt{2}\beta_{1}\sqrt{d}LC\eta}{(1-\beta_{1})^{2}\sqrt{t}}. ∎

Now we will complete the proof of Theorem1.

Proof of Theorem1.

By the descent lemma, since∇ℒ\nabla\mathcal{L}is Lipschitz with constantLL, we have

ℒ​(θt)−ℒ​(θt−1)\displaystyle\mathcal{L}(\theta_{t})-\mathcal{L}(\theta_{t-1})≤∇ℒ​(θt−1)⊤​(θt−θt−1)+L2​‖θt−θt−1‖22\displaystyle\leq\nabla\mathcal{L}(\theta_{t-1})^{\top}(\theta_{t}-\theta_{t-1})+\frac{L}{2}\|\theta_{t}-\theta_{t-1}\|_{2}^{2}(14)≤gt⊤​(θt−θt−1)+L2​‖θt−θt−1‖22.\displaystyle\leq g_{t}^{\top}(\theta_{t}-\theta_{t-1})+\frac{L}{2}\|\theta_{t}-\theta_{t-1}\|_{2}^{2}.By Lemma3and Lemma5, we have

ℒ​(θt)−ℒ​(θt−1)≤RHS of​(14)≤−ηt​(∑i=1dgi,t​m^i,tv^i,t⋅FLTα​(mt)i)+L​C2​ηt22\displaystyle\quad\ \mathcal{L}(\theta_{t})-\mathcal{L}(\theta_{t-1})\leq\text{RHS of }(\ref{09042001})\leq-\eta_{t}\left(\sum_{i=1}^{d}g_{i,t}\frac{\hat{m}_{i,t}}{\sqrt{\hat{v}_{i,t}}}\cdot\texttt{FLT}_{\alpha}(m_{t})_{i}\right)+\frac{LC^{2}\eta_{t}^{2}}{2}(15)≤L​C2​η22​t−ηt​∑i=1d1−β2​(|gi,t|−[2​2​β1(1−β1)2+41−β2]​L​C​ηt)⋅FLTα​(mt)i\displaystyle\leq\frac{LC^{2}\eta^{2}}{2t}-\frac{\eta}{\sqrt{t}}\sum_{i=1}^{d}\sqrt{1-\beta_{2}}\left(|g_{i,t}|-\left[\frac{2\sqrt{2}\beta_{1}}{(1-\beta_{1})^{2}}+\frac{4}{1-\beta_{2}}\right]\frac{LC\eta}{\sqrt{t}}\right)\cdot\texttt{FLT}_{\alpha}(m_{t})_{i}=−1−β2⋅ηt​‖gt⊙FLTα​(mt)‖1+[2​2​β1​1−β2(1−β1)2+41−β2+C2]​L​C​η2t⋅‖FLTα​(mt)‖1\displaystyle=-\frac{\sqrt{1-\beta_{2}}\cdot\eta}{\sqrt{t}}\|g_{t}\odot\texttt{FLT}_{\alpha}(m_{t})\|_{1}+\left[\frac{2\sqrt{2}\beta_{1}\sqrt{1-\beta_{2}}}{(1-\beta_{1})^{2}}+\frac{4}{\sqrt{1-\beta_{2}}}+\frac{C}{2}\right]\frac{LC\eta^{2}}{t}\cdot\|\texttt{FLT}_{\alpha}(m_{t})\|_{1}≤−1−β2⋅ηt​‖gt⊙FLTα​(mt)‖1+[2​2​β1​1−β2(1−β1)2+41−β2+C2]​L​C​η2​(d​α+B)t.\displaystyle\leq-\frac{\sqrt{1-\beta_{2}}\cdot\eta}{\sqrt{t}}\|g_{t}\odot\texttt{FLT}_{\alpha}(m_{t})\|_{1}+\left[\frac{2\sqrt{2}\beta_{1}\sqrt{1-\beta_{2}}}{(1-\beta_{1})^{2}}+\frac{4}{\sqrt{1-\beta_{2}}}+\frac{C}{2}\right]\frac{LC\eta^{2}(d\alpha+B)}{t}. By Lemma1and Lemma6, we have

‖gt⊙FLTα​(gt)‖1−‖gt⊙FLTα​(mt)‖1\displaystyle\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\|_{1}-\|g_{t}\odot\texttt{FLT}_{\alpha}(m_{t})\|_{1}=‖gt⊙FLTα​(gt)‖1−‖gt⊙FLTα​(mt1−β1t)‖1\displaystyle=\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\|_{1}-\left\|g_{t}\odot\texttt{FLT}_{\alpha}\left(\frac{m_{t}}{1-\beta_{1}^{t}}\right)\right\|_{1}=‖gt⊙FLTα​(m^t)‖1\displaystyle=\|g_{t}\odot\texttt{FLT}_{\alpha}(\hat{m}_{t})\|_{1}≤2​‖gt−m^t‖1\displaystyle\leq 2\left\|g_{t}-\hat{m}_{t}\right\|_{1}≤4​2​β1​d​L​C​η(1−β2)2​t.\displaystyle\leq\frac{4\sqrt{2}\beta_{1}\sqrt{d}LC\eta}{(1-\beta_{2})^{2}\sqrt{t}}. Thus,

ℒ​(θt)−ℒ​(θt−1)≤RHS of​(15)\displaystyle\quad\ \mathcal{L}(\theta_{t})-\mathcal{L}(\theta_{t-1})\leq\text{RHS of }(\ref{09051021})(16)≤−1−β2⋅ηt​‖gt⊙FLTα​(gt)‖1+[2​2​β1​1−β2(1−β1)2+41−β2+C2]​L​C​η2​(d​α+B)t\displaystyle\leq-\frac{\sqrt{1-\beta_{2}}\cdot\eta}{\sqrt{t}}\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\|_{1}+\left[\frac{2\sqrt{2}\beta_{1}\sqrt{1-\beta_{2}}}{(1-\beta_{1})^{2}}+\frac{4}{\sqrt{1-\beta_{2}}}+\frac{C}{2}\right]\frac{LC\eta^{2}(d\alpha+B)}{t}+4​2​β1​d​L​C​η2(1−β2)32​t\displaystyle\qquad\qquad\qquad\qquad\qquad+\frac{4\sqrt{2}\beta_{1}\sqrt{d}LC\eta^{2}}{(1-\beta_{2})^{\frac{3}{2}}t}=−C1t​‖gt‖1,top-​α+C2t≤−C1t​min1≤t≤T⁡‖gt‖1,top-​α+C2t,\displaystyle=-\frac{C_{1}}{\sqrt{t}}\|{g_{t}}\|_{1,\text{top-}\alpha}+\frac{C_{2}}{t}\leq-\frac{C_{1}}{\sqrt{t}}\min_{1\leq t\leq T}\|{g_{t}}\|_{1,\text{top-}\alpha}+\frac{C_{2}}{t},where

C1\displaystyle C_{1}=1−β2⋅η,\displaystyle=\sqrt{1-\beta_{2}}\cdot\eta,C2\displaystyle C_{2}=L​C​η2⋅{[2​2​β1​1−β2(1−β1)2+41−β2+C2]​(d​α+B)+4​2​β1​d(1−β2)32}.\displaystyle=LC\eta^{2}\cdot\left\{\left[\frac{2\sqrt{2}\beta_{1}\sqrt{1-\beta_{2}}}{(1-\beta_{1})^{2}}+\frac{4}{\sqrt{1-\beta_{2}}}+\frac{C}{2}\right](d\alpha+B)+\frac{4\sqrt{2}\beta_{1}\sqrt{d}}{(1-\beta_{2})^{\frac{3}{2}}}\right\}. Taking the summation of (15) from11toTT, we get

ℒ∗−ℒ​(θ0)\displaystyle\mathcal{L}^{*}-\mathcal{L}(\theta_{0})≤ℒ​(θT)−ℒ​(θ0)=∑t=1Tℒ​(θt)−ℒ​(θt−1)\displaystyle\leq\mathcal{L}(\theta_{T})-\mathcal{L}(\theta_{0})=\sum_{t=1}^{T}\mathcal{L}(\theta_{t})-\mathcal{L}(\theta_{t-1})≤−C1​(∑t=1T1t)⋅min1≤t≤T⁡‖gt⊙FLTα​(gt)‖1+C2​∑t=1T1t.\displaystyle\leq-C_{1}\left(\sum_{t=1}^{T}\frac{1}{\sqrt{t}}\right)\cdot\min_{1\leq t\leq T}\|g_{t}\odot\texttt{FLT}_{\alpha}(g_{t})\|_{1}+C_{2}\sum_{t=1}^{T}\frac{1}{t}. Since

∑t=1T1t\displaystyle\sum_{t=1}^{T}\frac{1}{\sqrt{t}}≥∑t=1T2t+t+1=∑t=1T2​(t+1−t)=2​(T+1−1),\displaystyle\geq\sum_{t=1}^{T}\frac{2}{\sqrt{t}+\sqrt{t+1}}=\sum_{t=1}^{T}2(\sqrt{t+1}-\sqrt{t})=2(\sqrt{T+1}-1),∑t=1T1t\displaystyle\sum_{t=1}^{T}\frac{1}{t}=1+∑t=1T−11t+1≤1+∑t=1T−1∫tt+11u​𝑑u≤1+∫1T1u​𝑑u=1+log⁡T,\displaystyle=1+\sum_{t=1}^{T-1}\frac{1}{t+1}\leq 1+\sum_{t=1}^{T-1}\int_{t}^{t+1}\frac{1}{u}\,du\leq 1+\int_{1}^{T}\frac{1}{u}\,du=1+\log T,we get

min0≤t≤T−1⁡‖∇ℒ​(θt)‖1,top-​α\displaystyle\min_{0\leq t\leq T-1}\|{\nabla\mathcal{L}(\theta_{t})}\|_{1,\text{top-}\alpha}=min1≤t≤T⁡‖gt‖1,top-​α=min1≤t≤T⁡‖gt⊙FLTα​(gt)‖1\displaystyle=\min_{1\leq t\leq T}\|{g_{t}}\|_{1,\text{top-}\alpha}=\min_{1\leq t\leq T}\|g_{t}\odot\texttt{FLT}_{\alpha}{(g_{t})}\|_{1}≤ℒ​(θ0)−ℒ∗+C2​∑t=1T1tC1​∑t=1T1t\displaystyle\leq\frac{\mathcal{L}(\theta_{0})-\mathcal{L}^{*}+C_{2}\sum_{t=1}^{T}\frac{1}{t}}{C_{1}\sum_{t=1}^{T}\frac{1}{\sqrt{t}}}≤ℒ​(θ0)−ℒ∗+C2​(1+log⁡T)2​C1​(T+1−1).\displaystyle\leq\frac{\mathcal{L}(\theta_{0})-\mathcal{L}^{*}+C_{2}(1+\log T)}{2C_{1}(\sqrt{T+1}-1)}. Thus, we have

min0≤t≤T−1⁡‖∇ℒ​(θt)‖1,top-​α=min1≤t≤T⁡‖ℒ​(θt)⊙FLTα​(ℒ​(θt))‖1=𝒪​(log⁡TT).\min_{0\leq t\leq T-1}\|{\nabla\mathcal{L}(\theta_{t})}\|_{1,\text{top-}\alpha}=\min_{1\leq t\leq T}\|\mathcal{L}(\theta_{t})\odot\texttt{FLT}_{\alpha}{(\mathcal{L}(\theta_{t}))}\|_{1}=\mathcal{O}\left(\frac{\log T}{\sqrt{T}}\right). By the relationship betweenL1,top-​αL_{1,\text{top-}\alpha}norm andLpL_{p}norm in Lemma2, for anyp∈[1,+∞]p\in[1,+\infty],

min0≤t≤T−1⁡‖∇ℒ​(θt)‖p≤1α​min0≤t≤T−1⁡‖∇ℒ​(θt)‖1,top-​α.\min_{0\leq t\leq T-1}\|\nabla\mathcal{L}(\theta_{t})\|_{p}\leq\frac{1}{\alpha}\min_{0\leq t\leq T-1}\|{\nabla\mathcal{L}(\theta_{t})}\|_{1,\text{top-}\alpha}.Therefore,

min0≤t≤T−1⁡‖∇ℒ​(θt)‖p=𝒪​(log⁡TT).\min_{0\leq t\leq T-1}\|\nabla\mathcal{L}(\theta_{t})\|_{p}=\mathcal{O}\left(\frac{\log T}{\sqrt{T}}\right). ∎

A.3Proof of Theorem2(Illustrative example: forgetting mitigation of MoFO)

Proof of Theorem2.

Let the setU:={θ∈ℝd:θi<bi/ai,∀1≤i≤d}U:=\{\theta\in\mathbb{R}^{d}:\theta_{i}<b_{i}/a_{i},\ \forall 1\leq i\leq d\}. We note that:

1. 1.The boundary ofUUis the subset ofS=∪i=1dSiS=\cup_{i=1}^{d}S_{i}, which is the collection all global minima of the fine-tuning loss.
2. 2.The pre-training stateθpretrain=(0,0,…,0)\theta_{\rm pretrain}=(0,0,\dots,0), which is also the starting point of fine-tuning, lies inUU.

We may as well assume that with proper learning rates, the parameterθ\thetaremains withinUUduring training, unless it converges to a minimum on the boundary. If it goes across the boundary at a certain iteration before converging, the learning rate can be adjusted to ensure that it remains withinUU. For anyθ∈U\theta\in Uand coordinatei∈{1,2,…,d}i\in\{1,2,\dots,d\}, we have

∂ℒ∂θi=2​ai​(ai​θi−bi)​∏j≠i(aj​θj−bj)2=2​ℒ​(θ)θi−biai<0.\displaystyle\frac{\partial\mathcal{L}}{\partial\theta_{i}}=2a_{i}(a_{i}\theta_{i}-b_{i})\prod_{j\neq i}(a_{j}\theta_{j}-b_{j})^{2}=\frac{2\mathcal{L}(\theta)}{\theta_{i}-\frac{b_{i}}{a_{i}}}<0.(17) For clarity of definition, we letθt\theta_{t}, represent the parameter at iterationtt. We letθi,t\theta_{i,t},gi,tg_{i,t},mi,tm_{i,t}denote theii-th coordinate ofθt\theta_{t},gtg_{t},mtm_{t}at iterationtt, respectively.

Analysis of MoFO.

We will use mathematical induction to show the following results:If MoFO selects the coordinatei0i_{0}at the first iteration step, it will always selecti0i_{0}at any iterationtt. Moreover, for any coordinatei≠i0i\neq i_{0}, we have

• •mi0,1≤mi,1≤0m_{i_{0},1}\leq m_{i,1}\leq 0, andmi0,t<mi,t≤0m_{i_{0},t}<m_{i,t}\leq 0for any iteration stept≥2t\geq 2.
• •0<bi0ai0−θi0,t<biai−θi,t0<\frac{b_{i_{0}}}{a_{i_{0}}}-\theta_{i_{0},t}<\frac{b_{i}}{a_{i}}-\theta_{i,t}for any iteration stept≥1t\geq 1if the algorithm has not reached the minimum.

Base case (1st iteration step).At the first iteration step, the momentumm1=(1−β1)​g1m_{1}=(1-\beta_{1})g_{1}. So for any1≤i≤d1\leq i\leq d, we get

m1,t=(1−β1)​g1,t=(1−β1)​∂ℒ∂θi|θpretrain<0.m_{1,t}=(1-\beta_{1})g_{1,t}=(1-\beta_{1})\frac{\partial\mathcal{L}}{\partial\theta_{i}}\bigg|_{\theta_{\rm pretrain}}<0.According to the momentum filtering mechanism of MoFO, we have

i0∈arg⁡max1≤i≤d⁡|mi,t|=arg⁡max1≤i≤d⁡|gi,t|=arg⁡max1≤i≤d⁡|∂ℒ∂θi|=arg⁡max1≤i≤d⁡1biai−θi,pretrain=arg⁡min1≤i≤d⁡{biai}.\displaystyle i_{0}\in\arg\max_{1\leq i\leq d}|m_{i,t}|=\arg\max_{1\leq i\leq d}|g_{i,t}|=\arg\max_{1\leq i\leq d}\bigg|\frac{\partial\mathcal{L}}{\partial\theta_{i}}\bigg|=\arg\max_{1\leq i\leq d}\frac{1}{\frac{b_{i}}{a_{i}}-\theta_{i,{\rm pretrain}}}=\arg\min_{1\leq i\leq d}\bigg\{\frac{b_{i}}{a_{i}}\bigg\}.Obviously, we getmi0,1≤mi,1≤0m_{i_{0},1}\leq m_{i,1}\leq 0for any coordinatei≠i0i\neq i_{0}.

The parameter updates at the first iteration are:

θi0,1\displaystyle\theta_{i_{0},1}=θi0,0−η1​1−β2​mi0,1(1−β1)​vi0,t=θi0,0−η1​gi,0|gi,0|=θi0,0−η1​sign​(∂ℒ∂θi​(θpretrain))>θi0,0,\displaystyle=\theta_{i_{0},0}-\frac{\eta_{1}\sqrt{1-\beta_{2}}m_{i_{0},1}}{(1-\beta_{1})\sqrt{v_{i_{0},t}}}=\theta_{i_{0},0}-\frac{\eta_{1}g_{i,0}}{|g_{i,0}|}=\theta_{i_{0},0}-\eta_{1}{\rm sign}\left(\frac{\partial\mathcal{L}}{\partial\theta_{i}}(\theta_{\rm pretrain})\right)>\theta_{i_{0},0},θi,1\displaystyle\theta_{i,1}=θi,0,∀i≠i0.\displaystyle=\theta_{i,0},\quad\forall i\neq i_{0}. If the algorithm has not converged at the first iteration, then we haveθi0,1<bi0/ai0\theta_{i_{0},1}<b_{i_{0}}/a_{i_{0}}. Moreover, for anyi≠i0i\neq i_{0},

0<biai−θi0,1<biai=biai−θi0,pretrain=biai−θi,pretrain=biai−θi,1.0<\frac{b_{i}}{a_{i}}-\theta_{i_{0},1}<\frac{b_{i}}{a_{i}}=\frac{b_{i}}{a_{i}}-\theta_{i_{0},{\rm pretrain}}=\frac{b_{i}}{a_{i}}-\theta_{i,{\rm pretrain}}=\frac{b_{i}}{a_{i}}-\theta_{i,1}. Induction step.Suppose that the induction hypothesis holds up to iterationtt. Then, for any coordinatei≠i0i\neq i_{0}, we have

• •mi0,t≤mi,t≤0m_{i_{0},t}\leq m_{i,t}\leq 0.
• •0≤bi0ai0−θi0,t<biai−θi,t0\leq\frac{b_{i_{0}}}{a_{i_{0}}}-\theta_{i_{0},t}<\frac{b_{i}}{a_{i}}-\theta_{i,t}.

So for the gradient,

gi0,t+1=2​ℒ​(θt)θi0,t−biai<2​ℒ​(θt)θi0,t−biai=gi,t+1<0,g_{i_{0},t+1}=\frac{2\mathcal{L}(\theta_{t})}{\theta_{i_{0},t}-\frac{b_{i}}{a_{i}}}<\frac{2\mathcal{L}(\theta_{t})}{\theta_{i_{0},t}-\frac{b_{i}}{a_{i}}}=g_{i,t+1}<0,and

mi0,t+1\displaystyle m_{i_{0},t+1}=β1​mi0,t+(1−β1)​gi0,t+1\displaystyle=\beta_{1}m_{i_{0},t}+(1-\beta_{1})g_{i_{0},t+1}<β1​mi,t+(1−β1)​gi,t+1=mi0,t+1<0.\displaystyle<\beta_{1}m_{i,t}+(1-\beta_{1})g_{i,t+1}=m_{i_{0},t+1}<0. Thus,i0i_{0}is the only coordinate inarg⁡max1≤i≤d⁡|mi,t+1|\arg\max_{1\leq i\leq d}|m_{i,t+1}|and MoFO still chooses the coordinatei0i_{0}to update. In addition,

θi0,t+1\displaystyle\theta_{i_{0},t+1}=θi0,t−ηt+1​m^i,t+1v^i,t+1>θi,t,\displaystyle=\theta_{i_{0},t}-\frac{\eta_{t+1}\hat{m}_{i,t+1}}{\sqrt{\hat{v}_{i,t+1}}}>\theta_{i,t},θi,t+1\displaystyle\theta_{i,t+1}=θi,t,∀i≠i0.\displaystyle=\theta_{i,t},\quad\forall i\neq i_{0}. If the algorithm has not converged at iteration stept+1t+1, then we haveθi0,t+1<bi0/ai0\theta_{i_{0},t+1}<b_{i_{0}}/a_{i_{0}}. Moreover, for anyi≠i0i\neq i_{0},

0<biai−θi0,t+1<biai−θi0,t≤biai−θi,t=biai−θi,t+1.0<\frac{b_{i}}{a_{i}}-\theta_{i_{0},t+1}<\frac{b_{i}}{a_{i}}-\theta_{i_{0},t}\leq\frac{b_{i}}{a_{i}}-\theta_{i,t}=\frac{b_{i}}{a_{i}}-\theta_{i,t+1}. Conclusion.MoFO consistently updatesθi0\theta_{i_{0}}and eventually converges toθMoFO∗=(0,…,0,bi0ai0,0,…,0)\theta^{*}_{\text{MoFO}}=(0,\dots,0,\frac{b_{i_{0}}}{a_{i_{0}}},0,\dots,0), with pre-training loss

ℒpretrain​(θMoFO∗)=bi022​ai02.\mathcal{L}_{\rm pretrain}(\theta^{*}_{\text{MoFO}})=\frac{b_{i_{0}}^{2}}{2a_{i_{0}}^{2}}.

Analysis of Adam.

Unlike MoFO, Adam updates all the parameters. By Inequality (17), we havegi,t<0g_{i,t}<0. By the momentum update rule of Adam:

mi,t+1=β1​mi,t+(1−β1)​gi,t+1,m_{i,t+1}=\beta_{1}m_{i,t}+(1-\beta_{1})g_{i,t+1},we get thatmi,t<0m_{i,t}<0for any1≤i≤d1\leq i\leq dand any iterationtt. Therefore, it holds for Adam that

θi,t+1=θi,t−ηt​ℒ​(θt)θi,t−biai>θi,t.\theta_{i,t+1}=\theta_{i,t}-\frac{\eta_{t}\mathcal{L}(\theta_{t})}{\theta_{i,t}-\frac{b_{i}}{a_{i}}}>\theta_{i,t}. Assuming that Adam converges toθAdam∗\theta^{*}_{\text{Adam}}, we have

• •θAdam,i∗>0\theta^{*}_{\text{Adam},i}>0for any1≤i≤d1\leq i\leq d,
• •There existsj0j_{0}such thatθGD,j0∗=bj0/aj0\theta^{*}_{\text{GD},j_{0}}=b_{j_{0}}/a_{j_{0}}.

Recall that at iteration 1, MoFO selects

i0∈arg⁡min1≤i≤d⁡{biai}.i_{0}\in\arg\min_{1\leq i\leq d}\bigg\{\frac{b_{i}}{a_{i}}\bigg\}. Thus, the pre-training loss for Adam is

ℒpretrain​(θAdam∗)=bj022​aj02+∑i≠j0θAdam,i∗2>bj022​aj02≥bi022​ai02=ℒpretrain​(θMoFO∗).\mathcal{L}_{\rm pretrain}(\theta^{*}_{\text{Adam}})=\frac{b_{j_{0}}^{2}}{2a_{j_{0}}^{2}}+\sum_{i\neq j_{0}}\theta^{*2}_{\text{Adam},i}>\frac{b_{j_{0}}^{2}}{2a_{j_{0}}^{2}}\geq\frac{b_{i_{0}}^{2}}{2a_{i_{0}}^{2}}=\mathcal{L}_{\rm pretrain}(\theta^{*}_{\text{MoFO}}).In other words,

‖θMoFO∗−θpretrain‖22=2​ℒpretrain​(θMoFO∗)<2​ℒpretrain​(θAdam∗)=‖θAdam∗−θpretrain‖22.\|\theta^{*}_{\rm MoFO}-\theta_{\rm pretrain}\|^{2}_{2}=2\mathcal{L}_{\rm pretrain}(\theta^{*}_{\text{MoFO}})<2\mathcal{L}_{\rm pretrain}(\theta^{*}_{\text{Adam}})=\|\theta^{*}_{\rm Adam}-\theta_{\rm pretrain}\|^{2}_{2}. In conclusion, MoFO converges to a minimum closer to the pre-training state than Adam, preserving a lower pre-training loss. ∎

A.4Challenges and Potential Extensions of Theorem1to Nonsmooth Objectives

This subsection outlines challenges and possible future directions for extending our convergence analysis to nonsmooth objectives; a complete extension is left for future work.

Why the extension is challenging.

Our Theorem1relies on the standardLL-smoothness assumption—i.e., the gradient ofℒ\mathcal{L}isLL-Lipschitz—to invoke the descent lemma and to derive an upper bound onmin0≤t≤T−1⁡‖∇ℒ​(θt)‖\min_{0\leq t\leq T-1}\|\nabla\mathcal{L}(\theta_{t})\|. Whenℒ\mathcal{L}is nonsmooth,

1. (i)∇ℒ​(θ)\nabla\mathcal{L}(\theta)may not exist at nondifferentiable points, which requires working with generalized gradients (e.g., subgradients) rather than classical gradients in the derivation.
2. (ii)The upper bound onmin0≤t≤T−1⁡‖∇ℒ​(θt)‖\min_{0\leq t\leq T-1}\|\nabla\mathcal{L}(\theta_{t})\|in the proof of Theorem1(AppendixA.2) scales with the smoothness constantLL; whenLLis unbounded (or effectively very large), these inequalities become non-informative.

Possible extensions.

We outline below several plausible directions; full development is left for future work.

1. (i)Subgradient analysis.Replace classical gradients with subgradients and assess convergence via a subgradient-based stationarity criterion.
2. (ii)Smoothing.Introduce a family of smoothed surrogatesℒμ{\mathcal{L}_{\mu}}, analyze the MoFO algorithm under these surrogates to obtainμ\mu-dependent bounds, and then letμ↓0\mu\downarrow 0to recover convergence results for the original lossℒ\mathcal{L}.
3. (iii)Algorithmic modifications.Incorporate techniques such as gradient clipping or stochastic subgradient steps(Xiao et al.,2024), and then analyze the convergence of the modified algorithm.

Appendix BSupplemental Figures and Explanations

B.1Reason for using RedPajama to approximate LLaMA-2’s training data

We note that original LLaMA-2 training dataset has not been publicly released. Thus, we can only rely on public datasets to approximate LLaMA-2’s original training data.

RedPajama project was explicitly designed as an open-source reproduction of the LLaMA training dataset(Weber et al.,2024; Together,2023). It closely mirrors the data sources outlined in the original LLaMA paper and adopts similar strategies for data collection, mixture, and preprocessing. We believe it serves as a reasonable proxy for approximating LLaMA-2’s training dataset.

B.2Supplementary Figures for Figure2(a)

(a)The loss landscapes of Pythia-160M after fine-tuning on a subset of the FLAN dataset using Adam and Lion. We plot the loss landscapes on (a) the fine-tuning dataset and (b) the pre-training dataset (Pile dataset(Gao et al.,2020)) and (c) the accuracies on CR tasks, including HellaSwag, ARC-c, and ARC-e. We visualize a 2D weight-space plane spanned by the vector from the pre-trained model to the Lion-tuned model (x-axis) and to the Adam-tuned model (y-axis). Axes are normalized so that one unit equals the length of the pre-trained→\toAdam vector. The color bar indicates the loss value—(a) fine-tuning loss and (b) pre-training loss. A logarithmic scale is applied to the loss values for better visualization. Two training methods converge to different minima with similar fine-tuning loss. Lion converges to a farther minimum from the pre-trained model and performs more forgetting than Adam.### B.3Supplementary Experiments on the Correlation between Distance and Forgetting

(b)Average accuracy on the MMLU benchmark (measuring preservation of factual knowledge) for Llama-2-7B after fine-tuning on MetaMathQA with Adam, Lion, and MoFO. Building on Figure2(b)(b), we add points for runs where Llama-2-7B was trained for 0.1, 0.2, and >3 epochs using both Adam and MoFO. The marker size encodes the number of training epochs (larger means more epochs). Red arrows indicate the points obtained after exactly 1 epoch for each optimizer.In this subsection, we augment Figure2(b)(b) by probing the relationship between a model’s parameter distance from its pre-trained state and evaluation accuracy under additional training budgets and optimizers. Concretely, under the same settings as Figure2(b)(b), we add runs at 0.1 and 0.2 epochs and extend training beyond 3 epochs, using Adam and MoFO. We report results on MMLU (as in the main text, measuring preservation of factual knowledge) and newly include HumanEval (measuring preservation of code-generation ability). Since the fine-tuning task is math, both can serve as forgetting mitigation metrics. The corresponding scatter plots are shown in Figure8(b)(MMLU) and Figure8(c)(HumanEval). When examining the pointsfor each optimizer separately, we make the following observations:

• •Observation 1 (sufficient training).Once training exceeds approximately 1 epoch, MMLU and HumanEval scores show a consistent strong negative correlation with the parameter distance to the pre-trained state.
• •Observation 2 (early training).For Adam or MoFO at less than 1 epoch, we may observe a mild correlation with the parameter distance to the pre-trained state. The correlation may be unstable and can be positive or negative.

We speculate that this short-lived positive trend may be related to benchmark alignment. The MMLU benchmark (used to measure the preservation of factual knowledge) might share partial overlap with the patterns of our math fine-tuning task (measured by GSM8K benchmark); in the early training steps, the model might incidentally acquire features that also benefit MMLU, resulting in a temporary gain. By contrast, as for HumanEval in Figure8(c), which measures code generation and differs from our math fine-tuning in both domain and output format, it may exhibit an unstable correlation whose sign can be either negative or positive. Another possible factor could stochasticity, since at less than 1 epoch the dataset has not yet been fully traversed.

Overall, the negative correlation becomes clear and consistent after sufficient training; while the early training presents a mild, benchmark-dependent relationship.

(c)Scores on HumanEval benchmark for Llama-2-7B after fine-tuning on MetaMathQA with Adam, Lion, and MoFO. Building on Figure2(b)(b), we add points for runs where Llama-2-7B was trained for 0.1, 0.2, and >3 epochs using both Adam and MoFO. The marker size encodes the number of training epochs (larger means more epochs). Red arrows indicate the points obtained after exactly 1 epoch for each optimizer.In addition, we emphasize an empirical point: the negative relationship between parameter distance and the preservation of pre-trained knowledge is evidentacross optimizers. In Figure2(b)(b),8(b), and8(c), the parameter distances roughly follow the orderingLion > Adam > MoFO, whereas the scores measuring preservation of pre-trained knowledge follow the inverse orderingMoFO > Adam > Lion. For a broader comparison, we evaluate five optimizers—MoFO, NAdam(Dozat,2016), Adam, RMSProp(Tieleman,2012), and Lion. After two epochs of training, we report (i) their parameter distance to the pre-trained state and (ii) their forgetting-mitigation performance on MMLU and HumanEval, shown in Figure8(d)(a) and Figure8(d)(b), respectively. The results show a consistent negative correlation between distance and performance. Notably, MoFO remains closer to the pre-trained state and achieves higher scores compared with the other optimizers.Therefore, this cross-optimizer rank-order correlation is sufficient to motivate our algorithm design: favor optimizers that converge closer to the pre-trained state, so as to better preserve the pre-trained knowledge.

(d)(a) MMLU accuracies and (b) HumanEval scores of Llama-2-7B after fine-tuning on the MetaMathQA dataset using different optimizers. All experiments are run for 2 epochs.### B.4Supplemental Explanation of Example1

(e)The loss landscapes of the example. We plot the landscapes on (a) the fine-tuning loss and (b) the pre-training loss. The color bar indicates the loss value—(a) fine-tuning loss and (b) pre-training loss. A logarithmic scale is applied to the loss values for better visualization. In Example1, MoFO converges to a minimum closest to the pre-trained model, with a low pre-training loss.In addition to Default FT and MoFO, we also analyze four other optimization methods in Example 1, namelyL2L_{2}regularization(Li et al.,2018), Half Fine-tuning (HFT)(Hui et al.,2024), Random BCD(Nesterov,2012), and Low-Rank Adaptation (LoRA)(Hu et al.,2022). These methods are also introduced in Section4.

For theL2L_{2}regularization method, we add a regularization termλ​‖θ−θpretrain‖22\lambda\|\theta-\theta_{\mathrm{pretrain}}\|^{2}_{2}to the fine-tuning loss to encourage the model to stay closer to the pre-trained state. As shown in Figure8(e)(b), theL2L_{2}regularization approach remains closer to the pre-trained model, thereby achieving a smaller pre-training loss. However, since the fine-tuning objective is modified, Figure8(e)(a) shows thatL2L_{2}regularization does not reach the minimum of the fine-tuning loss.

We note that HFT operates in a manner similar to Random BCD, which randomly selects a subset of coordinates (e.g.,θ1\theta_{1}orθ2\theta_{2}) to update. Figure8(e)further illustrates that both HFT and LoRA do not converge to minima as close to the pre-trained model as MoFO does, indicating that they may undergo higher levels of forgetting compared to MoFO.

For LoRA, we make the following modelling in the landscape visualization example. The core principle of LoRA (Low-Rank Adaptation) involves approximating the original training space by a low-rank subspace. Since we consider a two-dimensional training space for visualizing the landscape, we set the rank of LoRA space to 1. Specifically, the parametersθ1\theta_{1}andθ2\theta_{2}exhibit a linear relationship. Given that the pre-trained model is(0,0)(0,0), the parameters under LoRA are set to satisfyθ2=β​θ1\theta_{2}=\beta\theta_{1}, whereβ\betais a hyperparameter we set to 0.5. Figure8(e)shows that LoRA converges to a closer local minimum than Default FT.

B.5Synthetic Experiment for Example1

In this subsection, we conduct a synthetic experiment to provide a more concrete illustration of Example1. Specifically, we set the parameter dimension tod=10d=10, with parametersθ=(θ1,…,θ10)\theta=(\theta_{1},\dots,\theta_{10}). The pretraining loss is defined as the squaredL2L_{2}norm of the parameters:Lpretrain​(θ)=12​‖θ‖22L_{\rm{pretrain}}(\theta)=\frac{1}{2}\|\theta\|_{2}^{2}, with the pre-trained model given byθpretrain=(0,0,…,0)\theta_{\rm pretrain}=(0,0,\dots,0). Starting from the pre-trained model, we optimize the parameters with respect to the fine-tuning lossℒ​(θ)=∏i=1d(ai​θi−bi)2\mathcal{L}(\theta)=\prod_{i=1}^{d}(a_{i}\theta_{i}-b_{i})^{2}, whereai,bi>0a_{i},b_{i}>0for any1≤i≤d1\leq i\leq d. The coefficientsaia_{i}andbib_{i}are sampled from a standard normal distribution; to ensure positivity, we take their absolute values and add 0.3 and 0.1, respectively.

In this experiment, we compare two optimizers: Adam and MoFO. For each, we perform a grid search for the optimal learning rate over the set{10−2,10−3,10−4}\{10^{-2},10^{-3},10^{-4}\}. We consider the fine-tuning process to have converged to a minimum when the fine-tuning loss drops below10−810^{-8}within 10000 iterations. At convergence, we record two metrics: the Euclidean distance from the fine-tuned model to the original pre-trained model, and the value of the pretraining loss. The entire experiment is repeated across three different random seeds for robustness.

As presented in Table4, the results indicate that the minimum found by MoFO is at roughly half the distance from the pre-trained model compared to the one found by Adam. Concurrently, MoFO achieves a lower pre-training loss. These findings provide strong evidence that MoFO can converge to minima closer to the pre-trained model in Example1, thereby supporting Theorem2.

Table 4:The Euclidean distance from the fine-tuned model to the original pre-trained model, and the pretraining loss of parameters after optimizing the fine-tuning loss using Adam and MoFO. The results show that MoFO finds a fine-tuning minimum that is closer to the pre-trained model compared to the one found by Adam.AdamMoFODistance to pre-trained model0.5490.275Pre-training loss0.1510.038## Appendix CImplementation Details

C.1Datasets for Fine-Tuning.

MetaMathQA(Yu et al.,2024b). 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. We randomly select 10% of this dataset for training LLMs, which includes 39.5K question-answer pairs.

PMC-LLaMA-Instructions(Wu et al.,2024). This dataset comprises 514K instruction-response pairs. Fine-tuning LLMs on this dataset has been shown to enhance performance on medical NLP tasks, such as PubMedQA(Jin et al.,2019), MedMCQA(Pal et al.,2022), and MedQA(Jin et al.,2021). We randomly sampled 51K instances with prompt lengths less than 750 characters for training our models.

Magicoder-Evol-Instruct(Wei et al.,2023). This dataset comprises 110K instruction-response pairs related to coding, and it is decontaminated and redistributed from the Evol-CodeAlpaca-V1 dataset(Luo et al.,2023b)We randomly select 39.5K question-answer pairs from this dataset to fine-tune LLM and enhance their coding capability.

TRACE benchmark dataset(Wang et al.,2023b). 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.

C.2Evaluation Metrics for Instruction Fine-Tuning

We employ a comprehensive suite of widely used benchmarks to assess the performance and potential catastrophic forgetting effects on the general capabilities of LLMs after instruction fine-tuning. The benchmarks are as follows:

• •Factual knowledge (MMLU): We use the Massive Multitask Language Understanding (MMLU) benchmark(Hendrycks et al.,2021)to evaluate factual knowledge across 57 diverse subjects, ranging from STEM fields and the humanities to social sciences. Evaluations are performed using 8-bit precision with the open-instruct implementation, and by following the setup of(Hui et al.,2024), we report the 0-shot accuracy.
• •Common sense reasoning (CR): To measure the commonsense reasoning capabilities of LLMs, we employ the widely recognized benchmarks ARC-Challenge (ARC-C), ARC-Easy (ARC-E)(Clark et al.,2018), and HellaSwag(Zellers et al.,2019), collectively referred to as the Commonsense benchmark. We use the average of their metrics as the evaluation, conducting assessments using the LM Eval Harness framework(Gao et al.,2023)and reporting the 0-shot accuracy based on the “acc_norm, none” metric.
• •Mathematical Reasoning (GSM8K): We assess mathematical reasoning capability using GSM8K(Cobbe et al.,2021), which consists of 8.5K high-quality grade school math problems. Evaluations are conducted on the test set using the LM Eval Harness framework prompting in a 5-shot setting, reporting the “exact_match, flexible-extract” metric.
• •Code Generation (HumanEval): We adopt HumanEval(Chen et al.,2021), comprising 164 unique programming problems, to evaluate the coding capabilities of LLMs. For chat experiments, we report the pass@10 performance.
• •Medical Question Answering (MedQ): To assess medical knowledge, we utilize three benchmarks—PubMedQA(Jin et al.,2019), MedMCQA(Pal et al.,2022), and MedQA(Jin et al.,2021). Evaluations are performed using the LM Eval Harness framework. For PubMedQA, we report the “acc, none” metric; for MedMCQA and MedQA, we report the “acc_norm, none” metric.
• •Instruction Following (IFEval): We evaluate the instruction-following ability of LLMs using the IFeval benchmark. Evaluations are conducted with the LM Eval Harness implementation, and we report the “inst_level_strict_acc, none” metric.

All benchmarks—including CommonSense, GSM8K, PubMedQA, MedMCQA, MedQA, and IFeval—are evaluated using the LM Eval Harness framework(Gao et al.,2023), following their default settings unless specified otherwise.

C.3Hyperparameter Configurations

Instruction fine-tuning.In our instruction fine-tuning experiments, we follow the implementation ofIvison et al. (2023). For instruction fine-tuning, we set the maximum sequence length to 1024, the global batch size to 128, and we train the model for 2 epochs. For the Llama-2-7B model, we use a learning rate of 2e-5 and 0 warm-up ratio, with a cosine decay learning rate scheduler. The learning rate is set to 2e-5 for fine-tuning both the Llama-2-7B-Chat model on the MetaMathQA dataset and the Gemma-2B-IT model, while a learning rate of 1e-5 is used for fine-tuning the Llama-2-7B-Chat model on the PMC-LLaMA-Instruct dataset; all these settings employ a warm-up ratio of 0.03 and a cosine decay learning rate scheduler. For LoRA, we set the learning rate as 1e-4. The other hyperparameters in the experiments are as follows.

Fine-tuning Llama-2-7B on MetaMathQA.

• •Learning rate: 2e-5.
• •Update fraction of MoFO:α=15%\alpha=15\%.
• •LoRA:r=4,16,64,256.r=4,16,64,256.We report the best-performing hyperparameter configuration for the fine-tuning task in Table1, which, in this case, isr=256r=256.

Fine-tuning Llama-2-7B-Chat on PMC-LLaMA-Instruct.

• •Learning rate: 1e-5.
• •Update fraction of MoFO:α=15%\alpha=15\%.
• •LoRA:r=16,256.r=16,256.We report the best-performing hyperparameter configuration for the fine-tuning task in Table5, which, in this case, isr=256r=256.

Fine-tuning Llama-2-7B-Chat on MetaMathQA.

• •Learning rate: 2e-5.
• •Update fraction of MoFO:α=15%\alpha=15\%.
• •LoRA:r=16,256.r=16,256.We report the best-performing hyperparameter configuration for the fine-tuning task in Table7, which, in this case, isr=256r=256.

Fine-tuning Gemma-2B-IT on MetaMathQA.

• •Learning rate: 2e-5.
• •Update fraction of MoFO:α=5%\alpha=5\%.
• •LoRA:r=16,256,512.r=16,256,512.We report the best-performing hyperparameter configuration for the fine-tuning task in Table6, which, in this case, isr=512r=512.

Fine-tuning Llama-2-7B-Chat on Magicoder-Evol-Instruct.

• •Learning rate: 1e-5.
• •Update fraction of MoFO:α=20%\alpha=20\%.
• •LoRA:r=16,256.r=16,256.We report the best-performing hyperparameter configuration for the fine-tuning task in Table6, which, in this case, isr=256r=256.

Hyperparameters in the Pareto comparison.To provide a comprehensive comparison, we explore various hyperparameter settings forλ1\lambda_{1},λ2\lambda_{2}, LoRA’s rank, and the update fractionα\alphain MoFO in Figure5(a). Specifically, we setλ1\lambda_{1}as 1e-4, 1e-5, 1e-6, 1e-7, whileλ2\lambda_{2}is set as 1e-2, 5e-3, 1e-3, 5e-4, and 1e-4. The update fractionα\alphain MoFO is set as5%5\%,10%10\%,15%15\%,20%20\%,40%40\%,80%80\%. The rank of LoRA is set as 4, 16, 64, 256.

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 to5%5\%.

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

C.4More Explanation on the partitioning and Calculation of distance

Partitioning.We use the default partitioning scheme in PyTorch’s Transformer implementation. Different types of parameters within the Transformer, such as query (Q), key (K), value (V) weights for attention heads, and feed-forward network (FFN) weights, are divided into separate partitions. Notably, in the default PyTorch implementation, within a layer, the query (Q) weights of all attention heads are grouped into a single partition. The same applies to the key (K) and value (V) weights. Our momentum-based filtering mechanism is applied to each partition individually. A different parameter partition scheme, along with its corresponding experiments, is presented in AppendixF.6.

Calculation of distance.Following the notation in Section2.2, we suppose that the parameter parameters are partitioned into

θ=(θ(1),θ(2),…,θ(B)).\theta=(\theta^{(1)},\theta^{(2)},\dots,\theta^{(B)}).Denote the pre-trained model byθ0\theta_{0}and the fine-tuned model byθ\theta.

First, we calculate the relative change of parameters‖θ(k)−θ0(k)‖‖θ0(k)‖\frac{\|\theta^{(k)}-\theta_{0}^{(k)}\|}{\|\theta_{0}^{(k)}\|}in each partitionk∈{1,2,…,B}k\in\{1,2,\dots,B\}. Second, we compute the distance from the pre-trained modelθ0\theta_{0}to the fine-tuned modelθ\thetaby averaging the relative changes across all partitions, defined as:

D​(θ,θ0)=1B​∑k=1B‖θ(k)−θ0(k)‖‖θ0(k)‖.D(\theta,\theta_{0})=\frac{1}{B}\sum_{k=1}^{B}\frac{\|\theta^{(k)}-\theta^{(k)}_{0}\|}{\|\theta^{(k)}_{0}\|}.

Appendix DGuideline for Settingα\alpha

Given a pre-trained LLM and a dataset for fine-tuning, we recommend the following procedure:

1. 1.Random Sampling: Randomly sample a small subset of the dataset to serve as a proxy.
2. 2.Grid Search: Perform a grid search over candidate values ofα\alphausing this proxy subset.
3. 3.Selection: Choose theα\alphaconfiguration that strikes a good balance between fine-tuning performance on the target dataset and preserving the model’s general capability.

To illustrate this procedure, we use Llama-2-7B as an example. We randomly sample 10% of the instances from the current training set of MetaMathQA dataset (39.5k) for fine-tuning, and then perform a grid search overα\alphavalues of 5%, 10%, 15%, 20%, 40%, and 80%. As shown by the green line in Figure10(a)(a), the fine-tuning performance is relatively stable across these values ofα\alpha. However, the green line in Figure10(a)(b) indicates thatα=15%\alpha=15\%best preserves the model’s general capability. Therefore, we setα=15%\alpha=15\%.

(a)(a) Fine-tuning performance:Accuracy on the GSM8K math reasoning task for LLMs of different sizes, fine-tuned via MoFO with varying update fractions (α\alpha).(b) Preservation of pre-training knowledge:Average accuracy changes on MMLU, HumanEval, and commonsense reasoning benchmarks relative to the original pre-trained LLMs, illustrating how much pre-training knowledge is retained. All results are obtained by fine-tuning Llama2-7B on MetaMathQA and its proxy subset. The performance trends under different update fractions on the proxy subset align with those observed on the full dataset.Furthermore, by comparing the blue and green lines in Figure10(a), we observe that the trend of model performance with respect toα\alphaon the small proxy subset is consistent with the trend observed when fine-tuning on the full dataset. This implies that a small, randomly sampled subset is sufficient to guide the selection of a suitableα\alpha.

Empirically, the optimalα\alphaoften lies between 5% and 20%. A more fine-grained grid search within this range can be performed if needed. Designing more refined and efficient strategies for tuningα\alphais left for future work.

Appendix EAdditional Experiments on Instruction Fine-tuning

This section begins with a comparison of MoFO and baseline methods across additional datasets and models inE.1. InE.2, we explore the combination of LoRA and MoFO to assess their performance. Finally, inE.3, we compare MoFO with several algorithms designed to mitigate forgetting.

E.1More Experimental Results in Instruction Fine-Tuning

Table 5:The performance on the fine-tuning task (medical QA task), measured by MedQ, and general capability scores of Llama-2-7B-Chat after fine-tuning on the PMC-LLaMA-Instruct dataset. The figure on the right visualizes both MedQ accuracy and general capability scores. The results show that MoFO achieves comparable performance in the MedQ while significantly mitigating forgetting of general capabilities. Bold values denote the best results among these methods.MethodMedQGeneral CapabilityCRIFEvalHumanEvalAvg.Llama-2-7B-Chat49.865.641.424.343.8Default FT54.364.632.120.639.1HFT54.465.233.523.140.6LoRA54.264.433.923.540.6MoFO54.365.638.625.043.1

Results of fine-tuning on PMC-LLaMA-Instruct.

We fine-tune Llama-2-7B-Chat on the PMC-LLaMA-Instructions dataset using various baseline methods and present the experimental results on medical question answering (MedQ) and general capabilities in Table5. Since the MMLU benchmark already contains medical-related instances(Hendrycks et al.,2021), which may lead to improved performance after fine-tuning, we instead use IFEval to assess general capabilities.

MoFO performs well on the fine-tuning task of medical QA. It achieves compatible performance compared to Default FT and HFT. In terms of general capabilities, MoFO demonstrates the least degradation compared to other baselines, with an average accuracy reduction of only 0.2%. Specifically, on the IFEval benchmark, our method only exhibits a minor reduction of 0.3%, while Default FT, HFT, and LoRA experience significant degradations ranging from 7.5% to 9.3%. On code generation (HumanEval) tasks and commonsense reasoning (CR) benchmarks, our method also only exhibits a minor reduction less than 0.2%.

Table 6:The performance of the fine-tuning task (math), measured by GSM8K, and the general capability scores of Gemma-2B-IT after fine-tuning on the MetaMathQA dataset. The figure on the right visualizes both GSM8K accuracy and general capability scores. The results show that MoFO achieves comparable performance in the fine-tuning task, while significantly mitigating forgetting of general capabilities. Bold values denote the best results among these methods.MethodGSM8KGeneral CapabilityCRIFevalHumanEvalAvg.Gemma-2B-IT11.457.633.631.540.9Default FT42.052.124.320.632.3HFT41.553.924.121.233.1LoRA40.654.426.129.836.8MoFO42.155.028.729.137.6

Results of fine-tuning Gemma-2B-IT on MetaMathQA.

We also explore how MoFO performs in other LLMs. Specifically, we fine-tune Gemma-2B-IT on MetaMathQA using various baseline methods and present the experimental results on mathematical reasoning (GSM8K) and general capabilities in Table6. The experimental results demonstrate that MoFO achieves comparable performance of the fine-tuning task to Default FT and HFT across different models. In terms of general capabilities, MoFO exhibits significantly less forgetting compared to other baselines. This result demonstrates the versatility of the MoFO algorithm.

Table 7:The performance of the fine-tuning task (math), measured by GSM8K, and the general capability scores of Llama-2-7B-chat after fine-tuning on the MetaMathQA dataset. The figure on the right visualizes both GSM8K accuracy and general capability scores. The results show that MoFO achieves comparable performance in the fine-tuning task, while significantly mitigating forgetting of general capabilities. Bold values denote the best results among these methods.MethodGSM8KGeneral CapabilityCRIFevalHumanEvalAvg.Llama-2-7B-Chat13.765.641.424.343.8Default FT48.462.830.715.636.4HFT46.963.431.820.038.4LoRA45.363.935.621.040.2MoFO47.164.037.121.740.9

Results of fine-tuning Llama-2-7B-Chat on MetaMathQA.

We also fine-tune the Llama-2-7B-Chat on the MetaMathQA dataset. The results are presented in Table7. The results demonstrate that our approach achieves performance comparable to Default FT and HFT while exhibiting less forgetting compared to baseline methods.

Table 8:The performance of the fine-tuning task (coding), measured by HumanEval, and the general capability scores of Llama-2-7B-Chat after fine-tuning on the Magicoder-Evol-Instruct dataset. Here we choose the three benchmarks exhibiting the most significant forgetting. We set the rank of LoRA as 256, andα\alphaof MoFO is set as 20%. The results show that MoFO achieves comparable performance in the fine-tuning task, while mitigating forgetting of general capabilities. Bold values denote the best results among these methods.MethodHumanEvalGeneral CapabilityARC-EARC-CIFEvalAvg.Llama-2-7B-Chat24.274.546.341.154.0Default FT56.271.245.233.550.0HFT50.771.545.636.351.1LoRA48.672.145.133.450.2MoFO53.372.146.036.151.4

Results of fine-tuning Llama-2-7B-Chat on Magicoder-Evol-Instruct.

We also fine-tune the Llama-2-7B-Chat on the Magicoder-Evol-Instruct dataset. We use ARC-Easy (ARC-E) and ARC-Challenge (ARC-C) scores to measure general capability. The results in Table8demonstrate that our approach outperforms the baselines in the fine-tuning tasks and exhibits less forgetting compared to baseline methods.

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

E.2Experiment on the Combination of MoFO and LoRA

In addition to using LoRA as a baseline for comparison, we can also view it as an orthogonal method that can be integrated with MoFO. In the fine-tuning stage, LoRA restricts the trainable parameter space to a low-rank subspace. We note that the comparison of LoRA and MoFO (PEFT version) essentially evaluates MoFO and Adam within this same low-rank subspace defined by LoRA. To investigate this further, we conduct a comparative experiment following the setup in Table1of Section4.2. Table9implies that MoFO + LoRA effectively mitigates the forgetting issue that arises when using LoRA alone.

Table 9:The performance of the fine-tuning task (math), measured by GSM8K, and the general capability scores of Llama-2-7B after fine-tuning on the MetaMathQA dataset. The results show that MoFO + LoRA preserve more pre-training knowledge than using LoRA alone.MethodGSM8KGeneral CapabilityAvg.CRMMLUHumanEvalLoRA43.365.137.726.443.1LoRA + MoFO43.465.539.426.743.9

E.3Comparison with More Fine-Tuning Methods

Experiments on Heterogeneous Model Averaging (HMA)

We compare our proposed method with the Heterogeneous Model Averaging (HMA)(Lin et al.,2024a). HMA approach evenly divides the LLM into three parts—the input part, the middle part, and the output part—and averages these parts with different ratios. To facilitate a comprehensive comparison, following the setting in Section4.2, we evaluate the fine-tuning and forgetting mitigation performance for different HMA strategies. We select 15 different combinations of averaging ratios for different parts as follows: {(0.05, 0.2, 0.35), (0.1, 0.2, 0.3), (0.2, 0.2, 0.2), (0.3, 0.2, 0.1), (0.35, 0.2, 0.05), (0.3, 0.5, 0.7), (0.4, 0.5, 0.6), (0.5, 0.5, 0.5), (0.6, 0.5, 0.4), (0.7, 0.5, 0.3), (0.65, 0.8, 0.95), (0.7, 0.8, 0.9), (0.8, 0.8, 0.8), (0.9, 0.8, 0.7), (0.95, 0.8, 0.65)}. We plot the results to construct a Pareto front in Figure10.

Figure 10:The performance on the math task (GSM8K) and the scores in Commonsense Reasoning of Llama-2-7B after fine-tuning on the MetaMathQA dataset. The results show that the MoFO algorithm achieves a better Pareto front. The pink triangle represents the model obtained through HMA.Results show that our proposed method, MoFO achieves a more effective Pareto front compared to the baselines.

Experiments on CoFiTune and Soft-masking

Zhang et al. (2024a)introduces CoFiTune, a coarse-to-fine framework that balances specificity and versatility in LLMs by selectively updating specific modules and employing a soft-masking mechanism, which is introduced byKe et al. (2023b;a). We have compared MoFO with CoFiTune (with and without soft-masking) and the vanilla soft-masking method alone, following the setting in Table1of Section4.2. The results, presented in Table10below, demonstrate that MoFO outperforms these methods in both fine-tuning performance and mitigating forgetting. The results demonstrate that

• •CoFiTune achives similar forgetting mitigation performance as MoFO, but underperforms MoFO on fine-tuning tasks.
• •Vanilla Soft-masking exhibits slightly reduced performance in both fine-tuning tasks and mitigating forgetting than MoFO. These findings underscore the advantages of our proposed method.

Table 10:The performance of the fine-tuning task (math), measured by GSM8K, and the general capability scores of Llama-2-7B after fine-tuning on the MetaMathQA dataset. The results show that MoFO achieves comparable performance in the fine-tuning task, while significantly mitigating forgetting of general capabilities.MethodGSM8KGeneral CapabilityCRMMLUHumanEvalAvg.MoFO47.765.742.724.644.3Vanilla-SoftMask46.465.642.923.243.9CoFiTune w/o SoftMask37.765.442.125.844.4CoFiTune w/ SoftMask34.465.041.525.644.0From the results, we can see that MoFO achieves higher scores on the fine-tuning tasks while effectively reducing knowledge forgetting, demonstrating its superiority over these methods.

E.4Comparison with More Parameter-Efficient Fine-Tuning Methods

Parameter-Efficient Fine-Tuning (PEFT) encompasses a collection of fine-tuning methods designed to reduce the computational cost required for model training. In this subsection, we mainly focus on three famous PEFT methods:

• •Adapter(Houlsby et al.,2019): This approach involves inserting trainable “adapter” modules into every layer of model. During fine-tuning, the parameters of these adapter modules are updated for downstream tasks, while the majority of the original model’s parameters remain frozen.
• •BitFit(Zaken et al.,2021): This method reduces the number of trainable parameters by fine-tuning the bias terms of the model on a given downstream task, keeping other weights frozen.
• •LoRA(Hu et al.,2022): is a widely-used, parameter-efficient fine-tuning method. LoRA trains low-rank matrix adaptations on the base model’s weights. Recent work(Biderman et al.,2024)demonstrates that LoRA can mitigate forgetting.

We compared MoFO with Adapter, BitFit, and LoRA, following the setup in Table1of Section4.2. For the Adapter method, we set the learning rate to 1e-4 and performed a grid search for the adapter size over the values {16, 64, 128}. We report the best-performing hyperparameter configuration for each method on the fine-tuning task.

As shown in Table11, all three PEFT baselines (BitFit, Adapter, and LoRA) achieve higher average general capability scores than Default FT, indicating that they are indeed effective at mitigating catastrophic forgetting. However, they still lag behind MoFO in this regard.

Moreover, compared to MoFO, the three PEFT methods perform markedly worse on the fine-tuning task (GSM8K), with the gap being especially pronounced for Adapters and BitFit. We conjecture that this may stem from the fact that Adapters and BitFit were originally proposed and evaluated on relatively simple classification or QA benchmarks (e.g., GLUE) with masked language models. While these methods were effective in such settings, they may be less suited to today’s more challenging domain-specific LLM tasks, leading to the weaker performance observed here.

Table 11:The performance of the fine-tuning task (math), measured by GSM8K, and the general capability scores of Llama-2-7B after fine-tuning on the MetaMathQA dataset. The results show that MoFO outperforms three PEFT methods.MethodGSM8KGeneral CapabilityCRMMLUHumanEvalAvg.Llama-2-7B13.765.642.024.243.9Default FT49.462.336.616.138.3MoFO47.765.742.724.644.3BitFit15.164.836.124.441.8Adapter24.463.232.321.839.1LoRA43.365.137.726.443.1

Appendix FMore Explorations on MoFO

This section aims to provide a deeper understanding of MoFO through a series of experiments. In AppendixF.2, we conduct an efficiency analysis of MoFO. In AppendixF.3, we present additional comparative experiments on different filtering strategies. In AppendixF.4, we investigate why the momentum-filtered update rule in MoFO achieves optimal fine-tuning performance compared to other update strategies. In AppendixF.5, we present a preliminary analysis on why MoFO might compare favorably toL1/L2L_{1}/L_{2}regularization. In AppendixF.6, we investigate the performance of MoFO under an alternative parameter partitioning strategy. Finally, in AppendixF.7, we explore the application of MoFO to the Lion optimizer.

F.1Validating MoFO’s Impact on Preserving Pre-training Knowledge through Proximity

In this section, we empirically examine whether MoFO achieves its intended goal of converging to a minimum closer to the pre-trained model and mitigating forgetting mentioned in Section2.

Our exploratory experiment shows that MoFO indeed converges to a minimum closer to the pre-training model. As shown in Figure12(a)(a), both MoFO and the Adam optimizer achieve minimal fine-tuning loss, indicating that switching from Adam to MoFO does not lead to performance degradation. Moreover, the distance from the pre-trained model to the minimum reached by MoFO is approximately 20% of that reached by the default Adam optimizer.

(a)The loss landscapes of Pythia-160m after fine-tuning on a subset of the FLAN dataset using Adam optimizer and MoFO. We plot the loss landscapes on (a) the fine-tuning dataset and (b) the pre-training dataset (Pile). We visualize a 2D weight-space plane spanned by the vector from the pre-trained model to the MoFO-tuned model (x-axis) and to the Adam-tuned model (y-axis). Axes are normalized so that one unit equals the length of the pre-trained→\toAdam vector. The color bar indicates the loss value—(a) fine-tuning loss and (b) pre-training loss. A logarithmic scale is applied to the loss values for better visualization. We find that MoFO, reaching a closer point to the pre-trained model, has minimal fine-tuning loss and lower pre-training loss, compared to Adam.Table 12: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.HellaSwagARC-easyARC-challengeAveragePythia-160m30.139.623.831.2Adam28.337.422.129.3MoFO29.942.022.931.6 Our experiment demonstrates that the reduced parameter movement achieved by MoFO effectively mitigates the forgetting of pre-training knowledge. As shown in Figure12(a)(b), the fine-tuned model using MoFO experiences a smaller increase in pre-training loss. Additionally, Table12shows that MoFO achieves higher accuracy on commonsense reasoning tasks, indicating less forgetting.

F.2Efficiency Analysis on MoFO

We claim that MoFO does not lead to significant reduced fine-tuning efficiency. We provide an efficiency analysis by comparing the total training time between MoFO and Default FT on three LLMs with different sizes. The parameter update fraction of MoFO is set as 10%. The experimental results show that although MoFO requires computing a filter to select parameters with the largest momentum magnitudes, the additional computational overhead is minor. As shown in Table13, the additional training time incurred by MoFO is approximately 4%–5%, which is relatively minor and manageable in practical applications.

Table 13:Comparison of total training time for Default FT and MoFO on various LLaMA models and datasets. The additional time incurred by MoFO is around 4–5%, which is relatively minor in practical applications. LLaMA3-8B is fine-tuned on the UltraFeedback dataset(Cui et al.,2024).ModelDefault FTMoFOAdditional Training TimeLLaMA3.2-1B49m22s51m24s4.1%LLaMA3.2-3B1h30m18s1h34m24s4.5%LLaMA3-8B5h2m0.69s5h17m34.01s5.0%### F.3Further Comparative Experiments on Filtering Strategies

To further substantiate the claim in Section4.4, we conduct an additional comparison of filtering strategies—complementary to the results in Table3—using a different dataset and model. Specifically, we fine-tune Gemma-2B-IT on the IFEval-like dataset777The “filtered” subset fromhttps://huggingface.co/datasets/argilla/ifeval-like-data.using several BCD variants. The IFEval-like dataset contains instruction–response pairs in the style of the IFEval benchmark. For these experiments, we randomly sample 39.5k instances from the dataset for training, set the learning rate to1×10−51{\times}10^{-5}, and train the model for 2 epochs. We evaluate fine-tuning performance with the IFEval benchmark and assess general capabilities with CR (common-sense reasoning), HumanEval, and BBH (0-shot)(Suzgun et al.,2022).

As shown in Table14, all BCD variants mitigate forgetting: their average general-capability score exceeds that of Default fine-tuning by at least 2.6%. Although MoFO is lower than MV BCD by only 0.1% in average general capability, it achieves the strongest performance on the target fine-tuning task (IFEval), outperforming Random BCD, Grad BCD, and MV BCD by 5.8%, 2.6%, and 1.3%, respectively.

Taken together, these results indicate that while all filtering strategies are broadly effective at mitigating forgetting, they exhibit distinct differences in task-specific fine-tuning performance. In our setting, MoFO offers a favorable balance—matching the strongest methods on forgetting mitigation while delivering the highest IFEval score among the tested strategies.

Table 14:The performance on the instruction-following task (IFEval) and general capability scores of Gemma-2B-IT after fine-tuning on IFEval-like dataset using different updating strategies in MoFO. Here we choose the three benchmarks exhibiting the most significant forgetting. For all the update strategies, we set the parameter update fractionα\alphaas 10%. Bold values denote the best results among the BCD methods.MethodIFEvalGeneral CapabilityCRHumanEvalBBHAvg.Gemma-2B-IT33.657.631.532.740.6Default FT56.756.922.932.537.4Random BCD51.457.528.033.939.8Grad BCD54.657.328.134.139.8MV BCD55.957.528.434.640.2MoFO57.257.328.534.440.1### F.4Insights of the Choice of Filtering Strategy

In this section, we attempt to address the question:What makes the momentum-filtered update rule optimal among the candidates?

We hypothesize that the good performance of the momentum-filtered updating rule arises from its ability to promote more stable and consistent updates throughout training. Specifically, we hypothesize that:

1. 1.Utilizing momentum instead of gradient filtering leads to more stable updates.Momentum accumulates historical gradients, so it promotes stability by smoothing out fluctuations in the gradient updates. Thereby, during the training process, the momentum-filtering mechanism chooses updating parameters in a more stable manner.
2. 2.Excluding the introduction of in the filtering mechanism contributes to more stable updates.The 2nd-order momentvtv_{t}may normalize gradients based on their magnitudes, potentially averaging out the importance of individual parameters within the filtering mechanism. Thereby, during the training process, not incorporatingvvmay help choose updating parameters in a more stable manner.

For the ablation study, we add the GV-filtered BCD methods as a baseline, which replaces MoFO’s filter byFLTα​(gt/vt)\texttt{FLT}_{\alpha}(g_{t}/\sqrt{v_{t}}). Following the setting of experiments in AppendixB.2and Figure2(a), we run all four methods with the updating fractionα=3%\alpha=3\%over approximately 200 steps. To assess how many parameters change significantly during training, we calculate the percentage of weight parameters whose absolute change exceeds a threshold of 2e-6.

Table 15:Percentage of weight parameters with significant changes (absolute change >2×10−62\times 10^{-6}) during training.MethodPercentage of Significant UpdatesMoFO29.8%Gradient BCD35.7%MV-filtered BCD83.6%GV-filtered BCD87.1%Table15indicates that the parameter updating process of MoFO is more stable, which we think may contribute to its better fine-tuning performance.

F.5Preliminary Analysis on Why MoFO Might Compare Favorably toL1/L2L_{1}/L_{2}Regularization

As shown in Figure5(a)(Section4.2), when fine-tuning Llama-2-7B on MetaMathQA, MoFO yields a more favorable Pareto front—balancing fine-tuning performance and forgetting mitigation—thanL1L_{1}/L2L_{2}regularization and LoRA across hyperparameter configurations. This section presents a preliminary analysis offering one possible explanation for why MoFO may compare favorably toL1/L2L_{1}/L_{2}regularization. Importantly, we reuse the exact checkpoints from Figure5(a); no additional training runs are introduced.

Setup.We reuse the Llama-2-7B MetaMathQA sweeps underlying Figure5(a)—includingL1L_{1},L2L_{2}, and MoFO across the same hyperparameters. For each checkpoint, we compute the (unregularized) fine-tuning loss and theℓ2\ell_{2}norm of the gradient on 512 examples. To factor out severe forgetting, we report results for the subset whose CR scores fall within an acceptable range (same CR metric as in the main text). Figure14(a)(a) and Figure14(a)(b) plot fine-tuning loss and gradient norm versus CR, respectively, for this subset of checkpoints drawn from Figure5(a).

Observations.For comparable CR scores, MoFO checkpoints generally achieve lower fine-tuning losses than those trained withL1L_{1}orL2L_{2}; moreover, they also exhibit smaller gradient norms. Together, these appear to suggest that MoFO converges better thanL1L_{1}/L2L_{2}regularization.

One plausible explanation.Penalty methods promote proximity to the pretrained parameters, but they alsomodify the training objective and its gradient field, thereby possibly hindering convergence to a local minimum of the original fine-tuning loss. MoFO instead constrains how updates are applied (via momentum filtering) while continuing to optimize the original fine-tuning loss, which may allow the optimizer to follow task-relevant directions more effectively.

(a)(a) Fine-tuning loss and (b) gradient norm versus CR score for Llama2-7B fine-tuned on MetaMathQA with MoFO,L1L_{1}, andL2L_{2}regularization. The results show that MoFO generally achieves a lower fine-tuning loss and gradient norm thanL1L_{1}andL2L_{2}while effectively resisting forgetting.### F.6A Different Parameter Partition Scheme

In this section, we propose a different strategy for partitioning model parameters. Different from the default partitioning scheme in PyTorch’s Transformer implementation (AppendixC.4), our alternative approach partitions parameters at the granularity of individual attention heads. Recent studies(Zhang et al.,2024d;c)show that the Hessian matrix in Transformers is nearly block-diagonal, with several dense principal sub-blocks. In particular,Zhang et al. (2024c)finds that within the same layer, distinct Q (or K) heads form different blocks in the Hessian.

Motivated by this finding, we treat each head’s Q, K, and V weights as separate partitions and apply our momentum-based filtering mechanism to these finer-grained partitions individually. To evaluate our method, we conduct experiments on the LLaMA2-7B models, comparing the default partitioning approach against our alternative head-level partitioning scheme.

Table 16:Performance on the fine-tuning task (GSM8K) and general capabilities ofLlama-2-7Bafter fine-tuning. Bold values denote the best results among these methods. The updating fractionα\alphais15%15\%.MethodGSM8KGeneral CapabilityCRMMLUHumanEvalAvg.MoFO with default partition47.765.742.724.644.3MoFO with individual head partition46.665.741.624.143.8Table16shows that both partition schemes yield similar performance.

From the above initial experiment results, we believe that there is room for further exploration in better partitioning strategies.

F.7Momentum Filtering Mechanism on the Lion Optimizer

Algorithm 2Lion Optimizer1:Input: Number of partitions

BBwith the

kk-th partition of size

dkd_{k}, hyperparameters

β1,β2,λ\beta_{1},\beta_{2},\lambdaof Lion optimizer, learning rate schedule

{ηt}\{\eta_{t}\}.

2:Initialize

m0m_{0}as zero tensors.

3:foriteration

ttfrom

1,2,…1,2,\dotsuntil convergedo

4:forpartition

kkfrom

11to

BBdo

5:

gt(k)=∇θ(k)ℒf​i​n​e​t​u​n​e​(θt−1)g^{(k)}_{t}=\nabla_{\theta^{(k)}}\mathcal{L}_{finetune}(\theta_{t-1}) 6:

ct(k)=β1​mt−1(k)+(1−β1)​gt(k)c^{(k)}_{t}=\beta_{1}m^{(k)}_{t-1}+(1-\beta_{1})g^{(k)}_{t} 7:

mt(k)=β2​mt−1(k)+(1−β2)​gt(k)m^{(k)}_{t}=\beta_{2}m^{(k)}_{t-1}+(1-\beta_{2})g^{(k)}_{t} 8:endfor

9:

θt=Concat​(θt(1),…,θt(B))\theta_{t}=\texttt{Concat}(\theta_{t}^{(1)},\dots,\theta_{t}^{(B)}) 10:endfor

Algorithm 3MoFO + Lion1:Input: Filtering threshold

α\alpha, number of partitions

BBwith the

kk-th partition of size

dkd_{k}, hyperparameters

β1,β2,λ\beta_{1},\beta_{2},\lambdaof Lion optimizer, learning rate schedule

{ηt}\{\eta_{t}\}.

2:Initialize

m0m_{0}as zero tensors.

3:foriteration

ttfrom

1,2,…1,2,\dotsuntil convergedo

4:forpartition

kkfrom

11to

BBdo

5:

gt(k)=∇θ(k)ℒf​i​n​e​t​u​n​e​(θt−1)g^{(k)}_{t}=\nabla_{\theta^{(k)}}\mathcal{L}_{finetune}(\theta_{t-1}) 6:

ct(k)=β1​mt−1(k)+(1−β1)​gt(k)c^{(k)}_{t}=\beta_{1}m^{(k)}_{t-1}+(1-\beta_{1})g^{(k)}_{t} 7:forentry index

iifrom

11to

dkd_{k}do

8:

[FLTα(k)​(mt−1)]i=1[\texttt{FLT}^{(k)}_{\alpha}(m_{t-1})]_{i}=1if

|(mt−1(k))i||(m^{(k)}_{t-1})_{i}|is within the top-

α\alphaof

|mt−1(k)||m^{(k)}_{t-1}|’s valueselse0

9:endfor

10:

θt(k)=θt−1(k)−ηt​(sign⁡(ct(k)⊙FLTα(k)​(mt−1))+λ​θt−1(k))\theta^{(k)}_{t}=\theta^{(k)}_{t-1}-\eta_{t}(\operatorname{sign}(c^{(k)}_{t}\odot\texttt{FLT}^{(k)}_{\alpha}(m_{t-1}))+\lambda\theta^{(k)}_{t-1})# Momentum Filtering

11:

mt(k)=β2​mt−1(k)+(1−β2)​gt(k)m^{(k)}_{t}=\beta_{2}m^{(k)}_{t-1}+(1-\beta_{2})g^{(k)}_{t} 12:endfor

13:

θt=Concat​(θt(1),…,θt(B))\theta_{t}=\texttt{Concat}(\theta_{t}^{(1)},\dots,\theta_{t}^{(B)}) 14:endfor

In this section, we extend our investigation by integrating our proposed MoFO into the Lion optimizer(Chen et al.,2024). We first present the original formulation of the Lion optimizer in Algorithm2.

Building upon this foundation, we propose Algorithm3, where momentum filtering is applied to refine the parameter update in every iteration (Lines 7-10).

Following the experimental setup in AppendixB.2and Figure2(a), we conduct comparative experiments to evaluate the effectiveness of ‘MoFO + Lion’ in mitigating forgetting. As shown in Table17, incorporating MoFO into Lion improves the average accuracy of preserving old knowledge by 2.6% compared to using Lion alone. However, ‘MoFO + Lion’ still lags behind ‘MoFO + Adam’ by 2.5%.

Table 17:Performance comparison on HellaSwag, ARC-easy, and ARC-challenge, along with their average.MethodHellaSwagARC-easyARC-challengeAveragePythia-160m30.139.623.831.2Adam28.337.422.129.3MoFO (for Adam)29.942.022.931.6Lion26.529.024.126.5Lion + MoFO27.536.723.229.1