Source-linked AI summary

ReFT: Reasoning with Reinforced Fine-Tuning

Trung Quoc Luong, Xinbo Zhang, Zhanming Jie, Peng Sun, Xiaoran Jin, Hang Li

arXiv:2401.08967v3cs.CL

TL;DR

SFT can generalize weakly when it relies on one reasoning path per question. ReFT first warms up with SFT, then uses PPO-based online reinforcement learning to sample multiple paths and derive rewards from answer correctness, outperforming SFT across three datasets while remaining compatible with inference-time voting and reranking.

  • Problem

    SFT commonly uses one correct CoT reasoning path per training question, although multiple valid paths may exist, limiting generalization.

  • Method

    ReFT warmups with SFT, then applies online PPO reinforcement learning that repeatedly samples responses, evaluates answer correctness, and updates the policy.

  • Results

    ReFT outperforms SFT in performance and generalization across experiments on three datasets using two foundation models.

  • Takeaways & Limitations

    ReFT learns from the same training questions as SFT and is compatible with majority voting and reward-model reranking.

  • Takeaways & Limitations

    ReFT requires more epochs to converge, while larger learning rates can cause instability and larger batches increase computational costs.

Abstract

from arXiv · show

One way to enhance the reasoning capability of Large Language Models (LLMs) is to conduct Supervised Fine-Tuning (SFT) using Chain-of-Thought (CoT) annotations. This approach does not show sufficiently strong generalization ability, however, because the training only relies on the given CoT data. In math problem-solving, for example, there is usually only one annotated reasoning path for each question in the training data. Intuitively, it would be better for the algorithm to learn from multiple annotated reasoning paths given a question. To address this issue, we propose a simple yet effective approach called Reinforced Fine-Tuning (ReFT) to enhance the generalizability of learning LLMs for reasoning, with math problem-solving as an example. ReFT first warmups the model with SFT, and then employs on-line reinforcement learning, specifically the PPO algorithm in this paper, to further fine-tune the model, where an abundance of reasoning paths are automatically sampled given the question and the rewards are naturally derived from the ground-truth answers. Extensive experiments on GSM8K, MathQA, and SVAMP datasets show that ReFT significantly outperforms SFT, and the performance can be potentially further boosted by combining inference-time strategies such as majority voting and re-ranking. Note that ReFT obtains the improvement by learning from the same training questions as SFT, without relying on extra or augmented training questions. This indicates a superior generalization ability for ReFT.

1 Introduction

ReFT addresses the limited generalization of SFT from single Chain-of-Thought annotations by warm-starting with SFT and then using PPO to sample and learn from multiple reasoning paths. Experiments report improved performance and generalization on standard math datasets, including when combined with inference-time methods.

  • Motivation: SFT commonly trains on one correct CoT annotation per question, although multiple valid reasoning paths may exist.The paper identifies this mismatch as a source of relatively weak SFT generalization and motivates ReFT.
  • Approach: ReFT warm-starts with SFT before refining the model with online PPO-based reinforcement learning.The warm-up provides initial mathematical response accuracy, while PPO performs the subsequent policy refinement.
  • Approach: PPO samples multiple reasoning paths and learns from them using rewards derived from ground-truth answer correctness, without a separately trained reward model.The training data provide question-answer pairs, allowing golden rewards to be derived directly from the answers.
  • Results: ReFT improves generalization over SFT while using the same training questions and no extra or augmented training questions.The paper presents this as richer supervision obtained through automatically sampled reasoning paths rather than additional question data.
  • Results: Experiments with CodeLLAMA and Galactica on GSM8K, MathQA, and SVAMP report significantly improved ReFT performance and generalization across natural-language and program-based CoTs.The paper also reports benefits from majority voting and reward-model reranking at inference time.

2 Related Work

