Step-DPO (step-level preference optimization), explained

Step-DPO applies the DPO loss to single reasoning steps instead of whole answers. Each training example is a problem, a prefix of reasoning steps that are all correct, and two candidates for the next step: one that is correct and one that is the first mistake of a failed solution. The model is trained to prefer the correct step over the erroneous one, relative to a frozen reference, conditioned on the shared prefix. Everything else about DPO stays: no reward model at training time, no sampling inside the training loop, the same $\beta$-scaled log-ratio margin passed through a log-sigmoid. What changes is where the preference signal lands. Step-DPO sits in the post-training stack after SFT on reasoning data, as a cheap offline alternative or complement to process-reward-model-guided RL and to outcome-reward RL such as GRPO.

The problem and where it came from#

DPO works well for chat preferences but gives small gains on long mathematical reasoning. Lai, Tian, Chen, Yang, Peng and Jia (2024, The Chinese University of Hong Kong and Harbin Institute of Technology) diagnosed why: in a wrong solution the first error often appears midway, so rejecting the whole answer also pushes down the correct steps that precede it, and the model is never told which step was wrong. They proposed treating steps as the unit of preference, together with a data pipeline that needs no human-written corrections: generate solutions with the model itself, locate the first wrong step, and resample from just before it until a continuation reaches the right answer.

With about 10,000 step-level pairs and fewer than 500 training steps, they reported gains on MATH of 1.0 to 3.0 points over their SFT models (for example Qwen2-72B-SFT from 61.7% to 64.7%, Llama-3-70B-SFT from 56.9% to 59.5%, Qwen2-7B-SFT from 54.8% to 55.8%). Applied on top of Qwen2-72B-Instruct, which had already been through RLHF, it went from 69.4% to 70.8% on MATH and from 92.4% to 94.0% on GSM8K under their prompt. In a matched ablation with 5,000 pairs, vanilla DPO moved Qwen2-7B-SFT from 54.8% to 55.0% and Qwen2-72B-SFT from 61.7% to 62.5%, while Step-DPO reached 55.8% and 64.1%. They also found that the preferred step must be generated by the model being trained: using GPT-4 to write corrected steps gave 55.1% against 55.8% for self-generated ones.

The objective, with every symbol defined#

Let $x$ be a problem and $y=(s_1,\dots ,s_n)$ a solution split into steps. Let $s_{1\sim k-1}=(s_1,\dots ,s_{k-1})$ be a prefix of correct steps, $s_{\text{lose}}$ the incorrect step $k$ from a failed solution, and $s_{\text{win}}$ a correct step $k$ from a continuation that reached the right answer. ${\pi }_{\theta }$ is the policy being trained, ${\pi }_{\text{ref}}$ the frozen SFT model, $\sigma$ the logistic function and $\beta >0$ the strength of the implicit KL constraint. The Step-DPO loss is

Each log-probability is the sum of token log-probabilities over one step only, with the problem and prefix as context. Compare the DPO loss on full answers,

Write $\Delta (z\mid c)=\log {\pi }_{\theta }(z\mid c)-\log {\pi }_{\text{ref}}(z\mid c)$ and define the margin $m$ as the argument of $\sigma$. The gradient has the familiar DPO shape, a scalar weight times a contrast of log-likelihood gradients:

Why step-level pairs carry a sharper signal

Suppose a full-answer DPO pair shares a correct prefix $s_{1\sim k-1}$ and then diverges: $y_w=(s_{1\sim k-1},s_{\text{win}},t_w)$ and $y_l=(s_{1\sim k-1},s_{\text{lose}},t_l)$, where $t_w$ and $t_l$ are the tails after step $k$. The prefix terms cancel exactly in the margin, because both answers contain the same tokens in the same context. What remains is

The tail term contains everything after the mistake: the rejected solution's steps conditioned on an error the model will hopefully never make again. Pushing those down is cheap (they are often low-probability already) and inflates the margin, which shrinks the common weight $\sigma (-m)$ and therefore the gradient on the step that actually mattered. Step-DPO deletes the tail term. In practice DPO pairs are usually two independently sampled solutions with no shared prefix at all, so the margin also mixes in differences in phrasing of correct steps, which carry no information about correctness. Step-level pairs isolate the decision the model got wrong. Rafailov et al.'s token-level view of DPO ("From $r$ to $Q^{\ast }$") supports reading the step log-ratio as a credit assignment at that point in the sequence.

