Enhancing Multi-Step Reasoning Abilities of Language Models through Direct Q-Function Optimization

Q1: Fairness of Comparison

A1: We thank the reviewer for the suggestion and added more pairs to DPO training data. We ensure that under all setting the number of pairs of DPO is at least half of the numbers of trajectories for training DQO, which means that DPO is always using more trajectories than DQO. Specifically, the size of our new training set of DPO are shown as follows

The following table summarizes the results of new results on Qwen.

The following table summarizes the results of new results on Gemma

The results are also updated in our paper. We see that DQO still outperforms DPO and other baselines.

Q2: Lack of Experimental Details; How were generated responses evaluated for each augmented dataset? What is the distribution of positive and negative responses in the training and testing data?

A2: We thank the reviewer for the suggestions. In fact, as shown in Appendix A, our prompt requires the model to wrap its answer after "####" for GSM8K and in "\box{}" for MATH. Then for each response we use regular experession to extract the answer and use the official scripts to decide whether the extracted answer is correct. Following the suggestion of the reviewer, since the testing data only includes problems, we report the distribution of positive and negative responses in our generated training data as follows.

We have also rewrite related parts in our revision.

Q3: Hard to assess the benefits of MDP formulation

A3: We would like to remind the reviewer that importance sampling is a key component when applying soft-actor-critic to offline settings. Specifically, to estimate the $V$-fucntion of current policy $V^{\pi}(s)=E_{a\sim \pi_\theta(\cdot | s)}[Q(s,a)]$, we need to take expectation over all action $a$ with probability $\pi_\theta(a|s)$. However, since the training data is sampled from $\pi_\text{ref}$, taking expectation over $\mathcal{D}$ is essentially computing $E_{a\sim \pi_{\text{ref}}(\cdot | s))}[Q(s,a)]$, which actually equals to $V^{\pi_\text{ref}}$. Therefore, to bridge the mismatch in distribution, we have to incorporate important sampling to reweight the samples to make it align with the distribution generated from $\pi_{\theta}$. To summarize, importance sampling should be interpreted as part of the SAC algorithm rather than a stand-alone component of DQO.

Q8: For Qwen2-7B-Instruct, why does DPO consistently outperform KTO and DRO, while DRO outperforms DPO for Gemma-1.1-7B-it

A8: We believe that this is because the gemma is not performing well on both datasets, therefore there might be more prompt that all trajectories cannot solve the problem and these data cannot contribute to DPO training. However, these portions of trajectories can be utilized by DRO. As a result, DRO outperforms DPO on gemma. On the other hand, the Qwen base model is performing well on both datasets, with means that there are sufficient pairs for DPO and the benefit of utilizing imbalanced data is not prominent in this case. This makes DPO outperform DRO on Qwen.

Q9: The related work [2] reports pass rates of 82.3% and 49.6% for Qwen2-7B-Instruct on GSM8K and MATH, respectively. Why does this differ from the results in Table 3?

A9: We believe that this might be due to different prompting strategy. As we have shown in Appendix A, we wrap the questions in the following template to prompt the LLMs.

• GSM8K {{problem}} \n Please reason step by step, and produce a final answer following 4 \#', like #### 0'.
• MATH: {{problem}} \n Please reason step by step, and put your final answer within \boxed{}.

which might differ from the prompt template used in [2]. This discrepancy might introduce some discrepancy in perforamance.

Q10: Could the authors explain why DQO with lambda=1 outperforms lambda=0.95?

A10: We believe that this is because of the inaccurate initial estimation of V-function. Intuitively, when $\lambda=1.0$, the computed value target does not rely on the value of later states estimated by the value model, while $\lambda=0.95$ relies on such inaccurate value estimation. This introduces noisy signal to the training and causes instability, resulting to the phenomenon that $\lambda=1$ outperforms $\lambda=0.95$.

Q11: What is the correlation between the learned value function and constructed process scores during testing?

A11: We believe that the learned value function and constructed process scores are still fairly related on the test set. However, in our paper, the naive process score construction involves splitting each responses and sample 20 completions for each prefix of each responses, which is costly in both time and computational resources. Therefore due to time limit we are not able to report the results currently. We will report the exact results once they are available.

Q12: Minor Issues

A12: Thank you for pointing out these typos and we have revised them accordingly.

References:

[1] Xu, Ruijie, et al. "Benchmarking benchmark leakage in large language models." arXiv preprint arXiv:2404.18824 (2024).

