Pre-training large language model reasoning optimization method based on fine-grained early exit and dynamic word list cutting
By adding a confidence estimation module to each layer of the pre-trained large language model and performing dynamic vocabulary pruning, the early exit mechanism is optimized, the problem of high computational cost is solved, and efficient and low-cost model inference is achieved.
Patent Information
- Application Number
- CN202510750407.4
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-06-06
- Publication Date
- 2025-09-16
AI Technical Summary
Existing pre-trained large language models have high computational costs in the early exit mechanism, especially when the vocabulary is large, resulting in limited efficiency improvements and redundant parameters affecting model performance.
By adding a confidence estimation module after the attention sublayer and FFN sublayer of each layer, and utilizing a fine-grained early exit mechanism and dynamic vocabulary pruning, the model inference process is optimized, reducing the amount of computation and redundant parameters.
Without affecting model performance, it significantly improves model inference efficiency, reduces computing costs and memory overhead, and meets the needs of low-latency and resource-constrained scenarios.
Smart Images
Figure CN120654826A_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the technical field of pre-trained large language models, and more specifically relates to a pre-trained large language model reasoning optimization based on fine-grained early exit and dynamic vocabulary pruning. Background Art
[0002] As large language models (LLMs) continue to grow in size, their performance has significantly improved, but this has also led to slower inference speeds and increased computational costs. This problem is particularly prominent in scenarios with limited resources or low latency requirements. Early-exiting is a mechanism that accelerates inference by allowing the model to make predictions in intermediate layers. However, the vocabulary size of modern LLMs is typically very large (approximately 30-256K), resulting in the need to calculate the probability distribution of the entire vocabulary when performing confidence estimation in the early-exit mechanism. This incurs significant computational overhead, weakening the efficiency gains brought by the early-exit mechanism.
[0003] Existing early exit mechanisms typically rely on softmax confidence estimation, which calculates the predicted distribution by mapping the current hidden state to the full vocabulary. However, this approach is computationally expensive, especially when the vocabulary is large, resulting in limited improvement in the actual efficiency of early exit mechanisms. Furthermore, some research has shown that redundant parameters in pre-trained large language models have different importance on model performance at the attention sublayer and the FFN sublayer.
[0004] Therefore, how to perform confidence estimation on the fine-grained redundant parameters of pre-trained large language models and optimize the vocabulary mapping process in the confidence estimation module; providing users of pre-trained large language models with more efficient, fast, and high-precision model inference solutions are issues that researchers in this field urgently need to solve. Summary of the Invention
[0005] In view of this, the present invention provides a pre-trained large language model inference optimization method based on fine-grained early exit and dynamic vocabulary pruning. This method enables users of large language models to achieve fast inference of large language models with low memory overhead and accuracy loss without the need for additional training.
[0006] In order to achieve the above object, the present invention adopts the following technical solutions:
[0007] A method for optimizing inference of pre-trained large language models based on fine-grained early exit and dynamic vocabulary pruning, including:
[0008] S1. Obtain a pre-trained large language model and a calibration dataset, and perform fine-grained settings on the pre-trained large language model. That is, add a confidence estimation module after the attention sublayer and FFN sublayer in each layer of the large language model, and use the sublayer feature embedding as input. At the same time, set the exit mechanism of the confidence estimation module.
[0009] S2. Input the calibration dataset into the large language model after fine-grained setting for inference, and record the unnormalized score L of each attention layer and FFN layer l , the unnormalized score L l It is generated by the feature embedding of the attention sub-layer or FFN sub-layer through the confidence estimation module;
[0010] S3, based on the prediction results of the pre-trained model, unnormalized score L l And the preset cropping scale K, calculate the cropping starting layer p;
[0011] S4. Update the confidence assessment module using the cropped starting layer p and the cropping scale K to obtain the final pre-training model.
[0012] Optionally, the exit mechanism of the confidence estimation module of the present invention is:
[0013]
[0014] in, is the probability of the most likely word predicted by the current sub-layer, is the probability of the second most likely word;
[0015] when When a defined threshold is exceeded, an early exit is performed.
[0016] Optionally, the specific process of step S3 of the present invention is:
[0017] For each sub-layer l, calculate the unnormalized score ranking r of the final predicted vocabulary of the pre-trained large model in the sub-layer;
[0018] The method for selecting the cropping starting layer p is: select the earliest layer according to the score ranking r, so that in this layer, the probability of the final predicted word appearing in the top K most likely words exceeds a given threshold, and determine this layer as the cropping starting layer p.
[0019] Optionally, p in the present invention is expressed as:
[0020]
[0021] Among them, rank(x final |L l ) represents the final predicted vocabulary x finalThe ranking at level l is score ranking r.
[0022] Optionally, the cropping scale K described in the present invention is selected to be 512.
[0023] Optionally, the specific process of step S4 of the present invention is: according to the unnormalized score L of layer p l After softmax processing, the probability distribution on the vocabulary is obtained, and the top K words with the largest probability distribution are selected. An indicator vector is generated according to the index of the selected word. The corresponding rows in the LM head parameter matrix are retained according to the indicator vector, and the remaining rows are pruned to obtain the pruned parameter matrix as the new parameter matrix of the confidence assessment module.
[0024] Optionally, the pre-trained large language model supported by the present invention is a large language model with a decoder-only structure, including the Llama series, the GTP series, and the OPT series.
[0025] Optionally, the calibration data set used in step S2 of the present invention includes a data set that meets the usage scenario of the pre-trained model.
[0026] Beneficial effects:
[0027] Compared with the existing technology, the present invention provides a method for optimizing inference of pre-trained large language models based on fine-grained early exit and dynamic vocabulary pruning. It uses an early exit mechanism based on the attention layer and the FFN layer to ensure finer-grained identification of redundant layers in the pre-trained large language model. It also uses the dynamic vocabulary pruning method to dynamically reduce the size of the vocabulary applicable to confidence estimation during the inference process of the large language model. Without affecting the performance of the model, it significantly reduces the amount of computation, thereby improving the inference efficiency of the model. This meets the needs of users of pre-trained large language models for low-cost, low-precision loss, and high-efficiency deployment and inference of pre-trained large language models. BRIEF DESCRIPTION OF THE DRAWINGS
[0028] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the following briefly introduces the drawings required for use in the embodiments or the description of the prior art. Obviously, the drawings described below are merely embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on the provided drawings without paying any creative work.
[0029] Figure 1 This is a schematic diagram of the overall process structure of the method provided by the present invention.
[0030] Figure 2 This is a schematic diagram of the model structure after adding a confidence assessment module to the method provided by the present invention.
[0031] Figure 3 This is a schematic diagram of dynamic vocabulary clipping in the method provided by the present invention. DETAILED DESCRIPTION
[0032] The embodiments of the present invention are described in detail below with reference to the accompanying drawings.
[0033] It should be noted that, in the absence of conflict, the following embodiments and features in the embodiments may be combined with each other; and, based on the embodiments in this disclosure, all other embodiments obtained by persons of ordinary skill in the art without creative work are within the scope of protection of this disclosure.
[0034] It should be noted that various aspects of the embodiments within the scope of the appended claims are described below. It should be apparent that the aspects described herein can be embodied in a wide variety of forms, and any specific structure and / or function described herein is merely illustrative. Based on this disclosure, it should be understood by those skilled in the art that an aspect described herein can be implemented independently of any other aspect, and two or more of these aspects can be combined in various ways. For example, any number of aspects described herein can be used to implement an apparatus and / or practice a method. In addition, other structures and / or functionalities other than one or more of the aspects described herein can be used to implement this apparatus and / or practice this method.
[0035] like Figure 1 As shown, the inference optimization method for a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning in this embodiment includes the following steps:
[0036] S1. Obtain a pre-trained large language model and a calibration dataset, and add a fine-grained confidence estimation module to each layer of the pre-trained large model. Specifically, add a confidence estimation module after the attention sublayer and FFN sublayer in each layer, and set an exit mechanism for the confidence evaluation module. Its structure is as follows: Figure 2 shown.
[0037] The confidence evaluation module (LM head) of this embodiment takes the feature embedding of the sub-layer as input, performs softmax on the output of the LM head, and obtains a probability distribution. It then uses the top-2diff strategy to determine whether to perform early exit. The formula is as follows:
[0038]
[0039] in, is the probability of the most likely word predicted by the current sub-layer, is the probability of the second most likely word;
[0040] Set the confidence evaluation module to calculate the When a defined threshold is exceeded, an early exit is performed.
[0041] S2. Use the calibration dataset for inference and record the unnormalized score L for each attention layer and FFN layer l ; The unnormalized score L l It is generated by the feature embedding of the attention sub-layer or FFN sub-layer through the LM head without going through the softmax operation.
[0042] The calibration dataset used in this step is not limited to a specific dataset, as long as it meets the usage scenario of the pre-trained model.
[0043] S3, based on the prediction results of the pre-trained model, unnormalized score L l And the preset cropping scale K, calculate the cropping starting layer p; the specific process is:
[0044] For each sub-layer l, calculate the unnormalized score ranking r of the final predicted vocabulary of the pre-trained large model in the sub-layer;
[0045] The earliest layer p is selected so that in this layer, the probability of the final predicted word appearing in the top K most likely words exceeds a given threshold, where "earliest" means the layer with the smallest index that satisfies the condition that in this layer, the probability of the final predicted word appearing in the top K most likely words exceeds a given threshold.
[0046] It can be expressed by the following formula:
[0047]
[0048] Among them, rank(x final |L l ) represents the final predicted vocabulary x final In the ranking of the first level, K is the selected cropping scale, and in this embodiment, K is selected as 512.
[0049] S4. Update the confidence assessment module using the cropped starting layer p and the cropping scale K to obtain the final pre-training model.
[0050] The updated confidence assessment module will dynamically trim the vocabulary mapping matrix of the LM head during the inference process, thereby reducing the amount of calculation for confidence estimation. The schematic diagram is as follows Figure 3 As shown, specifically: according to the unnormalized score L of layer p lA softmax process is performed to obtain a probability distribution over the vocabulary. The top K words with the largest probability distributions are selected from these words. Indicator vectors are generated based on the indices of these words. The corresponding rows in the LM head parameter matrix are retained based on these indicator vectors, and the remaining rows are pruned. The resulting pruned parameter matrix of K × d_model serves as the new parameter matrix for the confidence assessment module. This dynamic pruning of the vocabulary reduces the original vocabulary-feature embedding mapping matrix from d_vocab × d_model dimensions to K × d_model dimensions, significantly reducing the computational effort during the confidence estimation process.
[0051] The dynamic vocabulary pruning update of the confidence assessment module in this step is completely post-processing and does not require fine-tuning of the pre-trained model.
[0052] The present invention provides a method for optimizing inference of a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning. It utilizes an early exit mechanism based on attention layers and FFN layers to ensure finer-grained identification of redundant layers in the pre-trained large language model. It also utilizes a dynamic vocabulary pruning method to dynamically reduce the size of the vocabulary applicable to confidence estimation during the inference process of the large language model, significantly reducing the amount of computation without affecting model performance. This method meets the low-cost, high-efficiency deployment and inference requirements of users of pre-trained large language models.
[0053] The various embodiments in this specification are described in a progressive manner, with each embodiment focusing on the differences from other embodiments. Reference can be made to the common and similar parts between the various embodiments. For the devices disclosed in the embodiments, since they correspond to the methods disclosed in the embodiments, the description is relatively simple, and the relevant parts can be referred to the method description.
[0054] The above description of the disclosed embodiments is intended to enable one skilled in the art to implement or use the present invention. Various modifications to these embodiments will be readily apparent to one skilled in the art, and the general principles defined herein may be implemented in other embodiments without departing from the spirit or scope of the present invention. Therefore, the present invention is not limited to the embodiments shown herein but is intended to conform to the widest scope consistent with the principles and novel features disclosed herein.
Claims
1. A method for optimizing inference of a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning, characterized in that: include: S1. Obtain a pre-trained large language model and a calibration dataset, and perform fine-grained settings on the pre-trained large language model. That is, add a confidence estimation module after the attention sublayer and FFN sublayer in each layer of the large language model, and use the sublayer feature embedding as input. At the same time, set the exit mechanism of the confidence estimation module. S2. Input the calibration dataset into the large language model after fine-grained setting for inference, and record the unnormalized score L of each attention layer and FFN layer l , the unnormalized score L l It is generated by the feature embedding of the attention sub-layer or FFN sub-layer through the confidence estimation module; S3, based on the prediction results of the pre-trained model, unnormalized score L l And the preset cropping scale K, calculate the cropping starting layer p; S4. Update the confidence assessment module using the cropped starting layer p and the cropping scale K to obtain the final pre-training model.
2. The method for optimizing inference of a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning according to claim 1 is characterized in that: The exit mechanism of the confidence estimation module is: in, is the probability of the most likely word predicted by the current sub-layer, is the probability of the second most likely word; when When a defined threshold is exceeded, an early exit is performed.
3. The method for optimizing inference of a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning according to claim 1 is characterized in that: The specific process of step S3 is: For each sub-layer l, calculate the unnormalized score ranking r of the final predicted vocabulary of the pre-trained large model in the sub-layer; The method for selecting the cropping starting layer p is: select the earliest layer according to the score ranking r, so that in this layer, the probability of the final predicted word appearing in the top K most likely words exceeds a given threshold, and determine this layer as the cropping starting layer p.
4. The method for optimizing inference of a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning according to claim 3 is characterized in that: The p is expressed as: Among them, rank(x final |L l ) represents the final predicted vocabulary x final The ranking at level l is score ranking r.
5. The method for optimizing inference of a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning according to claim 4 is characterized in that: The cropping scale K is selected as 512.
6. The method for optimizing inference of a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning according to claim 1, characterized in that: The specific process of step S4 is: according to the unnormalized score L of layer p l After softmax processing, the probability distribution on the vocabulary is obtained, and the top K words with the largest probability distribution are selected. An indicator vector is generated according to the index of the selected word. The corresponding rows in the LM head parameter matrix are retained according to the indicator vector, and the remaining rows are pruned to obtain the pruned parameter matrix as the new parameter matrix of the confidence assessment module.
7. The method for optimizing inference of a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning according to claim 1, characterized in that: The supported pre-trained large language models are decoder-only large language models, including the Llama series, GTP series, and OPT series.
8. The method for optimizing inference of a pre-trained large language model based on fine-grained early exit and dynamic vocabulary pruning according to claim 1, characterized in that: The calibration data set used in step S2 includes a data set that meets the usage scenario of the pre-trained model.