Elamur Research⌕
☰
home / papers / 2026 / Imbalanced Gradients in RL Post-Training of Multi-Task LLMs

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

Paper · 2026 · Runzhe Wu, Ankur Samanta, Ayush Jain et al.

Summary

When an LLM is post-trained on several tasks at once, the datasets are usually pooled and optimized together. This quietly assumes every task pushes the shared parameters with comparable force. The authors test that assumption for RL post-training using RLVR with GRPO on instruction-tuned Qwen2.5-3B, Qwen2.5-7B and Llama-3.2-3B. They use two task mixes: a heterogeneous one (Code, Countdown, MATH, FinQA) and a math-only one of varying difficulty (DeepScaleR, MATH, Arithmetic).

They find that gradient magnitudes differ widely between tasks, so a few tasks steer the joint update. This would be acceptable only if those tasks also improved faster, but they often do not. The authors check this in two ways: by measuring learning gain as smoothed growth in training reward, and by sampling tasks in proportion to their gradient size. Neither supports the idea that big gradients signal big gains. Training statistics such as reward, advantage and response length also fail to explain the gap. The work therefore warns against naive data mixing and argues for gradient-level balancing methods.

Key points

  • In the multi-domain mix, Code has the largest squared gradient norm throughout training, up to 15× that of MATH, yet it shows some of the smallest learning gains.
  • In the math mix, Arithmetic's squared gradient norm reaches up to 33× the others'. Under a single global learning rate, this acts like a learning rate about 5.7× higher, yet its gains largely disappear after the early phase of training.
  • Gradient-proportional sampling draws tasks via a temperature softmax over gradient norms, with each task guaranteed at least 10% probability. It helps some large-gradient tasks but generally lowers average accuracy. For Qwen-7B, the multi-domain average drops from 61.97 to 60.48 and the math average from 69.67 to 66.00; Qwen-3B on math is the only setting with a small gain (63.33 → 64.20).
  • Within a single task, gradient size does track learning gain; the relationship breaks down only when comparing across tasks. A convex-analysis argument in the appendix explains why this can happen.
  • The paper does not propose a fix. It also leaves open whether the same imbalance appears in pre-training or SFT, and suggests mirror-descent-style RL or multi-task gradient-surgery methods as possible remedies.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Columbia UniversityNew York · United States241 works · 54,053 citations🎓Cornell UniversityIthaca · United States120 works · 59,940 citations🎓Tel Aviv UniversityTel Aviv · Israel82 works · 10,655 citations🎓Technion – Israel Institute of TechnologyHaifa · Israel49 works · 5,306 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Runzhe Wu, Ankur Samanta, Ayush Jain, Scott Fujimoto, Jeongyeol Kwon, Ben Kretzu, Youliang Yu, Kaveh Hassani, Boris Vidolov and Yonathan Efroni. 2026. Imbalanced Gradients in RL Post-Training of Multi-Task LLMs. Findings of the Association for Computational Linguistics: EACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-eacl.164

