기술 요약: Max@K 정책 그래디언트를 위한 어드밴티지 추정치에 관하여

문제 정의

검증 가능한 보상(RLVR)을 활용한 강화 학습은 추론 작업에 대한 대규모 언어 모델(LLM) 사후 학습(post-training)의 지배적인 패러다임입니다. 그러나 이러한 설정은 일반적으로 희소한 결과 기반 보상(예: 정답/오답의 이진 신호)에 의존하며, 이는 거대한 조합 공간에서의 탐색을 어렵게 만듭니다. 엔트로피 정규화나 카운트 기반 보너스와 같은 방법들이 존재하지만, 이들은 세심한 튜닝이 필요한 부가적인 목적 함수를 도입합니다.

보완적인 접근 방식은 pass@K(K개의 샘플 중 적어도 하나가 정답일 확률) 및 그 연속적 일반화인 max@K(K개 샘의 샘플 중 기대 최대 보상)와 같은 추론 시점의 목적 함수를 직접 최적화하는 것입니다. 최근 여러 연구에서 이러한 목적 함수를 위한 비편향 정책 그래디언트(PG) 추정치를 제안했습니다. 그러나 이러한 추정치들 사이의 관계는 여전히 불분명합니다. 구체적으로, 기존 추정치들은 서로 다른 신호, 베이스라인 및 정규화를 사용합니다. 본 연구에서 식별된 핵심적인 문제는, 선도적인 추정치들(예: 기대 개선량(Expected Improvement) 기반)이 PG 추정기로서는 비편향적이지만, **비중심적 어드밴티지(non-centered advantages)**를 생성한다는 점입니다. 중심이 맞지 않는 어드밴티지(평균이 0이 아닌 경우)는 베이스라인이 과소 추정되었음을 의미하며, 이는 최적화 과정을 하위 최적의 행동으로 유도하고 그래디언트 분산을 증가시킬 수 있습니다.

방법론

본 논문은 max@k 목적 함수에 대한 어드밴티지 추정을 교정하기 위해 Leave-Two-Out (L2O) 베이스라인을 도입하는 방법인 MaxPO(Max@K Policy Optimization)를 제안합니다.

1. 기존 추정치 분석

저자들은 이전 연구(예: Walder 및 Karkhanis [54])에서 사용된 기대 개선량(EI) 추정치를 재검토합니다. 저자들은 EI 기반 추정치가 비편향적인 PG 추정기이기는 하지만, 결과적으로 발생하는 어드밴티지 신호 $s(a)$가 (EI 정의의 ReLU 연산으로 인해) 엄격하게 비음수(non-negative)임을 입증합니다. 결과적으로 $E[s(a)] > 0$이며, 이는 어드밴티지가 중심을 이루지 못함을 의미합니다. 이러한 중심 결여는 분산 감소를 위한 여지를 남깁니다.

2. Leave-Two-Out (L2O) 베이스라인

정확히 중심이 맞는 어드밴티지를 달성하면서 PG 비편향성을 유지하기 위해, 저자들은 L2O 베이스라인을 도입합니다.

• 과제: 표준적인 Leave-One-Out (L1O) 베이스라인은 max@K 설정에서 실패합니다. 왜냐하면 샘플 $i$의 보상이 다른 샘플 $j$의 어드밴티지 계산에 암묵적으로 영향을 미칠 수 있기 때문입니다 (샘прав $i$가 max 연산에서 잠재적인 비교 대상 역할을 함). 이러한 의존성은 비편향적인 PG 추정에 필요한 독립성 조건을 위반합니다.
• 해결책: L2O 베이스라인 $b_{-i}^{L2O}$는 타겟 샘플 $i$와 샘플 $j$ 자신을 모두 제외한 비교 집합에 대해 계산된 모든 다른 샘플 $j$의 기대 개선량 점수의 평균으로 계산됩니다.
• 특징:
  • 독립성: 모든 $j$에 대해 비교 집합에서 $i$를 제외함으로써, 베이스라인 $b_{-i}^{L2O}$는 $r_i$로부터 독립적이며, 이를 통해 PG 추정기의 비편향성을 보장합니다.
  • 중심화: 저자들은 실현된 배치 평균 어드밴티지가 모든 샘플링된 배치에 대해 정확히 0임을 증명합니다: $\frac{1}{B}\sum_{i=1}^B (s_i - b_{-i}^{L2O}) = 0$.
  • 효율성: 단순 계산 시 $O(B^3)$이 소요되지만, 저자들은 ReLU 기반 정식화를 사용하여 L2O 베이스라인을 $O(B^2)$ 시간에 계산하는 벡터화된 구현을 도출하여, GPU 가속 LLM 학습에 적합하도록 만들었습니다.

3. 정형적 유한 배치 어드밴티스 (Canonical Finite-Batch Advantage)

본 논문은 $u_i - v_i$로 정의되는 max@K를 위한 정형적 유한 배치 어드밴티스를 도출합니다:

• $u_i$: 응답 $i$를 포함하는 그룹의 조건부 기대 max@K 값.
• $v_i$: $i$를 제외한 그룹들로부터 계산된 leave-one-out 기대 max@K 값.
  저자들은 MaxPO 추정치(EI + L2O)가 이 정형적 신호($s_i - b_{-i}^{L2O} = u_i - v_i$)를 정확히 복구함을 보여줍니다. 이는 기존 추정치들을 "비중심적"(예: EI-only/PKPO)인 것과 "(정규화된) 정형적"(예: MaxPO 및 특정 분석적 pass@k 변형)인 것으로 분류하여 통합된 관점을 제공합니다.

4. LLM 학습으로의 통합

MaxPO는 그룹 기반 RL 알고리즘(예: GRPO)에 통합됩니다. 각 프롬프트에 대해 $G$개의 출력 그룹이 샘플링됩니다. L2O 베이스라인은 각 그룹 내에서 계산되어 시퀀스 수준의 어드밴티스를 생성하며, 이는 이후 토큰 수준의 손실(loss)로 브로드캐스트됩니다.

주요 기여

1. 비중심적 어드밴티스의 식별: 본 논문은 선도적인 max@K PG 추정치들(특히 EI 기반)이 그래디언트 추정기로서는 비편향적이지만, 중심이 맞지 않는 어드밴티스를 생성하여 최적의 분산 특성을 달성하지 못한다는 점을 식별했습니다.
2. Leave-Two-Out (L2O) 베이스라인: non-linear max@K 설정에서 PG 비편향성을 유지하면서도 실현된 배치 어드밴티스가 정확히 0의 평균을 갖도록 설계된 새로운 베이스라인을 제시했습니다.
3. 정형적 통합: 정형적 유한 배치 어드밴티스($u_i - v_i$)를 도출하고 MaxPO가 이 신호를 복구함을 증명함으로써, 기존의 pass@K/max@K 추정치들을 비교하고 분류할 수 있는 통합된 프레임워크를 제공했습니다.
4. 효율적 구현: L2O 베이스라인을 계산하기 위한 $O(B^2)$ 벡터화 알고리즘을 통해 LLM 사후 학습에서의 실질적인 적용을 가능하게 했습니다.

실험 결과

저자들은 통제된 밴딧(bandit)/미로(maze) 실험과 LLM 추론 작업을 통해 방법을 검증했습니다.

• 분산 감소: 밴딧 환경에서 L2O 베이스라인은 원시 EI 추정치에 비해 그래디언트 분산을 크게 감소시켰으며, 그 효과는 액션 공간의 크기가 커질수록 증가했습니다. LLM 학습(Llama-3.2-3B-Instruct)에서 이 방법은 PKPO(EI 전용 베이스라인) 대비 그래디언트 분산 프록시를 77.4% 감소시켰습니다.
• 성능 향상: 다섯 가지 수학 추론 벤치마크(AIME24, AIME25, AMC23, MATH500, Minerva)에서 MaxPO는 PKPO, GRPO, Entropy-Adv를 포함한 강력한 베이스라인들을 능가했습니다.
  • Qwen2.5-Math-7B에서 MaxPO는 PKPO 대비 pass@256을 5.2% 향상시켰습니다.
  • Llama-3.2-3B-Instruct에서 pass@256을 PKPO 대비 2.4% 향상시켰습니다.
• 추론 컴퓨팅 스케일링: MaxPO의 성능 향상은 더 높은 추론 컴퓨팅 수준(pass@k에서 더 큰 $k$)에서 가장 두드러졌으며, 이는 해당 방법이 정책이 추가적인 테스트 타임 샘플을 더 잘 활용하도록 돕는다는 것을 시사합니다.

의의 및 주장

본 논문은 MaxPO가 max@K와 같은 추론 시점 목적 함수를 위한 정책 그래디언트 추정치를 설계하고 비교할 수 있는 원칙적인 기초를 제공한다고 주장합니다. 기존 추정치의 편향을 교정하고 정형적 어드밴티스 형태를 식별함으로써, 본 연구는 추론 모델의 효율성과 안정성을 개선하는 방법을 제시합니다.

저자들은 본 분석이 i.i.d. 샘플링 하의 유한 배치, 온-폴리시(on-policy) 추정에 집중하고 있음을 겸허히 밝힙니다. 또한 오프-폴리시(off-policy) 학습, 적응형 샘플링, 상관관계가 있는 생성 환경과 관련된 한계를 향후 연구 과제로 인정합니다. 나아가, 반복 샘플링 추론 목적 함수의 최적화가 더 나은 최적화를 가져올 수 있다는 점을 강조하며, 이는 기만적인 도움이나 자동화된 문제 해결에 악용될 수 있는 모델의 역량을 강화할 수 있는 광범한 영향력을 가질 수 있음을 언급합니다.