A speculative decoding acceleration method based on multiple draft blocks and greedy verification

Abstract

This invention relates to the field of speculative decoding technology, specifically disclosing a method for accelerating speculative decoding based on multiple draft blocks and greedy verification. The method includes: S1, given a prefix sequence, generating K candidate draft sequences of length L independently and identically distributed using a draft model; S2, verifying the acceptance probability of each word sub-block in each candidate draft sequence, and taking the longest accepted block obtained during the verification process as the final accepted block; S3, sampling an additional word from the target model or a predefined residual distribution based on the length of the final accepted block; wherein, the distribution of the target model is corrected before the next round of speculative decoding; S4, appending the final accepted block and the sampled additional word to the prefix sequence, returning to S1, and continuing until the sequence end symbol is decoded, outputting the complete decoded sequence. This invention can achieve longer accepted sequences and reduce the number of target model calls, thus improving decoding efficiency.

Description

Technical Field

[0001] This invention relates to the field of speculative decoding technology, and more specifically to a speculative decoding acceleration method based on multiple draft blocks and greedy verification. Background Technology

[0002] Autoregressive language models have achieved state-of-the-art performance in various natural language processing tasks, but due to the sequential nature of their lexical generation, each lexical requires a complete forward propagation through the model, resulting in high inference latency. To alleviate this latency issue, researchers have proposed speculative decoding (SD) techniques. This method uses a small draft model to generate candidate lexicals, which are then selectively validated by a larger target model, thereby improving decoding efficiency through a rejection sampling process while maintaining distribution consistency. Standard SD typically processes only a single draft sequence, resulting in a limited number of acceptable lexicals in each iteration. To address this limitation, recent works have proposed multi-draft SD strategies. Among them, SpecTr is a representative method that utilizes optimal transport (OT) for multi-draft validation.

[0003] In SpecTr, given a prefix, the draft model generates multiple independent and identically distributed draft sequences. For each lexical position, SpecTr uses OT (Optimal Coherence Theory) to compute the optimal coupling between the draft and target distributions. Based on this coupling, the model determines which draft lexicals can be accepted, or whether all candidate lexicals at that position should be rejected. Theoretically, it can be proven that as the number of draft sequences increases, the probability of accepting draft lexicals also increases, thus increasing the average number of acceptable lexicals.

[0004] SpecTr employs a position-by-position verification mechanism, which means that the verification process iterates through each draft position until one of the following conditions is met: either all positions have one accepted word, in which case an additional word will be sampled based on the target distribution; or a certain position has no accepted words, in which case all subsequent words will be discarded and the word at that position will be resampled from the residual distribution.

[0005] However, this position-by-position verification process cannot achieve the optimal expected number of accepted tokens. How to achieve longer accepted sequences and reduce the number of calls to the target model to improve decoding efficiency is a problem that urgently needs to be solved by those skilled in the art. Summary of the Invention

[0006] In view of the above problems, the present invention proposes a speculative decoding acceleration method based on multiple draft blocks and greedy verification, so as to overcome the above problems or at least partially solve the above problems.

[0007] To achieve the above objectives, the present invention adopts the following technical solution:

[0008] This invention provides a speculative decoding acceleration method based on multiple draft blocks and greedy verification, comprising the following steps:

[0009] S1, given a prefix sequence Under the condition of independent and identically distributed generation through the draft model The length of the strip is Candidate draft sequence This forms a candidate set; the current longest accept block length is set. and initialize an empty rejection set. , used to record rejected lexical sub-blocks;

[0010] S2, Traversal A sequence of candidate drafts, for each candidate draft sequence, starting from position Begin validating each word sub-block in the candidate draft sequence, determining whether the word sub-block to be validated exists in the rejection set. If yes, then skip the verification of that word block; otherwise, calculate the acceptance probability of that word block, and reject or accept the word block according to the acceptance probability. If the word block is accepted, then update the longest accepted block length. If the word sub-block is rejected, it is added to the rejection set. The longest accepted block obtained during the verification process will be used as the final accepted block. ;

[0011] S3, based on the final received block The length of the term is determined by sampling an additional term from the target model or a predefined residual distribution. In each round of inference decoding, the distribution of the target model is corrected.

