Speculative decoding optimization method, electronic device and storage medium

Abstract

The embodiments of the present application disclose a speculative decoding optimization method, an electronic device and a storage medium, wherein a speculative decoding optimization method is used for a draft model and a target verification model, the method comprising: inputting the draft content generated by the draft model in the previous round into the target verification model; performing parallel verification on the tokens in the draft content in the target verification model; using a lightweight draft length predictor to analyze the latent vector of the last received token position and the last token generated, and predicting the optimal draft length for the next round; and autoregressively generating a corresponding number of drafts based on the optimal draft length predicted by the lightweight draft length predictor through the draft model.

Description

Technical Field

[0001] The present invention belongs to the technical field of speculative decoding optimization, and in particular relates to a speculative decoding optimization method, an electronic device and a storage medium. Background Art

[0002] Among related technologies, SD (Speculative Decoding) is a technology used to accelerate the decoding process in natural language processing and generation tasks. A technology used in natural language generation that aims to speed up generation and improve efficiency by predicting multiple outputs and computing them in parallel. The goal of speculative decoding is to increase the decoding speed of the model without significantly reducing the generation quality.

[0003] AR (Auto-regressive) is a method of generating models in which each generation step depends on what was generated previously. That is, when generating the next word or symbol, the model will gradually predict the next output based on the sequence that has been generated. This method is sequential and is often used in natural language generation tasks. For example, the GPT model generates a new word at each step until the end of the sentence.

[0004] The speculative sampling technique uses two independent models in series: a small model generates drafts in an autoregressive manner, and a large model verifies the drafts in parallel to speed up decoding.

[0005] EAGLE performs autoregression at the latent vector level through the hidden state of the last layer of the target model, avoiding the use of an additional draft model. DDD is based on EAGLE's speculative decoding architecture, uses the model output probability as the draft confidence, and uses a manually set threshold to control the draft generation length.

[0006] The inventors found that the speculative sampling technology has the problems of low accuracy in generating drafts and high response delay of the draft model. Among them, the EAGLE technology relies on a fixed draft structure. However, in actual applications, the length of the draft received varies significantly in different iteration steps, resulting in reduced draft calculation efficiency. The shortcoming of the DDD technology is that it cannot be further optimized for specific adaptive draft generation goals. Fixed thresholds have limitations when dealing with the starting position of drafts for different data sets, and it is difficult to ensure generalization. Both of the above technologies require manual selection of a fixed draft length to improve decoding efficiency, and it is impossible to dynamically adjust the draft structure during the decoding process to achieve adaptive optimization. Summary of the invention

[0007] The embodiments of the present invention provide a speculative decoding optimization method, an electronic device and a storage medium, which are used to solve at least one of the above-mentioned technical problems.

[0008] In a first aspect, an embodiment of the present invention provides a speculative decoding optimization method for a draft model and a target verification model, the method comprising: inputting the draft content generated by the draft model in the previous round into the target verification model; performing parallel verification on the tokens in the draft content in the target verification model; using a lightweight draft length predictor to analyze the latent vector of the last received token position and the last token generated, and predicting the optimal draft length for the next round; and autoregressively generating a corresponding number of drafts based on the optimal draft length predicted by the lightweight draft length predictor through the draft model.

[0009] In a second aspect, an embodiment of the present invention further provides a computer program product, wherein the computer program product comprises a computer program stored on a non-volatile computer-readable storage medium, wherein the computer program comprises program instructions, and when the program instructions are executed by a computer, the computer executes the steps of the speculative decoding optimization method of any embodiment of the present invention.

[0010] In a third aspect, an embodiment of the present invention further provides an electronic device comprising: at least one processor, and a memory communicatively connected to the at least one processor, wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor so that the at least one processor can perform the steps of the method described in the first aspect.

[0011] In a fourth aspect, an embodiment of the present invention further provides a storage medium having a computer program stored thereon, wherein the computer program, when executed by a processor, implements the steps of the method described in the first aspect.

[0012] In the method of the embodiment of the present application, our product models the draft structure through an explicit method, and innovatively uses a lightweight prediction module to predict the optimal draft generation length of the current position before generating the draft. This method avoids relying on the probability of the model output as the confidence level, and there is no need to manually set the threshold. We also implemented and verified this method based on the current optimal speculative decoding model EAGLE. By dynamically predicting the draft length, our system significantly improves the decoding efficiency without affecting the quality of model generation, and effectively solves various defects in traditional methods. BRIEF DESCRIPTION OF THE DRAWINGS

[0013] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the accompanying drawings required for use in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other accompanying drawings can be obtained based on these accompanying drawings without paying any creative work.

[0014] Figure 1A flowchart of a speculative decoding optimization method provided by one embodiment of the present invention; Figure 2 Motivation for the related art provided for an embodiment of the present invention: Finding the best draft stopping point to reduce the target forward pass and minimize computational waste; Figure 3 A comparison of the speculative decoding method with adaptive draft length provided by an embodiment of the present invention in terms of basic framework, draft boundary modeling and artificial threshold requirements; Figure 4 The maximum draft length accepted by EAGLE in different draft iterations provided by an embodiment of the present invention; Figure 5 Performance comparison of the EAGLE model with different fixed draft lengths (ranging from 2 to 6) and the adaptive length draft model provided by an embodiment of the present invention; Figure 6 AdaEAGLE framework schematic diagram provided for one embodiment of the present invention; Figure 7 Algorithm 1 provided for an embodiment of the present invention: Adaptive length draft decoding using LDLP; Figure 8 The average number of decoded blocks τ and throughput of different methods provided in one embodiment of the present invention in six benchmark tests; Fig. 9 Ablation experiment results on MT-Bench and Alpaca for training loss design provided in one embodiment of the present invention; Fig.10 Comparison of draft length distribution with lossy penalty and without lossy penalty provided by an embodiment of the present invention; Fig.11 Throughput results of different thresholds provided by an embodiment of the present invention; Fig.12 A schematic diagram of the structure of an electronic device provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0015] In order to make the purpose, technical solution and advantages of the embodiments of the present invention clearer, the technical solution in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0016] The inventors found that the defect of speculative sampling technology stems from its reliance on two independent large language models, usually a combination of large and small models in the same series. Although this configuration can improve the flexibility of the system, it also brings high optimization costs and low decoding efficiency. The shortcomings of EAGLE technology come from its use of a fixed draft structure and its failure to dynamically adjust the draft length of each step. This leads to a serious waste of draft computing resources when receiving a shorter length; and when the optimal receiving length exceeds the fixed draft length, the target model needs to increase the number of verifications, further affecting the overall efficiency. The defect of DDD technology is that it relies on the output probability of the model to directly serve as the confidence of draft reception. This implicit modeling method is difficult to effectively grasp the optimal draft length, which limits the efficiency improvement and fails to give full play to the potential of adaptation.

