Paper: Statistical Rejection Sampling Improves Preference Optimization (ICLR-2024)

Reviewer comment (Theory):
The unifying link between DPO and SLiC appears to be deliberately designed, particularly noticeable in the normalization term, i.e., $\pi_{sft}(y|x)$, in Eq 10.

Author response:
>Re1: The unifying link between DPO and SLiC appears to be deliberately designed, particularly noticeable in the normalization term, i.e., $\pi_\text{sft}(y|x)$, in Eq 10.

Answer: Your observation about the unifying link between DPO and SLiC in the context of Eq 10 is indeed insightful. However, I believe there are a few additional points to consider, which might offer a different perspective.

1. The purpose of integrating $\pi_\text{sft}(y|x)$ into Eq 10 goes beyond merely unifying DPO and SLiC. This integration is primarily aimed at establishing a more robust theoretical framework for the SLiC-HF approach. While it's true that SLiC-HF and DPO emerged around the same time, leading to some confusion about their connection, their starting points are fundamentally different. 

    * SLiC-HF is grounded in the concept of likelihood calibration, where a better response is expected to have a higher likelihood than a less desirable response, by a defined margin.

    * On the other hand, DPO builds on the Reinforcement Learning from Human Feedback (RLHF) objective and the Bradley-Terry model, leading to the derivation of the sigmoid loss function as a Maximum Likelihood Estimation (MLE).

2. Our contribution, as outlined in the paper, is to theoretically link these two approaches. By substituting the Bradley-Terry model with a Support Vector Machine (SVM) model, we obtain the hinge-norm of SLiC, effectively a normalized version. The original SLiC calibration loss overlooks the SFT policy normalization, relying solely on preference pairs labeled by the reward model. The hinge-norm loss, conversely, uses the SFT likelihood as a baseline to compute the relative likelihood ratio between the current policy and the SFT policy. This approach not only calibrates the likelihood ratio towards human preferences but also acts as an adaptive margin version of the SLiC loss. The margin, in this case, is influenced by the SFT likelihoods of the better and worse responses. Furthermore, this paper [1] highlights issues with the Bradley-Terry model and proposes the IPO loss, which aligns more closely with the hinge-norm loss than the original SLiC.

3. In terms of implementation, DPO and SLiC variants share similarities by the choice of link functions and normalization. There isn't conclusive evidence yet to prefer one over the other. The paper provides a comprehensive framework for choosing between these options, underscoring the need for further research in this area.

In conclusion, while the connection between DPO and SLiC might seem deliberate, especially in the context of Eq 10, it is part of a broader effort to establish a more theoretically sound approach to SLiC-HF, taking into account the nuances of different models and the importance of the SFT policy.


**References**

[1] Azar, Mohammad Gheshlaghi, et al. "A general theoretical paradigm to understand learning from human preferences." arXiv preprint arXiv:2310.12036 (2023).

Reviewer comment (Theory):
The theoretical conclusion that the policy induced from RSO in Eq 4 is optimal seems questionable. RSO continues to use the proxy reward model to label responses, regardless of whether they're from the pre-trained or learned policy. However, the reward model is learned from $D_p$ and not $D_p^*$ without further updates, which could introduce approximation errors and bias in subsequent labeling and learning steps. In other words, the policy induced from Eq 4 is not optimal because of the non-optimal reward model.

Author response:
>Re2: The theoretical conclusion that the policy induced from RSO in Eq 4 is optimal seems questionable. RSO continues to use the proxy reward model to label responses, regardless of whether they're from the pre-trained or learned policy. However, the reward model is learned from $D_p$ and not $D_p^∗$ without further updates, which could introduce approximation errors and bias in subsequent labeling and learning steps. In the other word, the policy induced from Eq 4 is not optimal because of the non-optimal reward model.

Answer: Thank you for pointing out this important issue. We agree that RSO in Eq 4 is not strictly optimal. We have addressed the questions one-by-one.

(1) “RSO continues to use the proxy reward model to label responses, regardless of whether they're from the pre-trained or learned policy.”

