사카나 AI, 추론 모델 증류 성능 높이는 '강화 학습 교사(RLT)' 소개

사카나 AI가 기존 증류(distillation) 방식보다 모델의 추론 능력을 크게 향상할 수 있는 새로운 학습 방식을 제안했다. 기존 방식이 학생(student) 모델의 ‘문제를 푸는 법’에 초점이 맞춰졌다면, 이번에는 교사(teacher) 모델의 ‘설명하는 법’을 강화하는 것이 핵심이다.

사카나 AI는 23일(현지시간) LLM의 추론 능력을 가르치기 위한 새로운 학습 프레임워크 ‘강화 학습 기반 교사 모델(RLT: Reinforcement-Learned Teachers)’에 관한 논문을 아카이브에 게재했다.

핵심은 학생 모델이 처음부터 문제를 해결하며 보상을 받는 것이 아니라, 교사 모델이 문제와 해답을 모두 제공받고 단계별 설명을 생성하도록 훈련하는 것이다. 즉, 교사가 더 자세하고 알아듣기 쉽게 설명하는 것이다.

'딥시크-R1'같이 증류를 활용한 추론 모델은 일반적으로 2단계 학습 방식을 따른다.

먼저 대형 교사(teacher) 모델이 RL을 통해 문제 해결 능력을 학습하고, 이후 해당 교사의 출력을 학생(student) 모델이 학습하는 방식이다.

두번째 단계는 증류, 즉 콜드 스타팅(cold-starting)이다. 그러나 기존 방식은 정답을 맞혀야만 보상을 받는 구조라 매우 느리고 비용이 많이 들어갈 뿐 아니라, 학생 모델의 학습 효과를 극대화하기 위해 출력을 필터링하고 후처리하는 추가 작업도 요구된다.

사카나 AI의 RLT 프레임워크는 이런 한계를 돌파하는 데 초점을 맞췄다.

RLT는 문제와 정답이 모두 주어지면, 교사 모델이 학생이 이해하기 쉽게 단계별 설명(Learning to Teach)을 생성하는 데 집중한다. 학생 모델은 이를 얼마나 잘 이해하는지에 따라 보상을 받는다.

이는 교사 모델의 훈련 목적을 학생 중심으로 재정렬, 기존보다 훨씬 작고 효율적인 모델도 효과적인 AI 교사로 활용될 수 있도록 한다.

또 RLT 프레임워크는 두가지 핵심 보상 지표를 중심으로 학습이 이뤄진다.

해결 점수(Solution Score)는 학생 모델이 주어진 문제와 교사 모델의 설명을 바탕으로 정답을 얼마나 정확하게 재현할 수 있는지를 평가한다. 설명 점수(Explanation Score)는 교사 모델이 생성한 설명이 논리적으로 일관되고 학생 입장에서 이해하기 쉬운지를 측정한다.

이를 통해 규모가 작은 모델도 효과적으로 학습할 수 있는 환경을 제공한다는 설명이다.

실험 결과도 주목할 만하다. 70억 매개변수의 RLT 모델은 6710억 매개변수의 딥시크-R1보다 더 뛰어난 학습 성능을 보였다.

7B 크기의 학생 모델을 훈련하는 경우, RLT를 사용한 성능은 26.3%, R1은 18.9%에 그쳤다. 또 32B 크기 학생 모델을 훈련하면 RLT가 37.6%, R1이 34.4%의 결과를 냈다.

사카나 AI는 70억 매개변수 RLT 모델이 'AIME 2024' 'MATH 500' 'GPQA 다이아몬드' 등 벤치마크에서 딥시크-R1과 '비스포크-7B', 심지어 대형 모델의 후처리 RL 결과물보다도 우수한 성능을 발휘했다고 밝혔다.

RLT는 RL의 초기 단계인 부트스트랩(Cold Start) 과정에서도 높은 효과를 나타냈다. 기존의 RL 기반 모델이 생성한 설명보다 RLT가 생성한 설명 데이터를 활용했을 때, 학생 모델의 학습 성과가 더 뛰어났다.

특히 제로샷(Zero-Shot) 전이 성능에서도 강점을 보였는데, 수학 문제에서 훈련된 RLT의 설명 능력이 전혀 다른 과제인 산술 기반 '카운트다운(Countdown)' 문제에서도 효과적으로 적용됐다. 이는 문제를 직접 푸는 능력보다 문제 해결 과정을 설명하는 능력이 더 높은 전이 가능성을 지닌다는 점을 보여준다.

사카나AI는 이번 연구의 모델과  코드를 각각 허깅페이스와 깃허브에 공개했다.

박찬 기자 cpark@aitimes.com