A downstream task-oriented LLM speculative decoding optimization method and system
By constructing a heterogeneous draft model pool and a task classification mechanism, the accuracy and efficiency issues of speculative decoding methods on general tasks are solved, achieving efficient and accurate multi-task reasoning optimization.
Patent Information
- Application Number
- CN202510098098.7
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2025-01-22
- Publication Date
- 2026-01-30
- Estimated Expiration
- 2045-01-22
AI Technical Summary
In existing technologies, speculative decoding methods struggle to achieve high accuracy on general tasks, and the reasoning process of large-scale language models consumes significant computational resources, resulting in low efficiency.
Multiple aligned datasets for different tasks are constructed to form a heterogeneous draft model pool. The optimal draft model is selected for parallel verification through a task classification mechanism. By combining fine-tuning of a small-scale language model and optimization of task-specific datasets, a high-quality sequence of guessed tokens is generated.
It significantly improves the inference speed and accuracy of large-scale pre-trained language models in multi-task environments, enhances the model's adaptability to diverse tasks, and improves the efficiency and accuracy of speculative decoding.
Smart Images
Figure CN120031128B_ABST
Abstract
Description
Technical Field
[0001] This invention belongs to the field of artificial intelligence large model technology, specifically relating to an LLM speculative decoding optimization method and system for downstream tasks. Background Technology
[0002] Speculative decoding optimization methods employ a small-scale draft model for multiple consecutive autoregressive decoding iterations, followed by parallel validation using a large language model (LLM). This approach is one of the mainstream methods for accelerating LLM inference. The decoding efficiency and guessing accuracy of the draft model are two crucial factors affecting the optimization performance of speculative decoding. Existing work primarily aims to improve the performance of speculative decoding methods on general tasks by refining the draft model's guessing mechanism or the parallel validation mechanism of the LLM. However, draft models have a limited number of parameters, making it difficult to achieve high accuracy on general tasks.
[0003] The scaling law states that the perplexity (PPL) of a language model scales power-lawically with model size, dataset size, and training computation. Currently, mainstream LLMs have hundreds of billions or even trillions of parameters, and depending on the user input, they may need to process long contexts, leading to significant memory and computational resource consumption during inference. For example, when the input prompt is 512 characters long, a GPT-2 model with 700 million parameters takes 10 seconds to generate 512 tokens on an RTX 2080ti. During LLM inference, the number of floating-point operations (FLOPS) or integer operations (OPS) is proportional to the square of the sequence length. Therefore, accelerating LLM inference is of great importance. Summary of the Invention
[0004] This invention aims to address the shortcomings of existing technologies and provides the following solutions:
[0005] An optimization method for LLM speculative decoding oriented towards downstream tasks includes the following steps:
[0006] Construct multiple aligned datasets for different tasks, and use the aligned datasets to build a heterogeneous draft model pool;
[0007] Obtain the prompt text of the downstream task, and classify the input task based on the prompt text using a task classification mechanism to obtain the classification result;
[0008] The optimal draft model is selected from the heterogeneous draft model pool based on the classification results.
[0009] The selected draft model is used to generate a sequence of guessed tokens, and the sequence of guessed tokens is input into the target model for parallel verification to complete the speculative decoding optimization.
[0010] Preferably, the method for constructing the aligned dataset includes:
[0011] Get a task set containing different tasks;
[0012] Extract a set of prompt words from the relevant dataset for each task in the task set;
[0013] The target model is used to infer and generate the prompt words, resulting in a corresponding output token sequence.
[0014] The prompt words and the output token sequence are used as the aligned dataset.
[0015] Preferably, the method for constructing the heterogeneous draft model pool includes:
[0016] For each task in the task set, based on the existing small-scale language model, the small-scale language model is fine-tuned using the alignment dataset to obtain a draft model applicable to all tasks.
[0017] The draft models corresponding to all tasks in the task set are integrated to obtain the heterogeneous draft model pool.
[0018] Preferably, the method for obtaining the classification result includes:
[0019] Obtain the input prompt text and delete the stop words in the prompt text to obtain the processed text;
[0020] A training dataset is constructed for all tasks in the task set. The training dataset consists of the prompt word texts belonging to different tasks. Each prompt word text is labeled with the correct task category. A public language text classification model is trained using the training dataset to obtain a fast classification model.
[0021] The processed text is input into the fast classification model for task classification to obtain the classification result.
[0022] The present invention also provides an LLM speculative decoding optimization system for downstream tasks, wherein the system applies the method described in any of the above-mentioned methods and includes: a model pool construction module, a task classification module, a model selection module, and a decoding optimization module;
[0023] The heterogeneous model pool construction module is used to construct multiple aligned datasets for different tasks, and to construct a heterogeneous draft model pool using the aligned datasets.
[0024] The task classification module is used to obtain the prompt text of downstream tasks, and to classify the input tasks based on the prompt text using a task classification mechanism to obtain the classification result;
[0025] The model selection module is used to select the optimal draft model from the heterogeneous draft model pool based on the classification results;
[0026] The decoding optimization module uses the selected draft model to generate a sequence of guessed tokens, and inputs the sequence of guessed tokens into the target model for parallel verification, thus completing the speculative decoding optimization.
[0027] Preferably, in the heterogeneous model pooling construction module, the method for constructing the aligned dataset includes:
[0028] Get a task set containing different tasks;
[0029] Extract a set of prompt words from the relevant dataset for each task in the task set;
[0030] The target model is used to infer and generate the prompt words, resulting in a corresponding output token sequence.
[0031] The prompt words and the output token sequence are used as the aligned dataset.
[0032] Preferably, in the heterogeneous model pool construction module, the method for constructing the heterogeneous draft model pool includes:
[0033] For each task in the task set, based on the existing small-scale language model, the small-scale language model is fine-tuned using the alignment dataset to obtain a draft model applicable to all tasks.
[0034] The draft models corresponding to all tasks in the task set are integrated to obtain the heterogeneous draft model pool.
[0035] Preferably, the task classification module includes: a text processing unit, a classification model training unit, and a task classification unit;
[0036] The text processing unit is used to acquire the input prompt text and delete the stop words in the prompt text to obtain the processed text;
[0037] The classification model training unit constructs a training dataset for all tasks in the task set. The training dataset consists of prompt word texts belonging to different tasks, and each prompt word text is labeled with the correct task category. The publicly available language text classification model is trained using the training dataset to obtain a fast classification model.
[0038] The task classification unit is used to input the processed text into the fast classification model for task classification and obtain the classification result.
[0039] Compared with the prior art, the beneficial effects of the present invention are as follows:
[0040] This invention combines a task classification model with a draft model pool, achieving efficient inference in a multi-task environment through a speculative decoding mechanism. It not only improves the inference performance of large-scale pre-trained language models across different tasks but also enhances the model's adaptability to diverse tasks through the combination of fine-tuning and task classification. Experimental results show that this method significantly improves inference speed compared to the original speculative decoding method on multiple tasks, demonstrating strong practical value and potential for widespread application. Attached Figure Description
[0041] To more clearly illustrate the technical solution of the present invention, the drawings used in the embodiments are briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.
[0042] Figure 1 This is a schematic diagram of the method flow according to an embodiment of the present invention;
[0043] Figure 2 This is a flowchart of a method according to an embodiment of the present invention. Detailed Implementation
[0044] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.
[0045] To make the above-mentioned objects, features and advantages of the present invention more apparent and understandable, the present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.
[0046] Example 1
[0047] In this embodiment, as Figure 1 , Figure 2 As shown, an LLM speculative decoding optimization method for downstream tasks includes the following steps:
[0048] S1. Construct multiple aligned datasets for different tasks, and use the aligned datasets to build a heterogeneous draft model pool.
[0049] The method for constructing the alignment dataset includes: obtaining a task set containing different tasks; extracting a set of prompt words from the relevant dataset of each task in the task set; using the target model to infer and generate the prompt words to obtain the corresponding output token sequence; and using the prompt words and the output token sequence as the alignment dataset.
[0050] In this embodiment, multiple aligned datasets for different tasks are constructed. The key to this stage is the quality and representativeness of the datasets. Through reasonable data sampling and labeling, it is ensured that the datasets can comprehensively reflect various task contexts, thus providing a sufficient data foundation for model fine-tuning. To improve the accuracy of the guessed tokens generated by the draft model and reduce the performance difference between the draft model and the target model, this paper proposes a task-based prediction alignment method. Specifically, this method fine-tunes task-specific datasets under different tasks, enabling the draft model to generate outputs more consistent with the target model on specific tasks, thereby improving the overall performance and efficiency of speculative decoding. In a multi-task learning environment, different tasks often have different input-output patterns, which may cause uneven performance of the draft model when processing multiple tasks. To solve this problem, the proposed solution focuses on task-specific optimization of the draft model based on the characteristics of the task. In the task set T = {t1, t2, t3, ...}, for each task t... i and its related dataset d j By fine-tuning a draft model using a task-specific dataset, the output generated by the draft model in a specific task is made as close as possible to the output of the target model. This significantly improves the generation quality of the draft model on the target task and effectively reduces the performance degradation caused by inaccurate predictions. In this approach, the first step is to start with the dataset related to task t... i Related dataset d j All prompts are extracted, and the target model is used to generate the prompt sequence. The target model generates a corresponding token sequence for this prompt sequence, thus obtaining a complete alignment dataset. This dataset contains the input of the original dataset and the output generated by the target model, providing data corpus for a draft model suitable for customized tasks.
[0051] The method for constructing a heterogeneous draft model pool includes: for each task in the task set, fine-tuning the existing small-scale language model using an aligned dataset to obtain a draft model applicable to all tasks; and integrating the draft models corresponding to all tasks in the task set to obtain a heterogeneous draft model pool.
[0052] Constructing a Heterogeneous Draft Model Pool Using Aligned Datasets: The construction of the draft model pool is the core of this research. Through fine-tuning multiple pre-trained models, a batch of draft models targeting different fine-grained downstream tasks were formed. These models exhibit varying capabilities and strengths / weaknesses on different tasks. The construction process of the draft model pool involves multi-stage fine-tuning and evaluation. First, the pre-trained models need to be fine-tuned according to different downstream tasks. Fine-tuning strategies can include various forms, the most common being Supervised Fine-Tuning (SFT). In SFT, the model adjusts its parameters through training on a specific task dataset to better fit the data distribution of the target task. This embodiment uses the SFT method to optimize the draft models. Based on the newly generated dataset d... j Fine-tuning the draft model to make it perform well on a given task t i The generated results are closer to the target model. This process can be achieved by minimizing the output difference between the draft model and the target model under the same task, thereby improving the adaptability and accuracy of the draft model in that task. The fine-tuned draft model can generate higher-quality guess tokens in different tasks in a multi-task environment, reducing the performance loss caused by low-quality tokens generated by the draft model, and thus improving inference efficiency. The advantage of the draft model pool is that it can combine the advantages of multiple models, providing multiple options for subsequent task execution.
[0053] S2. Obtain the prompt text of the downstream task, and use the task classification mechanism to classify the input task based on the prompt text to obtain the classification result.
[0054] The method for obtaining the classification results includes: acquiring the input task text and removing stop words from the text to obtain the processed text; acquiring the input prompt text and removing stop words from the prompt text to obtain the processed text; constructing a training dataset for all tasks in the task set, the training dataset consisting of prompt texts belonging to different tasks, and labeling each prompt text with the correct task category; training a public language text classification model using the training dataset to obtain a fast classification model; and inputting the processed text into the fast classification model for task classification to obtain the classification results.
[0055] In this embodiment, for the multi-task scenario, considering the different requirements of different tasks for the draft model, a task classification mechanism is proposed. This mechanism can classify the input task according to its characteristics, so as to divide the tasks into different draft models for speculative decoding. Through this task-driven model selection strategy, it can be ensured that each task can use the most suitable draft model for inference, thereby improving the decoding efficiency and generation quality. FastText is an efficient text classification tool with excellent training and inference performance. In this embodiment, the FastText model is retrained using task features to make it more adaptable to the classification requirements of specific tasks, and a classification model is obtained. In some vertical fields, text data often has specific language structures or distribution characteristics. By adjusting the selection range of n-grams and optimizing training parameters, FastText can capture these characteristics more accurately. In the text classification task of this embodiment, after loading the text, necessary preprocessing is performed to optimize the input quality of the model. Among them, removing stop words is a key step. Stop words are usually words with low semantic contribution but high frequency in the text, such as "de", "le", "he", etc. By removing these words, noise interference can be reduced and the attention of the classification model to core semantic information can be improved. After completing the stop word processing, the text is tokenized, and the tokenized text is input into the classification model for classification prediction, and finally efficient and accurate text classification is achieved.
[0056] S3. Select the optimal draft model from the heterogeneous draft model pool according to the classification result.
[0057] In this embodiment, based on the classification result, the most suitable draft model is quickly matched for the task during inference, so as to achieve more efficient and accurate speculative decoding.
[0058] S4. Use the selected draft model to generate draft tokens, and input the draft tokens into the LLM target model for parallel verification to complete the optimization of speculative decoding.
[0059] Embodiment 2
[0060] In this embodiment, an LLM speculative decoding optimization system for downstream tasks includes: a heterogeneous model pool construction module, a task classification module, a model selection module, and a decoding optimization module.
[0061] The heterogeneous model pool construction module is used to construct multiple aligned datasets for different tasks, and use the aligned datasets to construct a heterogeneous draft model pool.
[0062] In the heterogeneous model pooling module, the method for constructing the alignment dataset includes: obtaining a task set containing different tasks; extracting a set of prompt words from the relevant dataset of each task in the task set; using the target model to infer and generate the prompt words to obtain the corresponding output token sequence; and using the prompt words and the output token sequence as the alignment dataset.
[0063] In the heterogeneous model pool construction module, the method for constructing a heterogeneous draft model pool includes: for each task in the task set, based on the existing small-scale language model, fine-tuning the small-scale language model using an alignment dataset to obtain a draft model applicable to all tasks; and integrating the draft models corresponding to all tasks in the task set to obtain a heterogeneous draft model pool.
[0064] The task classification module is used to obtain the prompt text of downstream tasks, and uses the task classification mechanism to classify the input tasks based on the prompt text to obtain the classification results.
[0065] The task classification module includes a text processing unit, a classification model training unit, and a task classification unit. The text processing unit acquires the input prompt text and removes stop words to obtain the processed text. The classification model training unit constructs a training dataset for all tasks in the task set. The training dataset consists of prompt texts belonging to different tasks, and each prompt text is labeled with the correct task category. The publicly available language text classification model is trained using the training dataset to obtain a fast classification model. The task classification unit inputs the processed text into the fast classification model for task classification and obtains the classification result.
[0066] The model selection module is used to select the optimal draft model from the heterogeneous draft model pool based on the classification results.
[0067] The decoding optimization module uses the selected draft model to generate a sequence of guessed tokens, and inputs the sequence of guessed tokens into the target model for parallel verification, thus completing the speculative decoding optimization.
[0068] The embodiments described above are merely preferred embodiments of the present invention and are not intended to limit the scope of the present invention. Various modifications and improvements made to the technical solutions of the present invention by those skilled in the art without departing from the spirit of the present invention should fall within the protection scope defined by the claims of the present invention.
Claims
1. A method for LLM speculation decoding optimization oriented to downstream tasks, characterized in that, The method comprises the following steps: A plurality of alignment data sets for different tasks are constructed, and a heterogeneous draft model pool is constructed using the alignment data sets; Prompt word text of a downstream task is obtained, and a task classification mechanism is used to classify an input task based on the prompt word text to obtain a classification result; An optimal draft model is selected from the heterogeneous draft model pool according to the classification result; The selected draft model is used to generate a guessed tokens sequence, and the guessed tokens sequence is input into a target model for parallel verification to complete speculative decoding optimization; The method for constructing the alignment data set comprises: A task set containing different tasks is obtained; A set of prompt words is extracted from the relevant data set of each task in the task set; The target model is used to generate an output token sequence corresponding to the prompt words through reasoning; The prompt words and the output tokens sequence are used as the alignment data set; The method for constructing the heterogeneous draft model pool comprises: For each task in the task set, a small-scale language model is fine-tuned using the alignment data set to obtain a draft model suitable for all tasks; The draft models corresponding to all tasks in the task set are integrated to obtain the heterogeneous draft model pool.
2. The method of claim 1, wherein, The method for obtaining the classification result comprises: Input prompt word text is obtained, and stop words in the prompt word text are deleted to obtain processed text; A training data set is constructed for all tasks in the task set, the training data set is composed of prompt word texts belonging to different tasks, and each prompt word text is labeled with a correct task category, a public language text classification model is trained using the training data set to obtain a fast classification model; The processed text is input into the fast classification model for task classification to obtain the classification result.
3. A downstream task oriented LLM speculation decoding optimization system, the system applies the method of any one of claims 1-2, characterized in that, It comprises: A heterogeneous model pool construction module, a task classification module, a model selection module, and a decoding optimization module; The heterogeneous model pool construction module is used to construct a plurality of alignment data sets for different tasks, and a heterogeneous draft model pool is constructed using the alignment data sets; The task classification module is used to obtain prompt word text of a downstream task, and a task classification mechanism is used to classify an input task based on the prompt word text to obtain a classification result; The model selection module is used to select an optimal draft model from the heterogeneous draft model pool according to the classification result; The decoding optimization module uses the selected draft model to generate a guessed tokens sequence, and inputs the guessed tokens sequence into a target model for parallel verification to complete speculative decoding optimization; In the heterogeneous model pool construction module, the method for constructing the alignment data set comprises: A task set containing different tasks is obtained; A set of prompt words is extracted from the relevant data set of each task in the task set; The target model is used to generate an output token sequence corresponding to the prompt words through reasoning; The prompt words and the output tokens sequence are used as the alignment data set; The method for constructing the heterogeneous draft model pool in the heterogeneous model pool construction module comprises the following steps: For each task in the task set, a small-scale language model is fine-tuned based on existing alignment data sets to obtain a draft model suitable for all tasks; Integrate the draft models corresponding to all tasks in the task set to obtain the heterogeneous draft model pool.
4. The downstream task oriented LLM speculation decoding optimization system of claim 3, wherein, The task classification module comprises a text processing unit, a classification model training unit, and a task classification unit. The text processing unit is configured to obtain an input prompt word text and delete stop words in the prompt word text to obtain a processed text. The classification model training unit constructs a training data set for all tasks in the task set, the training data set is composed of prompt word texts belonging to different tasks, and each prompt word text is labeled with a correct task category. A public language text classification model is trained using the training data set to obtain a fast classification model. The task classification unit is configured to input the processed text into the fast classification model for task classification to obtain the classification result.