극단적인 액터‑정책 불일치를 위한 최적 예산 거부 샘플링

초록

Jackpot은 비용이 큰 LLM 롤아웃을 저비용 모델로 대체하면서 발생하는 액터‑정책 분포 불일치를 최적 예산 거부 샘플링(OBRS)으로 직접 줄이는 프레임워크이다. OBRS는 허용 가능한 수용률(예산) 하에서 제안 분포를 목표 분포에 가장 가깝게 변형시키며, 이를 정책 업데이트와 롤아웃 모델의 공동 학습 목표에 통합한다. 이론적 보증과 실험 결과는 수용률이 90 % 이상 유지되는 동시에 KL 발산을 크게 감소시켜, 300 step까지 on‑policy PPO와 동등한 성능을 달성함을 보여준다.

상세 분석

본 논문은 LLM‑RL에서 롤아웃 비용이 전체 학습 비용의 80 % 이상을 차지한다는 현실적인 문제를 출발점으로 삼는다. 기존의 온‑policy PPO는 롤아웃 모델과 정책 모델이 동일해야 한다는 전제를 갖지만, 저비용 모델(양자화, 희소화, 디스토일 등)로 롤아웃을 수행하면 두 분포 사이의 KL 발산이 급격히 커져 학습이 불안정해진다. 기존 연구들은 중요도 샘플링(IS)이나 트렁케이션(IS‑TIS)으로 사후 보정하려 했지만, 액터‑정책 간 KL이 수십 배 차이날 경우 이러한 보정은 수치적 불안정과 높은 분산을 초래한다.

Jackpot은 이러한 한계를 극복하기 위해 “Optimal Budget Rejection Sampling”(OBRS)을 도입한다. OBRS는 전통적인 거부 샘플링이 요구하는 λ≥max_i(p_i/q_i)라는 비현실적인 상한을 완화하고, 사용자가 지정한 평균 수용률 (\bar a)에 맞춰 λ를 조정한다. 수식적으로는 토큰 i에 대해 수용 확률 a_i = min(1, p_i/(λ q_i))를 적용하고, 이를 통해 얻어진 사후 분포 (\tilde q)는 KL(p‖\tilde q) ≤ KL(p‖q) 를 만족한다는 정리를 제시한다(정리 3.3). 또한, 동일한 수용률 제약 하에서 KL을 최소화하는 유일한 규칙이 바로 OBRS임을 증명한다(정리 3.4).

시스템 구현 측면에서는 vocab 규모가 100 k를 초과하는 LLM에서 전체 확률을 직접 계산하는 것이 메모리·연산 비용이 prohibitive하기에, 상위 k 토큰만을 추출해 근사 확률을 사용한다. 배치 수준에서 발생하는 편향을 보정하기 위해 역 KL 손실과 배치‑레벨 보정 계수를 도입해, 실제 학습 파이프라인에 원활히 통합한다.

실험에서는 Qwen3‑8B‑Base를 목표 정책으로, Qwen3‑1.7B‑Base를 롤아웃 모델로 사용해 300 step, 배치 사이즈 64까지 학습하였다. OBRS 기반 Jackpot은 수용률이 90 % 이상 유지되는 동시에 KL 발산을 평균 0.2 → 0.02 수준으로 감소시켰으며, 최종 성능은 온‑policy PPO와 통계적으로 유의미하게 차이가 없었다. 반면, 기존 TIS와 단순 IS는 학습 초기에 KL이 급증하고, 결국 수렴에 실패하거나 성능이 크게 저하되는 현상을 보였다.

핵심 인사이트는 (1) 액터‑정책 불일치를 사전에 완화하면 사후 IS 보정의 필요성이 크게 감소한다, (2) OBRS는 “예산”이라는 직관적인 파라미터 하나로 효율과 정밀도 사이의 트레이드‑오프를 명시적으로 제어할 수 있다, (3) top‑k 근사와 배치 보정을 결합하면 대규모 vocab에서도 실시간 거부 샘플링이 가능해진다. 이러한 설계는 향후 더욱 큰 LLM(수십억 파라미터)에서도 비용 효율적인 RL 파이프라인을 구축하는 데 중요한 토대를 제공한다.