TGDPO: Harnessing Token-Level Reward Guidance for Enhancing Direct Preference Optimization

Q1 Comparison with existing works

All [1], [2] and ours are based on the assumption of dense token-level reward to derive respective DPO loss functions. Eq. (8) in our work seems similar to Eq. (12) in [1], but our Eq. (8) is not solvable since $s_t\sim\mathcal{D}\_t$ is dependent on the policy $\pi\_{\theta}$ to be optimized. This has been demonstrated in line 219 (left) - line 168 (right) of the main text. In the final version we will give citations, discuss connections with [1, 2] and the differences highlighted as follows:

Motivation Differences: We intend to incorporate a learned token-level reward explicitly into DPO. While, they only focus on asigning more weight on earlier tokens.

Loss Function Differences: Our work explicitly incorporate an existing token-level reward in the loss function for guiding DPO. [1] uses the same loss as DPO and modifies the sampling probability for different diffusion timesteps. [2] adds a temporal decay factor for each token in the DPO loss. Their loss functions do not leverage token-level reward.

Method Differences: Integrating token-level reward in DPO's loss function explicitly presents distinctive difficulties, especially in the derivation. Thanks to all reviewers' comments, we have proposed a new approach for canceling partition functions in the Bradley-Terry model, please see Q7.

[1] A Dense Reward View on Aligning Text-to-Image Diffusion with Preference

[2] Earlier Tokens Contribute More: Learning Direct Preference Optimization From Temporal Decay Perspective

Q2 Accurate word choice

We will change to use "response-level reward" and "KL-regularized policy optimization" to improve accuracy.

Q3 Explain "sentence-level PPO with token-level reward guidance"

A sequence of dense token-level reward $r_{\phi}(s_t, a_t)$ is obtained before PPO, and then it can guide PPO finetuning in a more fine-grained way.

Q4 What is the corresponding dense token-level reward

Actually, $\pi_{\hat\theta}$ does not necessarily relate to $\hat{r}$ in Assumption 4.2. $\pi_{\hat\theta}$ is only used for sampling $s_t$ in Equation (10). Now, the first sentence of Assumption 4.2 is modified as: Suppose we have learned a dense token-level reward $\hat{r}$ using some effective learning approach. In the revision, we will improve the description and notation to reduce misunderstanding.

Q5 Difference between $\hat{r}$ and $r_{\phi}$

$\hat{r}$ is an existing dense token-level reward, used for shaping the reward $r_{\phi}$ in Equ. (10). Following the derivation of DPO, $r_{\phi}$ is expressed with $\hat r$, $\pi_{\theta}$ and $\pi_{\text{ref}}$ by solving Equ. (10). And finally, $\hat{r}$ will appear in Equ. (16) of our loss function, but $r_{\phi}$ does not.

If we optimize $\pi_\theta$ against $\hat r$, then Equ. (10) becomes

$$\max_{\pi_{{\bf\theta}}} \mathbb{E}\_{s_t\sim\hat{\mathcal{D}}\_t, a_t\sim \pi_{\theta} (\cdot|s_t) } \left[{r_{\phi}(s_t, a_t)} - {\beta f(\hat{r}(s_t, a_t))}\log \frac{ \pi_{\theta} (a_t|s_t)}{\pi_{\text{ref}}(a_t|s_t)}\right].$$

We tried and found it is difficult to solve the problem due to the expectation on the product term.

Q6 Requirement for $f$

In Assumption 4.2, the $f$ satisfying $f(0)=1$ and $|f(u)-1|\le\varepsilon$ is sufficient for deriving our TGDPO. Concrete $f(u)$ can be chosen by developers personally. As for us, Equ. (17) is adopted in our experiments. If $f(u)\equiv 1$ then our TGDPO degenerates to DPO, which is not interesting.

Q7 Strong assumption for Equation 15

Thank you very much for the comment! Now the assumptions have been completely removed, the approach is novel. The theorem and proof can be seen in response to Q3 of Reviewer F88J. We will provide revisions in the final version.

