Implementing and training computationally efficient neural network architecture utilizing layer skipping logic

Abstract

A method includes providing a first set of data objects to a first skip router of a neural network (NN). The NN includes a first NN layer and a second NN layer. The first set of data objects is subdivided into a first set of skipped objects and a first set of non-skipped objects based on first skipping logic implemented by the first skipping router and a first context for each data object in the first set of data objects. A first set of processed objects is generated based on the first set of non-skipped objects and first layer logic implemented by the first NN layer. A prediction is generated based on the second set of data objects and second layer logic implemented by the second NN layer. The second set of data objects includes a first set of processed objects and a first set of skipped objects.

Description

[0001] Priority Declaration

[0002] This application is based upon and claims priority from U.S. Provisional Application No. 63 / 488,901, filed on March 7, 2023, which is incorporated herein by reference. Technical Field

[0003] The present disclosure relates generally to neural networks and, more particularly, to implementing and training neural network architectures that utilize layer skipping logic for achieving computational efficiency. Background Art

[0004] In recent years, large-scale transformer-based language models trained on general-purpose corpora have demonstrated significant generalization improvements, particularly in in-context few-shot learning. Despite their impressive capabilities for text generation, training and deploying these massive models is challenging, even with recent advances in hardware and software. One of the key challenges is that processing each input requires activating all parameters of the model, often requiring trillions of floating-point operations (FLOPs) per prediction. This places a significant burden on both model training and inference, as there is no control over the amount of computation that can be assigned to each input example. Summary of the Invention

[0005] Aspects and advantages of embodiments of the present disclosure will be set forth in part in the following description, or may be obvious from the description, or may be learned through practice of the embodiments.

[0006] One example aspect of the present disclosure relates to a method comprising providing a first set of data objects to a first skip router of a neural network (NN). The NN further comprises a first NN layer and a second NN layer. The first set of data objects is subdivided into a first set of skip objects and a first set of non-skipped objects based on first skip logic implemented by the first skip router and a first context for each data object in the first set of data objects. A first set of processed objects is generated based on the first set of non-skipped objects and a first layer of logic implemented by the first NN layer. One or more predictions are generated based on a second set of data objects and a second layer of logic implemented by the second NN layer. The second set of data objects comprises the first set of processed objects and the first set of skipped objects.

[0007] Other aspects of the disclosure relate to various systems, apparatuses, non-transitory computer-readable media, user interfaces, and electronic devices.

[0008] These and other features, aspects and advantages of various embodiments of the present disclosure will be better understood with reference to the following description and appended claims.The accompanying drawings, which are incorporated in and constitute a part of this specification, illustrate example embodiments of the present disclosure and, together with the description, serve to explain the relevant principles. BRIEF DESCRIPTION OF THE DRAWINGS

[0009] Detailed discussions of embodiments for those skilled in the art are set forth in this specification with reference to the accompanying drawings, in which:

[0010] Figure 1A Depicted is a block diagram of an example computing system for implementing and training computationally efficient neural networks according to example embodiments of the present disclosure.

[0011] Figure 1B Depicted is a block diagram of an example computing device implementing and training a computationally efficient neural network according to example embodiments of the present disclosure.

[0012] Figure 1C Depicted is a block diagram of an example computing device implementing and training a computationally efficient neural network according to example embodiments of the present disclosure.

[0013] Figure 2 Depicted is a flowchart of an example method to be performed in accordance with an example embodiment of the present disclosure.

[0014] Figure 3A The overall framework of the embodiment is shown.

[0015] Figure 3B A softmax function for a router function is shown according to various embodiments.

[0016] Figure 4 An overview of an efficient implementation of an embodiment based on dynamic gather and scatter is shown.

[0017] Figure 5 An overview of applying skip layers to a transformer-based model is shown, according to various embodiments.

[0018] Figure 6 Pseudo code for a first algorithm for a forward pass skipping layers according to various embodiments is provided.

[0019] Figure 7 Pseudo code for a second algorithm for skipping layers per decoding step according to various embodiments is provided.

[0020] Reference numerals that are repeated across multiple figures are intended to identify like features in the various implementations. DETAILED DESCRIPTION

[0021] Overview

[0022] In general, the present disclosure relates to implementing and training neural network architectures that utilize layer skipping logic for computational efficiency. Various neural network architectures can implement transformer models. Transformer models can be used for natural language (NL) tasks. The proposed neural network architecture significantly reduces the time and space complexity of training the implemented models. The trained models can be used for various prediction tasks, such as natural language processing (NLP) tasks.

[0023] In recent years, large-scale transformer-based language models trained on general-purpose corpora have demonstrated dramatic improvements in generalization, particularly for contextual few-shot learning. Despite their impressive capabilities for text generation, training and deploying these massive models is challenging, even with recent advances in hardware and software. One of the key challenges is that processing each input requires activating all parameters of the model, often requiring trillions of floating-point operations (FLOPs) per prediction. This places a significant burden on both model training and inference, as there is no control over the amount of computation that can be assigned to each input example.

