결과 기반 강화학습으로 변환기 추론 유도 데이터가 핵심

초록

본 논문은 최종 정답만을 보상으로 주는 강화학습이 단일 층 변환기에게 체인‑오브‑생각(Chain of Thought) 형태의 중간 추론 단계를 스스로 학습하게 하는 메커니즘을 이론적으로 분석한다. 합성 그래프 탐색 과제를 통해 정책 그래디언트가 어떻게 단계별 탐색 알고리즘으로 수렴하는지 증명하고, 학습 데이터 분포에서 “단순 예시”(짧은 추론 단계가 필요한 경우)의 존재가 필수적임을 밝혀낸다. 실험은 합성 데이터와 실제 수학 추론 모델에 적용돼 이론적 예측을 검증한다.

상세 분석

논문은 먼저 결과 기반 강화학습, 즉 에피소드 전체의 최종 토큰만을 이용해 보상을 주는 설정을 정의하고, 이를 단일 층 변환기 모델에 적용한다. 모델은 입력으로 두 개의 독립적인 체인 그래프와 시작 정점을 받으며, 목표는 시작 정점이 속한 체인의 마지막 정점을 예측하는 것이다. 이 문제는 한 번에 답을 내놓을 수 없는 구조적 복잡성을 가지고 있어, 실제로는 정점‑정점 간의 전진 단계들을 순차적으로 수행하는 알고리즘이 필요하다.

저자들은 정책 그래디언트(PG) 업데이트가 어떻게 이러한 순차적 알고리즘을 선호하게 되는지를 수학적으로 증명한다. 핵심 아이디어는 정책이 최종 보상을 최대화하려면 올바른 정답에 도달하는 경로의 확률을 높여야 하는데, 이때 가장 효율적인 경로는 매 단계마다 올바른 전진 토큰을 선택하는 것이다. PG는 기대 보상의 기울기를 계산할 때 각 단계의 선택 확률에 대한 미분을 포함하므로, 전진 스텝을 지속적으로 강화한다. 반면, 뒤로 가는 단계나 체인 전환 단계는 보상에 기여하지 않으므로 그 확률은 점차 감소한다.

하지만 이러한 수렴이 보장되기 위해서는 학습 데이터 분포가 “단순 예시”에 충분한 질량을 할당해야 한다. 논문은 Q라는 분포를 도입해 시작 정점의 위치 k를 샘플링하고, k가 작을수록 필요한 전진 단계 수가 적어진다. 만약 Q가 작은 k에 충분한 확률을 부여하면, 초기 학습 단계에서 짧은 체인 탐색이 빈번히 발생하고, 이는 정책이 효율적인 전진 연산을 빠르게 학습하도록 만든다. 반대로 Q가 큰 k에만 집중하면 초기에는 긴 체인을 한 번에 맞추어야 하므로, 높은 변동성의 그래디언트가 발생하고 학습이 실질적으로 불가능해진다.

이론적 결과는 두 가지 중요한 함의를 가진다. 첫째, 정책 그래디언트는 암묵적으로 “효율성 편향”(efficient bias)을 가지고 있어, 동일한 최종 보상을 제공하는 여러 알고리즘 중에서 단계 수가 적은 알고리즘을 선호한다. 둘째, 데이터 커리큘럼이 학습 성공에 결정적이며, 특히 복잡한 예시보다 단순한 예시가 충분히 포함된 경우에만 일반화가 가능하다. 저자들은 또한 “단순 예시를 외부 분포에서 가져오는” 전략이 동일 분포 내 복잡한 예시를 학습하는 것보다 더 큰 성능 향상을 가져올 수 있음을 보인다.

실험에서는 합성 그래프 데이터와 실제 수학 문제(예: 산술 연산, 방정식 풀이)에서 이론을 검증한다. 합성 실험에서는 Q가 단순 예시에 충분한 질량을 둘 때 변환기가 정확히 전진 단계만을 사용해 체인을 탐색하고, 훈련된 모델은 훈련보다 긴 체인에도 높은 정확도로 일반화한다. 실제 언어 모델 실험에서는 Qwen 기반 모델을 수학 추론 데이터에 미세조정했을 때, 단순 문제만으로 학습한 모델이 복잡한 문제에서도 체인‑오브‑생각 형태의 중간 단계들을 생성하며 높은 정답률을 달성한다. 또한, 단순 예시를 제거하면 체인‑오브‑생각이 사라지고 성능이 급격히 저하되는 현상이 관찰된다.

결론적으로, 논문은 결과 기반 강화학습이 스파스 보상만으로도 변환기에게 체계적인 단계별 추론을 학습하게 할 수 있음을 증명하고, 이 과정에서 데이터 분포, 특히 단순 예시의 존재가 결정적인 역할을 한다는 점을 명확히 제시한다. 이는 대규모 언어 모델을 RL로 미세조정할 때 데이터 커리큘럼 설계가 얼마나 중요한지를 이론적·실험적으로 뒷받침한다.