[0012] S4, the block that will be accepted in the end and additional tokens of sampling Append to prefix sequence In the middle, return to S1, and continue decoding until the end-of-sequence symbol is reached, then output the complete decoded sequence.

[0013] Furthermore, in S2, after verifying a certain candidate draft sequence... Subsequently, if all lexical sub-blocks in the candidate draft sequence are accepted, the iterative verification process for subsequent sequences is terminated early; otherwise, the verification of the next candidate draft sequence continues. .

[0014] Furthermore, in S2, the length is Word element blocks The probability of acceptance is The specific calculation formula is as follows:

[0015]

[0016] in, This indicates the next word sub-block to be sampled. Indicates the current word element sub-block, This indicates that the target model generates word sub-blocks. Joint probability distribution at time, This indicates that the draft model generates lexical units. Joint probability distribution at time, This indicates that the target model generates lexical units. The joint probability distribution, This indicates that the draft model generates lexical units. The joint probability distribution.

[0017] Furthermore, in S2, a sub-block of length L The probability of acceptance is The specific calculation formula is as follows:

[0018]

[0019] in, This indicates that the target model generates lexical units. The joint probability distribution, This indicates that the draft model generates lexical units. The joint probability distribution.

[0020] Furthermore, in S3, if the block is ultimately accepted... The length is equal to the length of the candidate draft sequence. Then an additional lexical unit is sampled from the target model. : ,in, This indicates that the target model is in a given word. When, the conditional probability distribution for generating the next word element, Indicates a given prefix, Indicates length is The accepting word block If the block is eventually accepted The length is less than the length of the candidate draft sequence Then from the predefined residual distribution Sample an additional lexical ,in, The sequence of the final accepted blocks.

[0021] Furthermore, in S3, if :

[0022]

[0023] in, Indicates the first in the next round of speculation decoding Each word element, This indicates that the target model generates lexical units. Joint probability distribution at time; This indicates that the draft model generates lexical units. Joint probability distribution at time; Indicates length The accepting word element sub-block, This represents the additional tokens obtained through residual sampling. Indicates the previous round of speculation decoding Each word element; This indicates that the target model generates lexical units. Joint probability distribution at time, This indicates that the draft model generates lexical units. Joint probability distribution at time, Represents any candidate word in the vocabulary;

[0024] like , ;

[0025] in, Indicates that the target model generates the first each word element The probability distribution.

[0026] Furthermore, the predefined expression for the residual distribution is:

[0027]

[0028] in, This indicates that the target model generates lexical units. Joint probability distribution at time, This indicates that the draft model generates lexical units. Joint probability distribution at time, Indicates the preceding Each word element, Indicates the next word to be sampled.

[0029] Furthermore, the objective of the verification process is to maximize the final accepted block length or minimize the cost function, which is expressed as:

[0030]

[0031] in, Describes all that make the marginal distributions respectively and The effective coupling set, express The joint distribution of 1 independent and identically distributed candidate draft sequences. This indicates a concatenation or direct sum of independent and identically distributed components. This represents the block distribution of the target model. Indicates in coupling distribution Lower Acceptable Length Expected value This represents the length of the continuous word block accepted by the target model in a single validation. Indicates from coupling distribution Target sequence block obtained by mid-sampling .

[0032] As can be seen from the above technical solution, compared with the prior art, the present invention has the following beneficial effects:

[0033] This invention replaces the position-by-position verification mechanism in SpecTr with a greedy block verification method, generating multiple independent and identically distributed draft sequences. The verification process is described as an optimal transmission problem between draft lexical blocks and target lexical blocks. By unifying the multi-draft and block verification strategies, a longer receiving sequence is achieved, the number of target model calls is reduced, decoding efficiency is improved, and the greedy block verification is naturally extended from single-draft scenarios to multi-draft scenarios. This invention achieves the optimal expected receiving length in a single iteration with any fixed number of drafts; and this optimal expected value increases with the number of drafts. Attached Figure Description

[0034] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on the provided drawings without creative effort.

[0035] Figure 1 This is a flowchart of the speculative decoding acceleration method based on multiple draft blocks and greedy verification provided in an embodiment of the present invention.