Next to its siblings

The data pipeline#

The paper's pipeline has three stages, and the dataset is the method as much as the loss is.

Error collection. Take problems with known answers, prompt the SFT model with a step-wise chain-of-thought prefix ("Let's think step by step. Step 1:") so every step starts with "Step i:", and keep the solutions whose final answer is wrong.

Step localisation. Check each wrong solution step by step until the first error, either by hand or with a strong model as judge, and record its index $k$. That step is $s_{\text{lose}}$.

Rectification. Sample several continuations from the reference model given $x$ and $s_{1\sim k-1}$, keep those that reach the correct final answer, and take the first step of one of them as $s_{\text{win}}$. Because a correct final answer does not guarantee a correct step, the authors filter again for step correctness. Humans or a judge model only locate errors and check steps; they never write a correction, which keeps $s_{\text{win}}$ in the model's own distribution.

In the paper, 374K SFT solutions were generated from MetaMath and MMIQC problems; 299K were used for SFT, and the remainder plus a subset of AQuA fed this pipeline to produce about 10,000 pairs.

Worked example: the loss and what the tail does to it#

Take $\beta =0.4$, the value the paper used for models below 72B. After some training, suppose the chosen step's log-ratio is $\Delta (s_{\text{win}})=+0.5$ and the rejected step's is $\Delta (s_{\text{lose}})=-0.3$. The Step-DPO margin is $m=0.4\times (0.5+0.3)=0.32$, the loss is $\log (1+e^{-0.32})=0.546$, and the gradient weight is $\beta \sigma (-0.32)=0.4\times 0.421=0.168$.

Now the same pair as full-answer DPO, where the tails are also scored. Suppose the correct tail has log-ratio $+0.6$ and the failed tail $-1.8$, the latter easy to push down because those steps follow an error. The tail term adds $0.4\times (0.6+1.8)=0.96$, so $m=1.28$, the loss falls to 0.245 and the weight $\beta \sigma (-1.28)=0.4\times 0.218=0.087$. The model has made the full-answer loss look half-solved without making step $k$ any more likely to be right, and the gradient on step $k$ has halved. Repeat this over a dataset and the loss curve flattens for reasons unrelated to the error you wanted fixed.

Toy runs: where the gain comes from#

We simulated reasoning as a chain of 8 steps, each with 6 options: 3 correct phrasings and 3 wrong ones. Once a step is wrong the chain enters an "error state" with its own options and the final answer is wrong. Policies are tabular softmaxes (numpy, seed 0), the reference sets per-step probabilities of a correct step, and we trained with full-batch gradient ascent at learning rate 0.5 and $\beta =0.4$ for 300 epochs.

One problem, equal data. With per-step success probabilities $(0.97,0.95,0.9,0.62,0.93,0.95,0.75,0.97)$, reference accuracy is 0.331. From 2,000 sampled chains we built 657 full-answer DPO pairs and 1,343 Step-DPO pairs. Here DPO was slightly better per epoch: 0.598 against 0.553 after 50 epochs, 0.868 against 0.855 after 300. That is not a contradiction of the paper. In a tabular world every step of a correct chain is clean positive supervision for that step, so a full-answer pair delivers eight useful signals while a step pair delivers one. Two other measurements matter more. At the end of training, 16% of DPO's average margin came from pushing down post-error steps, the tail term above; in a harder version (reference accuracy 0.066) it was 28%. And Step-DPO's data concentrates on early first errors: in the harder version it left the late hard step (step 7 of 8) at 0.819 against DPO's 0.918, because an error there only counts as "first" when every earlier step was right. A single round of first-error data under-trains late steps; regenerate pairs from the improved model.

Many problems, realistic coverage. The paper's setting is many problems, some of which the model almost never solves. We drew 300 problems with per-step success probabilities from a Beta(6, 1.2) distribution (mean reference accuracy 0.213; 27 problems below 0.1) and sampled 4 chains per problem.

