The Crucial Role of Samplers in Online Direct Preference Optimization

Questions:

[A1] No, we have stated in the preliminary that "we will omit the prompts (contexts) and slightly abuse the notations throughout Sections 3 and 4," and thus, no $x$ needs to be involved. The results can be easily adapted to the contextual bandit setting (depending on $x$), so our assumption is without loss of generality.

[A2] (i) As a common notation in machine learning literature, $h(\theta) = \text{sg}(f(\theta)) \cdot g(\theta)$ leads to $\nabla_\theta h(\theta) = f(\theta) \cdot \nabla_\theta g(\theta)$.

(ii) We use $\pi^s(y,y')$ to simplify the expression of $\pi^{s1}(y)\pi^{s2}(y') + \pi^{s1}(y')\pi^{s2}(y)$, which is the probability that the pair $(y,y')$ or $(y',y)$ is sampled by the sampler pair $(\pi^{s1}, \pi^{s2})$. If your question is about the design of $(\pi^{s1}, \pi^{s2})$ in different sampling regimes (uniform, with known reward, and practical setting), please refer to our response to all reviewers: "Explanation of sampler design."

(iii) No, we don't mean to compare $\pi^{s1}$ and $\pi^{s2}$. $\pi^{s1}$ and $\pi^{s2}$ form a sampler pair, each for one response $y_1$ and $y_2$ in a data pair $(y_1,y_2)$, respectively. When we say "sampling regimes," we refer to DPO-Unif, DPO-Mix-R, and DPO-Mix-P, corresponding to uniform sampling, sampling with known reward, and the practical setting, respectively. The reason why we use different $\alpha$s can be found in our response to all reviewers.

(iv) Since we omit the prompts, Eq. (4) does not contain an expectation over $x$. As we stated, the results can be easily extended to the case with $x$, the complete form of Eq. (4) (and in practice) should contain such an expectation.

[A3] Regarding the notations, we will add explanations in the revision.

(i) Yes, $G^{(t)} \in \mathbb{R}^{|\mathcal{Y}|}$ and $G_y^{(t)}$ is the $y$-th entry of $G^{(t)}$. This can be inferred from Def. 2, which indicates that $G^{(t)}$ has the same shape as $\theta^{(t)}$.

(ii) Yes, $G_y^{(t)}$ is the true gradient plus a sub-Gaussian noise scaled by $\beta A$.

(iii) Gaussian noise is a special case of sub-Gaussian noise (please see https://en.wikipedia.org/wiki/Sub-Gaussian_distribution#Examples). We believe that it is better to have a more general result by modeling the noise under a broader class.

[A4] Please refer to our response to all reviewers: "Explanation of sampler design."

[A5] Thank you for this suggestion. We prove it here and will add this proof to the appendix in the revision.

From the rule of gradient descent Eq. (5), we know that $\theta^{(t+1)} = \theta^{(t)} - \eta\alpha\nabla_\theta\mathcal L(\theta^{(t)})$ (Eq. A). Note that $\delta(y,y';\theta^{(t+1)}) = r(y) - r(y') - \beta\theta_y^{(t+1)} + \beta\theta_y^{(0)} + \beta\theta_{y'}^{(t+1)} - \beta\theta_{y'}^{(0)}$ (Eq. B). Apply Eq. A to Eq. B, we get $\delta(y,y';\theta^{(t+1)}) = \delta(y,y';\theta^{(t)}) + \eta\beta\alpha\nabla_{\theta_y}\mathcal L(\theta^{(t)}) - \eta\beta\alpha\nabla_{\theta_{y'}}\mathcal L(\theta^{(t)})$. Then apply $\nabla_{\theta_y}\mathcal L(\theta) = -\beta\sum_{y'}\pi^{s}(y,y')\Delta(y,y')$, which is shown in the equation above Eq. (6), we finally get Eq. (6).

[A6] Thank you for pointing this out! We forgot to explain this in our main content. The $x$-axis is the number of gradient updates, and the $y$-axis is the total parameter difference $\sum_{y, y'} \delta(y, y'; \theta^{(t)})^2$.

You said $h(\theta)=sg\big[f(\theta)\big]\cdot g(\theta)$ means $\nabla_{\theta}h(\theta)=f(\theta)\cdot \nabla_{\theta} g(\theta)$.

Based on that, the definition of $\pi^s(y,y')={\rm sg}\big[\pi^{\rm s1}(y)\pi^{\rm s2}(y')+\pi^{\rm s1}(y')\pi^{\rm s2}(y)\big]$ in Definition (1) still looks vague.

I cannot find stopping gradient operator by both Google and AI search. They all refer to the criterion of when to stop (stochastic) gradient descent algorithm, which seems far from your definition.

Could you write down explicitly the definition of $\pi^s(y,y')={\rm sg}\big[\pi^{\rm s1}(y)\pi^{\rm s2}(y')+\pi^{\rm s1}(y')\pi^{\rm s2}(y)\big]$ in both comment and edited paper (can be uploaded now)? The most reasonable guess I can think of is $\pi^s(y,y')=\big[\pi^{\rm s1}(y)\pi^{\rm s2}(y')+\pi^{\rm s1}(y')\pi^{\rm s2}(y)\big]/2$ as the integral of each policy is 1, right? This is important as Definition 1 is the basis of this paper.

Thanks.

You can find the stopping gradient operator in various locations: TensorFlow (https://www.tensorflow.org/api_docs/python/tf/stop_gradient), JAX (https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.stop_gradient.html), and (using a different name, called detach) PyTorch (https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html). The definition of $\pi^s$ is to treat it as a static number instead of a function which might depend on $\theta$, which will be used in calculating the gradient of Equation (4), leading to the results in Section 4.1. The step using the stopping gradient is: \nabla_\theta \left\\{ -\sum_{y,y'\in \mathcal{Y}} \pi^s(y,y') p^\star(y\succ y')\log\sigma\left(\beta\log\frac{\pi_\theta(y)\pi_{\text{ref}}(y')}{\pi_{\text{ref}}(y)\pi_\theta(y')} \right) \right\\} = -\sum_{y,y'\in \mathcal{Y}} \pi^s(y,y') p^\star(y\succ y') \nabla_\theta \log\sigma\left(\beta\log\frac{\pi_\theta(y)\pi_{\text{ref}}(y')}{\pi_{\text{ref}}(y)\pi_\theta(y')} \right) , as $\pi^s$ has its gradient stopped.

The reason that we use stopping gradient operator is that it is usually used in existing online DPO approaches: for each iteration, they first sample a new dataset and then train on it, thus there is no gradient required from the sampler.

We will update our paper shortly once our new experiments are finished.

Weaknesses

[W1] While I agree that studying the stylized setting in the paper is a useful starting point, it would be useful to at least include some more discussion around the question of whether the insights in the paper extend.

[AW1] Yes, we do believe that "studying the stylized setting in the paper is a useful starting point"! We agree with your concern regarding small response spaces and bandit setting, and have already listed it as a future direction (limitation 2) in Section 6. A starting point would be log-linear parameterization, where the reward is parameterized as $r(y) = r^\top \phi(y)$, and the policy is parameterized as $\pi_\theta(y) \propto \exp(\theta^\top\phi(y))$. Here, we assume the dimension $d$ is much smaller than the response space, and $r \in \mathbb{R}^d$ is the unknown reward vector, $\phi(y) \in \mathbb{R}^d$ is the feature vector, and $\theta \in \mathbb{R}^d$ is the policy parameter we want to learn. We've found that, if the covariance matrix $\sum_{y,y'}(\phi(y)-\phi(y'))(\phi(y)-\phi(y'))^\top$ is full rank, then we can learn the optimal policy parameter $\theta^\star = \theta_{\text{ref}} + r/\beta$. Thus, we don't need to loop over all possible actions: if we have a small amount of responses $y_1, \ldots, y_m$ such that $\sum_{i=1}^m \sum_{j=1}^m (\phi(y_i)-\phi(y_j))(\phi(y_i)-\phi(y_j))^\top$ is full-rank, then it suffices for policy learning. Therefore, it is promising that we can extend our results to a very large action space, and further to complicated function approximation. We will add this discussion in the revision.

[W2] It would be useful to see some errors bars/confidence bounds to get a sense for whether the improvement the authors find is significant.

[AW2] We agree with this point, and will elaborate our tables in the revision. This needs to run additional experiments and may take some time. We will let the reviewer know when they are finished!

Questions

[AQ1] Please refer to our answer [AW1].

[AQ2] Please refer to our answer [AW2].

We've updated our tables in the revision, where we provide means and standard deviations. We are happy to answer any further questions!

Questions

[AQ1] Please refer to our response to all reviewers: Explanation of sampler design.

[AQ2] Please refer to our response to all reviewers: Explanation of sampler design.

[AQ3] In fact, your proposed version was the main motivation for this work: to directly study the convergence guarantee of on-policy DPO. However, if not considering posterior distribution, we cannot establish the quadratic convergence when both $\pi^{s1}$ and $\pi^{s2}$ are $\pi_\theta$ in Definition 3, as the coefficients of each $\Delta$ in the proof of Theorem 1 would be different, posing significant difficulty in converting $\Delta (y, y''; \theta) - \Delta (y', y''; \theta)$ to $\delta (y, y'; \theta)$, which is the central idea in our proof. The difficulties are similar when we change the distributions in Definitions 4 and 5. But these difficulties inspire us to carefully design the samplers.

Besides, if we view $\pi_\theta$ as a posterior distribution on $Y$, we can still have advantages. Please see our answer [AW2] for this point.

[AQ4] Our results indicate that $\delta(y,y';\theta^{(t)})$ for $y,y'$ in $Y$ are simultaneously optimized toward $0$. In Section 5, by the performance difference lemma, we show that $V^*-V^\theta\le\underset{y\sim\pi^\star,y'\sim\pi_\theta}{\mathbb E}\delta(y,y';\theta)\le\sqrt{\underset{y\sim\pi^\star,y'\sim\pi_\theta}{\mathbb E}\delta^2(y,y';\theta)}$, demonstrating that $\delta(y,y';\theta)$ contributes more to the performance when the joint probability $\pi^\star(y) \pi_\theta(y')$ is high. This inspires us to plug a posterior distribution on $Y$, as discussed in [AW2].

[AQ5] Please refer to our response to all reviewers: Explanation of evaluation.

[AQ6] Please refer to our answer [AW4].

[AW4] Thank you for your insightful questions! There are generally two points of view in RLHF when evaluating the final performance: 1) KL-regularization is only to stabilize the RLHF training; 2) The goal of RLHF is to balance the reward and KL-divergence from $\pi_{\text{ref}}$. For the first view, we provide Tables 2 and 3; and for the second, we provide Figure 2.

The proposed method outperforms baselines only in large KL divergence. Why?

There are two main reasons explaining the advantages that exist mostly in the large KL-divergence regime. 1) The closed-form solution of DPO is the solution of $\max_{\pi}\underset{y\sim \pi}{\mathbb E} r(y)-\beta\operatorname{KL}(\pi \Vert \pi_{\text{ref}})$, or equivalently, $\max_{\pi}\underset{y\sim\pi}{\mathbb{E}}r(y)$ subject to $\text{KL}(\pi\Vert\pi_{\text{ref}})\le C$, where $C$ is a constant induced by duality. When the KL-divergence is small, the policy space is relatively small, and thus the performance won't differ much. 2) Although in theory we are exponentially faster, there may still exist some large constants that slow down the convergence in the first two iterations.

The improvements in rewards and win-rate appear to be modest.

Small advantages are still meaningful. In Tables 2 and 3, the win-rate improvements might be relatively small but are still acceptable in LLM literature such as [5], especially taking our restricted computation into consideration. Our primary goal is not restricted to showing that our proposed sampler performs best, but is intended to show that all the cited DPO variants fit into our framework seamlessly. The results further align with our claims: 1) Vanilla DPO is DPO-Unif with the posterior distribution on $\mathcal{Y}$ set as $\pi_{\text{ref}}^{2\beta}$, on-policy DPO is DPO-Unif with the posterior distribution on $\mathcal{Y}$ set as $\pi_\theta^{2\beta}$ (which is closer to $\pi^\star$), thus on-policy DPO is better than vanilla DPO; 2) Hybrid-GSHF is approximately DPO-Mix-P (which is better than DPO-Unif) with the posterior distribution on $\mathcal{Y}$ set as $\pi_{\text{ref}}^{\beta}\pi_\theta^{\beta}$ (which is closer to $\pi^\star$), thus Hybrid-GSHF is better than vanilla DPO; 3) Our proposed sampler is DPO-Mix-P with the posterior distribution on $Y$ set as $\pi_\theta^{2\beta}$, and experiments validate that it performs better than all others. We thus believe the benefits indeed exist, as reflected in these experiments.

