Total Variation Distance Estimation in Autoregressive Models

2026-07-21

AI Summary

This work addresses the problem of estimating the total variation (TV) distance between output distributions of large language models that differ due to implementation-level discrepancies—such as batching, custom operators, or quantization—during deployment. Under three access models for autoregressive generation (sampling, exact logit access, and noisy logit access), the paper proposes efficient estimation algorithms that significantly reduce query complexity: achieving Õ(n²K/ε²) under sampling access, establishing the first tight bound of O(n/ε²) under exact logit access, and introducing a unified smoothing-based interpolation framework for the noisy logit setting. The theoretical guarantees improve upon existing results, and empirical validation on real-world inference engines like SGLang and vLLM demonstrates the method’s effectiveness—even in cases where the KL divergence is infinite, enabling robust TV distance estimation.

Abstract

Modern LLM deployments use a number of implementation choices and inference optimizations (e.g., batching, custom kernels, and quantization) on top of fixed weights, so two engines serving "the same model" can produce meaningfully different distributions. We study the problem of estimating the total variation (TV) distance between two length-$n$ autoregressive distributions to additive error $\varepsilon$, under three access models. (1) Under sample access, we use $\widetilde{O}(n^2 K/\varepsilon^2)$ queries, where $K$ is the maximum support of the next-token distribution. This improves upon the $\widetilde{O}(n^3 m/\varepsilon^5)$-query estimator of Meel et al. (2025), where $m \geq K$ is the total size of the token alphabet. (2) Under logit access, we use $O(n/\varepsilon^2)$ queries, and this is tight. (3) Under noisy logit access, we smoothly interpolate between the above two guarantees: if probability values are given to relative error $σ$, we use $\widetilde{O}((n+n^2σ^2)/\varepsilon^2)$ queries. We complement our theoretical results with an empirical evaluation of our algorithms, for example measuring the distance between SGLang and vLLM serving identical weights. Our experiments highlight the robustness and practicality of estimating the total variation distance, which remains estimable where the KL divergence is infinite. Our code is available at https://github.com/XunZhiyang/llm-tv-estimation.

Problem

• Total Variation Distance
• Autoregressive Models
• Distribution Estimation
• LLM Deployment
• Inference Discrepancy

Innovation

• Total Variation Distance
• Autoregressive Models
• Logit Access
• Sample Complexity
• Distribution Estimation