Techniques for compressing artificial neural networks

Abstract

The disclosure relates to techniques for compressing an artificial neural network. At least one of the various embodiments relates to a computer-implemented method for generating a trained artificial neural network. The method comprises: for each model layer included in a trained model, training one or more student model layers to simulate the model layer; for a first target device included in the plurality of target devices, generating one or more candidate architectures based on a constraint optimization problem and one or more trained student model layers; training one or more candidate architectures on the calibration data set; selecting a first candidate architecture associated with the minimum error amount included in the one or more candidate architectures; and performing a plurality of fine tuning training operations on the first candidate architecture to generate a first trained student model.

Description

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS

[0002] This application claims the benefit of U.S. Provisional Patent Application No. 63 / 554,538, entitled “Universal Model Compression for Transformers,” filed on February 16, 2024. The subject matter of this related application is hereby incorporated by reference. Technical Field

[0003] Embodiments of the present disclosure generally relate to large language model compression and, more particularly, to techniques for compressing artificial neural networks. Background Art

[0004] Large language models (LLMs) are artificial neural networks (ANNs) that have demonstrated remarkable performance on a wide range of natural language processing (NLP) tasks, including text generation and classification. However, as LLMs increase in size and complexity, the computational and memory costs and latency associated with training and deploying them for various end-user applications also increase. These increasing costs and latency can limit the overall effectiveness and practicality of LLMs. Consequently, various techniques have been developed to facilitate the training and deployment of LLMs.

[0005] One approach for improving the overall efficiency of LLMs involves building a set of smaller LLMs, rather than a single large one, where each smaller LLM is tailored to execute on specified hardware and within specified time constraints. Neural Architecture Search (NAS) is a technique commonly used to build sets of smaller LLMs subject to various model size, hardware, and memory constraints. NAS algorithms are categorized based on the three phases used to construct the model architecture: the search space, the search strategy, and the performance estimation strategy. The search space defines the set of architectures that can be used to represent the various smaller LLMs, including the number of layers, the layer types (e.g., multilayer perceptron, convolution, attention, etc.), and the number of parameters per layer (e.g., the number of neurons). The search strategy is used to explore the search space, select a given architecture based on various factors, and build smaller LLMs based on the selected architecture. Performance estimation strategies estimate how well the model architecture found in the search phase performs on new data. A common performance estimation strategy is to train the model found in the search phase on a training dataset and evaluate its performance on a validation dataset.

[0006] One disadvantage of using NAS to develop a set of smaller LLMs is that the NAS algorithm takes a significant amount of time to explore the search space and select architectures for the smaller LLMs, especially if the search space is large. Shrinking the search space can speed up execution, but if the search space is too small, the NAS algorithm is less likely to find optimized architectures for the smaller LLMs. Another disadvantage of NAS is that performance estimation strategies typically require training each smaller LLM from scratch and then evaluating it. Training a large number of LLMs (even smaller ones) from scratch can take a considerable amount of time and consume significant computational resources, which can make NAS impractical for many applications.

[0007] Another approach to improving the overall efficiency of LLMs involves model compression, which is the process of compressing a pre-trained LLM to produce a smaller LLM. Three common model compression techniques are pruning, quantization, and knowledge distillation. Pruning is the process of removing redundant parameters (such as neurons) from an existing model. Redundant parameters are generally considered to be parameters whose removal from the model has the least impact on the output of the model. Pruning can be unstructured, which is the removal of individual parameters from the model regardless of where they are located in the model; or it can be structured, which is the removal of groups of parameters from certain locations in the model. Quantization is the use of a reduced number of bits to represent different weights in the model. Using fewer bits to represent weights reduces the amount of storage resources consumed by the model and can also reduce the computational complexity of operations performed using the model. Knowledge distillation uses a larger pre-trained "teacher" model to train a smaller "student" model, and during the training process, the knowledge of the teacher model is transferred to the student model.

[0008] A drawback of model compression is the trade-off between improved model efficiency and a loss of model accuracy. While the aforementioned model compression techniques can produce smaller models that execute faster, extensive compression can lead to significant inaccuracies due to the removal of a significant amount of information from the model. Furthermore, model compression often requires manual tuning of the model's hyperparameters, a process that can be tedious, time-consuming, and error-prone.

[0009] As previously mentioned, there is a need in the art for more efficient techniques to compress LLMs and other artificial neural networks. Summary of the Invention

[0010] At least one of the various embodiments relates to a computer-implemented method for generating a trained artificial neural network. The method includes: for each model layer included in a trained model, training one or more student model layers to simulate the model layer; for a first target device included in a plurality of target devices, generating one or more candidate architectures based on a constrained optimization problem and the one or more trained student model layers; training the one or more candidate architectures on a calibration dataset; selecting a first candidate architecture included in the one or more candidate architectures that is associated with a minimum amount of error; and performing a plurality of fine-tuning training operations on the first candidate architecture to generate a first trained student model.

[0011] At least one technical advantage of the disclosed technology over the prior art is that the disclosed technology can greatly facilitate the training and deployment of LLMs across multiple different hardware implementations. In this regard, the disclosed technology can be used to generate different trained student LLMs for each different hardware implementation based on a single trained LLM. Therefore, using the disclosed technology, multiple different smaller trained student LLMs can be generated for multiple different hardware implementations without having to train each student LLM from scratch, thereby reducing the time and computing resources required to deploy new trained models. In addition, the student LLMs generated using the disclosed technology can improve the latency and memory usage of the original trained LLM without causing any significant loss of accuracy. The disclosed technology also implements a constrained optimization problem to automatically design different student LLM architectures, thereby eliminating the need for manual hyperparameter fine-tuning. These technical advantages provide one or more technical improvements over prior art methods. BRIEF DESCRIPTION OF THE DRAWINGS