[0036] Figure 2 This is a schematic diagram of the position-by-position verification process of the traditional method SpecTr.

[0037] Figure 3 This is a schematic diagram of the verification process of the method of the present invention. Detailed Implementation

[0038] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0039] like Figure 1 As shown, this invention discloses a speculative decoding acceleration method based on multiple draft blocks and greedy verification, comprising the following steps:

[0040] S1, given a prefix sequence Under the condition of independent and identically distributed generation through the draft model The length of the strip is Candidate draft sequence This forms a candidate set; the current longest accept block length is set. and initialize an empty rejection set. , used to record rejected lexical sub-blocks;

[0041] S2, Traversal A sequence of candidate drafts, for each candidate draft sequence, starting from position Begin validating each word sub-block in the candidate draft sequence, determining whether the word sub-block to be validated exists in the rejection set. If yes, then skip the verification of that word block; otherwise, calculate the acceptance probability of that word block, and reject or accept the word block according to the acceptance probability. If the word block is accepted, then update the longest accepted block length. If the word sub-block is rejected, it is added to the rejection set. The longest accepted block obtained during the verification process will be used as the final accepted block. ;

[0042] S3, based on the final received block The length of the term is determined by sampling an additional term from the target model or a predefined residual distribution. In each round of inference decoding, the distribution of the target model is corrected.

[0043] S4, the block that will be accepted in the end and additional tokens of sampling Append to prefix sequence In the middle, return to S1, and continue decoding until the end-of-sequence symbol is reached, then output the complete decoded sequence.

[0044] In this embodiment of the invention, the draft model can be DeepSeek-1.3B, and the corresponding target model can be DeepSeek-33B or DeepSeek-6.7B; or the draft model can be CodeLlama-7B, and the corresponding target model can be CodeLlama-13B.

[0045] The standard speculative decoding method SD uses only a single draft sequence. To improve acceptance rate, the SpecTr method extends this to a multi-draft setting, generating a sequence given a prefix. This expands the candidate space by generating independent and identically distributed draft sequences. For each position... (i.e., across each column of the draft), the goal is to be in candidate lexical units It accepts a valid token. SpecTr formalizes this verification step as an optimal transport (OT) problem.

[0046] This discrete OT problem can be solved using linear programming, but its running time varies. Exponential growth, therefore in large vocabularies or large This is not feasible under certain circumstances. Therefore, an efficient approximation algorithm is adopted. The algorithm iterates over the draft lexicon at each position and accepts it with the following probabilities:

[0047]

[0048] in The scaling factor is determined by solving the equation. Obtain, among which . Indicates scaling factor Then, the sum of the overlap probability mass between the target model distribution and the draft model distribution, Indicates the draft model in the given prefix At that time, generate the next word. The conditional probability, Indicates that the target model has the same prefix The next word is generated. The conditional probability;

[0049] For each location, Output the first accepted term, and retain only the draft sequence that matches it for the next step. If no candidate is accepted, all subsequent terms are discarded, and the sequence is resampled from the residual distribution.

[0050]

[0051] If all positions have accepted terms, then another term is sampled from the target model. Theoretically, when As the number increases, the acceptance probability rises, resulting in a greater acceleration effect.

[0052] However, the above position-by-position verification terminates when all candidates at a certain position are rejected, and this strategy fails to achieve the maximum expected number of accepted lexical units in a single iteration. Therefore, this invention introduces the Greedy Block Verification (GBV) method. Unlike the standard lexical unit-by-lexical verification method, GBV verifies the entire sub-block... That is, before the draft model is generated Each consecutive word element is evaluated as a whole.

[0053] The verification process of the method of this invention and the traditional SpecTr method is as follows: Figures 2-3 As shown.

[0054] For multiple independent and identically distributed (i.e., i.id) draft sequences, SpecTr uses a position-by-position verification method: such as Figure 2 As shown, at each position, it accepts a word (e.g., bounces, then like) and only retains sequences consistent with the accepted word. If no word is accepted, a new word is sampled from the residual distribution (e.g., replacing ball with toy).

