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

• DRO's bandit formulation neglects long-horizon reasoning and aggregates process rewards, losing their positional information.Source: §1 · pp. 2–3
• DQO formulates response generation as an MDP and uses soft actor-critic to fit Q- and V-values, with Q directly parameterized by the policy network.Source: §1 · p. 3
• DQO learns from offline negative or unbalanced samples, without requiring pairwise data.Source: §1 · p. 3
• The practical DQO instantiation adds λ-return and importance sampling, which the authors say stabilize training and ensure good performance.Source: §1 · p. 3
• DQO outperforms a wide range of alignment baselines on math problem-solving datasets, and available process rewards further boost its performance.Source: §1 · p. 3

• Canonical RLHF trains a reward model on preference data, then updates the language model with online PPO.Source: §2 · p. 3
• DPO-style methods directly parameterize reward models with the language model and train on offline preference data.Source: §2 · p. 3
• DRO combines offline policy learning with iterative value-function learning, while modeling language generation as a bandit.Source: §2 · pp. 3–4
• DQO models language generation as an MDP, enabling process supervision during training.Source: §2 · p. 4
• Multi-step methods use intermediate values, process feedback, or multi-turn preference learning, but some require pairwise data that may be difficult to obtain.Source: §2 · p. 4

• The token-level MDP treats the generated-token history as the state, the vocabulary as the action space, and token generation as deterministic transitions.Source: §3 · p. 4
• The formulation omits external-tool calls for notational simplicity; stochastic transitions can be approximated using samples.Source: §3 · p. 4
• The KL-regularized objective uses prompts from the dataset, token-level rewards, a regularization coefficient, and an initialization policy.Source: §3 · p. 5
• The framework can incorporate dense or sparse rewards from rule-based, turn-level, process-supervised, or outcome-supervised reward models.Source: §3 · p. 5
• Rewriting the KL term as cross-entropy plus entropy yields a maximum-entropy RL objective with a KL-regularized reward.Source: §3 · p. 5
• The optimal policy, soft Q-function, and soft V-function are interdependent, so knowing any two determines the third.Source: §3 · p. 5

4.1 The DQO objective

• DQO removes a separate Q-function parameter and directly infers the policy by parameterizing Q with the policy network.SAC learns the state-value and state-action-value functions from sampled states and actions.Source: §4.1 · p. 6
• DRO is a special case of DQO when a whole utterance is treated as one action and the horizon is set to H = 1.Source: §4.1 · p. 6

4.2 Mitigating Bias with λ-Return

• λ-return replaces one-step targets in the value and Q losses with a weighted average of n-step returns, addressing one-step TD bias and poor initialization.Source: §4.2 · p. 7

4.3 Reweighting Offline Data with Importance Sampling

• Importance sampling reweights pre-collected πref data toward πθ to mitigate distribution shift; the ratio is clipped at e to avoid gradient explosion, and weights are excluded from gradient computation.Source: §4.3 · pp. 7–8

5.3 Empirical Results and Case Studies

• On Qwen2.5-QwQ GSM8K, DQO raises greedy-decoding performance from 76.57% to 83.01%, surpassing all baselines; on Gemma, it improves the base model, while MATH greedy performance is comparable with DRO.On Gemma GSM8K, DQO surpasses all other baselines by a significant margin.Source: Table 2 · p. 8, Table 3 · p. 9, §5.3 · p. 10

5.1 Models and Dataset Constructions

• Qwen2.5-QwQ was trained on 4,096 sampled trajectories before Olympiad-level NuminaMATH; Gemma was trained on GSM8K and MATH.For each training problem, the base model generated 20 responses, labeled positive or negative to construct training data.Source: §5.1 · p. 9

5.2 Baselines and Evaluation

• SFT, reject sampling, DPO, KTO, and DRO were baselines; evaluation used GSM8K, MATH, and Qwen-only AIME24 with greedy and sampling decoding.GSM8K and MATH sampling averages five seeds; AIME24 reports average pass rate over 30 responses per problem.Source: §5.2 · pp. 9–10

5.3 Empirical Results and Case Studies

• The learned value estimate stays high through four correct segments, then drops after the erroneous x=y claim at the start of segment five.The case study uses a GSM8K training response divided into nine segments.Source: §5.3 · p. 10, Table 4 · p. 11

5.4 Ablation Studies

• More than 4.31% on GSM8K and 2.30% on MATH: lowering λ from 1.0 to 0.95 reduces Gemma performance, supporting λ-return's stabilizing role.The best performance was observed at λ = 1.Source: §5.4 · pp. 11–12, Table 5 · p. 12
• Over 3.80% on GSM8K and 2.16% on MATH: removing Q-loss importance sampling degrades performance; retaining it alone nearly preserves GSM8K but still costs on MATH.The ablation indicates that importance sampling in both Q-function and V-function losses contributes to DQO training.Source: Table 6 · p. 12, §5.4 · p. 12

5.5 DQO with Process Score

• 0.53% on GSM8K and 0.32% on MATH with greedy decoding: estimated process scores further improve DQO; sampling also rises on GSM8K and stays nearly unchanged on MATH.A prefix received a positive score when at least one of 20 generated completions was correct.Source: §5.5 · pp. 12–13, Table 7 · p. 13

