📊 statistics

Total Variation Distance Estimation in Autoregressive Models

This paper presents efficient algorithms for estimating the total variation distance between two length- $n$ autoregressive distributions under sample, logit, and noisy logit access models, offering significant improvements over prior methods and demonstrating practical utility in quantifying distributional differences between modern LLM inference engines.

Original authors: Eric Price, Kevin Tian, Zhiyang Xun, Yusong Zhu

Published 2026-07-23

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Eric Price, Kevin Tian, Zhiyang Xun, Yusong Zhu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Invisible Differences Between AI Twins

Imagine you have two identical twins. They were raised in the exact same house, eat the same food, and have the same DNA. If you ask them a question, you expect them to give the exact same answer, word for word. In the world of Artificial Intelligence, specifically Large Language Models (LLMs), the "DNA" is the model's weights—the mathematical brainpower stored in a file. But in the real world, things aren't that simple. Just like twins might react differently depending on whether they are tired, hungry, or talking to a friend, AI models can produce different results depending on the computer hardware they run on, how many questions they are answering at once, or even tiny, invisible glitches in the software code.

To measure how different two things are, scientists often use a tool called "Total Variation Distance" (TV). Think of this as a "mismatch meter." If you have two bags of marbles, and you pull one out of each, the TV distance tells you the maximum chance that the two marbles will be different colors. It's a perfect score for "how likely are these two things to disagree?" This is different from other tools that try to measure "how surprised" a model is, which can break down completely if the models disagree on even a single word. The big question researchers are asking is: If two companies claim to be running the exact same AI model, are they actually giving you the same experience, or are there hidden differences lurking in the code?

The Paper's Mission: Catching the Ghosts in the Machine

This paper, titled "Total Variation Distance Estimation in Autoregressive Models," is a detective story about finding those hidden differences. The authors, a team from the University of Texas at Austin, realized that while two AI engines might look identical on paper, the way they actually generate text can drift apart due to things like "batching" (answering many questions at once) or "quantization" (simplifying numbers to save space). They wanted to build a reliable way to measure the Total Variation Distance between two AI engines to see just how much they disagree.

The authors discovered that measuring this distance is tricky because AI models don't just spit out answers; they build them one word at a time, like a chain reaction. If you only get to see the final answer (like a sample), it's like trying to guess the plot of a movie by watching only the last scene. To solve this, they developed three different "superpowers" or ways to peek under the hood, and they found that the more you can see, the easier the job becomes.

1. The "Sample" Superpower (The Blindfolded Guess)
The most basic way to check an AI is to ask it a question and see what it says. This is called "sample access." The authors found that if you only have this view, you need to ask the AI a lot of questions to get a good measurement. Specifically, the number of questions you need grows with the square of the story length ( $n^2$ ) and the size of the vocabulary's "active" list ( $K$ ). It's like trying to map a huge forest by walking a single path; you have to walk it many, many times to be sure you haven't missed a hidden clearing. Their method improves on previous attempts, making it faster, but it still requires a massive number of queries.

2. The "Logit" Superpower (The X-Ray Vision)
Many AI systems also show their "logits," which are the raw numbers the model uses to decide what word comes next, before it turns them into a final choice. This is like seeing the AI's internal thought process. The authors showed that if you have access to these numbers, the job becomes incredibly easy. You only need a number of queries that grows linearly with the story length ( $n$ ). It's a huge leap forward—like switching from walking through the forest to flying over it in a helicopter. They proved that this is the absolute fastest way to do it; you can't do it any quicker than this.

3. The "Noisy Logit" Superpower (The Foggy Window)
Here is where it gets really interesting. In the real world, those "logit" numbers aren't always perfect. Sometimes the computer is busy, or the software is slightly buggy, and the numbers come back a little fuzzy or "noisy." The authors realized that if you treat this noise as a known quantity (like knowing how foggy your window is), you can still get a great measurement. They created a method that smoothly blends the "blindfolded" and "X-ray" approaches. If the noise is low, you act like you have X-ray vision. If the noise is high, you act more like you are blindfolded. This is the most practical tool because it works even when the AI isn't behaving perfectly.

Putting It to the Test: The Real-World Showdown