[0017] When addressing these defects, practitioners in this industry usually adopt the following methods: (1) selecting and optimizing different fixed draft structures according to different usage scenarios and reasoning tasks; (2) designing new speculative decoding models to improve the accuracy of draft generation; (3) studying hardware optimization solutions to improve the efficiency of speculative decoding models in utilizing hardware resources.

[0018] The reason why our solution is not easy to come up with is that the technical difficulty of explicitly modeling the draft structure is high. Explicit modeling usually brings additional computational costs, so the industry tends to use existing speculative decoding models to achieve efficiency improvements in specific domain tasks rather than developing new modeling strategies. This inertial thinking makes explicit modeling solutions less valued.

[0019] The embodiment of the present application models the draft structure through an explicit method, and innovatively uses a lightweight prediction module to predict the optimal draft generation length of the current position before generating the draft. This method avoids relying on the probability of the model output as the confidence level, and does not require manual setting of the threshold. We also implemented and verified this method based on the current optimal speculative decoding model EAGLE. By dynamically predicting the draft length, our system significantly improves the decoding efficiency without affecting the quality of model generation, and effectively solves various defects in traditional methods.

[0020] The product architecture of the embodiment of the present application consists of the following four steps: 1. Input the draft content generated in the previous round into the target verification model.

[0021] 2. The target model performs parallel validation on the tokens in the draft.

[0022] 3. The lightweight draft length predictor (LDLP) predicts the optimal draft length for the next round by analyzing the latent vector of the last received token position and the last token generated.

[0023] 4. The draft model autoregressively generates the corresponding number of drafts based on the draft length predicted by LDLP.

[0024] The training data construction process is divided into the following three steps: 1. Use the target model to generate standard output based on the input prompt.

[0025] 2. Perform processing at any position in the generated output and use the draft model to generate a certain upper limit of draft tokens.

[0026] 3. Compare the output of the draft model with the standard output of the target model and calculate the length of their longest common prefix, which is used as the optimal draft length target for the corresponding position.

[0027] This architecture makes the prediction and generation process of draft length more efficient and accurate, avoids manual threshold setting and inefficient verification process, and effectively improves decoding performance and generation quality.

[0028] In the process of implementing this application, the applicant has also adopted some alternative solutions and beta versions. Among them, the alternative solution is to directly collect training data and train a model to predict whether the draft token generated at each position will be received. In the inference stage of this model, if the predictor determines that the currently generated draft will not be received, the draft generation will be stopped. The advantage of this method is that it is relatively simple to implement and can directly judge the draft stop position by the quality of the draft generation. The disadvantage is that each draft generation step needs to be predicted, which generates a large additional computing overhead and reduces the acceleration effect. In the Beta version, compared with the final version, we combined the latent vector of the context to jointly predict the optimal draft length, and did not use a length-based loss penalty during the training process. Since the introduction of context brings a lot of computing overhead, the acceleration effect is insufficient. At the same time, the lack of loss penalty causes the model to predict a short length, making it difficult to achieve a good acceleration effect.

[0029] The direct effect of the solution of the embodiment of the present application is to significantly reduce the computational overhead of model training and reasoning. By using a lightweight prediction module to predict the optimal draft length before generating a draft, this solution reduces unnecessary calculation steps, improves decoding efficiency, and reduces resource overhead. It opens the door to modeling adaptive draft structures in an "end-to-end" manner, and can be integrated into a wide range of speculative decoding frameworks and sampling strategies. This can accelerate the application of AI technology in more industries. This will promote the improvement of social productivity and the popularization of technology, promote the development of broader and more sustainable intelligent systems, and ultimately change people's work and lifestyles.

[0030] Please refer to Figure 1, which shows a flow chart of an embodiment of a speculative decoding optimization method of the present application. The method is used for a draft model and a target verification model in speculative decoding.

[0031] like Figure 1 As shown, in step 101, the draft content generated by the draft model in the previous round is input into the target verification model; In step 102, tokens in the draft content are verified in parallel in the target verification model; In step 103, a lightweight draft length predictor is used to analyze the latent vector of the last received token position and the last token generated to predict the optimal draft length for the next round; In step 104, a corresponding number of drafts are autoregressively generated by the draft model based on the optimal draft length predicted by the lightweight draft length predictor.

[0032] The method of the embodiment of the present application models the draft structure through an explicit method, and innovatively uses a lightweight prediction module to predict the optimal draft generation length of the current position before generating the draft. This method avoids relying on the probability of the model output as the confidence level, and there is no need to manually set the threshold. Furthermore, we have also implemented and verified this method based on the current optimal speculative decoding model EAGLE. By dynamically predicting the draft length, our system significantly improves the decoding efficiency without affecting the quality of model generation, and effectively solves various defects in traditional methods.

[0033] In some optional embodiments, the process of constructing the training data of the lightweight draft length predictor includes: Using the target verification model to generate standard outputs based on input prompts; Processing at any position of the generated output, using the draft model to generate a certain upper limit of draft tokens; The output of the draft model is compared with the standard output of the target verification model, and the length of the longest common prefix is ​​calculated as the optimal draft length target for the corresponding position.

[0034] In some optional embodiments, the true values ​​and predicted values ​​in the training data are trained using L1 loss.

[0035] In some optional embodiments, when training by L1 loss, the method further includes: A penalty coefficient is introduced to amplify the L1 loss when the predicted value is smaller than the true value.

[0036] In some optional embodiments, a three-layer multilayer perceptron with residual connections is used as the model structure of the lightweight draft length predictor, and the output hidden state is mapped to a length scalar through a linear transformation.

[0037] In some optional embodiments, the method is based on sequential draft generation and greedy decoding, in which the optimal draft length is defined as the length that minimizes the number of forward passes required to verify the target model.

[0038] It should be noted that the above method steps are not used to limit the execution order of each step. In fact, some steps may be executed simultaneously or in the opposite order to the specified step order, and this application has no limitation on this.

[0039] The solution of the present application is described below through a specific embodiment so that those skilled in the art can better understand the solution of the present application.

[0040] Speculative Decoding (SD) is a popular technique for losslessly accelerating large language models (LLMs). Our research shows that the decoding speed of the static draft structure SD framework can be significantly improved by adopting a context-aware adaptive draft structure. However, current research on adaptive draft structures has limitations in performance, modeling methods, and applicability. In this application embodiment, we will propose AdaEAGLE, the first speculative decoding (SD) framework that explicitly models an adaptive draft structure. AdaEAGLE uses a lightweight draft length prediction (LDLP) module to explicitly predict the optimal number of draft tokens during inference to guide the draft model. It achieves comparable acceleration to the current state-of-the-art (SotA) without manual thresholds and allows for deeper and more specialized optimizations. In addition, combined with a threshold-based strategy, AdaEAGLE is 1.62 times faster than the original autoregressive (AR) AR decoding while maintaining output quality, and surpasses the fixed-length SotA baseline.

[0041] In the related art, the auto-regressive (AR) model is very effective in language modeling, but faces latency issues due to its sequential generation of tokens, especially in large language models (LLMs). Speculative decoding (SD) alleviates this problem by dividing decoding into two stages: a draft stage and a verification stage. A smaller draft model predicts multiple tokens, which are verified by the larger model in a forward pass. By iterating these two stages, speculative decoding (SD) speeds up the LLM while maintaining the output distribution of the original model. Most methods use static structures to generate draft tokens, such as fixed-length sequences or static trees.

[0042] Figure 2 : Motivation: finding the optimal draft stopping point to reduce targetforward passes and minimize wasted computation.

[0043] However, studies have shown that the accepted length of a draft token depends heavily on the generated context. This means that a draft model with a static draft structure cannot always accurately generate the optimal number of tokens, resulting in more target model forward pass costs and draft computation waste, such as Figure 2 shown. Figure 4 It is further shown that during sample inference, the accepted lengths of different draft iterations vary greatly. Therefore, to optimize decoding efficiency, an ideal draft model should use a context-aware adaptive draft structure instead of a static structure. In fact, in our preliminary experiments, we further demonstrate that a draft length oracle (telling the draft model the optimal draft length before generation) can improve the decoding throughput of fixed-length EAGLE by 29%. Therefore, it is meaningful and valuable to implement an adaptive draft structure.

[0044] Figure 3 : Comparison of speculative decoding methods with adaptive draft length in terms of base framework, draft boundary modeling, and manual threshold requirements.

[0045] Despite Figure 3As listed, multiple efforts have explored adaptive draft structures in different ways, but they lack comprehensive analysis and systematic methods. Their limitations can be summarized as follows: (1) Unsatisfied performance. Most of the current methods are based on intrinsic output and threshold control, which cannot effectively simulate the optimal draft length. Experiments also show that there is still a significant optimization gap compared with the top-line performance of EAGLE Oracle. (2) Hard to optimize. The implicit method of modeling adaptive draft structure relies on the LM head output of the draft model, which is naturally trained based on the original token sequence. This output cannot be further optimized for the specific goal of adaptive draft generation. In addition, the predefined fixed threshold cannot be applied to all datasets and draft starting points, which poses a challenge to generalization ability. (3) Obsolete applicability. Early methods are not suitable for the current SotA framework EAGLE (Li et al., 2024b; Zhou et al., 2024). For example, SpecDec++ assumes that draft generation follows a Markov decision process (MDP), while EAGLE does not meet this assumption. PEARL is only applicable when the cost of each inference step of the draft model is of the same order as the cost of each inference step of the target model. (4) Complex. Methods that require training involve complex data construction and complex training processes. In the present application embodiment, we propose a new method AdaEAGLE, which is the first speculative decoding (SD) framework that explicitly models adaptive draft structures through a lightweight draft length prophet module (LDLP). The advantages of AdaEAGLE include (1) Explicit modeling of adaptive draft lengths. We use LDLP to predict the optimal draft length before generating a draft with simple input. Thanks to this, LDLP does not rely on the output logarithm and does not require manual threshold setting. (2) Built on SotA SD framework EAGLE (EAGLE). We use hidden states and output tokens to predict the draft length. (3) Simple and easy. The simplified data structure and highly parallelized training make it easy to integrate into the existing speculative decoding (SD) framework.Our results show that AdaEAGLE achieves a speedup of 1.61 over vanilla AR decoding; a 2% improvement over fixed-length EAGLE decoding, and beats implicit modeling approaches (throughput 65.20 vs 64.43).

