PaperlayerView exact revision on arXiv

Source-linked AI summary

Likelihood-Based Reward Designs for General LLM Reasoning

Ariel Kwiatkowski, Natasha Butt, Ismail Labiad, Julia Kempe, Yann Ollivier

arXiv:2602.03979v1cs.CL

TL;DR

The paper addresses the need for reward functions that avoid task-specific verifiers and sparse binary feedback in CoT reinforcement learning. It systematically compares probability-based and log-probability rewards across verifiable and non-verifiable settings, finding that log-probability rewards are the only variant effective across all setups. They match standard RL success rates with better perplexity on verifiable tasks and match SFT on long-form tasks, while probability rewards fail in non-verifiable settings.

  • Problem

    CoT reinforcement learning commonly relies on benchmark-specific binary rewards, motivating reward signals that are less sparse and do not require task-specific verifiers.

  • Method

    The paper systematically evaluates probability and log-probability rewards against SFT and standard RL across verifiable and non-verifiable CoT-training domains.

  • Results

    Log-probability rewards perform well across all tested settings, match standard 0/1 RL success rates with substantially better perplexity on reasoning benchmarks, and match SFT on long-form proofs.

  • Takeaways & Limitations

    Log-probability rewards provide a single training criterion for CoT fine-tuning across short verifiable and long non-verifiable answers without relying on ground-truth correctness labels.

  • Takeaways & Limitations

    The paper leaves the causes of non-verifiable CoT shortening and the possible transition between short and long-answer regimes for future work.

Abstract

from arXiv · show

Fine-tuning large language models (LLMs) on reasoning benchmarks via reinforcement learning requires a specific reward function, often binary, for each benchmark. This comes with two potential limitations: the need to design the reward, and the potentially sparse nature of binary rewards. Here, we systematically investigate rewards derived from the probability or log-probability of emitting the reference answer (or any other prompt continuation present in the data), which have the advantage of not relying on specific verifiers and being available at scale. Several recent works have advocated for the use of similar rewards (e.g., VeriFree, JEPO, RLPR, NOVER). We systematically compare variants of likelihood-based rewards with standard baselines, testing performance both on standard mathematical reasoning benchmarks, and on long-form answers where no external verifier is available. We find that using the log-probability of the reference answer as the reward for chain-of-thought (CoT) learning is the only option that performs well in all setups. This reward is also consistent with the next-token log-likelihood loss used during pretraining. In verifiable settings, log-probability rewards bring comparable or better success rates than reinforcing with standard binary rewards, and yield much better perplexity. In non-verifiable settings, they perform on par with SFT. On the other hand, methods based on probability, such as VeriFree, flatline on non-verifiable settings due to vanishing probabilities of getting the correct answer. Overall, this establishes log-probability rewards as a viable method for CoT fine-tuning, bridging the short, verifiable and long, non-verifiable answer settings.

1 Introduction

The paper systematically evaluates likelihood-based rewards for CoT training across verifiable and non-verifiable domains. Log-probability rewards are the only tested variant that performs well across all settings, while also improving perplexity and bridging both domains.

  • Systematic evaluation: The study compares probability and log-probability rewards with SFT and standard RL across verifiable and non-verifiable CoT-training settings.Experiments cover MATH, DeepScaleR, Alpaca, and the non-verifiable proof portion of NuminaMath, alongside variants including VeriFree, RLPR, and JEPO.
  • Universality of log-probability rewards: Log-probability rewards perform well in every tested scenario, whereas other reward variants fail in one or more settings.The scenarios include short, verifiable answers and long, non-verifiable answers.
  • Verifiable domains: In verifiable domains, probability-based rewards slightly outperform base RL on greedy success rate and avoid sampling an answer during training.Log-probability rewards also combine strong success rates with substantially better perplexity than base RL and probability-based rewards.
  • Non-verifiable domains: In non-verifiable domains, base RL and pure probability rewards collapse because long-answer probabilities vanish, while log-probability rewards perform similarly to SFT.This contrasts with the viability of log-probability rewards across both answer types.
  • CoT length: Log-probability rewards initially shorten CoTs; lengths recover on verifiable tasks but remain short on non-verifiable tasks.Length rewards and KL penalties preserve CoTs but hurt performance, and base RL and VeriFree do not show the same shortening on verifiable domains.
  • Implication: The paper positions log-likelihood rewards as a single, broadly applicable criterion for bridging verifiable and non-verifiable LLM fine-tuning.The criterion is motivated by its consistency with the log-likelihood signal used during pretraining and its applicability without task-specific verifiers.

2 Method