To prove their methods work, the authors didn't just stay in the lab. They set up a real-world experiment comparing two popular AI engines, vllm and sglang, running the exact same model (Qwen3-0.6B). They wanted to see if these two engines, which are supposed to be identical, actually produced the same text.

They found that the engines did disagree. The Total Variation Distance between them was about 0.586 for a 500-word story. This number represents the maximum difference in probability between the two engines for any possible outcome. In practical terms, it means that if you were to run a specific test to tell the two engines apart, the best possible test would succeed roughly 59% of the time. It is a measure of how distinguishable the two engines are, rather than a direct probability that any single random sentence generated by one will be different from the other.

The study also revealed why they disagreed. Sometimes, one engine would pick a word that the other engine considered impossible (a "support mismatch"), causing the distance to spike. Other times, they picked the same words but assigned them slightly different probabilities. The authors also noticed that the "noise" in the system changed depending on how the engines were configured. For example, if you asked the engines to answer 256 questions at once, the noise increased, making the measurements fuzzier. But their new "multilevel" method was smart enough to handle this. By asking the same question many times and averaging the results, they could filter out the noise and find the true distance.

The Takeaway

The paper concludes that we can now reliably measure how different two AI engines are, even when they are running on different hardware or using different software tricks. They proved that while it's hard to do this if you only see the final answers, it becomes much easier if you can peek at the internal numbers. Most importantly, they showed that even when those numbers are a bit noisy, we can still get an accurate picture by using their new statistical tricks.

This matters because as AI becomes more common, companies need to know if their "cheap" version of a model is actually the same as the "expensive" one, or if a new, faster software update is secretly changing how the AI behaves. The authors have provided the ruler to measure these invisible differences, ensuring that when we say two AIs are the same, they really are. Their code is now available for anyone to use, turning this complex math into a practical tool for the future of AI.

Technical Summary: Total Variation Distance Estimation in Autoregressive Models

1. Problem Statement

Modern Large Language Model (LLM) deployments often involve implementation choices and inference optimizations (e.g., batching, custom kernels, quantization, speculative decoding) that can cause two engines serving "the same model" to produce meaningfully different output distributions. While prior work has focused on measuring these differences via Kullback-Leibler (KL) divergence, the paper argues that KL is often ill-suited for this task. KL divergence is unbounded when distributions have disjoint supports (a common occurrence due to top- $k$ truncation or bugs), making it impossible to estimate to $\epsilon$ precision from finite samples without support recovery.

In contrast, Total Variation (TV) distance is bounded in $[0, 1]$ , interpretable, and robust to small implementation differences that alter support. The paper addresses the problem of estimating the TV distance between two length- $n$ autoregressive distributions, $\pi$ and $\mu$ , over a token alphabet $\Sigma$ (size $m$ ), to an additive error $\epsilon$ . The study considers three distinct access models to these distributions:

Sample Access: Querying a prefix and receiving a sample from the next-token distribution.
Logit Access: Querying a prefix and receiving the full vector of next-token log-probabilities (exact).
Noisy Logit Access: Querying a prefix and receiving an unbiased estimator of the next-token distribution with bounded relative variance $\sigma^2$ .

2. Methodology and Theoretical Results

The paper provides algorithms and lower bounds for TV estimation under each access model, formalizing the trade-offs between query complexity, sequence length $n$ , support size $K$ , and noise level $\sigma$ .

A. Logit Access (Exact)

Methodology: The authors observe that $O(n)$ queries to prefix logit oracles allow for the exact computation of the likelihood ratio $\pi(x)/\mu(x)$ for any sequence $x$ . They adapt a likelihood-ratio estimator from [CR14], which estimates TV distance by sampling sequences from one distribution and evaluating the likelihood ratio under the other.
Result (Theorem 2): An algorithm exists that estimates TV distance with error $\epsilon$ using $O(n/\epsilon^2)$ queries. This bound is shown to be tight (up to constant factors) via a lower bound construction that reduces TV estimation to learning a hidden Bernoulli parameter $p$ , where revealing one sample of $p$ requires $\approx n$ queries.
Significance: This is the first formalization of TV distance estimation under logit access, offering a linear dependence on sequence length $n$ and eliminating dependence on the alphabet size $m$ or support size $K$ .

