Mitigating Reward Over-optimization in Direct Alignment Algorithms with Adaptive Importance Sampling

Thank you for your thoughtful review and valuable feedback. Please find our response to your comments below.

1. Can the authors theoretically justify that the estimator is unbiased? How does importance sampling estimator depend on the optimal policy $\pi^{*}$?

To show why importance sampling is unbiased. From [10], we have the objective

In RSO, they propose to use rejection sampling to approximate the optimal policy in order to construct preference data $\mathcal D\_p=\{(x,y\_w,y\_l)|y\_w,y\_l\sim\pi^*(\cdot|x)\}$ and then fitting a parametrized policy $\pi\_\theta$ with maximum likelihood objective, resulting in:

Instead of using RSO to sample from the optimal policy, one can directly parameterize the optimal policy as $\pi\_\theta$ for collecting preference data and reward learning (as shown in [10, 12]), resulting in:

This corresponding to Online-DPO objective, where the parametrized policy $\pi\_\theta$ is now serve as an approximation to $\pi^*$ for generating preference data and reward learning. This is also our desire objective.

Asuming that $\pi\_\theta$ is fully supported on $\pi\_{\text{ref}}$, the online-DPO objective can be written:

which shows that DPO with Importance sampling (our objective) is an unbiased estimator of Online-DPO.

2. Can this paper provide comparisons with online algorithms such as PPO, REINFORCE, and online DPO? Except for PPO, which requires extensive computational resources, other methods require nearly the same resources as DPO.

We conducted further experiments where we compare DAAs with online alignment methods. As PPO requries extensive computational resources, we consider an alternatives online algorithms which is REINFORCE Leave-One-Out (RLOO) [7]. Detail training setting and a figure presenting the result can be found in our updated manuscript (Section C in the Appendix). As expected, RLOO achives better win-rate compared to DPO and AIS-DPO and utilizing better KL budget. The result also shows that AIS-DPO help close the gap between offline and online algorithms.

3. Can the authors clarify the methods used for KL calculation and "gold win-rate" calculation in the experiments?

The calculation of KL divergence in our paper is based on [8] and the KL is estimated by taking on-policy samples from the current LM $\pi\_\theta$.

Specifically, we first sample $N$ input prompts $\\{\mathbf{x}\_i\\}\_{i = 1}^N$ from the evaluation set. For each input prompt $\mathbf{x}\_i$, we generate a response $\mathbf{y}\_i$ using the current policy $\pi\_\theta$. Let $T\_i$ be the length of the response $\mathbf{y}\_i$, we compute the KL divergence between $\pi\_\theta$ and $\pi\_{\text{ref}}$ as follows:

We set $N=512$ in our experiments.

For Gold win-rate calculation, we first use a well fine-tuned pythia 6.9b in [9]. The model achieve $\approx 70$% accuracy in evaluation set and achieving 76.7% agreement ratio with ChatGPT Evaluation using LLM as a judges prompting.

For a given prompt $x$, we first sample a response $y\sim\pi\_\theta(\cdot|x)$ and then use the golden reward model $r^{\text{gold}}$ to compare against reference summaries $y\_{\text{ref}}$ in evaluation set to determine the win-rate with the following calculation:

We have included the calculation of KL divergence and gold win-rate calculation in Section B, Appendix in the revised manuscript.

4. Can the authors discuss the issue of length bias in the importance sampling weight? The length bias affects reward estimation [5], so the reviewer wonders about its effect on the importance sampling estimator used in this paper.

The length of responses could affect the variance of the importance sampling estimators. Since our importance weights are computed as the ratio of probabilities from two distributions, accumulated over all tokens. The variance of importance sampling estimators can grow exponentially with the number of tokens in the output. To address this, we applied exponential smoothing of importance weights, using a smoothing factor $\alpha = \frac{1}{T}$. This adjustment reduces the variance significantly, preventing exponential growth with the response length. We provide a detailed analysis in the general response and in Sections E and F of the Appendix in our revised manuscript.

5. The superiority over other simple baselines is unclear. A straightforward way to address the distribution shift issue is to use a moving average of the reference policy that can ensure the policy moving beyond the KL contraint.

While it is straightforward to address the distribution shift issue is to periodically updating the reference policy [5] or using entropy regularization [6]. Both of these methods require online-sample from the current LM policy in the loop of training and and explicit reward model, incurring significant computational costs.

While in our paper, we working on an offline setting, we only have access to a static offline dataset. we propose to use importance sampling to mitigate distribution shift problem by adding a simple weighting to simulate Online RLHF without requiring generating samples and explicit reward model, which is computationally expensive.

