Feedback Distillation 소개
처음 어려운 수학 문제를 풀 때를 떠올려 봅시다. 정답만 통째로 외우기보다는, 선생님이 "이 부분에서 절댓값을 경우로 나눠봐", "여기서는 분모를 없애는 게 낫다" 같은 짧은 조언을 줄 때 우리는 훨씬 빠르게 성장합니다. 이 논문은 바로 이 직관을 LLM 사후학습(post-training)에 옮긴 연구입니다. 즉 Feedback Distillation 은 모델이 스스로 만든 풀이에 대해 또 다른 LLM이 남긴 피드백을 받아, 그 피드백이 반영된 분포를 토큰 단위로 자기 가중치에 증류(distillation)하는 학습 기법입니다. 저자들은 이 방법을 Lean 4 정리 증명(theorem proving) 환경에서 검증하여, 표준 강화 학습 기법인 GRPO보다 더 높은 다양성과 더 나은 test-time 확장성을 얻을 수 있음을 보였습니다.
이 연구는 Meta의 FAIR(Fundamental AI Research) 와 Inria, Sorbonne Université 등의 연구진이 함께 수행했습니다.
추론 모델 사후학습의 현재 지형과 GRPO의 한계
최근 인공지능은 코드 생성, 수학 경시 문제 풀이처럼 복잡한 추론(reasoning) 과제에서 빠르게 발전하고 있습니다. 그중에서도 수학적 추론은 탐색 공간이 방대하고 정답 여부가 명확하다는 점 때문에 추론 연구의 핵심 시험대가 되었습니다. 특히 Lean 같은 증명 검증기(verifier)가 출력을 자동으로 채점해주는 검증 가능한 보상(Verifiable Rewards) 환경에서는, 전문가 풀이로 지도 미세조정(SFT, Supervised Fine-Tuning)을 한 뒤 강화 학습(RL, Reinforcement Learning)으로 검증기 보상을 최대화하는 파이프라인이 사실상 표준으로 자리잡았습니다. 이때 쓰이는 RL 알고리즘이 바로 DeepSeekMath에서 제안된 GRPO 로, 하나의 문제에 대해 여러 출력을 뽑은 뒤 그 보상들을 서로 비교해 정책(policy)을 갱신합니다.
하지만 이 SFT와 GRPO를 결합한 패러다임에는 분명한 한계가 있습니다. 첫째, SFT 단계는 미세조정용 데이터셋을 만드는 비용과 난이도에 발목이 잡히고, off-policy 시연 데이터로 학습하기 때문에 좁은 미세조정 분포를 벗어나면 치명적 망각(catastrophic forgetting)을 일으킬 수 있습니다.
둘째, RL 학습이 새로운 추론 능력을 만들어낸다는 통념과 달리, 최근 연구들은 GRPO가 모델 답변의 다양성을 키우거나 탐색(exploration)을 촉진하지 않고 오히려 RL 이전부터 이미 존재하던 고보상 모드(mode)로 붕괴(collapse)한다는 증거를 제시합니다. 저자들은 이를 두고 "GRPO는 이미 베이스 모델에 있던 능력을 증폭할 뿐, 그 너머로 나아가지 못한다" 는 흐름의 연구들을 인용합니다. 이는 Lean의 Mathlib처럼 300{,}000 개가 넘는 보조정리(lemma) 라이브러리를 폭넓게 탐색해야 성공할 수 있는 형식 수학(formal mathematics)에서 특히 치명적입니다.
셋째, GRPO는 궤적(trajectory) 전체에 대해 하나의 보상만 주는 trajectory-level 보상이라서, 샘플이 전부 실패하면 학습 신호가 아예 없습니다. 좋은 토큰과 나쁜 토큰을 구분하지 못하고 보상을 궤적 전체에 균일하게 흩뿌리는 셈입니다.
Feedback Distillation의 핵심 아이디어
이러한 문제들을 해결하기 위해 저자들이 출발점으로 삼은 것은 최근 제안된 자기 증류(Self-Distillation) 계열의 알고리즘입니다. 자기 증류는 모델이 특권 정보(privileged information)를 받았을 때의 출력을 다시 자기 자신이 모방하도록 학습시켜, 프롬프팅 전략 자체를 모델 가중치 안으로 녹여 넣는 방식입니다. 이 방법은 보상을 궤적 전체에 균일하게 분배하는 대신 토큰 단위로 세밀하게 credit assignment를 제공한다는 점에서 매력적입니다.
이 논문의 발상 전환은 그 특권 정보를 무엇으로 둘 것인가 에 있습니다. 기존 자기 증류 연구들은 특권 정보로 환경의 원시 출력(Hübotter et al.)이나 정답 풀이(Shenfeld et al.)를 사용했습니다. 반면 저자들은 특권 정보를 모델의 시도에 대해 LLM이 생성한 피드백 으로 둡니다. 모델이 어떤 증명을 시도하면, 피드백 모델이 그 시도와 검증기 출력을 보고 "X를 해라", "Y를 하지 마라" 같은 실행 가능한 조언을 만들어내고, 학생 모델은 이 조언을 받은 버전의 분포를 모방하도록 학습합니다. 피드백이 더 강한 모델에서 나올 수 있으므로 저자들은 이를 Feedback Distillation 이라 부르며, 피드백 모델이 학습 대상 모델의 동결(frozen) 복사본인 경우는 Self-Feedback Distillation 이라 부릅니다.
위 그림은 이 아이디어를 한눈에 보여줍니다. MiniF2F 데이터셋의 Lean 문장에 대해 Qwen3-8B 학생이 증명을 시도하면, Claude가 피드백을 제공합니다. 그 뒤 학생의 답안을 토큰별로 \log p_{\text{teacher}} - \log p_{\text{student}} 값에 따라 색칠합니다. 초록색은 교사가 선호하는 토큰, 빨간색은 학생이 선호하는 토큰입니다. 예시에서 학생은 "π"라고 썼지만 피드백을 반영한 교사는 같은 자리에 98\% 확률로 "Real"을 쓰려고 합니다. KL 손실을 통해 학생은 이 위치에서 "Real"의 확률을 높이도록 학습되고, 결과적으로 더 나은 Lean 코드를 작성하게 됩니다.
핵심 결과를 미리 말하자면, Qwen3.5-9B를 LeanWorkbook으로 학습했을 때 GRPO 단독은 LeanWorkbook 테스트셋에서 59\% 의 pass@1 정확도에 그친 반면, Feedback Distillation으로 초기화한 뒤 GRPO를 이어 붙이면 75\% 까지 올라갔습니다.
Feedback Distillation의 작동 원리
학생과 교사, 그리고 특권 정보
Feedback Distillation은 가중치가 독립적인 동일 모델의 두 인스턴스를 사용합니다. 하나는 학습 대상인 학생(student) \pi_{\theta} 이고, 다른 하나는 프롬프트에 특권 정보를 받는 교사(teacher) 입니다. 둘 다 같은 사전학습 모델 \pi 에서 출발합니다. 학습에 쓰는 데이터셋 D 는 풀이가 없는 문제들의 모음(예: 증명이 없는 Lean 문장)입니다.
문제 x 에 대해 학생이 증명 시도 y 를 생성하면, 교사는 다음과 같이 정의됩니다.
여기서 F 는 시도 y 를 받아 텍스트 피드백을 돌려주는 함수, 즉 특권 정보입니다. 저자들은 일관성을 위해 \pi'_{\mu,y} 를 교사라 부르지만, 사실 "증강된 학생(augmented student)" 이라고 부르는 편이 더 정확하다고 솔직하게 밝힙니다. 교사는 결국 학생과 같은 모델에 피드백만 추가로 받은 존재이기 때문입니다.
토큰 단위 KL 손실
학습 손실은 학생이 생성한 궤적을 따라가며 교사 분포와 학생 분포 사이의 토큰별 KL 발산(Kullback-Leibler divergence)을 측정합니다.
여기서 y 의 샘플링 자체에는 그래디언트가 흐르지 않습니다(stop-gradient). 교사 분포가 \theta 에 의존하지 않으므로 KL 발산과 교차 엔트로피(cross-entropy)는 \theta 에 대한 그래디언트가 동일합니다. 따라서 실제 구현에서는 다음의 등가 손실을 사용합니다.
이 손실의 의미는 직관적입니다. 학생이 직접 만든 토큰 시퀀스 y 의 각 위치에서, 피드백을 본 교사라면 어떤 토큰에 얼마의 확률을 줄 것인가 를 목표로 삼아 학생의 분포를 끌어당기는 것입니다. GRPO처럼 궤적 전체에 보상 하나를 주는 대신, 모든 토큰 위치마다 풍부한 분포 형태의 지도(supervision)를 제공한다는 점이 가장 큰 차이입니다.
EMA 교사 갱신과 Top-K 절단
교사를 학습 내내 동결해두면 문제가 생깁니다. 모델이 여러 피드백을 거치며 학습할수록, 앞선 피드백에서 얻은 지식을 교사는 내재화하지 못한 채 낮은 확률을 매기게 되고, 결국 학생이 어렵게 익힌 내용에 벌점을 주는 셈이 됩니다. 저자들은 이를 막기 위해 교사 가중치를 지수 이동 평균(EMA, Exponential Moving Average) 으로 갱신합니다. 다섯 번의 그래디언트 스텝마다 \mu \leftarrow \alpha\,\mu + (1-\alpha)\,\theta 로 교사가 학생을 천천히 뒤따르게 만드는 것입니다. 이는 자기지도 학습(self-supervised learning)에서 움직이는 목표(moving target)로 인한 표현 붕괴를 막기 위해 BYOL이나 MoCo 같은 연구들이 EMA 목표 네트워크를 쓰는 것과 같은 맥락입니다.
또한 계산 비용을 줄이기 위해 어휘(vocabulary) 전체에 대한 합을 교사 \pi'_{\mu,y} 기준 확률 상위 K 개 토큰으로 절단(truncation)하며, 실제로는 K=25 를 사용합니다.
Lean 정리 증명 실험
Lean 4 환경과 도구 사용
실험 환경은 Lean 4 형식 수학 환경입니다. 저자들은 이 환경이 복잡 추론용 RL을 연구하기에 적합한 이유를 세 가지로 꼽습니다. 첫째, 베이스 모델인 Qwen3-8B와 Qwen3.5-9B는 초기 상태에서 MiniF2F 기준 각각 2\%, 11\% 밖에 풀지 못할 만큼 Lean에 서툴기 때문에, 의미 있는 향상은 곧 새로운 능력의 발견을 뜻합니다. 둘째, 모델은 적절한 tactic을 호출하고 300{,}000 개 이상의 선언(declaration)을 담은 Mathlib에서 보조정리를 참조해야 하므로 탐색 공간이 매우 넓습니다. 셋째, Lean 검증기가 증명의 정답 여부를 확정적으로 알려줍니다.
도구 사용(tool use)은 최신 형식 수학 모델의 핵심 요소입니다. 저자들은 모델에 세 가지 도구를 제공합니다.
lean_write_file: Lean 소스 파일을 작성합니다.lean_check_file: 작성한 파일을 컴파일하여 컴파일러 출력(오류, 경고, 성공)을 반환합니다.rg_in_mathlib: 정규식으로 Mathlib을 검색해 관련 보조정리를 찾습니다.
모델이 도구를 호출하면 생성이 멈추고, 도구가 실행된 뒤 그 출력이 맥락에 덧붙여진 다음 생성이 재개됩니다. 부록 실험에서 도구 사용이 양쪽 학습법 모두의 성능을 끌어올린다는 점도 확인되었습니다.
데이터셋과 피드백 모델
학습에는 LeanWorkbook에서 뽑은, 증명이 알려진 10{,}000 개의 문장을 사용합니다. 평가에는 LeanWorkbook 테스트셋 256 개 문장과 경시 수준 문제 244 개로 이루어진 MiniF2F 테스트 분할을 씁니다. 각 시도는 최대 8{,}192 토큰의 단일 생성으로 이루어지며, 그 안에서 모델은 여러 번 도구를 호출할 수 있습니다. 실험은 8장의 NVIDIA H200 GPU 한 노드에서 수행되었고, 학습 코드베이스로는 verl을, Lean 증명 채점에는 Kimina Lean Server를 사용했습니다. 한 스텝당 GRPO는 평균 550 초, Feedback Distillation은 300 초가 걸립니다.
피드백 모델로는 Claude Opus 4.6 을 사용합니다. 피드백 모델은 형식 문장 x, 모델의 시도 y(모든 도구 호출과 그 출력 포함), 그리고 최종 Lean 컴파일러 출력을 받습니다. 그리고 증명 전략의 오류를 짚거나, tactic 또는 보조정리의 오용을 지적하거나, 더 나은 도구 사용을 권하는 "X를 해라" 또는 "X를 하지 마라" 형태의 실행 가능한 지시 1개에서 3개를 생성하도록 프롬프트됩니다. 이 피드백은 같은 문제에 대한 새로운 시도에 도움이 되도록 작성됩니다.
실험 결과 및 성능 분석
저자들은 GRPO, Feedback Distillation, 그리고 GRPO over Feedback Distillation(Feedback Distillation 체크포인트에서 GRPO를 이어 학습) 세 가지를 비교합니다. 결과를 관통하는 핵심 메시지는, Feedback Distillation이 생성 궤적의 다양성을 유지하면서도 효율적으로 학습한다는 것입니다.
GRPO over Feedback Distillation이 가장 강력하다
저자들은 Feedback Distillation을 200, 300, 400 스텝까지 진행한 체크포인트에서 GRPO를 이어 학습했습니다. 위 그래프에서 보이듯 이 조합은 GRPO 단독을 뚜렷하게 능가합니다. 앞서 언급한 대로 Qwen3.5-9B 기준 LeanWorkbook 테스트 정확도는 GRPO 단독 59\% 에서 Feedback Distillation과 GRPO의 결합으로 75\% 까지 향상되었습니다. 두 방법이 서로 보완적인 이점을 제공한다는 의미입니다.
엔트로피와 pass@k 확장성
Feedback Distillation은 학습 내내 GRPO보다 훨씬 높은 정책 엔트로피(policy entropy)를 유지합니다. 이는 정책이 더 많은 다양성을 보존한다는 뜻입니다. 그리고 이 다양성은 더 나은 test-time 확장으로 이어집니다.
위 그림 오른쪽은 GRPO와 Feedback Distillation이 비슷한 pass@1을 보이는 200 스텝 체크포인트에서 샘플 수 k 를 늘렸을 때의 pass@k 를 비교합니다. 두 방법의 pass@1은 비슷하지만, k 가 커질수록 Feedback Distillation이 훨씬 잘 확장됩니다. 이는 높은 엔트로피가 단순한 균일 잡음이 아니라 실제 답변의 다양성을 반영한다는 점을 확인해 줍니다.
사용 tactic 수와 보조정리 발견
다양성은 모델이 성공적인 증명에서 새로운 Lean tactic과 Mathlib 보조정리를 얼마나 빨리 동원하는가로도 드러납니다. 위 그림 왼쪽처럼, Feedback Distillation은 GRPO에 비해 서로 다른 tactic과 보조정리를 더 빠르고 지속적으로 발견합니다. 넓은 라이브러리 탐색이 성패를 가르는 형식 수학에서 이는 중요한 이점입니다.
샘플 효율성과 학습 불안정성
샘플 효율성(sample efficiency) 측면에서도 Feedback Distillation은 고무적입니다. Qwen3.5-9B를 베이스로 할 때, 학습 초기에는 Feedback Distillation이 GRPO보다 훨씬 샘플 효율적이며, GRPO는 나중에 따라잡지만 곧 붕괴(crash)합니다. 여기서 GRPO의 배치 크기는 그룹 크기 5 때문에 Feedback Distillation의 5 배라는 점을 감안하면 더욱 인상적입니다.
다만 저자들은 GRPO와 GRPO over Feedback Distillation에서 학습 불안정성(training instability)을 관찰했다고 솔직히 보고합니다. 이는 기존 연구들에서도 지적된 현상으로, 학습률(learning rate)을 조정해 불안정성을 늦추는 것만으로는 GRPO over Feedback Distillation의 최대 성능을 GRPO 단독으로 재현할 수 없음을 확인했습니다. 또한 더 안정적으로 학습되는 Qwen3-8B에서도 "GRPO와 Feedback Distillation의 결합이 GRPO 단독을 능가한다" 는 결론이 동일하게 성립하여, 개선이 단순히 GRPO 불안정성의 부산물이 아님을 뒷받침합니다.
Feedback Distillation 심층 분석
그렇다면 왜 이 방법이 작동하는가, 그리고 어떤 요소가 결정적인가 도 궁금해집니다. 저자들은 피드백의 출처, EMA의 역할, 그리고 학습의 질적 특성을 차례로 분석합니다.
피드백 출처에 따른 차이
저자들은 Claude 외에 두 가지 피드백 출처를 비교합니다. 하나는 LeanWorkbook 문장의 정답 풀이를 그대로 주는 ground-truth 피드백, 다른 하나는 Lean 컴파일러의 원시 출력을 주는 Lean output 피드백 입니다. 비교 결과 Claude와 ground-truth 피드백이 Lean output보다 더 나은 성능을 보였는데, 둘 다 더 강한 정보원에 접근하기 때문에 자연스러운 결과입니다. 흥미로운 점은 Claude를 피드백 모델로 쓸 때가 ground-truth 풀이를 쓸 때보다 더 높은 엔트로피를 유도한다는 것입니다. 완성된 정답을 통째로 주기보다 힌트와 교정을 주기 때문이라고 저자들은 해석합니다.
Self-Feedback Distillation
피드백 모델을 더 강한 외부 모델이 아니라, 학습 대상과 동일한 모델의 동결된 복사본으로 두는 설정이 Self-Feedback Distillation 입니다. 이 설정은 외부 정보를 전혀 추가하지 않으므로 GRPO와의 공정한 비교를 가능하게 하며, 모델이 오직 자기 생성물에 대한 자신의 분석만으로 향상될 수 있는지를 분리해 보여줍니다. 실험 결과 GRPO over Self-Feedback Distillation은 GRPO 단독보다 더 높은 최고 성능을 달성했습니다. 모델이 GRPO만으로는 발견하지 못하는 유용한 학습 신호를 자기 피드백에서 끌어낼 수 있다는 뜻입니다.
피드백은 가중치에 내재화된다
저자들은 피드백이 정말 모델 가중치 안으로 녹아드는지 확인하기 위해, 피드백 F(y) 를 학습 내내 고정된 문자열로 둔 단순화 실험을 합니다. "간결하게 답하라(Be concise)" 와 "답하기 전에 신중히 생각하고 풀이가 맞는지 검증하라" 두 경우를 비교한 결과, 평균 응답 길이가 피드백에 맞게 변하고 학생과 교사 사이의 KL 발산이 0 으로 수렴했습니다. 학습이 끝나면 학생은 마치 모든 프롬프트에 피드백이 이미 붙어 있는 것처럼 행동합니다. 피드백이 가중치에 내재화(internalize)된다는 직접적인 증거입니다.
인식적 표현(Epistemic Verbalization)의 보존
한 가지 미묘한 발견도 있습니다. 정답 풀이처럼 풍부한 특권 정보로 자기 증류를 하면, wait, hmm, perhaps 같은 불확실성 표현, 즉 인식적 표현(epistemic verbalization) 이 억제되어 분포 밖(out-of-distribution) 성능이 떨어질 수 있다는 선행 연구가 있습니다. 그런데 저자들이 관찰한 바로는, 피드백을 정답이 아닌 비평(critique) 형태로 주는 Feedback Distillation에서는 이 억제가 훨씬 덜 일어났습니다. 조건화 신호의 정보 성격이 이 효과를 좌우할 수 있음을 시사합니다. 비평형 피드백은 정답형 피드백보다 더 긴 응답과도 연관되었는데, 이는 짧고 단정적인 증명으로 수렴하기보다 궤적 다양성을 유지하는 모습과 일관됩니다.
On-policy 지식 전이로서의 Feedback Distillation
피드백 모델이 학생보다 강할 때, Feedback Distillation은 일종의 간접 지식 증류로 작동합니다. 강한 모델이 완전한 풀이를 제공하지 않고도 정보를 전달한다는 점에서 표준 증류와 다릅니다. 또한 학습 신호가 학생의 생성물 위에서 계산되므로 본질적으로 on-policy이며, 이는 SFT 같은 off-policy 방법보다 일반화를 개선하고 치명적 망각을 줄이는 것으로 알려져 있습니다. 표준 on-policy 증류 대비 실용적 이점도 있습니다. 지식이 로짓 매칭(logit matching)이 아니라 자연어 피드백을 통해 전달되므로, 피드백 모델을 직접 호스팅하거나 토크나이저(tokenizer)를 공유할 필요가 없습니다.
한계점 및 향후 전망
저자들은 현재 결과가 대규모로 학습된 최신 SFT와 GRPO 파이프라인의 수준에는 아직 미치지 못한다는 점을 분명히 합니다. 그럼에도 Feedback Distillation을 외부 지식 주입, 다양성 유지, 세밀한 credit assignment라는 기존 접근법의 근본적 한계를 동시에 건드리는 유망한 메커니즘으로 평가합니다.
가장 큰 미해결 과제는 학습 안정성입니다. 큰 \alpha 를 써도 장기적으로 불안정성이 남기 때문에, 더 견고한 안정화 메커니즘을 개발하는 것이 더 긴 학습으로 확장하기 위한 자연스러운 다음 단계입니다. 또한 피드백 프롬프트 설계를 개선하면 Claude 기반 피드백의 성능이 더 올라갈 여지가 있으며, 원리적으로는 컴파일러 출력에 접근하고 LeanWorkbook 문제를 풀 만큼 유능한 Claude가 Lean output과 ground-truth를 모두 능가할 가능성도 있다고 봅니다.
넓게 보면 이 연구는, 현재의 LLM이 형식 수학 같은 grounded 환경에서 자신의 성공과 실패를 스스로 분석할 만큼의 추론 능력을 이미 갖추고 있으며, Feedback Distillation이 그 자기 성찰 능력을 모델 가중치 안으로 내재화하는 한 가지 길을 제시한다는 점에서 의미가 있습니다. Lean 정리 증명을 넘어 다른 형식적, 비형식적 추론 과제로 확장된다면 이 접근법의 일반성을 더 분명히 보여줄 수 있을 것입니다.
Distilling LLM Feedback for Lean Theorem Proving 논문
더 읽어보기
이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. 
파이토치 한국 사용자 모임
이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일
로 보내드립니다!
텔레그램(Telegram)이나 Slack/Discord/Teams/Dooray/GoogleChat 등으로도 새 글 알림을 받으실 수 있습니다. 
아래
쪽에 좋아요
를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ 