Paper: Reward Collapse in Aligning Large Language Models (ICLR-2024)

Reviewer comment (Theory):
The paper claims that we should use prompt-based utility function Eq.(3). However, all follow analysis is about the utility function without the prompt, which is inconsistent.

Author response:
Our theoretical analysis focuses on determining the reward distribution when a fixed utility function is chosen. Therefore, for the purposes of our analysis, we can fix a utility function.

Reviewer comment (Reproducibility):
The paper does not provide the implementation details of the prompt-aware training.

Author response:
In our experiment in Section 3.2, we randomly assigned each question as either open-ended or concrete. For open-ended questions, the utility function is selected as G(x) = -1/x. For concrete questions, their utility function is chosen as G(x) = x. This choice of utility functions serves the specific purpose of demonstrating that the prompt-aware approach can effectively prevent reward collapse. Additionally, in Appendix A.2, we delve into potential practical methods for assigning prompt types.

Reviewer comment (Theory):
How do the conclusions change if the assumption that the LLM is sufficiently overparameterized so that it can maximize the utility for all the prompts (discussions around equation (2))?

Author response:
The assumption that the LLM is sufficiently overparameterized is quite strong, but it’s essential for our conclusions. Since if the reward does not maximize the objective function exactly, reward collapse may not exists. For example, in the experiment in Section 3 (Figure 1 and 4), reward collapse does not show up in the early stage of training. This overparameterization assumption is satisfied in most real applications. For example, InstructGPT  se a 6B reward model [1].

Reviewer comment (Experiments):
While the story and theoretical analysis are sound, the evidences of this paper are limited. But one thing I believe can largely improve the paper is that we can further evaluate the quality of the reward model by best-of-n policy. Specifically, we can fix a LLM, and for each prompt, we sample n responses and then take the one with the highest reward as the final output. Then, we can compare the responses by either human evaluation or GPT4 evaluation. For more details, you may check [1].

Author response:
While we acknowledge the potential improvement of RLHF in current LLMs, the primary focus of this paper is to document and investigate the phenomenon of reward collapse during reward model training in LLMs. We present theoretical results on reward collapse and derive the reward distribution under a specific class of utility functions. Although intuitively our approach should enhance the performance of RLHF, we believe that the impact of using a prompt-aware utility function requires thorough further research. Numerous techniques implicitly avoid the harm of reward collapse, such as early-stop or adding regularizers, which may alleviate reward collapse. Investigating the extent to which the trained reward model enhances the capabilities of large language models, such as their ability to self-calibrate uncertainty, is an interesting direction for future work beyond the scope of this paper.

Reviewer comment (Experiments):
The experiments conducted are rather limited and simple. In particular, the utility function used for the response length seems to be hard to generalize to general practical applications. Could you give an example of the choice of utility function in practice, e.g., used for the HH-RLHF dataset (whose details can be found in huggingface).

Author response:
Aassigning prompt types is fundamental in the prompt-aware approach, as discussed in Appendix A.2. A straightforward method for determining the prompt type is to gather input from human labelers, who typically rank different responses, as seen in InstructGPT. Additionally, we can request them to assess how open-ended the prompt is using a scale ranging from -1 to 1. Automated annotation processes are also possible. For instance, one approach involves assessing the variability of responses to a given prompt. If the responses exhibit high variability, the prompt can be classified as open-ended. Conversely, if the responses show low variability, the prompt may be deemed close-ended. Determining the prompt type is indeed a complex and fascinating task, providing a promising avenue for future research.

Reviewer comment (Experiments):
Does not provide much detail on how the prompt-aware utility function U_prom adaptively selects between U(x) = x and U(x) = -1/x in the experiments mentioned in Section 3.2. Do you manually assign utility functions based on the question type?

Author response:
In Section 3.2, we specified that each question was randomly assigned as either open-ended or concrete. For open-ended questions, their utility function is chosen as G(x) = -1/x. For concrete questions, their utility function is chosen as G(x) = x. This is only for the purpose of our experiment. Additionally, in Appendix A.2, we delve into potential practical methods for assigning prompt types.

Reviewer comment (Experiments):
Experiments are done on only synthetic datasets where the word count is the ground-truth reward. Real-world ranking datasets would provide stronger validation.

Author response:
We understand your concern about the validation of our findings on real-world datasets. While Section 3.1 presented experiments on a real-world dataset to demonstrate the existence of reward collapse, we used synthetic datasets in Section 3.2 to illustrate how prompt-aware training can mitigate this issue. We acknowledge the potential for stronger validation with real-world ranking datasets and agree that it is an area for improvement. Nevertheless, our controlled experiments, combined with those in Section 3.1, effectively convey the key message: reward collapse can be avoided through prompt-aware training.

Reviewer comment (Experiments):
It is not clear how the prompt-dependent reward distribution will contribute to the performance increase for RLHF or other direct optimization methods or just the best of n sampling. It would be great to see how this reward distribution will increase performance.

Author response:
The primary focus of our paper is to document and investigate the phenomenon of reward collapse in the training of reward models for LLMs. We have developed theoretical results on reward collapse and derived the reward distribution under a specific class of utility functions. While we intuitively believe that our approach enhances the performance of RLHF, we recognize the need for further research on the effects of using prompt-aware utility functions. We also acknowledge the existence of other techniques that implicitly address reward collapse, such as early-stop or regularizers. Exploring how the trained reward model enhances the capabilities of large language models, including their ability to self-calibrate uncertainty, is an intriguing avenue for future research, albeit beyond the scope of this paper.

Reviewer comment (Presentation):
The evidence of the reward collapse in Figure 1: what does the dashed region represent? is it over the 128 prompt datasets? It would be clearer to present the reward collapse phenomenon on a close-end question.

Author response:
Yes, the dashed region signifies the randomness across 128 distinct prompts. We appreciate your suggestion, and in response, we will employ an alternative utility function to ensure a clearer presentation of the results.