THOR: 도구 통합 계층 강화학습으로 수학 추론 능력 극대화

초록

THOR는 다중 에이전트 기반 데이터 생성 파이프라인(TIRGen)과 계층형 강화학습을 결합해, LLM이 외부 코드 도구를 효율적으로 활용하도록 학습한다. 단계별 코드 실행 성공 여부를 중간 보상으로 사용하고, 실행 실패 시 즉시 자기 교정 메커니즘을 적용해 추론 과정을 동적으로 수정한다. 이 접근법은 다양한 수학·코드 벤치마크에서 동등 규모 모델 대비 최신 성능을 달성한다.

상세 분석

본 논문은 LLM이 고정밀 수학 연산과 형식적 기호 조작에 한계를 보이는 문제를 도구 통합(Tool‑Integrated Reasoning, TIR) 방식으로 해결하고자 한다. 핵심 기여는 세 가지로 요약된다. 첫째, TIRGen 파이프라인은 ‘Generator‑Refiner’ 구조를 채택한다. Generator는 자연어 기반 사고 단계만을 생성하고, Refiner는 해당 단계에서 코드로 구현 가능한 부분을 식별·변환한다. 변환된 코드는 샌드박스 실행기(S)에서 즉시 실행되어 관찰값(oₜ)을 반환하고, 이 관찰값은 다시 Generator에게 피드백으로 제공된다. 이렇게 순환적인 사고‑행동‑관찰 루프를 통해 생성된 데이터는 정책 모델(πθ)의 사고 스타일과 일치하므로 ‘policy‑aligned’ 특성을 갖는다. 또한, 대규모 LLM(GPT‑4o 등)에 의존하지 않고, Generator와 Refiner 각각이 상대적으로 간단한 능력만 요구되므로 데이터 구축 비용이 크게 감소한다.

둘째, 계층형 강화학습(Hierarchical RL) 설계는 에피소드‑레벨 최적화와 스텝‑레벨 최적화를 동시에 수행한다. 에피소드‑레벨에서는 GRPO(Generalized Reward‑Based Policy Optimization)를 이용해 최종 정답의 정확도를 직접 보상으로 사용한다. 스텝‑레벨에서는 각 코드 호출의 성공 여부를 별도 보상으로 정의함으로써, ‘중간 도구 호출 성공 → 최종 정답 정확도’라는 경험적 인사이트를 정량화한다. 이는 전통적인 에피소드‑단위 RL이 겪는 희소 보상 문제를 완화하고, 코드 생성 능력을 미세하게 조정할 수 있게 한다.

셋째, 추론 단계에서의 자기 교정(Self‑Correction) 메커니즘은 실행 실패를 감지하면 즉시 이전 사고 단계로 되돌아가 대안을 탐색한다. 이는 도구 호출이 실패했을 때 단순히 오류를 반환하는 기존 방식과 달리, LLM이 자체적으로 오류 원인을 추론하고 수정하도록 유도한다. 실험 결과, 이 메커니즘은 특히 복잡한 다중 단계 문제에서 성공률을 크게 끌어올렸다.

전체 실험에서는 Math500, AIME 2024/2025, AMC, Minerva Math, Olympiad Bench 등 다양한 수학 벤치마크와 HumanEval, MBPP, LiveCodeBench 등 코드 생성 벤치마크를 사용하였다. 동일 규모 모델(7B) 기준, THOR는 기존 Tool‑Integrated 방법들보다 평균 4~7%p 높은 정확도를 기록했으며, 비추론 모델에도 적용 가능함을 보였다. 또한, 추론 시 오버헤드가 크게 증가하지 않아 실용성도 확보했다.

이 논문은 (1) 정책 정렬된 고품질 TIR 데이터 자동 생성, (2) 중간 도구 호출 성공을 보상으로 활용한 계층형 RL, (3) 실행 피드백 기반 실시간 자기 교정이라는 세 축을 통해 LLM의 수학·코드 복합 추론 능력을 체계적으로 향상시켰다는 점에서 의미가 크다. 향후 연구는 더 다양한 도구(예: 정리 증명 엔진)와의 연계, 그리고 멀티모달 환경에서의 확장 가능성을 탐색할 여지가 있다.