Jackpot: 롤아웃 비용 절감과 극단적 정책-액터 분포 불일치 문제를 Optimal Budgeted Rejection Sampling으로 해결한 혁신적 강화학습 전략

안녕하세요, 여러분! 오늘은 최신 논문 “JACKPOT: Optimal Budgeted Rejection Sampling for Extreme Actor-Policy Mismatch Reinforcement Learning”를 소개해 드릴게요. RL(강화학습)과 LLM(대형 언어모델) 훈련에서 발생하는 고질적인 문제, 바로 ‘액터(rollout 모델)와 정책(policy 모델)의 분포 불일치 문제’를 해결하는 혁신적인 접근법이랍니다.

---

RL-LLM 훈련의 엄청난 비용, 주범은 ‘롤아웃’

지금까지 RL은 LLM의 후학습(post-training)에서 수학, 코딩, 에이전트 작업 등 다양한 분야에서 큰 성과를 냈어요. 하지만 RL 훈련 비용의 약 80% 이상이 바로 ‘롤아웃’ 단계에서 발생합니다. 롤아웃이란 LLM이 자기자신의 정책으로 토큰 시퀀스를 생성하는 과정인데, 이 과정이 너무나도 컴퓨팅 자원을 많이 소모해요.

 

그래서 ‘더 작은 모델을 사용해 롤아웃을 대체하면 어떨까?’는 아주 매력적인 아이디어지만, 문제는 ‘액터와 정책 모델의 분포가 너무 다르면 학습이 불안정해진다’는 겁니다. 기존 방법들(예: truncated importance sampling, TIS)은 이 분포 격차가 작은 경우엔 잘 작동하지만, 모델 크기 차이나 상태가 극단적으로 다르면 실패하기 일쑤였죠.

---

Jackpot: 분포 차이를 ‘뿌리부터’ 잡다

이 논문에서 제안한 Jackpot은 Optimal Budgeted Rejection Sampling (OBRS) 라는 기법을 활용해서 해결합니다. OBRS는

• 전통적인 rejection sampling이 ‘정확한 분포 일치’를 위해 샘플 수용률이 극히 낮아지는 문제를,
• 일정한 ‘budget(수용 확률)’ 내에서 분포 차이를 최적적으로 줄이는 방향으로 완화하였어요.

즉, 완벽한 분포 같음을 포기하되, 유용한 샘플은 많이 받아들이면서 분포차는 가능한 적게 만드는 일종의 ‘분포 조정 샘플링’인 셈입니다.

---

핵심 기술 포인트

1. 분포 격차 직접 축소
   이전 중요도 샘플링(IS) 방식들은 샘플링 후 학습시 보정하는 ‘사후 처리’인데 비해, OBRS는 ‘샘플링 시점’에 샘플을 거절(Reject)하는 방식. 이로써 분포 차이를 티끌부터 줄여버립니다.
2. 수용 확률 예산 설정 가능
   OBRS는 수용률이라는 예산을 설정하고, 일정 수준 이상의 데이터를 보장하면서 분포를 최대한 맞춰줍니다. 그래야 대형 어휘집(vocab 10만개 이상)에서 ‘특정 희귀 토큰 비율 폭주’ 문제를 효과적으로 막을 수 있어요.
3. 롤아웃 모델 동기화
   롤아웃 모델이 훈련 모델을 따라잡도록 역KL-loss로 주기적으로 조정해 분포 격차가 점점 벌어지는 걸 억제합니다.
4. 탑-k 확률 근사와 편향 보정
   어휘 전체에 대해 확률 계산해야 하는 데서 발생하는 메모리 병목문제를 top-k 토큰만 살펴보는 효율적 방법으로 극복하고, 통계적 보정으로 생기는 편향도 똑똑하게 보정합니다.
5. 기존 RL 시스템과도 호환
   특별한 운영체제나 커널 수정 없이, 일반 vLLM 등에 그대로 적용 가능하여 실사용에 용이합니다.

---

기존 연구와 비교해보면?

• 분포 차이 보정 방식의 차이
  기존 truncated importance sampling(TIS) 방식은 분포 차이가 적고, 모델간 차이가 미미할 때 효과적이지만, 이 논문은 RL 훈련에서 아예 ‘작은 크기 모델로 롤아웃’을 수행하며 ‘극심한 분포 차이’를 감당할 수 있게 합니다.
• 안정성과 효율성 측면
  ‘액터 정책 분리가 극단적’인 상황에서 기존 방법들은 학습이 곧 붕괴하지만, Jackpot은 300회 이상 업데이트에서도 안정적으로 훈련을 유지하며, on-policy 학습과 견줄 성능을 보여주었어요.
• 대형 언어모델 훈련의 ‘롤아웃 병목’ 문제 직접 타겟팅
  많은 논문과 시스템들이 롤아웃 속도 향상에 집중하지만, 여전히 정책과 액터 분포 불일치 문제를 ‘사후 보정’에 맡기는 경우가 많았습니다. Jackpot은 이 문제에 ‘근본적이고 이론적인’ 해법을 제시합니다.

---

간단 정리: Jackpot이 우리에게 주는 기술적 가치

항목	기술적 가치
롤아웃 비용 감소	작은 모델로 롤아웃하며 큰 정책 모델 훈련 가능, 비용 최대 절감
분포 맞춤 최적화	OBRS로 샘플링 과정에서 분포 간 차이 최소화
안정된 학습	분포 mismatch로 인한 학습 붕괴 방지, 안정성 극대화
시스템 호환성	기존 RL 프레임워크에 자연스럽게 통합 가능
메모리 효율	top-k 근사 도입으로 대규모 어휘에서 구현 가능

---

마치며

지금껏 RL-LLM 훈련에서 단골 병목이었던 액터-정책 분포 차이 문제를, ‘Optimal Budgeted Rejection Sampling’이라는 한 차원 높은 수학적 최적화 기법으로 풀어낸 점이 Jackpot의 진짜 큰 매력입니다.

 

더 적은 샘플 거절로 더 안정적인 학습을 보장하며, ‘완전 분리된 작은 롤아웃 모델’도 실제로 쓸 수 있다는 점은 LLM RL 훈련 효율화를 고민하는 연구자 및 실무자분들에게 매우 귀중한 기술적 인사이트가 될 거예요.

 

대형 RL-LLM 프로젝트에서 ‘롤아웃 비용 지옥’에 고민이라면 Jackpot 한번 꼭 시도해 보시는 걸 추천드려요!

---

논문 링크 및 코드: GitHub, 웹사이트

읽어주셔서 감사합니다! 다음에도 새로운 AI 연구 소식으로 만나요 :)