B. Sample Access

Methodology: Under sample access, the algorithm cannot directly compute likelihood ratios. The paper proposes an estimator that improves upon the state-of-the-art $O(n^3 m / \epsilon^5)$ bound by [MKP25]. The new approach avoids "taming" distributions (padding with uniform noise) which previously introduced dependence on the full alphabet size $m$ . Instead, it leverages the specific structure of autoregressive models and the maximum support size $K$ of next-token distributions.
Result (Theorem 1): The proposed algorithm estimates TV distance with error $\epsilon$ using $\tilde{O}(n^2 K / \epsilon^2)$ queries. A matching lower bound of $\Omega(n K / \epsilon^2 \log K)$ is provided.
Significance: This improves the dependence on $n$ from cubic to quadratic, on $\epsilon$ from $1/\epsilon^5$ to $1/\epsilon^2$ , and replaces the alphabet size $m$ with the support size $K$ , which is often much smaller in practice (e.g., top- $k$ sampling).

C. Noisy Logit Access

Methodology: Recognizing that real-world inference engines introduce noise (due to batching, kernel variations, or precision issues), the authors model access as an unbiased estimator with relative variance $\sigma^2$ . They employ a Multilevel Monte Carlo (MLMC) variance reduction scheme.

They construct a telescoping sum of estimators $Z_\ell$ at different levels of accuracy (controlled by the number of repetitions $r_\ell$ ).
By coupling the samples across levels (using the same trajectory $X$ for adjacent levels), they ensure the variance of the difference estimator decays geometrically.
This allows trading off cheaper, high-variance queries at coarse levels with expensive, low-variance queries at fine levels.
Result (Theorem 3): The algorithm estimates TV distance with error $\epsilon$ using $\tilde{O}((n + n^2 \sigma^2)/\epsilon^2)$ queries.
Significance: This result smoothly interpolates between the sample access bound (where $\sigma^2 \approx K$ ) and the logit access bound (where $\sigma^2 = 0$ ). It demonstrates that even with noisy oracles, TV distance remains estimable where KL divergence might be infinite.

3. Empirical Evaluation

The authors validate their theoretical findings through experiments comparing two production inference engines, vllm and sglang, serving identical Qwen3-0.6B weights.

Noise Characterization: They measured the relative standard deviation $\sigma$ of log-probabilities across different backends (cuDNN, FlashAttention-2) and batching conditions. They found that while some configurations are deterministic in isolation, serving conditions introduce significant noise.
Distance Estimation: The estimated TV distance between vllm and sglang was found to be substantial (e.g., $\approx 0.58$ at sequence length $n=500$ ). The distance was decomposed into a "shared-support" component (log-probability differences) and a "support mismatch" component (tokens present in one engine's top- $k$ but not the other's), with the latter contributing significantly to the total distance.
MLLM Efficiency: The experiments confirmed the utility of the multilevel schedule. Compared to fixed single-level repetition strategies, the pilot-chosen multilevel schedule achieved the same accuracy with up to $1.28\times$ fewer oracle queries. The results also highlighted that the estimated distance is sensitive to serving configurations (batch size) and prompt content.

4. Key Contributions and Significance

Theoretical Bounds: The paper establishes tight query complexity bounds for TV distance estimation in autoregressive models under three access models, resolving the dependence on sequence length $n$ and support size $K$ .
Robustness to Noise: It introduces a model for noisy logit access that captures real-world inference inconsistencies, showing that TV distance can be estimated even when logit values are non-deterministic.
Practical Relevance: The work provides a tool for auditing LLM providers and verifying that inference optimizations (like quantization or different backends) do not significantly alter the output distribution.
Limitations of KL: The paper reinforces the argument that TV distance is a more practical metric than KL divergence for comparing LLMs, particularly when support mismatches (common in top- $k$ sampling) render KL infinite or unestimable.

The authors conclude that while the TV distance is estimable, the query complexity depends heavily on the noise level and the specific access model, and that the separation between linear and quadratic dependence on $n$ for different access models remains an open question for future work.