위상 인식 디코딩과 Wasserstein 정규화 및 질량 페널티

초록

Top‑W는 토큰 임베딩 공간의 거리 정보를 활용해 Wasserstein‑1 거리 기반의 정규화와 엔트로피·질량 균형을 동시에 고려하는 새로운 truncation 규칙이다. 고정된 잠재 함수 f에 대해 최적의 토큰 집합 S는 확률 점수와 f를 합산한 값 ϕ에 따라 정렬한 뒤 접두(prefix) 형태로 찾을 수 있어 O(|V|) 시간에 구현된다. 실험 결과 GSM8K·GPQA·AlpacaEval·MT‑Bench 등 네 가지 베치마크와 3가지 instruction‑tuned 모델에서 기존 최첨단 디코딩 방식보다 최대 33.7% 향상된 정확도와 창의성을 보여준다.

상세 분석

본 논문은 기존의 확률 기반 truncation 기법이 토큰 간 의미적 구조를 무시한다는 점을 지적하고, 이를 보완하기 위해 토큰 임베딩에서 정의된 지상 메트릭(d) 위에 Wasserstein‑1 거리(W₁)를 도입한다. 모델이 출력하는 다음 토큰 분포 p에 대해, 선택된 토큰 집합 S를 기준으로 남은 질량 Γ_S와 재정규화된 분포 q_S를 정의하고, 다음 목적함수를 최소화한다: F_{λ,β}(S)=W₁(p,q_S)+λ·H(q_S)−β·log Γ_S. 여기서 λ는 엔트로피 조절, β는 질량 보존을 위한 페널티를 의미한다.

핵심 이론적 기여는 두 가지이다. 첫째, Lemma 3.1을 통해 W₁(p,q_S)를 (1−Γ_S)·W₁(p(·|S^c),p(·|S)) 형태로 정확히 분해함으로써, 제거된 질량과 남은 질량 사이의 거리 비용을 명시적으로 표현한다. 둘째, 고정된 1‑Lipschitz 잠재 함수 f에 대해 S‑step을 G_f(S)= (∑_{i∈S}p_i ϕ_i)/Γ_S + (β−λ)·log Γ_S 형태의 단일 스칼라 함수 최대화 문제로 변환한다. ϕ_i = f_i + λ·log p_i 로 정의되며, 이때 최적 S는 ϕ_i 내림차순 정렬 후 접두 집합(prefix) 형태가 된다(Theorem 3.4). β≥λ이면 접두 최적화, β≤λ이면 단일 토큰 선택으로 수렴한다는 두 가지 구조적 경우를 제시한다.

실제 디코딩에서는 KR‑dual을 직접 계산하는 것이 비용이 크므로, 저자는 “거리‑to‑set” 함수 dist(i,S)=min_{j∈S}d(i,j)를 이용해 f(i)=−dist(i,S) 로 정의한다. 이 함수는 1‑Lipschitz 조건을 만족하면서 현재 집합 S에 가장 가까운 토큰에 가장 큰 보상을 주어, 의미적으로 군집된 토큰을 보존하고 멀리 떨어진 토큰을 배제한다. 알고리즘 1은 (1) 초기 S를 확률 기반 규칙으로 설정하고, (2) f‑step에서 dist 기반 f를 업데이트한 뒤, (3) ϕ 기반 접두 스캔을 통해 최적 S를 찾는 과정을 T_alt 번 반복한다.

복잡도 측면에서, 토큰 정렬이 O(|V| log |V|) 정도 필요하지만, 실제 구현에서는 상위 몇 백 개 토큰만 고려하도록 k‑NN 프루닝을 적용해 실시간 디코딩에 충분히 빠른(ms/token 수준) 성능을 달성한다.

실험에서는 GSM8K(수학 문제), GPQA(다중 선택 과학 퀴즈), AlpacaEval(지시 따름 평가), MT‑Bench(다중 턴 대화) 네 가지 베치마크와 LLaMA‑2‑7B, LLaMA‑2‑13B, Mistral‑7B 모델을 대상으로 Top‑W와 기존 Top‑k, Top‑p, Top‑H, Min‑p, Locally‑Typical 등 7가지 디코딩 방식을 비교하였다. 정확도 기준에서는 Top‑W가 평균 4.2~7.9%p 상승을 보였으며, 특히 GSM8K와 GPQA에서 33.7%까지 개선했다. 창의성 평가에서는 LLM‑as‑judge가 제시한 주관적 점수에서도 Top‑W가 유의미하게 앞섰다.

한계점으로는 (i) 토큰 임베딩 품질에 크게 의존한다는 점, (ii) β와 λ의 하이퍼파라미터 튜닝이 필요하다는 점, (iii) 현재는 정규화된 거리만 사용했지만 더 복잡한 메트릭(예: 컨텍스트‑의존 거리) 도입 가능성이 남아 있다는 점을 언급한다. 향후 연구에서는 동적 β·λ 스케줄링, 멀티‑헤드 임베딩 결합, 그리고 대규모 멀티모달 모델에의 적용을 제안한다.