Q8 Benefit of introducing $f(\hat r)$ in Equ. (10)

The benefit of introducing $f(\hat r)$ in Equ. (10) is in Equ. (16) of the loss function. We adopt Equ. (17) in our experiments. Then take win response $y_w$ for an example, the other is similar. In Equ. (16),

$$f(\hat{r}([x, y_w^{<t}], y_w^t)) \log\frac {\pi_{\theta} (y_w^t|[x, y_w^{<t}])}{\pi_{\text{ref}}(y_w^t|[x, y_w^{<t}])} = (1 + \alpha \hat{r}([x, y_w^{<t}], y_w^t)) \log\frac {\pi_{\theta} (y_w^t|[x, y_w^{<t}])}{\pi_{\text{ref}}(y_w^t|[x, y_w^{<t}])}.$$

For $\hat{r}([x, y_w^{<t}], y_w^t)>0$, we emphasize more on this token, otherwise less or keep the weight unchanged as 1.

Q9 Improper description in L286, L406-408

Thank you very much for pointing out the issues. We modify them as:

1. L283-286 (right): "Then this token is optimized with less strength during the optimization of the loss function $\mathcal{L}\_{\text{TGDPO}}(\pi_{\theta})$, since $f(\hat{r}([x, y_w^{<t}], y_w^t)) <1$."
2. L296-299 (right): Similar to point 1.

3. L406-408 (left): "distinguish preferred tokens in chosen samples and dispreferred tokens in rejected ones, TGDPO takes care of them and enables ..."

Q1 Clarification regarding Equ. 8

By [1] (Thanks to Reviewer aioX), Equ. (8) has connections to an approximation approach common in prior RL works [2, 3, 4]. [1] pointed out this by providing 5 papers including [2, 3, 4], and followed this line to derive the loss function for their DPO, please see Appendix B.2.1 in [4]. Hence, this is a reasonable approach. We will give citations and provide related illustrations.

Moreover, it was pointed out in [5] that "[6] showed that under the Max-Entropy RL formulation, the token-level log-ratio $\log \frac{\pi_{\theta}(y|x)}{\pi_{\text{ref}}(y|x)}$ can be seen as an implicit token-level reward or advantage function (invariant under reward shaping)."

Notably, our derivation is mainly from the relaxation approach in optimization. Although we start from the state-action reward, the final derived loss function coincides with the streamline of DPO. Specially, when the $\hat{r}(s_t, a_t)$ is not available, we may simply set $\hat{r}(s_t, a_t)=0$ and by Assumption 4.2 our loss function in Equ. (16) is precisely the loss function of DPO, which demonstrates our approach is reasonable.

Q2 Explanation for $f$ in Assumption 4.2

Thanks. $f(u)$ is set such that $f(0)=1$ and $|f(u) - 1| \le \varepsilon$. Concrete $f(u)$ can be chosen by users for their personal usage. Please also check our response to Q8 of Reviewer aioX to see benefits of introducing $f$.

Q3 Strong assumption for Equ. 15

Thank you for the insightful comment! We have resolved this issue based on the finding that $\delta( f, \hat{r}; x, y_w, y_l)$ does not depend on the policy $\pi_{\theta}$ to be optimized, and the assumptions are removed. The approach is new for canceling partition functions in the BT model.

Now Equ. (15) and its proof have been reorganized as follows:

$$\Pr(y_w \succ y_l | x) = \sigma\left( \varphi(\pi_{\theta}, f, \hat{r}; x, y_w, y_l) + \delta( f, \hat{r}; x, y_w, y_l) \right), \qquad (1)$$

where $\delta( f, \hat{r}; x, y_w, y_l)$ does not depend on the policy $\pi_{\theta}$ to be optimized, but only on $f, \hat{r}, x, y_w, y_l$ and the partition function $Z(s_t)$ (does not depend on $\pi_{\theta}$, see Theorem 4.3). For briefness, let $h\triangleq (f, \hat{r}; x, y_w, y_l) .$