[0024] In particular, allocating the same computational power of a large model uniformly to process all examples is often wasteful and inefficient. These problems are even more severe when using a corpus of real-world data to train large models, as the redundancy of easy examples becomes more prominent as more data is used.

[0025] Compared to other transformer models, the example model architecture described herein can be operated to adjust the model's computational performance based on the 'difficulty' of the input. Specifically, the proposed architecture can process 'easy' examples (such as commonly used stop words, punctuation marks, patches in image backgrounds, etc.) with only a relatively low amount of computational cost. However, the architecture can be adjusted to allow for additional processing (e.g., more computational cost) for 'difficult' examples (e.g., rare abstract concepts) when they are truly needed.

[0026] More generally, aspects of this disclosure utilize conditional computation, a paradigm in which only a small subset of model parameters are activated based on the input representation, thereby reducing the amount of computation required per example. However, due to the discrete nature of the decisions made based on each input, training models with conditionally activated components end-to-end, differentiably and efficiently remains challenging, and thus, this paper provides a solution to this challenge.

[0027] In particular, in the present disclosure, a framework known as skip layer logic model (SLLM) is developed that allows an input context representation to conditionally skip any layer that can be wrapped inside the framework. More specifically, an example model employing SLLM can be trained end-to-end differentiably while still adhering to discrete decisions during forward pass, which enables the proposed model to precisely control the performance-compute tradeoff through external constraints.

[0028] Furthermore, because discrete decisions can be preserved during forward pass, an efficient implementation is also developed such that additional computation can be further saved in both pre-training and inference given a target budget. In one example application, SLLM is applied to transformer architecture to exhibit the efficiency of the method for decoding language model pre-training and decoding. Various methods of embodiments are evaluated based on a suite of NLP benchmarks from open-domain QA tasks, reading comprehension, common sense reasoning to natural language inference tasks. The model with SLLM has demonstrated strong one-shot performance compared to multiple competing baselines, and has controllable compute tradeoff between model quality and decoding efficiency.

[0029] While many of the embodiments are discussed in connection with applying SLLM to neural networks trained for various NLP tasks, embodiments are not limited thereto. Embodiments can be applied to any neural network application that receives input data that is essentially sequential. Word tokens in an NLP phrase (e.g., a paragraph, a sentence, or a sentence fragment) are typically provided as sequential input to a SLLM-based neural network architecture. Each word token can be embedded in the form of a vector embedding, and a sequence of word token embeddings can be fed to the neural network. Likewise, embodiments can be used for sequential genetic sequences. In such embodiments, certain fragments of a genetic sequence can be able to skip various layers of a SLLM-based neural network depending on the context of the atoms of the genetic sequence. Embodiments can be used for other sequential input data-based applications such as, but not limited to, signal processing, genomics, autonomous vehicles, computer vision, etc.

[0030] Aspects of the present disclosure provide a number of technical effects and benefits. As one example technical effect and benefit, via the skip logic, the computational efficiency of pre-training, training, and employing a predictive model is significantly reduced. That is, the spatial and temporal complexity of pre-training, training, and employing a predictive model is significantly reduced because, based on the context of a data object, the computation for each data object (e.g., vector embedding of an NL token) is significantly reduced. Thus, the use of computational resources can be reduced. For example, the number of processor cycles can be reduced, the use of computer memory can be reduced, and / or the use of network bandwidth can be reduced.

[0031] Example embodiments of the present disclosure will now be discussed in further detail with reference to the accompanying drawings.

[0032] Example devices and systems

[0033] Figure 1A A block diagram of an example computing system 100 for implementing and training a computationally efficient neural network according to an example embodiment of the present disclosure is depicted. The system 100 includes a user computing device 102, a server computing system 130, and a training computing system 150 communicatively coupled via a network 180.

[0034] The user computing device 102 may be any type of computing device, such as, for example, a personal computing device (e.g., a laptop or desktop computer), a mobile computing device (e.g., a smartphone or tablet computer), a gaming console or controller, a wearable computing device, an embedded computing device, or any other type of computing device.

[0035] The user computing device 102 includes one or more processors 112 and memory 114. The one or more processors 112 may be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.) and may be a single processor or multiple processors operatively connected. The memory 114 may include one or more non-transitory computer-readable storage media, such as RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, etc., and combinations thereof. The memory 114 may store data 116 and instructions 118 that are executed by the processor 112 to cause the user computing device 102 to perform operations.

[0036] In some implementations, the user computing device 102 may store or include one or more machine learning models 120. For example, the machine learning model 120 may be or may otherwise include various machine learning models, such as a neural network (e.g., a deep neural network) or other types of machine learning models, including nonlinear models and / or linear models. The neural network may include a feedforward neural network, a recurrent neural network (e.g., a long short-term memory recurrent neural network), a convolutional neural network, or other forms of neural networks. Some example machine learning models may utilize an attention mechanism, such as self-attention. For example, some example machine learning models may include a multi-head self-attention model (e.g., a transformer model).

[0037] In some implementations, one or more machine learning models 120 may be received from the server computing system 130 via the network 180, stored in the user computing device memory 114, and then used or otherwise implemented by the one or more processors 112. In some implementations, the user computing device 102 may implement multiple parallel instances of a single machine learning model 120.

