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

## Abstract

Abstract Reinforcement Learning (RL) plays a crucial role in aligning large language models (LLMs) with human preferences and improving their ability to perform complex tasks. However, current approaches either require significant computational resources due to the use of multiple models and extensive online sampling for training (e.g., PPO) or are framed as bandit problems (e.g., DPO, DRO), which often struggle with multi-step reasoning tasks, such as math problem-solving and complex reasoning that involve long chains of thought.
To overcome these limitations, we introduce Direct Q-function Optimization (DQO), which formulates the response generation process as a Markov Decision Process (MDP) and utilizes the soft actor-critic (SAC) framework to optimize a Q-function directly parameterized by the language model. The MDP formulation of DQO offers structural advantages over bandit-based methods, enabling more effective process supervision.
Experimental results on two math problem-solving datasets, GSM8K and MATH, demonstrate that DQO outperforms previous methods, establishing it as a promising offline reinforcement learning approach for aligning language models.

## 1 Introduction

Large language models (LLMs) have shown remarkable performance and potentials on a wide range of tasks including dialog generation (Han et al., ), general question answering (Alawwad et al., ), code generation (Jimenez et al., ; Chen et al., ), agent (Wang et al., ) and math problem solving (Yu et al., ; Shao et al., ).
To ensure good performance, one of the key procedures is to align the language models with human preference or task-specific requirements by reinforcement learning (RL) (Bai et al., ; Touvron et al., ). Canonically, the alignment training pipeline consists of two stages (Ouyang et al., ). In the first stage, a reward model under the Bradley-Terry model (Bradley & Terry, ) is trained on human or environment-labeled preference data. Then the language model is trained by online RL algorithms like Proximal Policy Optimization (PPO) (Schulman et al., ) with the reward signals provided by the reward model trained in stage one.

Despite the good performance achieved, the online RL methods usually involve sampling during training, which is both costly and unstable compared to offline methods (Choshen et al., ). These issues are overcome by offline preference learning methods, of which the representative is Direct Preference Optimization (DPO) (Rafailov et al., ). DPO and its follow-ups (e.g., Zhao et al. (); Azar et al. ()) treat the language model as the policy model and reward model simultaneously and train the model on offline pairwise preference data directly, therefore eliminating the need for a separate reward model. Though simple, direct preference learning has been shown effective and efficient in LLM alignment (Tunstall et al., ).

However, in practice, sometimes it is hard to acquire pairwise data required by the above methods. This issue becomes particularly severe under scenarios like math problem solving or code generation when generating high-quality data requires domain-specific expertise (Saunders et al., ; Bowman et al., ). This drawback of DPO has recently been circumvented by Direct Reward Optimization (DRO) (Richemond et al., ). DRO formulates the LLM generation task as a single-step MDP (i.e., bandit) and adopts the framework soft actor-critic (SAC) (Haarnoja et al., ), where the advantage is directly parameterized by the language model. Consequently, DRO inherits the advantage of offline policy gradient and gets rid of the dependency on pairwise data.

Nevertheless, DRO treats the process as a bandit problem, which neglects the intrinsic long-horizon nature of a wide spectrum of tasks that require complex multi-step reasoning like math problem solving and code generation (Kang et al., ; Miao et al., ), where an erroneous reasoning is almost fatal.
Previous RL research found that if rewards are only provided at the end of the episode, discovering this sparse reward signal is a hard exploration problem and sparse reward functions may not be able to meaningfully distinguish between a wide range of different policies, which makes the training inefficient (Riedmiller et al., ; Wilcox et al., ).
In the meanwhile, recent studies show that signals from process reward models (PRMs) can further boost the performance of language model (Zhang et al., ; Lightman et al., ). The positional information of PRM scores usually implies the critical mistakes in the reasoning and therefore provides stronger supervision signals.
However, if the whole generation process is simplified as a single action, the process reward will be aggregated and the positional information will be lost, implying that DRO cannot efficiently utilize process supervision.

**Table 1: A comparison between DQO, offline learning algorithms, including supervised fine-tuning (SFT), reject sampling (RS) (Dong et al., ), DPO (Rafailov et al., ), KTO (Ethayarajh et al., ), DRO (Richemond et al., ) and online algorithm PPO (Schulman et al., ). DQO enjoys all the benefits listed in the left-most column.**
|  | SFT | RS | DPO | KTO | DRO | PPO | DQO |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Free from online sampling during training | ✓ | ✓ | ✓ | ✓ | ✓ | ✗ | ✓ |
| Learn from negative samples | ✗ | ✗ | ✓ | ✓ | ✓ | ✓ | ✓ |
| Learn from unbalanced samples | ✗ | ✗ | ✗ | ✓ | ✓ | ✓ | ✓ |
| Ability to use process supervision | ✗ | ✗ | ✗ | ✗ | ✗ | ✓ | ✓ |

In order to overcome the aforementioned issues, in this paper, we propose Direct $Q$-function optimization (DQO), an offline RL algorithm for LLMs. In DQO, the responding procedure is formulated as a Markov Decision Process (MDP) and our goal is to learn an optimal policy under KL-regularization. Our algorithm adopts the framework of soft $Q$-learning, where the $Q$-function is directly parameterized by the language model. Then both the $Q$-function network and the value network are updated according to Soft Bellman Equation on offline data. The MDP formulation makes DQO a multi-step learning algorithm, and can therefore exploit process reward signals. A holistic comparison of our method and other alignment methods is shown in Table 1.
Unlike DPO or DRO, where all tokens (actions) within a response (trajectory) are evenly incentivized or punished, DQO can discover the correct reasoning steps in a incorrect response, see the case in Table 2.
Specifically, our contributions are summarized as follows

- •
We propose Direct $Q$-function optimization, or DQO, an offline RL algorithm for LLMs. Compared to previous methods, DQO can learn from offline and negative or unbalanced samples. Moreover, DQO is featured by step-wise learning, which is favorable for long-horizon tasks and able to exploit process rewards.
- •
We introduce a practical instantiation of DQO, which equips DQO with $\lambda$-return and importance sampling. These techniques stabilize the training process and ensure a good performance.
- •
We empirically compare DQO with a wide range of widely used alignment algorithms on math problem-solving tasks. Experiment results show that DQO outperforms these baselines on both GSM8K and MATH datasets. Moreover, as shown by our experiment, when process rewards are available, the performance of DQO can be further boosted, indicating that DQO can benefit from process rewards.

