기술 요약: 연속적 잠재 추론을 위한 Dropout-GRPO

1. 문제 정의

본 논문은 **그룹 상대 정책 최적화(Group Relative Policy Optimization, GRPO)**와 연속적 잠재 추론 모델(특히 COCONUT 아키텍처) 사이의 근본적인 구조적 불일치 문제를 다룹니다.

• GRPO의 요구사항: GRPO는 그룹 내 $K$개 롤아웃(rollout)의 다양성에 의존하여 비제로 그룹 평균 어드밴티지($A^{(k)} = r^{(k)} - \mu_r$)를 계산합니다. 만약 모든 롤아웃이 동일한 궤적을 생성하면, 어드밴티지가 0으로 수렴하여 최적화가 중단됩니다.
• 잠재 추론의 과제: 토큰 수준의 확률적 샘플링을 통해 다양성을 유도하는 표준 Chain-of-Thought(CoT) 모델과 달리, 연속적 잠재 추론 모델(COCONUT 등)은 결정론적인 은닉 상태 재귀(deterministic hidden state recurrence)를 통해 작동합니다. 주어진 프롬프트와 고정된 파라미터에 대해, 여러 번의 롤아웃은 동일한 잠재 궤적과 정답을 생성합니다.
• 결과: 이러한 모델에 표준 GRPO를 적용하면 제로 분산($\sigma_r = 0$)이 발생하여, 정책 그래디언트(policy gradient)가 소실되고 학습 신호가 생성되지 못합니다.

2. 방법론: Dropout-GRPO

저자들은 변분 추론(variational inference)의 이론적 보장을 유지하면서, **구조적 드롭아웃(structured dropout)**을 통해 잠재 추론 과정에 필요한 확률성을 주입하는 Dropout-GRPO를 제안합니다.

핵심 메커니즘: 공유 마스크 변분 드롭아웃 (Shared-Mask Variational Dropout)

토큰 수준의 샘플링 대신, 이 방법은 파라미터 수준에서 확률성을 도입합니다:

1. 마스크 생성: 각 롤아웃 $k$에 대해, 단일 베르누이 마스크 $\xi^{(k)}$를 추출합니다.
2. 일관성: 이 마스크는 해당 롤아웃 내의 모든 $T$개 잠재 재귀 단계 동안 일정하게 유지됩니다.
3. 섭동(Perturbation): 마스크는 네트워크 가중치에 구조적 섭동 $\tilde{\theta}(\xi)$를 유도합니다. 이에 따라 잠재 재귀는 $h_t = f_{\tilde{\theta}(\xi)}(h_{t-1})$가 됩니다.
4. 마스크 재생(Mask Replay): 결정적으로, 롤아웃 단계 중에 마스크를 위한 RNG 시드가 저장됩니다. 정책 업데이트 중에 동일한 마스크가 재생성되어 적용됩니다. 이를 통해 보상이 평가된 궤적이 $\theta = \theta_{old}$일 때 그래디언트 계산에 사용된 궤적과 **비트 단위로 동일함(bit-identical)**을 보장합니다.

이론적 해석

• 베이지안 모델 평균(Bayesian Model Average): Gal와 Ghahramani의 이론을 따라, 공유 마스크는 각 롤아웃을 파라미터에 대한 구조적 분포로부터 추출된 단일 사후 샘플(posterior sample)로 취급합니다. 한계 정책(marginal policy) $\bar{\pi}_\theta$는 베이지안 모델 평균으로 해석됩니다.
• 그래디언트 무편향성(Gradient Unbiasedness): 저자들은 평균 전용 어드밴티지($A^{(k)} = r^{(k)} - \mu_r$)를 사용하는 서로게이트 그래디언트가 $(1 - 1/K)$로 스케일링된 한계 정책의 기대 보상 그래디언트의 무편향 추정치임을 증명합니다.
• 분산 감소: 업데이트 시 롤아웃 단계에서 사용했던 것과 동일한 마스크를 재사용(Common Random Numbers)함으로써, 새로운 마스크를 다시 샘플링할 때보다 분산을 줄입니다.

구현 세부 사항

학습 안정화를 위해 네 가지 특정 개선 사항이 도입되었습니다:

