PreSTO가 LLM 파워 샘플링의 순차 호출을 선반입으로 줄인다

무엇인가

파워 샘플링은 베이스 LLM 분포를 π_α(x) ∝ p0(x)^α로 날카롭게 만들어 테스트 시점 추론 성능을 끌어올린다. 문제는 그 샘플러인 PowerMH의 속도다. Metropolis-Hastings(MH) 갱신이 본질적으로 순차적이어서, k번째 제안의 수락·거부 결과가 k+1번째 제안을 만들 상태를 결정한다. 논문이 인용한 수치에 따르면 MATH500에서 PowerMH는 표준 샘플링보다 최대 8.8배 많은 토큰을 생성하고, LCB V6와 Qwen3.5-9B 조합에서는 K=10, 20, 30일 때 저온 단일 샘플링 대비 프롬프트당 6.5배, 12.2배, 18.4배 오래 걸린다.

어떻게 동작하나

PreSTO는 이 순차성을 없애지 않고 숨긴다. 미래 MH 결과를 수락·거부 결정의 트리로 보고, 한 번의 모델 호출에서 미리 선반입한 서브트리의 제안들을 배치로 생성한 뒤, 실제로 실현된 경로를 원래 MH 수락 규칙으로 따라간다. 핵심은 Condition 1이다. 수락 자식 노드의 컷 인덱스 c′가 부모의 c보다 작거나 같으면, 접두부가 보존되므로 제안 분포가 달라지지 않고 같은 배치에서 생성할 수 있다. BFS로 최대 N_pf개의 실행 가능한 서픽스 리샘플링 요청을 모아 vLLM이나 SGLang 같은 서빙 엔진에 한 번에 보내고, 로그 확률을 받아 MH 로그 수락 확률을 계산한다. 긴 접두부를 공유하는 요청들은 프리픽스 캐싱으로 KV 캐시를 재사용한다. 이후 서브트리 루트에서 시작해 각 노드에서 확률 A(v+,v)로 왼쪽(수락), 아니면 오른쪽(거부) 간선을 따라가며 실현된 리프까지 상태를 갱신한다. 저자들은 Proposition 1로 PreSTO-PowerMH가 PowerMH의 MH 전이 커널을 바꾸지 않음을 증명한다.

무엇과 다른가

여기에 더해 저자들은 MH가 선반입 트리의 희소하고 거의 단일 경로에 가까운 부분만 방문한다는 관찰을 이용한다. 수락 확률이 1-ε 이상이면 ε-확실 수락, ε 이하면 ε-확실 거부(ε=0.01)로 정의하면, LCB V6에서 전체 간선의 약 95%가 ε-확실하다. 기본 BFS는 MH가 거의 방문하지 않는 노드에 예산을 낭비하므로, 컷 비율 c_j/T를 입력으로 하는 3-클래스 로지스틱 회귀가 생성 전에 ε-확실성 클래스 확률을 예측하고 이를 분기 확률로 변환한 뒤, 동적 계획법으로 완료될 전이 수의 기댓값을 최대화하는 요청 N_pf개를 고른다. 이 예측 순회 규칙은 지연된 노드 아래의 요청을 아예 선택하지 않는데, BFS는 N_pf=10, LCB V6, Qwen3.5-9B에서 발행한 요청의 19%를 이런 도달 불가능한 노드에 쓴다.

어떻게 쓰나