[0012] In order that the above features of the present invention may be understood in detail, a more particular description of the invention, briefly summarized above, may be obtained by reference to embodiments, some of which are illustrated in the accompanying drawings. It should be noted, however, that the drawings illustrate only typical embodiments of the invention and are therefore not to be considered limiting of its scope, as the invention may admit to other equally effective embodiments.

[0013] Figure 1 is a block diagram of a computer-based system configured to implement one or more aspects of various embodiments;

[0014] Figure 2 According to various embodiments Figure 1 A more detailed illustration of the architecture of the trained LLM;

[0015] Figure 3 According to various embodiments Figure 1 Block diagram of the LLM architecture optimizer;

[0016] Figure 4 According to various embodiments Figure 3 A more detailed diagram of the candidate architecture generator;

[0017] Figure 5 According to various embodiments Figure 3 A more detailed illustration of the student architecture selector;

[0018] Figure 6 is a flow chart of method steps for generating candidate architectures according to various embodiments; and

[0019] Figure 7 is a flow chart of method steps for generating a student LLM, according to various embodiments. DETAILED DESCRIPTION

[0020] In the following description, numerous specific details are set forth to provide a more thorough understanding of the present invention. However, it will be apparent to one skilled in the art that the present invention can be practiced without one or more of these specific details.

[0021] System Overview

[0022] Figure 1 A block diagram of a computer-based system 100 configured to implement one or more aspects of various embodiments is shown. As shown, system 100 includes, but is not limited to, a compression server 110, a data store 120, a network 130, and a computing device 140. Compression server 110 includes, but is not limited to, a processor 112 and system memory 114. System memory 114 includes, but is not limited to, an LLM architecture optimizer 116 and trained LLMs 118. Computing device 140 includes, but is not limited to, a processor 142 and memory 144. Memory 144 includes, but is not limited to, applications 145. Data store 120 includes, but is not limited to, student LLMs 122.

[0023] The compression server 110 shown herein is for illustrative purposes only and may be varied and modified without departing from the scope of the present disclosure. For example, the number and type of processors 112, the number of GPUs and / or other processing unit types, the number and type of system memory 114, and / or the number of applications contained in the system memory 114 may be modified as needed. Figure 1 In some embodiments, any combination of processor 112 and system memory 114, and / or GPU may be included in and / or replaced with any type of virtual computing system, distributed computing system, and / or cloud computing environment, such as a public, private, or hybrid cloud system.

[0024] The processor 112 receives user input from an input device (e.g., a keyboard or mouse). The processor 112 can be in the form of any technically feasible processing device that is configured to process data and execute program code. For example, any processor 112 can be a central processing unit (CPU), a graphics processing unit (GPU), an application-specific integrated circuit (ASIC), a field programmable gate array (FPGA), etc. In various embodiments, any operations and / or functions described herein can be performed by the processor 112 or any combination of these different processors, such as a CPU working in conjunction with one or more GPUs. In various embodiments, one or more GPUs perform parallel processing tasks, such as matrix multiplication in LLM model calculations. The processor 112 can also receive user input from an input device (e.g., a keyboard or mouse) and generate output on one or more displays.

[0025] The system memory 114 of the compression server 110 stores content, such as software applications and data, for use by the processor 112. The system memory 114 can be any type of memory capable of storing data and software applications, such as random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash ROM), or any suitable combination of the foregoing. In some embodiments, a storage device (not shown) can supplement or replace the system memory 114. The storage device can include any number and type of external memory accessible to the processor 112. For example, and without limitation, the memory can include a secure digital card, an external flash memory, a portable compact disc read-only memory (CD-ROM), an optical storage device, a magnetic storage device, and / or any suitable combination of the foregoing.

[0026] The LLM architecture optimizer 116 stored in the system memory 114 is configured to generate a student LLM 122 by compressing and optimizing the architecture of the trained LLM 118. More specifically, the LLM architecture optimizer 116 generates the student LLM 122 by replacing layers of the trained LLM 118 with different compressed or optimized layers according to various target constraints (e.g., hardware and memory constraints or latency and number of parameters). The LLM architecture optimizer 116 then stores the student LLM 122 in the data memory 120. The student LLM 122 can then be used in any suitable application (e.g., the application 145 executed on the computing device 140) to perform inference operations.

[0027] The trained LLM 118 can be any type of technically feasible machine learning model. For example, in various embodiments, the trained LLM 118 can be a transformer-based LLM model with any suitable architecture, such as a generative pre-trained transformer (GPT). Similarly, the student LLM 122 can be any type of technically feasible machine learning model. For example, in various embodiments, the student LLM 122 can be a transformer-based LLM with any suitable architecture, such as a GPT. Figure 2 The architecture of the trained LLM 118 is described in more detail. Figures 3 to 7 The operations performed by the LLM architecture optimizer 116 to generate the student LLM 122 by compressing and optimizing the trained LLM 118 are described in more detail.

[0028] Data storage 120 provides non-volatile storage for applications and data in compression server 110 and computing device 140. For example, but not limited to, training data, trained (or deployed) machine learning models, and / or application data (including trained LLM 118 and student LLM 122) can be stored in data storage 120. In some embodiments, data storage 120 can include fixed or removable hard drives, flash memory devices, and CD-ROM (Compact Disc Read Only Memory), DVD-ROM (Digital Versatile Disc ROM), Blu-ray, HD-DVD (High Definition DVD), or other magnetic, optical, or solid-state storage devices. Data storage 120 can be network attached storage (NAS) and / or storage area network (SAN). Although shown as being coupled to compression server 110 and computing device 140 via network 130, in various embodiments, compression server 110 or computing device 140 can include data storage 120.