1. Huberized $k_3$ KL: 표준 클램핑(clamping) 대신 Huber 스타일의 확장을 사용하여 KL 발산 추정치를 대체함으로써, 로그 비율이 큰 토큰에서 그래디언트가 소실되는 것을 방지하고 규제 없이 표류하는 "탈출한(escaped)" 토큰을 방지합니다.
2. 신뢰 영역 어닐링(Trust-Region Annealing): KL 계수 $\beta$를 학습률 $\eta_t$에 비례하여 어닐링하여, KL 인력이 정책 그래디언트를 상쇄하여 정책이 평형 상태에 고착되는 것을 방지합니다.
3. 그룹 레벨 정확도 필터: 경험적 정확도 $\mu_r$이 설정된 윈도우(예: $[0.20, 0.95]$) 밖에 있는 그룹은 제외됩니다. 이는 비사소한(non-trivial) 분산을 가진 프롬프트에 계산 자원을 집중시키는 창발적 커리큘럼 역할을 합니다.
4. DDP 협력 스킵(DDP-Coordinated Skip): 그룹을 제외할 때 분산 학습의 일관성을 보장합니다.

3. 주요 기여

1. 새로운 잠재 RL 패러다임: 본 논문은 결정론적 잠재 재귀가 GRPO 실패의 원인임을 식별하고, 공유 마스크 변분 드롭아웃을 최소하면서도 이론적으로 깔끔한 해결책으로 제안합니다. 이는 그룹 상대 RL을 연속적 잠재 추론 모델에 성공적으로 적용한 첫 번째 방법입니다.
2. 이론적 정당성:
   • 표준 트랜스포머 정칙성(regularity) 하에서 잠재 자코비안 $\partial h_T / \partial \theta$의 정의 가능성을 증명합니다.
   • 한계 정책에 대한 서로게이트 그래디언트의 무편향성을 증명합니다.
   • 마스크 재생이 공통 랜덤 숫자(common-random-numbers) 방식의 분산 감소를 제공함을 입증합니다.
3. 실증적 검증: GSM8K 벤치마크에서 COCONUT 모델에 GRPO를 성공적으로 적용하였으며, 이는 결정론적 GRPO가 완전히 실패하는 영역입니다.
4. 참조 구현: 전체 잠재 재귀 과정에서 비트 단위로 동일한 마스크 재생을 실현하는 코드를 공개합니다.

4. 실험 결과

Qwen2.5-1.5B를 백본 모델로 하고 잠재 깊이(latent depth)를 $T=6$으로 설정하여 실험을 진행했습니다.

• 성능 향상: Dropout-GRPO는 GSM8K pass@1 정확도를 COCONUT 베이스라인인 **27.29%**에서 **29.01%**로 향상시켰습니다.
• 능력 회복: 커리큘럼 학습으로 인해 발생한 SVAMP 데이터셋에서의 1.33포인트 성능 저하를 역전시켜, 베이스 모델의 분산과 일치하는 44.00%까지 성능을 회복했습니다.
• 절제 연구(Ablation Studies):
  • Deterministic GRPO: 제로 어드밴티지로 인해 학습 궤적이 SFT 초기값과 동일하게 평탄해지며 완전히 실패했습니다.
  • REINFORCE + Dropout (EMA): 높은 분산으로 인해 실패했습니다. 단일 롤아웃 그래디언트는 SFT 초기값을 안정화하기에 너무 노이즈가 심했습니다.
  • Dropout-GRPO: 유효한 학습 신호를 생성하는 유일한 구성이었으며, 이를 통해 그룹 상대 베이스라인이 이 환경에서 분산 감소에 필수적임을 확인했습니다.

5. 의의 및 주장

본 논문은 이 연구를 잠재 추론 대규모 언어 모델(LLM)의 실용적이고 이론적으로 근거 있는 포스트 트레이닝(post-training) 접근법으로 규정합니다.

• 실현 가능성: 본 연구는 롤아웃의 다양성 부족으로 인해 이전에 차단되었던, 은닉 상태에서 직접 추론하는 모델에 대한 RL 포스트 트레이닝이 가능하다는 것을 보여줍니다.
• 메커니즘: 드롭아웃을 단순한 규제 도구(현대적 LLM에서 종종 비활성화되는)가 아니라, 변분 정책 최적화에 필요한 **구조적 외생적 확률성(structured exogenous stochasticity)**의 원천으로 재정의합니다.
• 한계점: 저자들은 GSM8K에서의 1.72포인트 이득과 같은 향상 폭이 크지 않다는 점에 대해 겸허한 태도를 보이며, 이를 작은 규모의 백본(1.5B)과 적절하고 간결한 단계의 추론 데이터셋의 부족 때문으로 돌립니다. 또한, 배포된 정책(드롭아웃 없는 추론)이 최적화된 한계 정책과 다르며, 이는 립시츠 상수(Lipschitz constants)에 의존하는 편향을 유발할 수 있음을 언급합니다.

요약하자면, Dropout-GRPO는 그룹 상대 강화 학습과 연속적 잠재 추론 사이의 간극을 메워, 결정론적 내부 역학으로 인해 기존에 GRPO를 활용할 수 없었던 모델들의 최적화를 가능하게 합니다.