Reusing Embeddings: Reproducible Reward Model Research in Large Language Model Alignment without GPUs

TL;DR

This paper argues that reward models in RLHF are hindered by high computation and reproducibility barriers. It proposes embedding-based reward models, where inputs are embeddings rather than full LLM prompts, to reduce training/evaluation costs and improve reproducibility. Through a case study reproducing reward-model ensemble findings on CPU-only hardware, it demonstrates stable performance, scalable evaluation, and efficient inference-time optimization, while highlighting scenarios where embedding-based methods excel or underperform relative to LLM-based approaches. The work aims to democratize reward-model research, encourage public embedding assets, and spur developments in general-purpose reward embeddings for safer, more effective LLM deployments.

Abstract

Large Language Models (LLMs) have made substantial strides in structured tasks through Reinforcement Learning (RL), demonstrating proficiency in mathematical reasoning and code generation. However, applying RL in broader domains like chatbots and content generation -- through the process known as Reinforcement Learning from Human Feedback (RLHF) -- presents unique challenges. Reward models in RLHF are critical, acting as proxies that evaluate the alignment of LLM outputs with human intent. Despite advancements, the development of reward models is hindered by challenges such as computational heavy training, costly evaluation, and therefore poor reproducibility. We advocate for using embedding-based input in reward model research as an accelerated solution to those challenges. By leveraging embeddings for reward modeling, we can enhance reproducibility, reduce computational demands on hardware, improve training stability, and significantly reduce training and evaluation costs, hence facilitating fair and efficient comparisons in this active research area. We then show a case study of reproducing existing reward model ensemble research using embedding-based reward models. We discussed future avenues for research, aiming to contribute to safer and more effective LLM deployments.

Figures (9)

• Figure 1: In reward model research, using embeddings as input (i.e., focusing on the pink box) brings the following benefits: 1. there are much less parameters in those reward models; 2. it has a much lower training cost than using LLM-based reward models; 3. it has a much lower evaluation cost as compared to the LLM-based reward models; 4. it minimizes the inference-time cost by generating embeddings as by-products in language generation; 5. research using embedding-based reward models are highly reproducible due to the low computational demand, high training stability, and minimal hardware requirement.
• Figure 2: Comparing performances of Embedding-based RM with LLM-based RMs. The Embedding-based RMs demonstrate high learning stability and strong performance as compared to LLM-based RMs, but are much cheaper to train and evaluate, and more scalable in inference time. Results are from the Gemma 2B model. Additional results using the Gemma 7B and LLaMA3 8B models are presented in Appendix \ref{['appdx:more_results']}
• Figure 3: The inputs of embedding-based reward models are by-products of language model generation. Unlike conventional LLM-based reward models that require another LLM forward pass for inference time evaluations, embedding-based models alleviate the memory challenge and facilitate inference time optimization for LLM-free service providers. These providers, who rely on third-party LLM services via APIs rather than hosting large models locally, can efficiently perform inference time optimization using only embeddings.
• Figure 4: Using embeddings as inputs in a lightweight reward model ensemble practice to mitigate reward overoptimization. Reproduction of prior findings across over $12000$ configurations can be completed in less than 1 day using CPU-only resources.
• Figure 5: Reproduction of reward model ensemble papers using embedding-based reward models. Additional results using the Gemma 7B and LLaMA3 8B models are presented in Appendix \ref{['appdx:more_results']}