PKTP는 추론 가지치기에서 pass@k를 직접 최적화해 품질과 다양성을 맞춘다

무엇인가

이 논문은 테스트타임 스케일링에서 추론 모델이 한 문제에 대해 N개의 후보 롤아웃을 생성하고 그중 일부만 끝까지 실행하는 상황을 다룬다. 모든 후보를 완주시키는 비용이 크기 때문에, 짧은 접두사 T토큰만 생성한 뒤 k개를 남기고 나머지를 버리는 가지치기가 핵심이다. 기존 품질 기반 방법인 Linear Probe는 성공 가능성이 높은 후보를 고르지만 비슷한 후보끼리 함께 틀릴 수 있고, 다양성 기반 RepExp는 서로 다른 추론 경로를 고르지만 성공 가능성이 낮은 이상치를 고를 수 있다. 저자들은 이 둘을 탐험과 활용의 긴장으로 보고, 유지한 집합 S가 완주 후 적어도 하나의 정답을 낼 확률, 즉 pass@k를 직접 최대화하는 문제로 정식화한다.

어떻게 동작하나

제안 방법 PKTP는 후보의 최종 정답 여부가 접두사 표현과 어떻게 상관되는지를 잠재 가우시안 모델로 학습한다. 각 후보 i의 잠재 점수는 품질 점수 Qθ(h_i), 비슷한 표현을 가진 후보들이 함께 성공하거나 실패하게 만드는 공유 랜덤 효과 h_i^T f, 후보별 잡음 m_i의 합으로 모델링된다. 여기서 h_i는 접두사 숨은 상태의 지수이동평균(EMA) 표현이고, f와 m은 가우시안이다. 점수가 0 이상이면 정답으로 처리하며, 이 모델에서 두 후보의 점수 공분산은 h_i^T h_j가 되어 표현이 비슷한 후보의 결과가 양의 상관을 갖는다. Qθ는 선형 모델 θ^T h+θ0로 두고, 한 문제의 모든 후보 결과를 함께 설명하는 결합 우도를 최대화하도록 학습한다. 결합 우도의 기대값은 닫힌 형식이 없어 2차 근사로 대체한다.

무엇과 다른가

선택 단계에서는 k개 집합 S의 실패 확률 L(S), 즉 S의 모든 후보가 틀릴 확률을 최소화한다. 근사 로그 실패확률은 세 항으로 해석된다. 첫째는 후보별 품질 q_i가 높을수록 작아지는 품질 항이고, 둘째는 표현들이 여러 방향을 덮도록 장려하는 로그 행렬식 항이며, 셋째는 가중 벡터 λ_i h_i가 같은 방향으로 겹치는 집합을 벌하는 항이다. 정확한 최소화는 NP-난해지만, 근사 목적함수는 단조 부분모듈러임을 증명해 탐욕적 선택이 최적의 1-1/e 이상을 보장한다. 탐욕 단계는 현재 집합에 추가했을 때 로그 실패확률을 가장 많이 줄이는 후보를 고르며, 랭크-1 갱신으로 O(k(Nd+d^2)) 시간과 O(Nd+d^2) 메모리로 구현된다. 학습된 잠재 모델은 k와 무관하고 탐욕 선택이 중첩된 집합을 만들어, 추론 시 k를 온라인으로 바꿔도 재학습이 필요 없다. 상관 항을 끄면 Linear Probe, 품질이 같고 벌점이 무시되면 RepExp의 잔차 기준과 연결된다.

어떻게 쓰나

실험은 두 크기 계층의 오픈웨이트 추론 모델 여덟 개에서 진행됐다. 작은 계층은 Qwen3-0.6B, Qwen3-1.7B, OpenMath-Nemotron-1.5B, DeepSeek-R1-Distill-Qwen-1.5B이고, 큰 계층은 Qwen3-8B, OpenMath-Nemotron-7B, DeepSeek-R1-Distill-Qwen-7B, Skywork-OR1-Math-7B다. 두 계층 모두 AIME 2024·2025·2026, HMMT 2024·2025·2026, RIMO-N을 쓰고, 작은 계층은 BeyondAIME, OlymMATH Easy, Omni-MATH 레벨 7·8을, 큰 계층은 AMO-Bench, OlymMATH Hard, Omni-MATH 레벨 8·9를 추가한다. 모델당 약 600문제에서 128회 시도로 전부 맞거나 전부 틀린 문제를 제거했고, 8폴드 교차검증을 20회 반복해 평균과 표준오차를 보고했다. 정답 판정은 \boxed{} 답을 Qwen3-32B 심판이 참조 답과 동등한지 비교하는 방식이다. 주요 설정은 N=128, k∈{1,2,4,8,16}, 접두사 컷오프 T=4000토큰, 전체 예산 B=32768토큰이며, 온도 0.8, top-p 0.95, 표현은 EMA rate 10^-3, RepExp 정규화 τ=1을 썼다.