[0038] Additionally or alternatively, one or more machine learning models 140 may be included in or otherwise stored and implemented by a server computing system 130 that communicates with the user computing device 102 according to a client-server relationship. For example, the machine learning models 140 may be implemented as part of a web service by the server computing system 140. Thus, one or more models 120 may be stored and implemented at the user computing device 102, and / or one or more models 140 may be stored and implemented at the server computing system 130.

[0039] The user computing device 102 may also include one or more user input components 122 for receiving user input. For example, the user input component 122 may be a touch-sensitive component (e.g., a touch-sensitive display or touchpad) that is sensitive to the touch of a user input object (e.g., a finger or a stylus). A touch-sensitive component may be used to implement a virtual keyboard. Other example user input components include a microphone, a traditional keyboard, or other devices by which a user can provide user input.

[0040] The server computing system 130 includes one or more processors 132 and memory 134. The one or more processors 132 may be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.) and may be a single processor or multiple processors operatively connected. The memory 134 may include one or more non-transitory computer-readable storage media, such as RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, etc., and combinations thereof. The memory 134 may store data 136 and instructions 138, which are executed by the processor 132 to cause the server computing system 130 to perform operations.

[0041] In some implementations, the server computing system 130 includes or is otherwise implemented by one or more server computing devices. Where the server computing system 130 includes multiple server computing devices, such server computing devices may operate according to a sequential computing architecture, a parallel computing architecture, or some combination thereof.

[0042] As described above, the server computing system 130 may store or otherwise include one or more machine learning models 140. For example, the model 140 may be or may otherwise include various machine learning models. Example machine learning models include neural networks or other multi-layer nonlinear models. Example neural networks include feedforward neural networks, deep neural networks, recurrent neural networks, and convolutional neural networks. Some example machine learning models may utilize attention mechanisms, such as self-attention. For example, some example machine learning models may include a multi-head self-attention model (e.g., a transformer model).

[0043] User computing device 102 and / or server computing system 130 may train models 120 and / or 140 via interaction with training computing system 150 communicatively coupled via network 180. Training computing system 150 may be separate from server computing system 130 or may be part of server computing system 130.

[0044] The training computing system 150 includes one or more processors 152 and memory 154. The one or more processors 152 may be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.) and may be a single processor or multiple processors operatively connected. The memory 154 may include one or more non-transitory computer-readable storage media, such as RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, etc., and combinations thereof. The memory 154 may store data 156 and instructions 158, which are executed by the processor 152 to cause the training computing system 150 to perform operations. In some implementations, the training computing system 150 includes or is otherwise implemented by one or more server computing devices.

[0045] The training computing system 150 may include a model trainer 160 that trains the machine learning models 120 and / or 140 stored at the user computing device 102 and / or the server computing system 130 using various training or learning techniques, such as, for example, error backpropagation. For example, a loss function may be backpropagated through the model to update one or more parameters of the model (e.g., based on the gradient of the loss function). Various loss functions may be used, such as mean squared error, likelihood loss, cross entropy loss, hinge loss, and / or various other loss functions. Gradient descent techniques may be used to iteratively update the parameters over multiple training iterations.

[0046] In some implementations, performing error backpropagation may include performing truncated backpropagation through time.The model trainer 160 may perform various generalization techniques (eg, weight decay, dropout, etc.) to improve the generalization ability of the model being trained.

[0047] In some implementations, if the user has provided consent, the training examples may be provided by the user computing device 102. Thus, in such implementations, the model 120 provided to the user computing device 102 may be trained by the training computing system 150 based on user-specific data received from the user computing device 102. In some cases, this process may be referred to as personalizing the model.

[0048] The model trainer 160 includes computer logic for providing the desired functionality. The model trainer 160 can be implemented in hardware, firmware, and / or software that controls a general-purpose processor. For example, in some implementations, the model trainer 160 includes a program file stored on a storage device, loaded into a memory, and executed by one or more processors. In other implementations, the model trainer 160 includes one or more sets of computer-executable instructions stored in a tangible computer-readable storage medium (such as RAM, a hard disk, or an optical or magnetic medium).

[0049] The network 180 may be any type of communication network, such as a local area network (e.g., an intranet), a wide area network (e.g., the Internet), or some combination thereof, and may include any number of wired or wireless links. In general, communications over the network 180 may be conducted via any type of wired and / or wireless connection using a wide variety of communication protocols (e.g., TCP / IP, HTTP, SMTP, FTP), encodings or formats (e.g., HTML, XML), and / or protection schemes (e.g., VPN, secure HTTP, SSL).

[0050] Figure 1A An example computing system that can be used to implement the present disclosure is shown. Other computing systems may also be used. For example, in some implementations, the user computing device 102 may include a model trainer 160 and a training dataset 162. In such implementations, the model 120 may be both trained and used locally at the user computing device 102. In some of such implementations, the user computing device 102 may implement the model trainer 160 to personalize the model 120 based on user-specific data.

