확률 심플렉스 상의 최적화로서의 디코딩: Top-K, Top-P에서 Best-of-K 샘플러까지

기존 LLM 디코딩 방식들을 통합된 최적화 프레임워크로 재정의하여, 디코딩이 단순한 휴리스틱이 아닌 수학적 원리에 기반한 계층임을 입증했다. 이를 통해 설계된 Best-of-K 샘플러는 추가 학습 없이도 다중 샘플링 환경에서 모델의 추론 성능을 획기적으로 개선할 수 있는 실용적인 도구를 제공한다.

기존의 LLM 디코딩은 모델이 준 점수(Logit)를 바탕으로 토큰을 뽑는 '요리법' 같은 절차적 규칙으로 여겨졌다. 하지만 이 논문은 디코딩을 '모델 점수를 최대화하면서도 특정 구조적 선호도(다양성, 희소성 등)를 만족하는 확률 분포를 찾는 최적화 문제'로 바라본다. 이는 딥러닝에서 손실 함수를 최소화하는 것과 동일한 논리적 구조를 디코딩 단계에 도입한 것이다.

예를 들어, Softmax 샘플링은 단순히 지수 함수를 취하는 것이 아니라, 모델 점수와 엔트로피(무작위성) 사이의 균형을 맞추는 최적화 문제의 유일한 해가 된다. 이 관점에서는 '어떻게 샘플링할 것인가'라는 질문이 '어떤 목적 함수를 최적화할 것인가'로 바뀐다. 즉, 디코딩 방식의 차이는 알고리즘의 차이가 아니라 우리가 디코딩 결과물에 기대하는 수학적 목적의 차이에서 기인한다.

이러한 인식의 전환은 새로운 디코더 설계를 가능하게 한다. 단순히 한 번의 샘플링이 아니라 K번의 샘플링을 할 때, 우리는 '정답이 그 K개 안에 포함될 확률'을 높이고 싶어 한다. 논문은 이 '커버리지' 개념을 목적 함수에 직접 반영하여 Best-of-K라는 새로운 샘플러를 도출했다. 이는 기존의 무작위 샘플링보다 훨씬 전략적으로 후보군을 탐색하게 만든다.

전체 접근 방식은 디코딩을 확률 심플렉스 $\Delta(\mathcal{V})$ 상에서의 정규화된 목적 함수 최대화 문제로 정의하는 것이다. 마스터 목적 함수는 $\max_{q \in \Delta(\mathcal{V})} [\langle q, s_t \rangle - \lambda \Omega(q)]$ 형태를 취하며, 여기서 $s_t$는 모델의 로짓 점수, $\Omega(q)$는 정규화 항, $\lambda$는 정규화 강도를 조절한다.

핵심 메커니즘은 KKT(Karush-Kuhn-Tucker) 조건을 활용한 해의 유도이다. 확률 분포 $q$가 합이 1이고 각 원소가 0 이상이어야 한다는 제약 조건을 라그랑주 승수(Lagrangian Multiplier) $\eta$와 함께 수식화한다. [로짓 점수와 정규화 항의 미분값을 입력으로] → [라그랑주 승수를 포함한 평형 방정식을 계산하여] → [최적의 확률 $q^*(v)$를 산출하고] → [이 값이 특정 디코딩 규칙(예: Softmax)과 일치함을 확인]하는 과정을 거친다.

Best-of-K(BoK) 샘플러는 커버리지 유틸리티 $U_{K,t}(q) = \sum w_t(v) (1 - (1-q(v))^K)$를 정규화 항으로 사용한다. [토큰 $v$가 선택되지 않을 확률 $(1-q(v))^K$를 입력으로] → [1에서 이를 빼서 적어도 한 번 선택될 확률을 구하고 가중치 $w_t(v)$를 곱해 합산하여] → [전체 후보군에 대한 기대 커버리지 점수를 얻고] → [이 점수를 높이는 방향으로 분포를 조정]한다. BoK는 닫힌 형태의 해가 없으므로 Mirror Ascent를 사용하여 반복적으로 업데이트한다.

메인 벤치마크 결과, Qwen2.5-Math-7B 모델을 사용한 MATH500 테스트에서 BoK 샘플러는 높은 샘플링 온도($\tau=0.9$) 기준 정확도를 53.0%에서 71.6%로 18.6%p 향상시켰다. 이는 기존의 Top-K 샘플링(56.2%)보다도 15.4%p 높은 수치로, 다중 샘플링 환경에서 BoK의 압도적인 효율성을 입증한다.

GPQA-diamond와 HumanEval 벤치마크에서도 유사한 경향이 나타났다. GPQA에서는 기본 샘플링 대비 6.06%p, HumanEval에서는 14.64%p의 성능 향상을 기록했다. 특히 모델의 원래 분포가 분산되어 있어 신뢰도가 낮은 고온 샘플링 영역에서 BoK가 유망한 후보군을 효과적으로 포착함으로써 성능 하락을 방지하고 오히려 개선하는 효과를 보였다.

효율성 분석 측면에서, BoK는 토큰당 단 5단계의 Mirror Ascent 업데이트만으로도 충분한 수렴 성능을 보였다. MATH500 실행 시간 측정 결과, 기본 디코딩(15.84초) 대비 BoK(16.88초)는 약 1초 내외의 미미한 오버헤드만을 추가하면서도 성능은 비약적으로 상승하여 실무적인 유용성을 확보했다.

본 연구는 디코딩을 확률 심플렉스 상의 볼록 최적화(Convex Optimization) 문제로 정립한다. 핵심은 정규화 항 $\Omega(q)$의 선택에 따라 디코딩의 기하학적 특성이 결정된다는 점이다. Shannon Entropy를 사용하면 심플렉스 내부에서 부드러운 분포를 갖는 Softmax가 유도되고, Quadratic Penalty를 사용하면 경계면에 도달하여 일부 확률을 0으로 만드는 Sparsemax가 유도됨을 이론적으로 통합했다.

BoK 샘플러의 수학적 기반은 KL Divergence를 활용한 신뢰 영역(Trust Region) 제약과 커버리지 최대화의 결합이다. 목적 함수 $f(q) = \langle q, s_t \rangle - \lambda KL(q \| p_t) + \bar{\beta} U_{K,t}(q)$를 최대화하며, 여기서 $p_t$는 모델의 기본 분포이다. 이는 모델의 원래 지식을 크게 벗어나지 않으면서도 다중 샘플링 시의 효율성을 극대화하도록 설계되었다.

구현 측면에서는 Mirror Ascent 알고리즘을 사용한다. 이는 심플렉스 제약을 자연스럽게 만족시키는 업데이트 규칙 $q_{j+1} \propto q_j \odot \exp(\eta abla f(q_j))$를 제공한다. 수치적 안정성을 위해 Log-Sum-Exp 트릭을 적용하며, 초기값으로 모델의 기본 분포 $p_t$를 사용하여 빠른 수렴을 유도한다. 실험적으로 2~5회의 반복만으로도 최적해에 근접한 성능을 낼 수 있음이 확인되었다.

본 논문은 토큰별(Per-step) 디코딩 최적화에 집중하고 있으며, 전체 시퀀스 수준(Sequence-level)에서 커버리지나 제약 조건을 최적화하는 방향은 향후 과제로 남겨두고 있다.