[0046] In summary, the main contributions of our work are (1) a systematic analysis of the huge potential gains of implementing adaptive draft length. (2) a simple but effective SotA-based approach, AdaEAGLE, is proposed, demonstrating that explicit modeling of draft length in adaptive draft generation control is feasible and efficient. (3) Explicit modeling facilitates deeper and more sophisticated optimizations and paves the way for improving adaptive draft structures in more complex situations such as tree-like decoding and non-greedy decoding.

[0047] 2 Analysis of Adaptive Draft Length In this section, we present a systematic analysis of the huge potential gains of integrating adaptive draft length control into speculative decoding (SD). We first propose Speculative Decoding (SD) and its SotA variant EAGLE. Then, we analyze the optimal draft length under the EAGLE framework. Finally, based on pilot experiments, we discuss the significant advantages of adaptive draft.

[0048] 2.1 Speculative Decoding Let T 1:j represents the sequence of tokens t1,t2,...,t generated so far by the AR model (usually LLM) j , p(t j ) is the target model assigned to token t j Speculative decoding introduces a smaller draft model that can guess candidate tokens i∈[1,k] Sequences and their probabilities The target model then computes the actual probabilities of these tokens in parallel , and decides in turn whether to accept these tokens. Each draft token The acceptance is determined by the following factors: If the token Rejected, all subsequent tokens will be cancelled and removed from distribution Resample a new token in . This process allows speculative decoding (SD) to verify multiple tokens at a time, greatly reducing the number of consecutive steps required for AR decoding. This approach ensures that the output distribution of speculative decoding (SD) is consistent with the vanilla AR decoding of the target LLM.

[0049] 2.2 EAGLE To improve draft quality, EAGLE (Extrapolation Algorithm for Greater LanguageModel Efficiency) no longer trains a separate small model, but reuses the input embedding and LM head of the target model, adding a trainable draft head in between. The first version of EAGLE used a static tree for a forward validation, while the second version dynamically selected candidates based on confidence values. As mentioned earlier, DDD (Brown et al., 2024) improved this approach, but still relied on confidence-based acceptance and manually set thresholds.

[0050] 2.3 EAGLE's Draft Length Oracle For ease of discussion, unless otherwise stated, the embodiments of the present application are based on sequential draft generation and greedy decoding (sample temperature is set to 0). We first illustrate the new EAGLE-Oracle model with ideal adaptability through the draft length Oracle. The EAGLE-Oracle model uses the same architecture and parameters as the original EAGLE model, but during reasoning, the draft length Oracle module can tell the optimal draft length. In greedy decoding, the optimal draft length is defined as the length that minimizes the number of forward passes required for the target model. In the case of greedy decoding and EAGLE, the optimal draft length is exactly the acceptance length (additional assumptions are required, see the proof in Appendix A). By telling the draft model the optimal length before generation, EAGLE-Oracle ensures that the draft model outputs the required number of tokens, thereby minimizing waste and maximizing efficiency.

[0051] 2.4 Benefit of Adaptive Draft Lengths By using the scratch length Oracle, EAGLE-Oracle achieves the theoretically optimal speedup under given model parameters. We evaluate the speedup effects of fixed-length EAGLE and EAGLE-Oracle. The experimental setup follows the examples below and the results are shown in Figure 5 The first five rows represent EAGLE with fixed-length draft token sequences of 2 to 6, respectively, while the sixth row corresponds to EAGLE-Oracle. For EAGLE-Oracle, N waste= 0, because it generates the optimal number of tokens at each draft start position. As a result, the average acceptance length and throughput of EAGLEOracle are much higher than those of the EAGLE model with fixed-length drafts (42% improvement), indicating that there is great potential value in pursuing the perfect adaptive draft length.

[0052] Figure 6 : Illustration of AdaEAGLE framework. (a) The architecture of AdaEAGLE. The draft length are predicted by an LDLP module. (b) A closer look of LDLP. (c) An example of collecting paired training data of LDLP.

[0053] 3 AdaEAGLE In this section, we will first present the overall architecture of AdaEAGLE, the first speculative decoding (SD) framework that explicitly models the adaptive draft structure. To this end, we carefully design a lightweight draft length predictor (LDLP), which functions like a prophet and can directly estimate the optimal draft length in advance during inference. More details of LDLP will be presented in subsequent examples. Finally, we will demonstrate how to train LDLP.

[0054] 3.1 Speculative Decoding with Adaptive Draft Length (SD) The target model is denoted as M T , the draft model is denoted as M D The draft model is responsible for automatically generating drafts, while the target model verifies the correctness of all tokens in the draft in one forward pass. The verified prefixes in the generated draft are fed back to the draft model to continue decoding until the entire decoding process is completed.

[0055] In form, for M T and M D In the interaction of the rth iteration, let T r = Indicates M T The input sequence is is the formal sequence generated until the last iteration, is the draft of the current iteration, k ∈ N+ is the draft length. The target model M T Verify the draft by the following formula:

[0056] in, is the number of correct tokens in the draft, Reward tokens for additional predictions during the validation process ( Figure 6 -(a)-①②). Then, M D The draft for the (r + 1)th iteration is generated as follows:

[0057] Among them, + represents the connection operation ( Figure 6 -(a)-④).

[0058] As shown in Formula 2, in standard speculative decoding, M D The length of the generated draft is a pre-fixed hyperparameter k. According to the above analysis, in M T and M D During each iteration of interaction, if k is equal to the number of correct tokens in the draft If the gap between is too large, it will reduce the decoding speed of the model. Inspired by this, we designed a lightweight draft length predictor (LDLP), which functions like a prophet and can estimate the optimal draft length ( Figure 6 -(a)-③), that is, predict k r ∈N + The value of is close to the optimal draft length Mathematically, k in formulas 1 and 2 is replaced by k r We will Figure 7 Detailed description in. Figure 7 Algorithm 1 for adaptive length draft decoding using LDLP is shown. The Chinese and English correspondences are as follows: Algorithm - algorithm, Require - input, Ensure - output, while - loop, do - execute, end - end.

