Advancing General-Purpose Reasoning Models with Modular Gradient Surgery

AI summary

Overview

• Research area: Natural language processing / LLM post-training — specifically reinforcement learning (RL) for large reasoning models (LRMs), multi-domain (multi-task) RL, and gradient-conflict manipulation.
• Technical level: Advanced. The paper assumes familiarity with policy-gradient RL for LLMs (GRPO/DAPO), transformer internals (attention, MLP, LayerNorm blocks), and gradient-surgery methods such as PCGrad.
• Scope: This paper diagnoses why training a single general-purpose reasoning model across math, chat, and instruction following fails under Sequential and Mixed RL, and proposes Modular Gradient Surgery (MGS), which resolves conflicting task gradients separately inside each transformer module.

What This Paper Is About

Practitioners want one reasoning model that is simultaneously good at math, open-ended chat, and instruction following, but naively combining these domains — either one after another (Sequential RL) or in the same training batch (Mixed RL) — produces limited gains because optimizing one domain damages another. The authors systematically characterize this interference and propose MGS, a method that detects and removes conflicting gradient components inside individual transformer modules rather than across the entire model at once.

Key Contributions

1. A systematic study of multi-domain RL for reasoning, identifying two distinct failure modes: Mode Interference in Sequential RL (Forgetting of earlier domains plus Rigidity on later ones) and gradient conflicts in Mixed RL.
2. Modular Gradient Surgery (MGS), a new method that partitions model parameters into disjoint modules and applies PCGrad-style gradient projection independently within each module, so that a conflict in one block does not force a projection across the whole network.
3. An extensive evaluation across two model families (Llama-3.1-8B and Qwen-2.5-7B) and up to three domains (Math, Chat, IF), showing MGS beats Naïve Mixing by 4.3 points (16.6%) on Llama and 4.5 points (11.1%) on Qwen, and outperforming Sequential RL, Global Gradient Surgery, Normalized Advantage, and Model Merging.
4. Evidence that the method scales: MGS improves further with doubled training steps (an additional 3% relative gain on the Math benchmark) and generalizes to a third task, with a 19.4% relative average improvement on Llama-3.1-8B in the three-task setting.

Main Findings

• Sequential RL suffers "Mode Interference" with two symptoms. Forgetting: training a second domain reduces performance on the first (e.g., continuing with Math after Chat training causes a sharp drop in Chat performance). Rigidity: prior training constrains later gains — Chat→Math fails to reach Math-only performance, and Math→Chat underperforms Chat-only.
• The interference is asymmetric. The gap between Math→Chat and Chat-only on Chat is -9.6, much larger than the -1.3 gap between Chat→Math and Math-only on Math. Ordering matters: Chat→Math performs slightly better on average than Math→Chat.
• Continuous sequencing keeps degrading. Training Chat→Math→IF further downgrades performance on both Chat and Math.
• Entropy dynamics explain Rigidity. Entropy typically decreases during Math training and increases during Chat training. When Math is trained first, entropy on Chat data stays below that of Chat-only training; when Chat is trained first, the resulting high-entropy prior helps exploration on the later Math stage.
• Mixed RL suffers gradient conflicts, not just data imbalance. Gradient cosine similarity between Math and Chat gradients (measured on Llama-3.1-8B) is frequently negative during training, and gradient norms are imbalanced across domains. Increasing the Math proportion in the mixture monotonically improves Math, but even a 90% Math / 10% Chat mixture still lags a Math-only model trained for 2 epochs. Compared with Sequential RL, Mixed RL checkpoints generally show higher Chat scores but lower Math scores.
• MGS wins on the two-domain and three-domain settings. On Llama with 1:1 Math/Chat, MGS reaches a total average of 30.2 versus 25.9 for Naïve Mixing (+4.3); on Qwen it reaches 44.9 versus 40.4 (+4.5). On the three-task Llama setting, MGS reaches an average of 32.6 versus 27.3 for Naïve Mixing and 30.3 for Global Surgery.
• Global Gradient Surgery (GGS) is overly conservative. Applying surgery globally is worse than MGS and sometimes worse than Naïve Mixing, apparently because a conflict in one less-critical module forces projection of the entire model's gradient and can stall learning in modules where no conflict existed.
• Module choice matters. Excluding different module families from the surgery harms different capabilities: excluding MLP hurts Chat most, while excluding Attn or LayerNorm hurts Math. Excluding LayerNorm causes a dramatic drop on both axes despite LayerNorm containing the fewest parameters.
• Gradient norm is a partial heuristic. Restricting MGS to the top 10% of modules by gradient norm shows a trend similar to module-family exclusion, but restricting to the top 1% boosts Chat while Math performance drops to only marginal improvement over the base model — suggesting Math relies on coordinated gradients across a broad set of modules.
• MGS surpasses prior general-purpose models. In Table 2 it outperforms RLMT and General Reasoner across Math, Chat, and IF.
• Overhead is small. The authors report negligible overhead when using high-performance parallel training frameworks such as FSDP (detailed analysis referenced in Appendix F).

Methodology in Plain English

