Why Does Feedback-Augmented Self-Distillation Fail to Improve Retrieval-Interleaved Search Agents?

기술 요약: 피드백 증강 자기 증류(Feedback-Augmented Self-Distillation)가 검색 인터리브형 탐색 에이전트의 성능 개선에 실패하는 이유

문제 정의

추론과 외부 도구 사용(검색)을 교대로 수행하는 검색 인터리브형 탐색 에이전트는 결과 수준의 보상(outcome-level rewards)이 희소하다는 문제로 인해 학습에 어려움을 겪습니다. 강화 학습(RL)과 조밀한 감독 신호(예: 과정 보상)가 제안되었으나, **온폴리시 자기 증류(On-Policy Self-Distillation, OPSD)**는 단일 대규모 언어 모델(LLM)을 학생과 교사로 모두 사용하는 유망한 대안을 제시합니다. OPSD에서 교사는 "특권 정보(privileged information)"(예: 성공적인 롤아웃)를 조건으로 하여, 무조건화된 학생에게 조밀한 토큰 수준의 감독을 제공합니다.

그러나 이 패러다임이 검색 인터리브형 에이전트에도 효과적으로 확장될 수 있는지는 불분명합니다. 표준 언어 모델링과 달리, 검색 에이전트는 모델의 행동이 후속적으로 검색된 관찰값(환경 출력)을 결정하는 궤적을 생성하므로, 표준 토큰 수준 증류와 효과적인 검색을 위한 장기적(long-horizon) 검색 결정 사이에는 미스매치가 발생합니다. 본 논문은 성공적인 롤아웃을 특권 데모로 사용하여 셀프-티처(self-teacher)를 조건화하는 **피드백 증강 자기 증류(Feedback-Augmented Self-Distillation, FA-SD)**가 이러한 에이전트에게 신뢰할 수 있는 감독을 제공할 수 있는지 조사합니다.

방법론

1. 피드백 증강 자기 증류 (FA-SD)

저자들은 FA-SD를 검색 에이전트에 맞게 구현했습니다. 과정은 다음과 같습니다:

• 롤아웃 생성: 주어진 질문에 대해 일련의 궤적을 샘플링합니다.
• 특권 피드백: 그룹 내에 성공적인 궤적이 존재하면, 이를 피드백($f$)으로 사용하여 "셀프-티처"의 프롬프트를 증강합니다.
• 증류 목적 함수: 무조건화된 학생 정책($\pi_\theta$)은 피드백으로 증강된 셀프-티처($\pi_\theta(\cdot | x, f)$)의 토큰 분포를 따르도록 학습됩니다. 손실 함수는 학생과 스톱-그레이디언트(stop-gradient) 셀프-티처 사이의 발산(예: KL 발산)을 최소화하며, 이때 검색된 관찰값은 마스킹하고 어시스턴트가 생성한 토큰에 집중합니다.

2. 실패 모드 진단

FA-SD의 불안정성을 이해하기 위해 저자들은 감독 신호의 불안정성을 두 가지 구성 요소로 분해합니다:

• 모델 불일치(Model Inconsistency): 셀프-티처와 학생은 동일하게 지속적으로 업데이트되는 파라미터를 공유하므로, 훈련 중에 티처 타겟이 표류(drift)하게 됩니다.
• 프롬프트 불일치(Prompt Inconsistency): 티처는 피드백 데모를 조건으로 하는 반면, 학생은 그렇지 않습니다. 이는 깔끔하게 전이되지 않을 수 있는 분포적 미스매치를 생성합니다.

3. 제안된 완화 방법

모델 불일치를 해결하기 위해 저자들은 셀프-티처를 위한 두 가지 정규화 전략을 평가합니다:

• 고정 참조 티처(Fixed Reference Teacher): 티처를 훈련 시작 시점의 학생의 초기 파라받터로 고정합니다.
• 지수 이동 평균(EMA) 티처: 티처 파라미터를 학생 파라미터의 이동 평균으로 업데이트하여($\theta_{EMA} = \alpha \theta_{EMA} + (1-\alpha)\theta_t$), 티처 타겟의 표류를 늦춥니다.

4. 외부 티처 확장

저자들은 피드백 증강 기법이 외부 티처를 사용하는 **멀티 롤아웃 온폴리시 증류(Multi-Rollout On-Policy Distillation, MOPD)**인 FA-MOPD로 일반화될 수 있는지 테스트합니다. 여기서 외부 티처는 성공적인 롤아웃을 조건으로 하며, 학생은 무조건화된 상태를 유지합니다.

주요 결과 및 발견

1. 표준 FA-SD의 실패

7개의 오픈 도메인 QA 벤치마크(NQ, HotpotQA 등) 전반에서 표준 FA-SD 변체(unclipped 및 PPO 스타일 clipped)는 개선 지속에 실패했습니다.

