🤖 AI

Reinforcement-aware Knowledge Distillation for LLM Reasoning

이 논문은 분포 불일치와 목적 함수 간섭을 극복하기 위해 신뢰 영역 비율 증류(Trust Region Ratio Distillation, TRRD) 목적 함수를 통해 선택적 모방을 강화 학습에 통합하는 방식인 RL-aware Distillation(RLAD)을 제안하며, 이를 통해 더 작은 언어 모델이 탐색과 활용의 균형을 맞추면서 더 큰 교사 모델로부터 긴 사고 사슬(chain-of-thought) 추론 능력을 효과적으로 상속받을 수 있도록 한다.

원저자: Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia, Stefano Soatto

게시일 2026-06-19

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia, Stefano Soatto

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신에게는 아주 뛰어나지만 믿을 수 없을 정도로 비싼 마스터 셰프(스승)가 있다고 상상해 보십시오. 이 스승은 복잡하고 다채로운 코스 요리를 완벽한 추론과 함께 만들어낼 수 있습니다. 당신은 이 마스터처럼 요리할 수 있는 더 작고, 빠르고, 저렴한 견습 요리사(제자)를 가르쳐서, 큰 비용을 들이지 않고도 더 많은 사람에게 훌륭한 음식을 대접하고자 합니다.

오랫동안 견습생을 가르치는 방법은 간단했습니다. 바로 마스터의 레시피 북을 복사하는 것이었습니다. 견습생은 그저 과거에 마스터가 수행했던 단계들을 그대로 암기하기만 하면 되었습니다. 이 방식은 어느 정도 효과가 있지만, 경직되어 있습니다. 만약 견습생이 약간 다른 요리를 시도하거나 새로운 식재료를 마주하게 되면, 단순히 과거의 기록을 맹목적으로 따르고 있을 뿐이라서 어떻게 '요리에 대해 생각해야 하는지' 배우지 못했기 때문에 막혀버리게 됩니다.

최근에는 셰프들이 새로운 방법을 사용하기 시작했습니다: 피드백을 통한 시행착오(Trial and Error with Feedback). 견습생은 요리를 시도하고, 그 요리가 얼마나 맛있는지에 대한 점수(보상, Reward)를 받은 뒤, 다음번에 더 높은 점수를 얻기 위해 자신의 기술을 조정합니다. 이 방법은 훌륭하지만, 느리고 비용이 많이 듭니다.

문제는 이 두 가지 방법을 결합하려고 할 때 발생합니다. 만약 당신이 견습생에게 "마스터의 단계를 복사하면서 동시에 높은 맛 점수를 얻도록 해라"라고 명령하면, 그들은 종가 혼란에 빠집니다.

때로는 마스터의 예전 단계들이 현재 상황에서 최고의 맛 점수로 이어지지 않을 수도 있습니다.
"복사하라"는 지시와 "높은 점수를 얻으라"는 지시가 서로 충돌하여, 견습생이 갈팡질팡하며 학습 효율이 떨어지게 됩니다.

새로운 솔루션: "스마트 모방" (RLAD)

논문의 저자들은 RLAD(Reinforcement-Aware Distillation, 강화학습 인지 증류)라고 불리는 새로운 교육 방식을 제안합니다. 단순히 견습생에게 항상 마스터를 100% 복사하라고 강요하는 대신, "스마트 모방" 규칙을 도입합니다.

GPS 비유를 통한 핵심 아이디어는 다음과 같습니다:

목표: 견습생은 가장 높은 "보상"(예: 최고의 전망)을 주는 목적지로 운전해서 가고 싶어 합니다.
기존 방식 (표준 증류): GPS(스승)는 도로가 막혀 있거나 더 좋은 경로가 존재하더라도 끊임없이 "좌회전하세요! 좌회전하세요!"라고 외칩니다. 견습생은 막다른 길로 향하게 되더라도 맹목적으로 따라갑니다.
새로운 방식 (RLAD): GPS는 자신의 현재 계획이 더 좋은 전망을 향해 나아가는 방향과 일치할 때만 지침을 내립니다.
- 견습생이 멋진 전망을 향해 운전하고 있고, GPS가 "네, 그 방향으로 회전하는 것이 좋습니다"라고 말한다면, 견습생은 GPS를 밀접하게 따릅니다.
- 견습생이 멋진 전망을 향해 운전하고 있지만, GPS가 "아니요, 우회전하세요"라고 말한다면, 견습생는 "맛 점수"(보상)가 현재의 경로가 더 낫다고 말해주기 때문에 GPS의 말을 무시합니다.
- 견습생이 길을 잃었을 때(낮은 보상), GPS가 개입하여 안전하고 알려진 경로로 돌아오도록 안내합니다.