Q3: Unfortunately, (1) the above explanation does not address my concerns. (2) I still believe that strong numerical evidence on large models is necessary to justify the technique or, alternatively, (3) a rigorous theoretical analysis should be provided as a valuable contribution. At present, (4) the standard importance sampling technique with a length normalization trick does not appear to be a significant technical contribution for me.

We annotate the Reviewer's points ((1), (2), (3), and (4)) for easier references.

As we extensively explained in our rebuttal (an incorporated in our revised paper), along with the response in Q2 of this response, the design of our Adaptive IS is rigorously analyzed, both theoretically and empirically, which we hope it could resolve concerns (1), (3), and (4) of the Reviewer. Additionally, AIS addresses the reward over-optimization in existing DAAs, where the weak regualrization in DAAs stemming from the mismatch between training data distribution and the current LM policy, which is a significant technical contribution in this domain (as also appreciated by Reviewers u6Dx and zEKh); we hope that for this reason, we answer the Reviewer's concern in (4).

Finally, we'd like to defend our paper, focusing on the Reviewer's comments:

• We closely follow the exact same experimental setup (including TL;DR dataset and the Pythia model) in many directly related works [1,3,11,5] that study reward over-optimization in RLHF and DAAs. The fact that AIS improves the performance of existing DAAs in this well-established experimental setup already demonstrates the significance of our contributions and the practicality of the proposed method.

  • While we understand the Reviewer's suggestion to evaluate "on large models", we sincerely think that our contributions are significant and our method is practical, without such evaluation, similar to the acceptance of the related works [1,3,11,5]. Additionally, evaluating on larger models requires high computational resource and laborous effort as we also need to replicate many existing baselines on these large models. Nevertheless, we followed the Reviewer's suggestion on larger models. In the limited Rebuttal window, we evaluate our method with Pythia-2.8b and present the win-rate evolution as training progress in the below tables.

  Method	Step-14400	Step-28800	Step-43200	Step-57600	Step-72000	Step-92820
  DPO	0.658	0.753	0.472	0.371	0.564	0.675
  Adaptive-IS DPO	0.625	0.673	0.763	0.703	0.718	0.75

  The table below shows the KL divergence evolution as training progress with Pythia-2.8b model.

  Method	Step-14400	Step-28800	Step-43200	Step-57600	Step-72000	Step-92820
  DPO	9.92	15.673	26.184	32.59	28.25	26.15
  Adaptive-IS DPO	9.158	12.662	16.595	22.193	24.1	22.463

  In this setting, Adaptive-IS still shows consistent performance as the training progress while DPO performance start to degrade. Moreover, Adaptive-IS also maintains better KL regularization.

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

1. The motivation of adaptive smoothing factor

In auto-regressive language models, our importance weights are computed as the ratio of probabilities from two distributions, accumulated over all tokens. If the model’s distribution $\pi_\theta$ differs greatly from the reference distribution $\pi_{\text{ref}}$, the variance of importance sampling estimators can grow exponentially with the number of tokens in the output. To address this, we applied exponential smoothing of importance weights, using a smoothing factor $\alpha = \frac{1}{T}$. This adjustment reduces the variance significantly, preventing exponential growth with the response length. We provide a detailed analysis in the general response and in Sections E and F of the Appendix in our revised manuscript.

2. Ablation over alpha in TL;DR task

We conducted additional experiments where the value of $\alpha$ is fixed instead of adaptive (Section F, Appendix). The value of $\alpha$ are vary in $\{0.0, 0.05, 0.1, 0.2, 0.4 \}$ in the AIS-DPO loss function on Reddit TL;DR dataset. We observed that increasing $\alpha$ helps increase the regularization effect and win rate. However, up to a specific point (around 0.1), the regularization effect starts to diminish due to high variance in the importance ratio. Moreover, Adaptive-IS with $\alpha=1/T$ achieves the best result.

We also provide an ablation of $\alpha =\frac{1}{\sqrt{|y|}}$. Although it achieves lower win-rate than DPO and AIS-DPO. We speculate that setting setting $\alpha=\frac{1}{\sqrt{y}}$ can still have high variance in the importance ratio, leading to a small number of samples having enormous weights can potentially dominate learning signals of other valuable samples [1]. Square root IS-DPO still shows better regularization effect than other 2 objectives. Thus, the value of $\alpha$ can be used to control regularization effect in DPOs.

3. How does AIS perform in Section 4.3 and when the preference dataset is not generated from the SFT policy

We presented new results of applying IS on the distribution shift experiment in Section G, Appendix. Our observations are: (1) when the data distribution is close to $\pi_{\text{ref}}$, Adaptive IS and DPO show similar performance in terms of win rate and KL divergence, but both still suffer from the distribution shift effect, and (2) as the data distribution shifts away from $\pi_{\text{ref}}$, AIS-DPO is shown to achieve better regularization and win rate compared to standard DPO even when the data is no longer from $\pi_{\text{ref}}$. This phenomenon is helpful in practice, where the preference data is usually generated from an unknown policy $\mu$, not from $\pi_{\text{ref}}$. AIS-DPO can still improve performance and regularization when $\pi_{\text{ref}}$ is not far from $\mu$.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback. We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

[1] Is Value Learning Really the Main Bottleneck in Offline RL?

Thank you once again for your valuable feedback. We are pleased that our responses addressed your concerns and appreciate your support for the paper.

We appreciate your constructive comments on our paper. Please find our response to your comments below.

1. Why do we need to estimate $E_{\pi_{\theta}}[\rho_\theta]$? What is the ideal optimization objective? Our motivation to derive the original equation comes from the fact that DAAs can be derived from vanilla policy gradient (VPG) ([4], [5]), an on-policy algorithm. We have shown the equivalence in Section H, Appendix, where DAAs can be seen as maximizing binarized reward using policy gradient. However, this equivalence only holds when we consider the online version of DPO or IPO. In off-policy setups, DAAs can suffer from the distribution shift problem, which has been well-studied in offline RL literature [3]. This also explains the ineffectiveness of regularization in DAAs when using offline data due to the sampling bias in the regularization objective [9, 10].

A method to reduce the effect of sampling bias is using importance sampling to estimate the loss function under $\pi_\theta$ distribution given samples from a reference distribution $\pi_{\text{ref}}$ [3, 11]. This leads to importance weights that we proposed in the paper.

We thank the reviewer for their feedback and will respond to the raised questions below.

1. This method is inherently offline and not online, so importance sampling is just for reducing the variance, right?

In this work, we follow the offline settings presented in [3,1], where only a fixed dataset is provided; we do not rely on model-generated samples or external rewards for learning.

The role of our importance sampling approach does not actually aim to reduce the variance but to adjust the distributional differences between the offline data and the current policy, for better approximation in the online alignment algorithms. Specifically, importance sampling will reweight samples from the offline data such that those with high likelihood under the current learned model, $\pi_\theta$, are given more "weights", and those with low likelihood will be down-weighted. This helps to prevent over-optimizing samples that already have low-likelihood under $\pi_\theta$.

2. The reward used is dependent on the policy generating the response? So, its dependent on $\theta$ right? In line 977, while taking the derivative that needs to be considered as well right? Can you please derive this considering the reward being considered in the paper and showing what the actual optimization problem is, its not clear still.

In standard RL, the reward is a fixed function that assigns a scalar reward to each state-action pair and it does not depend on the policy parameters; thus the gradient of the reward with respect to the policy parameters will be zero. In RLHF, there exists an additional KL divergence term $\text{KL}(\pi_\theta||\pi_{\text{ref}})$. One can define the additional KL regularization as maximizing the new reward $r(x,y)$: $\mathbb E_{y\sim\pi_\theta(\cdot|x)}\left[r(x,y)\right]=\mathbb E_{y\sim\pi_\theta(\cdot|x)}\left[r_\phi(x,y)-\beta\log\frac{\pi_\theta(y|x)}{\pi_{\text{ref}}(y|x)}\right]$

Although this reward depends on the policy's parameters due to the addition of $-\log\pi_\theta(y|x)$ term, the gradient of this new reward w.r.t the policy's parameters will still vanish using the log-derivative trick ([6]):

The property results show that if we replace the reward in policy gradient objective with a binary signal, which depends on which generation is preferred, we follow the same gradient as IPO. In this sense, policy gradient also subsumes direct alignment approaches (see in [5]).

3. Can you provide detailed explaination on expanding $\rho_\theta=0$ since, it doesn't only refers when $\pi_\theta$ and $\pi_{\text{ref}}$ are close or far, but also when $\pi_\theta(y_x|x)$ and $\pi_\theta(y_l|x)$ are far, this won't hold right?

Note that the log-ratio difference $\rho_\theta=\log\frac{\pi_\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)}-\log\frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)}$ depends on both $\pi_\theta$ and $\pi_{\text{ref}}$. If $\pi_\theta(y_w|x)$ and $\pi_\theta(y_l|x)$ are far apart (i.e the difference is high), when the reference model is close to $\pi_\theta$, it would similarly exhibit a large difference between $\pi_{\text{ref}}(y_w|x)$ and $\pi_{\text{ref}}(y_l|x)$. This would ensure that the log-ratio difference $\rho_\theta$ remains small, even when $\pi_\theta(y_w|x)$ and $\pi_\theta(y_l|x)$ are far apart.

