💬 NLP

Advancing General-Purpose Reasoning Models with Modular Gradient Surgery

To address the issue of cross-domain interference in multi-task reinforcement learning, this paper introduces **Modular Gradient Surgery (MGS)**, a method that resolves gradient conflicts at the transformer module level to significantly improve the performance of general-purpose large reasoning models across diverse domains.

Original authors: Min Cai, Yu Liang, Longzheng Wang, Yan Wang, Yueyang Zhang, Long Xia, Zhiyuan Sun, Xi Ye, Daiting Shi

Published 2026-02-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Min Cai, Yu Liang, Longzheng Wang, Yan Wang, Yueyang Zhang, Long Xia, Zhiyuan Sun, Xi Ye, Daiting Shi

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 are trying to train a world-class polymath—someone who needs to be a genius at high-level mathematics, a charming conversationalist, and a perfect assistant who follows complex instructions.

This paper, "Modular Gradient Surgery," tackles the exact problem researchers face when training "Large Reasoning Models" (the AI brains behind things like ChatGPT or DeepSeek).

Here is the breakdown of the problem and their brilliant solution using everyday analogies.

1. The Problem: The "Conflicting Coach" Dilemma

When researchers try to train an AI to do multiple things at once, they usually use one of two methods, both of which have major flaws:

The "One-at-a-Time" Method (Sequential RL): This is like training an athlete by spending all January on math, all February on social skills, and all March on following orders.
- The Flaw: By the time you get to March, the athlete has "forgotten" how to do math (this is called Forgetting), or they’ve become so rigid in their math habits that they can't learn social skills effectively (this is called Rigidity).
The "Chaos in the Classroom" Method (Mixed RL): This is like throwing a math textbook, a poetry book, and a rulebook into a room and telling the student, "Learn everything at the same time!"
- The Flaw: The instructions start to contradict each other. The "math" part of the brain wants to move in one direction, while the "chat" part wants to move in another. In AI terms, their gradients (the directions they learn in) are fighting each other. It’s like having two coaches screaming different directions at the same time; the student ends up confused and mediocre at everything.

2. The Discovery: The Brain is Modular

The researchers realized something crucial: An AI isn't just one giant, uniform blob of intelligence. It is built like a high-tech machine with different parts:

The "Memory" Modules (MLP layers): These act like the library, storing facts and knowledge.
The "Logic" Modules (Attention layers): These act like the reasoning engine, figuring out how to connect ideas.
The "Regulators" (LayerNorm): These act like the nervous system, keeping everything stable.

They found that when the AI gets "confused," the fight isn't happening everywhere. It’s usually localized. The "math" coach and the "chat" coach might be fighting over the Logic module, but they actually agree perfectly on the Memory module.

3. The Solution: "Modular Gradient Surgery" (MGS)

Instead of trying to force the entire brain to agree (which is too aggressive and slows everything down), the researchers invented Modular Gradient Surgery.

The Analogy: The Specialized Surgeon
Imagine a patient has a conflict in their body—their heart wants to beat fast, but their lungs want to breathe slowly. A "Global" doctor might try to sedate the entire body to stop the conflict, which makes the patient useless.

A "Modular Surgeon" is much smarter. They look at the body part-by-part.

They check the Heart: "Are the signals conflicting here?" If yes, they perform a tiny "surgery" to remove only the conflicting part of the signal.
They check the Lungs: "Are the signals conflicting here?" If the signals are already working together, they leave them alone and let them learn at full speed.

By performing "surgery" only on the specific modules where the "coaches" are fighting, the AI can learn math and chat simultaneously without one destroying the other.

4. The Result: A True Polymath

The results were impressive. When they applied this "surgery" to famous AI models (like Llama and Qwen):

Better at everything: The models didn't just get better at one thing; they became better at math, chat, and following instructions all at once.
No more forgetting: Unlike the "one-at-a-time" method, the AI didn't lose its math skills while learning to be a better conversationalist.
Scalable: The more they trained, the better it worked.