비밀 소스: "신뢰 구역" (TRRD)

이 방식이 견습생을 미치게 만들지 않도록 하기 위해, 논문에서는 TRRD(Trust Region Ratio Distillation, 신뢰 영역 비율 증류)라는 기술을 사용합니다.

이것은 견습생에게 연결된 안전 목줄이라고 생각하면 됩니다.

이 목줄은 견습생이 '방금 전까지 있었던 위치'와 '마스터가 갔을 법한 위치'의 혼합물에 고정되어 있습니다.
견습생은 새로운 경로를 탐색(Exploration)하거나 자신이 아는 것에 집중(Exploitation)하며 자유롭게 달릴 수 있습니다.
하지만 견습생이 마스터의 지혜로 정의된 "안전 구역"에서 너무 멀리 벗어나려 하면, 목줄이 부드럽게 잡아당깁니다.
결정적으로, 이 목줄은 마스터의 조언이 실제로 견습생이 더 높은 점수를 얻는 데 도움이 될 때만 조여집니다. 만약 현재 상황에서 마스터가 틀렸다면, 목줄은 느슨하게 유지되어 견습생이 더 나은 길을 찾을 수 있도록 해줍니다.

연구 결과는 어떠했는가?

연구진은 이를 두 가지 유형의 "요리 도전 과제"에 대해 테스트했습니다:

논리 퍼즐: 복잡한 미스터리를 풀거나 물류 문제를 해결하는 것과 같은 과제.
수학 문제: 어려운 방정식이나 경시대회 수학 문제를 푸는 것과 같은 과제.

결과:

더 높은 점수: 이 새로운 "스마트 모방" 방식으로 훈련된 견습생들은 단순히 마스터를 복사하거나 스스로 추측하며 학습한 모델들보다 현저히 높은 점수를 받았습니다.
어려운 문제일수록 효과적: 개선 폭은 가장 어려운 문제에서 가장 컸습니다. 쉬운 문제에서는 모두가 잘 해냈지만, "불가능한" 수준의 문제에서는 새로운 방식이 빛을 발했습니다.
안정성: 훈련 과정이 훨씬 매끄러웠습니다. 기존 방식은 견습생을 갈팡질팡하게 만들고 진전된 성과를 잃게 만드는 경우가 많았지만, 새로운 방식은 정상을 향한 안정적인 경로를 유지했습니다.
진정한 학습 vs 단순 복사: 기존 방식은 주로 견습생이 마스터의 특정 정답을 암기하게 만들었습니다(많은 답을 찍어내는 데는 좋지만, 최선의 단 하나의 답을 찾는 데는 나쁨). 반면, 새로운 방식은 첫 시도에 최선의 답을 찾아내는 견습생의 능력을 실제로 향상시켰습니다.

요약하자면

이 논문은 작은 AI 모델에게 추론하는 법을 가르치는 더 스마트한 방법을 소개합니다. 단순히 큰 규모의 똑똑한 스승을 맹목적으로 복사하도록 강요하는 대신, 스승의 조언이 승리에 도움이 될 때만 스승의 말을 듣도록 가르칩니다. 이를 통해 학생은 똑똑하면서도(스승처럼), 적응력이 뛰어난(더 나은 해결책을 찾을 수 있는) 모델이 되며, 이 모든 과정을 더 빠르고 안정적으로 학습할 수 있습니다.

기술 요약: LLM 추론을 위한 강화학습 인지 지식 증류 (RLAD)

1. 문제 정의

대규모 언어 모델(LLM)은 강화학습(RL) 사후 학습, 특히 긴 사고 사슬(Chain-of-Thought, CoT) 작업에서 상당한 추론 능력을 달성했습니다. 그러나 이러한 대형 모델의 높은 추론 비용으로 인해 더 작은 학생 모델로의 지식 증류가 필수적입니다. 기존의 지식 증류(Knowledge Distillation, KD) 방법들은 RL 사후 학습에 적용할 때 다음과 같은 결정적인 한계에 직면합니다:

분포 불일치(Distribution Mismatch): 표준 KD 방법(오프라인, 오프-폴리시 또는 온-폴리시)은 종종 고정된 교사 궤적이나 정적인 KL 발산 정규화에 의존합니다. RL에서 학생의 정책은 지속적으로 진화하며, 이로 인해 롤아웃 분포가 이동하게 됩니다. 정적인 증류 타겟은 학생이 교사의 초기 분포에서 멀어짐에 따라 최적이 아니게 되거나 불안정해집니다.
목적 함수 간섭(Objective Interference): KD를 RL 루프에 통합하는 방법(예: KDRL)은 일반적으로 학생과 교사 사이의 KL 발산을 정규화 페널티로서 추가합니다. 이는 보상 최대화 목적과 모방 목적이 서로 경쟁하는 다중 목적 최적화 문제를 생성합니다. 이러한 경쟁은 세심한 하이퍼파라미터 튜닝(손실 균형)을 요구하며, 교사의 감독이 학생의 현재 어드밴티지 추정치와 일치하지 않을 수 있기 때문에 불안정한 학습 역학을 초래할 수 있습니다.

2. 방법론: 강화학습 인지 증류 (RLAD)

저자들은 RL 사후 학습 중에 **선택적 모방(selective imitation)**을 수행하도록 설계된 프레임워크인 RLAD를 제안합니다. 핵심 원리는 학생이 RL 목적 함수 하에서 현재 정책 업데이트를 개선할 때만 교사에 의해 유도되어야 한다는 것입니다.

핵심 구성 요소: 신뢰 영역 비율 증류 (Trust Region Ratio Distillation, TRRD)
RLAD는 독립적인 KL 정규화 도구 대신, 그룹 상대 정책 최적화(GRPO)의 표준 중요도 샘플링 비율을 교사-학생 혼합물로 대체하는 TRRD를 도입합니다.

메커니즘: TRRD는 $i$ 번째 응답의 $t$ 번째 토큰에 대해 새로운 가능도 비율 $r^{TRRD}_{i,t}$ 를 정의합니다:
$r^{TRRD}_{i,t}(\theta_S) = \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_S,old}(y_t | x, y_{<t})} \right)^\alpha \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_T}(y_t | x, y_{<t})} \right)^{1-\alpha}$
여기서 $\pi_{\theta_S}$ 는 현재 학생 정책, $\pi_{\theta_S,old}$ 는 이전 학생 정책, $\pi_{\theta_T}$ 는 교사 정책입니다. $\alpha \in [0, 1]$ 는 혼합 계수입니다.
신뢰 영역 해석: 이 비율은 신뢰 영역이 단순히 이전 학생뿐만 아니라 이전 학생과 교사의 혼합물( $r_{\pi_{mix}}$ )에 닻을 내리는 PPO 스타일의 업데이트로 볼 수 있습니다.
어드밴티지 인지 선택 (Advantage-Aware Selection): 업데이트는 그룹 상대 어드밴티지 $\hat{A}_{i,t}$ $\hat{A}_{i, t}$ 에 의해 가중치가 부여됩니다.
- 양의 어드밴티지 (Positive Advantage): 어드밴티지가 양수이면 학생은 해당 토큰의 확률을 높이도록 권장됩니다. 만약 교사 또한 높은 확률을 할당한다면, 혼합 앵커(mixture anchor)가 높아져 학생이 클리핑 경계(clipping bound)에 도달하기 전까지 더 큰 업데이트를 허용하게 됩니다(더 큰 업데이트 허용). 만약 교사가 낮은 확률을 할당한다면, 앵커가 낮아져 업데이트 크기가 제한됩니다.
- 음의 어드밴티지 (Negative Advantage): 어드밴티지가 음수이면 학생은 해당 토큰의 확률을 낮추도록 권장됩니다. 만약 교사가 높은 확률을 할당한다면, 클리핑 경계가 학생이 토큰을 공격적으로 억제하는 것을 방지하여 교사와 정렬된 행동을 보존합니다.
- 제로에 가까운 어드밴티지 (Near-Zero Advantage): 보상 신호가 유익하지 않을 때 교사의 영향력이 무시되어 간섭을 방지합니다.

이 설계는 교사의 가이드가 보상 신호(어드밴티지)에 의해 조절되고 신뢰 영역 내로 제한되도록 보장함으로써 탐색, 착취, 모방 사이의 균형을 암시적으로 조절합니다.

3. 주요 기여

RLAD 프레임워크: 정적인 KL 정규화 대신 클리핑 및 어드밴티지 가중 가능도 비율 목적 함수를 통해 교사 가이드를 통합하는 TRRD를 사용하는 RL 사후 학습용 신규 증류 프레임워크입니다.
이론적 분석: TRRD가 보상 조건부의 선택적 모방을 가능하게 함을 입증하는 분석을 제공합니다. 이는 신뢰 영역 경계가 있는 업데이트를 보존하며, 세밀한 손실 균형 없이도 보상 최적화와 교사 추종 사이의 트레이드오프를 자연스럽게 조절합니다.
실증적 검증: 논리적 추론(K&K Logistics) 및 긴 문맥 수학 추론(AIME, AMC, MATH)에 대한 광범적인 실험을 통해 베이스라인 대비 일관된 개선을 보여주었습니다.

