SeRA: Self-Reviewing and Alignment of Large Language Models using Implicit Reward Margins

TL;DR

Direct Alignment Algorithms (DAAs) for RLHF offer efficiency but are vulnerable to off-policy data issues, including learning from spurious cues and poor alignment due to distribution drift when policy updates outpace offline preferences. SeRA introduces Implicit Reward Margin (IRM) to select informative off-policy samples and to bootstrap preference data using the updated policy, enabling cost-efficient, self-guided alignment across multiple DAAs and model families. The method demonstrates broad effectiveness on instruction-following tasks and diverse datasets, improving win rates and robustness to noisy annotations without external reward models. This approach provides a practical pathway to better alignment by leveraging self-generated signals and offline data, reducing reliance on expensive AI-Feedback or reward-model supervision.

Abstract

Direct alignment algorithms (DAAs), such as direct preference optimization (DPO), have become popular alternatives for Reinforcement Learning from Human Feedback (RLHF) due to their simplicity, efficiency, and stability. However, the preferences used in DAAs are usually collected before the alignment training begins and remain unchanged (off-policy). This can lead to two problems where the policy model (1) picks up on spurious correlations in the dataset (as opposed to learning the intended alignment expressed in the human preference labels), and (2) overfits to feedback on off-policy trajectories that have less likelihood of being generated by an updated policy model. To address these issues, we introduce Self-Reviewing and Alignment (SeRA), a cost-efficient and effective method that can be readily combined with existing DAAs. SeRA comprises of two components: (1) sample selection using implicit reward margins, which helps alleviate over-fitting to some undesired features, and (2) preference bootstrapping using implicit rewards to augment preference data with updated policy models in a cost-efficient manner. Extensive experimentation, including some on instruction-following tasks, demonstrate the effectiveness and generality of SeRA in training LLMs on offline preference datasets with DAAs.

Figures (11)

• Figure 1: SeRA uses an iterative policy-model ensemble to obtain reward margins between preferred and dis-preferred trajectories in off-policy and on-policy preference data. It then selects preference data with large reward margins and rejects the data with smaller rewards margins to update the policy models via Direct Alignment Algorithms (DAAs).
• Figure 2: Motivation of SeRA : The reward-margin based selection showed higher win-rates (measured by LLM-as-a-Judge zheng2024judging with Claude 3 claude3) against the base SFT model across a variety of settings-- (a) Model types: TinyLlama-1.1B zhang2024tinyllama, Pythia-2.8B biderman2023pythia, Phi-2 textbooks2, and Mistral-7B jiang2023mistral. (b) Model size: various model sizes of Pythia (410M, 1.0B, 2.8B, 6.9B). (c) DAAs: DPO rafailov2024direct, IPO azar2024general, SLiC zhao2023slic, and KTO-Pair ethayarajh2024kto and (d) training datasets.
• Figure 3: Comparison of correlations between the rewards of the trained policy model and features (e.g., GPT-4 score and response length) across different selection methods (i.e., no selection, GPT-4 score-based, log probability of the reference model, and IRM-based selection). [Row 1] Correlation between GPT-4 Score & implicit reward (i.e,$r(\mathbf{x}, \mathbf{y}_{w})$) for $\mathbf{y}_w$. [Row 2] Correlation between response length (i.e.$|\mathbf{y}_w|$) and the implicit reward for chosen responses. The model with IRM selection (i.e.[Column 4]) shows the highest $R^2$ score for the first row, but the lowest $R^2$ score for the second row. This indicates that IRM-based selection strategy can effectively mitigate the over-optimization on response length park2024disentangling.
• Figure 4: Comparison of performance on 20% (left) and 40% (right) of noisy preference
• Figure 5: Analyzing the effect of varying percentage (determined by $k$) of sampled off-policy preference data, and generated data (determined by $\tilde{k}$).

Theorems & Definitions (11)

• Remark 1: Informal Statement for \ref{['thm:main']}
• Lemma A.1: pmlr-v139-menon21a
• proof
• Definition 1: Covering Number
• Definition 2: Uniform Covering Number
• Lemma A.2: Theorem 6, maurer2009empirical
• proof
• Lemma A.3: Modification from pmlr-v139-menon21a
• proof
• Theorem 1: Formal Statement