We have provided a detailed explanation of expanding $\rho_\theta=0$ in Section J of the Appendix. Our analysis indicates that during training, there is a mismatch in data distribution: KL divergence is computed based on samples from the current LM, while Offline regularization in DAAs is calculated based on offline samples. This can result in cases where minimizing offline regularization objectives might not necessarily reduce KL divergence. The reason is that the offline samples may not accurately reflect those produced by the current policy, which can potentially lead to performance degradation (see [2], [4]).

In summary, we have answered the questions posed by the reviewers. We hope that these questions clarify all the reviewer's concerns about our work. Nevertheless, if there are any additional questions, please kindly let us know during the rebuttal period and we will promptly provide our responses.

[1] Understanding the performance gap between online and offline alignment algorithms.
[2] Generalized Preference Optimization: A Unified Approach to Offline Alignment
[3] Scaling Laws for Reward Model Overoptimization in Direct Alignment Algorithms
[4] Preference Learning Algorithms Do Not Learn Preference Rankings
[5] Contrastive Policy Gradient: Aligning LLMs on sequence-level scores in a supervised-friendly fashion
[6] Notes on policy gradients and the log derivative trick for reinforcement learning

Thank you for your review and positive feedback on our paper. We appreciate your acknowledgment of the clarity and organization of our work.

Please find our response to your comments below.

1. How was the KL computed? Was it estimated by taking on-policy samples from the current LM? The calculation of KL divergence in our paper is based on [1] and the KL is estimated by taking on-policy samples from the current LM $\pi_\theta$.

Specifically, we first sample $N$ input prompts $\\{\mathbf{x}\_i\\}\_{i = 1}^N$ from the evaluation set. For each input prompt $\mathbf{x}\_i$, we generate a response $\mathbf{y}\_i$ using the current policy $\pi\_\theta$. Let $T\_i$ be the length of the response $\mathbf{y}\_i$, we compute the KL divergence between $\pi_\theta$ and $\pi_{\text{ref}}$ as follows:

We set $N=512$ in our experiments. For clarity, we have included the calculation of KL divergence between $\pi_\theta$ and $\pi_{\text{ref}}$ in the appendix.

2. What is the result of online alignment algorithm such as PPO in the main evaluation in Section 4.2? How much of the online-offline gap can be explained by distribution shift and can be addressed by AIS?

We agree that answering this question could provide greater insight. We conducted further experiments where we compare DAAs with online alignment methods. As PPO requries extensive computational resources, we consider an alternatives online algorithms which is REINFORCE Leave-One-Out (RLOO) [2]. Detail training setting and a figure presenting the result can be found in our updated manuscript (Section C in the Appendix). As expected, RLOO achives better win-rate compared to DPO and AIS-DPO and utilizing better KL budget. The result also shows that AIS-DPO help close the gap between offline and online algorithms.

3. The necessity of the adaptive heuristic

In auto-regressive language models, our importance weights are computed as the ratio of probabilities from two distributions, accumulated over all tokens. If the model’s distribution $\pi_\theta$ differs greatly from the reference distribution $\pi_{\text{ref}}$, the variance of importance sampling estimators can grow exponentially with the number of tokens in the output. To address this, we applied exponential smoothing of importance weights, using a smoothing factor $\alpha = \frac{1}{T}$. This adjustment reduces the variance significantly, preventing exponential growth with the response length. We provide a detailed analysis in the general response and in Sections E and F of the Appendix in our revised manuscript.

4. How does AIS perform in the synthetic distribution shift experiment and when the preference dataset is not generated from the SFT policy?

We presented new results of applying IS on the distribution shift experiment in Section G, Appendix. Our observations are: (1) when the data distribution is close to $\pi_{\text{ref}}$, Adaptive IS and DPO show similar performance in terms of win rate and KL divergence, but both still suffer from the distribution shift effect, and (2) as the data distribution shifts away from $\pi_{\text{ref}}$, AIS-DPO is shown to achieve better regularization and win rate compared to standard DPO even when the data is no longer from $\pi_{\text{ref}}$. This phenomenon is helpful in practice, where the preference data is usually generated from an unknown policy $\mu$, not from $\pi_{\text{ref}}$. AIS-DPO can still improve performance and regularization when $\pi_{\text{ref}}$ is not far from $\mu$.

6. Incorrect citation formats and typos in writting Thank you for these remarks. We corrected citations and typos in our revised manuscript.

We hope that our response addresses the issues you raised. Please let us know if you have any additional concerns or questions.

[1] Understanding the performance gap between online and offline alignment algorithms.
[2] Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs