Reward-Augmented Data Enhances Direct Preference Alignment of LLMs

We thank the reviewer for identifying our work's technical contributions. The valuable comments have helped us improve our manuscript (marked blue in the revision). Below are our specific responses to the questions raised by the reviewer:

Weakness 1: Conditioning on low reward data is only empirically useful. Despite improvements indicated by the experiments, I find the method to be highly empirical and lacks theoretical guarantees. During sampling, the LLM is always prompted to generate only high-reward responses, which makes learning low-reward prompts during training useful only because it proves "generalization".

• Direct alignment on our augmented reward-relabeled data ensures that the high-quality rejected responses are not suppressed, since the LLM is trained to generate responses corresponding to different reward values, enabling it to become aware of and adapt to these reward distinctions.
• We also added theoretical results at the end of Section 4 and in Appendix A. Specifically, we showed that our method attains global convergence to the optimal policy and the suboptimality decays at the order of $N^{-1/2}$ ($N$ denotes the size of the reward-augmented preference dataset), which provides a theoretical justification for the strong empirical performance of the reward-augmented DPO introduced in this paper. Unlike prior works on goal-conditioned RL with supervised learning [4, 5], which typically establish weaker results such as local performance improvements or the optimization of a lower bound on $J(\pi)$, our analysis guarantees global convergence to the optimal policy. This distinction underscores the significance of integrating DPO-like methods with goal-conditioned approaches.

Weakness 2: Lack of comparison with other reward-augmented methods. One main claimed contribution of the paper is being able to learn from quantitative rewards, but it did not compare with any reward-augmented methods. Related prior works are [1] and [2]. We compared with reward-conditioned baselines in experiments including DPA [3], SteerLM [4], and RPO [7], and we thank the reviewer for pointing out the great missing related work [1, 2]. We have made two modifications in the revision:

• We added the following paragraph to the related work section: "Notably, Noise Contrastive Alignment [1] and Unified Language Model Alignment [2] introduce unified frameworks for alignment with binarized or reward datasets by leveraging (information) noise contrastive estimation and a hybrid of SFT with point-wise DPO, respectively. In contrast, our work focuses on addressing the limitations of direct alignment algorithms with data relabeling (on implicit-reward augmented binarized or reward datasets), and do not make algorithm changes."
• Experimentally, we add comparsions with Noise Contrastive Alignment [1]. Specifically, we compare the Zephyr-SFT models fine-tuned with our method and DPO, DPA [3], SteerLM [4], NCA-preference, NCA-reward, InfoNCA-preference, InfoNCA-reward, all on UltraFeedback. The results on the AlpacaEval 2.0 benchmark are shown as follows:

It can be observed that NCA and InfoNCA are strong baselines and significantly outperform SteerLM, DPO, and DPA. Besides, our method improves DPO by a considerable margin. We have incorporated the above results in our revised manuscript (Figure 4 and Appendix B.2).

We thank the reviewer for identifying our work's novelty, soundness, and technical contributions. The valuable comments have helped us improve our manuscript (marked blue in the revision). Below are our specific responses to the questions raised by the reviewer:

Weakness 1: The method relies heavily on reward labels, whether in directly constructing the reward-augmented dataset or through a multi-attribute rewards dataset. The actual impact of these rewards on data quality remains unexplored. It is unclear if the model would benefit more from high-reward rejected answers than from low-reward chosen answers, which warrants further investigation. Following the reviewer's suggestions, we conduct additional experiments to show the benefits of learning from the rejected high-reward samples. Specifically, using the UltraFeedback dataset, we construct two reward-augmented preference datasets by filtering out augmented data based on rejected responses with low and high reward values, respectively, which we denote as Ours-filter-low and Ours-filter-high. Compared to our method, these datasets isolate the impact of excluding low- and high-reward rejected responses as goals. The results on AlpacaEval 2.0 are shown as follows:

It can be observed that learning from rejected high-reward samples (Ours and Ours-filter-low) demonstrates superior performance compared to the approach that excludes these samples (Ours-filter-high), demonstrating the benefits of keeping the rejected high-reward samples.

