Post-Training is About States, Not Tokens: A State Distribution View of SFT, RL, and On-Policy Distillation

기술적 요약: 포스트-트레이닝을 상태 분포 형성으로 해석

문제 제기

대형 언어 모델 (LLM) 의 포스트-트레이닝 방법 (특히 지도 미세조정 (SFT), 강화 학습 (RL), 증류) 에 대한 현재 분석은 주로 최적화 목표 (예: 최대 가능도, 정책 경사, KL 발산) 에 초점을 맞추고 있습니다. 그러나 이러한 목표 중심의 관점은 몇 가지 경험적 현상을 설명하는 데 어려움을 겪습니다:

• 왜 온건한 SFT 는 최소한의 망각으로 성능을 향상시키는 반면, 공격적인 ("스트레스") SFT 는 치명적인 망각과 목표 성능 저하를 초래하는지.
• 왜 RL 이 희소하거나 노이즈가 많은 보상을 사용함에도 불구하고 종종 안정적인 개선을 이루는지.
• 왜 증류에서 학생 모델이 때로는 교사 모델이 유일한 감독 소스임에도 불구하고 열화된 교사보다 더 나은 성능을 발휘할 수 있는지.

이 논문은 이러한 불일치가 표준 분석이 감독이 적용되는 상태 분포라는 중요한 변이 축을 간과하기 때문에 발생한다고 주장합니다. 자기회귀 모델에서 "상태"는 프롬프트와 생성된 접두사로 정의됩니다 ($s_t = (x, y_{<t})$). 다양한 포스트-트레이닝 방법은 근본적으로 다른 상태 분포 (고정된 데이터셋 접두사 대 학습자가 유도한 롤아웃) 에 신호를 적용하며, 이는 모델의 행동, 국소성, 그리고 유지 능력에 상당한 영향을 미칩니다.

방법론

저자들은 포스트-트레이닝을 상태 분포 형성으로 보는 통합 프레임워크를 제안합니다. 그들은 모든 포스트-트레이닝 방법을 두 가지 명확한 축으로 분해합니다:

1. 상태 소스 ($q_k$): 훈련 컨텍스트 (상태) 가 기원하는 곳.
2. 신호 소스 ($S$): 해당 상태에서 제공되는 감독 (토큰, 보상, 로짓) 의 내용.

이론적 공식화:

• SFT: 오프-정책 상태 피팅. 모델은 고정된 데이터셋 상태 ($d_{data}$) 에서 토큰 손실을 최소화합니다. 데이터셋 궤적이 모델 자체의 롤아웃과 달라질 경우 불일치가 발생하여 노출 편향과 스트레스 하에서의 잠재적 불안정성을 초래합니다.
• RL: 온-정책 국소 개선. 모델은 현재 정책 ($d_{\pi_\theta}$) 에서 궤적을 샘플링하고 보상을 기반으로 업데이트합니다. 업데이트는 모델이 실제로 방문하는 상태에 자연스럽게 국소화됩니다.
• 온-정책 증류 (OPD): 학생이 상태 ($d_{\pi_S}$) 를 샘플링하지만 교사가 신호 (로짓 또는 연속어) 를 제공하는 하이브리드 방식입니다. 이는 상태 소스를 신호 소스로부터 분리하여, 학생이 교사의 전역 궤적 분포를 물려받지 않고 국소적 수리를 학습할 수 있게 합니다.

실험 설정:

• 모델: LoRA 어댑터를 적용한 Qwen3-0.6B-Base.
• 하드웨어: 단일 RTX 3090 GPU.
• 작업:
  • 목표: GSM8K (수학적 추론).
  • 유지: TruthfulQA 및 MMLU (치명적 망각 측정을 위해).
• 지표: 목표 정확도, 유지 비율, 망각 크기, 그리고 롤아웃 상태의 어휘적 특징에 대한 최대 평균 불일치 (MMD) 를 통해 측정한 상태 드리프트.

실험 조건:

1. 온건한 SFT: GSM8K 에 대한 표준 훈련.
2. 스트레스 SFT: 열화를 유도하기 위한 공격적인 훈련 (5 에포크, 높은 학습률, 높은 LoRA 랭크).
3. OPD: 신호 소스로는 열화된 "스트레스 SFT" 교사를 사용하지만, 상태는 학생의 자체 정책에서 샘플링하여 학생을 훈련.
4. 온-정책 RL: GSM8K 에서 그룹 상대적 정답 보상을 사용하는 경량 RL.

주요 결과

