Imbalanced Gradients in RL Post-Training of Multi-Task LLMs
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866911233602486272 |
|---|---|
| author | Wu, Runzhe Samanta, Ankur Jain, Ayush Fujimoto, Scott Kwon, Jeongyeol Kretzu, Ben Yu, Youliang Hassani, Kaveh Vidolov, Boris Efroni, Yonathan |
| author_facet | Wu, Runzhe Samanta, Ankur Jain, Ayush Fujimoto, Scott Kwon, Jeongyeol Kretzu, Ben Yu, Youliang Hassani, Kaveh Vidolov, Boris Efroni, Yonathan |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2510_19178 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Imbalanced Gradients in RL Post-Training of Multi-Task LLMs Wu, Runzhe Samanta, Ankur Jain, Ayush Fujimoto, Scott Kwon, Jeongyeol Kretzu, Ben Yu, Youliang Hassani, Kaveh Vidolov, Boris Efroni, Yonathan Machine Learning Artificial Intelligence 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. |
| title | Imbalanced Gradients in RL Post-Training of Multi-Task LLMs |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2510.19178 |