• 훈련 역학: 통과율(pass rate)은 훈련 중 대부분 거의 0에 가깝게 유지되다가, 일시적인 반등 후 다시 퇴보합니다.
• 디코딩 붕괴(Decoding Collapse): 저자들은 디코딩 붕해라고 명명된 특정 실패 모드를 식별했습니다. 모델이 반복되는 추론 및 검색 출력 템플릿에 의존하기 시작합니다. 이러한 궤적은 다양해 보이지만 실제로는 입력 질문과는 무관합니다. 결과적으로 KL 기반 증류 신호는 무의미해지는데, 이는 모델이 특정 질문을 해결하는 법을 배우는 것이 아니라 일반적인 템플릿을 흉내 내고 있기 때문입니다. 표준 집계 지표들은 입력 특이성(input-specificity)을 검증하지 않기 때문에 이를 놓치는 경우가 많습니다.

2. 셀프-티처는 능력이 있으나 내부화에 실패함

실험 결과, 피드백 증강된 셀프-티처는 추론 시 무조건화된 학생보다 종종 더 강력함을 보여주었습니다. 이는 실패의 원인이 약한 티처 때문이라는 가설을 배제합니다. 대신, 문제는 **내부화 실패(internalization failure)**입니다. 즉, 학생이 특권 피드백으로부터 도출된 토큰 수준의 증류 신호를 무조건화된 정책으로 안정적으로 흡수하지 못하는 것입니다.

3. 정규화의 영향

• 안정화: 고정 참조 및 EMA 티처 모두 티처 타겟이 너무 빠르게 표류하는 것을 방지함으로써 훈련을 안정화합니다.
• 성능: EMA 티처가 가장 좋은 결과를 나타냈습니다. 이는 모델이 일시적인 퇴보로부터 회복할 수 있게 하며, 베이스라인(0.114)에 비해 더 높은 최종 성능(벤치마크 평균 점수 0.206)을 달-성합니다.
• 웜업 단계: EMA 방식은 감독 신호가 안정화되고 개선되기 전까지 성능이 일시적으로 퇴보할 수 있는 웜업 단계를 필요로 합니다.

4. 외부 티처로의 일반화

피드백 증강 기법을 강한 외부 티처를 가진 **MOPD(FA-MOPD)**에 적용했을 때, 성능은 표준 MOPD에 비해 저하되었습니다. 이는 강력한 외부 티처가 있더라도 특권 피드백이 프롬프트 불일치 또는 불필요한 조건화를 유발하며, 외부 티처 기반의 OPD로 직결되게 일반화되지 않음을 시사합니다.

핵심 기여

1. FA-SD의 구현: 본 논문은 성공적인 롤아웃을 특권 데모로 활용하여, 검색 인터리브형 탐색 에이전트를 위한 피드백 증강 자기 증류를 정의하고 구현했습니다.
2. 디코딩 붕괴의 식별: 저자들은 모델이 반복되는 템플릿을 통해 입력과 무관한 궤적을 생성하여 증류 신호를 무의미하게 만드는 실패 모드인 디코딩 붕괴를 규명했습니다. 이 현상은 단순한 성능 퇴보와는 다르며, 표준 지표로는 흔히 간과될 수 있습니다.
3. 감독 불일치 분석: 연구는 FA-SD의 불안정성을 모델 불일치(표류하는 티처)와 프롬프트 불일치(조건화 미스매치)로 분해합니다. 또한 EMA나 고정 참조를 통해 티처를 안정화하는 것이 효과적인 훈련에 필수적임을 입증했습니다.
4. 일반화에 대한 부정적 결과: 본 논문은 강한 외부 티처가 MOPD에서 안정적인 초기 감독을 제공하지만, 여기에 특권 피드백을 추가하는 것(FA-MOPD)은 성능을 오히려 저하시킨다는 것을 보여줌으로써, 피드백 증강이 검색 에이전트 설정에서 특정한 불일치를 유발함을 밝혀냈습니다.

의의

본 논문은 조밀한 감독을 위한 유망한 기술인 자기 증류가 복잡한 검색 인터리브형 에이전트 환경에서 왜 실패하는지에 대한 비판적인 진단을 제공합니다. 저자들은 실패의 원인이 단순히 강한 티처의 부재가 아니라, 디코딩 붕괴와 감독 불일치에 있음을 강조합니다. EMA 정규화가 셀프-티처를 안정화하고 성능을 회복할 수 있음을 입증함으로써, 본 연구는 에이전트 작업에 자기 증류를 적용하기 위한 실질적인 경로를 제시합니다. 나아가, 피드백 증강을 외부 티처 프레임워크에 나이브하게 적용하는 것에 대해 경고하며, 특권 정보의 이점이 증류 설정의 특정 아키텍처와 조건화에 매우 민감하다는 점을 시사합니다.