Distributional DAgger를 이용한 풍부한 피드백 기반의 강화학습 (RL)

추론 모델 (Reasoning models)은 급격히 발전해 왔으나, 검증 가능한 보상 기반의 강화학습 (RLVR, Reinforcement Learning from Verifiable Rewards) 방식은 놀라울 정도로 협소한 상태로 남아 있습니다. 즉, 많은 응답을 샘플링하고 최종 답변의 정답 여부를 나타내는 단일 비트(single bit)로 각 응답에 보상을 주는 방식입니다. 하지만 실행 추적 (execution traces), 도구 출력 (tool outputs), 전문가의 수정 (expert corrections), 모델의 자기 평가 (model self-evaluations)를 포함하여 많은 설정이 풍부한 피드백을 제공합니다. 본 연구에서는 고전적인 모방 학습 (imitation learning) 알고리즘인 DAgger의 분포 변형 (distributional variant)을 통해 이러한 피드백을 사용하는 방법을 연구합니다. 여기서 학습자는 현재 정책 (policy)에 의해 방문된 상태들에 대한 전문가 분포 (expert distribution)에 지역적으로 접근할 수 있습니다. 이는 블랙박스 전문가를 허용하는 단순한 순방향 교차 엔트로피 (forward cross-entropy) 목적 함수를 생성하며, 이 함수의 시퀀스 수준 그래디언트 (sequence-level gradient)는 미래의 전문가-학생 간 불일치 (expert-student disagreement)를 이전의 결정으로 전파함으로써 풍부한 신용 할당 (credit assignment)을 수행합니다. 우리는 역방향 KL (reverse KL) 또는 젠슨-샤논 (Jensen-Shannon) 기반의 자기 증류 (self-distillation) 목적 함수를 사용하는 기존의 강화학습 (RL)이 단조적 정책 개선 (monotonic policy improvement)을 보장하지 못함을 보여줍니다. 즉, 전문가가 더 높은 보상을 가짐에도 불구하고, 업데이트 과정에서 더 나쁜 행동의 확률을 높일 수 있습니다. 이와 대조적으로, 우리는 순방향 교차 엔트로피가 단조적 정책 개선을 허용하며 후회 (regret)에 대한 보장을 누림을 보여줍니다. 나아가 우리의 목적 함수가 교사 가중 성공 가능도 (teacher-weighted likelihood of success)의 하한 (lower bound)을 최적화하여 Pass@N을 향상시킨다는 것을 보여줍니다. 실험적으로 우리의 접근 방식인 DistIL은 과학적 추론 (scientific reasoning), 코딩 (coding), 어려운 수학 문제 해결 등 다양한 도메인에서 RLVR 및 자기 증류 기반 강화학습 (RL) 베이스라인보다 뛰어난 성능을 보입니다.