Only 170 of the 300 problems produced at least one correct and one wrong chain, so full-answer DPO had 206 pairs, and only 14 of the 27 hardest problems were represented. Step-DPO's rectification, which only needs the remaining steps to go right from a correct prefix, produced 679 pairs covering 287 problems, including 23 of the 27 hardest. After 300 epochs mean accuracy was 0.379 with Step-DPO against 0.285 with DPO on the same 4 samples per problem. Rectification is not free: it generated 14,389 steps on top of the 9,600 in the initial samples. Giving DPO the same budget (10 samples per problem, 24,000 steps) raised it to 613 pairs on 246 problems and a mean accuracy of 0.392, slightly above Step-DPO. The difference is where the gains landed. On the 27 hardest problems Step-DPO reached 0.161 against 0.133 for budget-matched DPO (from 0.076); on problems the reference already solved at least half the time, DPO reached 0.751 against 0.619.

Read together, the toys say what Step-DPO is for. Its loss removes the tail term and puts all the gradient on the decision that went wrong, and its pipeline creates training signal on problems the model rarely solves outright, because a correct continuation from a correct prefix is far more likely than a correct solution from scratch. On easy problems, where correct full solutions are plentiful, outcome-level pairs are as good or better. The paper's largest gains were on 70B-class models and on competition problems, which fits that picture.

The algorithm step by step#

# data construction (offline, once per round)
for x, answer in problems:
    for y in sft_model.sample(x, n=4, prefix="Let's think step by step. Step 1:"):
        if final_answer(y) == answer: continue
        steps = split_on("Step i:", y)
        k = first_wrong_step(x, steps)                  # human or judge model
        prefix, s_lose = steps[:k], steps[k]
        for cont in sft_model.sample(x + prefix, n=8):  # rectification
            if final_answer(cont) == answer and step_is_correct(x, prefix, first_step(cont)):
                D.append((x, prefix, first_step(cont), s_lose)); break

# training: standard DPO machinery with step-level spans
ref = frozen copy of sft_model
for batch in D:                                         # e.g. 8 epochs (7B) at batch 128
    ctx = x + prefix
    m = beta * ((logp(pi, s_win | ctx) - logp(ref, s_win | ctx)) - (logp(pi, s_lose | ctx) - logp(ref, s_lose | ctx)))
    loss = -logsigmoid(m).mean()                       # loss only on step tokens; context masked
    loss.backward(); opt.step(); opt.zero_grad()
log: margin, accuracy of m > 0 on held-out pairs, chosen and rejected log-ratios, MATH/GSM8K accuracy

Any DPO trainer can run Step-DPO if you pack each example as prompt = problem plus prefix, chosen = $s_{\text{win}}$, rejected = $s_{\text{lose}}$, because DPO trainers already mask the prompt from the loss. The engineering is all in the data: step segmentation, first-error localisation and rectification sampling. About 10,000 pairs was enough in the paper. Expect rectification to dominate generation cost on hard problems, since each failed solution needs several continuations to find one that succeeds.

Hyperparameters that matter#

If you use LoRA instead of full weights, raise the learning rate by roughly 10 times as for any DPO run. Tune $\beta$ by watching the held-out pair accuracy and the chosen log-ratio: if the chosen log-ratio goes strongly negative while accuracy rises, the model is winning by pushing everything down.

Compute and memory cost#

Training costs the same as DPO per token, and less per example, because only one step per side is scored and the shared prefix is context. With a frozen reference you hold two copies of the model (policy with optimiser state at about 16 bytes per parameter, reference at 2 bytes), or one plus an adapter switch with LoRA. Sequences are prefix plus one step, and the prefix's activations still need a forward pass on both policy and reference. The paper ran 7B to 72B models with DeepSpeed ZeRO-3 and CPU offload. Data construction is the larger bill: sampling solutions, judging steps (a judge-model call per step until the first error) and rectification rollouts. In our multi-problem toy, rectification consumed 1.5 times as many generated steps as the initial sampling.

Failure modes and how to detect them#

Wrong first-error labels

If the judge marks a correct step as wrong, Step-DPO pushes down good reasoning with full force, because all the gradient lands on that step. Audit a random sample of localisations by hand before training, and measure judge agreement on a small labelled set (PRM800K from Lightman et al. is a public reference for step labels).

Correct answer, wrong step

A continuation can reach the right answer despite a flawed first step. The paper filters for this explicitly. Without the filter, the model learns to prefer lucky steps.

Out-of-distribution winners

Steps written by a stronger model have low probability under the reference, and the paper found them less effective (55.1% against 55.8%). Symptom: the chosen log-ratio climbs very slowly while the rejected one falls. Keep $s_{\text{win}}$ self-generated.

Everything goes down

As in DPO, both chosen and rejected log-probabilities can fall while the margin grows. Track both log-ratios separately; if the chosen one falls by more than a few nats, lower the learning rate, raise $\beta$, or add a small NLL term on $s_{\text{win}}$.

Late steps under-trained

First-error data over-represents early mistakes, as the toy showed. Look at the distribution of $k$ in your dataset, and regenerate pairs from the trained model for a second round.

When to use Step-DPO, and what to use instead#

Minimal implementation#

A Step-DPO loss for causal LMs. The only difference from a DPO loss is how the label mask is built: context is problem plus prefix, and only the step tokens are scored. With the TRL library's DPO trainer you get the same effect by putting problem plus prefix in the prompt field.

import torch, torch.nn.functional as F

def span_logp(model, ctx_ids, step_ids):
    ids = torch.cat([ctx_ids, step_ids], dim=1)
    logits = model(input_ids=ids).logits[:, :-1].float()
    lp = torch.log_softmax(logits, -1).gather(-1, ids[:, 1:, None]).squeeze(-1)
    mask = torch.zeros_like(lp); mask[:, ctx_ids.shape[1] - 1:] = 1.0   # score step tokens only
    return (lp * mask).sum(-1)

def step_dpo_loss(policy, ref, ctx, s_win, s_lose, beta=0.4):
    # ctx: problem + correct prefix; s_win / s_lose: next-step token ids (pad and mask in real code)
    pw, pl = span_logp(policy, ctx, s_win), span_logp(policy, ctx, s_lose)
    with torch.no_grad():
        rw, rl = span_logp(ref, ctx, s_win), span_logp(ref, ctx, s_lose)
    m = beta * ((pw - rw) - (pl - rl))
    loss = -F.logsigmoid(m).mean()
    stats = {"margin": m.mean().item(), "acc": (m > 0).float().mean().item(),
             "chosen_logratio": (pw - rw).mean().item(), "rejected_logratio": (pl - rl).mean().item()}
    return loss, stats

# TRL: build rows {"prompt": problem + prefix, "chosen": s_win, "rejected": s_lose}
# and train with DPOTrainer(beta=0.4); prompt tokens are masked from the loss by default.

In POFE, Step-DPO is a recipe whose first stage is the data pipeline: sampling, judge-based first-error localisation and rectification, with the pair distribution over $k$ shown before training starts.

Further reading#

Lai et al. (2024) is the paper, with the pipeline in Section 3.2 and the DPO and OOD ablations in Section 4.3. Rafailov et al. (2023) is the DPO paper and Rafailov et al. (2024, "From $r$ to $Q^{\ast }$") gives the token-level reading that justifies scoring spans. Lightman et al. (2023, "Let's Verify Step by Step") and Wang et al. (2023, Math-Shepherd) are the human-labelled and automatic routes to step-level supervision used by process reward models. Pang et al. (2024) is the outcome-level iterative alternative for reasoning.

Frequently asked questions#

References#

1. Xin Lai, Zhuotao Tian, Yukang Chen, Senqiao Yang, Xiangru Peng, Jiaya Jia (2024). Step-DPO: Step-wise Preference Optimization for Long-chain Reasoning of LLMs. arxiv.org/abs/2406.18629
2. Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, Chelsea Finn (2023). Direct Preference Optimization: Your Language Model is Secretly a Reward Model. arxiv.org/abs/2305.18290
3. Rafael Rafailov, Joey Hejna, Ryan Park, Chelsea Finn (2024). From r to Q*: Your Language Model is Secretly a Q-Function. arxiv.org/abs/2404.12358
4. Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, Karl Cobbe (2023). Let's Verify Step by Step. arxiv.org/abs/2305.20050
5. Peiyi Wang, Lei Li, Zhihong Shao, R. X. Xu, Damai Dai, Yifei Li, Deli Chen, Y. Wu, Zhifang Sui (2023). Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations. arxiv.org/abs/2312.08935
6. Richard Yuanzhe Pang, Weizhe Yuan, Kyunghyun Cho, He He, Sainbayar Sukhbaatar, Jason Weston (2024). Iterative Reasoning Preference Optimization. arxiv.org/abs/2404.19733

Related reading