Omni-Thinker: Scaling Multi-Task RL in LLMs with Hybrid Reward and Task Scheduling

TL;DR

Omni-Thinker introduces a unified RL framework that enhances multi-task learning in LLMs by combining hybrid rewards with task scheduling based on backward transfer, leading to improved performance across diverse domains.

Contribution

The paper presents a novel RL framework that integrates hybrid rewards and BWT-guided scheduling to scale LLMs for multi-task learning, addressing forgetting and performance issues.

Findings

Achieved 6.2% gain over joint training

Achieved 12.4% gain over model merging

Accurate prediction of curriculum outcomes using BWT assumptions

Abstract

The pursuit of general-purpose artificial intelligence depends on large language models (LLMs) that can handle both structured reasoning and open-ended generation. We present Omni-Thinker, a unified reinforcement learning (RL) framework that scales LLMs across diverse tasks by combining hybrid rewards with backward-transfer-guided scheduling. Hybrid rewards integrate rule-based verifiable signals with preference-based evaluations from an LLM-as-a-Judge, enabling learning in both deterministic and subjective domains. Our scheduler orders tasks according to accuracy backward transfer (BWT), reducing forgetting and improving multi-task performance. Experiments across four domains show gains of 6.2% over joint training and 12.4% over model merging. Moreover, we demonstrate that simple assumptions on accuracy transfer yield accurate predictions of curriculum outcomes, with entropy dynamics explaining deviations due to generative tasks. These findings underscore the importance of BWT-aware scheduling and hybrid supervision for scaling RL-based post-training toward general-purpose LLMs.