The paper formulates chain-of-thought reinforcement learning around reference-answer rewards, comparing binary correctness with probability- and log-probability-based alternatives. It evaluates these methods across verifiable and non-verifiable settings, with log-probability rewards aligned to pretraining and computable without sampling an answer.

  • RL setup: The framework has the model generate a chain of thought z and answer a for each prompt p, then optimize a reward R based on the answer.Base RL samples both z and a before checking correctness, whereas probability-based rewards use a reference answer a⋆.
  • Reward designs: Probability-based rewards directly compute the reference answer’s probability or log-likelihood given the generated chain of thought.Log-prob rewards can be computed in one transformer pass and do not require sampling an answer after z.
  • Reward designs: The compared reward variants include Base RL, Probability (VeriFree), Average Probability, Log-prob, Average Log-prob, and JEPO.Average Log-prob divides the summed answer-token log-probability by answer length, changing how dataset answers are weighted.
  • Reward designs: Log-probability rewards use the reference answer’s token log-probability, while Average Log-prob rescales it by the number of reference-answer tokens.The unnormalized reward has larger magnitude for longer answers; length normalization produces different sample weighting.
  • Evaluation: The study reports success rate and log-probability metrics, including greedy and temperature-T = 1 sampling success, with per-token and per-answer averaging schemes.Log-probability metrics are difficult to estimate because the expectation over chains of thought lies inside the logarithm, so Monte Carlo estimates are used.
  • Evaluation: Monte Carlo estimates approximate the expectation over sampled chains of thought using N samples, including logprob-MC1 and the less frequent logprob-MC32 estimate.For SFT without CoT, no expectation over z exists, so log π(a⋆|p) can be reported directly.

3 Experimental Results

Experiments compare likelihood-based reward variants with SFT and standard RL across verifiable and non-verifiable domains. Log-probability rewards generally combine strong greedy success or SFT-level performance with substantially better perplexity, while sampling, probability rewards, and long-CoT interventions expose important trade-offs.

  • Setup: Experiments use Llama-3.2-3B-Instruct and Qwen-2.5-3B-Instruct across verifiable and non-verifiable datasets, comparing SFT, Base RL, probability, log-probability, and JEPO rewards.Verifiable experiments use G = 4 and G = 32, while non-verifiable experiments use G = 4; verifiable runs include KL regularization, whereas the main non-verifiable results do not.
  • Verifiable domains: At G = 32, all probability- and log-probability-based variants outperform Base RL on greedy success in verifiable domains.The comparison concerns deterministically decoded answers; sampled-answer rankings can differ.
  • Verifiable domains: T = 1 answer sampling generally worsens performance and causes log-probability variants to underperform Base RL and Probability, unlike greedy decoding.The authors find no strong greedy-sampling difference between JEPO and simple Logprob, while JEPO is more computationally heavy.
  • Verifiable domains: Only Logprob, AvgLogprob, and JEPO achieve good perplexities, significantly improving SFT, whereas Base RL and Probability yield poor perplexity despite competitive success rates.The authors attribute this pattern to log-probability training assigning nonzero probability to plausible alternatives rather than producing sharply confident predictions.
  • Non-verifiable domains: In non-verifiable domains, log-probability variants consistently match SFT, Probability fails to improve, and Average Probability is noisier but remains close to the log-probability family.The log-probability family also undergoes a rapid CoT-length collapse, reducing the resulting behavior toward SFT.
  • CoT length: Length penalties prevented the CoT-length dip but reduced performance, while warm-start training stabilized length yet only matched SFT under a reasonable compute budget.A related JEPO study exceeded SFT after training roughly an order of magnitude longer, with larger batches and a lower learning rate, suggesting a substantially higher compute requirement.

4 Conclusion

The paper presents log-probability rewards as a unified CoT training signal across verifiable and non-verifiable domains, without ground-truth correctness labels. They match standard RL success rates on reasoning benchmarks, improve perplexity, and match SFT on long-form proofs.

  • Log-probability rewards work in both verifiable and non-verifiable domains without relying on ground-truth correctness labels.
  • On MATH and DeepScaleR, log-probability rewards match standard 0/1 RL success rates while substantially improving perplexity.
  • On long-form proofs, log-probability rewards match supervised fine-tuning, whereas other probability-based variants perform much worse.
  • The same criterion can be carried across short, verifiable answers and long, non-verifiable answers.

A Losses and Advantages for the Rewards Considered

The reward analysis formulates likelihood-based rewards within a policy-gradient loss, showing how z-independent control terms can be subtracted without changing the expected gradient.

  • The analysis considers rewards that depend on the chain-of-thought and possibly on the policy, including log πθ(a⋆|z).
  • The resulting loss has the same gradients, up to sign, as the original objective.
  • A z-independent expression cθ can be subtracted from the reward-weighted policy gradient without changing its expectation.
  • In RLOO, cθ is the average reward over independently sampled completions.

B Experimental details

