RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment

Hanze Dong; Wei Xiong; Deepanshu Goyal; Yihan Zhang; Winnie Chow; Rui Pan; Shizhe Diao; Jipeng Zhang; Kashun Shum; Tong Zhang

Source-linked AI summary

RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment

Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, Tong Zhang

arXiv:2304.06767v4Apr 13, 2023cs.LGcs.AIcs.CLcs.CVstat.ML

TL;DR

Generative models can inherit bias from large unsupervised datasets, while RLHF alignment often relies on reinforcement learning methods with stability and efficiency challenges. RAFT instead ranks generated samples by reward and fine-tunes on the selected subset. The paper reports improved reward performance and automated metrics across language and diffusion models, while noting that sampling and reward modeling remain open areas.

Problem
Large unsupervised datasets can introduce biases into generative-model outputs, while RLHF methods face instability and inefficiency challenges during alignment.
Method
RAFT samples multiple outputs, ranks them with a reward function, filters high-reward samples, and fine-tunes the generative model on the filtered subset.
Results
RAFT improves reward performance and other automated metrics across large language models and diffusion models.
Takeaways & Limitations
RAFT offers a simple, SFT-like alignment framework with robust convergence, reduced memory requirements, flexible data sources, and an interpretable filtered-data objective.
Takeaways & Limitations
Refining RAFT’s sampling and reward modeling remains an open area of exploration.

Abstract

from arXiv · show

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 large language models and diffusion models.

1 Introduction

Generative models can inherit biases from large unsupervised datasets, motivating alignment with human preferences. RAFT addresses limitations of RLHF by ranking generated samples with a reward model and fine-tuning on the selected high-quality subset.

Large unsupervised datasets can embed biases that produce inaccurate or unfair model outputs.
RLHF aligns models with human preferences through reward-guided reinforcement learning, but PPO is less stable and efficient than supervised learning.PPO training can also impose substantial memory demands when multiple models are loaded.
RAFT alternates sampling model outputs, scoring and filtering them with a reward function, and fine-tuning on the high-reward subset.
RAFT uses SFT-like training with limited hyper-parameters, offering enhanced stability and robustness compared with conventional PPO-based reinforcement learning.
RAFT decouples data generation from model fine-tuning, reducing memory burden and allowing flexible data resources and processing.
Reward ranking rather than absolute reward values makes RAFT less sensitive to reward scaling and supports monitoring selected samples to mitigate reward hacking.The reward model’s imperfection can otherwise be exploited by algorithms seeking high reward.

2 Related Work

Prior alignment work uses human feedback and reward optimization for language and visual generative models. Related approaches also highlight the challenge of balancing preference alignment with image fidelity and include contemporaneous sample-filtering methods.

Alignment methods learn reward functions from human preferences and optimize language models with reinforcement learning such as PPO.
Visual alignment studies target control signals such as colors, counts, and backgrounds, but balancing human preferences with high-fidelity images remains challenging.
RRHF is a contemporaneous method that filters high-reward samples for generative-model alignment, while using diverse data sources before fine-tuning.

3 Algorithm

RAFT combines reward-based sample ranking with supervised fine-tuning: it generates multiple responses, retains the highest-reward response for each prompt, and iteratively trains on those samples. Its ranking objective, optional KL regularization, and decoupled stages support flexible alignment and reduced sensitivity to reward scale.

RAFT: Reward rAnked FineTuning: RAFT generates K responses per prompt, selects the response with the highest reward, and fine-tunes the current model on the resulting subset.The process repeats across stages until reward converges.
RAFT: Reward rAnked FineTuning: The algorithm alternates data collection, reward-based data ranking, and model fine-tuning in an iterative process.
RAFT’s data filtering uses reward ranking instead of absolute reward values, making it less sensitive to reward scale.The authors hypothesize greater robustness to reward noise and provide evidence in an appendix.
Extension: RAFT can incorporate a KL-divergence regularizer that ranks samples using reward adjusted by divergence from the initial model.The coefficient β balances reward learning with maintaining a low KL divergence.
Data collection and model improvement are decoupled, allowing the stages to be implemented separately and only one model to be loaded at a time.

4 LLM Experiments

The LLM experiments evaluate RAFT on HH-RLHF against PPO, finding that reward-ranked fine-tuning improves reward while preserving stronger perplexity and diversity behavior. Additional analyses examine training dynamics, computation, K, temperature, and KL regularization.

