Advancing General-Purpose Reasoning Models with Modular Gradient Surgery

TL;DR

The paper tackles the challenge of training a single general-purpose reasoning LLM across heterogeneous domains (math, instruction-following, and open-ended chat). It shows that naive strategies like Sequential RL and Mixed RL suffer from cross-domain interference at both behavioral and gradient levels, hindering performance. The authors introduce Modular Gradient Surgery (MGS), which performs gradient projection at the module level within transformers to resolve conflicts while preserving compatible updates, achieving consistent gains over baselines on Llama and Qwen across three domains and longer training. The work clarifies sources of interference in multi-domain RL and provides a scalable, efficient recipe for advancing general-purpose reasoning models with practical implications for RL post-training of large language models.

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.

Figures (11)

• Figure 1: Effectiveness of different ways for training reasoning models on multiple domains. Naive strategies, such as sequential RL training (Sequential RL), or mixing different domains in the same batch (Mixed RL), often result in limited performance across domains. We propose Modular Gradient Surgery (MGS), which resolves conflicting gradients at the module level and achieves the best multi-domain performance.
• Figure 2: Chat vs. Math capabilities for Sequential RL training, with different training steps in the first stage. We identify Forgetting and Rigidity in this type of approach.
• Figure 3: Entropy analysis on Reasoning Rigidity. Left (Chat training): entropy dynamics comparison of Math$\rightarrow$Chat and Chat-only. Right (Math training): entropy dynamics comparison of Chat$\rightarrow$Math and Math-only.
• Figure 4: Mixed RL training with different data proportions on Qwen-2.5-7B. Math performance is positively correlated with the number of Math data, while the correlation on Chat is more complex. Compared to Sequential RL, Mixed RL checkpoints generally have higher Chat scores with lower Math scores.
• Figure 5: Gradient cosine similarity between Math , Chat in Llama-3.1-8B during training. Gradient norm ratio between Math and Chat in is also shown. Gradient conflict and imbalance exist between different tasks across different training steps.