[0055] This invention validates lexical sub-blocks on all draft sequences, such as Figure 3 As shown, the process first jointly validates the sub-blocks starting from the first position (e.g., bounces, bounces when, ..., bounces when I am), and selects the longest accepted sub-block (e.g., bounces). Then, validation continues from the next position in the next sequence (e.g., wraps around, ..., wraps around my arm). If no sub-block is accepted, the process continues to the next sequence, progressively selecting the longest accepted sub-block (e.g., bounces like a). After the longest accepted sub-block, a new term is sampled from the residual distribution (e.g., replacing ball with toy).

[0056] The steps described above in this invention will be further explained below.

[0057] S1. Generate multiple draft sequences: These are generated by a lightweight draft model from a given prefix. Under the condition of independent and identically distributed generation The length of the strip is Candidate draft sequence This forms a candidate set. The current longest accepted block length is set. Initialize an empty collection. This is used to record rejected lexical sub-blocks to avoid redundant subsequent calculations. The generation process of any candidate draft sequence is represented as follows:

[0058]

[0059] in, This represents the conditional distribution of the draft model in a given context. For the first One generated word unit, and These are the previously generated word sequences. These candidates are then processed by the target model. Parallel validation, the target model provides the true conditional distribution. To simplify the notation, let's call it... , , Indicates the draft model in the given prefix When, the conditional probability distribution for generating the next word element, Indicates the target model in a given prefix When generating the next word, the conditional probability distribution is determined.

[0060] S2. Perform joint verification on the word sub-blocks of the draft sequence:

[0061] sequentially for the first arrive A sequence of independent and identically distributed drafts To perform verification, for each sequence, start from position... Start validating the word sub-block. If the word sub-block to be validated... It already exists in the rejection set. If the probability calculation and verification of the word sub-block are skipped, it is considered a rejection, and the verification of the next sub-block continues. .

[0062] The reception probability is calculated using a specific calculation method for the position of each word sub-block. The core reason for using different reception probability calculation formulas for word sub-blocks at different positions is to meet the requirements of optimal transmission, so as to maximize the expected reception length while ensuring the consistency of output distribution.

[0063] For length of The word element sub-block, The probability of acceptance is The specific calculation formula is as follows:

[0064]

[0065] in, This indicates the next word sub-block to be sampled. Indicates the current word element sub-block, This indicates that the target model generates word sub-blocks. Joint probability distribution at time, This indicates that the draft model generates lexical units. Joint probability distribution at time, This indicates that the target model generates lexical units. The joint probability distribution, This indicates that the draft model generates lexical units. The joint probability distribution.

[0066] This pair of sub-blocks of length L The probability of acceptance is The specific calculation formula is as follows:

[0067]

[0068] in, This indicates that the target model generates lexical units. The joint probability distribution, This indicates that the draft model generates lexical units. The joint probability distribution.

[0069] The current word block is determined to be accepted or rejected based on the calculated acceptance probability. If rejected, the word block is recorded in the set. If accepted, update the longest accepted block. After completing After verifying all sub-blocks, if the entire word block... If all sequences are accepted, the iteration over the different sequences terminates early; otherwise, the next draft sequence will be verified. The final accepted block is the longest accepted block obtained through this process. .

[0070] Specifically, the standard rejection sampling method is used to determine whether the current word segment is accepted or rejected. A random number is sampled from a uniform distribution in the interval [0,1]. If the sampled random number is less than the calculated acceptance probability and the segment is not in the rejection set H, then the word segment is accepted; otherwise, the word segment is rejected.

[0071] S3. Residual Sampling and Distribution Correction:

[0072] If the block is ultimately accepted The length is equal to the length of the candidate draft sequence. Then an additional lexical unit is sampled from the target model. : ,in, This indicates that the target model is in a given word. When, the conditional probability distribution for generating the next word element, Indicates a given prefix, Indicates length is The accepting word block If the block is eventually accepted The length is less than the length of the candidate draft sequence Then from the predefined residual distribution Sample an additional lexical ,in, The sequence of the final accepted blocks.

[0073] To maintain consistency between the generated sequence distribution and the target model, the target model distribution is adjusted before the next round of inference decoding iteration. After making corrections, a new distribution is obtained. This is used for the next round of decoding. The following... Each lexical unit needs to be sampled from a modified distribution, denoted as . This distribution is different from the original target distribution. different.

[0074] like :

[0075]