4. 실험 결과

저자들은 0.6B에서 8B 파라미터 규모의 학생 모델과 최대 32B 규모의 교사를 사용하여 논리 추론 및 복잡한 수학 추론 작업에 대해 RLAD를 평가했습니다.

논리 추론 (K&K Logistics):
- Qwen3-8B 교사를 사용할 때, RLAD는 표준 GRPO와 비교하여 Qwen3-0.6B(8K 컨텍스트에서 0.76에서 0.94로) 및 Qwen3-1.7B(8K 컨텍스트에서 0.95에서 0.99로)의 평균 정확도를 크게 향상시켰습니다.
- 이득은 가장 어려운 서브셋(PPL7 및 PPL8)에서 가장 두드et게 나타났으며, 여기서 RLAD는 KDRL에 대해 명확한 격차를 유지했습니다.
- 학습 역학은 GRPO 및 KDRL 모두에 비해 더 빠른 수렴 속도와 더 높은 최종 보상을 보여주었습니다.
수학 추론:
- 베이스 모델 (R1-Zero 스타일): Qwen3-8B-Base에서 RLAD는 긴 문맥(30K) 작업에 대해 GRPO보다 평균 성능을 +5.5 포인트 향상시켰습니다.
- 사후 학습 모델: RLAD는 명시적 KL 증류(KDRL) 및 오프라인 증류보다 우수한 성능을 보였습니다.
- 지표 특이사항: RLAD는 KDRL(KDRL은 더 다양한 교사 솔루션을 순수하게 모방하여 상대적으로 더 큰 이득을 보이는 경향이 있음)에 비해 Pass@1 (Mean@32)에서 더 큰 이득을 보였습니다. 저자들은 이를 RLAD가 성공적으로 보상 정렬된 정책 개선(착취)을 이끌어낸 반면, KDRL의 이득은 순수한 교사 솔루션 모방에 의해 주도되었다고 해석합니다.
- 안정성: RLAD는 목적 함수 결합 문제로 어려움을 겪은 KDRL에 비해 진동이 적고 더 안정적인 학습 역학을 보였습니다.
강건성 및 효율성:
- 약한 교사 (Weak Teacher): 교사가 학생보다 약할 때, KDRL은 성능이 크게 저하된 반면 RLAD는 GRPO 성능에 근접하게 유지되었습니다.
- 하이퍼파라미터: RL-AD는 KL 가중치 $\lambda$ 에 민-감한 KDRL과 달리, 혼합 계수 $\alpha$ (0.5로 고정)에 대해 매우 둔감한 것으로 나타났습니다.
- 효율성: RLAD를 통합하면 표준 GRPO 대비 학습 지연 시간이 약 12%만 증가하는데, 이는 교사가 궤적 생성이 아닌 로그 확률 계산에만 사용되기 때문입니다.

5. 의의 및 주장

본 논문은 RLAD가 정적인 교사 감독과 동적인 RL 정책 업데이트 사이의 근본적인 불일치를 해결한다고 주장합니다. 교사 가이드를 신뢰 영역 업데이트 메커니즘에 포함시킴으로써, RLAD는 다중 목적 KL 정규화에 필요한 "취약한 튜닝(brittle tuning)"을 피합니다.

저자들은 RLAD가 학생이 앵커 분포 내에서 높은 어드밴티지를 가진 행동을 향해 이동하면서도 교사의 가이드를 기회적으로 활용할 수 있게 한다고 단언합니다. 이는 다음과 같은 결과를 가져옵니다:

우수한 성능: 도전적인 벤치마크에서 GRPO, KDRL 및 오프라인 증류를 일관되게 능가합니다.
안정성: 더 안정적인 학습 역학 및 더 나은 수렴을 제공합니다.
진정한 RL 최적화: 모방에 크게 의존하는 방식(Pass@32를 높일 수는 있지만 Pass@1을 높이지는 못함)과 달리, RLAD는 보상 정렬된 행동을 우선시하여 모델의 가장 확률 높은 정답(Pass@1)에서 더 강력한 이득을 이끌어냅니다.

결론적으로, 이 접근 방식은 로짓 접근을 위해 잘 보정된 오픈 웨이트 교사 모델에 의존하지만, RL 사후 학습의 동적인 특성에 대한 증류를 가능하게 하는 중요한 진전임을 시사합니다.