Since $\sigma(t)$ is the sigmoid function with $\sigma'(t)>0$, we have:

Theorem 1. The $\Pr(y_w \succ y_l | x)$ in Equ. (1) has the same maximal solution and the same ascent direction as the function $\sigma\left( \varphi(\pi_{\theta}, h)\right)$ with respect to $\pi_{\theta}$.

Proof. Note that, $\delta( h)$ is not dependent on the policy $\pi_{\theta}$ and $\sigma'(t)>0$. $d$ is an ascent direction of function (1) if $d^T \nabla_{\pi_{\theta}} \sigma\left( \varphi(\pi_{\theta}, h) + \delta(h) \right) >0,$ which is equivalent to

$$\begin{aligned} & d^T \sigma'\left( \varphi(\pi_{\theta}, h)+ \delta(h)\right) \nabla_{\pi_{\theta}} \varphi(\pi_{\theta}, h) >0 \\ & \Longleftrightarrow d^T \sigma'\left( \varphi(\pi_{\theta}, h)\right) \nabla_{\pi_{\theta}} \varphi(\pi_{\theta}, h) >0 \\ & \Longleftrightarrow d^T \nabla_{\pi_{\theta}} \sigma\left( \varphi(\pi_{\theta}, h)\right) >0. \end{aligned}$$

Hence function (1) has the same ascent direction as the function $\sigma\left( \varphi(\pi_{\theta}, h)\right)$. Similarly,

$$\begin{aligned} & \nabla_{\pi_{\theta}} \sigma\left( \varphi(\pi_{\theta}, h) + \delta( h) \right) =0 \\ & \Longleftrightarrow \sigma'\left( \varphi(\pi_{\theta}, h)\right) \nabla_{\pi_{\theta}} \varphi(\pi_{\theta}, h) = 0 \\ & \Longleftrightarrow \nabla_{\pi_{\theta}} \sigma\left( \varphi(\pi_{\theta}, h)\right) = 0. \end{aligned}$$

So Theorem holds.

Thus by Thm 1, since we focus only on optimal $\pi_{\theta}$ of function (1), we may redefine $\Pr(y_w \succ y_l | x) \triangleq \sigma\left( \varphi(\pi_{\theta}, f, \hat{r}; x, y_w, y_l)\right),$ and use it for constructing the loss function in Equ. (16).

We will update the statement and proof of Equ. (15). Thanks!

Q4 Why $r(s_t,a_t) = \beta \log \pi_{\theta} / \pi_{\text{ref}}$

It was shown in [6] that $r(s_t,a_t) = \beta \log \pi_{\theta} / \pi_{\text{ref}}$ under the definition of equivalent state-action reward class and invariant re-parameterization, which does not require $V(s_{t}) = V(s_{t+1})$. Please see Theorem 1 there. Moreover, the final loss function in [6] is equivalent to that in [7]. Hence we adopt the equation directly. It is a common practice in many works [5, 8].

[1] A dense reward view on aligning text-to-image diffusion with preference

[2] Approximately optimal approximate reinforcement learning

[3] Relative entropy policy search

[4] Trust Region Policy Optimization

[5] Self-Play Preference Optimization for Language Model Alignment

[6] From r to Q*: Your language model is secretly a Q-function

[7] Direct preference optimization: Your language model is secretly a reward model

[8] Free process rewards without process labels

Q1 TGDPO's dependence on the choice of $\hat{r}$

Trained token-level rewards have shown effectiveness for PPO [1, 2]. For DPO, it is interesting to ask if there exists a framework that can incorporate a trained token-level reward explicitly for better performance. Our TGDPO fills this nontrivial gap.

Q2 Strong assumption for Equ. 15

Thanks for the insightful comment! Now the assumptions have been removed, the approach is novel. The theorem and proof can be seen in Response to Q3 of Reviewer F88J. We will make revisions in the final version.