[0051] Figure 1B Depicted is a block diagram of an example computing device 10 performing in accordance with an example embodiment of the present disclosure. Computing device 10 may be a user computing device or a server computing device.

[0052] Computing device 10 includes multiple applications (e.g., Application 1 through Application N). Each application includes its own machine learning library and machine learning model. For example, each application may include a machine learning model. Example applications include text messaging applications, email applications, dictation applications, virtual keyboard applications, browser applications, etc.

[0053] like Figure 1B As shown, each application can communicate with multiple other components of the computing device (such as, for example, one or more sensors, a context manager, a device state component, and / or additional components). In some implementations, each application can communicate with each device component using an API (e.g., a public API). In some implementations, the API used by each application is specific to that application.

[0054] Figure 1C Depicted is a block diagram of an example computing device 50 performing in accordance with an example embodiment of the present disclosure. Computing device 50 may be a user computing device or a server computing device.

[0055] The computing device 50 includes a plurality of applications (e.g., Application 1 through Application N). Each application communicates with a central intelligence layer. Example applications include a text messaging application, an email application, a dictation application, a virtual keyboard application, a browser application, and the like. In some implementations, each application can communicate with the central intelligence layer (and the models stored therein) using an API (e.g., a public API across all applications).

[0056] The central intelligence layer includes multiple machine learning models. For example, Figure 1C As shown, a corresponding machine learning model can be provided for each application, and the corresponding machine learning model can be managed by the central intelligence layer. In other implementations, two or more applications can share a single machine learning model. For example, in some implementations, the central intelligence layer can provide a single model for all applications. In some implementations, the central intelligence layer is included in the operating system of the computing device 50 or is otherwise implemented by the operating system.

[0057] The central intelligence layer can communicate with the central device data layer. The central device data layer can be a centralized data repository for the computing device 50. Figure 1C As shown, the central device data layer can communicate with multiple other components of the computing device (such as, for example, one or more sensors, a context manager, a device state component, and / or additional components). In some implementations, the central device data layer can communicate with each device component using an API (e.g., a private API).

[0058] Example methods

[0059] Figure 2A flow chart of an example method 200 is depicted for use in performing according to an example embodiment of the present disclosure. Figure 2 The steps of method 600 are described in a specific order for the purpose of illustration and discussion, but the method of the present disclosure is not limited to the specific order or arrangement shown. The steps of method 600 can be omitted, rearranged, combined and / or adjusted in various ways without departing from the scope of the present disclosure.

[0060] Method 200 may be performed by a computing system including one or more processors and one or more non-transitory computer-readable media. The media may store instructions that, when executed by one or more processors, cause the computing system to perform operations including those blocks of method 200.

[0061] At block 202, a computing system may provide a first set of data objects to a first set of router inputs of a first skip router of a neural network (NN) implemented by the computing system. The NN may include a set of NN layers and a set of skip routers. The set of NN layers may include at least a first NN layer, the first NN layer including a first set of layer inputs and a first set of layer outputs. The set of skip routers may include at least a first skip router, the first skip router further including a first set of router outputs that are fed forward to the first set of layer inputs and a second set of router outputs that are fed forward to the first set of layer outputs.

[0062] At block 204, the computing system may subdivide the first set of data objects into a first set of skipped objects and a first set of non-skipped objects that is disjoint from the first set of skipped objects. Subdividing the first set of data objects may be based on first skip logic implemented by a first skip router and a first context for each data object in the first set of data objects.

[0063] At block 206, the computing system may provide the first set of non-skipped objects to the first NN layer via the first router output set and the first layer input set. The first NN layer may process the first set of non-skipped objects based on the first layer logic implemented by the first NN layer to generate a first set of processed objects. The first set of processed objects may be provided to the first layer output set.

[0064] At block 208, the computing system may provide the first set of skipped objects to the first set of tier outputs via the second set of router outputs.The first set of skipped objects and the first set of processed objects may form a second set of data objects.

[0065] At block 210 , the computing system may determine one or more predictions based on the second set of data objects and a second NN layer in the set of NN layers.

[0066] In some embodiments, the second NN layer includes a second layer input set and a second layer output set. The set of skip routers further includes a second skip router including a second router input set, a third router output set fed forward to the second layer input set, and a fourth router output set fed forward to the second layer output set. The first layer output set is fed forward to the second router input set.

[0067] The computing system may provide a second set of data objects to a second skip router via a first layer output set and a second router input set. The computing system may further subdivide the second set of data objects into a second set of skip objects and a second set of non-skipped objects that is disjoint from the second set of skip objects. The computing system may perform the subdivision based on second skip logic implemented by the second skip router and a second context for each data object in the second set of data objects. The computing system may provide the second set of non-skipped objects to a second NN layer via a third router output set and a second layer input set. The second NN layer processes the second set of non-skipped objects based on the second layer logic implemented by the second NN layer to generate a second set of processed objects. The second set of processed objects is provided to a second layer output set. The computing system may provide the second set of skipped objects to a second layer output set via a fourth router output set. The second set of skipped objects and the second set of processed objects form a third set of data objects. The computing system may determine one or more predictions based on the third set of data objects and a third NN layer in the set of NN layers.

[0068] In various embodiments, a first NN layer is a first feed-forward layer (FFL) in a set of NN layers. A computing system may subdivide a first set of non-skipped objects into a first subset of non-skipped objects and a second subset of non-skipped objects that is disjoint from the first set of non-skipped objects. The computing system may provide the first subset of non-skipped objects to the first FFL. The second subset of non-skipped objects may be provided to the first FFL. The first subset of non-skipped objects is provided to the first FFL in series with the second subset of non-skipped objects being provided to the first FFL. The first subset of non-skipped objects includes at least a first non-skipped object and a second non-skipped object. The second subset of non-skipped objects includes at least a third non-skipped object and a fourth non-skipped object. The computing system provides the first and third non-skipped objects to a first implementation of the first FFL. The computing system provides the second and fourth non-skipped objects to a second implementation of the first FFL.

[0069] The first non-skipped object and the second non-skipped object are provided to the first implementation of the first FFL in parallel with the second implementation of the first FFL. The third non-skipped object and the fourth non-skipped object are provided to the second implementation of the first FFL in parallel with the first implementation of the first FFL. The first non-skipped object and the third non-skipped object are provided to the first implementation of the first FFL in series with the first implementation of the first FFL. The second non-skipped object and the fourth non-skipped object are provided to the second implementation of the first FFL in series with the second implementation of the first FFL. A first cardinality of the first subset of non-skipped objects and a second cardinality of the second subset of non-skipped objects are indicated by a hyperparameter of the NN.

[0070] The first set of data objects is a first ordered set having a first order. The first set of skipped objects is a second ordered set having a second order based on the first order. The first set of non-skipped objects is a third ordered set having a third order based on the first order. The computing system can form a second set of data objects based on the first order of the first set of data objects, the second order of the first set of skipped objects, and the third order of the first set of non-skipped objects, such that the second set of data objects is a fourth ordered set having the first order.

[0071] In various embodiments, the first NN layer is a first feed forward layer (FFL) of a set of NN layers. The set of NN layers further includes a self-attention layer and a first layer normalization (LN) layer. The computing system can provide the first set of objects to the first LN layer. A key-value pair for each data object in the first set of data objects can be determined based on first LN logic implemented by the first LN layer. The key-value pair for each data object in the first set of data objects is provided to the self-attention layer. The computing system can provide the first set of non-skipped objects to the self-attention layer. A self-attention output is generated based on self-attention logic implemented by the self-attention layer, the key-value pair for each data object in the first set of data objects, and the set of non-skipped objects. The first set of processed objects is generated based on the self-attention output and the first layer logic.

[0072] The computing system may determine a first LN output based on the first LN logic and the first set of data objects. The first LN output may be provided to a first skip router. The computing system may further subdivide the first set of data objects into a first set of skipped objects and a first set of non-skipped objects based on the first LN output. The computing system may determine a query for each non-skipped object in the first set of non-skipped objects. The query for each non-skipped object in the first set of non-skipped objects may be provided to a self-attention layer. A self-attention output may be generated further based on the query for each non-skipped object in the first set of non-skipped objects. The self-attention output may be provided to a second LN layer. A second LN output may be determined based on second LN logic implemented by the second LN layer and the self-attention output. The second LN output may be provided to the first FFL. A first processed object set may be generated further based on the second LN output and the first layer logic.

[0073] Each data object in the first set of data objects may be a vector embedding of a natural language (NL) word. The NN may implement an NL transformer model. One or more predictions may be associated with one or more NL tasks performed by the NL transformer model. The NL transformer model may be a decoder model. The first skip logic and the first layer logic may be jointly trained via a loss function that includes a first term associated with the accuracy of the one or more predictions and a second term associated with the cardinality of the first set of non-skipped objects. When the cardinality of the first set of non-skipped objects is large, the second term of the loss function generates a larger penalty.

[0074] System architecture

[0075] We will now discuss skipping layers. represents a parameterized layer (or module) of a neural network, where the input and output is given, and the optional set of weights is given by For example, a normal feed-forward layer (FFN) can be represented by Indicates that and are the input weight and output weight respectively.

[0076] Skip layers can be designed to wrap existing layers so that

[0077] ,

[0078] in Has learnable weights The router function. Figure 3A 300 of the embodiment is shown. More specifically, Figure 3AIn

[15] , a router (e.g., a skip router) can determine whether to activate or skip the embedding layer logic based on the input context.

[0079] For each word input, given a batch of B sequences , each sequence has a length of T and an embedding dimension of d.

[0080] , making

[0081]

[0082] Therefore, if Figure 3A As shown, any existing layer (e.g., FFN) that is applied to the input in a point-by-point manner can be embedded inside the skip layer. Based on the context, if the router determines to skip, the input will be connected directly to the output, otherwise the input will go through the embedded layer logic.

[0083] Various functionalities of a router (e.g., a skip router). The routers of various embodiments implement the router function G(X|W G ), the router function can be learned to assign only a subset of the input to the embedding layer to achieve the best model performance under a given budget. For an input token batch, the router can output a binary mask matrix

[0084] .

