Speculative sampling inference acceleration method, apparatus, and electronic device
By combining small and large models in a speculative sampling method, and utilizing the Monte Carlo search tree algorithm and reinforcement learning to evaluate candidate paths, the memory access bottleneck in the autoregressive decoding stage of large language models is solved, the acceptance rate and diversity of candidate path sequences are improved, and speculative sampling inference is accelerated.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- BEIJING CENTURY TAL EDUCATION TECH CO LTD
- Filing Date
- 2026-01-09
- Publication Date
- 2026-04-28
AI Technical Summary
Large language models suffer from a delay in token generation due to memory access bandwidth bottlenecks during the autoregressive decoding stage. Existing small models lack path diversity and dynamic adaptation when generating candidate path sequences, which affects the candidate acceptance rate.
A small model is used to generate candidate path sequences. By combining the Monte Carlo search tree algorithm and reinforcement learning, the small model is dynamically updated to improve the acceptance rate of candidate path sequences by evaluating the comprehensive potential value, basic reward value, diversity reward value and exploration reward value of the candidate paths.
It significantly improves the acceptance rate of candidate path sequences generated by small models, reduces the delay of large language models in the autoregressive decoding stage, and enhances path diversity and adaptability.
Smart Images

Figure CN121503699B_ABST
Abstract
Description
Technical Field
[0001] This application relates to artificial intelligence technology, and in particular to methods, apparatus and electronic devices for accelerating speculative sampling inference. Background Technology
[0002] Large language models generate tokens (tokens / lexicals) sequentially during the autoregressive decoding stage. When generating each new token, they need to access historical tokens stored in memory to ensure that the generation of each new token depends on previous historical tokens (equivalent to integrating the entire historical context) to guarantee the semantic coherence of the generated content. However, in practical applications, memory access bandwidth becomes a bottleneck, which delays the token generation time of the large language model during the autoregressive decoding stage. Summary of the Invention
[0003] This application provides a method, apparatus, and electronic device for accelerating speculative sampling inference, so as to avoid delays in the generation of tokens by large language models during the autoregressive decoding stage due to bandwidth bottlenecks in memory access.
[0004] This application provides a method for accelerating speculative sampling inference. This method is applied to speculative sampling scenarios and includes:
[0005] Obtain a candidate path sequence generated by the small model based on user input; each candidate path in the candidate path sequence is a path that meets the requirements of the comprehensive potential value of the path, and the comprehensive potential value of any candidate path is determined based at least on the prediction probability of each token predicted by the small model in the candidate path, the confidence of the candidate path, and the length of the candidate path;
[0006] The reward parameters for each candidate path in the candidate path sequence are determined. The reward parameters for any candidate path are determined based on its base reward value, diversity reward value, and exploration reward value. Specifically, the base reward value for any candidate path is determined based on the acceptance of the candidate path by the large model, the diversity reward value for any candidate path is determined based on the difference between the candidate path and existing historical paths, the exploration reward value for any candidate path is determined based on the frequency of the candidate path's occurrence in existing historical paths and the potential value of the candidate path, and the potential value of any candidate path is determined based on the predicted probabilities of each token in the candidate path output by the small model.
[0007] The small model is updated based on each candidate path in the candidate path sequence and the reward parameters of each candidate path; if the small model does not meet the optimal conditions, when user input is received again, the obtained small model is returned to generate a candidate path sequence based on user input until the small model meets the optimal conditions.
[0008] A speculative sampling inference acceleration device, applied in speculative sampling scenarios, the device comprising:
[0009] A path unit is used to obtain a candidate path sequence generated by a small model based on user input; each candidate path in the candidate path sequence is a path that meets the requirements of the comprehensive potential value of the path, and the comprehensive potential value of any candidate path is determined based at least on the prediction probability of each token predicted by the small model in the candidate path, the confidence of the candidate path, and the length of the candidate path.
[0010] A determining unit is used to determine the reward parameters of each candidate path in the candidate path sequence; the reward parameters of any candidate path are determined based on the basic reward value, diversity reward value, and exploration reward value of the candidate path; wherein, the basic reward value of any candidate path is determined based on the acceptance of the candidate path by the large model, the diversity reward value of any candidate path is determined based on the difference between the candidate path and existing historical paths, the exploration reward value of any candidate path is determined based on the frequency of the candidate path appearing in existing historical paths and the potential value of the candidate path, and the potential value of any candidate path is determined based on the predicted probability of each token in the candidate path output by the small model;
[0011] The processing unit is used to update the small model based on each candidate path in the candidate path sequence and the reward parameters of each candidate path; if the small model does not meet the optimal conditions, when user input is received again, the obtained small model is used to generate a candidate path sequence based on user input until the small model meets the optimal conditions.
[0012] This application also provides an electronic device. The electronic device includes: a processor and a machine-readable storage medium;
[0013] The machine-readable storage medium stores machine-executable instructions that can be executed by the processor;
[0014] The processor is used to execute machine-executable instructions to implement the steps of the disclosed method.
[0015] As can be seen from the above technical solutions, in this application, the small model and the large model work together to perform speculative sampling, avoiding the delay in the large language model generating tokens in the autoregressive decoding stage due to the bandwidth bottleneck of memory access.
[0016] Furthermore, in this embodiment, the small model evaluates the comprehensive potential value of the candidate path based on several dimensions, including the prediction probability of each predicted token in the candidate path, the confidence level of the candidate path, and the length of the candidate path. The candidate path sequence is generated based on the comprehensive potential value of each candidate path, which can improve the acceptance rate of the candidate path sequence by the large model.
[0017] Furthermore, this embodiment updates the small model by using each candidate path in the candidate path sequence and the reward parameters of each candidate path. Since the reward parameters of the candidate path are determined based on the candidate path's base reward value (determined based on the large model's acceptance of the candidate path), diversity reward value (determined based on the difference between the candidate path and existing historical paths), and exploration reward value (determined based on the frequency of the candidate path's appearance in existing historical paths and the candidate path's potential value), this ensures that the updated small model adapts to the large model's preferences (or adjusts the small model based on the large model's distribution), significantly improving the small model's acceptance rate in the next round of candidate path sequences generated based on user input (i.e., significantly improving the acceptance rate of speculative sampling), which is equivalent to accelerating speculative sampling inference.
[0018] Furthermore, in specific applications, this embodiment will autonomously optimize and update the small model based on the candidate paths generated by the small model in the application and the reward parameters of each candidate path. This dynamic self-updating method of the small model will enable the small model to make up for the shortcomings of traditional small models such as N-gram models that generate candidate path sequences based on static features, improve path diversity, and tend to generate candidate paths with higher diversity reward values and exploration reward values, thereby improving the acceptance rate of candidate path sequences. Attached Figure Description
[0019] The accompanying drawings, which are incorporated in and form a part of this specification, illustrate embodiments consistent with this disclosure and, together with the description, serve to explain the principles of this disclosure.
[0020] Figure 1 A flowchart illustrating the method provided in this application embodiment;
[0021] Figure 2 This is a schematic diagram of speculative sampling provided for an embodiment of this application;
[0022] Figure 3 A flowchart of the small model self-optimization process provided in the embodiments of this application;
[0023] Figure 4 A flowchart of the small model initialization process provided in the embodiments of this application;
[0024] Figure 5 This is a structural diagram of the device provided in the embodiments of this application;
[0025] Figure 6 This is a structural diagram of an electronic device provided in an embodiment of this application. Detailed Implementation
[0026] Speculative sampling is a sampling strategy used in fields such as natural language processing for text generation. In practice, speculative sampling takes a tricky approach, considering multiple possible next words at once and sampling and predicting in advance along these different possibilities.
[0027] Applying speculative sampling to the autoregressive decoding stage can fundamentally avoid delays in the generation of tokens (also known as candidate paths) by large language models during the autoregressive decoding stage due to memory access bandwidth bottlenecks.
[0028] In speculative sampling methods, a small model with lightweight and fast inference capabilities is often used to generate candidate token sequences (candidate output tokens) in an autoregressive manner. A large model is then used to validate and refine these candidate token sequences to determine which tokens to include or exclude from the candidate token sequences output by the small model, thereby reducing inference steps and increasing token generation speed.
[0029] Here, the differences between small and large models in the field of artificial intelligence are mainly reflected in multiple dimensions such as parameter scale, training data requirements, computational resource consumption, application scenarios, and capability characteristics. For example, large models typically have billions to trillions of parameters (such as GPT-3's 175 billion parameters), require massive amounts of data for training, and are suitable for general tasks; small models have a parameter scale of millions to billions (such as BERT-base's 110 million parameters), have lower data requirements, and are specialized for specific scenarios. Here, small models are also called draft models, such as N-gram models, while large models are also called target models, such as large language models (LLM).
[0030] In practical applications, small models in speculative sampling methods, such as traditional N-gram models, often suffer from the following problems:
[0031] Insufficient path diversity: The statistical characteristics of small models may lead to overly conservative paths for generating candidate sequences, making it impossible to effectively explore candidate sequences with low probability but high semantic consistency.
[0032] Lack of dynamic adaptation: Small models rely heavily on fixed rules (such as Beam Search) when generating candidate sequences, making it difficult to dynamically adjust based on task context or feedback from large models.
[0033] The aforementioned issues directly impact the core metric in speculative sampling methods: the acceptance rate. Here, the acceptance rate refers to the proportion of tokens in the candidate token sequence generated by the smaller model that are accepted by the larger model. A high acceptance rate indicates high prediction quality from the smaller model, consistent with the output distribution of the larger model, thus reducing the number of rejections during the validation phase and improving speedup.
[0034] To address the aforementioned technical problems, this embodiment provides a method for accelerating speculative sampling inference. To enable those skilled in the art to better understand the technical solutions provided in this application's embodiments, and to make the above-mentioned objectives, features, and advantages of this application's embodiments more apparent and understandable, the technical solutions in this application's embodiments are further described in detail below with reference to the accompanying drawings.
[0035] See Figure 1 , Figure 1 This is a flowchart illustrating a method provided in an embodiment of this application. This method is applied to speculative sampling scenarios. Figure 1 As shown, the method includes:
[0036] Step 101: Obtain a candidate path sequence generated by the small model based on user input; each candidate path in the candidate path sequence is a path that meets the requirements of the comprehensive potential value of the path, and the comprehensive potential value of any candidate path is determined based at least on the prediction probability of each token predicted by the small model in the candidate path, the confidence level of the candidate path, and the length of the candidate path.
[0037] In this embodiment, the small model utilizes the Monte Carlo Tree Search (MCTS) algorithm to generate candidate paths. Specifically, based on a heuristic decision-making algorithm combining stochastic simulation and tree search, the small model dynamically balances the exploration of optimal candidate paths under limited computing power through a four-step iterative process: selection (choosing the optimal path from the root of the search tree) → expansion (expanding the search tree, such as exploring new nodes) → simulation (evaluating the path) → backpropagation (updating the visit counts of all nodes on the path, etc.), ultimately yielding a sequence of candidate paths. Specifically, based on the current context token sequence, the small model selects child nodes in the search tree according to the UCT formula (Upper Confidence Bounds for Trees), then expands the Top-k candidate nodes according to the policy network, and uses lightweight Rollout (simulation / deduction) to predict the path acceptance rate. Finally, it updates the node visit counts and cumulative rewards, outputting the Top-M (the top M candidate paths with the highest rewards).
[0038] In this embodiment, the small model uses MCTS search to generate candidate path sequences, which can explore the diversity of paths, escape the local optimum trap, start from the globally better path, and increase the probability that the candidate paths generated by the small model are accepted by the large model such as LLM.
[0039] In addition, this embodiment uses a small model, such as an N-gram model, to generate candidate path sequences. By leveraging the characteristics of small models, such as speed and lightweightness, lightweight path evaluation can be achieved.
[0040] It should be noted that in this embodiment, the small model does not blindly generate candidate path sequences. The aforementioned candidate path sequences are generated by the small model through pruning strategies, such as dynamically removing low-potential candidate paths. As an example, the potential of any candidate path can be characterized by its comprehensive potential value. That is, each candidate path in the aforementioned candidate path sequence is a path that meets the comprehensive potential value requirement (e.g., the comprehensive potential value is greater than a set threshold).
[0041] As an example, the overall path potential value of any candidate path is determined at least based on the predicted probability of the tokens predicted by the small model in the candidate path, the confidence level of the candidate path, and the length of the candidate path. Optionally, the overall path potential value of any candidate path is positively correlated with the inverse of the predicted probability of each token predicted by the small model, the confidence level of the candidate path, and the length of the candidate path.
[0042] In this embodiment, any candidate path may include tokens predicted by the small model. When predicting each token, the small model also predicts the probability of that token.
[0043] Therefore, for any candidate path, when calculating the overall potential value of the candidate path based on the predicted probabilities of tokens in that candidate path predicted by the aforementioned small model, the local confidence of the candidate path is first determined based on the predicted probabilities of the tokens predicted by the small model. As an example, the local confidence of the candidate path can be the geometric mean of the product of the predicted probabilities of the tokens predicted by the small model in that candidate path. For instance, if a candidate path is "I'm going to the mall," meaning the candidate path consists of 6 consecutive tokens with predicted probabilities of 0.6, 0.3, 0.8, 0.9, 0.2, and 0.5 respectively, then the product of 0.6, 0.3, 0.8, 0.9, 0.2, and 0.5 (0.6*0.3*0.8*0.9*0.2*0.5) is calculated, and the square root of this product yields the aforementioned geometric mean. The reason for calculating the geometric mean of the product of the predicted probabilities of each token is to prevent the overall path potential value of candidate paths from changing due to the size of N. Here, the overall path potential value of any candidate path is positively correlated with the local confidence of that candidate path.
[0044] In this embodiment, the confidence score of a candidate path is also calculated based on the predicted probability of the token in that candidate path, for example, according to the following formula: . This represents the predicted probability of the i-th token in the candidate path. A higher confidence level for a candidate path indicates a higher acceptance rate by the larger model. The overall potential value of any candidate path is positively correlated with its confidence level.
[0045] In this embodiment, when determining the comprehensive potential value of a candidate path, the length of the candidate path is further considered to avoid increased performance overhead during large model validation due to excessive length. The comprehensive potential value of any candidate path is positively correlated with the inverse of its length. Here, considering the length of the candidate path when determining its comprehensive potential value is equivalent to introducing a path length penalty.
[0046] By employing the pruning method described above, the resulting candidate path sequence can be used as training data for subsequent reinforcement learning-based updates of the small model. The updating of the small model will be described below.
[0047] Step 102: Determine the reward parameters for each candidate path in the candidate path sequence. The reward parameters for any candidate path are determined based on its base reward value, diversity reward value, and exploration reward value. Specifically, the base reward value for any candidate path is determined based on the acceptance of the candidate path by the large model, the diversity reward value for any candidate path is determined based on the difference between the candidate path and existing historical paths, the exploration reward value for any candidate path is determined based on the frequency of the candidate path's occurrence in existing historical paths and the potential value of the candidate path, and the potential value of any candidate path is determined based on the predicted probabilities of each token in the candidate path output by the small model.
[0048] In this embodiment, the base reward value for any candidate path is determined based on the acceptance of that candidate path by the larger model. For example, such as Figure 2 As shown, the candidate path sequence generated by the small model is input into the large model, such as an LLM, so that the large model can select and reject candidate paths in the candidate path sequence. For example, based on its own probability distribution, the large model determines whether the candidate path generated by the small model, such as a token, conforms to the large model's generation intention. If it does, the candidate path is accepted; otherwise, the candidate path and all paths following it are rejected. Optionally, if the candidate path is accepted by the large model, the base reward value of the candidate path can be a first set value, such as 1; otherwise, it can be a second set value, such as 0.
[0049] In this embodiment, the diversity reward value of any candidate path is negatively correlated with the frequency of the candidate path appearing in existing historical paths. The less frequently a candidate path appears in existing historical paths, the higher its diversity reward value; conversely, the more frequently a candidate path appears in existing historical paths, the lower its diversity reward value.
[0050] In this embodiment, the exploration reward value of any candidate path is negatively correlated with the frequency of its occurrence in existing historical paths and positively correlated with its potential value. Here, the potential value of the candidate path is negatively correlated with information entropy. Information entropy is determined based on the predicted probabilities of each token within the candidate path as predicted by the aforementioned small model.
[0051] Step 103: Update the small model based on each candidate path in the candidate path sequence and the reward parameters of each candidate path; if the small model does not meet the optimal conditions, return to step 101 when user input is received again, until the small model meets the optimal conditions.
[0052] In step 101 above, when the small model generates any candidate path based on the MCTS algorithm, it also generates at least one reference path that does not meet the path comprehensive potential value requirement; the path comprehensive potential value of the reference path is lower than that of the candidate path. The candidate path and the reference path together form a path pair.
[0053] As an example, the reward parameters of the reference path can be further determined in the same manner as the reward parameters of the candidate path determined in step 102 above. Here, the reward parameters of any candidate path are superior to the reward parameters of the reference path corresponding to that candidate path.
[0054] Based on this, in step 103, the small model can be updated using reinforcement learning based on each path in the path pair and its reward parameters. Since there are candidate paths with optimal reward parameters and reference paths with non-optimal reward parameters in the path pair, by using the reward parameters of the candidate paths and the reference paths in the path pair, and combining the specific determination method of the reward parameters, the updated small model can be made to adapt to the preferences of the large model (or adjust the small model based on the distribution of the large model), significantly improving the acceptance rate of the small model in the next round of candidate path sequences generated based on user input (i.e., significantly improving the acceptance rate of speculative sampling), which is equivalent to accelerating speculative sampling inference. Figure 3 An example illustrates the process of model optimization.
[0055] It should be noted that once the small model meets the optimal conditions, there is no need to adjust the small model further. Instead, the small model can be used to generate candidate path sequences based on subsequent user inputs, and the large model can be used to verify the candidate path sequences.
[0056] Optionally, in this embodiment, the small model is the deep learning model of the Decoder architecture. The small model is based on the KVCache mechanism, which caches the key and value of the token each time a token is generated; and generates a token based on the cached key and value of the token to reduce the latency of token generation.
[0057] This concludes the process. Figure 1 The process is shown below.
[0058] pass Figure 1 As can be seen from the process shown, speculative sampling is performed by combining small and large models to avoid delays in the generation of tokens by the large language model during the autoregressive decoding stage due to bandwidth bottlenecks in memory access.
[0059] Furthermore, in this embodiment, the small model evaluates the comprehensive potential value of the candidate path based on several dimensions, including the prediction probability of each predicted token in the candidate path, the confidence level of the candidate path, and the length of the candidate path. The candidate path sequence is generated based on the comprehensive potential value of each candidate path, which can improve the acceptance rate of the candidate path sequence by the large model.
[0060] Furthermore, this embodiment updates the small model by using each candidate path in the candidate path sequence and the reward parameters of each candidate path. Since the reward parameters of the candidate path are determined based on the candidate path's base reward value (determined based on the large model's acceptance of the candidate path), diversity reward value (determined based on the difference between the candidate path and existing historical paths), and exploration reward value (determined based on the frequency of the candidate path's appearance in existing historical paths and the candidate path's potential value), this ensures that the updated small model adapts to the large model's preferences (or adjusts the small model based on the large model's distribution), significantly improving the small model's acceptance rate in the next round of candidate path sequences generated based on user input (i.e., significantly improving the acceptance rate of speculative sampling), which is equivalent to accelerating speculative sampling inference.
[0061] Furthermore, in specific applications, this embodiment will autonomously optimize and update the small model based on the candidate paths generated by the small model in the application and the reward parameters of each candidate path. This dynamic self-updating method of the small model will enable the small model to make up for the shortcomings of traditional small models such as N-gram models that generate candidate path sequences based on static features, improve path diversity, and tend to generate candidate paths with higher diversity reward values and exploration reward values, thereby improving the acceptance rate of candidate path sequences.
[0062] It should be noted that in this embodiment, the small model is initially initialized using knowledge distillation to allow it to learn from the large model and achieve performance close to that of the large model. Knowledge distillation is a model optimization technique that transfers the knowledge (such as complex feature representations and soft probability distributions between categories) contained in a large model with many parameters and high performance to a small model with a simpler structure and fewer parameters through a specific loss function design. This improves the performance of the small model while simultaneously achieving model compression and accelerated inference. Figure 4 An example is provided illustrating the process of initializing a small model.
[0063] like Figure 4 As shown, the process may include the following steps:
[0064] Step 401, Data Acquisition.
[0065] Optionally, this embodiment can generate multi-turn dialogue data based on a database in a specified domain, such as a code repository, and use a mainstream pre-trained large model (such as Llama 3.1-7B) to generate high-quality responses for each turn of dialogue data. Optionally, this embodiment can also further record the intermediate states of the pre-trained large model during the response generation process (such as attention weights, hidden layer vectors of the last layer), and finally construct an offline training dataset. Each training dataset may include, for example,: ([one turn of dialogue data, the large model's response to that turn of dialogue data, the intermediate states of the large model during the response generation process]).
[0066] Step 402: Train the small model using the knowledge distillation method and the specified loss function and the training data mentioned above, so that the small model learns the knowledge of the large model and its performance is close to that of the large model.
[0067] In step 402, the loss function introduced during small model training is set according to actual needs. For example, distillation loss such as KL divergence is used to measure the difference in output distribution between small and large models based on the same dialogue data; and / or reinforcement learning (RL) loss determined based on the proximal policy optimization (PPO) algorithm, etc. This embodiment does not specifically limit the loss function.
[0068] Thus, through Figure 4 The process shown demonstrates how to initialize a small model based on knowledge distillation, enabling the small model to learn the knowledge of the large model and achieve performance close to that of the large model.
[0069] The methods provided in the embodiments of this application have been described above. The apparatus provided in the embodiments of this application is described below:
[0070] See Figure 5 , Figure 5 This is a structural diagram of the device provided in an embodiment of this application. Figure 5 As shown, the device may include:
[0071] A path unit is used to obtain a candidate path sequence generated by a small model based on user input; each candidate path in the candidate path sequence is a path that meets the requirements of the comprehensive potential value of the path, and the comprehensive potential value of any candidate path is determined based at least on the prediction probability of each token predicted by the small model in the candidate path, the confidence of the candidate path, and the length of the candidate path.
[0072] A determining unit is used to determine the reward parameters of each candidate path in the candidate path sequence; the reward parameters of any candidate path are determined based on the basic reward value, diversity reward value, and exploration reward value of the candidate path; wherein, the basic reward value of any candidate path is determined based on the acceptance of the candidate path by the large model, the diversity reward value of any candidate path is determined based on the difference between the candidate path and existing historical paths, the exploration reward value of any candidate path is determined based on the frequency of the candidate path appearing in existing historical paths and the potential value of the candidate path, and the potential value of any candidate path is determined based on the predicted probability of each token in the candidate path output by the small model;
[0073] The processing unit is used to update the small model based on each candidate path in the candidate path sequence and the reward parameters of each candidate path; if the small model does not meet the optimal conditions, when user input is received again, the obtained small model is used to generate a candidate path sequence based on user input until the small model meets the optimal conditions.
[0074] Optionally, the processing unit further includes:
[0075] Initially, the small model is initialized based on knowledge distillation so that it learns the knowledge of the large model and its performance approaches that of the large model.
[0076] Optionally, the overall potential value of any candidate path is positively correlated with the predicted probability of each token predicted by the small model, the confidence level of the candidate path, and the inverse of the length of the candidate path.
[0077] Optionally, the diversity reward value of any candidate path is negatively correlated with the frequency of the candidate path appearing in existing historical paths, wherein the less frequently a candidate path appears in existing historical paths, the higher the diversity reward value of the candidate path.
[0078] And / or, the exploration reward value of any candidate path is negatively correlated with the frequency of the candidate path appearing in existing historical paths and positively correlated with the potential value of the candidate path;
[0079] And / or, the potential value of any candidate path is negatively correlated with the entropy determined based on the predicted probability of each token predicted by the small model in that candidate path.
[0080] Optionally, the small model determines the candidate path sequence based on the Monte Carlo Tree Search (MCTS) algorithm; when generating any candidate path in the candidate path sequence based on the MCTS algorithm, the small model also generates at least one reference path that does not meet the path comprehensive potential value requirement; the path comprehensive potential value of the reference path is lower than the path comprehensive potential value of the candidate path.
[0081] The determining unit is further configured to: determine the reward parameters of the reference path for each candidate path in accordance with the method for determining the reward parameters of each candidate path in the candidate path sequence;
[0082] The processing unit updates the small model based on each candidate path in the candidate path sequence and the reward parameters of each candidate path, including:
[0083] The small model is updated based on each path in the path pair and the reward parameters of each path; the path pair includes a candidate path and a reference path for the candidate path.
[0084] Optionally, the small model is based on a KV Cache mechanism, which caches the key and value matrix of the token each time a token is generated;
[0085] Additionally, tokens are generated based on the key and value matrix of cached tokens.
[0086] This concludes the process. Figure 5 Structural description of the device shown.
[0087] This application also provides embodiments that... Figure 5 The hardware structure of the device shown. See also Figure 6 , Figure 6 This is a structural diagram of an electronic device provided in an embodiment of this application. Figure 6 As shown, the hardware structure may include: a processor and a machine-readable storage medium, the machine-readable storage medium storing machine-executable instructions that can be executed by the processor; the processor is used to execute the machine-executable instructions to implement the method disclosed in the above example of this application.
[0088] Based on the same application concept as the above method, this application embodiment also provides a machine-readable storage medium storing a plurality of computer instructions, which, when executed by a processor, can implement the method disclosed in the above examples of this application.
[0089] For example, the aforementioned machine-readable storage medium can be any electronic, magnetic, optical, or other physical storage device that can contain or store information such as executable instructions, data, etc. For instance, machine-readable storage media can be: RAM (Random Access Memory), volatile memory, non-volatile memory, flash memory, storage drives (such as hard disk drives), solid-state drives, any type of storage disk (such as optical discs, DVDs, etc.), or similar storage media, or combinations thereof.
[0090] The above description is merely an embodiment of this application and is not intended to limit this application. Various modifications and variations can be made to this application by those skilled in the art. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principle of this application should be included within the scope of the claims of this application.
Claims
1. A method for accelerating speculative sampling inference for text generation, characterized in that, This method is applied to speculative sampling scenarios, and includes: A candidate path sequence generated by a small model based on user input is obtained; the candidate path sequence represents possible words after the last word in the user input; each candidate path in the candidate path sequence is a path that meets the path comprehensive potential value requirement; the path comprehensive potential value requirement is used to ensure that the tokens in the candidate path sequence predicted by the small model are accepted by the large model, and the path comprehensive potential value of any candidate path is determined based at least on the prediction probability of each token predicted by the small model in the candidate path, the confidence of the candidate path, and the length of the candidate path; The reward parameters for each candidate path in the candidate path sequence are determined. The reward parameters for any candidate path are determined based on its base reward value, diversity reward value, and exploration reward value. Specifically, the base reward value for any candidate path is determined based on the acceptance of the candidate path by the large model, the diversity reward value for any candidate path is determined based on the difference between the candidate path and existing historical paths, the exploration reward value for any candidate path is determined based on the frequency of the candidate path's occurrence in existing historical paths and the potential value of the candidate path, and the potential value of any candidate path is determined based on the predicted probabilities of each token in the candidate path output by the small model. The small model is updated based on each candidate path in the candidate path sequence and the reward parameters of each candidate path; if the small model does not meet the optimal conditions, when user input is received again, the obtained small model is returned to generate a candidate path sequence based on user input until the small model meets the optimal conditions.
2. The method according to claim 1, characterized in that, The method further includes: The small model is initialized using a knowledge distillation method, so that the small model learns the knowledge of the large model and its performance approaches that of the large model.
3. The method according to claim 1, characterized in that, The overall potential value of any candidate path is positively correlated with the predicted probability of each token predicted by the small model, the confidence level of the candidate path, and the inverse of the length of the candidate path.
4. The method according to claim 1, characterized in that, The diversity reward value of any candidate path is negatively correlated with the frequency of its occurrence in existing historical paths. Specifically, the less frequently a candidate path appears in existing historical paths, the higher its diversity reward value.
5. The method according to claim 1, characterized in that, The exploration reward value of any candidate path is negatively correlated with the frequency of its occurrence in existing historical paths and positively correlated with its potential value.
6. The method according to claim 1 or 5, characterized in that, The potential value of any candidate path is negatively correlated with information entropy, which is determined based on the predicted probabilities of each token in the candidate path as predicted by the small model.
7. The method according to claim 1, characterized in that, The small model determines the candidate path sequence based on the Monte Carlo Tree Search (MCTS) algorithm; when generating any candidate path in the candidate path sequence based on the MCTS algorithm, the small model also generates at least one reference path that does not meet the requirements of the comprehensive potential value of the path. The overall path potential value of the reference path is lower than that of the candidate path. The method further includes: determining the reward parameters of the reference path for each candidate path in the same manner as determining the reward parameters of each candidate path in the candidate path sequence; The step of updating the small model based on each candidate path in the candidate path sequence and the reward parameters of each candidate path includes: The small model is updated based on each path in the path pair and the reward parameters of each path; the path pair includes a candidate path and a reference path for the candidate path.
8. The method according to claim 1, characterized in that, The small model is based on the KV Cache mechanism, which caches the key and value of the token each time a token is generated; Additionally, tokens are generated based on the key and value of the cached token.
9. A speculative sampling inference acceleration device for text generation, characterized in that, This device is used in speculative sampling scenarios and includes: A path unit is used to obtain a candidate path sequence generated by a small model based on user input; the candidate path sequence represents possible words after the last word in the user input; each candidate path in the candidate path sequence is a path that meets the path comprehensive potential value requirement; the path comprehensive potential value requirement is used to ensure that the tokens in the candidate path sequence predicted by the small model are accepted by the large model; the path comprehensive potential value of any candidate path is determined based at least on the prediction probability of each token predicted by the small model in the candidate path, the confidence of the candidate path, and the length of the candidate path; A determining unit is used to determine the reward parameters of each candidate path in the candidate path sequence; the reward parameters of any candidate path are determined based on the basic reward value, diversity reward value, and exploration reward value of the candidate path; wherein, the basic reward value of any candidate path is determined based on the acceptance of the candidate path by the large model, the diversity reward value of any candidate path is determined based on the difference between the candidate path and existing historical paths, the exploration reward value of any candidate path is determined based on the frequency of the candidate path appearing in existing historical paths and the potential value of the candidate path, and the potential value of any candidate path is determined based on the predicted probability of each token in the candidate path output by the small model; The processing unit is used to update the small model based on each candidate path in the candidate path sequence and the reward parameters of each candidate path; if the small model does not meet the optimal conditions, when user input is received again, the obtained small model is used to generate a candidate path sequence based on user input until the small model meets the optimal conditions.
10. The apparatus according to claim 9, characterized in that, The processing unit further includes: Initially, the small model is initialized based on knowledge distillation so that it learns the knowledge of the large model and its performance approaches that of the large model.
11. The apparatus according to claim 9, characterized in that, The overall potential value of any candidate path is positively correlated with the predicted probability of each token predicted by the small model, the confidence level of the candidate path, and the inverse of the length of the candidate path.
12. The apparatus according to claim 9, characterized in that, The diversity reward value of any candidate path is negatively correlated with the frequency of its occurrence in existing historical paths. Specifically, the less frequently a candidate path appears in existing historical paths, the higher its diversity reward value. And / or, the exploration reward value of any candidate path is negatively correlated with the frequency of the candidate path appearing in existing historical paths and positively correlated with the potential value of the candidate path; And / or, the potential value of any candidate path is negatively correlated with the entropy determined based on the predicted probability of each token predicted by the small model in that candidate path.
13. The apparatus according to claim 9, characterized in that, The small model determines the candidate path sequence based on the Monte Carlo Tree Search (MCTS) algorithm; when generating any candidate path in the candidate path sequence based on the MCTS algorithm, the small model also generates at least one reference path that does not meet the path comprehensive potential value requirement; the path comprehensive potential value of the reference path is lower than the path comprehensive potential value of the candidate path. The determining unit is further configured to: determine the reward parameters of the reference path for each candidate path in accordance with the method for determining the reward parameters of each candidate path in the candidate path sequence; The processing unit updates the small model based on each candidate path in the candidate path sequence and the reward parameters of each candidate path, including: The small model is updated based on each path in the path pair and the reward parameters of each path; the path pair includes a candidate path and a reference path for the candidate path.
14. The apparatus according to claim 9, characterized in that, The small model is based on the KV Cache mechanism, which caches the key and value matrix of the token each time a token is generated; Additionally, tokens are generated based on the key and value matrix of cached tokens.
15. An electronic device, characterized in that, The electronic device includes: a processor and a machine-readable storage medium; The machine-readable storage medium stores machine-executable instructions that can be executed by the processor; The processor is configured to execute machine-executable instructions to implement the method steps of any one of claims 1-8.
Citation Information
Patent Citations
Large language model reasoning acceleration method and related device
CN118333172A
Optimized training, data processing methods, devices, equipment and media for large models
CN119761417A