Q3 How to get the Equation in Lines 311-313 (left)? How to get $\hat{r}$ in practice?

Thm 1 of [3] shows $r(s_t,a_t) = \beta \log \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\text{ref}} (a_t|s_t)}$ under equivalent state-action reward class and invariant re-parameterization. The final loss function in [3] is equivalent to that of DPO [4]. So we adopt the result by setting $s_t= [x, y^{<t}]$ and $a_t= y^t$ and get the Equation. This is a common practice in many works [5].

To obtain a token-level reward $\hat{r}$ in practice, we can use off-the-shelf open-sourced models trained by DPO, other methods [1, 2], or run DPO by ourselves.

Thanks. We will illustrate more for this equation in the final version.

Q4 Clarification on win rate and score

The win rate is evaluated by calling a judge model (e.g., gpt-4o-2024-11-20) to pairwise compare the responses of the model and a baseline model (e.g., gpt-4-0314) and determine which one is better [6]. The MT-Bench score is evaluated by calling a judge model to directly assign a score to the model's response [7]. Many prior works observed the scores of different methods being similar [6, 8]. This is likely due to the single-instance scoring protocol.

Q5 Comparison with TDPO

Below is the experiment result with TDPO under the Llama3-8B-Instruct PairRM setting.

Result with TDPO using the SFT model OpenRLHF/Llama-3-8b-sft-mixture is in Q1 of Reviewer GJSH. The result table with TDPO under other settings is in this link https://anonymous.4open.science/r/tgdpo_rebuttal/results_with_tdpo.pdf

Q6 Will the choice of $f$ affect the results

$f(u)$ may be adjusted by parameter $\alpha$ as in Equ. (17). With this, experiment results in Fig. 1 reveal in some degree the effects of different choices of $f(u)$. The outcome benchmark difference presented in Fig. 1 and Tab. 3 is subtle.

Our paper proposes a framework for incorporating existing token-level reward into DPO explicitly, with Equ. (17) as an example. Many choices of $f(u)$ can satisfy Assumption 4.2, which one is the best is left as an open problem for stimulating further research.

Q7 Clarification on convergence and checkpoint selection

In Sec. 5.3, we consider loss moving average below a certain threshold (e.g., 0.1) as convergence. In other sections, we train all methods (DPO, SimPO, TGDPO) for 1 epoch and select the checkpoint at the training end.

Q8 Stability of results

Below are the Arena-Hard win rate of Llama3-8B-Instruct PairRM and the 95% confidence interval. From the table we can see all methods have similar levels of stability.

Q9 Would error in token-level reward affect results

By Assumption 4.2, $|1-f(\hat{r}(s_t, a_t))|\le\varepsilon$ where $\varepsilon$ is small. Moreover, by the choice of $f(\hat{r}(s_t, a_t))$ in Equ. (17), the error can be reduced by the parameter $\alpha$. Thus mild errors in token-level reward may not affect the outcome greatly, as demonstrated in Fig. 1 and Tab. 3.

Q10 Why does $\alpha$ affect the convergence rate

Take win response $y_w$ for an example, the other case is similar. In Equ. (16), for $\hat{r}([x, y_w^{<t}], y_w^t)>0$ and larger $\alpha$, the gradient w.r.t. this token is larger, and the convergence would be faster.

Q11 Code release

The code will be released upon acceptance.

[1] Preference-grounded token-level guidance for language model fine-tuning

[2] Segmenting text and learning their rewards for improved RLHF in language model

[3] From r to Q*: Your language model is secretly a Q-function

[4] Direct preference optimization: Your language model is secretly a reward model

[5] Free process rewards without process labels

[6] From Crowdsourced Data to High-Quality Benchmarks: Arena-Hard and BenchBuilder Pipeline

[7] Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena

[8] SimPO: Simple Preference Optimization with a Reference-Free Reward

Q1 Experiment on SFT model