[2] Lai, Xin, et al. "Step-dpo: Step-wise preference optimization for long-chain reasoning of llms." arXiv preprint arXiv:2406.18629 (2024).

Thank you for your constructive feedback! We answer your questions as follows

Q1: It seems that the author just apply SAC in the context of LLM and rewrite the objectives.

A1: We believe that this is a misunderstanding and would like to reiterate our contributions as follows. First, in order to overcome the sparse reward issue, instead of formulating the problem as a bandit, we consider formulating it as a MDP problem. To learn the MDP, we do not simply apply SAC but adapt it to a direct-preference-learning algorithm to eliminate the Q-function model, which can directly learn the policy. Besides, we also introduce importance sampling to tackle the distribution shift issue in offline learning and incorporate lambda-return to address the instability due to V-model initialization. Therefore, our contribution goes far beyond simply applying SAC to the context of LLM and rewriting the objectives.

Q2: The excessive number of formulas has made the text lengthy and difficult to follow.

A2: Thank you very much for your valuable suggestions! We have revised our paper according to your suggestions. For the second point, we think that the SAC framework, whose objective function is equations (6) and (7), is one of many approaches to learning the optimal $Q^*$ and $V^*$. Therefore, it might not be appropriate to place (6) and (7) in Section 2. In the meanwhile, the objective functions are derived based on (6) and (7), which means that deferring them to the appendix might cause the main text not self-contained.

Q3: Additionally, minor issues in the formulas require proofreading to correct

A3: Thank you very much for pointing out the errors! We have fixed them in our revision.

Q4: The experimental results over Qwen2-7B-Instruct model are not convincing

A4: Thank you for your suggestions. For sampling decoding, we sample 5 time with difference seeds and calculated the mean and standard deviation. The results are directly updated in the paper. Specifically, we summarize the updated results for Qwen here. Due to the version of vllm the results might various from previous.

The results show that DQO outperforms all other baselines significantly on MATH dataset and GSM8K greedy decoding, and comparable with DPO when doing sampling on GSM8K, which is consistant with our claims in the paper. Due to time limit, our results on other models are not availble now and we will show the result once it is available.

Thank you for your comment. We answer your questions as follows

Q1: I believe the method can heavily rely on the accuracy of the process value. but the difficulties should be analyzed in your experiment.

A1: This is a very good point. Indeed, thanks to our MDP formulation, our method DQO can naturally incorporate process reward and achieve even better performance. Empirically, in Section 4.4, we show that while DQO itself can already outperform all baselines, it can be further improved by integrating even the simple process rewards. Therefore, we believe that DQO can perform even better when more accurate process rewards are available. However, since the main focus of this paper is the offline reinforcement learning algorithm DQO, a detailed discussion on acquiring better process value lies out of the scope of this paper and we leave it as our future works.

Q2: There is a lack of analysis regarding whether the process value is fairly accessible or measurable.

A2: This is a good point. Actually, in this paper, we formulate the problem of LLM generation as a Markov Decision Process. Then, to solve this MDP, we adopt the framework of soft-actor-critic (SAC), which includes training a value function that can serve as the process value. Therefore, even if external process values are not available, the SAC framework itself already provides process supervision. Our experiment results, as shown in Table 3 and Table 4, consolidate our claim and show that DQO outperforms other baselines without process value and is thus still valuable when process value is inaccessible. As for the measurability of process value, we believe that we can construct a preference dataset and see whether the prediction of process value aligns with the ground-truth labels. However, this lies out of the scope of this paper and we leave this interesting problem as future work.

Q3: The motivation appears to be aligned with the process-supervised reward model approach. Could you clarify and demonstrate the key differences between your method and theirs?

A3: We believe that this is a misunderstanding. In this paper, our major contribution is an RL algorithm for LLM that does not rely on process-reward model. Specifically, we propose a DQO, which adopts the framework of soft-actor-critic and direct-preference learning, and equipped with $\lambda$-return and importance sampling to tackle the issue of inaccurate V-function initialization and offline learning. Our method indeed can further benefit from an off-the-shelf process reward models. However, as shown in Table 3 and Table 4, even without external process reward, DQO can still outperform all baseline methods.

Q4: Could you provide a more in-depth analysis of your strengths when it comes to handling long-horizon tasks?