**Table 2: A case study of learned $V$-function. We select a prompt-response pair from the GSM8K training set and split the response into nine segments by double line breaks. We extract the $V$-function value at each line break. The model makes a mistake at the beginning of the fifth segment (where $$x=y$$ is erroneous). The model predicts high $V$-values from the first segment to the fourth segment. After that, the predictions drop significantly. The value trends preciously reflect the correctness of the reasoning and thus can serve as an informative process supervision.**
|  | Text | Process Value | Correctness |
| --- | --- | --- | --- |
| Prompt | Jorge has an equal number of baseball cards as Matias, who has 6 fewer cards than Carlos.<br>If Carlos has 20 baseball cards, what is the total number of baseball cards the three have?<br>Please reason step by step, and produce a final answer following 4 ’#’, like ’#### 0’. | - | - |
| Response | **Step 1: Identify the information provided** | 0.6016 | ✓ |
|  | - Jorge has an equal number of baseball cards as Matias.<br>- Matias has 6 fewer cards than Carlos.<br>- Carlos has 20 baseball cards. | 0.6719 | ✓ |
|  | **Step 2: Translate the information into equations** | 0.7969 | ✓ |
|  | Let x be the number of baseball cards Jorge and Matias each have.<br>Let y be the number of baseball cards Carlos has. | 0.6680 | ✓ |
|  | $$x = y$$ $$y = x + 6$$ $$y = 20$$ | 0.1016 | ✗ |
|  | **Step 3: Substitute the equations to find x** | 0.0013 | ✗ |
|  | $$x = x + 6$$ $$0 = 6$$ This is a contradiction, so the given information is not consistent. | 0.4688 | ✗ |
|  | **Therefore, there is no solution to the problem as the given constraints are contradictory.** | 0.1768 | ✗ |
|  | #### 0 $\langle$eos$\rangle$ | - | ✗ |

## 2 Preliminaries

In this section, we introduce the foundational concepts and notations that underpin our proposed algorithm. We first review the basic framework of modeling language generation as a reinforcement learning task, followed by a KL-regularized reinforcement learning objective.

Modeling Language Generation as Token-Level MDP
Reinforcement Learning (RL) is concerned with learning a policy that maximizes the cumulative reward for an agent interacting with an environment.
In this work, we formalize language generation tasks as a Markov decision process (MDP).
We denote prompt as $x$ and a response to the prompt as $y$, which can each individually be broken down into a sequence of tokens, for example, $x=(x_{0},\dots,x_{m})$, from a fixed discrete vocabulary $\mathcal{A}$.
We define the token-level MDP as a tuple $\mathcal{M}=(\mathcal{S},\mathcal{A},\mathbb{P},r,d_{0},\omega)$. In the defined MDP, $\mathcal{S}$ is the space of the state consisting of all tokens generated so far, i.e., $s_{t}=(x_{0},\dots,x_{m},y_{0},\dots,y_{t})$. The action space $\mathcal{A}$ is the fixed discrete vocabulary. The dynamics $\mathbb{P}$ are the deterministic transition model between tokens, i.e., $\mathbb{P}(s_{t+1}|s_{t},a)=1$ for $s_{t}=(x_{0},\dots,x_{m},y_{0},\dots,y_{t})$, $a=y_{t+1}$ and $s_{t+1}=(x_{0},\dots,x_{m},y_{0},\dots,y_{t},y_{t+1})$(^1^11For notational simplicity, we ignore the case that LLM can call an external tool. If the tool does not introduce randomness, the state transitions are also deterministic. Even if the state transitions are random, we can use samples to approximate the state transition probabilities.). The generation process will terminate once the terminal action $\omega$ (usually end-of-sentence token) is taken.
The reward function $r(s,a)$ provides scalar feedback for the agent’s performance after taking action $a$ in state $s$. In RLHF, the reward function is usually learned from human feedback over preferences or given by a series of rules depending on the specific tasks. The initial state distribution $d_{0}$ is a distribution over prompts $x$, where an initial state $s_{0}$ is comprised of the tokens from $x$.

KL-Regularized Reinforcement Learning Objective We formulate the optimization objective as a KL-regularized RL problem. Our goal is to approximate the optimal KL-regularized policy

$$ $\displaystyle\pi^{*}=\mathop{\mathrm{argmax}}_{\pi}\mathbb{E}_{\pi,s_{0}\sim d _{0}}\left[\sum_{h=0}^{H}\left(r(s_{h},a_{h})-\beta\text{KL}\big{(}\pi(\cdot|s _{h})\|\pi_{\text{ref}}(\cdot|s_{h})\right)\big{)}\right],$ (2.1) $$

where $H$ is the total number of decision steps, $s_{0}$ is a prompt sampled from the dataset, $r(s_{h},a_{h})$ is the token-level reward from the reward function, $\beta$ is the coefficient controlling the magnitude of KL-regularization and $\pi_{\text{ref}}$ is the initialisation policy.
In classic RLHF and most LLM-related tasks, the reward is sparse and is only applied at the terminal action $\omega$, i.e. the end-of-sentence token <eos>. However, our structure is flexible enough to incorporate both dense and sparse rewards from ruled-based reward models, turn-level reward models, process-supervised reward models (PRM), or just outcome-supervised reward models.

We consider rewriting our objective function (2.1) under the framework of max-entropy reinforcement learning. Specifically, we decompose the KL-regularization term $\mathrm{KL}(\pi(\cdot|s_{h})\|\pi_{\text{ref}}(\cdot|s_{h}))$ into cross-entropy and entropy, leading to the following objective:

$$ $\displaystyle\pi^{*}$ $\displaystyle=\mathop{\mathrm{argmax}}_{\pi}\mathbb{E}_{\pi,s_{0}\sim d_{0}} \left[\sum_{h=0}^{H}\left(r(s_{h},a_{h})+\log\pi_{\text{ref}}(a_{h}|s_{h})+ \beta\mathcal{H}(\pi(\cdot|s_{h}))\right)\right]$ $\displaystyle=\mathop{\mathrm{argmax}}_{\pi}\mathbb{E}_{\pi,s_{0}\sim d_{0}} \left[\sum_{h=0}^{H}\left(\overline{r}(s_{h},a_{h})+\beta\mathcal{H}(\pi(\cdot |s_{h}))\right)\right],$ (2.2) $$

where $\mathcal{H}(\pi(\cdot|s_{h}))=-\mathbb{E}_{a\sim\pi}\log\pi(a|s_{h})$ denotes the entropy of the policy at state $s_{h}$ and the KL-regularized reward $\overline{r}$ is defined as
$\overline{r}(s_{h},a_{h})=\beta\log\pi_{\text{ref}}(a_{h}|s_{h})+r(s_{h},a_{h})$.
Equation 2.2 leads to a maximum entropy reinforcement learning problem, which enjoys the well-known closed-form solution (Haarnoja et al., ) as follows:

$$ $\displaystyle\pi^{*}(a|s_{h})=\exp{\left(\frac{Q^{*}(s_{h},a)-V^{*}(s_{h})}{ \beta}\right)},$ (2.3) $$

where the Soft $Q$-function is defined as