In short: Instead of forcing an AI to choose between being a mathematician or a poet, Modular Gradient Surgery allows it to be both, by surgically resolving the "arguments" happening inside its digital brain.

Technical Summary: Advancing General-Purpose Reasoning Models with Modular Gradient Surgery (MGS)

1. Problem Statement

The recent success of Large Reasoning Models (LRMs) has been largely driven by Reinforcement Learning (RL), particularly Reinforcement Learning with Verifiable Rewards (RLVR) in domains like mathematics and coding. However, a significant challenge remains: training a single general-purpose LRM that excels across diverse, heterogeneous domains (e.g., combining rule-based math rewards with preference-based chat rewards).

The authors identify two primary failure modes in current multi-domain RL strategies:

Sequential RL (Mode Interference): Training on domains one after another leads to Forgetting (performance degradation on previously learned domains) and Reasoning Rigidity (the model becomes too constrained by the first domain to learn the second effectively).
Mixed RL (Gradient Conflicts): Mixing multiple domains in a single training batch leads to gradient conflicts, where the update direction for one task directly opposes and cancels out the update for another, preventing the model from reaching the performance levels of single-domain experts.

2. Methodology: Modular Gradient Surgery (MGS)

To resolve these conflicts, the authors propose Modular Gradient Surgery (MGS). While traditional gradient surgery methods (like PCGrad) treat the entire model as a single monolithic vector, MGS leverages the inherent modularity of the Transformer architecture.

Key technical components include:

Module-Level Granularity: The model parameters ( $\theta$ ) are partitioned into disjoint modules ( $\Phi$ ), such as Attention layers, MLP layers, and LayerNorm.
Localized Conflict Resolution: Instead of applying a global projection to the entire gradient vector, MGS applies gradient surgery independently to each module. If two task gradients ( $v_i, v_j$ ) within a specific module have a negative cosine similarity ( $v_i \cdot v_j < 0$ ), the conflicting component is projected onto the orthogonal complement of the other.
Preserving Compatible Updates: By operating at the module level, MGS ensures that if a conflict exists in the MLP layers but not in the Attention layers, only the MLP gradients are modified. This prevents the "overly conservative" behavior of Global Gradient Surgery (GGS), which might stall learning in non-conflicting modules.

3. Key Contributions

Systematic Analysis of Multi-Domain RL: The paper provides a rigorous study of why Sequential and Mixed RL fail, identifying "Mode Interference," "Forgetting," and "Rigidity" as core obstacles.
Introduction of MGS: A novel, module-aware gradient manipulation technique specifically designed for the architectural structure of Transformers.
Empirical Validation: Extensive testing across multiple model families (Llama-3.1, Qwen-2.5) and diverse task sets (Math, General Chat, and Instruction Following).

4. Results

The experimental results demonstrate that MGS significantly outperforms existing baselines:

Performance Gains: MGS achieved average improvements of 4.3 points (16.6%) on Llama and 4.5 points (11.1%) on Qwen over standard multi-task RL across three domains.
Superiority over Baselines: MGS consistently outperformed Sequential RL, Naïve Mixing, Model Merging, and Global Gradient Surgery (GGS). Notably, while GGS often proved too conservative, MGS maintained high performance across all tasks.
Scalability: The effectiveness of MGS scales with increased training steps and the number of tasks. When expanded to three tasks (Math, Chat, and Instruction Following), MGS continued to provide the most balanced and highest average performance.
Mechanistic Insights: Ablation studies revealed that LayerNorm is a critical module for resolving conflicts, and that high-norm modules are vital for chat, while math requires coordinated alignment across a broader range of modules.

5. Significance

This work is significant because it provides a principled, efficient, and scalable solution for building general-purpose reasoning models. By moving from global to modular optimization, it bridges the gap between specialized "expert" models and versatile "generalist" models. Furthermore, it represents the first successful application of gradient manipulation techniques to the post-training RL phase of Large Language Models, offering a new direction for the development of highly capable, multi-functional AI agents.