🤖 machine learning

Muon-OGD: Muon-based Spectral Orthogonal Gradient Projection for LLM Continual Learning

The paper proposes Muon-OGD, a continual learning framework for large language models that mitigates catastrophic forgetting by integrating Muon-style spectral-norm geometry with orthogonal gradient projection, thereby improving the stability-plasticity trade-off and outperforming existing Euclidean-norm-based methods on diverse benchmarks.

Original authors: Binghang Lu, Zheyuan Deng, Runyu Zhang, Bing Hu, Yunhan Zhao, Yuan Tian, Changhong Mou, Guang Lin, Xiaomin Li

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Binghang Lu, Zheyuan Deng, Runyu Zhang, Bing Hu, Yunhan Zhao, Yuan Tian, Changhong Mou, Guang Lin, Xiaomin Li

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

The Big Problem: The "Sponge" Effect

Imagine a student (a Large Language Model, or LLM) who is incredibly smart. They learn to speak French, then they learn to speak Spanish. But here's the catch: as soon as they start learning Spanish, they start forgetting French. This is called catastrophic forgetting.

In the world of AI, when a model learns a new task, it often overwrites the "neural pathways" it used for the old task. The model becomes a sponge that soaks up new information but squeezes out the old stuff.

The Old Solution: The "Traffic Cop"

Scientists have tried to fix this using a method called Orthogonal Gradient Descent (OGD).

Think of the model's knowledge as a giant, multi-dimensional map.

Old Tasks: The "French" knowledge occupies a specific lane on this map.
New Task: The "Spanish" learning wants to drive a car (the update) onto the map.

The old method acts like a Traffic Cop. It says, "You can drive anywhere you want to learn Spanish, but you are strictly forbidden from driving into the 'French' lane." It does this by projecting the new learning path so it stays perfectly perpendicular (at a 90-degree angle) to the old path.

The Flaw: This Traffic Cop assumes the map is flat and uniform, like a standard grid (Euclidean geometry). It measures distance using the "Frobenius norm," which is like measuring distance with a standard ruler. However, the paper argues that the internal structure of these AI models isn't a flat grid; it's more like a complex, curved landscape where some directions are "heavier" or more important than others.

The New Idea: The "Mountain Guide" (Muon-OGD)

The authors noticed a new optimizer called Muon that works differently. Instead of using a standard ruler, Muon uses a spectral norm.

The Analogy:
Imagine you are hiking up a mountain.

The Old Way (Frobenius): You measure your steps with a flat tape measure on the ground. You think you are moving efficiently, but you might be walking up a steep cliff face because you didn't account for the terrain's shape.
The Muon Way (Spectral Norm): You have a Mountain Guide who understands the 3D shape of the terrain. The guide knows that some paths are "steep" (high impact) and some are "gentle" (low impact). The guide tells you to take steps that respect the shape of the mountain, not just the flat distance. This keeps you stable and prevents you from sliding down the wrong side.

Muon-OGD combines these two ideas:

It keeps the Traffic Cop (to protect the old French knowledge).
But it hires the Mountain Guide (Muon) to decide how to take the steps.

Instead of just saying "don't go into the French lane," the new method says, "Don't go into the French lane, and also, take steps that respect the 3D shape of the mountain so you don't slip."

How It Works (The "Dual" Dance)

The paper describes a mathematical process to make this happen efficiently:

The Setup: The model identifies the "protected lanes" (the French knowledge) that must not be touched.
The Correction: Before taking a step, the model calculates a "corrected path." It asks: "If I take this step, will I accidentally bump into the French lane?"
The Iteration: It doesn't just guess once. It runs a quick, internal loop (called dual iterations) to fine-tune the step. It's like a dancer adjusting their foot placement repeatedly in a split second to ensure they don't step on their partner's toes while moving in the most efficient direction possible.
The Sign: Finally, it uses a mathematical trick (Newton-Schulz iteration) to snap the step into the perfect "orthogonal" shape, ensuring it's the most efficient move possible without breaking the rules.

The Results: A Better Student

The authors tested this new method (Muon-OGD) against the old methods on various "school subjects" (benchmarks):

Standard Tests: Text classification (like sorting news articles).
Harder Tests: A 15-task stream (learning 15 different topics in a row).
Specialized Tests: Coding, Math, and Medical reasoning.

The Outcome:
In every test, the Muon-OGD student performed better.

They learned the new subjects (plasticity) just as well as the others.
But they forgot much less of the old subjects (stability).
They were particularly good at keeping their skills in Coding, Math, and Medicine intact while learning new things, whereas the old methods started to lose their previous skills.

Summary

The paper claims that by changing how we measure "distance" in the AI's brain—from a flat ruler to a 3D mountain guide—and combining that with strict rules to protect old knowledge, we can teach AI new things without making it forget what it already knows. It's a more stable, efficient way to keep learning forever.

Technical Summary: Muon-OGD

Problem Statement
The central challenge addressed in this work is catastrophic forgetting in Large Language Models (LLMs) during continual learning. When LLMs are sequentially fine-tuned on new tasks, their performance on previously learned tasks often degrades rapidly. Existing projection-based continual learning methods (e.g., OGD, GEM, O-LoRA) mitigate this by restricting parameter updates to subspaces orthogonal to directions associated with past tasks. However, these methods are typically formulated under Euclidean parameter geometry, where update magnitudes and orthogonality are governed by the Frobenius norm.

