기술 요약: 양자화 강화학습을 위한 적응형 중요도 샘플링 (AIS)

1. 문제 제기

대형 언어 모델 (LLM) 에 대한 강화학습 (RL) 은 현재 트래젝토리 롤아웃 생성의 계산 비용에 의해 병목 현상을 겪고 있습니다. 이를 완화하기 위해 최근 산업계 관행 (예: DeepSeek V3.2, veRL) 은 롤아웃 엔진에는 저정밀도 양자화 (예: FP8) 를 적용하고, 더 높은 정밀도 (BF16) 의 트레이너를 유지합니다. 이로 인해 롤아웃–트레이닝 불일치가 발생합니다: 양자화된 롤아웃 정책 ($\hat{\pi}_{\theta_{old}}$) 이 생성하는 트래젝토리는 그래디언트 계산에 사용되는 전체 정밀도 학습자 정책 ($\pi_{\theta_{old}}$) 과 분포상에서 차이가 납니다.

본 논문은 이러한 불일치를 **비정상적 (non-stationary)**이며 "양날의 검"으로 작용한다고 규명합니다:

• 초기 학습: 양자화 노이즈는 확률적 탐색 보너스로 작용하여, 트레이너가 그렇지 않으면 과소 표집했을 트래젝토리에 그래디언트를 노출시킵니다.
• 후기 학습: 정책이 정제됨에 따라, 동일한 노이즈는 불안정성을 유발하는 편향의 원천이 되어 추론 벤치마크에서 학습 붕괴를 초래합니다 (예: Qwen3-8B 의 AIME25 정확도가 6.63% 하락).

기존의 중요도 샘플링 (IS) 은 이러한 편향을 보정하려 시도하지만, 긴 시퀀스에서 과도하게 높은 분산으로 인해 종종 실패합니다. 반면, 절단된 중요도 샘플링 (TIS) 은 학습을 안정화시키지만, 진화하는 불일치의 특성에 적응할 수 없는 고정된 보정 강도를 사용하여 과도한 보정 (유용한 탐색 억제) 또는 보정 부족 (불안정성 허용) 으로 이어집니다.

2. 방법론: 적응형 중요도 샘플링 (AIS)

저자들은 **적응형 중요도 샘플링 (AIS)**을 제안합니다. 이는 배치 단위로 중요도 샘플링 보정의 강도를 동적으로 조절하는 프레임워크입니다. 고정된 임계값 대신, AIS 는 보정되지 않은 그래디언트와 완전히 중요도 가중된 그래디언트 사이를 보간하는 단일 혼합 계수 $\alpha(x) \in [0, 1]$를 계산합니다.

2.1 혼합 계수

AIS 의 핵심은 정책 업데이트 중 이미 계산된 통계량에서 도출된 세 가지 실시간 진단을 결합하는 인수분해된 게이트 메커니즘입니다:

1. 가중치 신뢰도 ($\alpha_{ess}$): 절단된 중요도 가중치의 유효 샘플 크기 (ESS) 비율을 측정합니다. 가중치 분포가 집중되어 (낮은 ESS) 있으면 보정은 신뢰할 수 없으며 하향 조정됩니다.
   $$\alpha_{ess} = (1 + CV^2(\bar{w}))^{-1/2}$$
2. 발산 심각도 ($\alpha_{mis}$): 롤아웃 정책과 트레이닝 정책 간의 평균 절대 로그 확률 불일치를 측정합니다. 보정은 이 불일치가 포화 임계값 $\delta$를 초과할 때만 활성화됩니다.
   $$\alpha_{mis} = \min\left(1, \frac{\bar{d}}{\delta}\right)$$
3. 분산 증폭 ($\alpha_{var}$): 큰 중요도 가중치가 큰 크기의 어드밴티지와 일치하여 허용 가능한 한계를 초과하는 그래디언트 분산을 증폭시키는지 모니터링합니다. 분산 증폭 비율 $\Delta\sigma$가 허용 오차 $\gamma$를 초과하면 페널티가 적용됩니다.
   $$\alpha_{var} = \max\left(0, \frac{\Delta\sigma - \gamma}{\gamma}\right)$$

최종 혼합 계수는 다음과 같이 계산됩니다:
$$\alpha = \text{clip}(\alpha_{ess} - \beta \cdot \alpha_{var}, 0, 1) \cdot \alpha_{mis}$$

2.2 GRPO 와의 통합

AIS 는 그룹 상대 정책 최적화 (GRPO) 에 통합됩니다. 이 방법은 GRPO 목적 함수의 어드밴티지 항을 조정하는 정정 가중치 $\tilde{w}_i$를 도입합니다:
$$\tilde{w}_i = 1 + \alpha \cdot (\min(\rho_{rt}, C) - 1)$$
여기서 $\rho_{rt}$는 롤아웃 - 트레이닝 비율이고 $C$는 절단 임계값입니다. 이 공식은 불일치가 무시할 수 있을 때 ($\alpha \to 0$) 목적 함수가 정확히 일반 GRPO 로 축소되도록 하여 하향 호환성을 유지합니다.