6 Conclusion

• DQO frames language-model generation as an MDP and uses a soft actor-critic framework to learn V- and Q-functions, with Q directly parameterized by the language model.The paper describes this as an offline reinforcement learning algorithm for enhancing multi-step reasoning.Source: §6 · p. 13
• Unlike online methods such as PPO, DQO's offline approach bypasses the need for an extra reward model and online sampling during training.Source: §6 · p. 13
• Experiments on math problem-solving datasets report that DQO outperforms all of the paper's baselines.Source: §6 · p. 13

Dataset Construction

• For each dataset, the authors sample 20 responses per problem, format GSM8K and MATH prompts with guides, and feed Numina and AIME24 problems directly to the model.Prompts are then wrapped with the tokenizer's default chat template; Figure 1 illustrates prompts and sample answers from GSM8K, MATH, and Numina.Source: §A.1 · p. 13, Figure 1 · p. 14

Process Score Construction

• For failed responses, the process-score procedure samples 20 continuations from prefixes and assigns reward 1/n to preceding steps when a prefix has a correct continuation.The response is split into segments, and the procedure searches backward from the final segment; Figure 4 illustrates the construction.Source: §A.1 · p. 14, Figure 4 · p. 16

Dataset Construction

• Figure 2 plots each prompt's proportion of positive responses against prompt count and reports that negative responses dominate many problems, especially for Gemma-MATH and Qwen2.5-QwQ-Numina.Figure 3 plots response length against response count and reports longer Qwen2.5-QwQ generations than Gemma; Qwen2.5-QwQ responses trimmed at 4096 tokens are excluded.Source: Figure 2 · p. 15, Figure 3 · p. 15

Evaluation

• Evaluation uses greedy decoding and sampling; sampling uses top p=0.9, top k=16, threshold=0.01, temperature=0.7, and five runs, with mathematical equivalence determining correctness.GSM8K answers are extracted after “####”; MATH and AIME answers are parsed from boxed expressions.Source: §A.1 · p. 15

Additional Training Details

• DQO training took approximately 8 hours for Qwen2 and Gemma and 12 hours for Qwen2.5-QwQ on 8 NVIDIA A100 GPUs.Source: §A.2 · p. 16

SFT and Reject Sampling

• SFT and reject sampling used three epochs and a global batch size of 64, with final learning rates selected separately for each model.SFT used 2e-5 for Qwen and 5e-6 for Gemma; reject sampling used 2e-5 for Qwen2 and Qwen2.5-QwQ and 1e-6 for Gemma.Source: §A.2 · p. 16

DPO and KTO

• DPO used β=0.1, one epoch, and model-specific learning rates; KTO used the same hyperparameters as DPO.The learning rates were 5e-8 for Gemma and 1e-7 for Qwen2 and Qwen2.5-QwQ.Source: §A.2 · p. 17

DRO and DQO

• DRO and DQO used β=0.03, global batch size 256, and up to five epochs, selecting checkpoints from the training curve.The reported final checkpoints were from epoch one for DRO and Qwen2.5-QwQ DQO, and epoch two for Qwen2 and Gemma DQO.Source: §A.2 · p. 17

B.1 Results on Qwen2-7B-Instruct

• Across GSM8K, MATH, and AIME24, Qwen2-7B results show DQO comparable with DPO in some settings and exceeding the reported baselines in others.The passage reports results consistent with those for Qwen2.5-QwQ and Gemma.Source: §B.1 · p. 17
• GSM8K DQO improved from 72.77% to 87.95% for greedy generation and from 60.77% to 85.13% for sampling, surpassing other baselines by 1.22% and 0.70%, respectively.Source: §B.1 · p. 17
• On MATH, DQO reached 49.36% for sampling, 1.12% above DPO; on AIME24, it improved from 8.55% to 10.56% and exceeded all baselines by more than 0.59%.For greedy MATH decoding, DQO was comparable with DPO and DRO and exceeded other baselines by 1.64%.Source: §B.1 · p. 17

B.2 Additional Case Studies

• In relatively simple GSM8K and MATH cases generated by Gemma base, learned V-values reflect the correctness of each reasoning step.In one correct GSM8K response, predicted V-values gradually increase.Source: §B.2 · p. 17, Table 11 · p. 20
• DQO significantly improves the Qwen2-7B-Instruct base model, surpassing all baselines with greedy decoding and matching DPO while surpassing other baselines with sampling.Table 9 reports these comparisons without supplying numerical values in the provided passage.Source: Table 9 · p. 18

Response

• Zhang et al. (2024a) list Rest-MCTS*: LLM self-training via process reward-guided tree search.The entry identifies the work as an arXiv preprint, arXiv:2406.03816.Source: Response · p. 23
• Zhang et al. (2024b) list Iterative Nash Policy Optimization: Aligning LLMs with general preferences via no-regret learning.The entry identifies the work as an arXiv preprint, arXiv:2407.00617.Source: Response · p. 23
• Zhao et al. (2023) list SLIC-HF: Sequence likelihood calibration with human feedback.The entry identifies the work as an arXiv preprint, arXiv:2305.10425.Source: Response · p. 23