🔢 mathematics

Optimizer-Model Consistency: Full Finetuning with the Same Optimizer as Pretraining Forgets Less

This article introduces the concept of "Optimizer-Model Consistency" and demonstrates that using the same optimizer for both pre-training and full fine-tuning reduces catastrophic forgetting and improves the learning-forgetting trade-off compared to other optimizers or LoRA, while also showing that certain optimizers such as Muon may underperform on reasoning tasks due to a tendency toward rote memorization.

Original authors: Yuxing Liu, Jianyu Wang, Tong Zhang

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuxing Liu, Jianyu Wang, Tong Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant student who has spent years reading every book in a massive library (this is Pretraining). They possess general knowledge, grammar, and have learned how to think about the world. Now, you wish to teach them a specific new skill, such as solving math problems or writing computer code (this is Supervised Finetuning or SFT).

The major challenge is: How do you teach them this new skill without them forgetting everything they learned in the library? If they focus too much on the new mathematics, they might begin to forget how to write a simple story. This is known as the "Learning-Forgetting-Tradeoff".

This work discovers a surprising rule for how best to teach this student, and it has to do with how you teach them, not just what you teach them.

The Main Discovery: "The Same Teacher, The Same Style"

The authors found that the best way to teach the new skill is to use the exact same teaching style (Optimizer) that was employed during the years of library reading.

The Analogy: Imagine the student learned reading with a specific method, say "Method A" (like a certain way of highlighting text and taking notes). If you try to teach them math with "Method B" (a completely different way of highlighting and note-taking), the student becomes confused. They may learn the math, but they begin to forget their reading skills because the new method clashes with the way their brain has been wired through years of practice.
The Insight: If you stick with "Method A" for the math instruction, the student learns the math and retains their reading skills intact. The work calls this "Optimizer-Model Consistency".

Why Does This Happen? (The "Brain Shape" Theory)

The work explains that different teaching methods actually shape the student's brain (the model's internal weights) in different ways.

Different Shapes: Some teaching methods (like AdamW, the most common one) train the brain to be "sparse" or "efficient," like a library where books are neatly organized with empty spaces in between. Other methods (like Muon, a newer, sophisticated method) train the brain to be "dense," like a library where books are packed tightly together.
The Discrepancy: If you train the brain for years to be "sparse" and then suddenly switch to a "dense" teaching method for math, the brain must panic to reorganize its entire structure. This chaos causes it to drop old books (forgetting).
The Match: If you continue using the "sparse" method for math as well, the brain simply adds new books to the existing shelves. It doesn't need to rebuild the library, so it forgets less.

The LoRA Surprise

There is a popular shortcut called LoRA (Low-Rank Adaptation). Imagine LoRA as a small, separate notebook for the student, where they write math answers without touching their main library books. Many people thought this was the best way to avoid forgetting.

The Work's Twist: The authors found that while LoRA is good, full retraining (rewriting the entire brain) using the same teaching method as before is actually even better.

The Analogy: LoRA is like carrying a separate notebook. It works, but if you use the same teaching style to rewrite your main brain, you can preserve the new math and the old reading skills even more effectively than with just a side notebook.

The Case of the "Rote Learner" (Muon)

The work also examined a specific, advanced teaching method called Muon.

The Good: Muon is excellent in the library phase (Pretraining). It helps the student memorize facts incredibly well.
The Bad: When it comes to learning new thinking skills (like math) with very little data, Muon tends to become a "rote learner." It tries to memorize the specific math problems it sees, rather than understanding the underlying patterns.
The Result: If you use Muon for the entire journey (Library + Math), the student might be able to recite facts brilliantly but struggle to solve new, unseen math problems. They have memorized the examples but haven't understood the logic.

Summary in Simple Language

Consistency is Key: To learn a new skill without forgetting old ones, continue using the same learning algorithm (Optimizer) with which the model was originally trained.
Do Not Change the Style: Changing the learning algorithm halfway forces the model to reorganize its "brain," causing it to lose old knowledge.
Full Retraining > Shortcuts: Sometimes, full retraining with the correct algorithm is better than using shortcuts like LoRA, because it aligns better with the original structure of the model's brain.
Beware of Over-Memorization: Some advanced algorithms (like Muon) are great at memorization but may be worse at learning new patterns when data is scarce.

The work concludes that if you want the best balance between learning new things and remembering old things, you should stick with the same teacher you started with.

Technical Summary: Consistency between Optimizer and Model

Problem Statement

In the training paradigm of Large Language Models (LLMs), models undergo a pretraining phase on massive corpora, followed by a Supervised Fine-Tuning (SFT) phase on smaller, high-quality datasets to acquire specific capabilities. A critical challenge in SFT is the trade-off between learning and forgetting: the model must learn new, task-specific knowledge while preserving the general language capabilities acquired during pretraining.

Although various optimizers (e.g., AdamW, Muon) and parameter-efficient methods (e.g., LoRA) are employed, it remains unclear which combination offers the optimal balance between learning new tasks and minimizing catastrophic forgetting. Existing literature provides contradictory intuitions; some studies suggest that LoRA inherently forgets less due to its low-rank nature, while others emphasize the importance of optimizer choice. This work investigates whether the optimizer choice for SFT should be decoupled from or aligned with the optimizer used during pretraining.

Methodology

The authors employ a combination of empirical experiments and theoretical analysis to investigate the relationship between optimizer choice and model performance during SFT.

Empirical Experiments