In our revised approach, we acknowledge the impracticality of directly sampling from an optimal policy. To address this, we've shifted from using the term 'optimal policy' to 'estimated optimal policy' in our documentation, specifically when referencing the sourcing of preference pairs from the rejection sampled policy. This change underlines our strategy to approximate the optimal policy via a rejection sampling method. The distribution shift of reward model training data also exists in RLHF, even though it conducts sampling from online policy. The errors in the proxy reward model can result in reward hacking. To mitigate the issue, we can collect preference data from the current best policy iteratively. In Llama2 [1], they collect human preference data from the current best policy on a weekly basis, as the policy keeps improving. Our work can follow the same approach as we improve the policy iteratively using rso-sample-rank.

(2) “However, the reward model is learned from $D_p$ and not $D_p^*$ without further updates, which could introduce approximation errors and bias in subsequent labeling and learning steps”

We acknowledge the concerns regarding approximation errors and bias in our reward model learned from $D_p$. However, we assert that our pairwise reward model holds significant advantages over the implicit pointwise reward model (Equation (5)) induced by the Bradley-Terry model as per DPO, for several reasons:

1. **Pointwise vs Pairwise**: Our pairwise comparison approach, structured as “[Context] {context} [Response A] {response a} [Response B] {response b}”, aligns more naturally with human cognitive processes, compared with assigning a real-valued reward score for a response given a context.  The pairwise approach not only simplifies the task but also leverages the extensive knowledge transfer from pre-trained language models in classification and comparison tasks. Unlike the Bradley-Terry model, which imposes a rank-1 approximation thus potentially losing information, our model directly represents the complexity of pairwise interactions without such assumptions.

2. **Implicit vs Explicit**: Our model is an explicit text-to-text classification model, a task that is straightforward and less prone to the complexities associated with language generation. In contrast, the implicit reward model in DPO, based on likelihood scoring or generation, navigates a more intricate language output space. This complexity can hinder the effective transfer of knowledge from pre-trained models. Our approach, grounded in the principle that classification tasks are inherently more manageable than generation tasks, offers a more robust and reliable framework.

To substantiate our claims, our comprehensive experimental studies demonstrate the superiority of our model. For instance, rso-sample-rank shows to be significantly better than DPO. These results highlight not just the theoretical but also the practical advantages of our approach.


**References**

[1] Touvron, Hugo, et al. "Llama 2: Open foundation and fine-tuned chat models." arXiv preprint arXiv:2307.09288 (2023).