[0076] in, Indicates the first in the next round of speculation decoding Each word element, This indicates that the target model generates lexical units. Joint probability distribution at time; This indicates that the draft model generates lexical units. Joint probability distribution at time; Indicates length The accepting word element sub-block, This represents the additional tokens obtained through residual sampling. Indicates the previous round of speculation decoding Each word element; This indicates that the target model generates lexical units. Joint probability distribution at time, This indicates that the draft model generates lexical units. Joint probability distribution at time, It represents any candidate word in the vocabulary.

[0077] like , ;

[0078] in, Indicates that the target model generates the first each word element The probability distribution.

[0079] The word position i is decomposed using L-τ-1 as the decomposition point. This decomposition point is determined by the initially set draft length L and the length already accepted in this round. The difference between them determines the result. In the mathematical proof of speculative decoding, to ensure that the final sequence distribution generated in multiple rounds is strictly equivalent to the target model distribution, each round of verification is theoretically considered as processing a complete block of length L. (Accepted) 1 lexical unit was sampled, and an additional lexical unit y was sampled (whether from the residual or the target model). Therefore, this round actually determined 1 lexical unit y. +1 word unit. The original plan length is L. Because in The segment (with the additional lexical y) is truncated early and moves to the next round, so the remaining part of the original length L... The conditional probability distribution at each position is affected by the truncation operation in this round (i.e., posterior bias from rejecting sampling). Because the "greedy truncation" in this round introduces bias, if the original target model q is used directly for verification in the next round, the final generated long sequence distribution will deviate from the true distribution. This will be addressed in the following... Use the modified distribution in the step Using this as a benchmark to calculate the acceptance probability can mathematically offset the bias caused by the previous round of truncation, thus ensuring that the entire algorithm is theoretically lossless and that the generated results are completely consistent with simply running a large model.

[0080] The expression for the predefined residual distribution is:

[0081]

[0082] in, This indicates that the target model generates lexical units. Joint probability distribution at time, This indicates that the draft model generates lexical units. Joint probability distribution at time, Indicates the preceding Each word element, Indicates the next word to be sampled.

[0083] S4, Prefix Update and Loop Execution: Append the received word block to the prefix sequence; if the termination flag is not reached, return to S1 and repeat the draft generation and verification process. When the end-of-sequence symbol is decoded, output the complete decoded sequence.

[0084] The generation of terminators follows the same validation rules as ordinary lexical units. When generating candidate sequences, the draft model may generate terminators in the same way as other lexical units. When the draft model deems the sequence should end, it outputs a terminator as a candidate lexical unit at a certain position. If the terminator is accepted by the target model (i.e., the acceptance probability condition is met), the sequence officially ends. If the terminator is rejected, subsequent lexical units are generated.

[0085] The goal of the entire verification process is to maximize the final accepted block length or minimize the cost function, which is expressed as:

[0086]

[0087] in, Describes all that make the marginal distributions respectively and The effective coupling set, express The joint distribution of 1 independent and identically distributed candidate draft sequences. This indicates a concatenation or direct sum of independent and identically distributed components. This represents the block distribution of the target model. Indicates in coupling distribution Lower Acceptable Length Expected value This represents the length of the continuous word block accepted by the target model in a single validation. Indicates from coupling distribution Target sequence block obtained by mid-sampling .

[0088] The present invention, through the above steps, enables the desired receiving length. Strictly adhere to the theoretical value of the maximum expected accepted length And solve OT problems efficiently with provable guarantees.

[0089] The various embodiments in this specification are described in a progressive manner, with each embodiment focusing on its differences from other embodiments. Similar or identical parts between embodiments can be referred to interchangeably. For the apparatus disclosed in the embodiments, since they correspond to the methods disclosed in the embodiments, the description is relatively simple; relevant parts can be referred to the method section.

[0090] The above description of the disclosed embodiments enables those skilled in the art to make or use the invention. Various modifications to these embodiments will be readily apparent to those skilled in the art, and the general principles defined herein may be implemented in other embodiments without departing from the spirit or scope of the invention. Therefore, the invention is not to be limited to the embodiments shown herein, but is to be accorded the widest scope consistent with the principles and novel features disclosed herein.

Claims