Whether the model overfits to the reward model?

For the reward overfitting issue, please refer to our response to all reviewers: Explanation of evaluation.

Whether the comparison is fair?

We want to clarify that theoretically all methods should converge to the same optimal solution (though they may still differ in reality), and thus we want our algorithm to obtain a good policy quickly. Now let us rethink the comparison, vanilla DPO cannot obtain a larger (but acceptable) KL quickly, and then the reward increases slowly, while ours quickly achieves a higher reward at a low cost of KL-divergence, indicating that the comparison is reasonable: it means ours converges faster than vanilla DPO.

[5] Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF. https://arxiv.org/abs/2405.21046.

Weaknesses

[AW1] Please refer to our response to all reviewers: Explanation of sampler design. If we didn't understand your problem accurately, please don't hesistate to reach out.

[AW2] Thank you for the suggestion! Bridging the gap between theory and practice is one of this paper's goals. To clarify this point, we have slightly abused the notation $\mathcal{Y}$, by viewing it as an action set with a posterior distribution.

It means that we don't need to modify the theory and only need to assume a different posterior distribution on $\mathcal{Y}$. Then, the theoretical samplers would naturally change as shown in Section 5, "Setting the Posterior." Therefore, $(\pi_{\text{ref}}, \pi_{\text{ref}})$ is represented by $(\text{Uniform}(\mathcal{Y}), \text{Uniform}(\mathcal{Y}))$ with a posterior of $\pi_{\text{ref}}$. It is similar for other regimes. In other words, our theories mainly focus on the policy difference $\log\frac{\pi^{s1}}{\pi^{s2}}$ between the heterogeneous samplers (since we can eliminate the $\text{Uniform}(\mathcal{Y})$ in policy difference like $\log\frac{\text{Uniform}(\mathcal{Y})}{\text{Uniform}(\mathcal{Y})}=\log\frac{\pi_{\text{ref}}}{\pi_{\text{ref}}}=0$). To align our theory to practice, a concern would be that which posterior distribution is more useful and practical? In Section 5, we show that $\pi^\star$ would be a good one, but $\pi^\star$ is too costly to obtain, as shown in [1], we thus plug $\pi_\theta^{2\beta}$ as a compromise (this may lose a bit theoretical soundness since $\pi_\theta$ is not fixed, but it works well in practice).

