Optimizer-Model Consistency: Why Your Fine-tuning Should Respect its Pretraining Lineage
Optimizer-Model Consistency: Full Finetuning with the Same Optimizer as Pretraining Forgets Less
The paper uncovers "Optimizer-Model Consistency," a phenomenon where full fine-tuning with the same optimizer used during pretraining (e.g., AdamW-AdamW or Muon-Muon) achieves a superior learning-forgetting tradeoff compared to switching optimizers or using LoRA. It demonstrates that this consistency allows LLMs to acquire new tasks (SFT) while preserving significantly more pretrained knowledge.
TL;DR
A common assumption in LLM development is that once a model is pretrained, the optimizer used to get there doesn't matter for downstream tasks. This paper shatters that notion by introducing Optimizer-Model Consistency. The core finding: if you pretrain with AdamW, you should fine-tune with AdamW; if you pretrain with Muon, use Muon. Doing so allows the model to learn new skills (SFT) while forgetting significantly less of its original knowledge, often outperforming the industry-standard LoRA in the process.
The Problem: The Hidden "Shape" of Models
When we train an LLM, the optimizer doesn't just find a set of weights; it "sculpts" the model's internal activations. Previous research often treated SFT as a fresh start, frequently switching to different optimizers or sticking to parameter-efficient methods like LoRA to prevent "catastrophic forgetting."
However, the authors argue that different optimizers leave unique fingerprints—specifically regularization effects on activations. If the fine-tuning optimizer ignores this "shape," the updates move the model into regions that aggressively destroy pretrained features.
The Hypothesis: Convergence of Geometry
The authors propose that optimizers are not just search algorithms but builders that enforce specific norms.
- AdamW (and SignSGD) acts like an optimizer, regularizing activations toward sparsity.
- Muon acts like an optimizer, regularizing activations via the norm and pushing weights toward a higher-rank spectrum.
The "Why" is simple: If an optimizer has already shaped the model's manifold to be sparse (in the case of AdamW), using the same optimizer during SFT ensures that updates "fit" the existing curvature, leading to a much better learning-forgetting tradeoff.
Figure: Empirical evidence showing that different families of optimizers lead to drastically different activation sparsity levels.
Methodology: Proving Consistency works
The authors conducted a massive sweep across GPT-2 and Llama-2-7B models, testing various optimizers (AdamW, Muon, SGD) and LoRA. They plotted Pareto Frontiers where:
- X-axis: Forgetting (Pretraining validation loss - lower is better).
- Y-axis: Learning (SFT task validation loss - lower is better).
Key Insights:
- Consistency is King: In every task—Math, Code, or Instructions—the model that used the same optimizer as its pretraining stage outperformed others on the learning-forgetting frontier.
- The LoRA Myth: While LoRA is parameter-efficient, the authors found that full fine-tuning with the same optimizer can actually forget LESS than LoRA while learning MORE, provided the learning rate is correctly swept.
Figure: The Pareto frontier shows Full Finetuning with AdamW (the pretraining optimizer) outperforming LoRA and Muon in the math task.
A Case Study on Muon: Rote Memorization vs. Reasoning
Muon has recently gained fame for its speed in pretraining (used in models like DeepSeek-v4 and Kimi k2.5). However, this paper reveals a potential "dark side."
Through synthetic experiments with corrupted (shuffled) text, the authors found that Muon tends to rote-memorize data faster than AdamW. While this helps in pretraining and simple instruction following (Alpaca), it can be a disadvantage in reasoning tasks (Math/Code) where the model needs to learn patterns rather than just memorizing facts.
Critical Analysis & Conclusion
Takeaways
- For Developers: If you are using a pretrained base model, identify its pretraining optimizer. Avoid the temptation to automatically default to LoRA or a "fast" new optimizer if you want to preserve the model's general intelligence.
- For Researchers: The choice of SFT optimizer is a hyperparameter that must be aligned with the model's genealogy.
Limitations
The study was primarily performed on smaller scales (up to 7B parameters) due to the compute required for full pretraining sweeps. Whether these dynamics change as models reach 70B+ or 400B+ parameters (where the "basin" of the pretrained checkpoint might be deeper) remains an open question.
Final Thought
This work shifts the focus of fine-tuning from "efficiency of parameters" to "consistency of optimization." It suggests that we should treat an LLM not as a static set of weights, but as a dynamic system with a specific geometric history.
