Self-Improving Robust Preference Optimization (2406.01660v4)

Abstract: Online and offline RLHF methods, such as PPO and DPO, have been highly successful in aligning AI with human preferences. Despite their success, however, these methods suffer from fundamental limitations: (a) Models trained with RLHF can learn from mistakes or negative examples through RL mechanism or contrastive loss during training. However, at inference time, they lack an innate self-improvement mechanism for error corrections. (b) The optimal solution of existing methods is highly task-dependent, making it difficult for them to generalize to new tasks. To address these challenges, we propose Self-Improving Robust Preference Optimization (SRPO), a practical and mathematically principled offline RLHF framework. The key idea behind SRPO is to cast the problem of learning from human preferences as a self-improvement process, mathematically formulated as a min-max objective that jointly optimizes a self-improvement policy and a generative policy in an adversarial fashion. Crucially, the solution for this optimization problem is independent of the training task, which makes it robust to its changes. We then show that this objective can be reformulated as a non-adversarial offline loss, which can be efficiently optimized using standard supervised learning techniques at scale. To demonstrate SRPO's effectiveness, we evaluate it using AI Win-Rate (WR) against human (GOLD) completions. When tested on the XSum dataset, SRPO outperforms DPO by a margin of 15% after 5 self revisions, achieving an impressive 90% WR. Moreover, on the challenging Arena-Hard prompts, SRPO outperforms both DPO and IPO (by 4% without revision and 6% after a single revision), reaching a 56% WR against against Llama-3.1-8B-Instruct.

• The paper introduces a novel SRPO framework that employs a two-step self-improvement process to enhance RLHF robustness across diverse task distributions.
• It leverages a min-max optimization approach transformed into joint supervised loss, eliminating the need for reward models and facilitating scalable training.
• Numerical evaluations reveal a 15% improvement in AI win-rate on out-of-distribution tasks, underscoring SRPO’s superior performance over existing methods.

Self-Improving Robust Preference Optimization: An Overview

Introduction

The paper "Self-Improving Robust Preference Optimization" introduces a new offline framework for Reinforcement Learning from Human Feedback (RLHF), focusing on mitigating the sensitivity of current RLHF methods to out-of-distribution (OOD) tasks. Unlike existing methods such as Proximal Policy Optimization (PPO) and Direct Preference Optimization (DPO), which are task-dependent, the proposed Self-Improving Robust Preference Optimization (SRPO) framework aims to be robust against variations in the task distribution. The proposed method aligns AI preferences with human preferences using a theoretically grounded and practical min-max optimization approach. This essay provides a comprehensive overview of the SRPO method, strong numerical results, and its implications for future AI development.

Key Ideas and Contributions

The Problem with Existing Methods

Current RLHF methods face significant limitations due to their dependence on the training task's distribution. When the evaluation distribution deviates significantly from the training distribution, the performance of these methods degrades. This dependency makes it challenging to generalize and apply these models to OOD tasks.

SRPO Framework

The SRPO framework addresses these challenges by introducing a two-step self-improvement process:

1. In-Context Self-Improving Preference Optimization: This step involves learning an in-context self-improvement model $\pi_{\dagger}$, which generates improved outputs iteratively based on the completions and contexts from the initial model.
2. Robust Preference Optimization of the Generative Model: Utilizing the self-improvement policy learned in the first step, the framework learns a robust generative LLM (Language Learning Model) $\pi$. This model is optimized so that its outputs require minimal improvement, ensuring robustness across different distributions.

By recasting the optimization problem into a joint supervised optimization process, SRPO circumvents the need for a reward model and online inference.

The Mathematical Foundation of SRPO

The SRPO framework is formally expressed as a min-max optimization problem:

$$J^*(x) = \min_{\pi} \max_{\pi_{\dagger}} \mathbb{E}\left[p(y_2 \succ y_1 \mid x) - \beta \text{KL}(\pi_{\dagger} \mid\mid \text{ref} \mid y_1, x) + \beta \text{KL}(\pi \mid\mid \text{ref} \mid x)\right]$$

The inner maximization can be solved in closed form for $\pi_{\dagger}$:

$$\pi^*_{\dagger}(y_2 \mid y_1, x) = \frac{\exp\left(\frac{p(y_2 \succ y_1 \mid x)}{\beta}\right) \pi_{\text{ref}}(y_2 \mid y_1, x)}{Z^*(y_1, x)}$$

The solution to this problem is then translated into a non-adversarial offline supervised loss, allowing the joint optimization of both $\pi_{\dagger}$ and $\pi$.

Numerical Results

The paper presents a robust evaluation of the SRPO framework against well-established baselines, DPO and IPO. Notably, SRPO demonstrates substantial improvements in AI Win-Rate (WR) against human completions. For instance, when evaluated on the OOD XSUM dataset, SRPO achieves a significant improvement, outperforming DPO by 15% after 5 self-revisions, achieving a WR of 90%.

Implications and Future Research

Practical Implications:

1. Robustness to Task Distribution: SRPO's independence from the behavior policy ($\mu$) ensures that it remains robust to distribution shifts, making it suitable for deployment across various tasks without task-specific retraining.
2. Scalability: The transformation of the min-max optimization problem into a joint supervised loss facilitates scalable, large-scale implementation.

Theoretical Implications:

1. Generalization of Preference Models: Unlike previous models restricted to the Bradley-Terry framework, SRPO's formulation holds across all preference models, enhancing its applicability to diverse scenarios.
2. Self-Improvement Mechanism: The self-improvement policy embedded within SRPO introduces a novel paradigm in LLM training, focusing on iterative refinement of completions.

Future Developments:

1. Application to Complex Multi-Task Benchmarks: Testing SRPO on more complex multi-task benchmarks could validate its robustness and scalability further.
2. Improving Algorithms for General AI: By leveraging the robustness and self-improvement capabilities of SRPO, future research could focus on enhancing general AI's ability to perform consistently across varied tasks and distributions.

Conclusion

The "Self-Improving Robust Preference Optimization" framework represents a significant step towards making RLHF methods more robust and scalable. Through its innovative approach to self-improvement and robust preference optimization, SRPO addresses the limitations of task dependency in current methods. The strong numerical results and theoretical foundations pave the way for more resilient AI systems capable of maintaining performance across diverse and unforeseen tasks. Future research will likely build on these foundations, exploring broader applications and further refining robust AI training methodologies.