GRAM: A Generative Foundation Reward Model for Reward Generalization

TL;DR

GRAM presents a generative foundation reward model that leverages large-scale unlabeled data for pre-training and minimal labeled data for fine-tuning to achieve broad reward generalization. By formulating a two-stage training process and introducing label smoothing that regularizes a Bradley-Terry-like objective, GRAM unifies generative and discriminative perspectives and delivers improved OOD performance across response ranking, RLHF, and task adaptation. Extensive experiments show GRAM outperforms strong baselines on RewardBench and maintains robust performance under best-of-n sampling and adaptation tasks, often with far less task-specific labeled data. This approach enables scalable, generalizable reward modeling for aligning LLMs in diverse settings and demonstrates the practical benefits of unlabeled data and domain-aligned pretraining.

Abstract

In aligning large language models (LLMs), reward models have played an important role, but are standardly trained as discriminative models and rely only on labeled human preference data. In this paper, we explore methods that train reward models using both unlabeled and labeled data. Building on the generative models in LLMs, we develop a generative reward model that is first trained via large-scale unsupervised learning and then fine-tuned via supervised learning. We also show that by using label smoothing, we are in fact optimizing a regularized pairwise ranking loss. This result, in turn, provides a new view of training reward models, which links generative models and discriminative models under the same class of training objectives. The outcome of these techniques is a foundation reward model, which can be applied to a wide range of tasks with little or no further fine-tuning effort. Extensive experiments show that this model generalizes well across several tasks, including response ranking, reinforcement learning from human feedback, and task adaptation with fine-tuning, achieving significant performance improvements over several strong baseline models.

Figures (11)

• Figure 1: Architectures of discriminative and generative reward models. In discriminative models, the reward model is a scoring function that is trained to minimize the pairwise ranking loss between two responses. In generative models, we use an LLM to predict the label token given a prompt, an input, and a pair of responses. This model can be trained in the same way as standard LLMs.
• Figure 2: Accuracies of discriminative and generative reward models on the ID and OOD test sets.
• Figure 3: Illustration of the two-stage training method. In the first stage, we pre-train the model via response generation, which is an unsupervised task. In the second stage, we fine-tune the model to generate preferences in a standard supervised manner.
• Figure 4: Performance of GRAM and its baselines on BoN sampling. We use proxy scores to assess preference learning and oracle scores to evaluate the generalization capability. "D-" and "G-" denote that the reward model is trained using discriminative and generative reward modeling frameworks.
• Figure 5: The performance of reward models fine-tuned with varying amounts of task-specific preference data (summarization and harmlessness). Please refer to Figure \ref{['fig:all_results_adaptation']} for the results on the four remaining baselines, including D-Feeze, D-Regularization, G-Freeze, and G-Label Smoothing.