Weakness 2: Learning from preference usually involves guiding the model towards specific output styles, such as aligning to certain type of instructions or adapting to varied language tones. The authors used previously unseen instructions in DPO training that match existing data, potentially limiting the model's adaptability. The experimental results require additional examples and detailed explanations to clarify this impact. And in appendix $\mathcal{D}_N^i$ training prompts should be added.

• In RLHF or RLAIF, it is common practice to use instructions that are different from the ones in SFT or in previous alignment (e.g., DPO) iterations. Since our method does not alter the direct alignment algorithms, it is likely that it will share similarities with DPO, such as guiding the model towards specific output styles to align with human or AI preferences.
• Following the reviewer's suggestion, we report the performance of the DPO baseline and our method fine-tuned from the Zephyr-SFT model, where a large portion of instructions during alignment also appear in the SFT stage. The results on the AlpacaEval 2.0 benchmark are as follows. It can be observed that applying DPO on both the original preference dataset and our reward-augmented dataset improves the performance with longer output length, which might be a sign of learning generation styles preferred by judge models.

• In fact, it is widely observed that alignment can considerably improve scores judged by LLMs or humans but can decrease the reasoning abilities of the LLM, which is known as the alignment tax. In our experiments (Table 3), we observed that our method significantly and consistently improves the base models on instruction-following benchmarks that use LLM as a judge, and moderately improves most of the academic benchmarks. This indicates that our method does not suffer from severe alignment tax.

We thank the reviewer for identifying our work's novelty, soundness, and technical contributions. The valuable comments have helped us improve our manuscript (marked blue in the revision). Below are our specific responses to the questions raised by the reviewer:

Weakness 1.1: I am confused about Tables 1 and 2. How did you compute $π^*$? What are the assumptions used? If you assume a standard setting like the DPO paper, you would never get a policy that always generates one response over the other (because Bradley-Terry inherently assumes some noise). Also, just more realistically, LLMs are softmax parametrized, and the probability of a response can never be exactly 0.

• $\pi^*$ is defined as the policy that maximizes the reward learned from the preference dataset $\mathcal{D}$, i.e., $\pi^*(y\mid x) = argmax_\pi E_{x\sim\mathcal{D}, y\sim\pi(\cdot\mid x)}[r(x, y)]$. We may follow standard RLHF and assume that the true preference distribution can be modeled by Bradley-Terry, i.e., $r(x, y) = argmax_r E_{(x,y_w,y_l)\sim\mathcal{D}}[\log\sigma(r(x, y_w) - r(x, y_l))]$, where $\sigma(\cdot)$ is the logistic function.
• As an example, in Table 1, $\mathcal{D}$ contains a single preference pair $y_1 > y_2$. The fitted reward thus assigns a higher value of $y_1$ than $y_2$, i.e., $r(x, y_1) > r(x, y_2)$. To maximize $r$, the optimal policy is $\pi^*(y_1\mid x) = 1$.
• Since we studied tabular settings in Section 3, for clarity, we considered nonparametric policies without KL regularization. The reviewer is correct that in the KL-regularized formula in DPO or when the policy is softmax parameterized, $\pi^*(y_1\mid x) < 1$. Nevertheless, even if we follow the exact DPO formulation by defining the reference policy as a uniform distribution $\pi_{\text{ref}}(y_1\mid x)=\pi_{\text{ref}}(y_2\mid x)=0.5$, $\pi^*(y_1\mid x)$ is still arbitrarily larger than $\pi^*(y_2\mid x)$.

Weakness 1.2: Overall, I am confused about Section 3.1, where, for example, "reward sparsity" is used to describe something different than what it traditionally means (sparsity over trajectory time, not over the dataset) and the logic is unclear. This makes the motivation for the algorithm unclear -- Section 3.2 is titled way too strongly for a method that has no clear theoretical backing. This weakness is the reason for giving a 6 instead of an 8. It is hard to justify accepting a paper with such vague and possibly incorrect text, even if the idea and method are performant.