$$ $\displaystyle Q^{*}(s_{h},a)=\overline{r}(s_{h},a)+\mathbb{E}_{\pi^{*}}\left[ \sum_{j=h+1}^{H}\left(\overline{r}(s_{j},a_{j})+\beta\mathcal{H}(\pi^{*}(\cdot |s_{j}))\right)\right],$ (2.4) $$

and the Soft $V$-function is given by:

$$ $\displaystyle V^{*}(s_{h})=\beta\log\sum_{a}\exp{\left(\frac{Q^{*}(s_{h},a)}{ \beta}\right)}=\mathbb{E}_{\pi^{*}}\left[\sum_{t=h}^{H}\left(\overline{r}(s_{t },a_{t})+\beta\mathcal{H}(\pi^{*}(\cdot|s_{t}))\right)\right].$ (2.5) $$

Equation (2.3) reveals that the optimal policy $\pi^{*}$, soft $Q$-function $Q^{*}$, and soft V-function $V^{*}$ are interdependent, which means that knowing any two of them allows us to compute the third one.

## 3 Direct Q 𝑄 Q italic_Q -function Optimization (DQO)

### 3.1 The DQO objective

We adopt the Soft Actor-Critic (SAC) learning framework to learn the state value function $V$ and state-action value function $Q$. In SAC, the $Q$-function and $V$-function, which are parameterized by $\theta$ and $\phi$ respectively, are updated by minimizing the following squared residuals:

$$ $\displaystyle L_{V}(\phi)$ $\displaystyle=\mathbb{E}_{s_{h}\sim\mathcal{D}}\bigg{[}\bigg{(}V_{\phi}(s_{h}) -\mathbb{E}_{a_{t}\sim\pi}\Big{[}Q_{\theta}(s_{h},a_{h})+\beta\mathcal{H}\big{ (}\pi(\cdot|s_{h})\big{)}\Big{]}\bigg{)}^{2}\bigg{]},$ (3.1) $\displaystyle L_{Q}(\theta)$ $\displaystyle=\mathbb{E}_{(s_{h},a_{h})\sim\mathcal{D}}\left[\left(Q_{\theta}( s_{h},a_{h})-\overline{r}(s_{h},a_{h})-V_{\phi}(s_{h+1})\right)^{2}\right],$ (3.2) $$

where $\mathcal{D}$ is the distribution of previously sampled states and
actions, or a replay buffer.
As shown in (2.3), the optimal policy $\pi^{*}$, optimal $Q$-function $Q^{*}$, and optimal value function $V^{*}$ are tightly interconnected. Specifically, they satisfy the following relationship

$$ $\displaystyle Q^{*}(s_{h},a_{h})=\beta\log\pi^{*}(a_{h}|s_{h})+V^{*}(s_{h}).$ (3.3) $$

Inspired by this, we parameterize the $Q$-value-network with the policy as follows

$$ $\displaystyle Q_{\theta}(s_{h},a_{h})=\beta\log\pi_{\theta}(a_{h}|s_{h})+V_{ \phi}(s_{h}),$ (3.4) $$

where $\pi_{\theta}(\cdot|\cdot)$ is the policy network, or the language model in this paper. In equation (3.4), instead of using an additional model to parameterize $Q$-value and learning the optimal policy from the optimal $Q$-function $Q^{*}$, we directly infer the policy from the $Q$-function by parameterizing it with $\pi$. Therefore, we name our algorithm as Direct $Q$-function optimization (DQO).

By plugging in (3.4) to 3.2, we can rewrite the loss function for the policy as:

$$ $\displaystyle L_{\pi}(\theta)=\mathbb{E}_{(s_{h},a_{h})\sim\mathcal{D}}\left[ \overline{r}(s_{h},a_{h})+V_{\phi}(s_{h+1})-V_{\phi}(s_{h})-\beta\log\pi_{ \theta}(a_{h}|s_{h})\right]^{2}.$ (3.5) $$

Substituting $\overline{r}(s_{h},a_{h})$ by its definition $\beta\log\pi_{\text{ref}}(a_{h}|s_{h})+r(s_{h},a_{h})$, we obtain the objective function of the policy network as follows:

$$ $\displaystyle L_{\pi}(\theta)=\mathbb{E}_{(s_{h},a_{h})\sim\mathcal{D}}\bigg{[ }\bigg{(}\beta\log\frac{\pi_{\theta}(a_{h}|s_{h})}{\pi_{\text{ref}}(a_{h}|s_{h })}-\Big{(}r(s_{h},a_{h})+V_{\phi}(s_{h+1})-V_{\phi}(s_{h})\Big{)}\bigg{)}^{2} \bigg{]}.$ (3.6) $$

To eliminate the $Q$-function in the $V$-function objective (3.1), we consider the soft Bellman equation:

$$ $\displaystyle Q_{\theta}(s_{h},a_{h})=\overline{r}_{h}(s_{h},a_{h})+\mathbb{E} _{s_{h+1}\sim\mathbb{P}(\cdot|s_{h},a_{h})}[V^{\pi}(s_{h+1})].$ (3.7) $$

When $\mathcal{D}$ is generated online by the current policy $\pi_{\theta}$, we can estimate the expectation in state transitions by sampling and therefore substitute $\mathbb{E}_{s_{h+1}\sim\mathbb{P}(\cdot|s_{h},a_{h})}[V^{\pi}(s_{h+1})]$ with $V_{\phi}(s_{h+1})$. Consequently, we obtain the loss for the value function $V$ in the following form:

$$ $\displaystyle L_{V}(\phi)=\mathbb{E}_{(s_{h},a_{h})\sim\mathcal{D}}\Big{[}\big {(}\overline{r}(s_{h},a_{h})+V_{\phi}(s_{h+1})+\beta\mathcal{H}(\pi_{\theta}( \cdot|s_{h}))-V_{\phi}(s_{h})\big{)}^{2}\Big{]}.$ (3.8) $$

When $\mathcal{D}$ is composed of pre-generated offline data, we employ importance sampling to reweight the offline data, ensuring that the offline dataset can be used effectively. We defer the detailed discussion to Section 3.3.
It is worth highlighting that in our formulation of DQO, we consider generating each single token as an action. If we consider generating the whole utterance as a single action and set the horizon length $H=1$, then equation (3.6) and equation (3.8) degenerates to the loss used by DQO (Richemond et al., ). This means that DRO can be viewed as a special case of the learning framework of DQO.

### 3.2 Mitigating Bias with λ 𝜆 \lambda italic_λ -Return

One-step temporal difference (TD) errors have high bias and perform poorly when the value function is not well-initialized, resulting in inefficient learning. To address this, we incorporate $\lambda$-return (Schulman et al., ) to improve the updates for $Q$-function and $V$-function. By definition, we know that $V(s_{h})$ is the sum of reward gained by next $n$ action and $V(s_{h+n})$, or formally,