[0085] Can exist for design There are multiple choices for . One choice is the sigmoid function , the sigmoid function normalizes each value independently to be in the continuous range (0, 1) as a soft approximation of the binary mask. Although this approximation can be differentiated, it may require additional thresholds to produce the binary decision rule of the above equation.

[0086] Another design choice for the routing function is Top-K (K=1) routing, e.g. To solve the non-differentiability of the argmax operator, for each input token , can be Normalize the point integral score.

[0087]

[0088] This allows the gradient to be back-propagated through the coefficient g.

[0089] Figure 3B , a softmax function for a router function according to various embodiments is shown. Figure 3BIn [1], the router function uses the Gumbel-Softmax function. During the forward pass, binary variables are sampled. During the backward pass, gradients are backpropagated to update the routers. More specifically, during the forward pass, as shown above, the sampled binary values ​​are returned for the gating function G(X[b,t]). During the backward pass, soft probabilities are used as g (as shown above) to propagate gradients back to update the router weights. Because the percentage of tokens not skipped during the forward pass is calculated based on the binary mask shown above, the density of the model can be controlled.

[0090] We will now discuss the router capacity. The binary mask output by the router in the above equation is the assignment of a subset of the tokens in the batch to the embedding layer inside the skip layer. For simplicity, assume that each sequence in a batch of size B has the same sequence length T. Then, the ratio is the percentage (or probability) of a word being assigned to a layer, which can be called the capacity. Think of P as a global budget for how many words can be assigned to a layer. An auxiliary loss term can be introduced ,in It is a layer The capacity of each layer will obey the budget constraint. The overall loss function of the model can be ,in is the negative log-likelihood of the average predicted next word. By optimizing On the one hand, the layer capacity can be pushed closer to the target probability P. On the other hand, The item can continuously improve the prediction accuracy of the model. The term will force only P% of the token batch to pass through the layer, thus reducing the first term , “hard” examples that lead to marginal reduction on average can be prioritized, while “easy” examples that have already reached low perplexity can be skipped in order to save FLOPs.

[0091] As discussed below, these embodiments can be implemented efficiently. Advantages of the embodiments include that the number of inputs computed by each layer varies across the entire layer stack and changes continuously during training. At the same time, this dynamic nature is also challenging to implement on TPUs, where computations on tensors with static shapes are typically prioritized.

[0092] Non-limiting examples include first Figure 3A The given layer logic in is applied to the entire batch, and the output batch is then multiplied by the mask given above, so that the skipped tokens will not be used in this layer. However, this masking mechanism is computationally expensive, as embodiments should not spend the same computation on skipped inputs as on non-skipped inputs, especially when the skip ratio is high. Figure 4An overview 400 of an efficient implementation of embodiments based on dynamic collection and dispersion is shown. Figure 4 The overview 400 in FIG. 4 focuses on sparse computation of FFN layers, as FFN layers are widely used components and are typically computationally intensive. More specifically, non-skipped tokens are collected based on the indices generated by the router, and then fed as a group into the FFN. The results can be dispersed to the final output. The overall algorithm includes three main steps. First, based on the results of the router as shown above, all inputs are labeled as skipped or non-skipped. Second, all non-skipped inputs (in the shaded rectangle) can be collected and evenly divided into groups. Although each group can be sequentially fed into the FFN for computation, all elements in the same group can be collected, computed, and dispersed in parallel. Third, the computation results from the non-skipped inputs can be dispersed back to the final output of the FFN layer, while the skipped inputs can be directly written into the final output without any computation.

[0093] The group size (the number of inputs in a group), which can be denoted as Gsize, is a hyperparameter that controls how many tokens the FFN can process in parallel. Because the number of non-skipped inputs in a batch is dynamic and unknown in advance, Gsize can affect the training efficiency. When Gsize is too large (e.g., a single group), the group can include too many skipped inputs, resulting in suboptimal performance. When Gsize is too small, the group can produce too many small-size groups, and the computation can be close to sequential. Thus, there can be little parallelism, and the overhead can even be larger compared to the basic mask implementation. In practice, embodiments can set Gsize ∝ PBT, where P is the density, B is the batch size, and T is the sequence length.

[0094] Embodiments employing transformer-based models will now be discussed. More specifically, the following discussion focuses on applying the skip layer to a transformer-based decoder-only language model in a setup of contextual learning. A transformer layer mainly includes self-attention, layer normalization, and FFN as sublayers, and can be denoted as:

[0095]

[0096] .

[0097] Figure 5 An overview 500 of applying the skip layer to transformer-based models according to various embodiments is shown. More specifically, in Figure 5 the skip layer is applied to a single layer in a NN implementing a transformer-based model. In Figure 5 LN denotes a layer normalization layer. Also in Figure 5In , query, key, and value refer to the computation of query, key, and value projections in the self-attention layer. Attn refers to the attention computation. For simplicity, Figure 5 The residual connections in the self-attention layers and FFN layers are ignored. As noted above, embodiments (e.g., skip layers) can be applied to a single transformer layer. Some embodiments include wrapping the entire transformer-based model in skip layers to preserve the self-attention atomicity of the structure. However, Layer and Layers can have different skipping implementations. The layer can be the most computationally intensive component of a transformer model, but it can be applied to a given batch of tokens in a point-by-point manner. Therefore, each input token in the batch can be activated independently with probability P. layer. Because The layer may consume a significant amount of the total computation in the Transformer layer, so when the activation probability is small, embodiments may achieve significant savings in FLOPs.