[0059] 3.2 Lightweight Draft Length Predictor When designing LDLP, we followed the "KISS" (Keep It Simple and Stupid) principle. Our goal is to ensure that its output is very close to the correct number of tokens in the draft while keeping it as lightweight as possible. Heavy LDLP increases the computational overhead and data transfer between the CPU and GPU, which offsets the expected acceleration effect of speculative decoding. Based on this design principle, we adopted a three-layer multi-layer perceptron (MLP) as the model structure of LDLP. Previous research on LLM future token prediction pointed out that the hidden state of the last generated token contains a "plan" to generate subsequent future tokens. Inspired by this, in the rth iteration, LDLP only needs to replace the last token in the draft with the last token in the draft. Embedding in a validated prefix as input and put it in M T The last hidden state after the last layer of normalization As input (see Figure 6 -(b)). At the same time, LDLP outputs a scalar in R (the set of real numbers). Mathematically:

[0060] Among them, k max is a very loose hyperparameter that controls the upper limit of possible draft length.

[0061] 3.3 LDLP Training The training of LDLP involves two main issues: training data collection and optimization criterion design.

[0062] Training data collection As mentioned before, the training data of LDLP consists of the following data pairs , where e j and f j They represent the embedding and last hidden state of the jth token in the target model after normalization in the last layer. j represents the optimal draft length starting from the jth token. For a given input prompt T with m tokens in , we collect such token pairs through the following three steps: 1) Target model M TAccording to T in Generate output sequence T out (n tokens) ( Figure 6 -(c)-①).

[0063] 2) Based on T in +T :j out Draft model M D Continue drafting max Token (denoted as ) (For any case 1≤j≤n, see Figure 6 -(c)-②).

[0064] 3) Comparison and T out , calculate the length k of the longest common prefix between them j ( Figure 6 -(c)-③).

[0065] Optimization criteria design let and Represent the ground truth and prediction results of LDLP respectively. On the one hand, we hope and as close as possible and use L1 loss to achieve this goal. On the other hand, we observed that one forward pass of the target model takes about 20 times as long as the draft model. Therefore, we hope As far as possible not less than To this end, we introduce a penalty coefficient λ>1, when The L1 loss is expanded when . Formally, the criterion is defined as:

[0066] 3.4 Combining AdaEAGLE with other adaptive drafting techniques Our explicit draft length modeling via LDLP can be easily combined with other approaches that leverage intrinsic model outputs (e.g., log,entropy) for draft boundary modeling. By evaluating a condition at each step, the draft model terminates generation based on the predicted length and an intrinsic threshold criterion. In our experiments, we discuss a combination of AdaEAGLE and DDD (Brown et al., 2024), called AdaEAGLEDDD.

[0067] 4 Experiments 4.1 Experimental Setup Model We experiment with Vicuna7B-v1.3 (Chiang et al., 2023) and paired EAGLE (Li et al., 2024b) draft models. Our LDLP model is a three-layer MLP with residual connections, and the output hidden state is mapped to a length scalar via a linear transformation.

[0068] Training For a fair comparison, we train the length prediction model using 68,000 dialogue samples from the ShareGPT dataset. We first ask Vicuna to automatically regress responses and provide prompts for each dialogue sample. We then follow the data collection process detailed in the previous example. To train LDLP, we freeze the base target model and the corresponding Eagle Draft model. We set the starting learning rate to 5e-5 and use a cosine learning rate scheduler. LDLP is trained with a batch size of 128 on 8 A800 80GB GPUs, completing 5 epochs in 30 minutes, making AdaEAGLE very efficient to train.

[0069] Benchmarks Following the evaluation setting of EAGLE, we selected six representative text generation benchmarks in different domains to evaluate our method. These benchmarks include the multi-turn dialogue task MT-Bench (Zheng et al., 2024), the instruction following task Alpaca (Taori et al., 2024), the code generation task HumanEval (Chen et al., 2021), the mathematical reasoning task GSM8K (Cobbe et al., 2021), the summarization task CNN / Daily Mail (Nallapati et al., 2016), and the question answering task Natural Questions (Kwiatkowski et al., 2019). All experiments are conducted on a single A800-80G GPU with a batch size of 1, using greedy decoding (temperature set to 0.0) and FP32 precision. Metrics Since we did not modify the target model or the draft model, the generation quality is consistent with EAGLE and previous speculative decoding methods that achieve lossless acceleration. Therefore, we do not report the quality metrics of the benchmarks. Instead, we use the following two metrics to evaluate the speedup performance: (1) average acceptance length (τ), defined as the number of tokens accepted per decoding step; and (2) throughput (Tok / s), defined as the number of tokens generated per second.

[0070] 4.2 Main Results Baseline We choose the original EAGLE method with fixed draft length as the baseline. At each decoding step, EAGLE automatically generates a fixed number of draft tokens (ranging from a larger fixed draft length such as 5, which will reduce the target forward times). AdaEAGLE with adaptive draft length shows its adaptability in different tasks.

[0071] Figure 4 : Maximum draft lengths accepted by EAGLE in different draft iterations.

[0072] Figure 5 : Performance comparison of the EAGLE model with different fixed draft lengths (ranging from 2 to 6) and the adaptive length draft model. τ and Tok / s are the acceptance length and sample average throughput, respectively. T total represents the overall reasoning time, and T draft and T target Represents the time used for the draft model and the target model, respectively. N draft and N target Indicates the number of tokens generated during the decoding process, N waste Indicates the number of draft tokens rejected by the target model during validation.

