RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment (2304.06767v4)

Abstract: Generative foundation models are susceptible to implicit biases that can arise from extensive unsupervised training data. Such biases can produce suboptimal samples, skewed outcomes, and unfairness, with potentially serious consequences. Consequently, aligning these models with human ethics and preferences is an essential step toward ensuring their responsible and effective deployment in real-world applications. Prior research has primarily employed Reinforcement Learning from Human Feedback (RLHF) to address this problem, where generative models are fine-tuned with RL algorithms guided by a human-feedback-informed reward model. However, the inefficiencies and instabilities associated with RL algorithms frequently present substantial obstacles to the successful alignment, necessitating the development of a more robust and streamlined approach. To this end, we introduce a new framework, Reward rAnked FineTuning (RAFT), designed to align generative models effectively. Utilizing a reward model and a sufficient number of samples, our approach selects the high-quality samples, discarding those that exhibit undesired behavior, and subsequently enhancing the model by fine-tuning on these filtered samples. Our studies show that RAFT can effectively improve the model performance in both reward learning and other automated metrics in both LLMs and diffusion models.

• The paper introduces RAFT, a novel framework that fine-tunes generative models via iterative reward ranking.
• The methodology alternates between data sampling, ranking with a reward model, and fine-tuning to boost metrics like mean reward and diversity.
• The framework offers scalable, robust alignment that reduces memory load and mitigates issues such as reward hacking and overfitting.

Reward rAnked FineTuning for Generative Foundation Model Alignment (RAFT)

Introduction

Generative foundation models have shown proficiency in performing diverse tasks in various domains such as natural language processing and computer vision. These models include LLMs and diffusion models which are capable of generating high-quality, meaningful outputs. However, generative models trained on large-scale unsupervised datasets inherit biases which may lead to undesirable, skewed, and unfair outcomes. This paper addresses such concerns by introducing Reward rAnked FineTuning (RAFT), a novel alignment framework designed to align generative models with human ethics and preferences. RAFT seeks to enhance model performance via fine-tuning on high-quality samples which are selected based on a reward model (2304.06767).

Methodology

RAFT Framework

The RAFT framework operates in iterations, alternately performing data sampling, data ranking, and model fine-tuning. The process begins by generating candidate outputs from the model. A reward model quantifies the quality of each sample, and only those with high rewards are selected for fine-tuning. This ranks responses by quality, ensuring that iterative updates move models towards desired behaviors.

The RAFT process can be formalized as follows:

1. Data Collection: Sample a batch of prompts and generate multiple responses from the model.
2. Data Ranking: Assess each response using the reward model, keeping only the responses with the highest rewards.
3. Model Fine-tuning: Fine-tune the generative model using the high-reward responses to steer future outputs towards higher quality and ethically aligned results.

This methodology ensures a continuous improvement in model alignment by exposing it to and training it on high-caliber responses.

Figure 1: The left figure illustrates a typical RAFT training curve with different hyperparameters, highlighting improvements in reward over iterations.

Experimental Evaluation

LLMs Experiments

The RAFT was evaluated using the LLaMA-7B model, SFT (Supervised Fine-Tuning) baseline, and PPO (Proximal Policy Optimization). The performance of these models was assessed through mean rewards on a test set, perplexity, and diversity metrics like msttr, Distinct-1, and Distinct-2 scores. The LLaMA-7B-SFT model aligned using RAFT exhibited significant improvements in mean reward, outperforming the PPO-aligned model with better diversity metrics.

Diffusion Model Experiments

Beyond LLMs, RAFT was also applied to diffusion models enhancing resolution adaptation capabilities of Stable Diffusion (SD-1.5) models operating at reduced resolutions of 256x256 pixels. The metrics for assessing improvements included CLIP scores and aesthetic scores. RAFT demonstrated notable gains in both metrics for in-domain and out-of-domain samples.

Figure 2: The test reward variations during iterations for different settings of hyperparameter $K$.

Implications and Future Prospects

The RAFT framework offers enhanced stability and robustness over traditional reinforcement learning-based alignment approaches like PPO. By decoupling data generation and model fine-tuning, RAFT reduces memory burden and improves flexibility in model training. Its ranking-based data selection is more resistant to reward scale variations and noise, potentially mitigating issues like reward hacking and model overfitting.

Future work could explore integrating more diverse data sources and expert generators with RAFT. Investigating prompt engineering and post-processing could further optimize response quality for generative models. As a modular and adaptable framework, RAFT has the potential to be applied across different model architectures and domains.

Conclusion

RAFT provides a powerful, effective, and scalable solution for aligning generative foundation models with human preferences. Its iterative refinement of models via high-quality sample selection markedly enhances both ethical alignment and performance outcomes without the complexities and inefficiencies associated with reinforcement learning methods. This framework promises significant contributions to safe and responsible AI deployment.