Total Variation Distance Estimation in Autoregressive Models

AI summary

Overview

Research area: Machine learning theory / distribution testing, with direct application to auditing large language model (LLM) inference engines.

Technical level: Advanced. The paper is primarily a theoretical computer science contribution (query-complexity upper and lower bounds with proofs), but it also contains an empirical LLM-serving case study.

Scope in one sentence: The paper characterizes how many oracle queries are needed to estimate the total variation distance between two length-n autoregressive distributions under sample, logit, and noisy-logit access, and validates the approach by measuring the distance between two production inference engines serving identical weights.

Authors are Eric Price, Kevin Tian, Zhiyang Xun, and Yusong Zhu (University of Texas at Austin); authors are listed alphabetically. Code is available at https://github.com/XunZhiyang/llm-tv-estimation.

What This Paper Is About

Two LLM serving stacks with the same weights — different GPUs, batching, kernels, quantization, or speculative decoding — can produce meaningfully different output distributions, and providers sometimes even serve with bugs. The paper asks how to measure the difference between two such models in a way that is statistically meaningful, using a chosen distance metric, and how many queries to each model are required to do so to a target accuracy.

The authors focus on the total variation (TV) distance rather than the KL divergence, because KL is unbounded when supports are disjoint (small implementation differences such as top-k = 20 versus top-k = 30 or censoring swear words can make it diverge), it lacks a uniform approximation guarantee without pseudocount hyperparameters, and it has no chain-rule-free decomposition. TV is bounded in [0, 1] and remains estimable in exactly these cases.

Key Contributions

1. Sample-access estimator with improved query complexity. An algorithm estimating TV(π, μ) to additive error ε with probability at least 9/10 using O((n²K/ε²) · log²(1/ε)) prefix sampling queries, where K is the maximum support size of any next-token distribution. This improves on the O(n³m/ε⁵) estimator of Meel et al. (2025) in sequence length n, in accuracy ε, and by an m/K factor for sparse support (m ≥ K is the alphabet size).

2. First formalization and tight characterization of the logit-access setting. The authors state they are the first to formalize TV distance estimation under prefix logit access, giving an O(n/ε²) upper bound and a matching Ω(n/ε²) lower bound, tight up to constant factors.

3. A noisy-logit model that interpolates between the other two. A relative-variance model (parameter σ²) for unreliable logits, with an upper bound of O(((n + n²σ²)/ε²) · log²(1/ε)) queries. Setting σ = 0 recovers the logit bound, and the sample-only model is the special case σ² < K, so the noisy-logit result recovers the sampling result and matches the logit result up to a log²(1/ε) factor.

4. An empirical study of real serving stacks. Measurements of log-probability noise across attention backends and an evaluation of TV distance between vLLM and SGLang serving identical Qwen3-0.6B weights, plus synthetic validation with known ground truth.

Main Findings

• Sample access (Theorem 1, informal): O((n²K/ε²)·log²(1/ε)) prefix sampling queries to both oracles suffice for error ε with probability ≥ 9/10; any algorithm achieving the same guarantee requires Ω(nK/(ε² log K)) queries.
• Logit access (Theorem 2, informal): O(n/ε²) queries suffice, and this is tight up to constant factors for any algorithm. The upper bound uses the observation that O(n) prefix logit queries allow exact computation of the likelihood ratio π(x)/μ(x) for any sequence x, letting the authors reuse a TV estimator from Canonne and Rubinfeld (2014) that needs O(1/ε²) likelihood-ratio queries.
• Noisy logit access (Theorem 3, informal): O(((n + n²σ²)/ε²)·log²(1/ε)) queries suffice for error ε with probability ≥ 9/10; the stated lower bound is Ω(n · (1 ∨ σ²)/(ε² log K)). Table 1 lists the noisy-logit lower bound as (n + nσ²)/ε² and the sampling lower bound as nK/ε² with logarithmic factors omitted.
• Noise model is an interpolation (Lemma 4): σ = 0 is exactly the logit query model, and pure sample access corresponds to σ² = Σ_{a∈supp(p)}(1 − p_a) = |supp(p)| − 1 < K.
• One factor of n remains open: The authors state their lower and upper bounds differ by one additional factor of n in the sequence-length dependence.
• A representation change already helps prior work: Using an alternative TV representation (equation 6) avoids the "distribution taming" step caused by unbounded estimators in Meel et al. (2025); this simple change already improves that algorithm to approximately n²K/ε⁴.
• Measured engine noise (Figure 1): Across five Qwen models in bf16 with top-k = 20, cuDNN is noisy even in isolation, while FlashAttention-2 and math are deterministic in isolation and acquire noise only when served (batched with other, varying requests). σ² ≪ K throughout, so the noisy-logit bound (n + n²σ²)/ε² sits far below the sampling bound n²K/ε².
• Denoising separates real differences from oracle noise (Figure 2): TV(auto, cuDNN) decays with repetition count along the self-noise floor (auto dispatches to the cuDNN kernel), while TV(math, FlashAttention-2) converges to a positive value that further repetition does not reduce.
• Production comparison: Section 5 reports the TV distance between vLLM and SGLang serving identical Qwen3-0.6B weights, how the estimate depends on serving configuration, and validation of the multilevel schedule of Algorithm 3 on real noisy oracles. The specific numerical distance values are not reported in the available content.
• Support sparsity in practice: The total support K is somewhat larger than k but the authors find it well below 2k in practice (Figure 7, Appendix B.4). For reference, Qwen has m = 151,643 tokens and recommends sampling with k = 20.
• TV vs KL framing: Bounding TV directly implies bounds on the change to any distinguishing test, such as loss on a fixed evaluation metric, which the authors offer as a practical argument for their choice of metric.