The authors first ran controlled comparisons to understand the problem. They trained Qwen-2.5-7B on two domains (Math and Chat) in both orders, giving each domain 2 epochs and resetting the reference policy between stages, then measured how performance shifted. They also trained on mixed batches with 1:1, 1:9, and 9:1 Math-to-Chat ratios, keeping batch size and total optimization steps identical to the sequential runs so the comparison was fair. To find the cause of the mixed-training problem, they measured the cosine similarity between the gradients produced by Math samples and Chat samples, and the ratio of their norms, across training steps.

The fix builds on an existing idea from multi-task learning called gradient surgery (PCGrad): if two task gradients point in conflicting directions — detected by a negative dot product — each gradient is projected onto the orthogonal complement of the other, removing the component that would harm the other task. The authors' observation is that transformer components are specialized (MLP layers are associated with knowledge storage and memory; attention layers are associated with information routing and reasoning), so conflicts are localized rather than uniform. MGS therefore splits the parameters into disjoint modules (for example, attention and MLP blocks per layer, iterated over model.named_parameters), flattens gradients within each module, applies the PCGrad projection separately for each module where the within-module dot product is negative, and sums the resulting per-task gradients. This prevents a conflict in one block from triggering a projection of gradients in blocks that had no conflict at all.

Training itself uses a token-level GRPO objective as used in DAPO, with batch size 64, 8 samples per prompt (group size 8), max prompt length 1024 for Math and 896 for Chat/IF/Mixed, max response length 4096, 2 epochs, learning rate 1e-6, weight decay 0.01, a constant scheduler with no warmup, and KL coefficient 0.001. Math data comes from SimpleRL-Zoo (8,523 prompts) with SimpleRL-Hard for Qwen and SimpleRL-Easy for Llama; Chat data is Wildchat (7,544 prompts) scored by Skywork-Reward-V2-Llama-3.1-8B; IF data is sampled from Nemotron (7,930 of 56,339 original points, sampled by number of constraints) scored with IFEval scoring utilities. Evaluation spans MATH500, GSM8K, AIME24, WildBench, AlpacaEval V2, Creative Writing V3, IFEval, IFBench, MMLU redux, PopQA, and Zebra Logic.

Why This Matters

• Impact on research: The paper clarifies why multi-domain RL post-training resists naive scaling, separating behavioral interference (forgetting/rigidity) from gradient interference, and claims to be the first to apply gradient manipulation in LLM RL post-training. It provides a diagnostic toolkit (gradient cosine similarity, per-module conflict analysis) that other multi-task RL efforts can reuse.
• Real-world applications:
  • Deploying a single assistant model that handles quantitative reasoning and open-ended conversation without maintaining separate specialists.
  • Math and STEM tutoring systems that must also converse naturally and follow formatting constraints.
  • Instruction-following agents where long chain-of-thought reasoning must coexist with strict output-format compliance.
  • Reducing the cost of post-training pipelines by replacing complex sequential pipelines or distillation with a single mixed-domain RL run.
• Industry relevance: The method is reported to add negligible overhead on high-performance parallel frameworks such as FSDP, and it reduces the operational burden of training one generalist model instead of several domain experts. The work was conducted with Baidu Inc. and the University of Alberta, with code and a project website released.

Future Directions

1. Better module-selection heuristics. The gradient-norm ablation shows that top-1% norm modules help Chat but hurt Math; the paper explicitly states that norm magnitude as a selection heuristic "warrants further investigation."
2. Scaling to more domains and longer training. Only three domains (Math, Chat, IF) are tested, and prolonging training is tested once (doubled steps). Whether MGS continues to hold at larger domain counts, longer horizons, or larger models is not established.
3. Understanding why LayerNorm conflicts are so consequential. Excluding LayerNorm from surgery causes the largest drop in both capabilities despite it holding the fewest parameters, and the paper offers only a hypothesis about feature-statistic regulation rather than a mechanism.
4. Reconciling with alternative multi-task strategies. The paper positions MGS against model merging and advantage normalization but notes that other recent multi-domain LLM RL efforts rely on complex pipelines or distillation — how MGS composes with those pipelines is left open.

Target Audience

Researchers and engineers working on LLM post-training and reinforcement learning — particularly those building multi-task or multi-domain RL pipelines, studying gradient conflict and catastrophic forgetting, or trying to train a single generalist reasoning model instead of one model per domain. Readers without a background in policy-gradient RL and transformer module anatomy will find the method sections difficult, though the diagnostic findings about Sequential and Mixed RL are accessible to a broader applied-ML audience.

Authors’ abstract

Reinforcement learning (RL) has played a central role in recent advances in large reasoning models (LRMs), yielding strong gains in verifiable and open-ended reasoning. However, training a single general-purpose LRM across diverse domains remains challenging due to pronounced domain heterogeneity. Through a systematic study of two widely used strategies, Sequential RL and Mixed RL, we find that both incur substantial cross-domain interference at the behavioral and gradient levels, resulting in limited overall gains. To address these challenges, we introduce **M**odular **G**radient **S**urgery (**MGS**), which resolves gradient conflicts at the module level within the transformer. When applied to Llama and Qwen models, MGS achieves average improvements of 4.3 (16.6\%) and 4.5 (11.1\%) points, respectively, over standard multi-task RL across three representative domains (math, general chat, and instruction following). Further analysis demonstrates that MGS remains effective under prolonged training. Overall, our study clarifies the sources of interference in multi-domain RL and presents an effective solution for training general-purpose LRMs.