[0029] The network 130 includes any technically feasible type of communication network that allows data to be exchanged between the compression server 110, the computing device 140, the data storage 120, and an external entity or device (e.g., a web server or another networked computing device). For example, the network 130 may include a wide area network (WAN), a local area network (LAN), a cellular network, a wireless (WiFi) network, and / or the Internet.

[0030] The computing device 140 shown herein is for illustrative purposes only and may be varied and modified without departing from the scope of the present disclosure. For example, the number and type of processors 142, the number and type of system memory 144, and / or the number of applications contained in the system memory 144 may be modified as needed. Figure 1In some embodiments, any combination of processor 142 and / or system memory 144 may be included in and / or replaced by any type of virtual computing system, distributed computing system, and / or cloud computing environment, such as a public, private, or hybrid cloud system.

[0031] The processor 142 receives user input from an input device (e.g., a keyboard or mouse). The processor 142 can be in the form of any technically feasible processing device, configured to process data and execute program code. For example, any processor 142 can be a central processing unit (CPU), a graphics processing unit (GPU), an application-specific integrated circuit (ASIC), a field programmable gate array (FPGA), etc. In various embodiments, any operations and / or functions described herein can be performed by the processor 142 or any combination of these different processors, such as a CPU working in conjunction with one or more GPUs. In various embodiments, one or more GPUs perform parallel processing tasks, such as matrix multiplication in LLM model calculations. The processor 142 can also receive user input from an input device (e.g., a keyboard or mouse) and generate output on one or more displays.

[0032] Similar to the memory 114 of the compression server 110, the memory 144 of the computing device 140 stores content, such as software applications and data, for use by the processor 142. The system memory 144 can be any type of memory capable of storing data and software applications, such as RAM, ROM, EPROM, Flash ROM, or any suitable combination of the foregoing. In some embodiments, a storage device (not shown) can supplement or replace the system memory 144. The storage device can include any number and type of external memory accessible to the processor 142. For example, and without limitation, the memory can include a secure digital card, external flash memory, a portable CD-ROM, an optical storage device, a magnetic storage device, and / or any suitable combination of the foregoing.

[0033] To perform an inference operation, the application 145 stored in the memory 144 accesses the student LLM 122 in the data storage 120. The application 145 then provides input data to the student LLM 122 to generate output data.

[0034] Figure 2 According to various embodiments Figure 11 . A more detailed description of the architecture of the trained LLM 118 is provided. As shown, the trained LLM 118 includes, but is not limited to, an embedding layer 210, a plurality of layers 215(1)-(N), and a softmax layer 220. During execution, an input dataset 205 is input into the trained LLM 118, where the embedding layer 210, layers 215(1)-(N), and the softmax layer 220 then perform various operations to generate an LLM output 225.

[0035] In various embodiments, the trained LLM 118 comprises a transformer-based LLM configured to process an input dataset 205. In various embodiments, the input dataset 205 may be text data, such as words or sentences, or may be image or video data. More generally, the input dataset 205 may comprise any technically feasible data that can be processed by a transformer-based language model. Upon receiving the input dataset 205, an embedding layer 210 converts the elements of the input dataset 205 into numerical representations (called tokens) and encodes each token into a vector. The vectors generated by the embedding layer 210 are then passed through a plurality of layers 215 (1)-(N). Each layer 215 may comprise an attention layer or multiple layers of perceptrons (MLPs) layers having different numbers of internal parameters, including but not limited to the number of attention heads, key-value projection dimensions, number of neurons, or type of activation function. In various embodiments, each layer 215 may comprise a layer norm layer, a linear layer, a convolutional layer, a pooling layer, or any other type of feasible artificial neural network layer. Each layer 215 produces a vector or matrix that is the result of applying a weight matrix and an activation function to the vector or matrix output of its previous layer. The softmax layer 220 normalizes the output vector of the layer 215(N) to a probability distribution of the prediction results and generates an LLM output 225. In some embodiments, where the goal of the trained LLM 118 is question answering, next word / sentence prediction, word / sentence translation, or image generation, the LLM output 225 can be a probability distribution of the next word / sentence that appears after the input word / sentence, a translation of the input word / sentence, an answer to a question input, or an image generated in response to an image and text caption input.

[0036] Generate student large language model

[0037] Figure 3 According to various embodiments Figure 13. A more detailed illustration of the LLM architecture optimizer 116 of FIG. As shown, the LLM architecture optimizer 116 includes, but is not limited to, an operator database 310, a knowledge distillation engine 320, a candidate architecture generator 330, and a candidate architecture selector 340. In operation, the LLM architecture optimizer 116 receives a trained LLM 118 and control parameters 302 from a user via a user interface (not shown) and generates a student LLM 122. The control parameters 302 include, but are not limited to, student layers 304 and target devices 306. The student layers 304 list candidate layers for the student LLM 122. In various embodiments, the student layers 304 may include a copy of each layer 215 of the trained LLM 118, a pruned version of each layer 215 of the trained LLM 118, a quantized version of each layer 215 of the trained LLM 118, an identity layer, an attention layer, or an MLP layer. The target devices 306 are a list of devices on which the student LLM 122 may be deployed. The target device 306 may include, but is not limited to, a server machine, a desktop, a laptop, a mobile phone, or a GPU or other type of processor.

