Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

1. 문제 정의 (Problem Statement)

대형 언어 모델 (LLM) 은 하위 도메인 (Downstream) 작업에 적용하기 위해 지도 미세 조정 (Supervised Fine-Tuning, SFT) 을 거치는 것이 일반적입니다. 그러나 하위 도메인 데이터가 모델의 사전 학습 분포와 크게 다를 때 (Distribution Shift), SFT 는 재앙적 망각 (Catastrophic Forgetting) 을 유발하여 모델이 이전에 습득한 일반 능력을 상실하게 됩니다.

이 현상의 핵심 원인은 오프-폴리시 (Off-Policy) 학습의 불안정성에 있습니다.

• 분포 불일치: 기존 데이터 재작성 (Data Rewriting) 방법들은 주로 고정된 프롬프트 템플릿을 사용하여 전문가 해답 ($y^*$) 을 모델이 선호하는 QA 스타일 생성 분포 ($\pi_0(\cdot|x)$) 에 가깝게 재작성합니다.
• 한계: 그러나 이러한 재작성은 프롬프트 조건부 분포 ($\pi_0(\cdot|x, y^*, x_{prompt})$) 에만 최적화되어 있을 뿐, 모델의 자연스러운 QA 스타일 분포와는 여전히 괴리가 있을 수 있습니다. 또한 고정된 템플릿 사용은 출력 다양성 붕괴 (Diversity Collapse) 를 초래하여 학습 효율성을 저하시킵니다.

2. 제안 방법 (Methodology)

저자들은 데이터 재작성을 정책 학습 (Policy Learning) 문제로 재정의하고, 강화 학습 (RL) 을 기반으로 한 재작성 에이전트 (Rewriting Agent) 를 제안합니다. 이 에이전트는 백본 모델의 QA 스타일 생성 분포와 일치하면서도 다양성을 유지하는 고품질 재작성 데이터를 생성합니다.

2.1 프레임워크 개요

1. Stage I: RL 기반 재작성 에이전트 학습

   • 백본: 고정된 (Frozen) 지시 학습 (Instruction-tuned) 모델 $\pi_0$ 를 사용.
   • 파라미터 효율성: LoRA (Low-Rank Adaptation) 어댑터만 학습하여 경량화된 "패치 (Patch)"를 생성. 이는 백본의 분포를 크게 벗어나지 않도록 제어합니다.
   • 최적화 알고리즘: GRPO (Group Relative Policy Optimization) 사용. 별도의 가치 네트워크 없이 그룹 내 샘플 간 상대적 성능을 기반으로 정책을 업데이트합니다.

2. Stage II: 데이터 구축 및 하위 SFT

   • 학습된 에이전트로 재작성된 데이터셋 ($D_R$) 을 생성한 후, 이를 사용하여 목표 모델 $\pi_\theta$ 의 표준 SFT 를 수행합니다.

2.2 통합 목적 함수 (Unified Objective)

재작성 정책은 다음 세 가지 요소를 동시에 최적화하는 보상 신호를 받습니다.

• 1) 작업 일관성 (Task Consistency, $r_{task}$):
  • 하드 게이트 (Hard Gate): 재작성된 답변이 최종 정답과 추론 논리 모두에서 정확해야만 보상이 부여됩니다.
  • 검증 방식: 규칙 기반 검증 (정답 확인) + LLM 판사 (추론 타당성 평가) 의 2 단계 검증을 거칩니다.
• 2) 분포 정렬 (Distributional Alignment, $r_{dist}$):
  • 재작성된 텍스트가 백본 모델의 QA 스타일 조건부 분포 ($\pi_0(\cdot|x)$) 에서 높은 확률 (Likelihood) 을 가지도록 장려합니다.
  • 이를 위해 고정된 백본 모델에 대한 정규화된 음의 로그 가능도 (NLL) 를 보상으로 사용합니다.
• 3) 다양성 (Diversity, $r_{div}$):
  • 동일한 입력에 대해 생성된 여러 후보 간의 의미적 다양성을 장려하여 템플릿 기반의 획일화를 방지합니다.
  • 임베딩 공간에서의 쌍별 거리 (Pairwise distance) 를 기반으로 보상을 산정합니다.

