Training Language Models to Self-Correct via Reinforcement Learning

Thank you for your review and for a positive assessment of our paper. We are glad that you liked the paper. To address your concerns, we have made edits to the submission, clarified the question regarding the sign of the reward function below, and clarified experimental results comparing our method to single-turn approaches. Please let us know if your concerns are addressed and if so, we would be grateful if you are willing to increase your score. We would be happy to discuss further.

overall sign of first-turn being negative, which may lead the model to introduce minor errors in the first step to leave room for self-correction in subsequent steps.

This is a great question! You are right that the reward shaping term we introduce will often enforce a negative multiplier on the first-attempt response. However, this reward bonus term is only used to reward the tokens generated in the second-attempt. The first-attempt response is only rewarded with its own correctness reward (see Figure 9 in Appendix A.4 for an illustration). Since each reward value is used to only train tokens in the corresponding turn, we refer to our algorithm as multi-turn RL, in contrast to typical single-turn RLHF algorithms that only provide a single scalar reward at the end of the entire rollout.

As a result, the first-attempt response is never trained to be worse, but rather this reward-shaping term encourages correct second-attempt responses on prefixes generated by less-accurate first-attempt responses and discourages incorrect self-correct responses on prefixes with correct first-attempts. This is essential in learning self-correction as it allows us to improve from the first attempt while also preventing hurtful changes in correct solutions at the first turn.

More comparisons to single-turn baselines

We already include the performance of a single-turn RL baseline in “w/o multi-turn training” in Table 4, which achieves significantly worse accuracy@t2 as expected. We use this ablation as a representative single-turn method because REST-EM (also called STaR or Expert Iteration (EI)) is known to perform worse than well-tuned REINFORCE-based methods (Ahmadian et. al 2024), which we use in our paper. Additionally, Havrilla et. al. 2024 found that REST-EM/STaR/EI performs comparably to PPO, which in turn is worse than RLOO as shown in Ahmadian et. al. 2024. Our paper already includes other experiments using STaR/REST-EM (Tables 1, 4), which further confirm that these approaches underperform compared to policy gradient-based methods.

Q3: The authors claim that a policy which tries to give its best-guess attempt in turn 1 followed by no self-correction should generalize worse to new problems than self-correction policies, but don't substantiate this claim with theoretical arguments or empirical findings.

One theoretical intuition for why self-correction should improve performance over simply maximizing first-attempt performance is that model performance should increase as it is able to leverage more tokens (analogously to how LLM reasoning performance increases with higher depths (Ye. at al 2024)), i.e., self-correction is able to benefit from larger test-time token budgets.

Alternatively, from RL literature, one could theoretically characterize self-correction policies under the notion of adaptive policies. These adaptive policies condition action predictions not only on the current state but also on past attempts or previous episodes. It is known in the RL literature that such adaptive policies especially excel in generalization settings. For example, the benefits of adaptive policies are studied in this paper: https://arxiv.org/abs/2107.06277. The sequential classification setting in this paper is conceptually similar to self-correction (though not the same). We will add this discussion to the paper (Appendix A.5).

Please note that we also empirically demonstrate in Table 4 that a model trained to maximize only its first attempt performance (“w/o multi-turn training”) performs worse than our method.

Ye, Tian, et al. "Physics of language models: Part 2.1, grade-school math and the hidden reasoning process." arXiv preprint arXiv:2407.20311 (2024).

Thank you for your review and for a positive assessment of our paper. We are glad that you liked the paper. To address your concern regarding the “simplicity” of our method, we have now run new experiments to better understand the importance of certain specific design choices in our algorithm. We have also updated the paper to include a flowchart that explains our chain of logic guiding the inclusion of each component of SCoRe (Figure 11). We believe that this addition should help practitioners better understand the various considerations that went into these design choices. Furthermore, we have conducted new experiments to address the remaining questions as we discuss below. Please let us know if your questions are addressed, and if so, we would be grateful if you would be willing to raise your score.

Comparison to the simple baseline of RL directly against the final answer, but with the self-correction prompt inserted after the first turn. We’ve run an experiment on this where we apply the RL loss on the concatenated action of (turn 1 solution, self-correction instruction, turn 2 solution). Unfortunately, we found the performance of this variant to be quite unstable, with the performance of turn 1 dropping significantly. We do believe that joint training of multiple turns of self-correction is a fruitful avenue for future research!