[0038] The operator database 310 receives as input the student layers 304 and target devices 306 from the control parameters 302. The operator database 310 then generates a lookup table of performance metrics for each student layer 304 running on any target device 306 and the associated target deployment settings, such as the hardware characteristics of the target device 306 and the usage associated with the target device 306. The performance metrics included in the lookup table may include, but are not limited to, processing latency, processing throughput, and memory usage. Processing latency generally measures the total time required for each student layer 304 to generate output based on a given input prompt. Processing latency can be measured in multiple stages, including the time required for each student layer 304 to process input tokens (referred to as the pre-fill stage) and the time required for each student layer 304 to generate output tokens (referred to as the decoding stage). Processing throughput generally measures the number of tokens that each student layer 304 can process or generate in a certain amount of time. Memory usage generally measures the amount of memory required to store the parameters of each student layer 304. Candidate architecture generator 330 may then use the performance metrics contained in the lookup table of operator database 310 to generate one or more constraint equations for the constrained optimization problem, as described below.

[0039] For a given layer 215 of a trained LLM 118, the knowledge distillation engine 320 receives one or more student layers 304 having the same input and output dimensions as the given layer 215 of the trained LLM 118. The knowledge distillation engine 320 then trains the one or more student layers 304 to mimic the given layer 215 of the trained LLM 118. The knowledge distillation engine 320 may train the student layers 304 using any feasible training technique, such as stochastic gradient descent with backpropagation, adaptive moment estimation (Adam), or root mean square propagation (RMSprop). During training, the knowledge distillation engine 320 first calculates the layer-by-layer loss, referred to as the operator score, between the given layer 215 of the trained LLM 118 and the one or more student layers 304 according to formula (1):

[0040]

[0041] Where X is a set of training samples, t i is the operation of a given layer 215 of the trained LLM 118, x i is the input to a given layer 215 of the trained LLM 118, {s i,j} is a set of operations of one or more student layers 304, W = {w i,j} is a set of weights corresponding to a set of operations of one or more student layers 304, and Is the loss function. Suitable loss function Examples of include, but are not limited to, L1 norm, mean squared error (MSE), and normalized MSE. The knowledge distillation engine 320 then updates the weight set W of one or more student layers 304 according to the training technique to mimic the operation of a given layer 215 of the trained LLM 118. The knowledge distillation engine 320 repeats this process for each of the remaining layers 215 of the trained LLM 118.

[0042] The candidate architecture generator 330 receives the operator database 310, the operator scores and weights of the student layer 304 trained by the knowledge distillation engine 320, the trained LLM 118, and the target device 306 from the control parameters 302. For a given device in the target device 306, the candidate architecture generator 330 uses constrained optimization to determine a set of one or more candidate architectures for the student LLM 120. Figure 4 The operation of candidate architecture generator 330 is described in further detail.

[0043] The candidate architecture selector 340 receives the different sets of candidate architectures generated by the candidate architecture generator 330. The candidate architecture selector 340 then trains each received candidate architecture on various calibration data. The calibration data may include, but is not limited to, the original training data used to train the trained LLM 118, a subset of the original training data used to train the trained LLM 118, or any other data not presented to the trained LLM 118 during training. After training all the different sets of candidate architectures received from the candidate architecture generator 330, the candidate architecture selector 340 selects the candidate architecture with the smallest error between the predicted output and the true output with respect to the calibration data set as the student LLM 122. The operation of the candidate architecture selector 340 will be described below in conjunction with Figure 5 Described in further detail.

[0044] Figure 4 According to various embodiments Figure 3 3. A more detailed illustration of the candidate architecture generator 330 of FIG. As shown, the candidate architecture generator 330 includes, but is not limited to, a constrained optimization formulator 410 and an integer linear programming solver 415. As described above, in operation, the candidate architecture generator 330 receives the operator database 310, the operator scores and weights of the student layer 304 trained by the knowledge distillation engine 320, the trained LLM 118, and the target device 306 from the control parameters 302, and generates different candidate architectures 420 (1)-(K). More specifically, the target device 306, the operator database 310, and the operator scores and weights of the student layer 304 trained by the knowledge distillation engine 320 are input to the constrained optimization formulator 410. The constrained optimization formulator 410 then sets a constrained optimization problem to be solved by the integer linear programming solver 415, wherein the constrained optimization problem includes, but is not limited to, an objective function to be minimized with respect to specific variables and one or more constraint equations that set conditions on these specific variables. In this regard, for a given target device 306, the constrained optimization formulator 410 formulates the objective function contained in the optimization problem as a loss function The loss function estimates the error between the output of the candidate architecture 420 and the true value on a given dataset. In various embodiments, the loss function It can be determined according to formula (2):

[0045]

[0046] Where W is the weight set of the student layer 304 trained by the knowledge distillation engine 320, Z = {z i} is a set of binary vectors, where z i is a one-hot vector representing the selection of the student layer 304, is a candidate architecture defined by Z and W, and X is a dataset whose elements are labeled pairs (x, y). The constraint optimization formulator 410 then uses the lookup table of the operator database 310 to determine at least one constraint equation according to formula (3):

[0047]

[0048] in, is a function that defines the performance budget for a given target device 306, p i is a vector whose components are the performance indicators for each student layer 304 corresponding to a given target device 306 contained in the lookup table of the operator database 310, and is a user-defined constant.

[0049] The constrained optimization formulator 410 passes equations (2)-(3) to the integer linear programming solver 415. In turn, the integer linear programming solver 415 first approximates the loss function in equation (2) using a linear function according to equation (4):

