기술 요약: 자기 증류(Self-Distillation)에서 피드백 정렬의 역할

문제 정의

추론을 위한 대규모 언어 모델(LLM)의 사후 학습은 일반적으로 그룹 상대 정책 최적화(GRPO)와 같은 검증 가능한 보상 기반 강화 학습(RLVR)에 의존합니다. RLVR은 효과적이지만, 희소한 크레딧 할당(sparse credit assignment) 문제를 겪습니다. 즉, 최종 답변에 대한 단일 스칼라 보상은 모델이 추론 과정 중 어느 부분에서 오류를 범했는지 알려주지 못합니다. 반면, 지식 증류(Knowledge Distillation)는 조밀한 토큰 수준의 감독을 제공하지만, 더 강력한 교사 모델의 로짓(logits)에 접근해야 하며 이는 종종 불가능하거나 비용이 많이 듭니다.

자기 증류는 동일한 모델을 서로 다른 문맥 하에서 학생과 교사로 모두 사용하여 이 간극을 메우고자 합니다. 학생은 질문에만 의존하는 반면, "자기-교사(self-teacher)"는 질문에 추가적인 문맥(예: 참조 솔루션 또는 피드백)을 결합하여 조건화됩니다. 그러나 기존 연구들은 이러한 문맥을 고정된 선택지로 취급했을 뿐, 피드백의 구조가 학습 신호에 어떻게 영향을 미치는지에 대해서는 조사하지 않았습니다. 본 논문은 다음과 같은 질문을 던집니다: 어떤 형태의 피드백이 자기 증류를 위해 가장 효과적인 자기-교사를 생성하는가?

방법론

저자들은 수학적 추론 도메인에서 자기 증류를 연구하기 위해 솔버-크리틱(solver–critic) 설정을 구현합니다.

1. 설정: 학습 가능한 솔버($\pi_\theta$)가 수학 문제에 대한 추론 궤적을 생성합니다. 동결된 크리틱($\pi_{critic}$)은 이러한 궤적을 분석하고 피드백을 생성합니다. 이후 솔버는 질문에만 의존하는 출력 분포와 질문 및 크리틱의 피드백에 의존하는 자기-교사의 분포 사이의 발산(divergence)을 최소화함으로써 업데이트됩니다.
2. 실험 조건: 연구는 세 가지 뚜렷한 피드백 문맥을 비교합니다:
   • GRPO: 자기 증류 없이 이진 보상(정답/오답)을 사용하는 표준 RLVR 베이스라인입니다.
   • RefSol (참조 솔루션): 자기-교사가 더 강력한 모델으로부터 가져온 정답 참조 솔루션(사고의 사슬)에 조건화됩니다.
   • StepAlignFB (단계 정렬 비평): 자기-교사가 크리틱이 생성한 비평에 조건화됩니다. 결정적으로, 이 비평은 솔버의 특정 추론 궤적과 구조적으로 정렬되어 있습니다. 크리틱은 솔버의 올바른 단계들을 있는 그대로 복사하고, 오류가 발생하는 특정 단계만을 수정하도록 프롬프트되어, 올바른 부분에 대해서는 솔버의 스타일과 표기법을 보존합니다.
3. 메커니즘 분석: 저자들은 서로 다른 피드백 유형이 모델의 행동을 어떻게 변화시키는지 이해하기 위해 "토큰당 어드밴티지(per-token advantage)"(교사와 학생 간 확률의 로그 비율)를 분석합니다. 또한, 모델이 문맥 내 패턴을 복사하는 메커니즘인 "유도 헤드(induction heads)"가 피드백 내의 구절 반복을 어떻게 처리하는지 조사합니다.

주요 기여 및 결과

1. 단계 정렬 피드백이 참조 솔루션보다 우수함

StepAlignFB 조건은 RefSol 조건부 자기 증류와 GRPO보다 더 큰 성능 향상을 나타냈습니다.