[0073] Figure 5 A breakdown of the MT-Bench speedup factors is provided, showing the inference time taken by the draft model and the target model for each draft setting, as well as the number of tokens generated by each model. For EAGLE with a fixed draft length, the number of tokens generated by the draft model (N) increases as the draft length increases. draft ) increased from 47525 to 80112, thereby reducing the number of forward passes of the target model (N target ). However, if Figure 4 As shown in the figure, due to the drastic fluctuation of the maximum draft length at different positions, this also leads to the number of wasted tokens (N waste ) surge. This trend is directly reflected in the time spent on each component: as N draft The increase of T draft will also increase, and T target Then N target This poses a challenging trade-off problem. Our adaptive draft length prediction method achieves a better balance. Specifically, compared with EAGLE with a fixed draft length of 4, our method reduces the wasted draft tokens and target model forward passes, thereby reducing T draft and T target, ultimately increasing the overall throughput to 66.35 tok / s. We also tested a recently proposed acceptance rate modeling method DDD, which is based on the token probabilities generated by the model. Figure 6 As shown in Figure 2, the average speedup performance of AdaEAGLE is comparable to that of DDD, with throughputs of 64.44 tok / s and 64.43 tok / s, respectively. DDD has a higher average number of accepted tokens, which may be because our model is relatively conservative in predicting the draft length, while DDD is more aggressive, generating more draft tokens and more discarded tokens.

[0074] Our AdaEAGLE-DDD approach achieves additional gains in decoding throughput, reaching an average of 65.20 tokens per second. AdaEAGLE-DDD also achieves the highest throughput on four of the six benchmarks. Notably, AdaEAGLE-DDD has the best speedup performance on the code generation task HumanEval. This is because we retain the longer of the two lengths, thereby reducing the probability of prematurely exiting the draft and increasing the average accepted length.

[0075] 4.3 Ablation study In this section, we explore the effectiveness of different design schemes in our approach.

[0076] 4.3.1 Loss Penalty Figure 8 : Average number of decoded blocks τ and throughput Tok / s of different methods on six benchmarks, including the standard EAGLE speculative decoding method with a fixed draft length between 2 and 6 tokens.

[0077] Fig. 9 : Ablation experiment results on MT-Bench and Alpaca on training loss design. "w / o Penalty" means there is no L1 loss penalty when training the model. "w / Cls" means the model is trained with classification loss.

[0078] Fig.10 : Comparison of draft length distribution with and without lossy penalty.

[0079] We train AdaEAGLE using an L1 regression loss to predict the logarithm of the adaptive draft length. As mentioned before, shorter lengths have a larger tolerance than longer predictions. To address this issue, we employ a penalized L1 loss during training to penalize predictions that are shorter than the ground truth label. Fig. 9 The impact of adopting punitive loss is compared, showing that it improves the average acceptance length and throughput. Fig.10The distribution of draft lengths with and without loss penalties is further illustrated, highlighting that loss penalties encourage longer draft lengths.

[0080] 4.3.2 Regression and Classification We also explored other methods for predicting draft length, including regression and classification. In regression, the model maps hidden states to a continuous score, while in classification, the model maps hidden states to a set of predefined categories. Fig. 9 As shown in Figure 2, regression outperforms classification because regression can naturally model the sequential relationship between different labels.

[0081] 4.3.3 AdaEAGLE-DDD Threshold Fig.11 : Throughput results (Tok / s) for different thresholds. The threshold corresponds to the log probability and the result is negative. The ablation is performed on the MT-Bench dataset.

[0082] Fig.11 The performance of the DDD part of AdaEAGLE-DDD at different thresholds is shown, showing that the throughput remains stable and insensitive to different settings. Based on these results, we propose Figure 8 A threshold of -0.6 was chosen for the experiments reported in.

[0083] 5 Related works LLM is popular and powerful. However, its AR inference is slow and expensive. Important methods for efficient LLM inference have been proposed (Zhou et al., 2024). This section will focus on the speculative decoding (SD) technique and the adaptive draft structure for speculative decoding (SD).

[0084] Speculative Decoding Speculative decoding (SD) uses the draft-verify paradigm to achieve lossless acceleration. According to the type of draft model, speculative decoding (SD) frameworks can be roughly divided into the following categories (Xia et al., 2024): (1) Independent draft models. SpecDec (Xia et al., 2023) introduces a non-AR Transformer to generate multiple tokens simultaneously. In addition, many studies suggest using smaller pre-trained models from the same LLM family to accelerate inference, which avoids additional training and maintains consistency in prediction behavior due to shared tokenizers and training corpora. (2) Self-drafting eliminates the need for an external drafting model by leveraging the target LLM itself for drafting. Recent techniques adopt lightweight FFN heads to achieve parallel (Cai et al., 2024) or AR (Li et al., 2024b) token generation, thereby reducing computational overhead. Other studies (Yang et al., 2023; Zhang et al., 2024) explore early exit and layer skipping in the target model to improve drawing efficiency. However, most studies use static draft structures instead of adaptive draft structures. Adaptive draft structures refer to the fact that the draft model can dynamically adjust its generation structure (such as the length of the draft sequence or the depth / width / shape of the draft tree) according to the inferred context. Although some studies have recognized the need to make the draft structure dynamic, these methods still lack specialized and systematic methods, such as Figure 3 Regarding the modeling methods of draft boundaries, some works adopt intrinsic methods and use the confidence of the output distribution to determine the stopping point, such as (Zhang et al., 2024; Li et al., 2024a; Liu et al., 2024a; Brown et al., 2024; Wang et al., 2024), while other works (Huang et al., 2024; Mamou et al., 2024) adopt additional modules, such as additional trainable heads, to determine whether to stop generating drafts.