전제와 한계

주요 결과에서 PKTP는 Linear Probe와 RepExp를 모두 앞섰고, 특히 중간 k에서 우위가 컸다. DeepSeek-R1-Distill-Qwen-1.5B의 k=4에서 PKTP는 Linear Probe보다 3.0점, RepExp보다 3.2점 높았다. Qwen3-0.6B에서는 k≥2인 모든 경우에 두 베이스라인보다 2점 이상 앞섰고, DeepSeek-R1-Distill-Qwen-7B의 k=4에서는 Linear Probe보다 3.3점 높았다. PKTP는 모든 모델에서 k≥2일 때 Random보다도 높았다. 반면 Linear Probe는 k=1에서 강하지만 k가 커질수록 Random에 가까워지고 일부 모델에서는 큰 k에서 Random보다 나빠졌다. RepExp는 k=1에서 Random과 비슷하고 k=16에서 좋아지는 경향이 있었지만, 중간 k에서는 이상치를 선호해 Random보다 나쁜 경우도 있었다. 논문은 Qwen3-8B의 AMO-Bench 문제에서 Linear Probe 상위 16개가 서로 가깝고 모두 틀린 반면 무작위 16개 중 7개가 맞았고, Omni-MATH 문제에서는 RepExp의 4개 선택이 모두 틀린 반면 무작위 4개 중 1개가 맞았다는 예를 든다. 큰 k에서는 PKTP와 RepExp의 성능이 가까워진다.

컷오프 T를 바꾼 실험에서 T≥4000이면 Qwen3-1.7B와 DeepSeek-R1-Distill-Qwen-7B 모두 PKTP가 다른 방법을 앞섰다. 더 늦은 컷오프는 PKTP에 유리했다. Qwen3-1.7B의 pass@4는 T=2000에서 63.5%였고 T=8000에서 65.9%로 올랐으며, DeepSeek-R1-Distill-Qwen-7B는 57.4%에서 61.7%로 올랐다. 같은 실험에서 Qwen3-1.7B의 Linear Probe는 64.2%에서 62.3%로 내려가 Random 아래로 끝났다. 풀 크기 N을 바꾼 실험에서는 Qwen3-1.7B에서 PKTP가 모든 N에서 다른 방법을 앞섰고, DeepSeek-R1-Distill-Qwen-7B에서는 N=80부터 앞섰다. Qwen3-1.7B에서 RepExp는 N=32일 때 62.9%에서 N=128일 때 63.9%로, Linear Probe는 62.1%에서 63.1%로 변했다. DeepSeek-R1-Distill-Qwen-7B에서는 Linear Probe가 58.1%에서 56.4%로 내려가 모든 N에서 Random 아래였고, RepExp는 N=48의 59.5%에서 58.6%로 내려간 반면 PKTP는 58.8%에서 59.7%로 올랐다.

개발자 관점에서 이 방법은 추론 모델을 여러 번 샘플링하되 전체 롤아웃 비용을 줄여야 하는 서빙·평가 파이프라인에 직접 관련된다. 특히 k를 작게 고정하기 어렵고 중간 크기의 후보 풀을 유지해야 하는 상황에서, 품질 점수와 다양성 점수를 별도 가중치로 섞지 않고 pass@k라는 하나의 기준으로 후보를 고를 수 있다는 점이 실무적 의미다. 다만 접두사 표현이 최종 정답 가능성을 얼마나 잘 담는지가 성능을 좌우하므로, 자신의 모델과 과제에서 EMA 표현, 컷오프 T, 풀 크기 N, 심판/검증기 신뢰도를 먼저 확인해야 한다. 학습된 잠재 모델은 k와 무관해 k를 바꿔도 재학습하지 않는다는 장점이 있지만, 문제 분포가 바뀌면 결합 결과를 학습한 데이터의 대표성을 점검해야 한다.

저자들이 밝힌 한계도 분명하다. 큰 k에서는 PKTP와 RepExp가 가까워지는데, 저자들은 접두사 표현에 담긴 정답 신호가 약하기 때문이라고 보고, 부분 롤아웃에서 품질 표현을 어떻게 얻을 수 있는지, 그것이 가능한지 자체를 열린 질문으로 남긴다. 사용한 EMA rate 10^-3은 대략 최근 1000토큰을 평균하므로 초기 토큰의 신호를 잃을 수 있다. 데이터가 제한적이라 h_i는 고정하고 Qθ도 선형으로 두었으며, 학습된 인코더를 쓰는 확장은 데이터 양에 의존한다. 선택 목적함수는 2차 국소 근사라 정확한 항등식이나 경계가 아니고, 정확한 선택은 NP-난해이며 탐욕 보장은 근사 목적함수에 대한 것이다. 실험은 이진 정답 판정과 과제별 Verify를 전제로 하고, 오픈웨이트 수학 벤치마크와 모델당 약 600문제, Qwen3-32B 심판에 한정되어 있다.