확률 심플렉스 상의 최적화로서의 디코딩: Top-K부터 Top-P (뉴클리어스), 그리고 Best-of-K 샘플러까지

Decoding as Optimisation on the Probability Simplex: From Top-K to Top-P (Nucleus) to Best-of-K Samplers

디코딩은 언어 모델과 우리가 이를 활용하는 모든 작업 사이에 위치하지만, 여전히 휴리스틱한 매개변수 조정 작업 정도로 취급되고 있습니다. 우리는 디코딩이 원칙에 입각한 최적화 계층으로 이해되어야 한다고 주장합니다. 즉, 각 토큰에서 모델 점수와 구조적 선호도 및 제약 조건 사이의 균형을 맞추는 확률 심플렉스 상의 정규화된 문제를 푸는 것입니다. 이 단일 템플릿은 탐욕적 디코딩(greedy decoding), 소프트맥스 샘플링, Top-K, Top-P 및 Sparsemax 방식의 희소성(sparsity)을 특수한 경우로 도출해내며, 최적성 조건을 통해 이들의 공통된 구조를 설명합니다. 더 중요한 점은, 이 프레임워크를 통해 관행적 휴리스틱(folklore)에 의존하지 않고도 새로운 디코더를 쉽게 고안할 수 있다는 것입니다. 우리는 다중 샘플 파이프라인(자기 일관성, 재순위화, 검증자 선택)을 목표로 하는 KL 기반 커버리지 목적 함수인 Best-of-K(BoK)를 설계하여 이를 입증합니다. BoK는 고정된 K개의 샘플 예산 내에서 좋은 대안들을 포함할 확률을 최적화하여 경험적 성능을 향상시킵니다. 우리는 이러한 샘플링 기법이 높은 샘플링 온도 조건에서 MATH500 데이터셋을 수행하는 Qwen2.5-Math-7B 모델의 정확도를 예컨대 +18.6%까지 향상시킬 수 있음을 보여줍니다.

Decoding sits between a language model and everything we do with it, yet it is still treated as a heuristic knob-tuning exercise. We argue decoding should be understood as a principled optimisation layer: at each token, we solve a regularised problem over the probability simplex that trades off model score against structural preferences and constraints. This single template recovers greedy decoding, Softmax sampling, Top-K, Top-P, and Sparsemax-style sparsity as special cases, and explains their common structure through optimality conditions. More importantly, the framework makes it easy to invent new decoders without folklore. We demonstrate this by designing Best-of-K (BoK), a KL-anchored coverage objective aimed at multi-sample pipelines (self-consistency, reranking, verifier selection). BoK targets the probability of covering good alternatives within a fixed K-sample budget and improves empirical performance. We show that such samples can improve accuracy by, for example, +18.6% for Qwen2.5-Math-7B on MATH500 at high sampling temperatures.