1. A speculative decoding acceleration method based on multiple draft blocks and greedy verification, characterized in that, Includes the following steps: S1, given a prefix sequence Under the condition of independent and identically distributed generation through the draft model The length of the strip is Candidate draft sequence This forms a candidate set; the current longest accept block length is set. and initialize an empty rejection set. , used to record rejected lexical sub-blocks; S2, Traversal A sequence of candidate drafts, for each candidate draft sequence, starting from position Begin validating each word sub-block in the candidate draft sequence, determining whether the word sub-block to be validated exists in the rejection set. If so, then skip the verification of that word sub-block; Otherwise, calculate the acceptance probability of the sub-block and reject or accept the sub-block based on the acceptance probability. If the sub-block is accepted, update the longest accepted block length. If the word sub-block is rejected, it is added to the rejection set. The longest accepted block obtained during the verification process will be used as the final accepted block. ; S3, based on the final received block The length of the term is determined by sampling an additional term from the target model or a predefined residual distribution. In each round of inference decoding, the probability distribution of the target model is corrected. like : in, Indicates the first in the next round of speculation decoding Each word element, This indicates that the target model generates lexical units. Joint probability distribution at time; This indicates that the draft model generates lexical units. Joint probability distribution at time; Indicates length The accepting word element sub-block, This represents the additional tokens obtained through residual sampling. Indicates the previous round of speculation decoding Each word element; This indicates that the target model generates lexical units. Joint probability distribution at time, This indicates that the draft model generates lexical units. Joint probability distribution at time, Represents any candidate word in the vocabulary; like , ; in, Indicates that the target model generates the first each word element The probability distribution; S4, the block that will be accepted in the end and additional tokens of sampling Append to prefix sequence In the middle, return to S1, and continue decoding until the end-of-sequence symbol is reached, then output the complete decoded sequence.

2. The method as described in claim 1, characterized in that, In S2, after verifying a certain candidate draft sequence... Subsequently, if all lexical sub-blocks in the candidate draft sequence are accepted, the iterative verification process for subsequent sequences is terminated early; otherwise, the verification of the next candidate draft sequence continues. .

3. The method as described in claim 1, characterized in that, In S2, the length is Word element blocks The probability of acceptance is The specific calculation formula is as follows: in, This indicates the next word sub-block to be sampled. Indicates the current word element sub-block, This indicates that the target model generates word sub-blocks. Joint probability distribution at time, This indicates that the draft model generates lexical units. Joint probability distribution at time, This indicates that the target model generates lexical units. The joint probability distribution, This indicates that the draft model generates lexical units. The joint probability distribution.

4. The method as described in claim 1, characterized in that, In S2, a sub-block of length L The probability of acceptance is The specific calculation formula is as follows: in, This indicates that the target model generates lexical units. The joint probability distribution, This indicates that the draft model generates lexical units. The joint probability distribution.

5. The method as described in claim 1, characterized in that, In S3, if the block is eventually accepted The length is equal to the length of the candidate draft sequence. Then an additional lexical unit is sampled from the target model. : ,in, This indicates that the target model is in a given word. When, the conditional probability distribution for generating the next word element, Indicates a given prefix, Indicates length is The accepting word block If the block is eventually accepted The length is less than the length of the candidate draft sequence Then from the predefined residual distribution Sample an additional lexical ,in, The sequence of the final accepted blocks.

6. The method as described in claim 1, characterized in that, The expression for the predefined residual distribution is: in, This indicates that the target model generates lexical units. Joint probability distribution at time, This indicates that the draft model generates lexical units. Joint probability distribution at time, Indicates the preceding Each word element, Indicates the next word to be sampled.

7. The method as described in claim 1, characterized in that, The goal of the verification process is to maximize the final accepted block length or minimize the cost function, which is expressed as: in, Describes all that make the marginal distributions respectively and The effective coupling set, express The joint distribution of 1 independent and identically distributed candidate draft sequences. This indicates a concatenation or direct sum of independent and identically distributed components. This represents the block distribution of the target model. Indicates in coupling distribution Lower Acceptable Length Expected value This represents the length of the continuous word block accepted by the target model in a single validation. Indicates from coupling distribution Target sequence block obtained by mid-sampling .

Images