Q1: The natural way to do this would be to input first turns sampled from the base model and RL the model to give accurate second turns. Instead, this paper has the model generate both turns, with an RL training signal that rewards the model for high second-turn accuracy and with a large KL penalty against the base model for the first turn.

This is a great question and thanks for bringing this up! We did actually run this approach: generating first-attempt responses by sampling from the base model and then running single-turn RL to generate corrections on this fixed set of first-turn solutions. We’ve added this result to Table 4 in the paper (shown in blue), where it leads to only a 0.2% increase in self-correction $\Delta$(t1, t2), which is substantially lower than SCoRe (4.4% $\Delta$).

Our main finding is that while this approach is somewhat effective (outperforming simply prompting the base model for self-correction), it still suffers from distribution shift because training on the second turn still influences the model’s own distribution of first-attempt responses. While this issue may be resolved if the pre-trained base model does learn to decouple its first-attempt response from the second attempt, we found this was not the case. Hence, we applied a KL constraint that explicitly constrains the first-attempt response to not change much, which is the core idea behind the Stage I of our approach.

That said, you are right that the distinction between the large (stage I) and small (stage II) KL penalty can be a bit confusing. To clarify, while adaptations of REINFORCE for LLMs already come equipped with a KL penalty; the only modification we make is to explicitly incentivize stationarity of the first attempt response, and this is described as “large KL penalty” in the paper. We will add this clarification in Sections 5.1 and 5.2 to avoid any confusion.

Q.2: If so, then this effectively gives a large reward for second-turn correctness while actively penalizing first-term correctness. Is this intended? If so, why should this be better than just directly training on second-turn directness only?

This is a great question! This reward shaping term won’t negatively reinforce the first-turn correctness because we train each turn independently using its instantaneous reward (i.e. discount factor of 0). In other words, this negative term is only applied to turn 2’s reward and does not affect turn 1’s reward.

Therefore, the first-attempt response is never trained to be worse. Instead, this reward-shaping term encourages correct second-attempt responses on prefixes generated by less accurate first-attempt responses while discouraging incorrect self-corrections on prefixes stemming from correct first-attempts. This is essential in learning self-correction as it allows improving the first attempt while also preventing bad changes in already correct solutions at the first attempt.

We have already conducted an ablation study by removing this reward shaping term (see Table 4; “w/o reward shaping”), and found that it causes the method to perform worse than SCoRe.

Thank you for your review and the positive assessment of our paper. We are glad that you find our work to have commendable clarity and contain detailed explanations. To address the weaknesses and questions raised in the review, we have conducted additional experiments with the open-source 2B Gemma 2 model, showing that (1) our method similarly boosts the self-correction performance of the open source model, (2) improves multi-turn self-correction when trained with more than two attempts, and (3) generalizes to self-correction on completely held-out datasets (see results on Functional Math and MathOdyssey in Appendix A.1).

New experiments on open-source Gemma models. Beyond the models studied in the paper (Gemini 1.0 and Gemini 1.5 Flash), we have now added additional experiments on the open-source 2B Gemma v2 model, and found that SCoRe similarly boosts its self-correction performance. We have added these results to the paper in Appendix A.1, as well as below:

Multi-turn experiments. We now scale SCoRe to train for three attempts of self-correction and find a positive self-correction performance from attempt 2 to attempt 3. Full results are added in Appendix A.1 along with a summary of the results for Gemma 2 models shown above.

To extend SCoRe to multiple turns, we break Stage 1 into two sub-stages, say Stage 1(a) and Stage 1(b), with Stage 2 remaining unchanged. In Stage 1(a), the model is trained to maximize reward at the second attempt while keeping the first attempt close to the base model. Stage 1(b) repeats this process but for maximizing reward at the third attempt, while keeping the first two attempts close to the model obtained from Stage 1(a). Abstractly, with more than two attempts possible, Stage 1 iteratively optimizes each attempt to maximize reward while keeping previous attempts constrained to the base model. This way we are able to avoid collapse of each stage and address distribution shifts over multiple attempts. Stage 2 then proceeds as usual, optimizing the reward across all attempts and applying reward bonuses to incentivize the difference between rewards at a given attempt and the immediately previous attempt.

For the MATH dataset, why is a portion of the test data used for training? Could this make the evaluation less comprehensive? Thanks for the question. It is a common practice to use the MATH500 test set (which is uncontaminated) for evaluation while using the remaining MATH data for training (as described in Lightman et. al 2023). Several prior works studying reasoning on the MATH dataset have adopted a similar protocol for designing train/test splits ( Singh et al. 2024; Ying et. al. 2024; o1 blog post OpenAI 2024)

We also emphasize that all of our comparisons use identical splits, ensuring that none of our evaluations or comparisons are biased or unfair. As a result, these comparisons should still allow us to draw meaningful and functional conclusions about various approaches for training for self-correction.

Singh, Avi, et al. "Beyond human data: Scaling self-training for problem-solving with language models." arXiv preprint arXiv:2312.06585 (2023). Lightman, Hunter, et al. "Let's verify step by step." arXiv preprint arXiv:2305.20050 (2023). Ying, Huaiyuan, et al. "Internlm-math: Open math large language models toward verifiable reasoning." arXiv preprint arXiv:2402.06332 (2024). OpenAI, https://openai.com/index/openai-o1-mini-advancing-cost-efficient-reasoning/

Thanks to the authors for providing additional results and addressing my concerns. I will raise the score accordingly.

Thank you for your review and the positive assessment of our paper. We are glad that you find that our experiments are very solid and our paper to be in-depth. To address the weaknesses and questions raised in the review, we have conducted new experiments with the open-source 2B Gemma 2 model, demonstrating that SCoRe also enables positive self-correction performance with these open models. Additionally, we have added new results comparing inference-time of sequential sampling versus parallel sampling. Please let us know if your concerns and questions are addressed, and if so, we would be grateful if you would be willing to raise your score, thanks so much! We are happy to engage in further discussions.

New experiments on Gemma models. Beyond the models studied in the paper (Gemini 1.0 and Gemini 1.5 Flash), we have now added additional experiments on the open-source 2B Gemma v2 model, and found that SCoRe similarly boosts its self-correction performance of turn 2 accuracy from 16.8% to 24%. The results have been added to Appendix A.1 of the paper, and are included in the table :

Additional benchmarks. We absolutely agree that adding additional benchmarks and domains would be valuable, and this is one of the next steps of our research as well. However, we have been unable to find public datasets with large training splits suitable for this purpose (most benchmarks only provide evaluation splits, which are typically too small to be repurposed for training). Hence, we’ve added additional evaluations on held-out datasets (Functional Math and MathOdyssey) in Appendix A.1, showing that the self-correction abilities of our trained models generalized to out-of-distribution datasets. These results are also shown in the table above. It is worth noting that these datasets - especially MathOdyssey - are significantly harder than MATH.

In particular, with regards to AIME, we note that MATH already contains AIME problems, categorized as level 5 problems in the dataset (please see the discussion in Section 3.1 of the Hendrycks et al. MATH paper). To further analyze our method, we have added a breakdown of performance by difficulty level in Appendix A.2, which demonstrates the efficacy of our method across a spectrum of problem difficulties ranging from easy AMC problems (levels 1-2) to hard AIME ones. In particular, our method achieves a higher self-correction gap on AIME problems than even medium-difficulty ones (levels 3 and 4).

The suggestion to test on other domains (e.g., physics, chemistry) is also great. Would you have any specific recommendations for training sets or evaluation benchmarks that could be used to evaluate our method in these domains? We are absolutely happy to scale SCoRe up to these domains if you could point us to some train / test setups we could use. Currently, we are not aware of any specific public datasets for these domains that include non multiple-choice questions, which are essential for meaningful self-correction. If you have any suggestions, we would greatly appreciate them.

Inference-time of sequential sampling. We measured the inference time cost of sequential sampling and found that additional sequential inferences (i.e. turn > 1) have a constant additional latency that is significantly lower - around 2.5x faster - than that of the first attempt. This improvement is due to prefix caching during inference. As a result, although sequential self-correction cannot be parallelized, it is only 1 + c * (T - 1) times higher latency than fully parallel sampling, where T is the total number of attempts and c is a constant < 1.