Total Variation Distance Estimation in Autoregressive Models

Motivation

• Known: KL divergence는 sequence-level KL이 per-token KL의 합으로 chain rule 분해가 되어 logit 비교만으로 unbiased estimator를 얻을 수 있어, LLM 배포 간 분포 차이를 측정하는 practitioner의 기본 도구로 사용되어 왔다. Conditional property testing 관련 이론과 prefix sampling oracle을 이용한 distribution testing 문헌(Chakraborty et al., Acharya et al., Canonne et al. 등)이 존재하며, Meel et al.(2025)이 동일한 sample-access 문제를 다뤄 O(n^3 m/ε^5) query complexity를 제시한 바 있다.
• Gap: KL divergence는 두 분포의 support가 조금만 달라도(top-k 절단, 특정 단어 검열 등) 발산하여 유한 샘플로 ε 정밀도까지 추정이 불가능하고 pseudocount 등 hyperparameter에 의존적인 반면, TV distance는 [0,1]로 유계이고 해석 가능하지만 KL과 달리 chain rule 분해가 성립하지 않아 sequence 단위 TV를 per-token 통계량의 합으로 표현할 수 없다는 기술적 장벽이 존재했다. 또한 기존 Meel et al.(2025) 결과는 n, ε, 그리고 K 대신 m에 의존하는 등 비효율적이며, logit access 및 noisy logit access(실제 배포 환경에서 로짓이 non-deterministic한 상황)에 대한 TV 추정 문제는 이전에 공식화되거나 tight bound가 제시된 적이 없었다.
• Why: 실제 LLM 서비스는 배칭, 커스텀 커널, 양자화, lossy speculative decoding 등 다양한 구현 선택으로 인해 동일한 가중치를 사용하더라도 서로 다른 출력 분포를 가질 수 있어, 사용자가 저가 제공자와 고가 제공자를 비교하거나 모델 제공자가 새로운 추론 스택이 기존 스택과 얼마나 일치하는지 검증하는 등 실용적 필요가 크며, TV distance는 임의의 판별 테스트(예: 고정 평가지표의 손실 변화)에 직접적인 bound를 제공하는 해석 가능하고 강건한 지표이기 때문에 중요하다.
• Approach: Prefix likelihood-ratio 계산 및 Canonne & Rubinfeld(2014)의 TV 추정기를 재활용하는 logit-access 알고리즘, multilevel Monte Carlo(MLMC) 스타일의 분산 감소 기법을 적용한 noisy-logit 및 sample-access 알고리즘을 설계하고, 각각에 대해 matching lower bound를 갖는 hard instance를 구성하여 이론적 tightness를 증명한 뒤 sglang, vllm 등 실제 추론 엔진에 대한 실험으로 검증한다.

Achievement

1. Sample access 하 개선된 query complexity: prefix sampling oracle만으로 $\widetilde{O}(n^2K/\varepsilon^2)$ query로 TV distance를 추정하는 알고리즘을 제시하여, 기존 Meel et al.(2025)의 $O(n^3m/\varepsilon^5)$보다 n, ε 의존성과 m/K 팩터 측면에서 크게 개선하였다.
2. Logit access 하 tight bound 확립: $O(n/\varepsilon^2)$ query로 TV distance를 추정하는 알고리즘을 제시하고, 동시에 이 query complexity가 상수 배수까지 tight함을 보이는 matching lower bound(hidden parameter p로 인덱싱된 hard instance)를 구성하였다. 이는 logit access 하에서 TV 추정 문제를 최초로 공식화하고 완전히 해결한 결과이다.
3. Noisy logit access 모델 공식화 및 알고리즘: 실제 추론 엔진의 비결정성을 반영한 relative-variance σ 기반 noisy prefix distribution oracle 모델을 정의하고, $\widetilde{O}((n+n^2\sigma^2)/\varepsilon^2)$ query로 동작하며 σ=0일 때 logit access 결과를, K에 대응할 때 sample access 결과를 (log factor까지) 정확히 회복하는 알고리즘을 MLMC 기반 분산 감소로 설계하였다.
4. 실증적 검증: sglang과 vllm 등 실제 production 추론 엔진 간 TV distance를 측정하는 실험을 수행하여, 이론적 알고리즘의 견고성과 실용성을 입증하고, KL divergence 대비 TV distance 추정이 더 강건함을 보였다.

