Adaptive Segment-level Reward: Bridging the Gap Between Action and Reward Space in Alignment

TL;DR

This work tackles the credit-assignment gap in RLHF by introducing Adaptive Segment-level Reward, which segments sequences at semantically meaningful pivots where token rewards change. By masking and backpropagating only through segments whose rewards align with the overall sequence signal, the method can unify and improve upon PPO, DPO, and rejection sampling within a single framework. Theoretical analysis formalizes the reward-action gap and demonstrates how segment-level supervision reduces error compared to token- or step-level approaches. Empirically, ASR yields about a 10% gain on adversarial evaluations and roughly a 1.3% improvement on standard benchmarks, suggesting practical benefits for robust alignment across diverse tasks.

Abstract

Reinforcement Learning (RL) has proven highly effective in aligning Large Language Models (LLMs) with human preferences. Typical RL methods optimize under an overall sequence reward, which can lead to a suboptimal learning process. This reflects a key credit assignment problem: identifying which tokens to reinforce or suppress. To rectify these shortcomings, step-wise and token-wise methods have been proposed. However, step-wise methods rely on punctuation segmentation and still cannot accurately identify the key tokens. The token-level approach is too fine-grained, attending to many unimportant tokens and thus introducing a large amount of noise. To assign more accurate rewards to different tokens, improving credit assignment, we propose the "Adaptive Segment-wise Reward" method. We employ semantic meaning, rather than punctuation, to adaptively delineate segments. Experiments demonstrate that our method can be integrated into various training methods. Compared to training methods \textit{without} our approach, our method improves the success rate on adversarial samples by 10\%, and achieves a 1.3\% improvement on evaluation benchmarks such as MMLU, GSM8K, HumanEval, etc.

Figures (9)

• Figure 1: Adaptive Segment-level Reward: we divide sequence based on token rewards. We can train the model by masking certain sub-sequences or applying different rewards to various sub-sequences. This approach is closer to actual density than step-wise and token-wise methods.
• Figure 2: For general tasks, especially in low-information statement, long context(e.g., writing articles or RAG applications), step-wise supervision is significantly less accurate than sequence-wise supervision. Left and middle shows the performance of Outcome-supervised is better than Process-supervised for general tasks, even when well-trained PRMs(Appendix:\ref{['sec:more details to PRMs']} were used. ORM: Outcome supervised Reward Model. PRM:Process supervised Reward Model. $PRM_O$: Outcome supervised by PRMs. $PRM_P$: Process supervised by PRMs.
• Figure 3: Comparison of different ways of reward signals. Left: Green represents low-reward tokens, while orange represents high-reward tokens. Mid: A comparison of token-wise, step-wise, and sequence-wise reward signals. Token-wise rewards exhibit significant fluctuations and high noise levels, leading to unstable training. The lines highlighted by the yellow dashed box shows that in the step-wise approach, tokens at the same step can have completely different rewards, yet they are all represented by the same reward score, which can lead to errors. Right: This image presents the message-wise reward method proposed in this paper. In autoregressive generative transformer, each token represents an action, The size of the vocabulary is the size of the action space. This method significantly separates tokens with different reward scores into distinct subsequences, thereby preventing the same subsequence from containing both positive and negative actions.
• Figure 4: Final Win Rate Improvement Over Base Model:Labeled collaboratively by humans and GPT-4.
• Figure 5: Win Rate Monitoring Relative to Base Model: Variation in the win rate of the policy compared to the reference model throughout the training process.