Setup: RAFT and PPO are evaluated from LLaMA-7B-SFT on HH-RLHF using reward, perplexity, and diversity metrics.The hand-out test set contains 4608 samples; HH-RLHF provides 112K training and 12.5K test samples.
Main results: RAFT achieves the highest mean reward of 2.294 while maintaining perplexity of 4.031.The RAFT-aligned model also has better perplexity than the PPO-aligned model.
Main results: GPT-4 and human evaluations are consistent with the automatic metrics on 100 randomly sampled test prompts.Human evaluation uses seven experts, while GPT-4 evaluation controls for response order.
Training dynamics: RAFT maintains relatively stable perplexity during training, whereas PPO perplexity usually worsens rapidly as reward increases.For RAFT-K32-λ1.0, the reward–perplexity trade-off is better after reward exceeds 1.85.
Computation: RAFT takes 5, 6.05, and 7.05 hours for K values 8, 16, and 32, respectively, while the fastest PPO configuration takes approximately 8.7 hours.Larger K increases inference cost, but K=16 and K=32 generally converge in fewer iterations than K=8.
Hyper-parameters: Increasing K tends to improve test reward, while RAFT-K32 is never worse in diversity than K=8 or K=16.The authors recommend balancing computational cost and performance when selecting K.
Hyper-parameters: RAFT consistently improves reward across temperatures 0.7, 0.85, and 1.0, although final reward slightly decreases as temperature increases.Temperature controls the diversity of sampled responses.
Distillation: In the distillation experiment, RAFT-LLaMA-7B-K32 consistently outperforms training only on teacher outputs in reward learning and diversity metrics.The aligned model’s perplexity also improves relative to the GPT-Neo-2.7B starting checkpoint.

5 Diffusion Model Experiments

RAFT improves diffusion-model alignment across resolution adaptation and text-image alignment tasks. It restores 256 × 256 image quality with few samples and produces better-aligned outputs at 512 × 512, while incurring much lower training time than DDPO.

Diffusion-model alignment: RAFT uses CLIP-based reward scoring to filter generated images before fine-tuning the diffusion model.The framework ranks samples by reward and trains on the high-reward subset.
Resolution adaptation: 8.4 mins (RAFT) vs 415 mins (DDPO) training time on a single A40 for resolution adaptation.The reported comparison indicates substantially lower computational cost for RAFT in this experiment.
Resolution adaptation: RAFT restores SD-1.5’s ability to generate proper 256 × 256 samples despite its difficulty at that resolution.The experiment uses a small number of generated samples and a CLIP-based aesthetic predictor as the reward function.
Resolution adaptation: RAFT significantly improves 256 × 256 image quality and improves scores for both in-domain and out-of-domain prompts.The out-of-domain evaluation includes prompts such as CIFAR-100 labels.
Open directions: Refining RAFT’s sampling and reward modeling, including intermediate-state rewards and best-of-K selection, remains open.These are identified as potential directions rather than established components of the reported experiments.
Text-image alignment: At 512 × 512, RAFT facilitates well-aligned outputs by fine-tuning SD-1.5 against textual data.The section attributes observed generation bias partly to prompt presentation and inductive bias in the training data.

6 Discussion and Conclusion

The discussion presents RAFT as a simple, interpretable alternative for aligning generative models with human preferences. It emphasizes robust, efficient training and flexible data handling, while linking filtered-sample monitoring to mitigation of reward hacking.

Framework: RAFT aligns generative models to human preference using a reward function and a simple fine-tuning framework.The discussion characterizes the framework as simple but effective.
Training properties: RAFT typically converges more robustly and faster than PPO because its training feature resembles supervised fine-tuning.The comparison is made with the popular PPO algorithm.
Training properties: Decoupling data generation from fine-tuning reduces GPU-memory requirements and allows flexible data sources and collection strategies.The discussion identifies both implementation consequences explicitly.
Interpretability: RAFT’s interpretability comes from viewing it as iteratively learning from induced best-of-K policies.Performance depends heavily on the quality of the dataset derived from the best-of-K policy and on hyperparameter choices.
Interpretability: Monitoring the filtered dataset can mitigate reward-model imperfections and reduce reward hacking concerns.The authors connect this benefit to RAFT’s clear learning objective.

A.1 Reward Modeling Details

The reward models are trained from human-preference pairs using a Bradley–Terry likelihood, with separate Open-LLaMA configurations and evaluation procedures. The section also describes extensions for improving RAFT’s data quality and handling ranking, reward scaling, and reward hacking.

Dataset and evaluation: The HH-RLHF setup uses 112K pairwise samples for reward-model training and 6,275 test samples for evaluation, reserving the remainder as a handout set.The handout set supports a more reliable evaluation procedure for RAFT- and PPO-aligned models.
Reward-model training: Reward modeling uses the Bradley–Terry model to estimate preference probabilities from human-preferred and human-dispreferred response pairs.The predicted reward rθ(x, y) is learned by minimizing the pairwise logistic loss over Dtrain.
RAFT extensions: RAFT can use a larger reward model because data generation, ranking, and supervised fine-tuning are decoupled.This separation also permits expert generators, prompt engineering, and advanced search methods as alternative data sources or generation strategies.
Ranking and reward scaling: RAFT is based on ranking, making it less sensitive to fixed reward recentering and allowing the same reward-scale adjustment that stabilizes PPO without changing RAFT.The paper also describes local ranking as its main version because reward comparisons across prompts are generally meaningless, while global ranking can be more sample-efficient when cross-prompt comparisons are valid.
Model selection: The best Open-LLaMA-13B reward model reaches 81.73% accuracy on 6K validation samples, compared with 75.79% for the best Open-LLaMA-3B model.Despite the stronger 13B reward model, PPO’s simultaneous model and reward-model loading caused an out-of-memory error on 8×A40 GPUs, so the experiment used 3B for PPO.
Reward hacking: Postprocessing can remove samples exploited by an imperfect reward model, using reduced diversity in the filtered dataset to detect reward hacking.The reported example involved responses containing emoji and #, whose output probability rapidly collapsed.

