Jackpot 액터-정책이 안 맞을 때 예산 아끼는 RL

Jackpot 액터-정책이 안 맞을 때 예산 아끼는 RL

 

롤아웃 비용이 RL 학습비의 대부분을 차지하는 LLM 후학습에서, “더 싼 모델로 롤아웃하고 더 센 모델로 업데이트”는 매력적이지만 분포 미스매치로 학습이 쉽게 붕괴합니다. 2602.06107v1(Jackpot)은 사후 중요도보정(IS)만으로 버티기보다, 표본 생성 분포 자체를 OBRS로 p에 가깝게 만드는 정렬 관점을 전면에 둡니다.

OBRS로 분포를 ‘근원에서’ 당기는 접근

이 논문의 문제의식은 매우 현실적입니다. PPO/GRPO류 목적은 원래 rollouts의 추론 분포와 기준 분포가 같다는 가정(pinf=pref)에 기대는데, 실제 시스템에서는 stale/비동기/양자화/대체 모델 롤아웃 등으로 이 가정이 깨지며(actor–policy mismatch), 그 결과 advantage 추정이 틀어지고 학습이 무너질 수 있습니다. 논문은 특히 “완전히 다른 작은 모델로 롤아웃해서 큰 모델을 업데이트”하는 극단 세팅에서 KL 격차가 기존 비동기·양자화보다 훨씬 커져 기존 보정(TIS 등)이 한계에 부딪힌다고 정리합니다.

기존 접근의 중심은 ‘사후 보정’입니다. 예를 들어 truncated importance sampling(TIS)은 pref(x)/pinf(x) 같은 비율을 잘라내 분산 폭발을 막습니다. 그런데 KL이 커지는 구간에서는 “잘라내서 안정화”만으로는 충분하지 않고, 잘라낸 순간 학습 신호가 약해져 느리게 수렴하거나, 결국 붕괴를 막지 못할 수 있습니다. Jackpot은 여기서 방향을 틀어 “사후 보정은 보조로 두되, 미스매치를 원천에서 줄이자”를 제안합니다.

그 핵심이 OBRS(Optimal Budget Rejection Sampling)입니다. 표준 Rejection Sampling(RS)은 q에서 뽑은 토큰 i를 pi/(λqi)로 받아들이고, λ≥max(pi/qi) 조건 때문에 LLM의 10만 단위 vocab에서는 희귀 토큰에서 비율이 튀는 순간 λ가 과대해져 acceptance가 거의 0으로 무너집니다. 논문은 이 “샘플 효율 붕괴”가 LLM RS의 본질적 장애물이라고 짚고, OBRS는 λ를 사용자가 정한 ‘예산’으로 두어 받아들이는 규칙 ai=min(1,pi/(λqi))를 사용한다고 정의합니다. 즉 완전 정렬(정확히 p 샘플) 대신, “주어진 acceptance rate(예산) 내에서 p에 가장 가깝게” 가는 최적 규칙으로 목표를 재정의한 셈입니다.

논문이 설득력 있는 이유는 이론 보장이 전면에 있기 때문입니다. Theorem 3.3은 OBRS 적용 후 분포 q̃가 KL 기준으로 항상 p에 더 가까워진다고 말합니다(DKL(p∥q̃) ≤ DKL(p∥q), 비자명한 예산에서). 또한 Theorem 3.4는 고정된 평균 acceptance 예산 ā에서 KL을 최소화하는 acceptance rule이 OBRS이며, 그 예산을 만족시키는 λ가 유일하게 존재한다고 정리합니다. “우리가 하고 싶은 건 p를 완벽히 모사하는 게 아니라, 예산 안에서 제일 그럴듯한 정렬”이라는 메시지가 수학으로 고정되는 지점입니다.

여기까지가 ‘개념의 승리’라면, Jackpot은 그 개념을 RL 파이프라인에 얹어 “분포 정렬→업데이트 안정화”로 이어지게 만듭니다. 논문은 OBRS를 PPO 목적에 끼워 넣어, 샘플 x는 pinf에서 생성하지만 OBRS 마스크로 토큰을 선택적으로 학습에 반영하고, 동시에 OBRS로 바뀐 분포(p′inf)에 맞춰 reweighting을 수행합니다. 즉 “거절된 토큰은 학습에서 제거(Mask), 받아들인 토큰은 조정된 분포에 맞게 가중치로 보정”이라는 2단 구조입니다. 이는 단순히 ‘필터링’이 아니라, 필터링 후 남은 표본이 어떤 분포에서 왔는지를 명시하고 그에 맞춰 목적식을 다시 정렬한다는 점에서 일관적입니다.

다만 저는 사용자 비평과 같은 지점에서 이 논문의 다음 질문이 보인다고 봅니다. KL이 줄어든다는 보장은 “토큰 분포 수준”의 정렬 보장입니다. RL에서 중요한 것은 시퀀스 리턴과 크레딧 할당인데, 토큰 정렬이 그쪽 안정성(advantage 분산, 리턴 추정 편향)으로 어떻게 전이되는지는 추가 측정이 있어야 더 단단해집니다. 논문이 “OBRS는 안정적 PPO/GRPO 업데이트에 유리하다”라고 말할 때, 그 연결고리를 실험 지표(advantage 분산, ratio clipping 발생률, 업데이트 스텝당 KL drift 등)로 더 직접 보여주면 설득력이 더 올라갈 여지가 큽니다.

논문 핵심 주장	비평 기반 보완 포인트
OBRS로 q를 예산 내에서 p에 가깝게 정렬	KL 개선이 RL 동역학(advantage/return) 안정화로 이어지는 중간 메커니즘 측정이 필요합니다
Top-k 근사+배치 보정으로 Z 계산을 실용화	k·배치 크기·초기 불안정 구간에서의 실패 모드/민감도 표가 더 있으면 좋습니다
정책/롤아웃/증류를 결합한 Jackpot 목적	OBRS 자체 이득과 distill/actor PPO 이득을 분해하는 어블레이션이 더 강해야 합니다

마스킹이 만든 ‘토큰-시퀀스’ 간극을 어떻게 봐야 하나

사용자 비평에서 가장 날카로운 공격 포인트는 “토큰 마스킹이 시퀀스 수준에서 무엇을 최적화하는가”입니다. 논문은 Leviathan et al.(speculative decoding)과 대비해 Jackpot은 거절이 발생해도 suffix를 다시 샘플링하지 않고, “거절된 토큰만 마스크하고 나머지 트레이젝토리는 유지”한다고 강조합니다. 이는 시스템 효율 측면에서 큰 장점입니다. 재샘플링은 곧 추가 롤아웃 비용이기 때문입니다. 논문은 이 선택이 “추가 샘플링 없이” 분포 정렬 효과를 얻는 핵심 설계라고 말합니다.

하지만 효율적 설계가 곧 학습적으로 무해하다는 뜻은 아닙니다. RL에서 보상은 대개 시퀀스 전체(최종 정답, 완결된 코드 실행, 증명 종료 등)로 주어집니다. 그런데 시퀀스 중간에 ‘거절된 토큰’이 실제로는 환경에 그대로 노출되어 리턴을 결정하는데, 학습에서는 그 토큰이 gradient에서 제거된다면 크레딧 할당이 비정상적으로 보일 수 있습니다. 예컨대 수학·코딩처럼 한 토큰이 뒤 결론을 바꾸는 도메인에서는, 거절 토큰이 “실패의 원인”이더라도 학습이 그 원인을 직접 벌하지 못하고, 주변 토큰만 업데이트하는 형태가 될 수 있습니다. 이때 학습은 (1) 문제 토큰을 피하도록 유도되기보다, (2) 우연히 마스크되지 않은 다른 토큰에서 보상 신호를 억지로 설명하려는 방향으로 굴절될 위험도 있습니다.

논문이 제시하는 관점은 “거절된 토큰이 분포 정렬에 해로운 샘플이므로, 업데이트에서 제외하는 것이 오히려 안전하다”는 쪽입니다. 실제로 PPO 목적은 rollouts가 신뢰 가능한 분포에서 왔다는 전제에서 안정적으로 동작하므로, mismatch가 큰 샘플을 제거하면 ratio 폭발과 advantage 왜곡이 줄어들 수 있습니다. 이 논리는 타당합니다. 다만 이 논리가 시퀀스 리턴 문제에서도 안정적으로 성립하려면, 마스킹이 (a) 어디에서 얼마나 자주 발생하는지, (b) 마스킹이 발생한 위치가 성능에 치명적인지, (c) 마스크 비율이 높아질 때 학습이 ‘부분 학습’으로 퇴행하지 않는지 등을 함께 봐야 합니다. 사용자 비평이 요구한 “시퀀스 레벨 분석”이 바로 여기에서 필요합니다.

저는 이 논문이 더 강해지려면 마스킹을 다음처럼 ‘측정 가능한 현상’으로 바꿔 보여주면 좋다고 봅니다.

위치별 마스크 분포입니다: 프롬프트 직후/중간 추론/최종 답 구간 중 어디에서 거절이 집중되는지입니다.

마스크 비율-성능 곡선입니다: 배치 평균 마스크 비율이 높아질수록(=OBRS가 더 공격적으로 작동할수록) 성능과 안정성이 어떻게 변하는지입니다.

“치명 토큰” 민감도입니다: 수학/코드에서 특정 위치(예: 연산자, 괄호, 변수명) 거절이 성능을 더 크게 흔드는지입니다.

이런 지표가 들어가면, Jackpot의 효율 전략(“suffix 재샘플링 없음”)이 단순한 시스템 트릭이 아니라, RL 관점에서도 통제된 설계임을 보여줄 수 있습니다. 지금 논문은 효율성과 KL 정렬을 강하게 말하지만, 마스킹이 시퀀스 리턴 추정에 주는 영향은 상대적으로 짧게 다뤄져 ‘학습 동역학’이 독자의 상상에 남는 부분이 있습니다.

또 한 가지, 논문은 OBRS 후 분포 pOBRS(x)=min(pinf(x), ptarget(x)/λ)/Z 형태를 명시하고, Z가 acceptance rate와 같다는 점을 활용합니다(Z=Σ min(pinf, ptarget/λ)). 이 사실은 마스킹이 단순히 “샘플 제거”가 아니라 “분포를 재정의”한다는 근거입니다. 그래서 마스킹이 학습을 왜곡하는지 판단하려면, 마스크된 토큰이 발생한 시퀀스의 리턴을 어떻게 다루는지(예: 리턴은 그대로 쓰되 토큰별 logprob 항만 제외), 그리고 어드밴티지를 토큰 단위로 분해하는 방식이 어떤지까지 명시하면 좋습니다. 현재는 “마스크 토큰은 loss에서 제외”가 핵심으로 제시되며, 그때 생기는 리턴 신호의 토큰별 분배가 어떻게 작동하는지에 대한 직관을 더 제공할 여지가 있습니다.

실천 관점에서, Jackpot을 비슷한 파이프라인에 적용하려는 팀이라면 다음 체크리스트가 유용합니다.

마스크 비율 상한을 운영 지표로 두는 것이 좋습니다(예: 일정 구간 이상이면 λ를 완화)입니다.

수학/코드 도메인은 “치명 토큰”이 있을 수 있으므로, 토큰 위치/종류별 마스킹 통계 로그가 필요합니다.

마스킹이 과도해지면 학습 신호가 희박해지므로, distill 강도(λdistill)와 함께 폐루프 제어가 필요합니다(논문도 향후 방향으로 closed-loop를 언급합니다).

