Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Imbalanced Gradients in RL Post-Training of Multi-Task LLMs

Runzhe Wu, Ankur Samanta, Ayush Jain, Scott Fujimoto, Jeongyeol Kwon, Ben Kretzu, Youliang Yu, Kaveh Hassani, Boris Vidolov, Yonathan Efroni

TL;DR

The paper reveals a pronounced gradient imbalance across tasks during RL post-training of multi-task LLMs, where some tasks produce much larger gradients and bias updates, yet these large gradients do not guarantee greater learning gains. Through multi-domain and single-domain experiments, the authors show that gradient magnitude correlates with gains within a task but poorly across tasks. They test gradient-proportional sampling and find no reliable improvement, and analyze several training statistics (advantage, reward, length) without finding a cross-task explanation. The work highlights the need for gradient-level corrections and suggests exploring optimization geometry approaches to mitigate cross-task bias, with implications for safer and more effective multi-task RL training of LLMs.

Abstract

Multi-task post-training of large language models (LLMs) is typically performed by mixing datasets from different tasks and optimizing them jointly. This approach implicitly assumes that all tasks contribute gradients of similar magnitudes; when this assumption fails, optimization becomes biased toward large-gradient tasks. In this paper, however, we show that this assumption fails in RL post-training: certain tasks produce significantly larger gradients, thus biasing updates toward those tasks. Such gradient imbalance would be justified only if larger gradients implied larger learning gains on the tasks (i.e., larger performance improvements) -- but we find this is not true. Large-gradient tasks can achieve similar or even much lower learning gains than small-gradient ones. Further analyses reveal that these gradient imbalances cannot be explained by typical training statistics such as training rewards or advantages, suggesting that they arise from the inherent differences between tasks. This cautions against naive dataset mixing and calls for future work on principled gradient-level corrections for LLMs.

Imbalanced Gradients in RL Post-Training of Multi-Task LLMs

TL;DR

The paper reveals a pronounced gradient imbalance across tasks during RL post-training of multi-task LLMs, where some tasks produce much larger gradients and bias updates, yet these large gradients do not guarantee greater learning gains. Through multi-domain and single-domain experiments, the authors show that gradient magnitude correlates with gains within a task but poorly across tasks. They test gradient-proportional sampling and find no reliable improvement, and analyze several training statistics (advantage, reward, length) without finding a cross-task explanation. The work highlights the need for gradient-level corrections and suggests exploring optimization geometry approaches to mitigate cross-task bias, with implications for safer and more effective multi-task RL training of LLMs.

Abstract

Multi-task post-training of large language models (LLMs) is typically performed by mixing datasets from different tasks and optimizing them jointly. This approach implicitly assumes that all tasks contribute gradients of similar magnitudes; when this assumption fails, optimization becomes biased toward large-gradient tasks. In this paper, however, we show that this assumption fails in RL post-training: certain tasks produce significantly larger gradients, thus biasing updates toward those tasks. Such gradient imbalance would be justified only if larger gradients implied larger learning gains on the tasks (i.e., larger performance improvements) -- but we find this is not true. Large-gradient tasks can achieve similar or even much lower learning gains than small-gradient ones. Further analyses reveal that these gradient imbalances cannot be explained by typical training statistics such as training rewards or advantages, suggesting that they arise from the inherent differences between tasks. This cautions against naive dataset mixing and calls for future work on principled gradient-level corrections for LLMs.
Paper Structure (27 sections, 16 equations, 7 figures, 3 tables)

This paper contains 27 sections, 16 equations, 7 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Imbalanced Gradients
  3. Multi-domain tasks.
  4. Single-domain (math) tasks.
  5. Are Gradient Imbalances Explained by Learning Gains?
  6. Quantifying Learning Gains
  7. Gradient-Proportional Sampling
  8. Summary.
  9. Are Gradient Imbalances Explained by Other Training Statistics?
  10. Advantage Function.
  11. Training Reward (Accuracy).
  12. Prompt/Response Length.
  13. Conclusion
  14. Prior Works on Multi-Task LLMs
  15. Gradient Magnitude and Learning Gains Through the Lens of Convex Analysis
  16. ...and 12 more sections

Figures (7)

  • Figure 1: Gradient imbalance and its misalignment with learning gains. Top: Qwen-7B trained on four heterogeneous tasks. The left panel shows squared gradient norms, with Code dominating throughout training. However, the middle panel reveals that Code achieves the lowest learning gains---defined as the increase in training reward (\ref{['sec:quantifying-learning-gains']}). The right panel reports training rewards for reference. Bottom: Qwen-7B trained on three math tasks. Arithmetic exhibits up to $33\times$ larger squared gradients but major learning gains vanish after early training. Together, these results indicate an implicit training bias due to imbalanced gradients which cannot be attributed to learning gains. Plots of other models are provided in \ref{['sec:additional-plots']}.
  • Figure 2: Gradient norms versus absolute advantage. Scatter plots across models (from left to right: Qwen-3B, Qwen-7B, Llama-3B) and domains (from top to down: multi-domain and math-domain). Within each task, points tend to cluster and vary consistently. Within a task, larger gradients often align with larger absolute advantages, although unnecessarily linearly. Across tasks, however, this structure disappears: tasks with the largest gradients (e.g., FinQA) do not yield higher absolute advantages than other tasks. Thus, absolute advantage may approximate gradient magnitude within a task but fails across tasks.
  • Figure 3: Prompt and response lengths across tasks. Left: Average prompt length for each task. Middle and right: Curves of response length for Qwen-3B model through training, for multi-domain and single-domain tasks respectively.
  • Figure 4: Gradient norm of different attention blocks for Qwen-3B model on multi-domain tasks. There are a total of 36 attention blocks in the model. From left to right are the gradients of the last (36th) attention block, middle (18th) attention block, and the first attention block.
  • Figure 5: (Unsquared) Gradient norm across models. From left to right are Qwen-3B, Qwen-7B, and Llama-3B. Top row presents multi-domain tasks, and bottom row presents single-domain tasks. The pattern is identical to that of the squared norm in \ref{['fig:intro']} and \ref{['fig:intro-additional']}.
  • ...and 2 more figures

Theorems & Definitions (2)

  • Definition 1: Smoothness
  • Definition 2: Polyak-Łojasiewicz