TVP Loss: The Potential and Improvement of TV Distance in Speculative Decoding

Speculative Decoding accelerates large language model inference by having a draft model rapidly generate candidate tokens, which are then verified in parallel by the target model. Under lossless speculative sampling, the probability that a draft token is accepted equals 1 minus the total variation (TV) distance between the target distribution and the draft distribution. However, existing draft models are mostly trained with cross-entropy (CE) or KL divergence. A new study re-examines the role of TV distance in draft training and proposes an improved TVP loss function, achieving significant gains in acceptance length across multiple architectures.

The Potential of TV Distance in Draft Training

The study first shows that, under matched and well-tuned training conditions, TV distance can match or surpass several published draft training objectives during fine-tuning; when training from scratch, a fixed TV and KL mixture objective can achieve similar results. This means TV distance itself is an underrated loss function choice.

However, TV distance has two problems: first, gradients vanish for under-covered tokens; second, the prefix product in the acceptance length objective suppresses the contribution of later positions. These two issues limit the further application of TV distance.

Total Variation Distance is a measure of the difference between two probability distributions, defined as half the sum of the absolute differences in probability at each point between the two distributions, ranging from 0 to 1. In speculative sampling, it directly characterizes the degree of mismatch between the draft distribution and the target distribution: if the draft distribution exactly matches the target distribution, the TV distance is 0 and all draft tokens are accepted; the larger the TV distance, the higher the probability of rejection. Therefore, the acceptance probability equals exactly 1 minus the TV distance—a property that makes TV distance a natural optimization objective corresponding to the speculative decoding acceptance rate, in contrast to losses such as cross-entropy and KL divergence that only measure distribution discrepancy without directly corresponding to the acceptance rate.

Design of the TVP Loss Function

To address the above problems, the researchers derive TVP, a family of loss functions based on expected acceptance length. Its power parameter can interpolate between coverage-oriented and acceptance-oriented objectives. By annealing the power parameter from zero to one, TVP achieves the highest or tied-highest acceptance length across three draft model architectures and three target models, including longer drafts and block verification scenarios.

Experimental results show that during fine-tuning, TVP improves acceptance length by 9.8% over CE and by 6.2% to 12.9% over the publicly available DFlash2 training scheme. These improvements demonstrate that TVP can train draft models more effectively, thereby enhancing the overall speedup of speculative decoding.

Acceptance length refers to the number of tokens consecutively accepted and adopted by the target model after each draft generation in speculative decoding. Since the target model can verify multiple draft tokens in parallel in each forward pass, the longer the acceptance length, the more effective tokens are covered per verification, and the higher the end-to-end speedup ratio. Acceptance length is not simply linearly related to acceptance rate: it depends on the product of the probabilities that prefix tokens are accepted one by one, so later positions contribute only if all preceding positions are accepted. This prefix dependency structure is precisely the motivation for TVP introducing a power parameter to balance coverage-oriented and acceptance-oriented objectives.

Impact on Speculative Decoding Practice

The speedup of speculative decoding depends directly on the acceptance rate of draft tokens. By optimizing the acceptance length objective, TVP achieves better acceptance rates across multiple model combinations, providing a new training strategy for improving inference speed in practical deployment. The researchers note that TVP's power parameter annealing strategy is simple and effective, and easy to integrate into existing training pipelines.

In the future, as more draft architectures emerge, TVP is expected to become one of the standard loss functions for draft model training.

Background

Block verification is a variant of speculative decoding in which the draft model generates a fixed-length block of tokens at once, and the target model computes probabilities for all tokens in the block in parallel and performs verification and rejection sampling. Compared with token-by-token verification, block verification can make fuller use of the target model's parallel computing capability, but it places higher demands on the coherence among draft tokens, because rejection at any position within the block truncates subsequent acceptance. TVP also achieves the best acceptance length in block verification scenarios, indicating that its loss design can adapt to different verification granularities.

One-line Summary

A new study proposes the TVP loss function, which optimizes draft model acceptance length via power parameter annealing, significantly improving speculative decoding speedup.

Paragraph Summary

The study re-examines the role of total variation (TV) distance in training draft models for speculative decoding, finding that a tuned TV distance can match or surpass objectives such as cross-entropy and KL divergence, but it suffers from two defects: vanishing gradients for under-covered tokens and the prefix product suppressing the contribution of later positions. To this end, the authors derive the TVP loss family based on expected acceptance length, whose power parameter can interpolate between coverage-oriented and acceptance-oriented objectives, and train draft models via an annealing strategy from zero to one. Across three draft architectures and three target models, TVP achieves the highest or tied-highest acceptance length, improving by 9.8% over CE and by 6.2% to 12.9% over the DFlash2 scheme during fine-tuning, and it is easy to integrate into existing training pipelines.