Priority-aware method improves multi-goal language model training
Priority-Preserving Advantage Aggregation for Multi-Reward Policy Optimization
Machine LearningComputation and Language
Summary
Many language models learn by balancing several goals, like being correct and reasoning quickly. But combining feedback from these goals often mixes their priorities, leading to the less important goal sometimes overriding the main one. The authors present a new method called PAPO that merges feedback in a way that respects the main goal’s priority while still using helpful input from secondary goals. This approach improves primary goal performance without losing benefits from secondary ones.
What this means in practice
- •For machine learning engineers: Train language models that better respect key priorities like accuracy while still using helpful secondary signals during reinforcement learning.
- •For natural language processing teams: Optimize dialogue and reasoning models to maintain correctness as a priority without losing efficiency improvements from auxiliary rewards.
Authors
Xiangchi Yuan, Xingjian Diao, Zheyuan Liu, Chunhui Zhang, Dachuan Shi, Kejing Xia, Zhenbang Du, Shenglong Yao, Soroush Vosoughi, Wenke Lee
Abstract
LLM post-training often optimizes multiple objectives with unequal priorities, such as correctness and reasoning efficiency. However, the widely used additive reward aggregation in RL can violate these priorities: when primary and auxiliary advantages conflict, a sufficiently large auxiliary signal can overturn the preference induced by the primary objective. We term this problem \emph{advantage priority inversion}. Although reward-level gating can protect priority more directly, it still removes potentially useful auxiliary distinctions among rollout samples that fail the primary condition. To address this trade-off, we propose Primary Advantage Policy Optimization (\textbf{PAPO}), a priority-aware advantage aggregation method that asymmetrically combines a primary and an auxiliary advantage. PAPO retains auxiliary supervision when the objectives agree and smoothly attenuates conflicting auxiliary signals, preserving the primary sign during fusion for \emph{primary-consistent ordering}. This rule applies to sequence-level advantages in GRPO and token-level advantages in actor--critic PPO. Experiments on multiple domains and models show that PAPO consistently achieves the best primary objective while maintaining best or competitive auxiliary-objective performance across both GRPO- and PPO-based optimization. \href{https://xiangchi-yuan.github.io/papo-project/}{Code is available}.