Related work includes CoT prompt design, data engineering, and reinforcement-learning or reward-model approaches for mathematical reasoning. ReFT differs by improving the policy itself while remaining compatible with reranking methods.

  • Math Problem Solving: Recent math-solving work makes CoT prompts more comprehensive and fine-grained to expose step-by-step reasoning.Related approaches include natural-language CoT and Python programs used as CoT prompts.
  • Reinforcement Learning: Other work applies PPO and related preference-optimization algorithms to align language models with human preferences.The cited methods include DPO, IPO, and KTO alongside PPO-based reinforcement learning.
  • Math Problem Solving: For mathematical reasoning, outcome-based and process-based reward models have been used to rerank generated solutions.These approaches are described as achieving better performance than SFT and majority voting.
  • Reinforcement Learning: ReFT improves the policy model itself, while reward-model reranking can be integrated with the resulting policy.This positions ReFT as complementary to inference-time reranking approaches.

3 Method

ReFT combines supervised warm-up with online PPO-based reinforcement learning to sample and learn from multiple reasoning paths for mathematical problem solving.

  • 3 Method: ReFT has warm-up and reinforcement-learning stages, with the overall process specified in Algorithm 1.The warm-up stage uses question–CoT tuples, while reinforcement learning uses question–answer tuples.
  • 3 Method: During warm-up, supervised fine-tuning gives the model basic problem-solving ability and enables proper response generation.CoT generation is decomposed into next-token actions, ending when the <eos> token is produced.
  • 3 Method: At each timestep, the policy samples a vocabulary token from the current state, and the next state concatenates that token with the current state.The state contains the question tokens and all tokens generated so far.
  • 3 Method: ReFT uses online PPO to repeatedly sample responses, evaluate answer correctness, and update the policy using question–answer tuples.The value model is constructed by adding a linear value head to the policy model after warm-up.
  • 3 Method: Terminal rewards compare the extracted answer with the ground truth, returning 1 for correct answers and 0 otherwise, with optional 0.1 partial rewards for numeric answers.Non-terminal actions receive reward 0; the partial reward is intended to reduce sparse-reward effects.
  • 3 Method: The total reward combines the answer-based reward with a β-scaled KL-divergence term between the learned and initial policies.The method also uses generalized advantage estimation for advantage calculation.

6 // Reinforcement learning stage

The reinforcement-learning stage computes δt, Ât, and R̂t using the old policy and value model together with the question, sampled reasoning, and answers.

  • The stage computes δt using πθold, Vϕold, x, ê, ŷ, and y.
  • The stage computes Ât using πθold, Vϕold, x, ê, ŷ, and y.
  • The stage computes R̂t using πθold, Vϕold, x, ê, ŷ, and y.

15 return πθ

The return calculation uses temporal-difference quantities and generalized advantage estimation, while policy and value objectives are combined into a unified loss.

  • 15 return πθ: The temporal-difference definition uses a terminal-state value of V_ϕ(s_L+1) := 0, with λ and γ controlling return and TD discounting.λ ∈ (0, 1] is the reward discount factor, while γ ∈ [0, 1] is the TD discount factor.
  • 15 return πθ: The λ-return R̂_t is the sum of the generalized advantage estimate and the value estimate.Generalized advantage estimation is used to calculate the advantage.
  • 15 return πθ: The policy and value objectives are combined as a weighted unified loss, with α weighting the value objective.The old policy and value model are used for sampling CoTs and computing the estimated advantages and returns.

4 Experiments

Experiments across three math datasets and two foundation models compare ReFT with SFT and self-training. ReFT generally improves performance and generalization, supports inference-time voting and reranking, but encounters reward hacking on multiple-choice MathQA.

  • Experimental setup: ReFT is evaluated against SFT and self-training on GSM8K, SVAMP, and MathQA using Galactica and CodeLLAMA.The experiments include natural-language and program-based Chain-of-Thought representations.
  • Main results: ReFT improves over SFT across datasets, with CodeLLAMA gains of 6.7 points for N-CoT and 7.4 points for P-CoT on average.The reported exception is MathQAMCQ N-CoT.
  • Main results: ReFT outperforms offline and online self-training, indicating that on-policy exploration and reinforcement learning provide stronger improvements than these self-training baselines.The comparison reports that incorrect sampled instances also help guide exploration.
  • Limitations and analysis: On MathQAMCQ, reward hacking can assign positive rewards to incorrect reasoning when the final option prediction is correct despite an erroneous intermediate calculation.The paper identifies process-based rewards as a possible way to reduce this effect, while noting their development remains challenging.
  • Inference-time strategies: ReFT benefits from majority voting and reward-model reranking, with voting improving over SFT by 8.6 points on average and reranking by more than 3 points.Evaluation samples 100 Chain-of-Thought solutions per question for these strategies.
  • Inference-time strategies: The best P-CoT configuration reaches 81.2 accuracy on GSM8K and surpasses the compared open-source approaches and GPT-3.5-turbo using a 7B model.The configuration is CodeLLAMA with ReFT and reranking.

5 Analysis

ReFT’s training dynamics show improving reward and accuracy after exploration, while qualitative evaluation compares reasoning quality across logic, naming, and compactness. ReFT performs slightly better than SFT in the reported qualitative comparison.

  • Generalization: SFT converges and becomes overfitting when approaching the 40th epoch.
  • Generalization: ReFT’s mean reward reaches around 80% to 90% by the 40th epoch, while value accuracy also increases.
  • Generalization: KL divergence is initially large before stabilizing between 0 and 10, indicating exploration within a space containing appropriate programs.
  • Qualitative Evaluation: The qualitative evaluation scores reasoning paths on logic, naming, and compactness using 50 GSM8K questions and four annotators.Each criterion is scored from 0 to 1, with a perfect total score of 3.
  • Qualitative Evaluation: ReFT performs slightly better than SFT overall and outperforms the Warmup variant in the qualitative comparison.
  • Warm-up Analysis: ReFT policies initially decline after warm-up, with performance recovering as training reaches around the eighth epoch.

6 Conclusion

The conclusion presents ReFT as a fine-tuning method that explores multiple reasoning paths instead of relying on one annotation. Experiments report better performance and generalization than SFT, with compatibility with voting and reranking strategies.

  • ReFT explores multiple Chain-of-Thought annotations while searching for the correct answer, unlike SFT’s reliance on a single annotation.
  • Figure 5 compares SFT and ReFT accuracy across different numbers of warm-up epochs.
  • Experiments on three datasets and two foundation models show that ReFT outperforms SFT in performance and generalization ability.
  • ReFT is compatible with inference-time majority voting and reward model reranking.
  • ReFT shows superior performance to several publicly available open-source models of comparable sizes in math problem-solving.

7 Future Work

Future work extends ReFT beyond its current PPO-based setup and math problem-solving focus. Proposed directions include offline reinforcement learning, more efficient warm-up, process-based rewards, and broader reasoning tasks.

  • Future work proposes offline reinforcement learning techniques for ReFT.
  • The authors plan a warm-up-free method to improve training efficiency and performance and reduce the gap with reranking.
  • The authors identify process-based rewards as a direction because a well-trained process-based reward model may significantly enhance performance.
  • The authors intend to apply ReFT to more general reasoning tasks where inference can be formalized with Chain-of-Thought.

Limitations

The paper identifies training efficiency and reward design as limitations of ReFT. ReFT converges more slowly than SFT, while final-answer-only rewards can permit reward hacking when answer choices are limited.

  • Training Efficiency: ReFT requires more epochs to converge than SFT because it optimizes a non-differentiable objective and explores the generation space.
  • Training Efficiency: Increasing the learning rate may accelerate convergence but increases susceptibility to instability and policy collapse.
  • Training Efficiency: Larger batch sizes may help training but increase computational costs.
  • Reward Hacking: The reward function relies solely on the final answer, allowing reward hacking when the possible answers are limited to choices such as A, B, C, and D.
  • Reward Hacking: The paper suggests more detailed or process-based rewards to account for a broader range of factors.

C Case Study

Figure 7 compares SFT and ReFT responses to the same GSM8K question across epochs 1, 3, and 5. ReFT becomes shorter and correct at later checkpoints, while SFT remains incorrect; PPO was selected for its online exploration and direct reward use.

  • Response evolution: At epoch 1, ReFT is in warmup and produces a P-CoT similar to SFT.
  • Response evolution: At epochs 3 and 5, ReFT generates shorter, correct responses while SFT responses remain incorrect.
  • Algorithm choice: PPO was chosen because it explores CoT solution paths online and uses the reward function directly without reward-model training.
  • Response evolution: Figure 7 marks correct responses with green frames and incorrect responses with red frames.
Loading 2401.08967v3…