For the explanation of the mixing coefficients, please refer to our response to all reviewers: Explanation of sampler design.

[AW3] We agree with this point. As cited in Section 2, there are some works [2][3][4] studying policy gradient methods with access to exact gradients that have obtained impactful results. Inspired by this line of work, we believe our theoretical findings can serve as a useful starting point, motivating the community to further explore the empirical setting. Therefore, we provide Theorems 5 and 6 for an initial understanding of empirical DPO algorithms.

[1] Statistical Rejection Sampling Improves Preference Optimization. https://arxiv.org/abs/2309.06657.

[2] On the global convergence rates of softmax policy gradient methods. ICML 2020.

[3] Ordering-based conditions for global convergence of policy gradient methods. NeurIPS 2023.

[4] On the theory of policy gradient methods: Optimality, approximation, and distribution shift. JMLR 2021.

Weaknesses:

[AW1] Please refer to our response to all reviewers: Explanation of evaluation.

[AW2] Please refer to our response to all reviewers: Explanation of sampler design for the intuition. For now, we cannot prove that this mixed sampler is optimal by providing lower bounds, nor can we devise a scheme with faster convergence (like $\delta^{(t+1)} \le 0.9 |\delta^{(t)}|^{2.1}$). This is a valuable direction for future work. However, achieving quadratic convergence is already non-trivial in the optimization theory literature.