How

• Sample access: prefix sampling oracle에 대한 conditional property testing 기법을 확장하여 sequence-level TV를 추정하되, Meel et al.(2025)의 "distribution taming"(1/m-scaled uniform distribution으로 padding) 단계를 대체하는 대안적 TV 표현식(Lemma 1의 식 (6))을 사용해 unbounded estimator 문제를 회피함.
• Logit access: O(n) 번의 prefix logit query로 임의의 시퀀스 x에 대한 likelihood-ratio π(x)/µ(x)를 정확히 계산할 수 있다는 관찰을 이용해, Canonne & Rubinfeld(2014)의 O(1/ε²) likelihood-ratio 기반 TV estimator를 그대로 재사용함.
• Logit access lower bound: hidden parameter p∈[0,1]로 인덱싱된 hard instance 분포 쌍 π, µ를 구성하여 TV(π,µ) 추정을 p를 ε 정밀도로 학습하는 문제로 환원하고, 이때 Bern(p) 한 번의 toss를 시뮬레이션하려면 ≈n번의 prefix logit query가 필요함을 보여 lower bound를 도출함.
• Noisy logit access: 각 prefix에서 unbiased estimator p̂ (E[p̂]=p, bounded relative variance σ²)만 관측 가능하다고 가정하고, multilevel Monte Carlo(MLMC) 방식의 분산 감소 기법을 적용해 query complexity를 개선함.

Originality

• TV distance는 KL과 달리 chain rule 분해가 되지 않는다는 근본적 기술적 장벽을 정면으로 다루며, autoregressive 구조에서 sequence-level TV를 효율적으로 추정하는 최초의 체계적 프레임워크를 제시함.
• Logit access라는 새로운 접근 모델을 공식적으로 정의하고 이에 대해 상한과 하한이 모두 tight한 완전한 해를 제공한 최초의 연구임.
• 실제 배포 환경의 로짓 비결정성을 모델링하는 noisy prefix distribution oracle(상대 분산 σ 기반)이라는 참신한 access model을 도입하고, 이것이 sample access와 logit access를 매끄럽게 보간(interpolate)하는 통합 프레임워크임을 이론적으로 증명함.
• MLMC 기법을 TV distance 추정이라는 새로운 맥락에 적용하여 기존 Meel et al.(2025) 대비 이론적으로 개선된 query complexity를 달성함.

Limitation & Further Study

• Sample access 하의 상한 $\widetilde{O}(n^2K/\varepsilon^2)$이 tight한지에 대한 matching lower bound가 명시적으로 제공되었는지 불분명하며, logit access처럼 완전한 tightness 증명이 필요함.
• Noisy logit 모델은 각 prefix에서의 샘플이 독립적이라는 가정에 의존하는데, 실제 배포 환경에서 이 독립성 가정이 항상 성립하는지에 대한 검증이 제한적임.
• 실험은 sglang, vllm 등 특정 오픈소스 추론 엔진 조합에 국한되어 있어, 더 다양한 상용 API 제공자 및 다양한 모델 크기·아키텍처에 대한 일반화 검증이 추가로 필요함.
• σ (relative error) 값을 실제로 어떻게 사전에 추정하거나 보정하는지에 대한 practical guideline이 더 구체화될 필요가 있음.

Evaluation

총평: KL divergence의 실용적 한계를 명확히 지적하고 이를 대체할 TV distance 추정 문제를 세 가지 access 모델에서 이론적으로 엄밀하게(특히 logit access에서 tight bound까지) 해결하면서 실제 추론 엔진 실험으로 실용성을 입증한, 이론과 실무를 균형있게 연결한 견고한 연구이다.