실험은 7개 데이터셋과 5개 베이스 LLM, 총 105개 설정에서 PowerMH, EntropyCut, MultiTryMH 각각과 그 PreSTO 변형을 짝지어 진행했다. 각 데이터셋 첫 20개 프롬프트에서 K=100회의 실현된 MH 전이를 수행하고 1,024토큰 응답을 단일 블록으로 생성하며, 같은 vLLM 백엔드와 동일한 베이스 LLM·α·제안 규칙·초기화를 고정했다. PreSTO는 105개 설정 어디에서도 느리지 않았고 96개 설정에서 1.5배 이상을 기록했으며, 중앙값 기준 PowerMH 대비 1.71배, EntropyCut 대비 1.91배, MultiTryMH 대비 2.33배 빨랐다(범위는 각각 1.00~2.29배, 1.19~2.85배, 1.91~7.01배). 모델 호출 수는 1.77~4.92배 줄었다. 이득은 순차 스텝이 비쌀 때 가장 컸다. MBPP와 Qwen3.5-9B에서 EntropyCut 스텝이 7.62초로 PowerMH 스텝 4.97초보다 느려 PreSTO-EntropyCut이 2.85배를 낸 반면 PreSTO-PowerMH는 1.61배였다. 반대로 1.25배 미만인 5개 설정 중 4개(PreSTO-PowerMH)는 호출당 1.96~2.72회 전이만 실현하면서 호출 비용이 PowerMH 스텝의 약 2배였다.

전제와 한계

케이스 스터디는 이득의 출처를 분해한다. LCB V6와 Qwen3.5-9B에서 PreSTO-PowerMH는 539.24±30.92초를 375.12±35.08초로 줄여 1.44배 빨랐는데, 호출당 전이가 1회에서 4.96±0.47회로 늘어난 대신 호출 시간이 5.39±0.31초에서 18.45±0.51초로 길어졌다. 즉 호출 수 감소가 이득의 원천이다. 선반입 예산을 N_pf=8에서 20으로 늘리면 호출당 전이는 3.59회에서 4.91회, 속도 향상은 1.20배에서 1.44배로 커지지만 호출 시간도 16.2초에서 18.4초로 늘어난다. 출력 보존은 K=100 전이 후 터미널 트레이스를 p0로 재채점하고 TOST 동등성 검정(마진 0.2 기준 표준편차)으로 확인해 38개 비교 중 29개에서 동등성을 확립했고, 두 통계량의 통합 차이는 ±0.2 SD 마진 안에 들어갔다(p_eq<0.001). 예측 순회 규칙은 모든 예산에서 BFS보다 1.25~1.37배 빨랐고 100개 프롬프트-예산 쌍 중 98개에서 우세했으며, 호출당 전이가 19~33% 많아 N_pf=20에서 4.91회가 6.54회로 늘었다. KV 캐시 비용은 사전 할당 풀의 최대 3.1% 수준이다.

개발자 관점에서 PreSTO는 학습 없이 추론 지연을 줄이려는 경우, 특히 MH 기반 순차 샘플러를 이미 쓰고 있을 때 의미가 있다. 적용 조건은 두 가지다. 첫째, vLLM이나 SGLang처럼 배치 호출과 프리픽스 캐싱을 지원하는 서빙 엔진이 있어야 한다. 저자들도 접두부 캐시 재사용이 있을 때만 저장된 호출 수가 배치당 추가 비용을 넘어 총 지연이 줄어든다고 명시한다. 둘째, N_pf 예산을 튜닝해야 한다. 예산을 키우면 호출당 전이는 늘지만 실현되지 않은 제안도 모두 생성되므로 호출 시간이 함께 늘어난다. 또한 MultiTryMH의 속도 향상에는 선반입이 아닌 구현상 이득(세 온도에서 후보를 재채점하지 않고 샘플링된 로그 확률을 재사용)이 일부 포함되어 있어, 자기 파이프라인에 옮길 때는 이 부분을 분리해 봐야 한다.

저자들이 밝힌 전제와 한계도 분명하다. PreSTO는 105개 설정에서 결코 느리지 않았지만 하한이 1.00배로, 호출당 실현되는 전이가 적고 호출 비용이 큰 설정에서는 이득이 사라진다. 예산 스윕 표는 예산마다 시드 하나만 사용했고, ε-확실성 비율과 런타임의 연결은 실현된 수락이 아니라 채점된 모든 결정을 세므로 상관관계일 뿐 인과가 아니다. 지연된 요청은 라운드 간에 컷 인덱스를 유지해야 하고, 트리 깊이는 블록에 남은 K-k 전이로 제한되어 블록이 K회를 넘지 않는다는 제약도 따른다.