Experiments use synchronous RLOO with fixed optimization and batching choices across two verifiable and two non-verifiable datasets, using standardized prompting, completion, and answer-verification procedures.

  • Each batch contains 8 questions with G different chain-of-thought completions, and each batch constitutes one training step.
  • The datasets comprise MATH and DeepScaleR as verifiable benchmarks, plus Alpaca and another long-form dataset as non-verifiable settings.
  • Experiments use DeepSeek-R1–style prompts, model-specific instruct templates, and an assistant prefill of <think> to initiate reasoning.
  • Training generates G completions up to 1024 tokens, using G = 32 for verifiable domains and G = 4 for non-verifiable domains.
  • Base RL assigns rewards of 100 for correct answers, 10 for correctly formatted incorrect answers, and 0 when parsing fails.

C.1 Verifiable Domains

The verifiable-domain appendix provides learning curves and final results across model–dataset combinations, reward algorithms, metrics, and group sizes, including JEPO at G = 4.

  • G = 32: Figures 3–5 show verifiable-domain learning curves for Qwen 2.5 3B Instruct and Llama 3.2 3B Instruct on MATH and DeepScaleR with G = 32.
  • G = 4: JEPO was run only with G = 4 for efficiency reasons.
  • G = 4: Table 3 reports final performance across all algorithms, including JEPO, and across the evaluated metrics for verifiable domains at G = 4.
  • G = 4: The appendix states that the G = 4 conclusions mirror those of Table 1.
  • G = 4: Figures 6–9 show corresponding verifiable-domain results with G = 4 for Llama 3.2 and Qwen 2.5 on MATH and DeepScaleR.

C.2 Non-verifiable Domains

This section provides complementary non-verifiable-domain results across additional model and dataset combinations.

  • Figures 10–12 report non-verifiable results for Llama 3.2 3B Instruct on NuminaProof and Alpaca, and Qwen 2.5 3B Instruct on Alpaca.These figures complement Figure 2 across additional model/dataset combinations.

D Attempted regularization methods

The authors test KL regularization, length penalties, and warm-start initialization to stabilize chain-of-thought training in non-verifiable domains, but these interventions do not surpass SFT on perplexity.

  • Regularization methods: KL regularization keeps the model close to its initial parameters as one strategy for stabilizing non-verifiable chain-of-thought training.
  • Regularization methods: Length penalties apply a negative reward for each missing token below threshold l0, with thresholds of 100, 150, 300, and 500 tested.The penalty magnitude is calibrated to offset the initial reward increase accompanying chain-of-thought shortening.
  • Warm start: Warm-start training uses answer-only SFT on question, completion, and answer triples containing generated chain-of-thoughts, then initializes the various CoT training methods.The CoT tokens are masked during SFT so the model learns to answer in their presence without training the CoT itself.
  • Warm start: Warm-start initialization partly stabilizes CoT collapse, with lengths stabilizing around 100–200 tokens rather than 5 tokens in the cold-start case.The result supports the authors’ intuition that CoT presence initially affects performance negatively.
  • Outcome: Warm-start models remain close to the SFT baseline in perplexity and fail to beat it under the reported compute budget.The authors note that substantially more compute might be needed to reproduce results where JEPO exceeds SFT.

E Impact of the Marginal Log-Probability Estimate

Naive MC1 estimates can differ substantially from MC32 estimates of the true reference-answer log-probability. This discrepancy is especially large for base RL and probability rewards in verifiable domains.

  • Significant differences can occur between naive MC1 and MC32 estimates of the true reference-answer log-probability.
  • In verifiable domains, base RL and probability rewards show a large difference between the two estimates.
  • Figures 18 and 19 visualize these results for a subset of experimental settings because frequent MC32 evaluations are computationally costly.

F Correlation Analysis

The analysis compares global and per-question correlations between CoT length and the probability or log-probability of obtaining the correct answer. Local correlations are especially relevant for group-relative advantages, and their observed signs align with CoT degradation or stability across settings.

  • Correlation Analysis: For 100 random problems per dataset, the authors generate 1000 CoTs and correlate CoT length with the probability or log-probability of obtaining the correct answer.
  • Correlation Analysis: Global correlation pools all CoTs, whereas average local correlation computes per-prompt correlations and averages them across prompts.
  • Correlation Analysis: The two correlation variants can differ significantly, similarly to Simpson’s paradox.
  • Correlation Analysis: Local correlations are more consequential than global correlations for GRPO or RLOO because group-relative advantages compare trajectories for the same problem.
  • Correlation Analysis: For log-probabilities, negative local correlations coincide with CoT degradation on Numina and MATH, while slightly positive local correlation for probability on MATH coincides with no CoT collapse.
  • Correlation Analysis: Figures 20–22 use 100 questions with 1000 CoTs each, while displaying 20 questions for visual clarity; Figure 19 reports similar patterns for Qwen 2.5 3B Instruct on DeepScaleR as for the Llama model.
Loading 2602.03979v1…