@inproceedings{wu2026imbalanced,
  title = {{Imbalanced Gradients in RL Post-Training of Multi-Task LLMs}},
  author = {Runzhe Wu, Ankur Samanta, Ayush Jain, Scott Fujimoto, Jeongyeol Kwon, Ben Kretzu, Youliang Yu, Kaveh Hassani, Boris Vidolov and Yonathan Efroni},
  booktitle = {Findings of the Association for Computational Linguistics: EACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-eacl.164},
  url = {https://doi.org/10.18653/v1/2026.findings-eacl.164}
}

References

48 references, 11 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Jordan T Ash, Chicheng Zhang, Akshay Krishnamurthy, John Langford, and Alekh Agarwal. 2020. Deep batch active learning by diverse, uncertain gradient lower bounds. In International Conference on Learning Representations. find
  2. Meni Brief, Oded Ovadia, Gil Shenderovitz, Noga Ben Yoash, Rachel Lemberg, and Eitam Sheetrit. 2024. Mixing it up: The cocktail effect of multi-task fine-tuning on llm performance–a case study in finance. arXiv preprint arXiv:2410.01109. — [arXiv:2410.01109](https://arxiv.org/abs/2410.01109) find
  3. Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter,… — Language Models are Few-Shot Learners
  4. Xiaoyin Chen, Jiarui Lu, Minsu Kim, Dinghuai Zhang, Jian Tang, Alexandre Piché, Nicolas Gontier, Yoshua Bengio, and Ehsan Kamalloo. 2025. Self-evolving curriculum for llm reasoning. arXiv preprint arXiv:2505.14970. find
  5. Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. 2018. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International conference on machine learning, pages 794–803. PMLR. find
  6. Zhao Chen, Jiquan Ngiam, Yanping Huang, Thang Luong, Henrik Kretzschmar, Yuning Chai, and Dragomir Anguelov. 2020. Just pick a sign: Optimizing deep multitask models with gradient sign dropout. Advances in Neural Information Processing Systems, 33:2039–2050. find
  7. Zhiyu Chen, Wenhu Chen, Charese Smiley, Sameena Shah, Iana Borova, Dylan Langdon, Reema Moussa, Matt Beane, Ting-Hao Huang, Bryan R Routledge, and 1 others. 2021. Finqa: A dataset of numerical reasoning over financial data. In Proceedings of the 2021 Conference on Empirical Methods in Natural Langua…
  8. Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, and 1 others. 2024. The llama 3 herd of models. arXiv e-prints, pages arXiv–2407.
  9. Yonathan Efroni, Ben Kretzu, Daniel Jiang, Jalaj Bhandari, Zheqing Zhu, and Karen Ullrich. 2025. Aligned multi objective optimization. In Forty-second International Conference on Machine Learning. find
  10. Ahmed Elbakary, Chaouki Ben Issaid, Tamer ElBatt, Karim Seddik, and Mehdi Bennis. 2025. Mira: A method of federated multi-task learning for large language models. IEEE Networking Letters. find
  11. Wenfeng Feng, Chuzhan Hao, Yuewei Zhang, Yu Han, and Hao Wang. 2024. Mixture-of-loras: An efficient multitask tuning method for large language models. In Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources and Evaluation (LREC-COLING 2024), pages 1… find
  12. Zhaolin Gao, Jonathan Chang, Wenhao Zhan, Owen Oertell, Gokul Swamy, Kianté Brantley, Thorsten Joachims, Drew Bagnell, Jason D Lee, and Wen Sun. 2024. Rebel: Reinforcement learning via regressing relative rewards. Advances in Neural Information Processing Systems, 37:52354–52400. find
  13. Zhaolin Gao, Joongwon Kim, Wen Sun, Thorsten Joachims, Sid Wang, Richard Yuanzhe Pang, and Liang Tan. 2025. Prompt curriculum learning for efficient llm post-training. arXiv preprint arXiv:2510.01135. — [arXiv:2510.01135](https://arxiv.org/abs/2510.01135) find
  14. Zi Gong, Hang Yu, Cong Liao, Bingchang Liu, Chaoyu Chen, and Jianguo Li. 2024. Coba: Convergence balancer for multitask finetuning of large language models. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 8063–8077. find
  15. Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, and 1 others. 2025. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv:2501.12948. arXiv preprint Elad Hazan and 1 others. 2016. Introduction…
  16. Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the math dataset. NeurIPS.
  17. Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, and 1 others. 2024. Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06395.
  18. Sham M Kakade. 2001. A natural policy gradient. Advances in neural information processing systems, 14. find
  19. Ben Kretzu, Karen Ullrich, and Yonathan Efroni. 2025. Simple optimizers for convex aligned multi-objective optimization. arXiv preprint arXiv:2509.05811. — [arXiv:2509.05811](https://arxiv.org/abs/2509.05811) find
  20. Yiqing Liang, Jielin Qiu, Wenhao Ding, Zuxin Liu, James Tompkin, Mengdi Xu, Mengzhou Xia, Zhengzhong Tu, Laixi Shi, and Jiacheng Zhu. 2025. Modomodo: Multi-domain data mixtures for multimodal llm reinforcement learning. arXiv preprint arXiv:2505.24871. — [arXiv:2505.24871](https://arxiv.org/abs/2505.24871) find
  21. Bo Liu, Yihao Feng, Peter Stone, and Qiang Liu. 2023. Famo: Fast adaptive multitask optimization. Advances in Neural Information Processing Systems, 36:57226–57243. find
  22. Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. 2021a. Conflict-averse gradient descent for multi-task learning. Advances in Neural Information Processing Systems, 34:18878–18890. find
  23. Jiawei Liu and Lingming Zhang. 2025. Code-r1: Reproducing r1 for code with reliable rewards. find
  24. Liyang Liu, Yi Li, Zhanghui Kuang, Jing-Hao Xue, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. 2021b. Towards impartial multi-task learning. In International Conference on Learning Representations. find
  25. Michael Luo, Sijun Tan, Justin Wong, Xiaoxiang Shi, William Y. Tang, Manan Roongta, Colin Cai, Jeffrey Luo, Li Erran Li, Raluca Ada Popa, and Ion Stoica. 2025. Deepscaler: Surpassing o1-preview with a 1.5b model by scaling rl. Notion Blog. find
  26. Yuren Mao, Zekai Wang, Weiwei Liu, Xuemin Lin, and Pengtao Xie. 2022. MetaWeighting: Learning to weight tasks in multi-task learning. In Findings of the Association for Computational Linguistics: ACL 2022, pages 3436–3448, Dublin, Ireland. Association for Computational Linguistics. find
  27. Jiayi Pan, Junjie Zhang, Xingyao Wang, Lifan Yuan, Hao Peng, and Alane Suhr. 2025. Tinyzero. https://github.com/Jiayi-Pan/TinyZero. Accessed: 2025-01-24. ↗
  28. Shubham Parashar, Shurui Gui, Xiner Li, Hongyi Ling, Sushil Vemuri, Blake Olson, Eric Li, Yu Zhang, James Caverlee, Dileep Kalathil, and 1 others. 2025. Curriculum reinforcement learning from easy to hard tasks improves llm reasoning. arXiv preprint arXiv:2506.06632. find
  29. Boris Teodorovich Polyak. 1963. Gradient methods for minimizing functionals. Zhurnal vychislitel’noi matematiki i matematicheskoi fiziki, 3(4):643–653. find
  30. Zhen Qi, Jiajing Chen, Shuo Wang, Bingying Liu, Hongye Zheng, and Chihang Wang. 2024. Optimizing multi-task learning for enhanced performance in large language models. In 2024 4th International Conference on Electronic Information Engineering and Computer Communication (EIECC), pages 1179– 1183. IEE… find
  31. A Yang Qwen, Baosong Yang, B Zhang, B Hui, B Zheng, B Yu, Chengpeng Li, D Liu, F Huang, H Wei, and 1 others. 2024. Qwen2. 5 technical report. arXiv preprint. find
  32. Ozan Sener and Vladlen Koltun. 2018. Multi-task learning as multi-objective optimization. Advances in neural information processing systems, 31. find
  33. Burr Settles, Mark Craven, and Soumya Ray. 2007. Multiple-instance active learning. Advances in neural information processing systems, 20. find
  34. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, and 1 others. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300.
  35. Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. 2024. Hybridflow: A flexible and efficient rlhf framework. arXiv preprint arXiv: 2409.19256.
  36. Taiwei Shi, Yiyang Wu, Linxin Song, Tianyi Zhou, and Jieyu Zhao. 2025. Efficient reinforcement finetuning via adaptive curriculum learning. arXiv preprint arXiv:2504.05520. — [arXiv:2504.05520](https://arxiv.org/abs/2504.05520) find
  37. Richard S Sutton, Andrew G Barto, and 1 others. 1998. Reinforcement learning: An introduction, volume 1. MIT press Cambridge. find
  38. MosaicML NLP Team and 1 others. 2023. Introducing mpt-7b: A new standard for open-source, commercially usable llms. Retrieved December, 6:2023. find
  39. Kaiwen Wang, Rahul Kidambi, Ryan Sullivan, Alekh Agarwal, Christoph Dann, Andrea Michi, Marco Gelmi, Yunxuan Li, Raghav Gupta, Kumar Avinava Dubey, Alexandre Rame, Johan Ferret, Geoffrey Cideron, Le Hou, Hongkun Yu, Amr Ahmed, Aranyak Mehta, Leonard Hussenot, Olivier Bachem, and Edouard Leurent. 202… find
  40. Rongsheng Wang, Haoming Chen, Ruizhe Zhou, Yaofei Duan, Kunyan Cai, Han Ma, Jiaxi Cui, Jian Li, Patrick Cheong-Iao Pang, Yapeng Wang, and 1 others. 2023a. Aurora: Activating chinese chat capability for mixtral-8x7b sparse mixture-of-experts through instruction-tuning. arXiv preprint arXiv:2312.14557… — [arXiv:2312.14557](https://arxiv.org/abs/2312.14557) find
  41. Yizhong Wang, Hamish Ivison, Pradeep Dasigi, Jack Hessel, Tushar Khot, Khyathi Chandu, David Wadden, Kelsey MacMillan, Noah A Smith, Iz Beltagy, and 1 others. 2023b. How far can camels go? exploring the state of instruction tuning on open resources. Advances in Neural Information Processing Systems,…
  42. Zhenting Wang, Guofeng Cui, Yu-Jhe Li, Kun Wan, and Wentian Zhao. 2025. Dump: Automated distribution-level curriculum learning for rl-based llm post-training. arXiv preprint arXiv:2504.09710. — [arXiv:2504.09710](https://arxiv.org/abs/2504.09710) find
  43. Rui Yang, Xiaoman Pan, Feng Luo, Shuang Qiu, Han Zhong, Dong Yu, and Jianshu Chen. 2024. Rewards-in-context: Multi-objective alignment of foundation models with dynamic preference adjustment. In International Conference on Machine Learning, pages 56276–56297. PMLR. find
  44. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, and 1 others. 2025. Dapo: An open-source llm reinforcement learning system at scale. arXiv preprint arXiv:2503.14476.
  45. Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. 2020. Gradient surgery for multi-task learning. Advances in neural information processing systems, 33:5824– 5836.
  46. Ruiqi Zhang, Daman Arora, Song Mei, and Andrea Zanette. 2025. Speed-rl: Faster training of reasoning models via online curriculum learning. arXiv preprint arXiv:2506.09016. — [arXiv:2506.09016](https://arxiv.org/abs/2506.09016) find
  47. Tong Zhu, Daize Dong, Xiaoye Qu, Jiacheng Ruan, Wenliang Chen, and Yu Cheng. 2025. Dynamic data mixing maximizes instruction tuning for mixture-of-experts. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Te… find
  48. Stanisław Łojasiewicz. 1963. Une propriété topologique des sous-ensembles analytiques réels. In Les Équations aux Dérivées Partielles, Colloques Internationaux du CNRS, pages 87–89. Éditions du CNRS. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wu-2026-imbalanced-gradients-rl-post-training/