기술 요약: Pass@K 정책 최적화 (PKPO)

문제 정의

대규모 언어 모델(LLM)에 적용되는 강화 학습(RL) 알고리즘은 종종 단일 문제에 대해 여러 개의 솔루션 시도($n > 1$)를 샘플링하지만, 개별 샘플의 성능(pass@1)을 기준으로 최적화를 수행합니다. 이러한 접근 방식은 집합의 집합적 효용보다 단일 고보상 샘플의 효용을 우선시하며, 결과적으로 샘플링 능력을 충분히 활용하지 못하게 합니다. 따라서 표준 정책 경사(policy gradient) 방법은 특히 단일 정답 샘플의 확률이 낮은 어려운 과제에서 솔루션 공간을 충분히 탐색하지 못할 수 있습니다. 추론 시간 탐색(inference-time search) 방법들이 존재하기는 하지만, 추론과 모델 업데이트 사이의 결합 관계는 정답은 검증 가능하지만 솔루션이 사전에 알려져 있지 않은 과제에서 단순히 기대 단일 샘플 보상을 최적화하는 것이 하위 최적(sub-optimal)임을 시사합니다.

방법론

저자들은 **Pass-at-k 정책 최적화 (PKPO)**를 제안합니다. 이는 정책에서 추출된 $k$개의 독립적인 샘플에 대한 최대 보상의 기대값인 pass@k(또는 그 연속적 유사체인 maxg@k)를 직접 최적화하는 프레임워크입니다.

핵심 이론적 기여

주요 기여는 pass@k 목적 함수와 그 기울기(gradient)에 대한 **새롭고 분산이 낮으며 편향되지 않은 추정량(unbiased estimators)**을 도출한 것입니다. 이 추정량들은 $n$개의 샘플로 구성된 배치에서 크기가 $k$인 모든 $\binom{n}{k}$ 개의 부분 집합에 대해 단순 추정량들을 평균 내는 방식으로 작동합니다.

1. 이진 보상 (Section 2):

   • 이진 정답 레이블($f \in \{0, 1\}$)에 대해, 저자들은 $c$가 정답 샘플의 개수일 때 pass@k에 대한 편향되지 않은 추정량 $\rho(n, c, k)$를 도출합니다.
   • 그들은 $n \ge k$인 경우에만 편향되지 않은 추정량이 존재함을 증명합니다.
   • 이에 대응하는 편향되지 않은 기울기 추정량을 도출하여, 정답 여부와 성공적인 부분 집합에 포함될 조합론적 확률에 따라 샘플에 가중치를 부여합니다.

2. 연속 보상 (Section 3):

   • 프레임워크는 maxg@k 목적 함수 $E[\max_{i=1}^k g(x_i)]$를 통해 연속 보상 $g(x)$로 일반화됩니다.
   • 보상이 정렬되어 있다고 가정할 때($g_1 \le g_2 \le \dots \le g_n$), 저자들은 이항 계수를 사용하여 정렬된 보상에 가중치를 부여하는 편향되지 않은 추정량 $\rho^{(g)}(n, c, k)$를 도출합니다.
   • 샘플에 대한 가중치 $s_i$는 보상의 선형 결합으로 구성되며, $O(k + n \log n)$ 시간 내에 효율적으로 계산되는 편향되지 않은 기울기 추정량이 구축됩니다.

3. 분산 감소 (Section 4):

   • 편향을 도입하지 않으면서 기울기 추정치의 분산을 줄이기 위해, 저자들은 Leave-One-Out (LOO) 베이스라인을 제안합니다.
   • 저자들은 현재 샘플을 제외한 $k-1$ 크기의 부분 집합의 maxg@k로부터 $k$ 크기의 maxg@k를 빼는 특정 베이스라인 $s^{(loo-1)}_i$를 도입합니다. 이는 $n=k$인 경우로 제한되었던 기존 연구([TZSM25])를 일반화한 것입니다.
   • 이 공식화는 LLM 학습에서 흔히 발생하는 $n > k$인 시나리오에서도 견고한 최적화를 가능하게 합니다.

구현