• 결과: OpenMathReasoning 데이터셋(Avg@12 지표)에서 StepAlignFB는 35.83을 기록하여, RefSol 조건부 자기 증류의 30.56과 GRPO의 19.72를 상회했습니다.
• 의미: 정답 도출 과정을 전혀 보지 못했음에도 불구하고 StepAlignFB가 RefSol보다 높은 성능을 보였습니다. 이는 피드백의 단순한 정답 존재 여부보다, 솔버의 특정 궤적에 대한 피드백의 정렬이 더 중요하다는 것을 시사합니다.

2. 구조적 정렬을 통한 암시적 프로세스 감독

토큰당 어드밴티지 분석은 StepAlignFB의 성공 메커니즘을 밝혀줍니다:

• 국소화(Localization): 단계 정렬 피드백은 분포 변화(음의 어드밴티지)를 추론이 실패하는 지점의 토큰에 집중시킵니다. 또한 올바른 단계를 그대로 보존함으로써 이를 강화합니다(양의 어드밴티지 생성). 이는 별도의 보상 모델을 훈련하지 않고도 자연어 피드백으로부터 프로세스 보상 모델(PRM)의 동작을 모방합니다.
• RefSol에서의 확산: 반면, 참조 솔루션에 조건화되는 경우 "확산된 억제 신호(diffuse suppressive signal)"가 발생합니다. 솔버의 단계가 정답일지라도, 참조 솔루션은 종종 다른 문구나 표기법을 사용합니다. 이로 인해 자기-교사는 단순히 참조 솔루션과 다르다는 이유만으로 올바른 토큰을 억제하게 되어, 스타일적 불일치로 인해 학습 신호가 희석됩니다.

3. 유도 헤드 복사의 역할

논문은 크리틱의 피드백을 위한 핵심 설계 패턴인 **"충실한 필기자 관례(faithful-scribe convention)"**를 식별합니다:

• 전체 구절 반복: 만약 크리틱이 솔버의 전체 시도(오류 포함)를 그대로 복사하면, 모델의 유도 헤드가 잘못된 토큰에 고착되어 교사가 오류를 재현하게 만듭니다.
• 생략: 만약 크리틱이 솔버의 올바른 단계를 완전히 생략하면, 교사의 분포가 솔버의 올바른 추론으로부터 멀어집니다.
• 부분 반복 (해결책): 최적의 전략은 솔버의 올바른 단계를 오류가 발생하는 지점까지 그대로 복사하고, 해당 오류 단계를 수정 사항으로 대체한 뒤 계속 진행하는 것입니다. 이는 유도 헤드의 복사 기능을 활용하여 올바른 접두사(prefix)를 강화하는 동시에, 수정 사항이 오류 위치를 제어할 수 있도록 합니다.

의의 및 주장

본 논문은 자기 증류의 주요 이점이 단순히 추가적인 문맥에 접근하는 것이 아니라, 구조가 학생의 추론 궤적과 일치하는 문맥에 접근하는 데 있다고 주장합니다.

• 구조적 정렬 > 피드백 품질: 완전하고 올바른 참조 도출 과정은 강력한 신호이지만, 자기 증류에서는 표면적 형태가 솔버의 궤적과 달라지기 때문에 실패합니다. 단계 정렬 비평은 솔버의 실제 궤적을 다룸으로써, 오류가 발생하는 정확한 지점에 신호를 집중시키기 때문에 성공합니다.
• 암시적 PRM: 단계 정렬 피드백은 별도의 단계별 보상 모델을 훈련하는 계산 비용 없이도 오류를 국소화하는 일종의 암시적 프로세스 감독을 제공합니다.
• 실질적 트레이드오프: StepAlignFB가 우수하지만, 저자들은 고품질의 단계 정렬 피드백을 생성하기 위해 역량 있는 비평 모델이 필요하며, 이는 단순한 참조 솔루션 조건화나 이진 보상에 비해 훈련 비용과 시스템 복잡도를 증가시킨다고 언급합니다.

이 연구는 OpenMathReasoning 데이터셋과 수학적 추론에 국한되어 있으며, 저자들은 이러한 이점이 다른 도메인, 모델 크기 또는 솔버-크리틱 구성으로 전이될 수 있는지 확인하기 위해 추가 연구가 필요함을 명시했습니다.