[0050]

[0051] in, is the set of operations performed by the trained LLM 118 in response to input x, is a constant representing the error between the output of the trained LLM 118 and the true value y on the dataset X, and e i is a vector where each component is the difference between the output of the trained LLM 118 on the small annotated dataset and the output of the trained LLM 118 on the same small annotated dataset when a given student layer 304 contained in the student layers 304 replaces one of the layers 215 of the trained LLM 118 and all other layers 215 of the trained LLM 118 remain unchanged (repeated for all different student layers 304 to obtain the vector e i Then, the integer linear programming solver 415 generates a candidate architecture 420 by solving the following integer linear minimization problem given by equations (5)-(6):

[0052]

[0053] The integer linear programming solver 415 can solve equations (5)-(6) using any feasible integer linear optimization technique, such as a cutting plane algorithm or a branch-and-bound algorithm. The integer linear programming solver 415 generates a set of candidate architectures 420 by solving the following linear minimization problem given by equations (7)-(9):

[0054]

[0055] Wherein, equation (9) acts as a constraint for maximum overlap with any other solution of equations (7)-(8). After completing these operations, integer linear programming solver 415 passes candidate architecture 420 to candidate architecture selector 340.

[0056] Figure 5 According to various embodiments Figure 3 Detailed description of candidate architecture selector 340. As shown, candidate architecture selector 340 includes, but is not limited to, a candidate architecture trainer 510, a selected candidate architecture 520, and a tuner 530. In operation, candidate architecture selector 340 receives candidate architectures 420 from candidate architecture generator 330 and generates a student LLM 122. In this regard, different sets of candidate architectures 420 are first input to candidate architecture trainer 510, which trains the different sets of candidate architectures 420 on a calibration dataset. The calibration data may include, but is not limited to, the original training data used to train the trained LLM 118, a subset of the original training data used to train the trained LLM 118, or any other data not provided to the trained LLM 118 during training. Subsequently, candidate architecture trainer 510 selects, from each set of candidate architectures 420, the candidate architecture 420 with the smallest error between the predicted output and the true output on the calibration dataset as the selected candidate architecture 520. Each selected candidate architecture 520 is then input to tuner 530 for further training. The fine-tuner 530 trains each selected candidate architecture 520 on the same dataset used to train the trained LLM 118, using the same learning rate schedule as used to train the trained LLM 118. A learning rate schedule is a technique for adjusting the learning rate between iterations during training. Techniques implemented by a learning rate schedule can include, but are not limited to, step decay, exponential decay, and cosine annealing. After training, the fine-tuner 530 outputs the student LLM 122.

[0057] Figure 6 is a flow chart of method steps for generating candidate architectures according to various embodiments. Figures 1 to 5 Although the method steps are described with respect to a system, those skilled in the art will understand that any system configured to perform the method steps, in any order, falls within the scope of the various embodiments.

[0058] As shown, method 600 begins at step 602, where LLM architecture optimizer 116 receives control parameters 302 from a user via a user interface (not shown). Examples of different control parameters 302 that may be input by a user include, but are not limited to, student layer 304 and target device 306.

[0059] At step 604, the LLM architecture optimizer 116 generates an operator database based on the received control parameters 302. More specifically, the control parameters 302 are input into the operator database 310. The operator database 310 then generates a lookup table of performance metrics for each student layer 304 included in the control parameters 302, which is run on any target device 306 included in the control parameters 302 in conjunction with the associated target deployment settings. The target deployment settings may include, but are not limited to, hardware characteristics of the target device 306 and a usage regime associated with the target device 306. The performance metrics included in the lookup table may include, but are not limited to, processing latency, processing throughput, and memory usage.

[0060] At step 606, the LLM architecture optimizer 116 receives the trained LLM model 118, which can be any type of machine learning model. For example, in various embodiments, the trained LLM 118 can be a transformer-based LLM with any suitable architecture, such as GPT. The LLM architecture optimizer 116 can receive the trained LLM 118 from any storage device, such as the data storage 120.

[0061] At step 608, for each given layer 215 of the trained LLM 118, the knowledge distillation engine 320 trains the student layer 304 contained in the control parameters 302 to mimic the given layer 215 of the trained LLM 118. The knowledge distillation engine 320 trains the different student layers 304 by calculating the layer-by-layer loss between the given layer 215 of the trained LLM 118 and the different student layers 304. The knowledge distillation engine 320 can use any loss function during these training operations, such as the L1 norm, MSE, and normalized MSE. Similarly, the knowledge distillation engine 320 can use any feasible training technique to train the student layers 304, such as stochastic gradient descent with backpropagation, Adam, or RMSprop.

[0062] At step 610, for each target device 306 included in the control parameters 302, the constrained optimization formulator 410 generates a constrained optimization problem using the operator database 310 and the trained student layer generated by the knowledge distillation engine 320. The constrained optimization problem includes, but is not limited to, an objective function to be minimized with respect to specific variables and one or more constraint equations that condition these specific variables. The constrained optimization formulator 410 uses the trained student layer generated by the knowledge distillation engine 320 to generate the objective function of the constrained optimization problem and uses the performance indicators included in the lookup table of the operator database 310 to generate the one or more constraint equations of the constrained optimization problem.

[0063] At step 612 , for each constrained optimization problem generated by the constrained optimization formulator 410 , the integer linear programming solver 415 approximates the objective function contained in the constrained optimization problem as a linear function to generate a linear constrained optimization problem.

[0064] In step 614, the integer linear programming solver 415 solves different linear constrained optimization problems to generate candidate architectures 420. The integer linear programming solver 415 can use any feasible integer linear optimization technique to solve the linear constrained optimization problem, such as a cutting plane algorithm or a branch and bound algorithm.

[0065] Figure 7 is a flow chart of method steps for generating a student LLM according to various embodiments. Figures 1 to 5 Although the method steps are described with respect to a system, those skilled in the art will understand that any system configured to perform the method steps, in any order, falls within the scope of the various embodiments.

[0066] As shown, method 700 begins at step 702 , where candidate architecture selector 340 receives different sets of candidate architectures 420 from candidate architecture generator 330 .

[0067] At step 704, the candidate architecture selector 340 trains each set of candidate architectures 420 on calibration data. The calibration data may include, but is not limited to, the original training data used to train the trained LLM 118, a subset of the original training data used to train the trained LLM 118, or any other data not provided to the trained LLM 118 during training.

[0068] At step 706, for each set of candidate architectures 420, the candidate architecture trainer 510 selects the candidate architecture with the smallest error as the selected candidate architecture 520. As previously described herein, in various embodiments, the candidate architecture trainer 510 selects the candidate architecture 420 with the smallest error between the predicted output and the true output with respect to the calibration dataset as the selected candidate architecture 520.

[0069] At step 708, the fine-tuner 530 trains the selected candidate architectures 520 to generate the student LLM 122. The fine-tuner 530 trains each selected candidate architecture 520 on the same dataset used to train the trained LLM 118 and uses the same learning rate schedule as used when training the trained LLM 118. The fine-tuner can use any feasible learning rate schedule during this fine-tuning phase, such as step decay, exponential decay, and cosine annealing. At step 710, the candidate architecture selector 340 outputs the student LLM 122.

[0070] In summary, the architecture of a trained LLM is optimized for execution on a specific hardware and compressed based on target constraints (such as latency and number of parameters) to build a smaller "student" LLM. First, a database of potential layer types (called student layers) is built for the student LLM. Next, an operator score is calculated for each student layer of the student LLM based on how well the student layer mimics the corresponding layer in the original trained LLM. The operator score and hardware-specific latency and memory constraints are used to build a set of candidate architectures for the student LLM. Each candidate architecture is then trained on a calibration dataset, and the architecture with the smallest loss in accuracy at execution is selected for deployment and subsequently fine-tuned. The result is a smaller LLM with improved execution latency, reduced memory usage, and minimized loss in accuracy compared to the original, trained LLM.

[0071] At least one technical advantage of the disclosed technology relative to the prior art is that the disclosed technology can greatly facilitate the training and deployment of LLMs across multiple different hardware implementations. In this regard, the disclosed technology can be used to generate different trained student LLMs for each different hardware implementation based on a single trained LLM. Therefore, using the disclosed technology, multiple different smaller trained student LLMs can be generated for multiple different hardware implementations without having to train each student LLM from scratch, thereby reducing the time and computing resources required to deploy new trained models. In addition, the student LLMs generated using the disclosed technology can improve the latency and memory usage of the original trained LLM without causing any significant loss of accuracy. The disclosed technology also implements constrained optimization problems to automatically design different student LLM architectures, thereby eliminating the need for manual hyperparameter tuning. These technical advantages provide one or more technical improvements over prior art methods.

[0072] 1. Some embodiments relate to a computer-implemented method for generating a trained artificial neural network, wherein the method includes: for each model layer included in a trained model, training one or more student model layers to imitate the model layer; for a first target device included in a plurality of target devices, generating one or more candidate architectures based on a constrained optimization problem and the one or more trained student model layers; training the one or more candidate architectures on a calibration dataset; selecting a first candidate architecture included in the one or more candidate architectures that is associated with a minimum amount of error; and performing a plurality of fine-tuning training operations on the first candidate architecture to generate a first trained student model.

[0073] 2. A computer-implemented method according to clause 1, wherein generating the one or more candidate architectures includes: generating a linear constrained optimization problem based on the objective function contained in the constrained optimization problem; and calculating a solution to the linear constrained optimization problem to generate at least one of the one or more candidate architectures.

[0074] 3. A computer-implemented method according to clause 1 or 2, wherein the linear constrained optimization problem comprises a linear function, the linear function comprising an approximation to the objective function involved in the constrained optimization problem.

[0075] 4. The computer-implemented method of any of clauses 1 to 3, wherein the one or more student model layers and the plurality of target devices include user-defined control parameters.

[0076] 5. A computer-implemented method according to any one of clauses 1 to 4, wherein the first candidate architecture has a smaller error between a predicted output and a true output generated using the calibration dataset than any other candidate architecture included in the one or more candidate architectures.

[0077] 6. A computer-implemented method according to any one of clauses 1 to 5, wherein performing the plurality of fine-tuning operations on the first candidate architecture comprises: training the first candidate architecture on a dataset used to train the trained model.

[0078] 7. A computer-implemented method according to any one of clauses 1 to 6, wherein a learning rate schedule used in training the trained model is implemented when performing the plurality of fine-tuning operations on the first candidate architecture.

[0079] 8. A computer-implemented method according to any one of clauses 1 to 7, wherein the one or more student layers include a copy of each layer included in the trained model, a pruned version of each layer included in the trained model, at least one identity layer, at least one attention layer, or at least one multi-layer perceptron layer.

[0080] 9. A computer-implemented method according to any one of clauses 1 to 8, wherein the one or more student model layers and the model layer have the same input dimension and the same output dimension.

[0081] 10. A computer-implemented method according to any one of clauses 1 to 9, wherein the first trained student model has less execution latency relative to the execution latency associated with the trained model.

[0082] 11. One or more non-transitory computer-readable media comprising instructions that, when executed by one or more processors, cause the one or more processors to perform the following steps: for each model layer included in a trained model, train one or more student model layers to mimic the model layer; for a first target device included in a plurality of target devices, generate one or more candidate architectures based on a constrained optimization problem and the one or more trained student model layers; train the one or more candidate architectures on a calibration dataset; select a first candidate architecture included in the one or more candidate architectures that is associated with the least amount of error; and perform multiple fine-tuning training operations on the first candidate architecture to generate a first trained student model.

[0083] 12. The one or more non-transitory computer-readable media of clause 11, wherein the plurality of target devices comprises at least one of a server machine, a desktop machine, a graphics processing unit, a laptop machine, or a mobile phone.

[0084] 13. One or more non-transitory computer-readable media according to clause 11 or 12, wherein the first trained student model has a memory footprint that is smaller than a memory footprint associated with the trained model.

[0085] 14. One or more non-transitory computer-readable media according to any one of clauses 11-13, wherein generating the one or more candidate architectures comprises: generating a linear constrained optimization problem based on an objective function included in the constrained optimization problem; and calculating a solution to the linear constrained optimization problem to generate at least one of the one or more candidate architectures.

[0086] 15. One or more non-transitory computer-readable media according to any of clauses 11-14, wherein the linear constrained optimization problem comprises a linear function comprising an approximation to the objective function involved in the constrained optimization problem.

[0087] 16. One or more non-transitory computer-readable media as recited in any of clauses 11-15, wherein the one or more student model layers and the plurality of target devices include user-defined control parameters.

[0088] 17. One or more non-transitory computer-readable media according to any of clauses 11-16, wherein the first candidate architecture has a smaller error between a predicted output and a true output generated using the calibration data set than any other candidate architecture included in the one or more candidate architectures.

[0089] 18. One or more non-transitory computer-readable media according to any of clauses 11-17, wherein performing the plurality of fine-tuning operations on the first candidate architecture comprises: training the first candidate architecture on a dataset used to train the trained model.

[0090] 19. One or more non-transitory computer-readable media according to any of clauses 11-18, wherein a learning rate schedule used in training the trained model is implemented when performing the plurality of fine-tuning operations on the first candidate architecture.

[0091] 20. Some relate to a computer system comprising: one or more memories comprising instructions; and one or more processors coupled to the one or more memories and, when executing the instructions, configured to perform the following steps: for each model layer included in a trained model, train one or more student model layers to mimic the model layer; for a first target device included in a plurality of target devices, generate one or more candidate architectures based on a constrained optimization problem and the one or more trained student model layers; train the one or more candidate architectures on a calibration dataset; select a first candidate architecture included in the one or more candidate architectures that is associated with the least amount of error; and perform multiple fine-tuning training operations on the first candidate architecture to generate a first trained student model.

[0092] Any and all combinations of any claim elements recited in any claim and / or any elements described in this application, in any manner, are within the intended scope of the invention and protection.

[0093] The description of the various embodiments is presented for illustrative purposes only and is not intended to be exhaustive or limited to the disclosed embodiments. It will be apparent to those skilled in the art that many modifications and variations can be made without departing from the scope and spirit of the described embodiments.

[0094] Aspects of the present embodiment can be embodied as a system, method or computer program product. Therefore, various aspects of the present disclosure can adopt the form of a complete hardware embodiment, a complete software embodiment (including firmware, resident software, microcode, etc.) or an embodiment in combination with software and hardware, which can be generally referred to as a "module", "system" or "computer" in this article. In addition, any hardware and / or software technology, process, function, component, engine, module or system described in this disclosure can be implemented as a circuit or circuit group. In addition, various aspects of the present disclosure can adopt the form of a computer program product embodied in one or more computer-readable media, wherein the computer-readable medium has a computer-readable program code embodied therein.

[0095] Any combination of one or more computer-readable media may be used. A computer-readable medium may be a computer-readable signal medium or a computer-readable storage medium. A computer-readable storage medium may be, for example, but not limited to, an electronic, magnetic, optical, electromagnetic, infrared, or semiconductor system, device, or apparatus, or any suitable combination thereof. More specific examples (non-exhaustive list) of computer-readable storage media would include the following: an electrical connection with one or more wires, a portable computer disk, a hard disk, a random access memory (RAM), a read-only memory (ROM), an erasable programmable read-only memory (EPROM or flash memory), an optical fiber, a portable compact disc read-only memory (CD-ROM), an optical storage device, a magnetic storage device, or any suitable combination thereof. In the context of this document, a computer-readable storage medium may be any tangible medium that can contain or store a program for use by or in connection with an instruction execution system, apparatus, or device.

[0096] Aspects of the present disclosure have been described above with reference to the flowcharts and / or block diagrams of the methods, devices (systems) and computer program products according to embodiments of the present disclosure. It should be understood that each box in the flowchart and / or block diagram and the combination of boxes in the flowchart and / or block diagram can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer or other programmable data processing device to generate a machine. When the instructions are executed by the processor of a computer or other programmable data processing device, the functions / actions specified in the flowchart and / or block diagram box can be implemented. Such a processor can be, but is not limited to, a general-purpose processor, a special-purpose processor, a specific application processor or a field programmable gate array.

[0097] The flowcharts and block diagrams in the figures illustrate the architecture, functions and operations that may be implemented according to the systems, methods and computer program products of various embodiments of the present disclosure. In this regard, each box in the flowchart or block diagram may represent a part of a module, segment or code, which includes one or more executable instructions for implementing a specified logical function. It should also be noted that in some alternative implementations, the functions indicated in the box may not occur in the order indicated in the figure. For example, two boxes shown in succession can actually be executed substantially simultaneously, or sometimes in reverse order, depending on the functions involved. It should also be noted that each box in the block diagram and / or flowchart and the combination of boxes in the block diagram and / or flowchart can be implemented by a special hardware-based system or a combination of special hardware and computer instructions that performs the specified function or action.

[0098] While the foregoing is directed to embodiments of the present disclosure, other and further embodiments of the disclosure may be devised without departing from the basic scope thereof, and the scope of the disclosure is determined by the claims that follow.

Claims

1. A computer-implemented method for generating a trained artificial neural network, the method comprising: For each model layer included in the trained model, training one or more student model layers to imitate the model layer; For a first target device included in the plurality of target devices, generating one or more candidate architectures based on a constrained optimization problem and one or more trained student model layers; training the one or more candidate architectures on a calibration dataset; selecting a first candidate architecture included in the one or more candidate architectures that is associated with a minimum amount of error; as well as A plurality of fine-tuning training operations are performed on the first candidate architecture to generate a first trained student model.

2. The computer-implemented method of claim 1 , wherein: Generating the one or more candidate architectures includes: generating a linear constrained optimization problem based on an objective function included in the constrained optimization problem; and calculating a solution to the linear constrained optimization problem to generate at least one of the one or more candidate architectures.

3. The computer-implemented method of claim 2, wherein: The linear constrained optimization problem includes a linear function that includes an approximation to the objective function included in the constrained optimization problem.

4. The computer-implemented method of claim 1 , wherein: The one or more student model layers and the plurality of target devices include user-defined control parameters.

5. The computer-implemented method of claim 1 , wherein: The first candidate architecture has a smaller error between a predicted output generated using the calibration dataset and a true output than any other candidate architecture included in the one or more candidate architectures.

6. The computer-implemented method of claim 1 , wherein: Performing the plurality of fine-tuning operations on the first candidate architecture includes training the first candidate architecture on a dataset used to train the trained model.

7. The computer-implemented method of claim 6, wherein: A learning rate schedule used in training the trained model is implemented when performing the plurality of fine-tuning operations on the first candidate architecture.

8. The computer-implemented method of claim 1 , wherein: The one or more student layers include a copy of each layer included in the trained model, a pruned version of each layer included in the trained model, at least one identity layer, at least one attention layer, or at least one multilayer perceptron layer.

9. The computer-implemented method of claim 1 , wherein: The one or more student model layers and the model layer have the same input dimension and the same output dimension.

10. The computer-implemented method of claim 1 , wherein: The first trained student model has an execution latency that is less than an execution latency associated with the trained model.

11. One or more non-transitory computer-readable media comprising instructions that, when executed by one or more processors, cause the one or more processors to perform the following steps: For each model layer included in the trained model, training one or more student model layers to imitate the model layer; For a first target device included in the plurality of target devices, generating one or more candidate architectures based on a constrained optimization problem and one or more trained student model layers; training the one or more candidate architectures on a calibration dataset; selecting a first candidate architecture included in the one or more candidate architectures that is associated with a minimum amount of error; as well as A plurality of fine-tuning training operations are performed on the first candidate architecture to generate a first trained student model.

12. The one or more non-transitory computer-readable media of claim 11, wherein: The plurality of target devices include at least one of a server machine, a desktop computer, a graphics processing unit, a laptop computer, or a mobile phone.

13. The one or more non-transitory computer-readable media of claim 11, wherein: The first trained student model has a memory footprint that is smaller than a memory footprint associated with the trained model.

14. The one or more non-transitory computer-readable media of claim 11, wherein: Generating the one or more candidate architectures includes: generating a linear constrained optimization problem based on an objective function included in the constrained optimization problem; and calculating a solution to the linear constrained optimization problem to generate at least one of the one or more candidate architectures.

15. The one or more non-transitory computer-readable media of claim 14, wherein: The linear constrained optimization problem includes a linear function that includes an approximation to the objective function included in the constrained optimization problem.

16. The one or more non-transitory computer-readable media of claim 11, wherein: The one or more student model layers and the plurality of target devices include user-defined control parameters.

17. The one or more non-transitory computer-readable media of claim 11, wherein: The first candidate architecture has a smaller error between a predicted output generated using the calibration dataset and a true output than any other candidate architecture included in the one or more candidate architectures.

18. The one or more non-transitory computer-readable media of claim 11, wherein: Performing the plurality of fine-tuning operations on the first candidate architecture includes training the first candidate architecture on a dataset used to train the trained model.

19. The one or more non-transitory computer-readable media of claim 18, wherein: A learning rate schedule used in training the trained model is implemented when performing the plurality of fine-tuning operations on the first candidate architecture.

20. A computer system comprising: one or more memories comprising instructions; and One or more processors, coupled to the one or more memories and, when executing the instructions, configured to perform the following steps: For each model layer included in the trained model, training one or more student model layers to imitate the model layer; For a first target device included in the plurality of target devices, generating one or more candidate architectures based on a constrained optimization problem and one or more trained student model layers; training the one or more candidate architectures on a calibration dataset; selecting a first candidate architecture included in the one or more candidate architectures that is associated with a minimum amount of error; as well as A plurality of fine-tuning training operations are performed on the first candidate architecture to generate a first trained student model.

Images