$$ $\displaystyle V^{\pi}(s_{h})=\mathbb{E}_{\pi}\left[V^{\pi}(s_{h+n})+\sum_{l=0} ^{n-1}\left(\beta\log\frac{\pi_{\text{ref}}(a_{h+l}|s_{h+l})}{\pi(a_{h+l}|s_{h +l})}+r(a_{h},s_{h})\right)\right].$ (3.9) $$

Given a trajectory $\{s_{0},a_{0},r_{0},\cdots,s_{H},a_{H},r_{H}\}$ where $r_{h}=r(a_{h},s_{h})$, we use the empirical sample to estimate the n-step return and define the empirical $n$-step return as:

$$ $\displaystyle G_{\phi,\theta}^{(n)}(s_{h})=V_{\phi}^{\pi}(s_{h+n})+\sum_{l=0}^ {n-1}\left(\beta\log\frac{\pi_{\text{ref}}(a_{h+l}|s_{h+l})}{\pi_{\theta}(a_{h +l}|s_{h+l})}+r_{h+l}\right).$ (3.10) $$

Now we are able to define $\lambda$-return, which is the weighted average of all $n$-step returns:

$$ $\displaystyle G_{\phi,\theta}^{\lambda}(s_{h})=\begin{cases}(1-\lambda)\sum_{n =1}^{H-h}\lambda^{n-1}G_{\phi,\theta}^{(n)}(s_{h}),&\text{if }\lambda<1\\ G_{\phi,\theta}^{(H-h)}(s_{H}),&\text{if }\lambda=1\end{cases}.$ (3.11) $$

We replace the target for value updates from one-step return $V_{\phi}(s_{h+1})$ to the $\lambda$-return $G_{\overline{\phi},\overline{\theta}}^{\lambda}(s_{h})$, where $\overline{\phi}$ is the copy of $\phi$ but does not counted into the back-propagation gradients. Now we have the loss function for the value network as follows:

$$ $\displaystyle L_{V}(\phi)=\mathbb{E}_{s_{h}\in\mathcal{D}}\left[(G_{\overline{ \phi},\overline{\theta}}^{\lambda}(s_{h})-V_{\phi}(s_{h}))^{2}\right].$ (3.12) $$

Similarly, the loss for $Q$-function ($\pi$) using $\lambda$-return is:

$$ $\displaystyle L_{Q}(\theta)=\mathbb{E}_{(s_{h},a_{h})\sim\mathcal{D}}\left[ \left(\beta\log\frac{\pi_{\theta}(a_{h}|s_{h})}{\pi_{\text{ref}}(a_{h}|s_{h})} -\left(r(s_{h},a_{h})+G_{\overline{\phi},\overline{\theta}}^{\lambda}(s_{h+1}) -V_{\phi}(s_{h})\right)\right)^{2}\right].$ (3.13) $$

### 3.3 Reweighting Offline Data with Importance Sampling

Offline RL, also known as batch RL, focuses on learning a policy from a pre-collected, fixed dataset without further interaction with the environment. A key challenge in offline RL is the distributional shift between the behavior policy $\mu$, which generated the data, and the target policy $\pi$.
In order to mitigate this mismatch, offline RL algorithms often incorporate regularization techniques, such as importance sampling or constraints on the learned policy, to remain close to the behavior policy and avoid overestimating the Q-values for out-of-distribution actions. In this work, we first use a KL-constrained RL objective and then employ importance sampling to reweight the offline data to match the distribution of trajectories generated by the current optimized policy.

Here, we introduce importance sampling to help correct this mismatch, ensuring that the offline dataset can be used effectively for policy updates. Let $\mu$ represent the behavior policy under which the offline data $\mathcal{D}$ was generated and $\pi$ be the current online policy.
The probability of a trajectory $\tau$ under $\mu$ and $\pi$ are computed as follows:

$$ $\displaystyle\mu(\tau|s_{1})=\prod_{h=0}^{H}\mu(a_{h}|s_{h}),\quad\pi(\tau|s_{ 1})=\prod_{h=0}^{H}\pi(a_{h}|s_{h}).$ $$

Therefore, we know that when the offline dataset $\mathcal{D}$ is sampled from $\tau$, we have

$$ $\displaystyle\mathbb{E}_{\tau\sim\pi}[f(\tau)]=\mathbb{E}_{\tau\sim\mathcal{D} }\bigg{[}\frac{\pi(\tau|s_{h})}{\mu(\tau|s_{h})}f(\tau)\bigg{]}=\mathbb{E}_{ \tau\sim\mathcal{D}}\bigg{[}\frac{\prod_{h=0}^{H}\pi(a_{h}|s_{h})}{\prod_{h=0} ^{H}\mu(a_{h}|s_{h})}f(\tau)\bigg{]},$ $$

where $f(\tau)$ is any function of trajectory $\tau$. This indicates that we can use the importance ratio $\pi(\tau|s_{h})/\mu(\tau|s_{h})$ to adjust the loss. Empirically, we truncate the importance sampling rate to avoid gradient explosion caused by extreme values. The final ratio we apply is shown as follows

$$ $\displaystyle w(\tau)=\min\left(e,\prod_{h=1}^{H}\frac{\pi(a_{h}|s_{h})}{\mu(a _{h}|s_{h})}\right)=\exp\left(\min\left(1,\sum_{h=1}^{H}\log\frac{\pi(a_{h}|s_ {h})}{\mu(a_{h}|s_{h})}\right)\right),$ (3.14) $$

Now we plug in the importance ratio in (3.14) to the loss functions (3.12) and (3.13) and then obtain our final loss functions for offline learning.

$$ $\displaystyle L_{V}(\phi)$ $\displaystyle=\mathbb{E}_{\tau\sim\mathcal{D}}\left[w(\tau)\cdot\sum_{h=1}^{H} (G_{\overline{\phi},\overline{\theta}}^{\lambda}(s_{h})-V_{\phi}(s_{h}))^{2} \right],$ $\displaystyle L_{\pi}(\theta)$ $\displaystyle=\mathbb{E}_{\tau\sim\mathcal{D}}\left[w(\tau)\cdot\sum_{h=1}^{H} \left(\beta\log\frac{\pi_{\theta}(a_{h}|s_{h})}{\pi_{\text{ref}}(a_{h}|s_{h})} -\left(r(s_{h},a_{h})+G_{\overline{\phi},\overline{\theta}}^{\lambda}(s_{h+1}) -V_{\phi}(s_{h})\right)\right)^{2}\right].$ $$

What notable is that the importance sampling weight $w(\tau)$ is not involved in the gradient backward computation. The introduction of importance ration enables us to leverage offline datasets in an online RL framework, ensuring that the updated policy remains consistent with the distribution of trajectories it would encounter during online interaction.

## 4 Experiments

In this section, we conduct extensive experiments to demonstrate the effectiveness of our proposed method. Moreover, we show that our method can be further augmented by utilizing process rewards.

### 4.1 Datasets

