[논문 리뷰] Multi-Task GRPO: Reliable LLM Reasoning Across Tasks

Shyam Sundhar Ramesh, Xiaotong Ji외 6명|arXiv (Cornell University)|2026. 02. 05.

Multimodal Machine Learning Applications인용 수 0

한 줄 요약

MT-GRPO는 GRPO 기반 RL 사후 학습에서 여러 작업에 걸친 견고하고 균형 잡힌 추론을 달성하기 위해 개선 인식형 작업 재가중과 비율 보존 샘플러를 도입하여, 최악의 작업 정확도에서 기준선보다 우수하게 성능을 발휘하면서 평균 성능도 경쟁력을 유지합니다.

ABSTRACT

RL-based post-training with GRPO is widely used to improve large language models on individual reasoning tasks. However, real-world deployment requires reliable performance across diverse tasks. A straightforward multi-task adaptation of GRPO often leads to imbalanced outcomes, with some tasks dominating optimization while others stagnate. Moreover, tasks can vary widely in how frequently prompts yield zero advantages (and thus zero gradients), which further distorts their effective contribution to the optimization signal. To address these issues, we propose a novel Multi-Task GRPO (MT-GRPO) algorithm that (i) dynamically adapts task weights to explicitly optimize worst-task performance and promote balanced progress across tasks, and (ii) introduces a ratio-preserving sampler to ensure task-wise policy gradients reflect the adapted weights. Experiments on both 3-task and 9-task settings show that MT-GRPO consistently outperforms baselines in worst-task accuracy. In particular, MT-GRPO achieves 16-28% and 6% absolute improvement on worst-task performance over standard GRPO and DAPO, respectively, while maintaining competitive average accuracy. Moreover, MT-GRPO requires 50% fewer training steps to reach 50% worst-task accuracy in the 3-task setting, demonstrating substantially improved efficiency in achieving reliable performance across tasks.

연구 동기 및 목표

• RL 사후 학습에서 다양한 추론 작업 간 균형 있는 역량을 증진한다.
• 평균 작업 성능을 보존하면서 최악의 작업 성능을 직접 최적화한다.
• 단순 다중 작업 GRPO에서 발생하는 제로 그래디언트 프롬프트 및 작업 간섭을 해결한다.
• 학습된 작업 가중치를 실제 그래디언트 기여와 정렬시키는 메커니즘을 도입한다.

제안 방법

• 두 가지 핵심 아이디어: (i) 약하거나 천천히 개선되는 작업에 우선순위를 두는 개선 인식형 작업 재가중 및 (ii) 재가중된 작업 가중치를 그래디언트에 반영하도록 비율 보존 배치 구성 메커니즘.
• 평가 목표를 평균 성능과 강건성 사이의 균형을 맞추기 위한 제약된 최소-최대 문제로 형식화하고 그 라그랑주 RELAXATION(Eq. 5)을 제시한다.
• 정책 최적화와 적응형 작업 가중치를 교차 결합하는 업데이트 규칙은 교대 단계로 구성된다: z-가중 GRPO 그래디언트를 이용한 theta 업데이트(Eq. 6)와 개선 신호를 통해 z를 조정하는 xi 업데이트(Eq. 7).
• 개선 인식형 가중치 업데이트(IWU)는 작업 개선 I_k^(t)와 작업 보상의 결합 신호를 사용하여 재가중화를 안정화한다(서브루틴 1).
• 비율 보존(RP) 샘플러는 배치의 목표 작업 비율을 학습된 가중치와 일치시키도록 사후 필터링된 배치에서 비율을 강제하여 제로-그래디언트 샘플링 문제를 완화한다(알고리즘 2 및 Sec. 5의 논의).

Figure 1: GRPO assigns uniform task weights and samples without regard to task difficulty or zero-gradient rates. Consequently, easy tasks (Countdown) dominate while harder tasks (ARC, Zebra) lag, and effective gradient flow is skewed by varying zero-gradient rates ( $\otimes$ marks high zero-gradie

실험 결과

연구 질문

• RQ1 강건성 인식 다중 작업 목표가 평균 성능을 희생하지 않고 최악의 작업 성능을 개선할 수 있는가?
• RQ2 작업 가중치를 현재 성능과 개선 궤적 모두를 반영하도록 업데이트할 수 있는가?
• RQ3 다양한 제로-그래디언트 비율을 고려할 때 목표 작업 비율에 충실하도록 배치를 구성할 수 있는가?
• RQ4 개선 인식형 재가중 및 비율 보존 샘플링이 더 큰 작업 집합에서도 신뢰성을 유지하며 확장될 수 있는가?

주요 결과

• MT-GRPO는 baselines(GRPO, DAPO, SEC-DAPO)에 비해 실험 전반에서 최악의 작업 정확도를 일관되게 향상시켰다.
• 3-task 설정에서 MT-GRPO는 표준 GRPO 대비 최악의 작업 성능에서 절대 16–28% 개선했으며, DAPO 대비 6% 개선된 평균 정확도를 유지했다.
• MT-GRPO는 3-task 설정에서 학습 단계의 대략 절반 수준에서 최악의 작업 정확도 50%에 도달한다.
• 개선 인식형 재가중 체계가 단일 최악 작업에 대한 가중치 붕괴를 감소시키고 성능 향상이 부진한 작업으로의 최적화를 유도한다.
• 비율 보존 샘플러는 realized 배치 비율을 학습된 작업 가중치와 일치시키며 각 작업으로부터의 효과적인 그래디언트 기여를 보장한다.
• 9개 작업에 대한 실험은 더 큰 람다(lambda)가 최악의 작업 개선을 강화하는 반면 평균 성능은 다소 감소할 수 있음을 보여주어 조절 가능한 트레이드오프를 시사한다.

Figure 2: In strict worst-task optimization ( $\varepsilon=0$ ), task weights rapidly collapse to the current worst task and oscillate as the worst task shifts, resulting in near-zero weighting of Countdown.