Reward-Wise Value Estimation for Multi-Reward Optimization in Large Language Models

Value functions play a central role in critic-based reinforcement learning (RL) algorithms such as proximal policy optimization (PPO) and in inference-time scaling for large language models (LLMs). In complex LLM tasks, multiple reward signals are often leveraged to guide RL training and inference. However, existing value function estimation methods either aggregate all rewards into a single scalar, limiting flexibility across reward components, or rely on separate critic models for each reward, which is computationally inefficient. To address these limitations, we propose an efficient and flexible value function estimation method for multi-reward optimization in LLMs. Specifically, we introduce a multi-head critic model architecture that jointly estimates reward-wise value functions within a single model. This design enables fine-grained credit assignment across reward components while maintaining computational efficiency. We further prove that the proposed multi-head critic loss yields lower irreducible error compared to the vanilla single-head critic loss. Experimental results on multi-turn agentic search and mathematical reasoning tasks demonstrate that our reward-wise value estimation approach consistently improves both policy optimization and critic learning in PPO training.