Full-Step-DPO: Self-Supervised Preference Optimization with Step-wise Rewards for Mathematical Reasoning

TL;DR

Full-Step-DPO addresses the limitation of prior preference-based methods in mathematical reasoning by optimizing every step in a reasoning chain using step-wise rewards. It introduces a self-supervised Process Reward Model (PRM) to score each step and a Step-wise DPO Loss that dynamically weights per-step gradients by these rewards, enabling true step-wise optimization. The method is trained via a three-stage pipeline and evaluated across multiple open-source backbones on GSM8K, MATH, and out-of-domain datasets GK2023 and OCW, consistently surpassing DPO and Step-DPO baselines and demonstrating stronger, more robust reasoning capabilities. The results suggest significant practical impact for improving reasoning in LLMs and indicate avenues for extending step-wise reward frameworks to other complex reasoning tasks.

Abstract

Direct Preference Optimization (DPO) often struggles with long-chain mathematical reasoning. Existing approaches, such as Step-DPO, typically improve this by focusing on the first erroneous step in the reasoning chain. However, they overlook all other steps and rely heavily on humans or GPT-4 to identify erroneous steps. To address these issues, we propose Full-Step-DPO, a novel DPO framework tailored for mathematical reasoning. Instead of optimizing only the first erroneous step, it leverages step-wise rewards from the entire reasoning chain. This is achieved by training a self-supervised process reward model, which automatically scores each step, providing rewards while avoiding reliance on external signals. Furthermore, we introduce a novel step-wise DPO loss, which dynamically updates gradients based on these step-wise rewards. This endows stronger reasoning capabilities to language models. Extensive evaluations on both in-domain and out-of-domain mathematical reasoning benchmarks across various base language models, demonstrate that Full-Step-DPO achieves superior performance compared to state-of-the-art baselines.

Figures (5)

• Figure 1: Comparison between DPO, Step-DPO, and our Full-Step-DPO. DPO operates on solution-wise preference data. Step-DPO advances to step-wise data but optimizes only a single step. Full-Step-DPO optimizes all steps with a novel step-wise DPO loss, effectively enhancing the model's reasoning capability.
• Figure 2: The overall framework of Full-Step-DPO consists of three steps: (1) Training the PRM using the model itself and generated solutions. (2) Using the PRM to score and filter solutions to form preference data with step-wise rewards. (3) Training the policy model with the proposed step-wise DPO loss.
• Figure 3: Performance comparison of various verifications on GSM8K, with all models trained using our Full-Step-DPO.
• Figure 4: Accuracy of MetaMath-Mistral-7B-Full-Step-DPO using BoN decoding with $K=15$. The PRM uses a fixed sampling number $M=32$, while the simulation number $N$ varies. Purple bars indicate the A100-Hours cost for constructing PRM training data.
• Figure 5: Accuracy of MetaMath-Mistral-7B-Full-Step-DPO with different reward temperature $\gamma$.