Thank you for the advice! Below we show the experiment results on UltraFeedback using the open-sourced SFT model OpenRLHF/Llama-3-8b-sft-mixture, which has not been trained by RLHF. Our TGDPO using SimPO's token-level reward achieves much better performance than baselines. Specifically, it achieves win rate gains of 10.5 on AlpacaEval 2 and 3.9 on Arena-Hard compared to best-performing baselines.

Q2 Setting difference between SimPO (Meng et al., 2024) and ours

The key difference in the experiment setting is that we use the latest and more powerful gpt-4o-2024-11-20 as the judge model, while SimPO uses gpt-4-1106-preview (gpt-4 turbo), which was released in Nov 2023.

Below we compare the Arena-Hard win rate of Llama3-8B-Instruct PairRM using these two judge models. The win rate judged by gpt-4-1106-preview is generally consistent with the SimPO paper, while the result is different from the latest gpt-4o-2024-11-20. This is the reason for SimPO underperforming DPO in Table 1.

Q3 Starting point of TGDPO training

We would like to clarify that DPO, SimPO, and our proposed TGDPO have the same starting point for training, which are Instruct or SFT models. As described in Equation (16), TGDPO is designed to leverage any token-level reward models, including pre-trained DPO or SimPO models. We can use off-the-shelf open-sourced token-level reward models. So it is not necessary to train a token-level reward model by ourselves before starting TGDPO. Furthermore, TGDPO can enjoy faster convergence speed with satisfactory performances, as demonstrated in Figure 1 and Table 3 of this paper. Also, naively increasing the training budget for DPO or SimPO usually does not improve the performance. As demonstrated in Table 2, training DPO and SimPO to convergence leads to worse performance. We will clarify the description to avoid misunderstanding.

Q4 Clarification of Equ. 15

We followed the standard way of conducting experiments and TGDPO consistently outperforms baselines. We believe this empirical validation is a key strength of our approach, as it shows that the theoretical insights lead to practical improvements.

From theory aspect, now the related assumptions have been completely removed, the approach is novel, and the intuition is outlined below:

Let

$$$$

The Bradley-Terry model in Equ. (15) is

$$\Pr(y_w \succ y_l | x) = \sigma\left( \varphi(\pi_{\theta}, h) + \delta(h) \right).$$

We find that $\delta( h)$ does not depend on the policy $\pi_{\theta}$ to be optimized, but only on $f, \hat{r}, x, y_w, y_l$ and the partition function $Z(s_t)$ (does not depend on $\pi_{\theta}$, see Theorem 4.3 in the main text). Since $\sigma(t)$ is the sigmoid function with $\sigma'(t)>0$, the above $\Pr(y_w \succ y_l | x)$ has the same maximal solution and the same ascent direction as the function $\sigma\left( \varphi(\pi_{\theta},h)\right)$ with respect to $\pi_{\theta}$.

Hence, since we focus only on the maximal solution $\pi_{\theta}$, we may redefine $\Pr(y_w \succ y_l | x) \triangleq \sigma\left( \varphi(\pi_{\theta}, h)\right),$ and use it for constructing the loss function of our preference optimization.

The detail can be seen in response to Q3 of Reviewer F88J. We will provide the intuition in the final version.

Q5 Difference with TDPO (Zeng et al., 2024)

Our work aims to leverage an existing token-level reward to guide DPO training at the token level. Whereas, TDPO aims to enhance the regulation of KL-divergence by constraining each token with forward KL-divergence. It is not guided by a token-level reward.

We will add more discussions on the differences with TDPO in the final version.

Q6 Clarification on line 215

The paragraph demonstrates that it is possible to obtain a token-level reward, using the approach in [1, 2] or DPO. And, the token-level reward will be adopted for shaping the reward in PPO and subsequently the loss function of DPO.

We will make the description better.

[1] Preference-grounded token-level guidance for language model fine-tuning

[2] Segmenting text and learning their rewards for improved RLHF in language model