[AW3] Our main contribution lies in the theoretical part (we are the first to show a theoretical gap in using samplers in RLHF from the perspective of optimization, a significant topic), and the experiments are designed to demonstrate the potential of our framework, as restricted by limited computing resources.

Questions:

[AQ1] Although we have not provided a theoretical lower bound on uniform DPO under the empirical setting, which we would like to leave as an open problem, we can offer an intuitive explanation for this point. Refer to Section 4.1.1 and Appendix A.1.1, where we demonstrate the linear convergence of DPO-Unif. To do so, we need to establish a lower bound $\sigma_{\min}'$ on $\sigma'(\log\frac{\pi_\theta(y)\pi_{\text{ref}}(y')}{\pi_{\text{ref}}(y)\pi_\theta(y')})$, after which the convergence rate becomes $2-8\sigma_{\min}'$. Note that $\sigma'(x)$ decreases as $\vert x\vert$ increases, thus we need to upper bound $\vert\log\frac{\pi_\theta(y)\pi_{\text{ref}}(y')}{\pi_{\text{ref}}(y)\pi_\theta(y')}\vert$. However, when faced with noisy gradients, $\vert\log\frac{\pi_\theta(y)\pi_{\text{ref}}(y')}{\pi_{\text{ref}}(y)\pi_\theta(y')}\vert$ might deviate significantly when $\eta=\frac{1}{\beta A}$, and then it cannot converge as fast as DPO-Mix-R or DPO-Mix-P*. Note that the approach to circumvent $\sigma'_{\min}$ in the proof for DPO-Mix-P* is not applicable here. Moreover, in numerical simulation experiments shown in Figures 1 and 4 (Appendix D.1), we find that our proposed samplers consistently outperform DPO-Unif significantly under empirical settings. Thus, we believe our method is more efficient than DPO-Unif empirically.