• The reward sparsity in Section 3.1 refers to the sparsity of the highest-reward data. The goal-conditioned reward (defined in Sec. 4) $R(x, y, g=r_{\text{max}})$ is $1$ only for responses that achieve $r_{\text{max}}$, and $0$ for all other responses, which corresponds to reward sparsity.
• To address your concerns regarding the too-strong title and unclear theoretical backing, we have made the following two modifications:
  1. We changed the title of Section 3.2 from "Reward-Conditioned Policies Resolve the Limitations" to "Reward-Conditioned Policies Learn From the Full Spectrum of Response Quality". This new title reflects the property of the reward-conditioned policy, which leads to a natural fix for the limitations of direct alignment algorithms in Section 3.1.
  2. We also added theoretical results at the end of Section 4 and in Appendix A. Specifically, we showed that our method attains global convergence to the optimal policy and the suboptimality decays at the order of $N^{-1/2}$ ($N$ denotes the size of the reward-augmented preference dataset), which provides a theoretical justification for the strong empirical performance of the reward-augmented DPO introduced in this paper. Unlike prior works on goal-conditioned RL with supervised learning [4, 5], which typically establish weaker results such as local performance improvements or the optimization of a lower bound on $J(\pi)$, our analysis guarantees global convergence to the optimal policy. This distinction underscores the significance of integrating DPO-like methods with goal-conditioned approaches.

Thank you to the authors for their response.

I would recommend modifying the writing around Table 1 to clearly indicate what the setting is and to also describe that you are not training anything but just conceptually formulating what the policy would be. Overall, the writing quality of the paper is not good, and I found a lot of things confusing (as I mentioned in my review). I recommend authors invest substantial time into the writing to clarify the settings in each section as well as what is a concrete result and what is merely speculation.

I appreciate the authors' effort to run the additional ablation that I requested. I think this more clearly illustrates that the goal-conditioning is actually playing a role in the alignment procedure. I also appreciate the effort undertaken to add a new theorem and proof, both of which were interesting to read. However, I recommend the authors polish the presentation of the theorem more for the next revision.

With all this in mind, I am keeping my score. I would have increased it to a 7 if the option was available, but I feel an 8 is too high for a paper that is not well-written.

We thank the reviewer for identifying our work's novelty, soundness, and technical contributions. The valuable comments have helped us improve our manuscript (marked blue in the revision). Below are our specific responses to the questions raised by the reviewer:

Weakness 1: While the paper uses a 1-10 scale, there's no analysis of how different scoring scales (e.g., 1-5, 1-100) might affect performance or what the optimal scale might be. We follow the reviewer's suggestions to conduct ablations on the impact of reward scales. Specifically, for the UltraFeedback dataset that contains response rewards in the range of 1-10, we relabel them to be in the range of 1-5 and 1-100 with linear transformation. Our method is then applied to these different scaled datasets. The results are shown in the following table:

It can be observed that our method is robust to the reward scales. Since our paper uses the default 1-10 scale as in UltraFeedback, it is likely that the performance can be further boosted, such as the 1-100 scale. The above results and discussion have been added to Appendix B.3 in the revision.

Weakness 2: The paper does not include human evaluation studies to validate the effectiveness of reward-augmented responses.

Our work mainly studied RL from AI feedback, and in experiments, we follow previous works and test our models on various benchmarks, including instruction-following tasks that use LLM as a judge, as well as academic benchmarks. We are not aware of benchmarks that use human evaluations and are cheap to run. If the reviewer has any such benchmarks in mind, please let us know and we will include the results in our manuscript.

Weakness 3: The paper would benefit from a discussion of limitations and potential future directions.

We thank the reviewer for the suggestion. The following paragraph has been added to the last section of our revision: "Although the proposed method consistently shows considerable improvements on instruction-following benchmarks, the improvements on academic benchmarks are marginal, which are also observed in previous direct alignment algorithms. Besides, instruction-following benchmarks adopt LLM-as-a-judge and have the risk of favoring specific output styles. Human evaluations might be needed to determine the benefit of our method. Since our method is compatible with any direct alignment algorithm, future works also include applying our method to other algorithms beyond DPO."