3. 주요 기여

• 비정상성의 특성 규명: 본 논문은 롤아웃 - 트레이닝 불일치가 균일하게 해롭지 않음을 입증합니다. 이는 초기 학습 단계에서는 탐색 보너스를 제공하지만 후기에는 편향의 원천이 됩니다. 이는 양자화 불일치를 단순히 적대적 요소로만 보는 관점에 도전합니다.
• 적응형 보정 프레임워크: AIS 는 정적 임계값을 사용하는 대신 실시간 진단 (신뢰도, 발산, 분산) 에 기반하여 중요도 샘플링 보정의 강도를 동적으로 조절하는 최초의 방법입니다.
• 확산 모델 적용: 이 연구는 확산 기반 LLM(특히 LLaDA-8B-Instruct) 에 대한 RL 의 저정밀도 롤아웃에 대한 최초의 연구를 제시하며, 불일치 문제와 AIS 해결책이 자기회귀 아키텍처를 넘어 일반화됨을 입증합니다.
• 효율성 - 정확성 트레이드오프: 이 방법은 전체 정밀도 BF16 학습의 정확성을 회복하면서도 FP8 롤아웃의 처리량 및 메모리 이점을 유지하는 균형을 달성합니다.

4. 실험 결과

저자들은 수학 추론 (GSM8K, MATH500, AIME25, AMC, Olympiad) 및 계획 (Sudoku, Countdown) 벤치마크에 걸쳐 Qwen3-8B, Qwen3.5-9B(자기회귀), LLaDA-8B-Instruct(확산) 모델에서 AIS 를 평가했습니다.

• 정확도 회복:
  • 자기회귀 모델: Qwen3-8B 에서 단순 FP8 롤아웃은 AIME25 정확도를 BF16 의 36.70% 에서 26.81% 로 떨어뜨렸습니다. AIS 는 이를 **43.33%**까지 회복하여 BF16 기준치를 초과했습니다. Qwen3.5-9B 에서도 유사한 개선이 관찰되어, AIS 는 12 개 벤치마크 쌍 중 11 개에서 BF16 을 일치하거나 초과했습니다.
  • 확산 모델: LLaDA-8B 에서 FP8 롤아웃은 심각한 저하를 초래했습니다 (예: 512 토큰에서 Countdown -19.00%). AIS 는 다양한 시퀀스 길이와 파인튜닝 방식 (Full FT 및 LoRA) 에서 BF16 수준과 일치하거나 초과하는 성능을 회복했습니다.
• 일반 능력: AIS 는 일반 벤치마크 (MMLU, PIQA, ARC 등) 에서 성능을 유지하여 RL 보상 신호에 대한 과적합이 없음을 나타냈습니다.
• 효율성 향상:
  • 가속화: FP8 롤아웃은 배치 크기와 모델 아키텍처에 따라 엔드 - 투 - 엔드 롤아웃 시간을 1.5 배에서 2.76 배 가속화했습니다 (텐서 코어 활용으로 인해 확산 모델에서 더 큰 향상 관찰).
  • 메모리: FP8 롤아웃은 롤아웃 모델 풋프린트를 절반으로 줄였습니다 (예: LLaDA-8B 의 경우 ~7.95 GB 대 ~15.9 GB), 이로 인해 BF16 하에서 이전에는 메모리 부족 (OOM) 이었던 더 큰 배치 크기를 가능하게 했습니다.

5. 중요성 및 주장

본 논문은 AIS 가 확산 LLM 에 대한 RL 의 저정밀도 롤아웃에 대한 최초의 연구라고 주장합니다. 그 주요 중요성은 양자화의 효율성 이점을 희생하지 않고 "롤아웃 - 트레이닝 불일치"를 해결하는 데 있습니다.

저자들은 다음과 같이 주장합니다:

1. 제거가 아닌 관리: 불일치는 제거해야 할 것이 아니라 관리해야 합니다. AIS 는 후기 학습 단계의 불안정화 편향을 성공적으로 억제하면서도 초기 단계의 양자화 노이즈가 제공하는 탐색 이점을 보존합니다.
2. 확장성: 이 방법은 계산적으로 효율적이며, 진단이 기존 통계량에 대한 스칼라 축소이므로 추가 오버헤드가 무시할 수 있을 정도로 작습니다.
3. 견고성: AIS 는 다양한 모델 계열 (자기회귀 대 확산) 과 학습 패러다임 (Full FT 대 LoRA) 에서 일반화됩니다.

본 논문은 AIS 가 전체 정밀도 기준치를 일치하거나 초과하는 고통처리량 저정밀도 RL 학습을 가능하게 하여, 리소스가 제한된 환경에서 RL 을 확장하기 위한 실현 가능한 솔루션이라고 결론지었습니다.