Methodology in Plain English

The authors treat each model as an oracle. In the weakest ("sample") model, you can give a prefix and get one sampled next token. In the strongest ("logit") model, you give a prefix and see the whole next-token probability vector. The middle model, which they argue resembles real LLM APIs, returns an unbiased but noisy estimate of that vector, with noise controlled by a single relative-variance parameter σ².

Their algorithm for the logit case samples whole sequences from one model, scores each sequence under both models by summing n per-token log-probabilities, and averages a clipped likelihood-ratio statistic (1 − μ(x)/π(x))₊, whose expectation equals the TV distance. Because each term lies in [0, 1], Hoeffding's inequality gives the high-probability guarantee. For the sample-only case they use an alternative TV representation that avoids the unbounded estimators requiring "distribution taming," and then layer on a variance-reduction scheme patterned after multilevel Monte Carlo estimation, trading cheaper but noisier queries at coarse levels against costlier, more accurate queries at fine levels in a telescoping sum.

For hardness, they build adversarial pairs of distributions on a binary alphabet in which the answer reduces to learning a hidden Bernoulli parameter p to accuracy ε, arranged so that roughly n prefix queries are needed to simulate a single coin toss, after which roughly 1/ε² tosses are necessary. A second construction embeds a hard finite-domain TV instance into autoregressive distributions that emit deterministic "stay" tokens and escape with probability Θ(1/n); this preserves the TV distance up to constants while diluting information per query, and combined with a classical finite-domain lower bound (Jiao et al., 2018) shows the dependence on support size or noise is unavoidable.

Why This Matters

Research impact: The paper converts an important but loosely specified practical question — "is this inference stack the same model?" — into a distribution-testing problem with matching upper and lower bounds, extends conditional-sampling oracle theory to a new logit and noisy-logit regime, and shows that a widely used practitioner metric (KL) is theoretically the wrong tool in exactly the settings that matter for deployment.

Real-world applications:

• A model provider can test whether a new, faster inference stack matches its previous stack before shipping it.
• A user or enterprise can check how close a cheaper third-party provider is to an expensive reference model.
• Teams applying quantization, KV-cache quantization, or lossy speculative decoding can quantify the distributional drift those optimizations introduce.
• Auditors can detect model substitution or serving bugs even when downstream task accuracy looks unchanged.

Industry relevance: The experiments target production artifacts directly — cuDNN, FlashAttention-2, SDPA's automatic backend selection, vLLM, SGLang, and Qwen weights — and the paper's observation that nondeterminism in real serving makes even a model-vs-itself comparison look different until noise is averaged away is a concrete engineering concern.

Future Directions

• Whether the extra factor of n that separates the lower and upper bounds in the sequence-length dependence is inherent, or whether stronger lower bounds or better algorithms close the gap.
• Whether the separation between sampling access and logit access is fundamental: the authors explicitly leave open whether superlinear-in-n lower bounds or subquadratic-in-n upper bounds can be proven for the prefix sampling oracle model.
• Extending the lower-bound approach to the prefix structure itself, since the finite-domain lower bounds of Canonne and Rubinfeld (2014) do not apply because their hard instances lack prefix structure.
• Removing assumptions in neighboring auditing work, such as bounded estimator variance in KL estimation (Amini et al., 2025) and log-rank assumptions in rank-based detection (Zhu et al., 2025).

Target Audience

Theoretical computer scientists working on distribution testing and conditional sampling oracles; machine learning researchers studying LLM inference and evaluation; and ML infrastructure or evaluation engineers who need a principled, bounded metric for deciding whether two serving stacks produce the same distribution. Readers without a background in query-complexity analysis will find the empirical sections and the TV-versus-KL argument accessible, while the theorems and lower-bound constructions require comfort with probability and statistical estimation.

Authors’ 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.