Models and Datasets: Experiments were conducted on GPT-2-small (pretrained from scratch) and Llama-2-7B. Pretraining utilized the OpenWebText dataset. SFT was performed on three distinct tasks: MetaMathQA (mathematical reasoning), Alpaca (instruction following), and Magicoder (code generation).
Compared Optimizers: The study compares AdamW, Muon (a matrix-structured optimizer), and SGD. Additionally, LoRA (Low-Rank Adaptation) is evaluated with various ranks ( $r=4, 8, 32, 64, 128, 256$ ).
Evaluation Metric: The authors define a Pareto Frontier to visualize the trade-off between learning and forgetting.
- Forgetting: Measured by the validation loss on the pretraining dataset (OpenWebText) or the decline in general capability benchmarks. A lower loss indicates less forgetting.
- Learning: Measured by the validation loss on the new SFT task or task-specific accuracy. A lower loss/higher accuracy indicates better learning.
- Protocol: For each training algorithm, a grid search over learning rates was conducted to identify the Pareto Frontier, ensuring a fair comparison of the algorithm's limits rather than considering only a single hyperparameter setting.
Synthetic Analysis: To understand why certain optimizers behave differently, the authors conducted a synthetic experiment using a corrupted language corpus. They randomly shuffled token sequences in clean data to destroy linguistic patterns and force the model to rely on memorization. This enabled them to distinguish between an optimizer's ability to learn patterns and its tendency to memorize data.

Theoretical Analysis

The work provides a theoretical framework based on matrix-induced norms (following Pethick et al., 2025).

Optimizers as Regularizers: The authors postulate that optimizers do not merely converge to a solution but shape the model's weights and activations. Specifically, an optimizer derived from a norm $\|\cdot\|_{\alpha, \beta}$ $∥ \cdot ∥_{α, β}$ implicitly regularizes input activations via $\|\cdot\|_\alpha$ $∥ \cdot ∥_{α}$ and output activations via $\|\cdot\|_\beta$ $∥ \cdot ∥_{β}$ .
- AdamW (derived from $A_{1,\infty}$ ) regularizes inputs via $\|\cdot\|_1$ , leading to sparse activations.
- Muon (derived from $A_{2,2}$ ) regularizes inputs via $\|\cdot\|_2$ , leading to denser activations.
Approximation of Forgetting: Using a Taylor series expansion, the increase in pretraining loss (forgetting) after SFT is approximated as proportional to the expected squared norm of the weight update applied to pretraining inputs ( $\mathbb{E}[\|\Delta W x\|_2^2]$ ).
Consistency Hypothesis: The theory suggests that to minimize forgetting, the SFT update ( $\Delta W$ ) should align with the curvature of the loss landscape established by pretraining. This alignment is achieved when the SFT optimizer belongs to the same family (same $\alpha, \beta$ ) as the pretraining optimizer.

Key Findings

1. Consistency between Optimizer and Model

The central finding is the consistency between optimizer and model: Full fine-tuning with the same optimizer (or the same family of optimizers) used during pretraining consistently achieves the best trade-off between learning and forgetting.

AdamW Pretraining: When AdamW is used for pretraining, full fine-tuning with AdamW outperforms full fine-tuning with Muon or SGD and also performs better than LoRA across all tasks (math, code, instructions).
Muon Pretraining: Similarly, full fine-tuning with Muon, when Muon was used for pretraining, achieves the best trade-off compared to using AdamW or SGD for SFT.
Comparison with LoRA: Contrary to the intuition that LoRA forgets less, the Pareto Frontier analysis shows that for every LoRA checkpoint, a full fine-tuning checkpoint (using the same optimizer) exists that achieves both better learning (lower task loss) and less forgetting (lower pretraining loss). This holds when an appropriate learning rate sweep is considered.

2. Performance of Muon on Reasoning Tasks

While Muon often produces superior pretrained checkpoints (lower pretraining loss), its performance degrades relative to AdamW when fine-tuned on reasoning tasks (e.g., MetaMathQA).

Memorization vs. Pattern Learning: The synthetic experiment revealed that Muon memorizes corrupted (pattern-less) data faster than AdamW or SGD.
Implication: In SFT, where data is scarce, Muon's tendency to prioritize memorization over pattern acquisition can hinder the learning of complex reasoning capabilities, leading to poorer final performance on tasks requiring generalization, despite its strong pretraining metrics.

3. The Role of Learning Rates

The work highlights that the trade-off between learning and forgetting is heavily dependent on the learning rate. Previous studies suggesting that LoRA forgets less often fixed the learning rate to the one optimal for the task. When the learning rate is searched, full fine-tuning can achieve a superior trade-off simply by lowering the learning rate, a strategy that LoRA cannot replicate as effectively in this specific context.

Significance and Claims

The work claims to establish consistency between optimizer and model as a fundamental principle for LLM training:

Unified Metric: The trade-off between learning and forgetting across the entire pipeline from pretraining to SFT is a more comprehensive metric for evaluating an optimizer than pretraining performance alone.
Practical Guidance: To minimize forgetting while maximizing task learning, practitioners should use the same family of optimizers for SFT as for pretraining.
Re-evaluation of Muon: Although Muon is a powerful pretraining optimizer, its update mechanism (Full-Rank, $L_2$ regularization) may tend to lead to excessive memorization in SFT regimes with limited data, particularly on reasoning tasks. This suggests that a dedicated, potentially different fine-tuning strategy might be required for Muon-pretrained models to unlock their full potential.
Theoretical Insight: The work bridges the gap between optimizer design (norms) and model behavior (sparsity of activations), providing a theoretical explanation for why consistency in optimizer choice is important.

The authors acknowledge limitations, noting that experiments were primarily conducted on smaller models (GPT-2, Llama-2-7B) due to computational constraints and that theoretical proofs rely on simplifying assumptions (e.g., block-diagonal Hessian matrices). They suggest that future work should scale these findings to larger models and investigate optimizers specifically designed to mitigate the memorization tendencies of matrix-structured optimizers during SFT.