We evaluate the models using two widely established mathematical problem-solving datasets: MATH (Hendrycks et al., ) and GSM8K (Cobbe et al., ). The MATH dataset consists of 5,000 challenging problems in its test set and 7500 in its train set. The problems cover various fields like algebra, geometry, probability, and calculus. The GSM8K dataset consists of a train set of 7473 problems and a test set of 1319 problems. The problems are mostly simpler grade-school math word problems. The problems in both datasets usually require multi-step reasoning and complex arithmetic to derive the correct answers.

We use the 7.5K training problems from the MATH dataset and 7.47K training problems from the GSM8K dataset to generate the training corpus of our baselines and DQO.
We use our base model and sample 20 responses for each prompt in the training set and then label all these responses as positive and negative responses. The detailed usage of these samples is discussed in the next section and please refer to the Appendix A for a more detailed discussion of our dataset construction.

### 4.2 Models, Baselines and Evaluation

In our experiments, we select two pretrained models, Gemma-1.1-it-7B(^2^22https://huggingface.co/google/gemma-1.1-7b-it) (Gemma) (Team et al., ) and Qwen2-7B-Instruct(^3^33https://huggingface.co/Qwen/Qwen2-7B-Instruct) (Qwen) (Yang et al., ) as our base model. We implement our method based on HybridFlow (Pezeshkpour et al., ). The baselines and our methods are implemented as follows and we defer detailed hyperparameters to Appendix B

- •
SFT and reject sampling: We use the canonical responses provided in the training set as the target for our SFT and all the collected positive responses as the target for reject sampling. We train the model for 3 epochs.
- •
DPO: For each prompt, we pair up all generated positive and negative responses. Due to a large number of pairs which makes the training inefficient, we sample from these pairs to make our training set size moderate. We train the model for 1 epochs.
- •
KTO, DRO and DQO: We use all the generated responses for training. We assign reward 1 to those correct responses and reward 0 to those incorrect response. For KTO, we train the models for 1 epochs. For DRO and DQO, We select the best checkpoint according to the training curves for evaluation.

We evaluate the models generated by our method and baseline methods on the test set of GSM8K and MATH. We consider two inference strategies, greedy generation and sampling. We set the sampling parameters the same as when we generated the training corpus.

### 4.3 Empirical Results and Ablation Studies

**Table 3: Experiment results for Qwen2-7B-Instruct model. We use bold for the best and underline for the second best. DQO significantly improves the base model’s performance. This improvement surpass all the baselines when doing sampling. As for greedy generation, DQO is comparable to DPO when doing greedy generation and surpass all other baselines.**
| Data | GSM8K | MATH |  |  |
| --- | --- | --- | --- | --- |
| Generation Method | Greedy | Sample | Greedy | Sample |
| Qwen2-7B-Instruct | 59.06 | 53.30 | 37.44 | 35.42 |
| SFT | 85.06 | 84.15 | 44.56 | 41.34 |
| Reject Sampling | 84.15 | 83.47 | 49.82 | 48.02 |
| DPO | 87.26 | 84.23 | 51.66 | 49.00 |
| KTO | 86.35 | 83.55 | 50.32 | 46.30 |
| DRO | 86.73 | 83.39 | 51.18 | 48.40 |
| DQO | 87.26 | 84.69 | 51.72 | 50.18 |

Here we show our main results on both the GSM8K and MATH datasets. The results are shown in Table 3 and Table 4 respectively. From Table 3, we see that all the methods improve the performance of the base models by a significant margin. Particularly, on GSM8K, DQO improves the performance from 59.06% to 87.26% for greedy generation and 53.30% to 84.69% for sampling. This improvement is comparable with DPO and surpasses DRO and other baselines by a margin of 1.30% for sampling and 0.53% for greedy generation.
On MATH, we also see a significant performance improvement from DQO. The performance of DQO, while comparable with DPO and DRO, surpasses all other baselines by a margin of 1.40%. As for sampling, DQO reaches a performance of 50.18%, which surpasses the performance of the best baseline method DPO by a margin of 1.18%. These results indicate that DQO achieves a comparable performance of DPO and surpasses other baselines by a considerable margin.

**Table 4: Experiment results for Gemma-1.1-7B-it model. We use bold for the best performance and underline for the second best performance. DQO significantly improves the base model’s performance. On GSM8K, DQO surpasses all other baselines by a significant margin. On MATH dataset, DQO achieves a comparable performance with DRO when doing greedy sampling and outperforms all the baseline when doing sampling at inference.**
| Data | GSM8K | MATH |  |  |
| --- | --- | --- | --- | --- |
| Generation Method | Greedy | Sample | Greedy | Sample |
| Gemma-1.1-7B-it | 39.65 | 38.89 | 17.04 | 16.06 |
| SFT | 53.45 | 45.46 | 21.64 | 18.00 |
| Reject Sampling | 53.60 | 53.37 | 21.74 | 20.42 |
| DPO | 57.85 | 59.41 | 22.62 | 22.66 |
| KTO | 50.49 | 50.19 | 18.56 | 18.56 |
| DRO | 60.58 | 60.05 | 24.56 | 23.70 |
| DQO | 63.91 | 64.06 | 24.90 | 24.68 |

When it turns to the results on Gemma, we see that DQO enjoys larger advantages. As demonstrated in Table 4, we see that all considered methods result in significant improvement. Specifically, on GSM8K, DQO improves the base model’s performance by a margin of 24.36% for greedy generation and 25.17% for sampling. This significantly surpasses the improvement obtained by DRO by margins of 3.33% and 4.01% for the two generation strategies respectively, and even more compared to other baseline methods. On the MATH dataset, we see that DQO also improves the model’s performance by a prominent margin of 7.86% and 8.62% for greedy generation and sampling, respectively. This improvement is comparable with DRO and surpasses DPO and other baseline methods by a margin of at least 2.28% for greedy generation and 2.02% for sampling. In summary, DQO results in promising improvement over the base models under all the scenarios and outperforms all our baseline methods.

#### 4.3.1 Importance sampling

**Table 5: The impact of importance sampling rate on both $Q$-function loss and $V$-function loss. The experiments are conducted on Gemma. When training without an importance sampling ratio on $Q$-function loss, the performances degenerate significantly on both GSM8K and MATH. When keeping the importance ratio only on $Q$-function loss, there is also a moderate performance loss on MATH.**
| Data | GSM8K | MATH |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Generation Method | Greedy | Sample | Greedy | Sample |  |
| $Q$-loss w/o IS | $V$-loss w/o IS | 58.68 | 60.20 | 21.96 | 22.68 |
| $V$-loss w/ IS | 56.03 | 56.48 | 20.82 | 20.94 |  |
| $Q$-loss w/ IS | $V$-loss w/o IS | 63.53 | 64.06 | 22.28 | 23.18 |
| $V$-loss w/ IS | 63.91 | 64.06 | 24.90 | 24.68 |  |

To demonstrate the impact of the importance sampling ratio in DQO, we train DQO on Gemma without the importance sampling ratio for $Q$-function loss, $V$-function loss, and both. We present the results in Table 5. The results show that, without adding importance sampling, the performance will be significantly deteriorated. Specifically, on the GSM8K dataset, when importance sampling is not introduced to $Q$-function loss, the performance degenerates by a margin over 3.86%. Similarly, on the MATH dataset, we see that when we exclude the importance sampling ratio from $Q$-function loss, the performance decreases by a margin of about 2%. When we keep the importance sampling ratio only on $Q$-function loss, the performance on GSM8K almost maintains but we still see a moderate performance loss on MATH. These results show that the importance sampling ratio, on both $Q$-function and $V$-function loss, plays important roles in DQO training.

#### 4.3.2 λ 𝜆 \lambda italic_λ -return

**Table 6: The impact of $\lambda$-return on Gemma. When decreasing $\lambda$ from 1.0 to 0.95, we observe a significant performance dropping more than 3.71% on GSM8K and 2.30% on MATH.**
| Data | GSM8K | MATH |  |  |
| --- | --- | --- | --- | --- |
| Generation Method | Greedy | Sample | Greedy | Sample |
| $\lambda=0.95$ | 60.20 | 59.21 | 22.60 | 22.26 |
| $\lambda=1.0$ | 63.91 | 64.06 | 24.90 | 24.68 |

In order to demonstrate the impact of $\lambda$-return, we vary the value of $\lambda$ and evaluate the training results on Gemma. Empirically, we find that the best performance is obtained at $\lambda=1$ and quickly degenerates when decreasing $\lambda$. Therefore we pick $\lambda=0.95$ to make the comparison. The results are shown in Table 6. When switching $\lambda$ to 0.95, we observe that the performance on GSM8K decreases by a margin of more than 3.71% for greedy generation and almost 5% for sampling. The results on MATH demonstrate a similar pattern and the performances of $\lambda=0.95$ dropped by a margin of 2.3% on both inference strategies. The results indicate that $\lambda$-return is the key component in the target of policy training.

Figure: Figure 1: A visual demonstration of our process reward construction. We split all the responses to segments. For correct responses we assign all process reward to one. For negative responses, we start from each prefix and generate 20 samples. We then find the longest prefix where the best of 20 samples is correct and assign all the process rewards before to 1.
Refer to caption: x1.png

### 4.4 DQO with Process Score

In this section, we show that when process scores at intermediate steps are available, the performance of DQO can be further improved. Here, we use synthetic process scores. In order to obtain a synthetic process score, we consider using an empirical passing rate to estimate the quality of a given response prefix. Specifically, given a prompt string $x$, for each failed response $y$, we first split the response into several segments $y[0:n]$, where $n$ is the number of segments and we use $y[0:i]$ to denote the concatenation of first $i$ segments. Beginning from $i=n-1$, we randomly sample 20 trajectories given prefix $\mathrm{contat}(x,y[0:i])$. If there is at least one correct completion, we assume that the reasoning process in $y[0:i]$ is correct and all the process rewards for the previous step will be set to $1/n$. We combine these process reward scores with the original rewards. The process is summarized in Figure 1.

**Table 7: Experiment results for DQO augemented by process scores. With process rewards, the performance of DQO is further enhanced by 1.13% on GSM8K and 0.64% on MATH. The performance when doing sampling also increase on GSM8K and maintains almost the same on MATH.**
| Method | GSM8K | MATH |  |  |
| --- | --- | --- | --- | --- |
| Generation Method | Greedy | Sample | Greedy | Sample |
| Gemma-1.1-7b-it | 39.65 | 38.89 | 17.04 | 16.06 |
| DQO | 63.91 | 64.06 | 24.90 | 24.68 |
| DQO w/ process scores | 65.04 | 65.35 | 25.54 | 25.10 |

We conduct the experiments on Gemma, and the results are summarized in Table 7. Equipped with our estimated process scores, we see a further improvement. Specifically, on GSM8K, using our process scores further increases the performance by 1.13% for greedy generation and 1.29% for sampling. On MATH, process scores also boost the model’s performance by a further 0.64% when doing greedy generation and 0.42% for sampling. The results imply that DQO can be further enhanced by utilizing process scores. Additionally, we discovered that even without process scores, DQO is capable of identifying correct reasoning steps. The trained values demonstrate a growing correlation with the synthetic process scores, see in Figure 2, suggesting that DQO effectively learns to recognize correct reasoning steps, enhancing the reasoning process over time.

Figure: Figure 2: The correlation between the trained values (trained without process scores) and our constructed process scores. As the training proceeds, we observe a growth in the correlation between these two values, indicating that the value model in DQO learns the correctness of the reasoning process.
Refer to caption: x2.png

## 5 Related Work

Reinforcement Learning for Language Model Alignment Aligning language models with human preferences, or reinforcement learning with human feedback (RLHF), dates back to the work of Wirth et al. () and Christiano et al. (). It has been widely applied to a bunch of recent models including GPT-4 (Achiam et al., ), Gemini (Team et al., ), and Llama (Touvron et al., ), leading to the surprising performance of these models. The alignment procedure usually takes place after supervised finetuning (SFT). In the canonical approaches of RLHF (Ouyang et al., ; Bai et al., ; Munos et al., ), a reward model is first trained with preference data and then the model is updated with Proximal Policy Optimization (PPO). Another line of works, initiating from Direct Preference Optimization (DPO) (Rafailov et al., ), include SLiC (Zhao et al., ), IPO (Azar et al., ), KTO (Ethayarajh et al., ) and so on. These approaches are featured by directly parameterizing the reward models with the language model and then training on offline preference data. Following DPO, one branch of works, including GSHF (Xiong et al., ), SPPO (Wu et al., ) and INPO (Zhang et al., ), adapts DPO or its variant to online samples and iterative training and resulted to state-of-the-art models.
On the other hand, Richemond et al. () adapted offline reinforcement learning algorithm to direct preference learning and proposed Direct Reward Optimization (DRO), which combined offline policy learning with a value function learning and updated policy network and value network iteratively. Our work has a similar structure to DRO, but models the language generation as an MDP rather than a bandit and can utilize process supervision to facilitate training.

Multi-step and Long Horizon RL for LLM alignment Many tasks for LLMs require LLMs to reason step by step or interact with the environment turn by turn. However, the rewards are usually sparse since they are only provided at the end of a long horizon of reasoning or interactions. In traditional RL literature, one approach towards breaking the curse of lone horizon and sparse reward is to train or estimate an intermediate value function or process reward (Park et al., ) and use the process reward to guided searching (Torne et al., ; Zhang et al., ) and RL training. The utilization of process reward has also led to better performance for LLM reasoning (Zhang et al., ; Lightman et al., ).
Most straightforwardly, Snell et al. () proposed ILQL, which employed implicit Q-learning to train a Q-function network and V-function network. Then, at inference time, ILQL uses learned value functions to perturb the log probabilities of the initial policy towards utility-maximizing behavior.
The success of direct preference learning also stimulates a lot of work adapting DPO to multi-turn scenarios. To estimate process reward and utilize the information provided,  Chen et al. (); Lai et al. (); Xie et al. () leverages process reward signals or AI feedback to construct preference pairs for intermediate steps and update the model with original DPO. On the other hand, Xiong et al. (); Shani et al. () extends the vanilla DPO to accommodate the multi-turn structure. However, these approaches require pairwise data, which might not be available or easy to obtain on some specific occasions. Our work, while following the approach of direct preference learning, eliminates the need for pairwise data and can be boosted by process rewards.

## 6 Conclusion

In this work, we propose DQO, an offline reinforcement learning algorithm for enhancing the language model’s ability in multi-step reasoning. Compared to previous online methods like PPO, the offline nature of DQO bypasses the requirement of an extra reward model and online sampling during training. Previous offline methods usually formulate the LLMs’ responding process as a bandit problem, which usually fails to capture the implicit long-horizon and multi-step nature of those tasks requiring a long chain of thought. In contrast, DQO frames the tasks as a Markov decision process and further employs a soft actor-critic framework to learn the $V$-function and the $Q$-function, which is directly parameterized by the language model. To verify the effectiveness of DQO, we conduct extensive experiments on two math-problem-solving datasets, GSM8K and MATH, and empirical results show that DQO outperforms all our baseline. Currently, our experiment results are limited to two base models due to time constraints and we leave it as future work.

## Appendix A Dataset construction

In this section, we provide details about the generating process of our training data. For all problems, we add format guide to make to language models’ generation follows the same format as the solution provided by the dataset. Specifically, we use
the following templates for GSM8K and MATH respectively

- •
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{}.

The prompt is further wrapped by chat-ML template. We then sample 20 responses for each prompts with sampling parameter top_p=0.9, top_k=16 and threshold=0.01. For each response, we use regular expression to unwrap the answer and match it with the reference to generate the label. For reject sampling, we collect all correct response as the training target. For DPO, we first pairs up all positive and negative responses and then randomly sample at most 12 pairs from all possible pair to control the size of the training set. We summarize the dataset size for each baselines in Table 8.

**Table 8: The size of datasets**
| Model | Dataset | SFT | RS | DPO | KTO/DRO/DQO |
| --- | --- | --- | --- | --- | --- |
|  | MATH | 7500 | 71776 | 65568 | 150000 |
| Qwen | GSM8K | 7473 | 94889 | 87996 | 149460 |
|  | MATH | 7500 | 22600 | 35184 | 150000 |
| Gemma | GSM8K | 7473 | 46062 | 69540 | 149460 |

## Appendix B Detailed Hyperparameters for our baselines

We provide the detailed hyperparameters of the baselines and DQO in this section. For SFT and reject sampling, we select the best learning rate from {2e-5, 1e-5, 5e-6, 1e-6} and the best epoch from {1,2,3}. For SFT, the final learning rate is set to 2e-5 for Qwen and 5e-6 for Gemma. For reject sampling, the final learning rate is set to 2e-5 for Qwen and 1e-6 for Gemma. Both SFT and reject sampling are trained for 3 epochs. For DPO, we tried $\beta$ from {0.1, 0.01} and learning rate from {5e-7, 1e-7, 5e-8} and select the best. Finally, for both Qwen and Gemma $\beta$ is set to 0.1 and learning rate is set to 5e-8 and we evaluate the checkpoint after 1 epoch. We adapt all the hyperparameters of DPO to KTO. For all the baselines above, we set global batch size to 64. As for DRO, We tried KL-regularization coefficient in $\{0.01,0.03,0.1,0.3,1\}$ and learning rate to 5e-7 for Qwen and 1e-7 for Gemma. For our method DQO, we set the KL-regularization to $0.03$ and learning rate to 5e-7 for Qwen and 1e-7 for Gemma. We select the best checkpoints on the training curve for evaluation.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. 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. In particular, In Lines 149-151, the authors presuppose the existence of a ground-truth reward function. Given this, why do we still need log \pi^*(a|s) as the representation of r (See Eq 8 and Eq 9)? 
2. The transition from Eq(9) to Eq(11) is questionable. Eq(7) represents the value function fitting target, not the policy optimization target. It is unclear how substituting Eq(9) into Eq(7) achieves policy optimization, despite the introduction of a reward term reparameterized by \pi_\theta. The authors must elaborate on how the substitution of Eq(9) into Eq(7) facilitates policy optimization, especially considering that Eq(7) is a value function target, not a policy optimization target.
3. There is a conflict between the definitions of reward in Line 223 and Eq 14, with the former being log \pi_ref + r and the latter being \log \pi_ref/pi + r. This inconsistency needs to be resolved.
4. 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. Importance sampling is typically required when the sampling policy and the current policy are inconsistent during policy optimization, not during value function learning as implied by Eq 6 and Eq 7.
soundness: 2
strengths: 1. The paper addresses a significant issue in the field of language model alignment and multi-step reasoning, which is a valuable contribution to the advancement of LLMs.
2. The presentation of the material is coherent and the paper is generally easy to follow, which aids in the understanding of the proposed DQO algorithm.
summary: The paper presents a novel offline reinforcement learning (RL) algorithm, Direct Q-function Optimization (DQO), aimed at improving the multi-step reasoning capabilities of large language models (LLMs). The authors propose formulating the response generation process as a Markov Decision Process (MDP) and utilize the soft actor-critic (SAC) framework to optimize a Q-function parameterized by the language model. The paper claims that DQO outperforms previous methods on math problem-solving datasets, GSM8K and MATH, establishing it as a promising approach for aligning language models.
weaknesses: A primary concern is the rationality behind the derivation of the proposed method. Please see the question section below.

---

confidence: 4
contribution: 3
initial_rating: 3
presentation: 3
questions: * Could you provide a more in-depth analysis of your strengths when it comes to handling long-horizon tasks?
* 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.
* What methods or strategies can be employed to effectively generate process value?
soundness: 3
strengths: * The theorem is presented in a straightforward manner.
* The experiments demonstrate promising results, particularly in the context of small-scale open-source LLMs.
summary: Direct Q-function Optimization (DQO) improves language model performance by treating text generation as a Markov Decision Process and using soft actor-critic methods to directly optimize the Q-function, demonstrating superior results on math problem-solving tasks.
weaknesses: * I believe the method can heavily rely on the accuracy of the process value. but the difficulties should be analyzed in your experiment.
* There is a lack of analysis regarding whether the process value is fairly accessible or measurable.
* 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?

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: In Eq. (2), it appears $\beta$ is missing.

In Eq. (7), should it be $r$ instead of $\bar{r}$? 

In Line 220, "By plugging in (9) to (7)" -> "By plugging in Eq. (9) into Eq. (7)."

If I understand correctly, there is a missing minus sign in Eq. (11) which causing the RHS of Eq. (10) and Eq. (11) sum to zero. 

In Eq. (12), could you clarify the meaning of $ s_{h+1} \sim P(\cdot \mid s_h, a_h) $? Does this imply a deterministic process where $s_{h+1} = \text{Concat}(s_h, a_h) $?

In Eq. (14), should it be $ r(a_{h+l}, s_{h+l}) $?

If I’ve misunderstood any part of this, please feel free to correct me.
soundness: 3
strengths: The problem is significant, as sparse rewards cause the instability of learning process in RL.
summary: This paper focus on token-level optimization in LLMs. They rewrite the learning objective in LLMs and utilize the soft actor-critic (SAC) framework to optimize a Q-function directly parameterized by the language model. Experimental results on GSM8K and MATH  demonstrate the performance of proposed method.
weaknesses: It seems that the author just apply SAC in the context of LLM and rewrite the objectives. 

The excessive number of formulas has made the text lengthy and difficult to follow; the paper should be written in a more concise way. Below are some suggestions.
- Eq 9 can be removed and mention to parameterize the following modules $Q_\theta$, $\pi_{\theta}$, $V_{\phi}$, as it is almost the same with Eq 8. 
- Could it be more reasonable to move Eq. (6) and Eq. (7) to the preliminaries section or move them into appendix. Since Eq 6 and Eq 7 come from the previous work and it can be regarded to replace the original reward $r$ with $\bar{r}$.
- If $\pi^*$, $Q^*$ and $V^*$ are not used after the definition in Eq 3 - 5, also consider to move them into Appendix.
- Section 3.2 and Section 3.3 also can be more concise by referring to the original paper and discuss more about the challenges or special design in order to apply them in LLMs.

Additionally, minor issues in the formulas require proofreading to correct (see questions for specifics). 

The experimental results over Qwen2-7B-Instruct model are not convincing:  the proposed method achieves an improvement of less than 1% on the two datasets, it is not sure if it comes from randomness. Could you please provide statistical significance tests or report the average performance/standard error over different seeds? Alternatively, other experiments over additional models or tasks to demonstrate stronger improvement?

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. How were generated responses evaluated for each augmented dataset? What is the distribution of positive and negative responses in the training and testing data?

2. For Qwen2-7B-Instruct, why does DPO consistently outperform KTO and DRO, while DRO outperforms DPO for Gemma-1.1-7B-it? Could the authors provide insights into this discrepancy??
3. The related work [A1] 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?
4. Could the authors explain why DQO with lambda=1 outperforms lambda=0.95?
5. What is the correlation between the learned value function and constructed process scores during testing?

Minor issue:
1. Undefined/Misleading Notations: In eq (6), the expectation should be taken w.r.t. a_h. In Line 247, it should be V^{\pi}. The update equation (16) should be G^{\lambda}_{\phi, \theta}(s_h) = G^{(H-h)}_{\phi, \theta}(s_h) when \lambda=1. 
1. Grammar and Typos: The Line 240-241 should reference DRO instead of DQO.



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

[A2] Wei, Jason, et al. "Chain-of-thought prompting elicits reasoning in large language models." Advances in neural information processing systems 35 (2022): 24824-24837.
soundness: 2
strengths: 1. The proposed algorithm frames response generation as a MDP, making it more suitable for long-horizon, step-by-step reasoning tasks compared to DRO.
2. The approach demonstrates potential in effectively leveraging process rewards to enhance performance through feedback at each step.
3. DQO incorporates lambda-return and importance sampling to ensure efficient use of offline data.
4. Experimental results highlight the superior performance of DQO on widely used GSM8K and MATH datasets.
5. The paper is well-written and easy to follow.
summary: This paper introduces the Direct Q-Function Optimization (DQO) approach, which leverages a Markov Decision Process (MDP) formulation to enhance the performance of large language models (LLMs) on reasoning tasks. By adopting the multi-step nature of MDPs, DQO addresses limitations found in bandit-based offline methods, such as DRO, and remains effective even when preference-based datasets are unavailable.
weaknesses: 1. Fairness of Comparison:  According to Table 8, the size of datasets for each model differs significantly.  For example, DQO uses datasets up to four times larger than those for DPO. Although DPO's data subset is sampled from DQO's dataset, this disparity raises concerns about fair comparisons, especially given the relatively small performance margin of DQO over DPO for the Qwen2-7B-Instruct model when dataset differences are minimal.

2. Lack of Experimental Details: The paper lacks sufficient discussion on the evaluation of generated/augmented responses, the distribution of positive and negative responses in the training and testing data, and detailed training parameters (e.g., number of epochs for DRO and DQO). This makes it difficult to directly explain performance discrepancies, such as those seen between this work and [A1] for Qwen2-7B-Instruct (see Question 3).
3. Hard to assess the benefits of  MDP formulation: Without importance sampling, DQO’s performance (as shown in Table 5) is worse than DRO’s performance in Table 4. This raises the question of whether the key factor in DQO's superior performance is the MDP formulation itself or the use of importance sampling.
4. Risk of overfitting: The dependency on offline data and importance sampling introduces the risk of overfitting to this data, especially if it does not represent the diverse scenarios encountered in real-world applications.
5. Potentially high computational costs: DQO involves complex training procedures, including learning Q and V functions and using  \lambda-return and importance sampling. This complexity may lead to higher computational costs compared to DRO, especially for long-horizon tasks. The paper does not discuss computational efficiency; for example, chain-of-thought prompting [A1] can also provide intermediate checks without additional training. A more thorough comparison with CoT and [A2] is recommended.
6. Clarification of Specific Challenges: While this work effectively incorporates SAC, lambda_return, and importance sampling for improving reasoning performance, there is insufficient discussion on the challenges encountered during this integration. Highlighting these  challenges would strengthen the work.
7. Insufficient discussion on unbalanced data and process reward: While Table 1 notes that DQO can learn from unbalanced samples, no experiments substantiate this claim. In Section 4.4, a synthetic process reward mechanism is presented, but more explanation and experimentation with different process score designs are needed.

Meta-review: The paper tackles an important problem in language model alignment and multi-step reasoning, an area of growing significance for advancing the capabilities of Large Language Models (LLMs). Reviewers found the paper generally well-written and easy to follow, indicating a commendable clarity of presentation. However, a critical weakness was identified concerning the fundamental justification and derivation of the proposed DQO algorithm. This issue raises serious concerns about the validity and soundness of the core contribution.
While the paper possesses merits in addressing a relevant problem and presenting the material clearly, the identified weakness regarding the rationale behind the proposed method is substantial enough to warrant rejection.

Decision: Reject