이 방법은 모든 정책 경사 알고리즘(예: PPO)에서 표준 보상 변환의 "드롭인 교체(drop-in replacement)"로 설계되었습니다. $n$개의 샘플과 그들의 원시 보상이 주어진 배치에 대해, 정책 업데이트 단계에 입력되기 전에 도출된 추정량(제공된 코드의 sloo_minus_one 함수)을 사용하여 보상을 변환합니다.

주요 결과

저자들은 GEMMA2 (2B 및 9B 파라미터)와 LLAMA3.1 (8B 파라미터)를 사용한 토이 실험 및 실제 벤치마크를 통해 PKPO를 검증합니다.

1. 토이 실험:

   • 1D 가우시안 정책 실험에서, PKPO는 높은 $k$에 대해 최적 정책의 평균을 더 "위험 감수적(risk-tolerant)"으로 이동시키는 데 성공했습니다. 이는 모델이 확률은 낮지만 잠재적 보상이 높은 영역을 탐색할 수 있게 합니다.
   • 분산 분석 결과, 제안된 sloo_minus_one 추정량이 단순 파티셔닝이나 베이스라인이 없는 추정량에 비해 기울기 분산을 유의미하게 감소시킴을 확인했습니다.

2. 실제 벤치마크 (MATH, Coding, ARC-AGI-1):

   • 선택적 최적화: 최적화 대상인 $k_{opt}$를 평가 지표인 $k_{eval}$과 동일하게 설정하면 최상의 성능을 냅니다. 예를 들어, pass@8을 위해 최적화하는 것은 pass@1을 최적화하는 것보다 pass@8 평가 점수를 크게 향상시킵니다.
   • 더 어려운 문제 해결: 도전적인 ARC-AGI-1 벤치마크에서 기존의 pass@1 최적화는 누적 해결율을 개선하지 못한 채 정체됩니다. 반면, $k_{opt} > 1$인 PKPO는 학습을 "언블록(unblock)"하여 모든 $k$에 대해 더 높은 해결율과 pass@k 점수를 달성합니다.
   • 애닐링 전략 (Annealing Strategy): 저자들은 $k$를 학습 중에 애닐링하는 것(탐색을 장려하기 위해 높은 $k_{opt}$에서 시작하여 $k_{opt}=1$로 감소)이 pass@1 성능을 강력하게 확보하면서 동시에 상당한 pass@k 이득을 얻는 것을 보여줍니다. 이는 pass@k를 개선하면 pass@1이 저하되는 전형적인 트레이드오프를 피합니다.
   • 비교: PKPO는 $n=k$인 특수 사례([TZSM25]) 및 엔트로피 정규화 베이스라인보다 특히 탐색이 중요한 어려운 과제에서 우수한 성능을 보입니다.

의의 및 주장

본 논문은 PKPO가 단일 샘플 또는 평균이 아닌, 샘플 집합의 집합적 효용을 최적화하기 위한 견고하고 효율적이며 이론적으로 근거 있는 방법을 제공한다고 주장합니다.

• 탐색의 언블로킹 (Unblocking Exploration): 저자들은 개별 샘플의 효용보다 공동 효용(즉, $k$개의 샘플 중 적어도 하나가 성공할 확률)을 우선시함으로써 탐색을 개선한다고 주장합니다. 이는 솔루션이 검색 공간에서 희소한(sparse) 더 어려운 문제를 해결하는 데 매우 중요합니다.
• 유연성: $k=n$(전체 배치 중 최고를 최적화) 또는 $k=1$에 제한되었던 이전 방법들과 달리, PKPO는 임의의 $k \le n$에 대해 견고한 최적화를 가능하게 합니다.
• 실질적 영향: 이 방법은 실무자가 $k_{opt}$를 통해 탐색-활용(exploration-exploitation) 트레이드오프를 조정할 수 있게 하여, 표준 pass@1 최적화로는 접근할 수 없는 더 복잡한 추론 및 코딩 과제를 모델이 해결할 수 있도록 합니다.

이 연구는 기존의 순서 통계량 정책 경사(order-statistic policy gradient) 방법의 일반화이자 분산 감소 기법으로서, 심층적인 탐색이 필요한 작업에 대해 LLM의 RL 학습을 확장할 수 있는 실질적인 경로를 제시합니다.