[0085] In addition, there are some works orthogonal to the above, such as Sequoia (Chen et al., 2024), which is hardware-aware adaptation rather than context-aware adaptation. In summary, the current models that support adaptive draft structures lack performance, timeliness, and simplicity. The method proposed in the embodiment of this application is the first speculative decoding (SD) framework based on EAGLE to directly model the draft structure and achieve optimal performance.

[0086] 6 Conclusion In this application example, we propose a novel speculative decoding framework AdaEAGLE, which is the first framework to achieve adaptive draft control by explicitly modeling the optimal draft structure. By combining the lightweight draft length prophet (LDLP) module, AdaEAGLE predicts the optimal draft length to guide token generation, thereby minimizing the number of forward passes for large target models and maximizing the overall decoding speed. Experimental results show that the decoding throughput achieved by AdaEAGLE is closer to the oracle top line compared to existing speculative decoding (SD) methods. It has a lossless speed improvement of 1.61% over vanilla AR decoding, exceeds 2% of fixed-length SotA, and shows versatility and robustness in various test sets. Our method opens the door to modeling adaptive draft structures in an explicit and dedicated "end-to-end" manner, and can be integrated into a wide range of speculative decoding (SD) frameworks and sampling strategies.

[0087] limitation Unlike EAGLE (Li et al., 2024b) and EAGLE-2 (Li et al., 2024a) which use tree-based decoding, the method proposed in the embodiment of this application is currently limited to the sequential draft structure, so we only provide results under the greedy (top-1) decoding strategy. While the sequential structure allows us to clearly demonstrate our core ideas with a single goal (length), the tree-based structure requires more complex and fine-grained control, such as the width and depth of each subtree, which we will leave to future work to solve.

[0088] In addition, we can further explore how LDLP can improve accuracy. Although the examples in this application are based on EAGLE improvements, in fact, the concept of directly modeling adaptive draft structures can be applied to many other frameworks, such as (Zhang et al., 2024). Therefore, future research needs to optimize LDLP, evaluate different speculative decoding (SD) frameworks, and examine the interactions between these elements.

[0089] Appendix A Analysis of Optimal Draft Length Guided by Oracle Assume the input is a sequence of tokens , the output hidden states of the second to top layers of the target model and the AR head of the draft model are and . Let the token sequence generated so far be ,based on The draft token sequence of . Under the greedy decoding strategy, given the sequence history , no matter what the draft model outputs, the output of the target model is certain. We call the token sequence generated by the original target model (without the draft) the formal token sequence.

[0090] In SpecDec++ (Huang et al., 2024), the optimal draft length is defined as the length that minimizes the number of forward passes required by the target model. It can be inferred that the target model accepts , but refused This happens because in the SpecDec++ framework, a new draft token is generated Depends only on token history . Generate more than The candidate tokens of is a suboptimal candidate token because the rejected token will still be generated in the next draft iteration, which requires an additional forward pass of the target model.

[0091] However, in the EAGLE framework, the generation of the draft depends not only on the token history but also on the hidden states produced by the target model and the draft model, i.e. and If fewer than 1 candidate drafts are generated ,For example , then in the next draft iteration, the hidden history will change, which means that different draft tokens may be generated at the same position , rather than being rejected . This may be equivalent to a token , is then accepted by the target model, thus providing an opportunity to reduce the forward pass of the target model. This means that even if we know that the next draft token is likely to be accepted, greedily generating it is no longer guaranteed to be optimal.

[0092] To address this counterintuitive situation, this application introduces a reasonable assumption to solve this difficulty.

[0093] Assumption 1. Assumption and is the length of the history of generated formal tokens. The current draft model is generating a draft token at position m. In the EAGLE framework, we believe that:

[0094] In summary, the more hidden states the target has, the more accurate the draft's guess will be. This is reasonable because the draft's hidden states and tokens are estimates of the target. Since the target model has a longer history, the draft can generate an output distribution that is closer to the target model. With this assumption, under the framework of EAGLE, we get the following theorem: Theorem 1. Assume that the generated formal token is In the EAGLE framework, the optimal draft length is still , so that the target model accepts , but refused .

[0095] Proof Obviously, the draft length is greater than Still not optimal. = In the next draft iteration, the draft model is Start with End, that is . Call this draft plan Plan A.

[0096] When the generated candidate information is less than When k′ token is .because , so according to assumption 1:

[0097] This means that in the next draft iteration, the decoding progress of Scheme B cannot exceed that of Scheme A. Similarly, after the same number of iterations, Scheme B is still left behind by Scheme A. Similarly, when generating the same number of formal tokens, Scheme A requires fewer forward passes of the target model than Scheme B.

[0098] Theorem 1 can also be easily extended to the expected value of non-free decodings.

[0099] In other embodiments, embodiments of the present invention further provide a non-volatile computer storage medium storing computer executable instructions, the computer executable instructions being capable of executing the speculative decoding optimization method in any of the above method embodiments for a draft model and a target verification model; As an implementation mode, the non-volatile computer storage medium of the present invention stores computer executable instructions, and the computer executable instructions are configured as follows: Inputting the draft content generated by the previous round of the draft model into the target verification model; performing parallel validation on tokens in the draft content in the target validation model; A lightweight draft length predictor is used to analyze the latent vector of the last received token position and the last token generated to predict the optimal draft length for the next round. A corresponding number of drafts is autoregressively generated by the draft model based on the optimal draft length predicted by the lightweight draft length predictor.