top-k 근사와 배치 보정, 그리고 ‘OBRS의 순수 효과’ 분해

Jackpot의 설계가 깔끔한 이유는 “이론-알고리즘-시스템-실험”이 한 줄로 이어지기 때문입니다. 시스템 병목은 명확합니다. OBRS의 정규화 상수 Z는 vocab 전체 합이어서(|V|>100k), 배치×시퀀스×보캡 크기의 로그잇 텐서를 저장하는 순간 메모리가 무너집니다. 논문은 이를 정면으로 인정하고, top-k 근사로 Z를 계산합니다. 구체적으로 Vk=top-k(pinf) ∪ top-k(pθnew)로 합집합을 만들고, Zapprox=Σ_{a∈Vk} min(pinf(a|st), ptarget(a|st)/λ)로 근사합니다. 합집합을 쓰는 이유도 타당합니다. 한 분포에서만 큰 토큰이 min 연산에 중요한 영향을 줄 수 있기 때문입니다.

그런데 top-k는 본질적으로 편향이 있습니다. 누락된 토큰은 항이 모두 비음이 아니므로 Zapprox는 항상 과소추정(E[Zapprox] ≤ Z)입니다. 논문은 여기서 매우 영리한 보정을 붙입니다. Z 자체가 평균 acceptance rate(ᾱ)와 같다는 성질을 이용해, 관측된 acceptance rate(α̂¯=accepted/proposed)를 unbiased 추정량으로 얻고, κ=α̂¯ /(배치 평균 Zapprox)로 스칼라 보정 계수를 만들어 각 토큰의 Zapprox에 곱합니다. 즉 “편향이 있지만 낮은 분산의 Zapprox”와 “무편향이지만 분산이 큰 α̂¯”를 결합해 de-biased, low-variance 추정치를 만든다는 설명입니다.

사용자 비평(2)처럼, 여기에는 명확한 실패 모드가 있습니다. 분포가 heavy-tail로 퍼지는 붕괴 직전/초기 불안정 구간에서는 top-k가 tail을 크게 놓칠 수 있고, 그때 Zapprox 과소추정이 커지면 κ가 배치 통계에 과도하게 의존하면서 weight 스케일이 흔들릴 수 있습니다. 특히 배치 크기가 작거나, acceptance event 자체가 희귀해지면 α̂¯의 분산이 커져 κ가 요동칠 수 있습니다. 논문은 “낮은 분산(편향) + 높은 분산(무편향)” 조합이라는 직관을 제시하지만, 독자가 안심하려면 k, 배치 크기, λ(accept budget) 변화에 따른 안정성 민감도 표가 더 있으면 좋습니다. 예를 들어 “k를 늘리면 메모리는 얼마나 늘고, κ의 분산은 얼마나 줄고, 학습 안정성은 얼마나 좋아지는지” 같은 곡선이 있으면, 이 시스템 트릭이 재현 가능한 운영 가이드가 됩니다.

또 하나의 핵심 비평은 (3) “성능 이득이 OBRS인가, distill+actor PPO인가”입니다. 논문은 Jackpot 목적을 세 항의 합으로 둡니다: (i) 정책 모델의 PPO-OBRS, (ii) 롤아웃 모델의 PPO, (iii) 롤아웃이 정책을 따라오게 하는 forward-KL distillation입니다. 수식으로도 LJackpot(θ,ω)=LPPO-OBRS(θ)+LPPO(ω)+λdistill Ldistill(ω)로 명시합니다. 이 구성은 ‘widening gap’을 막는 데 합리적입니다. 롤아웃 모델을 고정하면 정책이 좋아질수록 격차가 커져 mismatch가 더 심해지므로, rollout을 함께 업데이트하고 distill로 붙잡는 것은 자연스럽습니다.