[0098] Relative to layer, self-attention layer can consume significantly less computation. However, It may not be applied to token batches in a point-by-point manner, because tokens need to pay attention to each other to compute their own attention outputs. If most of the tokens in a P-hour batch are skipped, the remaining non-skipped tokens may lose most of the context of the corresponding sequence they belong to.

[0099] Figure 5 Pseudo code for a first algorithm 600 for a forward pass skipping layers according to various embodiments is provided. Figure 6 Pseudo code for a second algorithm 700 for skipping layers per decoding step according to various embodiments is provided. Some embodiments may follow a partial skip mechanism. Figure 7 As shown, when input tokens are skipped, their key and value projections can still be preserved (see, for example, Figure 5 2 of the first algorithm 600), because they are part of the context and are needed for further attention to the remaining non-skipped tokens. However, since the skipped tokens themselves do not require attention computation, we can still omit their query projection. Note that lines 3-11 of the first algorithm 600 can be computed in parallel using the efficient implementation discussed above to increase training speed.

[0100] Figure 6The second algorithm 700 shows the greedy decoding logic of a transformer layer based on a skip layer. The router makes the skip decision by picking the most likely result. The key and value projections are calculated and saved in the decoding cache (e.g., see line 2 of the second algorithm 700). When the router activates the current layer, the query projection can be calculated for the current token, and then the decoding cache K and V containing the key and value projections of the previous decoding step can be used to calculate the self-attention. Otherwise, the layer does not perform further calculations.

[0101] Figure 7 Additional disclosure

[0102] The technology discussed herein relates to servers, databases, software applications, and other computer-based systems, as well as the actions taken and information sent to and from such systems. The inherent flexibility of computer-based systems allows for a variety of possible configurations, combinations, and partitioning of tasks and functions between and within components. For example, the processes discussed herein can be implemented using a single device or component, or multiple devices or components working in combination. Databases and applications can be implemented on a single system or distributed across multiple systems. Distributed components can operate sequentially or in parallel.

[0103] Although the present invention has been described in detail with respect to various specific example embodiments thereof, each example is provided by way of illustration and not limitation of the present invention. Those skilled in the art may readily produce modifications, variations, and equivalents to such embodiments upon understanding the foregoing. Therefore, the present invention does not exclude such modifications, variations, and / or additions to the present invention as would be readily apparent to those skilled in the art. For example, a feature shown or described as part of one embodiment may be used in conjunction with another embodiment to produce yet further embodiments. Therefore, the present invention is intended to encompass such modifications, variations, and equivalents.

Claims

1. A computing system comprising: one or more processors; as well as One or more non-transitory computer-readable media storing instructions that, when executed by the one or more processors, cause the computing system to perform operations comprising: providing a first set of data objects to a first skip router of a neural network NN implemented by the computing system, wherein the NN further comprises a first NN layer and a second NN layer; subdividing the first set of data objects into a first set of skipped objects and a first set of non-skipped objects based on first skip logic implemented by the first skip router and a first context of each data object in the first set of data objects; generating a first set of processed objects based on the first set of non-skipped objects and a first layer of logic implemented by the first NN layer; One or more predictions are determined based on a second set of data objects and second layer logic implemented by the second NN layer, wherein the second set of data objects includes the first set of processed objects and the first set of skipped objects.

2. The computing system of claim 1 , wherein the NN further comprises a second skip router and a third NN layer, and determining the one or more predictions comprises: providing the second set of data objects to the second skip router; subdividing the second set of data objects into a second set of skipped objects and a second set of non-skipped objects based on second skip logic implemented by the second skip router and a second context for each data object in the second set of data objects; generating a second set of processed objects based on the second set of non-skipped objects and a second layer of logic implemented by the second NN layer; providing the second set of skipped objects to a second set of layer outputs via a fourth set of router outputs, wherein the second set of skipped objects and the second set of processed objects form a third set of data objects; as well as The one or more predictions are determined based on a third set of data objects and a third layer of logic implemented by the third NN layer, wherein the third set of data objects includes the second set of processed objects and the second set of skipped objects.

3. The computing system of claim 1, wherein the first NN layer is a first feed-forward layer (FFL) in a set of NN layers of the NN.

4. The computing system of claim 3 , wherein the first set of non-skipped objects is an ordered set, and the operations further comprise: subdividing the first set of non-skipped objects into a first subset of non-skipped objects and a second subset of non-skipped objects based on an order of the first set of non-skipped objects; providing the first non-skipped object subset to the first FFL; as well as The second subset of non-skipped objects is provided to the first FFL. 5 . The computing system of claim 4 , wherein the first subset of non-skipped objects is provided to the first FFL in series with the second subset of non-skipped objects being provided to the first FFL.