A.3 Reward Imperfection and Reward Over-optimization

This section examines how imperfect or noisy reward models can affect alignment and how RAFT compares with PPO under such conditions. It highlights RAFT’s stability under several noise patterns while emphasizing that reward-model quality and calibration remain important.

Reward over-optimization: Reward-model imperfections can cause reward hacking, motivating experiments that compare RAFT under alternative reward models and assess its sensitivity to over-optimization.The study treats this as an initial investigation and identifies more comprehensive analysis as future work.
Noisy rewards: The experiments add independent, prompt-level, or response-level Gaussian noise to rewards to represent different sources of reward-model or classifier error.Prompt-level noise is shared across responses for a prompt, whereas response-level noise is sampled independently for each response.
Noise robustness: RAFT-K32-λ1.0 is more stable than its no-noise counterpart when rewards are noisy, whereas noise produces noisy PPO training in both critic and actor.The paper attributes RAFT’s robustness to selecting higher-reward samples and invariance to linear reward transformations.
Scope boundary: Reward-model accuracy and calibration are central to both RLHF and RAFT, but a thorough study of calibration effects and improved reward-model training is outside this paper’s scope.The authors identify this analysis as future work.

A.4 GPT-4 and Human Evaluation

The evaluation compares model outputs using human experts and GPT-4, while training-curve figures examine RAFT and PPO under different reward-model and noise conditions. The supplied passages specify the evaluators and comparison setup but do not report a figure-level outcome beyond the reward-model caption elsewhere.

Evaluation protocol: Seven human experts evaluate output pairs without seeing labels, with response order shuffled, while GPT evaluation uses the GPT-4-0613 API.Both evaluators judge which assistant better follows instructions and answers the user’s question.
Noise experiments: Figure 9 compares RAFT-K32-λ1.0 and PPO training curves with or without random noise, reporting rewards without baseline recentering and after noise removal.Because PPO samples one response per prompt, its second and third noise cases coincide.
Evaluation protocol: The GPT evaluation prompt asks judges to compare two assistant responses for helpfulness, relevance, accuracy, depth, creativity, and detail.The template presents the user question and the two answers in labeled sections.
Reward-model experiments: Figure 10 compares RAFT-K16-λ1.0 and PPO training curves under different reward models, using an Open-LLaMA-3B reward model as the gold model.

B.1 LLM Samples

The LLM examples compare responses from LLaMA, LLaMA-SFT, LLaMA-PPO, and LLaMA-RAFT on randomly sampled HH-RLHF prompts. The examples include ordinary advice and a request involving theft, where the responses differ in helpfulness and safety framing.

Representative examples: The displayed examples are qualitative response comparisons rather than a quantitative metric report.
Representative examples: For a request about taking goods without paying, LLaMA-RAFT explicitly discourages stealing and recommends finding lawful ways to obtain needed items.The comparison includes responses from the base, supervised-fine-tuned, and PPO models.

B.2 Diffusion Model Samples

The diffusion-model examples show RAFT improving generated-sample quality, including on prompts outside those used for training and in text-image alignment examples.

RAFT significantly improves the quality of diffusion-model samples during resolution adaptation without cherry-picking.Figure 11 presents random 256 × 256 SD-1.5 generated results, with black samples indicating failures.
Improvements observed at 256×256 extend from training prompts to CIFAR-100 labels outside the CIFAR-10 training samples.The authors identify this cross-prompt improvement as evidence of RAFT’s generalization in sample-quality enhancement.
Additional text-image alignment examples further demonstrate RAFT’s role in aligning diffusion models.Figure 13 contains the cited text-image alignment examples.
The diffusion demonstrations are included in the LMFlow package alongside an executable RAFT alignment example.The default example aligns GPT-2 on IMDB, while the package also includes diffusion demos.

D Parameter Settings

The appendix reports separate hyper-parameter settings for fine-tuning LLaMA-7B on HH-RLHF and SD-1.5.

Table 14 lists hyper-parameters for fine-tuning LLaMA-7B on HH-RLHF.
Table 15 lists hyper-parameters for fine-tuning SD-1.5.