Question 1: What led to choosing the 1-10 scale for quality scores, and have you explored whether other scoring ranges might be more effective for reward conditioning?

• The 1-10 scale is chosen because our main experiments use the UltraFeedback dataset, which has a default scale of 1-10 for the quality scores.
• As suggested by the reviewer, we have ablated other scoring ranges, including 1-5 and 1-100. We observed that the method is robust to the scoring ranges and it is likely that the 1-10 range is not the optimal one. Please see our responses to your Weakness 1 and Appendix B.3 in the revision for details.

We hope the reviewer could consider raising the score if we resolved the reviewer's concerns. We would be happy to have further discussions if the reviewer has any additional questions or comments.

We thank the reviewer for identifying our work's novelty, soundness, and technical contributions. The valuable comments have helped us improve our manuscript (marked blue in the revision). Below are our specific responses to the questions raised by the reviewer:

Weakness 1: While the method is innovative, it builds heavily on existing direct preference alignment algorithms.

We appreciate the reviewer for finding our method innovative. We would like to emphasize that the key motivation of this paper is to address the limitations of direct preference alignment algorithms, such as the indiscrimination between responses of varying quality and the difficulty of generalizing effectively to the sparse optimal responses. On the contrary, for PPO-style alignment algorithms, preference data is used to fit the reward value, which directly plays a role during the RL optimization and thus avoiding drawbacks inherent to direct alignment methods (see Section 5 for a detailed discussion).

Weakness 2: The experiments are extensive but primarily focus on instruction-following and academic benchmarks. Including additional real-world applications, diverse datasets could strengthen the generalizability claims.

The goal of RLHF and RLAIF is to enable LLMs to follow instructions aligned with human intent. Accordingly, our experiments adhere to standard setups commonly used in alignment research, focusing primarily on instruction-following benchmarks such as AlpacaEval, MT-Bench, and Arena-Hard-Auto. Additionally, we evaluate our method on a variety of academic benchmarks. These evaluations encompass diverse real-world use cases, including generating human-preferred responses and solving mathematical problems. If the reviewer has specific benchmarks for real-world applications in mind, please let us know and we would be happy to incorporate evaluations.

Question 1: What are the benefits of keeping the rejected high-reward samples? Can the author provide evidence for the consequential effects of the unlearning problem when LLM finetuned with the original DPO? The authors might need to start with a model that is not DPO fine-tuned and start to fine-tune using DPO and the proposed method. We thank the reviewer for the great question. The following experiment results and discussions have been added to Appendix B.3 in the revision.

• Following the reviewer's suggestions, we conduct additional experiments to show the benefits of learning from the rejected high-reward samples. Specifically, using the UltraFeedback dataset, we construct two reward-augmented preference datasets by filtering out augmented data based on rejected responses with low and high reward values, respectively, which we denote as Ours-filter-low and Ours-filter-high. Compared to our method, these datasets isolate the impact of excluding low- and high-reward rejected responses as goals. The results on AlpacaEval 2.0 are shown as follows:

It can be observed that learning from rejected high-reward samples (Ours and Ours-filter-low) demonstrates superior performance compared to the approach that excludes these samples (Ours-filter-high), demonstrating the benefits of keeping the rejected high-reward samples.

• Following the reviewer's suggestions, we conducted an additional ablation study to show the effects of the unlearning problem on the LLM that is not DPO fine-tuned. Specifically, we fine-tune Zephyr-SFT [1] with DPO and our method. We report the log probability of the high-quality rejected responses in the test set of UltraFeedback as follows:

• It can be observed from the above table that the probability of high-reward responses after DPO significantly decreases compared to the SFT model and our method. Besides, the resulting model of our method also outperforms the DPO model by a considerable margin on the AlpacaEval 2.0 benchmark: