Training Language Models to Self-Correct via Reinforcement Learning

Self-correction is a highly desirable capability of large language models (LLMs), yet it has consistently been found to be largely ineffective in modern LLMs. Current methods for training self-correction typically depend on either multiple models, a more advanced model, or additional forms of supervision. To address these shortcomings, we develop a multi-turn online reinforcement learning (RL) approach, SCoRe, that significantly improves an LLM's self-correction ability using entirely self-generated data. To build SCoRe, we first show that variants of supervised fine-tuning (SFT) on offline model-generated correction traces are often insufficient for instilling self-correction behavior. In particular, we observe that training via SFT falls prey to either a distribution mismatch between mistakes made by the data-collection policy and the model's own responses, or to behavior collapse, where learning implicitly prefers only a certain mode of correction behavior that is often not effective at self-correction on test problems. SCoRe addresses these challenges by training under the model's own distribution of self-generated correction traces and using appropriate regularization to steer the learning process into learning a self-correction behavior that is effective at test time as opposed to fitting high-reward responses for a given prompt. This regularization process includes an initial phase of multi-turn RL on a base model to generate a policy initialization that is less susceptible to collapse, followed by using a reward bonus to amplify self-correction. With Gemini 1.0 Pro and 1.5 Flash models, we find that SCoRe achieves state-of-the-art self-correction performance, improving the base models' self-correction by 15.6% and 9.1% respectively on MATH and HumanEval.

한 줄 요약

SCoRe는 강화학습을 활용해 LLM이 자체적으로 오류를 수정하는 능력을 향상시키는 다단계 온라인 학습 방법이다.

핵심 기여도

SCoRe는 기존의 여러 모델이나 외부 지도 없이 **단일 모델**로 자체 생성 데이터를 기반으로 자기 수정을 학습한다.
기존의 SFT 기반 접근법이 **분포 이동**(distribution mismatch)과 **행동 붕괴**(behavior collapse) 문제를 일으킨다는 것을 실험적으로 입증한다.
SCoRe는 **두 단계의 강화학습**(Stage I 및 Stage II)을 통해 자기 수정 능력을 안정적으로 학습시킨다.
Gemini 1.0 Pro와 1.5 Flash 모델에서 각각 **MATH 15.6%**, **HumanEval 9.1%**의 자기 수정 성능 향상을 달성한다.

핵심 아이디어

기존의 자기 수정 학습 방법은 여러 모델을 사용하거나 외부 지도가 필요하다는 한계가 있다. SCoRe는 이 문제를 해결하기 위해 **단일 모델**이 **자체 생성 데이터**를 기반으로 자기 수정을 학습하는 방식을 제안한다. 핵심 아이디어는 **강화학습**(RL)을 활용하여 모델이 자신의 오류를 인식하고 수정하는 능력을 훈련시키는 것이다.

기존의 **SFT**(Supervised Fine-Tuning)는 학습 데이터가 기존 모델의 오류를 기반으로 생성되기 때문에, 학습된 모델이 자신의 오류를 수정하는 능력을 습득하지 못한다는 문제가 있다. 이는 **분포 이동**(distribution mismatch)이나 **행동 붕괴**(behavior collapse)를 유발한다. SCoRe는 이를 해결하기 위해 **온정책**(on-policy) 기반의 **다단계 강화학습**을 도입한다.

첫 번째 단계(Stage I)에서는 기존 모델의 분포에 가까운 정책을 초기화하여 **붕괴 가능성**을 줄이고, 두 번째 단계(Stage II)에서는 **보상 형성**(reward shaping)을 통해 자기 수정 과정 자체를 보상으로 삼아 학습한다. 이는 단순히 최종 응답의 정확도를 최적화하는 기존 RL 접근법과 구분된다.

기술적 접근법

**SCoRe**는 **REINFORCE** 기반의 **정책 경사**(policy gradient) 학습을 사용하며, **KL 발산**(KL divergence) 페널티를 도입하여 기존 모델과의 분포 차이를 제어한다.
학습은 **두 단계**로 구성된다:
**Stage I**: 기존 모델의 분포에 가까운 정책을 초기화하여 **붕괴**(behavior collapse)를 방지.
**Stage II**: **보상 형성**(reward shaping)을 통해 자기 수정 과정을 보상으로 삼아 학습.
**보상 함수**는 최종 응답의 정확도가 아닌, **수정 과정 자체**(progress)를 기준으로 정의된다.
학습 대상은 **두 번의 시도**(t1, t2)를 통해 문제를 해결하는 정책이며, **정확도@t1**, **정확도@t2**, **Δ(t1, t2)** 등의 지표를 사용하여 자기 수정 효과를 평가한다.

주요 결과

**MATH 데이터셋**에서 SCoRe는 Gemini 1.0 Pro 기반 모델의 자기 수정 성능을 **15.6%** 향상시킴.
**HumanEval 데이터셋**에서 SCoRe는 Gemini 1.5 Flash 기반 모델의 자기 수정 성능을 **9.1%** 향상시킴.
**단일 시도**(t1) 정확도는 개선되지 않았으나, **두 번째 시도**(t2) 정확도가 크게 증가함.
**Δ(t1, t2)** 지표는 SCoRe가 **자기 수정 효과**가 있음을 입증하며, **Δ(i→c)**는 새로운 문제를 해결하는 능력을 측정.

의의 및 한계

SCoRe는 기존의 자기 수정 학습이 여러 모델이나 외부 지도에 의존하는 문제를 해결하고, **단일 모델**이 **자체 생성 데이터**로 자기 수정을 학습할 수 있음을 보여준다. 이는 LLM이 **테스트 시 자기 개선**을 수행하는 능력을 향상시키는 데 중요한 기여를 한다. 또한, **보상 형성**(reward shaping)과 **온정책 학습**이 자기 수정 학습에 필수적임을 실험적으로 입증한다.

그러나 SCoRe는 **강화학습의 계산 비용**이 높고, **보상 함수 설계**에 따라 성능이 크게 달라질 수 있다는 한계가 있다. 또한, **복잡한 문제**에서는 자기 수정 과정이 제한적일 수 있으며, **다중 단계**(multi-step) 자기 수정을 확장하는 연구가 필요하다.

실용적 활용

SCoRe는 **수학적 추론**, **코딩**, **문서 생성** 등에서 모델이 자신의 오류를 인식하고 수정하는 능력을 향상시킬 수 있어, **교육**, **소프트웨어 개발**, **법률 분석** 등 여러 산업 분야에서 활용 가능하다. 특히, **외부 피드백 없이도 자기 개선**을 수행하는 모델은 **실시간 의사결정**, **자동화 시스템**, **자율 학습 시스템** 등에 유용할 수 있다.