A method and system for large model inference optimization based on lightweight gating mechanism

Abstract

The application discloses a large model inference optimization method and system based on a lightweight gating mechanism, aiming to improve the efficiency of large-scale pre-training models in inference tasks and significantly reduce the consumption of computing resources. The method first initializes the pre-training model, introduces a dynamic routing mechanism during the inference process, and combines a lightweight gating and threshold judgment mechanism to effectively filter the key layer output in the inference process. By dynamically evaluating the output of each layer, the gating LSTM unit generates a gating output value, and the importance score of the generated layer is activated through the sigmoid function. If the score is lower than the preset threshold, the layer calculation is skipped to reduce unnecessary computational overhead. This optimization strategy can adaptively judge the importance of each layer output, reduce redundant calculation, and improve inference speed under the premise of ensuring the quality of model output. The application is suitable for various task scenarios, including natural language processing, image generation, etc., and has wide application value.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence, and in particular to a method and system for optimizing large-scale model inference based on a lightweight gating mechanism, which aims to improve the inference speed of large-scale pre-trained language models and reduce the consumption of computing resources. Background Technology

[0002] Current large-scale language models (such as GPT and BERT) demonstrate outstanding performance on many tasks, but due to their numerous layers, the computational resource requirements for model inference are extremely high, resulting in low inference efficiency in practical applications and affecting response speed. Existing optimization techniques for large-scale model inference, such as pruning and quantization, can partially improve inference speed, but still suffer from high computational cost, long inference time, and performance degradation. Therefore, optimizing inference speed without significantly affecting model generation quality has become a current research hotspot. Summary of the Invention

[0003] The purpose of this invention is to address the shortcomings of existing technologies by providing a method and system for optimizing large model inference based on a lightweight gating mechanism. It is particularly suitable for improving inference efficiency and speed by utilizing dynamic routing mechanisms during the inference process, and optimizing the large model inference process by selectively skipping certain network layers. This effectively reduces computational load, improves inference speed and resource utilization, while maintaining model performance and accuracy.

[0004] To achieve the above objectives, this invention provides a large model inference optimization method based on a lightweight gating mechanism, comprising the following steps:

[0005] (1) Use a large dataset to pre-train a large model to generate a pre-trained model suitable for inference tasks; input the data into the pre-trained model;

[0006] (2) Use a lightweight gating mechanism to quantitatively evaluate the importance score of each layer output of the pre-trained model to obtain the confidence or importance score of each layer output.

[0007] (3) Based on the confidence or importance score of each layer's output, dynamically decide whether to skip the calculation of that layer;

[0008] (4) The output of each layer is evaluated by gating LSTM unit to select key output features; if the output importance score is lower than the set threshold, the calculation of that layer and subsequent parts is skipped.

[0009] (5) Perform reasoning tasks based on the selected key output features to improve reasoning speed.

[0010] Furthermore, the pre-trained model is implemented through a causal language model with a Transformer structure, including BERT, GPT, and QWen model structures.

[0011] Further, step (3) includes: defining specific layers processed by the dynamic routing mechanism, and selecting the layers to be dynamically routed through the selected_layers parameter; the dynamic routing mechanism allows users or developers to select specific processing layers at different levels of the model; through the dynamic routing mechanism, the model can selectively skip certain layers or retain specific layers to ensure the accuracy of the output during the inference process according to the needs of the task.

[0012] Furthermore, in step (4), the layer output undergoes a linear transformation to adjust its dimensions to suit the input requirements of the LSTM; specifically including:

[0013] A. Use linear transformation to adjust the dimension of each layer's output to match the input size of the LSTM;

[0014] B. Linear transformation is implemented in the form of matrix multiplication, that is, the layer output is transformed through a weight matrix; the parameters of the linear layer are automatically updated during the training process to ensure its adaptability to the input.

[0015] Furthermore, in step (4), the gating mechanism adopts a lightweight neural network structure and uses a lightweight gated LSTM to determine whether to skip the calculation of the current layer. This LSTM generates a gate value based on the input features, which is between 0 and 1 and is processed by the torch.sigmoid() function. If the gate value is less than the set threshold, the calculation of the layer is skipped and None is returned; otherwise, the calculation of the layer continues.

[0016] Furthermore, during the layer-by-layer forward propagation process, the output of each layer is processed by the dynamic_routing function; if the output of a certain layer is judged to be of low importance, that is, the gate value is less than the set threshold, the calculation of that layer will be skipped and the output of the previous layer will continue to be used; the threshold is dynamically adjusted according to the complexity of the inference task.

[0017] Furthermore, when the reasoning task is a generative text task, the diversity and fluency of the generated text are evaluated based on the screening results of each layer.

[0018] To achieve the above objectives, the present invention also provides a large model inference optimization system based on a lightweight gating mechanism, the system comprising:

[0019] Pre-training module: Used to pre-train large models using large datasets;

[0020] Importance assessment module: used to embed a lightweight gating mechanism to quantitatively assess the importance of each layer's output;

[0021] Dynamic skip module: Used to dynamically skip the calculation of certain layers based on the confidence or importance score of the output;

[0022] Gating module: Evaluates the output of each layer through a lightweight gating mechanism to determine whether to skip the calculation of that layer;

[0023] Inference Execution Module: Executes inference tasks based on the selected key output information.

[0024] To achieve the above objectives, the present invention also provides an electronic device, including a memory and a processor, wherein the memory is coupled to the processor; wherein the memory is used to store program data, and the processor is used to execute the program data to implement the above-described large model inference optimization method based on a lightweight gating mechanism.

[0025] To achieve the above objectives, the present invention also provides a computer-readable storage medium having a computer program stored thereon, wherein the program, when executed by a processor, implements the above-described large model inference optimization method based on a lightweight gating mechanism.

[0026] The beneficial effects of this invention are: by introducing an information bottleneck mechanism and a selective activation layer design, this invention can significantly reduce unnecessary computations during large model inference and improve inference speed. Experimental results show that this scheme can significantly reduce inference time while maintaining the quality of generated text, especially when dealing with complex tasks, where inference efficiency is significantly improved. Attached Figure Description

[0027] To more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings used in the description of the embodiments will be 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.

[0028] Figure 1 This is a schematic diagram of a large model inference optimization process based on a lightweight gating mechanism, provided in an embodiment of the present invention.

[0029] Figure 2 This is a structural block diagram of a lightweight gating mechanism provided in an embodiment of the present invention;

[0030] Figure 3 This is a schematic diagram of an electronic device provided according to an embodiment of the present invention. Detailed Implementation

[0031] Exemplary embodiments will now be described in detail, examples of which are illustrated in the accompanying drawings. When the following description relates to the drawings, unless otherwise indicated, the same numbers in different drawings denote the same or similar elements. The embodiments described in the following exemplary embodiments do not represent all embodiments consistent with the present invention. Rather, they are merely examples of systems and methods consistent with some aspects of the invention as detailed in the appended claims.

[0032] The terminology used herein is for the purpose of describing particular embodiments only and is not intended to be limiting of the invention. The singular forms “a,” “the,” and “the” used in this invention and the appended claims are also intended to include the plural forms unless the context clearly indicates otherwise. It should also be understood that the term “and / or” as used herein refers to and includes any or all possible combinations of one or more of the associated listed items.

[0033] It should be understood that although the terms first, second, third, etc., may be used in this invention to describe various information, this information should not be limited to these terms. These terms are only used to distinguish information of the same type from one another. For example, first information may also be referred to as second information without departing from the scope of this invention, and similarly, second information may also be referred to as first information. Depending on the context, the word "if" as used herein may be interpreted as "when," "when," or "in response to a determination."

[0034] The present invention will now be described in detail with reference to the accompanying drawings. Unless otherwise specified, the features of the following embodiments and implementations can be combined with each other.

[0035] Based on the information bottleneck theory and dynamic routing mechanism, this invention proposes the following specific implementation methods:

[0036] Model Structure: An LSTM (Long Short-Term Memory) gating mechanism is introduced. The output of each layer is compressed into features through a linear layer and then fed into the LSTM to obtain a layer importance score. This score determines whether the layer is activated. If the score is below a threshold, the layer is skipped.

[0037] Accelerated Inference: During inference, the model decides whether to skip certain layers based on the information content and importance score of each layer. This reduces unnecessary computation while preserving the information flow that has a significant impact on the final result, thereby effectively improving inference speed.

[0038] Evaluation metrics: Generation quality is assessed through perplexity and diversity of generated text, ensuring that inference acceleration does not negatively impact the overall quality of text generation. Perplexity measures the quality of generated text, while repetition rate assesses text diversity.

[0039] See Figure 1 This invention provides a large model inference optimization method based on a lightweight gating mechanism, comprising the following steps:

[0040] (1) Pre-trained large model: A large model is pre-trained using a large dataset to generate a pre-trained model suitable for inference tasks. Input data passes through each layer of the pre-trained model sequentially, and each layer generates intermediate features, which are then passed to the next layer. The input data can be datasets used to train large language models, such as Pile, BookCorpus, PubMed, and CommonCrawl, or text extracted from HTML web pages. For example, when building a dialogue system (such as a chatbot), mixed data containing social media conversations, news reports, etc., can be used to enable the model to understand language of different styles and contexts. In the pre-trained large model, each layer typically refers to a layer in the Transformer architecture; these layers include the Self-Attention Layer and the Feed-Forward Layer.

[0041] Specifically, the pre-trained large model is implemented through a causal language model with a Transformer structure, including but not limited to model structures such as BERT, GPT, and QWen.

[0042] (2) Dynamic importance assessment: Based on the pre-trained model, a lightweight gating mechanism is used to quantitatively assess the importance score of each layer output, so as to obtain the confidence or importance score of each layer output.

[0043] Specifically, a lightweight LSTM gate unit (gate_lstm) is used to process the input vector. The gate_lstm generates a hidden state, which represents the importance of that layer in the current context. The generated hidden state is then activated by a sigmoid function, outputting a value in the range (0,1). This value can be considered the importance score of the current layer, i.e., the probability of "should this layer be skipped". The set gating parameter serves as an "importance threshold". When the importance score of a layer is less than this threshold, the layer is considered unimportant, its output is ignored, and the computation of the current layer is skipped; otherwise, the layer output is retained and passed to the next layer.

[0044] (3) Dynamic skipping of layers: Based on the confidence or importance score of each layer's output, the system dynamically decides whether to skip the calculation of that layer to reduce unnecessary computation. The pre-selected feature layers are filtered by the dynamic routing module to determine whether the layer should be skipped.

[0045] Specifically, the dynamic routing mechanism defines specific layers to be processed, and selects the layers to be dynamically routed using the `selected_layers` parameter. This mechanism allows users or developers to select specific processing layers at different levels of the model. Through dynamic routing, the model can selectively skip certain layers or retain specific layers during inference to ensure output accuracy, depending on task requirements. For example, for natural language processing tasks, layers 6, 12, 18, and 24 can be set as dynamic processing layers to reduce computational burden while retaining necessary contextual processing capabilities. Using the `selected_layers` parameter, developers can customize the model's computational path based on the complexity of the inference task or actual needs, further improving flexibility and efficiency during inference.

[0046] (4) Gating mechanism: The output of each layer is evaluated by gating LSTM units to select key output features; if the output importance score is lower than the set threshold, the calculation of that layer and subsequent parts is skipped. Information is propagated layer by layer until the output of the last layer.

[0047] The layer's output undergoes a linear transformation (self.linear) to adjust its dimensions to fit the input requirements of the LSTM. Specific operations include:

[0048] A. The output of each layer of the model usually has high-dimensional features. Directly inputting these features into the LSTM may lead to low computational efficiency due to dimensionality mismatch. Therefore, a linear transformation is needed to adjust the dimension of each layer's output to match the LSTM input size.

[0049] B. Linear transformations are usually implemented in the form of matrix multiplication, that is, the layer output is transformed through a weight matrix; the parameters of the linear layer can be automatically updated during the training process to ensure its adaptability to the input.

[0050] That is, before entering the LSTM, the input undergoes a linear transformation (self.linear) to adjust the dimensions to fit the input requirements of the LSTM.

[0051] Furthermore, the gating mechanism employs a lightweight neural network structure, using a lightweight gated LSTM (gate_lstm) to determine whether to skip the computation of the current layer. This LSTM generates a gate value (gate_output) based on the input features; its value is between 0 and 1 and is processed by the torch.sigmoid() function. If the gate value is less than the set threshold, the computation of that layer is skipped, and None is returned; otherwise, the computation of the layer continues.

[0052] During the layer-by-layer forward propagation, the output of each layer can be processed by the dynamic_routing function. If the output of a layer is determined to be of "low importance" (i.e., gate_output is less than the threshold), the calculation of that layer will be skipped, and the output of the previous layer will continue to be used.

[0053] This method determines whether to skip the calculation of subsequent layers based on the comparison result between the calculated output and a preset threshold, which can be dynamically adjusted according to the complexity of the reasoning task.

[0054] (5) Inference task execution: The information after layer filtering and dynamic routing is passed to the output layer. The inference task is executed according to the key output features selected, and the inference result of the task is generated, thereby improving the inference speed and reducing the computational overhead.

[0055] Furthermore, reasoning tasks can include generative text tasks, which evaluate the diversity and fluency of the generated text based on the selection results at each layer.

[0056] See Figure 2 This invention provides a large model inference optimization system based on a lightweight gating mechanism, the system comprising:

[0057] Pre-training module: Used to pre-train large models using large datasets;

[0058] Importance assessment module: used to embed a lightweight gating mechanism to quantitatively assess the importance of each layer's output;

[0059] Dynamic skip module: Used to dynamically skip the calculation of certain layers based on the confidence or importance score of the output;

[0060] Gating module: Evaluates the output of each layer through a lightweight gating mechanism to determine whether to skip the calculation of that layer;

[0061] Inference Execution Module: Executes subsequent inference tasks based on the selected important output information to improve inference speed and reduce computational burden.

[0062] The system reduces unnecessary computation and maintains the accuracy and diversity of inference results by quantifying the importance of the output information at each layer.

[0063] Corresponding to the aforementioned embodiments of the large model inference optimization method based on lightweight gating mechanisms, this application also provides an electronic device, including: one or more processors; a memory for storing one or more programs; when the one or more programs are executed by the one or more processors, the one or more processors implement the large model inference optimization method based on lightweight gating mechanisms as described above. Figure 3The diagram shown is a hardware structure diagram of any device with data processing capabilities on which the large model inference optimization method based on a lightweight gating mechanism provided in this application embodiment can be implemented. (Except for...) Figure 3 In addition to the processor, memory, DMA controller, disk, and non-volatile memory shown, any data processing device in the embodiment may also include other hardware depending on the actual function of the data processing device, which will not be described in detail here.

[0064] Corresponding to the aforementioned embodiments of the large model inference optimization method based on lightweight gating mechanism, this embodiment of the invention also provides a computer-readable storage medium storing a program thereon, which, when executed by a processor, implements the large model inference optimization method based on lightweight gating mechanism described in the above embodiments.

[0065] The computer-readable storage medium can be an internal storage unit of any data processing device described in any of the foregoing embodiments, such as a hard disk or memory. The computer-readable storage medium can also be any data processing device, such as a plug-in hard disk, smart media card (SMC), SD card, flash card, etc., equipped on the device. Furthermore, the computer-readable storage medium can include both internal storage units of any data processing device and external storage devices. The computer-readable storage medium is used to store the computer program and other programs and data required by the data processing device, and can also be used to temporarily store data that has been output or will be output.

[0066] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

[0067] The above embodiments are only used to illustrate the design concept and features of the present invention, and their purpose is to enable those skilled in the art to understand the content of the present invention and implement it accordingly. The protection scope of the present invention is not limited to the above embodiments. Therefore, all equivalent changes or modifications made based on the principles and design ideas disclosed in the present invention are within the protection scope of the present invention.

Claims

1. A large model inference optimization method based on a lightweight gating mechanism, characterized in that, Includes the following steps: (1) Use a large dataset to pre-train a large model to generate a pre-trained model suitable for inference tasks; input the data into the pre-trained model; (2) Use a lightweight gating mechanism to quantitatively evaluate the importance score of each layer output of the pre-trained model to obtain the confidence or importance score of each layer output. (3) Based on the confidence or importance score of each layer's output, dynamically decide whether to skip the calculation of that layer; (4) By using a gated LSTM unit, the output of each layer is gated and evaluated to select key output features; If the output importance score is lower than a set threshold, the calculation of that layer and subsequent parts is skipped. The gating mechanism adopts a lightweight neural network structure, using a lightweight gated LSTM to determine whether to skip the calculation of the current layer. This LSTM generates a gate value based on the input features, which is between 0 and 1 and is processed by the torch.sigmoid() function. If the gate value is less than the set threshold, the calculation of that layer is skipped and None is returned; otherwise, the calculation of the layer continues. (5) Perform reasoning tasks based on the selected key output features to improve reasoning speed.

2. The large model inference optimization method based on lightweight gating mechanism as described in claim 1, characterized in that, The pre-trained model is implemented through a causal language model with a Transformer structure, including BERT, GPT, and QWen model structures.

3. The large model inference optimization method based on lightweight gating mechanism as described in claim 1, characterized in that, Step (3) includes: defining specific layers to be processed by the dynamic routing mechanism, and selecting the layers to be dynamically routed through the selected_layers parameter; the dynamic routing mechanism allows users or developers to select specific processing layers at different levels of the model; through the dynamic routing mechanism, the model can selectively skip certain layers or retain specific layers to ensure the accuracy of the output during the inference process according to the needs of the task.

4. The large model inference optimization method based on lightweight gating mechanism as described in claim 1, characterized in that, In step (4), the layer output undergoes a linear transformation to adjust its dimensions to suit the input requirements of the LSTM; specifically, this includes: A. Use linear transformation to adjust the dimension of each layer's output to match the input size of the LSTM; B. Linear transformation is implemented in the form of matrix multiplication, that is, the layer output is transformed through a weight matrix; the parameters of the linear layer are automatically updated during the training process to ensure its adaptability to the input.

5. The large model inference optimization method based on lightweight gating mechanism as described in claim 1, characterized in that, During the layer-by-layer forward propagation process, the output of each layer is processed by the dynamic_routing function; If the output of a certain layer is judged to be of low importance, that is, the gating value is less than the set threshold, the calculation of that layer will be skipped and the output of the previous layer will continue to be used; the threshold is dynamically adjusted according to the complexity of the inference task.

6. The large model inference optimization method based on lightweight gating mechanism as described in claim 1, characterized in that, When the reasoning task is a generative text task, the diversity and fluency of the generated text are evaluated based on the screening results of each layer.

7. A large model inference optimization system based on a lightweight gating mechanism that implements the method described in claim 1, characterized in that, The system includes: Pre-training module: Used to pre-train large models using large datasets; Importance assessment module: used to embed a lightweight gating mechanism to quantitatively assess the importance of each layer's output; Dynamic skip module: Used to dynamically skip the calculation of certain layers based on the confidence or importance score of the output; Gating module: Evaluates the output of each layer through a lightweight gating mechanism to determine whether to skip the calculation of that layer; Inference Execution Module: Executes inference tasks based on the selected key output information.

8. An electronic device comprising a memory and a processor, characterized in that, The memory is coupled to the processor; wherein the memory is used to store program data, and the processor is used to execute the program data to implement the large model inference optimization method based on a lightweight gating mechanism as described in any one of claims 1-6.

9. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the program is executed by the processor, it implements the large model inference optimization method based on a lightweight gating mechanism as described in any one of claims 1-6.

Images