Reviewer comment (Experiments):
Works with a similar focus on rejection sampling, such as ReST (http://arxiv.org/abs/2308.08998), ALMoST (https://arxiv.org/abs/2305.13735), and RAFT (https://arxiv.org/abs/2304.06767), should be involved in comparison.

Author response:
>Re3: Works with a similar focus on rejection sampling, such as ReST(http://arxiv.org/abs/2308.08998), ALMoST (https://arxiv.org/abs/2305.13735) RAFT (https://arxiv.org/abs/2304.06767), should be involved in comparison.

Answer: Thank you for your suggestions regarding our baseline comparisons. We acknowledge the significance of ALMoST in the context of iterative LLMs and PPO algorithms. However, we opted not to include it in our baseline due to its distinct approach, particularly its reliance on synthetic preference pairs, which diverges significantly from RSO's methodology. This fundamental difference in approach could lead to an unfair comparison.

In contrast, RAFT and ReST share more similarities with RSO in terms of objectives and methodologies. For RAFT, we continue to fit the SFT model using cross-entropy loss on the best response selected by tournament ranking using a pairwise reward model. For ReST, we fix $\tau=0.7, I=1, G=1$ as suggested by the ReST paper. We normalize the reward scores between 0 and 1 and choose the responses with reward greater than 0.7 as new SFT targets. We did not include multiple rounds of grow and improve stages because of a fair comparison consideration. RSO can also be done in multiple rounds to generate better rso-sample-rank preference pairs.

Our results demonstrate that RSO outperforms these robust baselines across various metrics, highlighting its effectiveness. While we are confident in our findings, we remain open to future studies that may explore these comparisons further, potentially integrating models like ALMoST under a framework that allows for a fair comparison.  

| Task          | Approach           | Proxy Reward (%) | Gold Reward (%) | AutoSxS (%) |
| ------------- | ------------------ | ---------------- | --------------- | ----------- |
| Reddit TL; DR | RAFT               | 74.84            | 68.51           | 53.77       |
|               | ReST               | 49.03            | 46.17           | 34.36       |
|               | $RSO_{sigmoid-norm}$ | 92.37            | 82.22           | **71.86**       |
|               | $RSO_{hinge-norm}$   | **92.80**           | **83.45**           | 70.84       |
| AnthropicHH   | RAFT               | 58.21            | 40.00           | 24.99       |
|               | ReST               | 43.48            | 30.33           | 15.58       |
|               | $RSO_{sigmoid-norm}$ | **86.94**            | **59.15**           | **40.98**       |
|               | $RSO_{hinge-norm}$   | 84.44            | 57.75           | 38.58       |

Reviewer comment (Experiments):
Why do not directly sample responses from $\pi_{r_{\psi}}(y|x)$ (with a reward threshold like ReST) rather than sampling from $\pi_{sft}(y|x)$ with the statistical rejection sampling to approximate the samples from $\pi_{r_{\psi}}(y|x)$?

Author response:
>Re4: Why do not directly sample responses from $\pi_{r_\psi}(y|x)$ (with a reward threshold like ReST) rather than sampling from $\pi_\text{sft}(y|x)$ with the statistical rejection sampling to approximate the samples from $\pi_{r_\psi}(y|x)$?

Answer: Thank you for raising this important question. There are two reasons that we don’t directly sample responses from $\pi_{r_\psi}(y|x)$:

1. **Online vs Offline**: ReST is an online on-policy algorithm, while RSO is an offline algorithm. Online algorithm allows the model to sample responses from the current policy, while the offline algorithm only has access to the initial policy (SFT policy) to generate responses. RSO bridges the gap by sampling responses from a better policy that is closer to the optimal policy. Online algorithms like ReST and PPO are not time and memory efficient. ReST needs policy network and value network to co-exist during training, and sampling from the online policy is not parallelizable. PPO needs two more networks (reference policy network and value network). On the other hand, RSO can generate rso-sample-rank preference pairs from SFT policy with a rejection sampling algorithm in a fully parallelizable way. The training of the RSO policy also only needs the policy network, which is much more memory efficient than ReST/PPO.

2. **Connection with ReST**: in our paper, we demonstrate that the prevalent best-of-N or top-k-over-N rejection sampling methods are merely specific instances of our more comprehensive statistical rejection sampling algorithm, particularly when the beta parameter is set to zero. The ReST paper is very similar to an online version of the top-k-over-N. In Figure 3(b), we show RSO is better than the best-k-over-N. We also verify that in our updated draft in comparison with ReST. If beta approaches infinity, our method aligns seamlessly with sampling from an SFT policy. This crucial insight underpins our proposed method, which adeptly balances reliance on the SFT policy with the accuracy of the reward model. We use the hyper-parameter beta to control the balance between the two and thus offer a more robust and adaptable solution for practical applications.

Reviewer comment (Novelty):
The core idea of utilizing rejection sampling to filter the response feels somewhat trivial.

Author response:
Answer: Your feedback about the perceived triviality of using rejection sampling in our approach is appreciated. We understand that at first glance, rejection sampling might seem simplistic. However, the key contribution of our work lies in the context of its application and the nuanced improvements we have implemented.
1. **Bridging a Crucial Gap**: Our primary aim is to bridge the gap between offline off-policy algorithms and approximately on-policy algorithms. Current approaches like DPO and SLiC-HF have not effectively addressed the distribution-shift issue. Our work is pioneering in conducting a systematic study to approximate the optimal policy among offline human preference alignment algorithms.
2. **Connect with Existing Best-of-N**: In our paper, we demonstrate that the prevalent best-of-N or top-k-over-N rejection sampling methods are merely specific instances of our more comprehensive statistical rejection sampling algorithm, particularly when the beta parameter is set to zero. More importantly, we establish that as beta approaches infinity, our method aligns seamlessly with sampling from an SFT policy. This crucial insight underpins our proposed method, which adeptly balances reliance on the SFT policy with the accuracy of the reward model. Unlike existing methods, ours provides a nuanced approach that adapts dynamically to varying degrees of confidence in either the SFT policy or the reward model, thereby offering a more robust and adaptable solution for practical applications.
3. **Efficiency Enhancements**: We acknowledge the concerns raised by reviewer hBgS regarding the computational expense of calculating $M$. To address this, we have innovated by not computing
$M$ directly. Instead, we estimate $\frac{\pi_{r_{\psi}}(y|x)}{M\pi_{\text{sft}}(y|x)}$ in its entirety using 64 sequences sampled through the SFT policy. This methodological change, detailed in Appendix A.1, significantly enhances efficiency.
4. **Improved Rejection Sampling Method**: Our method diverges from traditional statistical rejection sampling by sampling from the proposal distribution without replacement. This approach is more efficient in high-dimensional settings where the probability of rejecting many proposals is high. By recalculating the maximum reward at each round’s start among unselected candidates, we ensure the selection of at least one candidate—the one with the highest reward. This modification guarantees efficiency, particularly when the target is a limited number of selected responses.
5. **Theoretical Backing and Simplicity**: Our algorithm, while simple in its execution as demonstrated in Algorithm 1 in Appendix A.1, is grounded in robust theoretical principles. The calculation of the threshold for a uniform random sample is simplified to computing exp((reward-max_reward) / beta), a result of deriving a closed-form solution for the optimal policy.

In summary, our approach, although founded on a basic concept like rejection sampling, is a novel and effective solution for the challenges in offline human preference alignment algorithms. The improvements we have made are not just in terms of algorithmic efficiency, but also in bridging a significant research gap. The simplicity of the method should not undermine its effectiveness and the novelty of its application in this context.

Reviewer comment (Experiments):
There might be flaws in the design of the experimental part: the reward for training and testing is the same. Although there are evaluations from GPT/human, generally speaking, these are not very convincing. You can refer to [1, 2] for Synthetic Data Setup, Recalibration, gold reward, and proxy reward to have a fair comparison.

Author response:
Answer: Thank you for emphasizing the need for a robust evaluation framework in our experimental design. Your suggestions regarding the utilization of distinct rewards for training and testing phases, as well as the incorporation of a gold reward model, are indeed pertinent.

In response, we have revised our approach following the methodologies outlined in [1]. Specifically, we have now implemented larger reward models (PaLM-2 S) to serve as our Gold Reward model. This model is not only used as an additional metric but also helps in differentiating the rewards between the training and testing phases, thereby addressing your concern about the potential overlap in rewards.

Moreover, to ensure transparency and thoroughness in our evaluation, we have included detailed metrics on the evaluation accuracy of the Gold Reward model in Section 5 of our updated draft. This additional layer of evaluation not only bolsters the validity of our findings but also aligns our methodology more closely with current best practices in the field.
We believe these modifications greatly enhance the rigor of our experimental setup and appreciate your valuable input in guiding these improvements.

**References**

[1] Gao, Leo, John Schulman, and Jacob Hilton. "Scaling laws for reward model overoptimization." International Conference on Machine Learning. PMLR, 2023.

Reviewer comment (Writing):
There are still rooms for improving the presentation. For example, it would be better if the authors can explain a bit more on $\rho_{\psi}(x,y,y_b)$ when first introducing it. The authors could provide some examples or details about what would be an instantiation of $\rho_{\psi}(x,y,y_b)$.

Author response:
Answer: Thanks for highlighting the need for additional clarity. We admit that the presentation is not very clear at this part. In response, we have expanded the explanations of $\rho_{\psi}(x,y,y_b)$ in our revised draft. Furthermore, to illustrate its practical application, we have included specific examples of inputs and outputs, ensuring a clearer understanding of its implementation.

Reviewer comment (Experiments):
Step 3 in section 3.2 might be overly time-consuming and computationally expensive. Especially the set $M = \min\\{m | m\pi_{sft}(y|x) \geq \pi_{r_{\psi}}(y|x) \text{ for all } y \notin \mathcal{Y}\\}$ is expensive to compute. Rejection sampling is also slow.

Author response:
Answer: Thank you for your insightful observation regarding the computational aspects. The reviewer is correct that directly computing $M$ can be quite expensive and rejection sampling can be slow. We clarify the following two points regarding the issue:

1. **Computing $M$**: To address this, it's important to clarify that we do not directly compute $M$. Instead, our approach involves estimating $\frac{\pi_{r_{\psi}}(y|x)}{M\pi_{\text{sft}}(y|x)}$ using 64 sequences sampled through the SFT policy, which offers a practical solution to the computational challenge you mentioned. For detailed insight into this methodology, including both the Python implementation and the algorithmic derivation, please refer to Appendix A.1 of our paper. We added a footnote to clarify this in our updated draft.

2. **Rejection sampling efficiency**:
    * Concerning the speed of the rejection sampling process, our method begins by sampling 64 sequences via the SFT policy, followed by rejection sampling. This initial SFT sampling step is notably efficient and highly parallelizable, particularly when prompts are shared, thereby enhancing the overall computational efficiency of the process. 
    * Regarding the statistical rejection sampling algorithm, as described in Algorithm 1, it exhibits enhanced efficiency by employing a sampling-without-replacement strategy. This is achieved by excluding the selected sequences subsequent to each sampling round. Furthermore, at the commencement of each round, the maximum reward is recalculated. This recalibration ensures that, in every round, at least one sequence is invariably chosen. Specifically, the sequence whose reward is equivalent to the maximum reward is selected with a probability of one, thereby guaranteeing the selection of at least one optimal sequence in each round. This approach not only optimizes the selection process but also maintains the algorithm's effectiveness throughout its execution.

We have added a section discussing the computational cost and efficiency of our approach in Appendix A.10 of the updated draft.

Reviewer comment (Evaluation):
The results in Figure 3 (b) is not statistically signifcant and the confidence intervals are overlapping.

Author response:
Answer: Thank you for your insightful feedback regarding the statistical significance and confidence intervals in Figure 3(b). We appreciate the opportunity to clarify and reinforce the findings of our study.

1. **Purpose of the Plot**: The primary objective of Figure 3(b) is to illustrate our process of hyper-parameter selection, focusing exclusively on the proxy reward model. This model's win rate against the SFT targets is depicted on the y-axis. It's crucial to note that the plot is to guide hyper-parameter choices. 

2. **Comparison among loss functions**: We don’t claim the preferred choice between sigmoid-norm and hinge-norm loss functions. Regarding hinge loss, it performs the best (not statistically significantly) because hinge loss ignores the SFT policy completely and fully trusts the preference order determined by the proxy reward model. That explains why the proxy reward model win rate is high.

3. **Significance of rso-sample-rank**: Despite the overlapping confidence intervals, a critical observation from the plot is that the rso-sample-rank consistently outperforms sft-sample-rank across all three loss functions. This is visually represented by the dotted lines (sft-sample-rank) remaining outside the shaded regions of all rso-sample-rank curves. 

4. **Evaluating top-k-over-N and Beta Selection**: The plot also challenges the notion that top-k-over-N is universally optimal. When beta equals zero, rso-sample-rank aligns with top-k-over-N. Interestingly, we observe that a beta of 0.5 achieves a higher proxy reward model win rate than beta=0. Although this difference is not statistically significant, it indicates that a nuanced approach to beta selection could yield better results than the default top-k-over-N strategy.

5. **Inclusion of Additional Baselines**: To robustly validate our findings and address the concern of statistical significance, we have incorporated additional baselines, namely RAFT [1] and ReST [2], into our revised draft. These inclusions will provide a broader comparative framework, enhancing the validity and reliability of our conclusions.

**References**

[1] Dong, Hanze, et al. "Raft: Reward ranked finetuning for generative foundation model alignment." arXiv preprint arXiv:2304.06767 (2023).

[2] Gulcehre, Caglar, et al. "Reinforced self-training (rest) for language modeling." arXiv preprint arXiv:2308.08998 (2023).

Reviewer comment (Experiments):
How does the proposed method compare against other methods, such as DPO, in terms computational efficiency?

Author response:
Answer: Thank you for raising the computational complexity question. In terms of computational efficiency, our approach is closer to SLiC-HF-sample-rank than DPO. We list the comparisons with other algorithms as follows: 

* Compared to DPO and SLiC-HF-direct, our approach requires an additional trained pairwise reward model, and also sample and rank stages. The pairwise reward model training is a standard text-to-text model without Bradley-Terry model assumption. The sample and rank are scalable and efficient with increased number of servers, especially with shared prompts in SFT sample and short decoding length with reward model.

* Compared to SLiC-HF-sample-rank, we need to sample more sequences from SFT policy as candidates, followed by a rejection sampling algorithm. Both steps are scalable and efficient as described in response to reviewer hBgS.

* Compared to RLHF, RSO is offline with efficient memory usage (only one copy of policy network instead of four copies in PPO: policy, reference, value, reward), and the sample-rank stage is fully parallelizable (no need to conduct the online sample, which is not parallelizable).

To clarify this, we added Appendix A.10 for thorough discussion in our updated draft.