The authors identify a fundamental mismatch: modern optimizers for matrix-valued parameters, such as the recently proposed Muon optimizer, operate under spectral-norm (operator-norm) geometry. Muon applies orthogonalized matrix updates (approximated via Newton–Schulz iterations) and can be interpreted as performing steepest descent under a spectral-norm constraint. The paper argues that relying on Frobenius-norm geometry for continual learning constraints is suboptimal for matrix-valued LLM parameters, as it does not align with the geometry of the underlying optimization dynamics.

Methodology: Muon-OGD
The paper proposes Muon-OGD (Muon-based Spectral Orthogonal Gradient Projection), a framework that integrates spectral-norm-aware optimization with projection-based continual learning constraints.

Formulation: Instead of the standard Frobenius-norm constrained problem, Muon-OGD formulates the update $\Delta$ as a spectral-norm-constrained optimization problem:
$\min_{\Delta \in \mathbb{R}^{m \times n}} \langle G, \Delta \rangle \quad \text{s.t.} \quad \|\Delta\|_2 \leq \eta, \quad \langle C_i, \Delta \rangle = 0, \forall i=1,\dots,k$
Here, $G$ is the gradient (or momentum) matrix, $\|\cdot\|_2$ is the spectral norm, $\eta$ is the step size, and $\{C_i\}$ are constraint directions (e.g., gradients from past tasks) that must remain orthogonal to the update.
Dual Optimization: Since the problem lacks a closed-form solution, the authors derive a dual formulation. The dual objective involves minimizing the nuclear norm of a shifted matrix $H = G + \sum \lambda_i C_i$ . The optimal update is recovered via the matrix sign function: $\Delta^* = -\eta \cdot \text{msgn}(H)$ .
Algorithm: The method solves the dual problem iteratively using subgradient descent on the dual variables $\lambda$ .
- Inner Loop: Iteratively updates $\lambda$ to enforce orthogonality constraints.
- Matrix Sign Approximation: The matrix sign function $\text{msgn}(H)$ is efficiently approximated using a small number of Newton–Schulz iterations (specifically 5 iterations, denoted as $NS_5$ ), avoiding the computational cost of explicit SVD.
- Low-Rank Instantiation: For efficiency, the authors propose a specific instantiation where constraints are derived from the top- $k$ singular vectors of past tasks, allowing the dual variables to be organized into a low-rank matrix $\Lambda$ , reducing the dimensionality of the dual problem.

Key Contributions

Spectral-Norm Geometry for CL: The paper formulates projection-based continual learning under spectral-norm geometry, aligning the stability constraints of continual learning with the operator-norm geometry of modern optimizers like Muon.
Practical Update Rule: It derives a practical algorithm (Muon-OGD) that combines past-task orthogonality constraints with Muon-style matrix-sign updates, solving the constrained problem efficiently via a dual formulation and Newton–Schulz approximations.
Empirical Validation: The framework is instantiated for LLM continual fine-tuning and evaluated across diverse benchmarks, demonstrating consistent improvements over Frobenius-norm-based counterparts.

Experimental Results
The authors evaluate Muon-OGD on:

Standard Benchmarks: 5-task and 15-task continual learning benchmarks using T5-Large.
TRACE Benchmark: A benchmark for instruction-tuned LLMs (using LLaMA-2-7B-Chat) covering diverse skills like coding, math, and multilingual understanding.
Domain-Specific Curricula: Sequential learning on Coding (BigCodeBench) $\to$ Math (GSM8K) $\to$ Medical (HuatuoGPT-o1) using Qwen2.5 and LLaMA3.2 models.

Findings:

Performance: Muon-OGD consistently outperforms sequential fine-tuning (SeqFT) and competitive orthogonal-gradient baselines (e.g., O-LoRA, Sculpting Subspaces, OSFT) in terms of Average Accuracy (AA) and Backward Transfer (BT).
- On the T5-Large 15-task benchmark, Muon-OGD achieved 72.0% average accuracy, surpassing O-LoRA (69.6%) and OSFT (71.3%).
- On the TRACE benchmark, Muon-OGD achieved the highest AA (49.4%) and lowest forgetting rate (-4.7%) among compared methods.
Scalability: The method remains computationally scalable, leveraging the efficiency of Newton–Schulz iterations and low-rank constraint representations.
Architecture Agnosticism: The approach is effective across both encoder-decoder (T5) and decoder-only (LLaMA, Qwen) architectures.

Significance and Claims
The paper claims that spectral-norm-aware update geometry provides a practical and effective alternative to Frobenius-norm projection for continual learning in LLMs. By bridging the gap between the geometry of modern optimizers (like Muon) and the stability requirements of continual learning, Muon-OGD improves the stability-plasticity trade-off. The authors suggest that existing projection-based methods can be systematically upgraded to operate under spectral-norm geometry to better preserve the conditioning and stability benefits observed in Muon while preventing catastrophic forgetting. The work does not claim to solve all continual learning challenges but offers a principled geometric alternative that yields empirical gains in retention and adaptation across various domains and model scales.