[0100] The non-volatile computer-readable storage medium may include a program storage area and a data storage area, wherein the program storage area may store an operating system, an application required for at least one function; the data storage area may store data created according to the use of the speculative decoding optimization method and system, etc. In addition, the non-volatile computer-readable storage medium may include a high-speed random access memory, and may also include a non-volatile memory, such as at least one disk storage device, a flash memory device, or other non-volatile solid-state storage device. In some embodiments, the non-volatile computer-readable storage medium may optionally include a memory remotely arranged relative to the processor, and these remote memories may be connected to the speculative decoding optimization method via a network. Examples of the above-mentioned network include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0101] An embodiment of the present invention also provides a computer program product, which includes a computer program stored on a non-volatile computer-readable storage medium, and the computer program includes program instructions. When the program instructions are executed by a computer, the computer executes any of the above-mentioned speculative decoding optimization methods.

[0102] Fig.12 is a schematic diagram of the structure of an electronic device provided by an embodiment of the present invention, such as Fig.12 As shown, the device includes: one or more processors 1210 and a memory 1220, Fig.12 A processor 1210 is taken as an example. The device of the speculative decoding optimization method and system may also include: an input device 1230 and an output device 1240. The processor 1210, the memory 1220, the input device 1230 and the output device 1240 may be connected via a bus or other means. Fig.12 In the example, the bus connection is used. The memory 1220 is the above-mentioned non-volatile computer-readable storage medium. The processor 1210 executes various functional applications and data processing of the server by running the non-volatile software programs, instructions and modules stored in the memory 1220, that is, the speculative decoding optimization method of the above-mentioned method embodiment is implemented. The input device 1230 can receive input digital or character information, and generate key signal input related to user settings and function control of the speculative decoding optimization device. The output device 1240 may include a display device such as a display screen.

[0103] The above product can execute the method provided by the embodiment of the present invention, and has the functional modules and beneficial effects corresponding to the execution method. For technical details not described in detail in this embodiment, please refer to the method provided by the embodiment of the present invention.

[0104] As an implementation mode, the electronic device is applied to a speculative decoding optimization device, comprising: at least one processor; and a memory connected to the at least one processor in communication; wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor so that the at least one processor can: Inputting the draft content generated by the previous round of the draft model into the target verification model; performing parallel validation on tokens in the draft content in the target validation model; A lightweight draft length predictor is used to analyze the latent vector of the last received token position and the last token generated to predict the optimal draft length for the next round. A corresponding number of drafts is autoregressively generated by the draft model based on the optimal draft length predicted by the lightweight draft length predictor.

[0105] The electronic device of the embodiment of the present application exists in various forms, including but not limited to: (1) Mobile communication equipment: This type of equipment is characterized by having mobile communication functions and its main purpose is to provide voice and data communications. This type of terminal includes: smart phones, multimedia phones, functional phones, and low-end phones.

[0106] (2) Ultra-mobile personal computer devices: These devices fall into the category of personal computers, have computing and processing functions, and generally also have mobile Internet access features. These terminals include: PDA, MID and UMPC devices, etc.

[0107] (3) Portable entertainment devices: These devices can display and play multimedia content. They include audio and video players, handheld game consoles, e-books, smart toys, and portable car navigation devices.

[0108] (4) Server: A device that provides computing services. The server consists of a processor, hard disk, memory, system bus, etc. The server is similar to a general computer architecture, but because it needs to provide highly reliable services, it has higher requirements in terms of processing power, stability, reliability, security, scalability, and manageability.

[0109] (5) Other electronic devices with data interaction functions.

[0110] The device embodiments described above are merely illustrative, wherein the units described as separate components may or may not be physically separated, and the components shown as units may or may not be physical units, i.e., they may be located in one place, or they may be distributed on multiple network units. Some or all of the modules may be selected according to actual needs to achieve the purpose of the scheme of this embodiment. Those of ordinary skill in the art may understand and implement it without creative effort.

[0111] Through the description of the above implementation modes, those skilled in the art can clearly understand that each implementation mode can be implemented by means of software plus a necessary general hardware platform, or of course by hardware. Based on such an understanding, the above technical solution can essentially or in other words be embodied in the form of a software product that contributes to the prior art. The computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, a magnetic disk, an optical disk, etc., and includes a number of instructions for a computer device (which can be a personal computer, a server, or a network device, etc.) to execute the methods of each embodiment or some parts of the embodiment.

[0112] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A speculative decoding optimization method for a draft model and a target verification model, the method comprising: Inputting the draft content generated by the previous round of the draft model into the target verification model; performing parallel validation on tokens in the draft content in the target validation model; A lightweight draft length predictor is used to analyze the latent vector of the last received token position and the last token generated to predict the optimal draft length for the next round. A corresponding number of drafts is autoregressively generated by the draft model based on the optimal draft length predicted by the lightweight draft length predictor.

2. The method according to claim 1, wherein: The process of constructing the training data of the lightweight draft length predictor includes: Using the target verification model to generate standard outputs based on input prompts; Processing at any position of the generated output, using the draft model to generate a certain upper limit of draft tokens; The output of the draft model is compared with the standard output of the target verification model, and the length of the longest common prefix is ​​calculated as the optimal draft length target for the corresponding position.

3. The method according to claim 2, wherein: The true values ​​and predicted values ​​in the training data are trained using L1 loss.

4. The method according to claim 3, wherein: When training by L1 loss, the method further comprises: A penalty coefficient is introduced to amplify the L1 loss when the predicted value is smaller than the true value.

5. The method according to any one of claims 1 to 4, wherein: A three-layer multilayer perceptron with residual connections is used as the model structure of the lightweight draft length predictor, and the output hidden state is mapped to a length scalar through a linear transformation.

6. The method according to any one of claims 1 to 5, wherein: The method is based on sequential scratch generation and greedy decoding, where the optimal scratch length is defined as the length that minimizes the number of forward passes required to verify the target model.

7. An electronic device, comprising: At least one processor, and a memory communicatively connected to the at least one processor, wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor to enable the at least one processor to perform the steps of the method described in any one of claims 1 to 6.

8. A storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 6 are implemented.

Images