보상 함수:
$$r(x, y^*, \tilde{y}) = r_{task} + r_{task} \cdot (\lambda_{dist} r_{dist} + \lambda_{div} r_{div})$$

• $r_{task}=0$(실패) 인 경우, 정렬 및 다양성 보상은 0 으로 처리되어 잘못된 데이터가 학습에 영향을 미치지 않도록 합니다.

3. 주요 기여 (Key Contributions)

1. 분포 불일치 통찰: 기존 데이터 재작성 파이프라인이 모델의 자연스러운 QA 생성 분포와 불일치한다는 점을 규명하고, 이를 해결하기 위한 RL 기반 프레임워크를 제안했습니다.
2. 통합 목적 함수: 작업 일관성을 강제하면서 분포 정렬과 다양성을 동시에 최적화하는 새로운 목적 함수를 설계하여, 고품질의 재작성 데이터셋을 구축했습니다.
3. 실험적 검증: 다양한 백본 모델 (Llama-3.2-1B/3B, Mistral-7B) 에서 표준 SFT 와 비교하여 하위 작업 성능은 유지하면서 재앙적 망각을 획기적으로 줄임을 입증했습니다.

4. 실험 결과 (Results)

실험은 수학 추론 (Math Reasoning) 을 하위 작업으로, 일반 도메인 벤치마크 (MMLU 등) 를 재앙적 망각 측정 지표로 설정하여 수행되었습니다.

• 성능 향상: 제안된 Rewriting-agent는 표준 SFT 와 유사하거나 더 높은 수학 문제 해결 능력 (MathAvg) 을 보여주었습니다.
  • 예: Mistral-7B 에서 SFT 대비 +55.23% 향상 (SFT 는 +54.19%).
• 망각 감소 (Catastrophic Forgetting Mitigation):
  • 표준 SFT 는 일반 도메인 성능이 19.17% 하락했으나, 제안 방법은 7.35% 하락으로 크게 줄였습니다.
  • 평균적으로 12.34% 더 적은 망각을 보였습니다.
• 학습 안정성: 재작성된 데이터로 학습한 모델은 SFT 학습 손실 (Training Loss) 이 더 낮고 안정적으로 수렴하는 것을 확인했습니다. 이는 재작성 데이터가 모델의 자연스러운 분포에 더 잘 맞기 때문입니다.
• Ablation Study:
  • 분포 정렬 보상 ($r_{dist}$) 을 제거할 경우 일반화 성능이 가장 크게 저하됨.
  • 다양성 보상 ($r_{div}$) 이 도메인 내 성능 유지에 중요함.
  • 'Hard gating' 방식이 'Soft shaping'보다 학습 안정성과 성능 면에서 우월함.

5. 의의 및 결론 (Significance)

이 논문은 LLM 의 하위 도메인 적응 과정에서 발생하는 오프-폴리시 불안정성을 데이터 소스 (Data Source) 에서 해결하는 새로운 패러다임을 제시합니다.

• 데이터 중심 접근: 모델 구조를 변경하거나 복잡한 정규화 기법을 적용하는 대신, 학습 데이터의 분포를 모델의 자연스러운 생성 분포에 맞춰 재구성함으로써 안정성을 확보했습니다.
• RL 의 효과적 활용: 비미분 가능한 검증 과정 (Task Consistency) 을 하드 게이트로 처리하고, 분포 정렬과 다양성과 같은 복잡한 목표를 RL 보상으로 최적화하여 효과적으로 해결했습니다.
• 실용성: LoRA 를 활용한 경량화 접근으로 계산 비용을 낮추면서도, 재앙적 망각을 방지하고 하위 작업 성능을 동시에 향상시킬 수 있는 실용적인 솔루션을 제공합니다.

결론적으로, 이 연구는 재작성 에이전트를 통해 "모델이 배우기 쉬운 (In-distribution)" 데이터를 생성함으로써, SFT 의 안정성을 높이고 일반 능력을 보존하는 효과적인 전략을 제시합니다.