A4: From a mathematical perspective, As we have demonstrated in Section 1, we show that previous works usually formulate generation tasks as a bandit problem, which makes it hard to capture the implicit long-horizon structures. Motivated by this, we formulate the task as a Markov Decision Process where each action is a token. We then adopt a soft-actor-critic framework, which includes training a value function that serves as a process-supervision. This process supervision addresses the challenge of sparse reward in long-horizon tasks and therefore benefits the training process. Empirically, as shown in Tables 3 and 4, DQO outperforms the baselines without external process supervision, and Table 7 further shows that DQO can be further boosted when stronger process-supervisions is available. This shows the DQO's advantage in dealing with such long-horizon tasks.

Q5: Could you expand on how experiments can contribute to evaluating the process value? And how to produce these process values to make it could scale in practice.

A5: We would like to emphasize that the primary goal of our algorithm is to output a well-performing policy. The output value is only a by-product of our approach and is not used during inference. Therefore, conducting experiments on evaluating the value model lies out of the scope of our paper and we leave it as our future work.

Thank you for your feedback! We answer your questions as follows

Q1: The paper assumes the existence of a ground-truth reward function, which undermines the necessity of the $\log \pi^*(a|s)$ representation used later in the paper.

A1: We believe that this is a misunderstanding. Actually, our setup does not require a set of known rewards $r^*$. Given any unknown reward function, we can derive the optimal $Q$-function and $V$ function as in equation (4) and (5) and therefore further derive the corresponding optimal policy $\pi^*$. Since the optimal reward $r^*$ is unknown, the optimal policy $\pi^*$ is also unknown. To estimate the optimal policy $\pi^*$, we optimize the SAC objective functions defined in (6) and (7) to learn the environment. However, if we directly follow the original framework of SAC and model the Q-function and the policy separately, at inference time it requires extra combination of $Q$-value model and the policy model which is the language model (like the method in ILQL, Offline RL for Natural Language Generation with Implicit Language Q Learning), which is inefficient. Thus, inspired by direct-preference learning approach, we directly parameterize $Q$-value function by the policy $\pi$. This approach enables us to directy obtain the optimal policy $\pi^*$ if we can learn the optimal Q-function $Q^*$.

Q2: The transition from Eq(9) to Eq(11) is questionable.

A2: We kindly remind the reviewer that equation (7) is not value function fitting target. Instead, (7) should be viewed as a function of $\theta$ and is the Q-function fitting target. In DQO, $\theta$ is actually the parameter of the policy (i.e., the language model). Therefore, optimizing the objective (7) is indeed optimizing the policy. By substituting the $Q_\theta$ in (7) we get (11). We have add interpretion in this part in our revision to avoid further misunderstanding.

Q3: There is a conflict between the definitions of reward in Line 223 and Eq (14)

A3: We believe that this is a misunderstanding. Actually, as shown in equation (5), our value function is defined as $E_{\pi}[\sum_{t=h}^H \bar{r}(s_t, a_t) - \beta \sum_{t=h}^H \log \pi(a_t|s_t)]$, which is the expection of total adjusted reward $\bar{r}$ and entropy of current policy. Furthermore, $\bar{r}(s_t, a_t)$ can be decomposed into the actual reward $r(s_t, a_t)$ and the negative entropy of the reference policy $\log \pi_{\text{ref}}(a_t|s_t)$. Therefore, in the right hand side of (14), besides $V^{\pi}(s_{h+g})$, we also need all three terms above, that is, the actual reward $r(s_{h+l}, a_{h+l})$, the negative entropy of the reference policy $\log \pi_{\text{ref}}(a_{h+l}|s_{h+l})$ and the entropy of current policy $-\log \pi_{\theta}(a_{h+l}|s_{h+l})$. We apologize for the confusion and revised the definition of soft Q-function and V-function in our revision.

Q4: The rationale behind the use of importance sampling in Section 3.3 needs further justification, particularly in the context of value function learning where it is not typically necessary.

A4: This is a very good point. Actually, when applying to offline setting, in principle it is necessary to incorporate importance. Specifically, to estimate the $V$-fucntion of current policy $V^{\pi}(s)=E_{a\sim \pi_\theta(\cdot | s))}[Q(s,a)]$, we need to take expectation over all action $a$ with probability $\pi_\theta(a|s)$. However, since the training data is sampled from the reference policy $\pi_\text{ref}$, taking expectation over $\mathcal{D}$ is essentially computing $E_{a\sim \pi_{\text{ref}}(\cdot | s))}[Q(s,a)]$, which actually equals to $V^{\pi_\text{ref}}$. Therefore, to bridge the distributional mismatch here, we have to incorporate important sampling to reweight the samples to make it align with the distribution generated from $\pi_{\theta}$. This motivates the introduction of importance sampling.