그러나 리뷰어 관점에서 “그럼 distill이 본체 아닌가”라는 질문은 거의 확정적으로 나옵니다. distill은 gap을 직접 줄이기 때문입니다. 논문이 이를 방어하려면, 다음 어블레이션이 매우 강력합니다.

OBRS만 켠 경우(정책 PPO-OBRS만, actor/distill 최소화)입니다.

distill만 켠 경우(OBRS 없이 rollout이 따라오게만)입니다.

actor PPO만 켠 경우(OBRS 없이 rollout을 RL로 업데이트)입니다.

조합별로 KL(rollout–policy), 학습 붕괴 시점, 성능(Mean@k/Pass@k) 곡선을 함께 제시하는 것입니다.

논문은 “Jackpot이 300 step까지 안정” 같은 큰 그림을 보여주고, Table 1에서 극단적 1.7B→8B에서도 Jackpot이 TIS+Reverse KL보다 성능이 좋거나 비슷하다고 제시합니다. 또한 “작은 shift에서는 이득이 작다”는 솔직한 섹션(When Jackpot isn’t Effective)도 제공합니다. 그럼에도 “OBRS 자체의 기여”가 distill/actor 업데이트와 분리되어 얼마나 남는지는 독자가 추정해야 하는 부분이 있습니다. 이 분해가 들어가면 논문이 훨씬 단단해질 것입니다.

마지막으로 “완전 분리의 실용화”에 대한 스코프도 중요합니다. 논문은 Limitations에서 300 step 이후에는 여전히 붕괴할 수 있고, 32B 같은 더 큰 모델 검증은 없다고 명시합니다. 또한 향후 방향으로 distill 강도와 RL 강도를 KL 측정에 따라 조절하는 closed-loop control을 제안합니다. 저는 이 정직함이 오히려 논문의 신뢰를 올린다고 봅니다. 다만 제목과 메시지가 “완전한 decoupling”처럼 읽힐 수 있으니, 독자에게는 “많이 가까워졌지만 아직 불완전”이라는 경계가 더 분명히 각인되도록, 붕괴가 재발하는 조건(예: KL 임계, 마스크 비율 임계, κ 분산 임계)을 정량화해 주면 실용성이 크게 오를 것입니다.

(결론: Jackpot은 OBRS로 오프폴리시 미스매치를 ‘사후 IS’가 아니라 ‘표본 분포 정렬’로 줄이는 관점이 핵심입니다. 다만 토큰 마스킹의 시퀀스 동역학, top-k+배치 보정의 안정성, OBRS·distill·actor PPO의 기여 분해가 보강되면 결론이 훨씬 단단해질 것입니다.)

자주 묻는 질문 (FAQ)

Q. OBRS는 표준 Rejection Sampling과 무엇이 다른가요 A. 표준 RS는 λ≥max(pi/qi) 조건 때문에 LLM의 큰 vocab에서 acceptance가 거의 0으로 무너질 수 있습니다. OBRS는 사용자가 정한 λ(예산)로 ai=min(1,pi/(λqi))를 적용해, “고정 예산에서 KL을 최소화”하는 최적 정렬을 합니다.

Q. Jackpot에서 마스킹은 왜 필요한가요
A. OBRS는 토큰별로 받아들일지 결정하고, 거절된 토큰은 loss/gradient에서 제외해 mismatch가 큰 샘플이 업데이트를 망치지 않도록 합니다. 동시에 남은 토큰은 OBRS 후 분포(p′inf)에 맞게 reweighting해 목적식을 일관되게 만듭니다.

Q. top-k 근사와 배치 보정은 어떤 리스크가 있나요
A. top-k는 Z를 과소추정하는 편향이 있고, 이를 acceptance rate로 κ를 구해 보정합니다. 분포가 heavy-tail로 퍼지거나 배치가 작아 acceptance 통계가 흔들리면 κ가 요동할 수 있어, k·배치 크기·λ 민감도 분석과 실패 모드 보고가 있으면 더 안전합니다.

[출처]
https://arxiv.org/html/2602.06107v1