1. SFT 는 컨텍스트에 의존적임: 온건한 SFT 는 유지 손실이 거의 없는 상태에서 GSM8K 를 향상시켰습니다 (0.448 $\to$ 0.512). 반면, 스트레스 SFT 는 목표 작업 (0.420) 과 유지 작업 (TruthfulQA/MMLU 가 크게 하락) 모두를 열화시켰으며, 이는 SFT 가 본질적으로 파괴적인 것이 아니라 모델의 자연스러운 상태 분포와 불일치하는 오프-정책 압력 하에서 파괴적으로 변한다는 것을 보여줍니다.
2. OPD 는 열화된 교사를 능가함: "스트레스 SFT" 모델을 교사로 사용하여 OPD 로 훈련된 학생은 모든 지표에서 교사를 능가했습니다 (GSM8K: 0.466 대 0.420; MMLU: 0.430 대 0.364). 이는 학생이 자신의 상태에서 훈련되고 국소적 조건부 지도만 받을 경우, 교사의 전역 궤적 실패를 물려받지 않는다는 것을 확인시켜 줍니다.
3. RL 은 국소성을 유지함: 경량 온-정책 RL 은 GSM8K 를 향상시켰으며 (0.472), 베이스 수준의 유지 능력을 거의 유지했습니다. 이는 RL 이 학습자가 도달 가능한 상태에 기반한 국소 개선 절차로 작용한다는 관점을 지지합니다.
4. 드리프트 지표는 불충분함: 거의 동일한 스칼라 상태 드리프트 (MMD) 를 가진 방법들이 완전히 다른 유지 결과를 보였습니다. 예를 들어, 스트레스 SFT 와 스트레스 교사로부터의 OPD 는 유사한 MMD 값 (~0.0109) 을 보였지만 유지 비율은 크게 달랐습니다 (0.82 대 0.95). 이는 분포 이동의 크기보다 훈련 상태의 소스와 국소성이 더 중요하다는 것을 나타냅니다.

기여

1. 통합 프레임워크: SFT, RL, OPD 를 공통의 "상태 분포 관점" 하에 공식화하여 상태 소스를 신호 소스로부터 명시적으로 분리했습니다.
2. 통제된 경험적 파이프라인: 상태 분포의 효과를 분리하기 위해 목표 정확도, 유지, 망각, 롤아웃 상태 드리프트를 측정하는 단일 GPU 실험 설정을 구현했습니다.
3. 상태 중심 주장에 대한 증거: 다음에 대한 경험적 증거를 제공했습니다:
   • OPD 는 상태 샘플링과 신호 생성을 분리함으로써 성능을 회복하고 열화된 교사를 능가할 수 있습니다.
   • 온-정책 RL 은 국소성으로 인해 최소한의 유지 손실로 목표 성능을 향상시킵니다.
   • 스칼라 드리프트 지표 (MMD 등) 는 상태 소스를 고려하지 않으면 포스트-트레이닝 역동성을 완전히 설명할 수 없습니다.

의의 및 주장

이 논문은 포스트-트레이닝 행동이 감독 신호와 훈련 상태 분포 간의 상호작용에 의해 지배된다고 주장합니다. 저자들은 훈련 상태의 "소스, 국소성, 그리고 학습자 의존성"이 감독 신호의 형태 자체만큼이나 중요하다고 논합니다.

• SFT 재정의: 치명적인 망각은 단순히 최대 가능도의 함수가 아니라, 밀집된 오프-정책 업데이트가 모델을 자신의 미래 상태 분포와 양립할 수 없는 영역으로 밀어 넣을 때 발생합니다.
• 증류 재정의: 학생이 교사를 능가할 수 있는 이유는 교사의 오류가 종종 교사의 자체 궤적 분포와 결합되어 있기 때문입니다. 학생의 상태에 대해 교사에게 질의함으로써 학생은 전역 궤적 실패를 물려받지 않고 유용한 국소적 수리에 접근할 수 있습니다.
• RL 재정의: RL 의 안정성은 부분적으로 온-정책 성질에 기인하며, 이는 업데이트가 모델이 실제로 도달하고 회복할 수 있는 상태에만 적용되도록 보장합니다.

이 논문은 목표 수준 분석만으로는 설명할 수 없는 안정성과 성능 역전을 설명하기 위해 상태 중심 관점이 필요하다고 결론지었습니다. 저자들은 이러한 발견이 소규모 연구에 기반한 것이므로 결정적인 벤치마크 주장보다는 기계적 가설에 대한 증거로 간주되어야 한다고 강조합니다.