6. The computing system of claim 4 , wherein the first subset of non-skipped objects includes at least a first non-skipped object and a second non-skipped object, the second subset of non-skipped objects includes at least a third non-skipped object and a fourth non-skipped object, and the operations further comprise: providing the first non-skipped object and the third non-skipped object to a first implementation of the first FFL; as well as The second non-skipped object and the fourth non-skipped object are provided to a second implementation of the first FFL.

7. The computing system of claim 4 , wherein the first non-skipped object is provided to the first implementation of the first FFL in parallel with the second non-skipped object being provided to the second implementation of the first FFL, and the third non-skipped object is provided to the first implementation of the first FFL in parallel with the fourth non-skipped object being provided to the second implementation of the first FFL.

8. The computing system of claim 7 , wherein the first non-skipped object is provided to the first implementation of the first FFL in series with the first implementation of the third non-skipped object being provided to the first FFL, and the second non-skipped object is provided to the second implementation of the first FFL in series with the second implementation of the fourth non-skipped object being provided to the first FFL.

9. The computing system of claim 4, wherein a first cardinality of the first subset of non-skipped objects and a second cardinality of the second subset of non-skipped objects are indicated by hyperparameters of the NN.

10. The computing system of claim 1 , wherein the first data object set is a first ordered set having a first order, the first skipped object set is a second ordered set having a second order based on the first order, and the first non-skipped object set is a third ordered set having a third order based on the first order.

11. The computing system of claim 11, wherein the operations further comprise: The second data object set is formed based on the first order of the first data object set, the second order of the first skipped object set, and the third order of the first non-skipped object set, such that the second data object set is a fourth ordered set having the first order.

12. The computing system of claim 1 , wherein the first NN layer is a first feed-forward layer (FFL), and the NN further comprises a self-attention layer and a first normalization layer (LN), and the operations further comprise: providing the first set of objects to the first NN layer; determining a key-value pair for each data object in the first set of data objects based on first LN logic implemented by a first LN layer; providing the key-value pair of each data object in the first set of data objects to the self-attention layer; providing the first set of non-skipped objects to the self-attention layer; generating a self-attention output based on the self-attention logic implemented by the self-attention layer, the key-value pairs for each data object in the first set of data objects, and the set of non-skipped objects; as well as The first set of processed objects is generated based on the self-attention output and the first layer of logic.

13. The computing system of claim 12, wherein the operations further comprise: determining a first LN output based on the first LN logic and the first set of data objects; providing the first LN output to the first skip router; as well as The first set of data objects is further subdivided into the first set of skipped objects and the first set of non-skipped objects based on the first LN output.

14. The computing system of claim 12, wherein the operations further comprise: determining a query for each non-skipped object in the first set of non-skipped objects; providing the query for each non-skipped object in the first set of non-skipped objects to the self-attention layer; as well as The self-attention output is further generated based on the query for each non-skipped object in the first set of non-skipped objects.

15. The computing system of claim 14, wherein the set of NN layers further includes a second LN layer, and the operations further include: Providing the self-attention output to the second LN layer; determining a second LN output based on a second LN logic implemented by the second LN layer and the self-attention output; providing the second LN output to the first FFL; as well as The first set of processed objects is generated further based on the second LN output and the first layer logic.

16. The computing system of claim 1 , wherein each data object in the first set of data objects is a vector embedding of a natural language (NL) word-gram, the NN implements an NL transformer model, and the one or more predictions are associated with one or more NL tasks performed by the NL transformer model.

17. The computing system of claim 16, wherein the NL transformer model is a decoder model.

18. The computing system of claim 1 , wherein the first skip logic and the first layer logic are jointly trained via a loss function comprising a first term associated with an accuracy of the one or more predictions and a second term associated with a cardinality of the first set of non-skipped objects.

19. The computing system of claim 18, wherein the second term of the loss function generates a larger penalty when the cardinality of the first set of non-skipped objects is larger.

20. A computer-implemented method comprising: providing a first set of data objects to a first set of router inputs of a first skip router of a neural network NN, the NN comprising a set of NN layers and a set of skip routers, wherein the set of NN layers comprises at least a first NN layer, the first NN layer comprises a first set of layer inputs and a first set of layer outputs, and the set of skip routers comprises at least the first skip router, the first skip router further comprising a first set of router outputs fed forward to the first set of layer inputs and a second set of router outputs fed forward to the first set of layer outputs; subdividing the first set of data objects into a first skip object set and a first non-skipped object set disjoint from the first set of skip objects based on first skip logic implemented by the first skip router and a first context of each data object in the first set of data objects; providing the first set of non-skipped objects to the first NN layer via the first router output set and the first layer input set, such that the first NN layer processes the first set of non-skipped objects based on first layer logic implemented by the first NN layer to generate a first set of processed objects, wherein the first set of processed objects is provided to the first layer output set; as well as providing the first set of skipped objects to the first set of layer outputs via the second set of router outputs, wherein the first set of skipped objects and the first set of processed objects form a second set of data objects; One or more predictions are determined based on the second set of data objects and a second NN layer in the set of NN layers.

Images