Attention Neural Network with Gated Attention Units

Abstract

Methods, systems, and apparatus for performing machine learning tasks on network inputs to generate network outputs, including a computer program encoded on a computer storage medium. In one aspect, one of the systems includes a neural network configured to perform the machine learning task, the neural network including one or more attentive layers, each of which includes an attention unit.

Description

[Technical field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims the benefit of priority to U.S. Provisional Application No. 63 / 304,559, filed January 28, 2022, which is incorporated by reference in its entirety herein. [Background technology]

[0002] This specification relates to performing machine learning tasks of network inputs using neural networks.

[0003] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict an output for a received input. Some neural networks contain one or more hidden layers in addition to an output layer. The output of each hidden layer is used as the input for the next layer in the network, i.e. the next hidden layer and the output layer. Each layer of the network generates an output from the received input depending on the current values ​​of its respective set of parameters. Summary of the Invention

[0004] Described herein is a system implemented as a computer program on one or more computers at one or more locations that uses a neural network including one or more attentive layers to generate network outputs for received network inputs. The one or more computers may include one or more machine learning accelerators, such as one or more TPUs, GPUs, and / or other ASICs.

[0005] Instead of first applying an attention mechanism and then processing the output of the attention mechanism using a feedforward neural network, each attentive layer contains a gated attention unit that "gates" the output of the attention mechanism applied to the input sequence to the attentive layer and the outputs of one or more feedforward layers that are applied to the input sequence to the attentive layer.

[0006] Particular embodiments of the subject matter described in this specification can be implemented to realize one or more of the following advantages.

[0007] Global attention-based neural networks, e.g., those with transformer architectures, have delivered impressive results in many machine learning tasks. Despite the rapid expansion of model sizes for these neural networks, many of the existing systems are limited to short context sizes by the quadratic complexity of the transformer with respect to the input length. This limitation fundamentally prevents the models from processing long-term information, a property that is important for many complex machine learning tasks.

[0008] In particular, existing systems that apply self-attention to input sequences are constrained in the length of the input sequences they can process, i.e., they are constrained to processing input sequences with at most a certain number of elements, where the input sequence can refer to the input sequence of the task being performed, or to an already-generated portion of the output sequence that the neural network needs to generate to perform the task.

[0009] In particular, both the computational and memory requirements for applying self-attention over the entire input sequence grow quadratically with the number of elements in the input sequence (i.e., the computational and memory complexity is O(n 2), where n is the number of elements in the input sequence. Thus, when n is too large, applying self-attention to the input sequence may be infeasible, and these existing systems are therefore limited to small problem domains.

[0010] Herein, we describe a new attentive layer design that naturally enables higher quality approximations. Specifically, the attentive layer includes a gating mechanism that alleviates the burden of self-attention, so that the attentive layer is computationally less expensive and its quality is less dependent on the accuracy of the attention mechanism. Thus, in some cases, the attentive layer can use an attention mechanism that approximates the second-order attention mechanism of the Transformer, resulting in a variant with linear complexity in context size without a memory bottleneck.

[0011] In particular, the new attentive layers can be efficiently implemented on machine learning accelerators, e.g., TPUs, GPUs, or other ASICs, resulting in improved performance both in theory and when the neural network is deployed on one or more of these accelerators.

[0012] Furthermore, for a given quality (e.g., as measured by a confusion score), a system with an attentive layer as described herein can be trained in less time (e.g., with one or more machine learning accelerators) compared to known Transformer models, especially when longer sequence lengths (e.g., sequence lengths of 512 or more, such as sequence lengths in the range of 512 to 8192) are used.

[0013] The details of one or more embodiments of the subject matter herein are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, drawings, and claims. [Brief description of the drawings]

[0014] [Figure 1] 1 shows an example of a neural network system. [Diagram 2] 1 illustrates an exemplary architecture of a gated attention unit in the attentive layer. [Diagram 3] FIG. 1 is a flow diagram of an example process for processing input using an attentive layer. [Figure 4] FIG. 13 is a flow diagram of an example process for performing mixed chunk attention. DETAILED DESCRIPTION OF THE PREFERRED EMBODIMENTS

[0015] Like reference symbols and designations in the various drawings indicate like elements.

[0016] Described herein is a system, implemented as a computer program on one or more computers at one or more locations, that performs a machine learning task on a network input to generate a network output for the machine learning task.

[0017] The machine learning task can be any machine learning task that (i) operates on a network input, which is a sequence of inputs; (ii) produces a network output, which is a sequence of outputs; or (iii) does both.

[0018] Here are some examples of machine learning tasks that the system can be configured to perform:

[0019] As an example, the task may be a neural machine translation task. For example, if the input to the neural network is a sequence of text in one language, e.g., a sequence of words, phrases, characters, or word fragments, the output generated by the neural network may be a translation of the text sequence into another language, i.e., a translation of the text sequence in the other language that is the translation of the input sequence of text. As a particular example, the task may be a multilingual machine translation task, in which a single neural network is configured to translate between multiple different source and target language pairs. In this example, the source language text may be augmented with an identifier indicating the target language into which the neural network should translate the source language text.

[0020] As another example, the task may be a speech processing task. For example, if the input to the neural network is a sequence representing an utterance, the output generated by the neural network may be a score for each of a set of fragments of text, each score representing an estimated likelihood that the fragment of text is a correct transcript of the utterance. As another example, if the input to the neural network is a sequence representing an utterance, the output generated by the neural network may indicate whether a particular word or phrase (a "hot word") was spoken within the utterance. As another example, if the input to the neural network is a sequence representing an utterance, the output generated by the neural network may identify the natural language in which the utterance was spoken. For a speech processing task, the input to the neural network may include audio data (e.g., an audio signal), for example in the form of a sequence of audio data frames, and may be processed to perform a speech processing task (e.g., speech recognition).

[0021] As another example, the task can be a natural language processing or understanding task that operates on a text sequence in some natural language, such as an entailment task, a paraphrase task, a text similarity task, a sentiment task, a sentence completion task, a grammar task, etc.

[0022] As another example, the task can be a text-to-speech task, where the input is text in a natural language, or features of the text in a natural language, and the network output is a spectrogram, waveform, or other data defining the audio of the text being spoken in the natural language.

[0023] As another example, the task can be a health prediction task, where the input is a sequence derived from the patient's electronic health record data and the output is a prediction related to the patient's future health, such as a predicted treatment to be prescribed for the patient, the likelihood that the patient will experience an adverse health event, or a predicted diagnosis for the patient.

[0024] As another example, the task can be a text generation task, where the input is a text sequence and the output is another text sequence, e.g., a completion of the input sequence of text, a response to a question posed in the input sequence, or a text sequence on a topic specified by the first text sequence. As another example, the input to a text generation task can be non-textual input, e.g., an image, and the output sequence can be text that describes the input.

[0025] As another example, the task can be an image generation task, where the input is a conditioning input, e.g., text, a low-resolution image, or a subimage, and the output is an input sequence of intensity values ​​of image pixels. The goal of the image generation task can be to generate an image from a distribution that is the same as the distribution of samples in the training data conditioned on the conditioning input (i.e., the conditional distribution).

[0026] As another example, the task can be an image processing task. For example, the input can be intensity values ​​of image pixels, or an encoded representation of the intensity values ​​of pixels generated by an encoder neural network, and the network output can be (i) an image classification output that classifies the input image into one of multiple object categories, (ii) an object detection output, i.e., a sequence specifying the coordinates of one or more bounding boxes in the image that are predicted to contain an object, or (iii) a segmentation output that classifies each pixel of the input image into one of multiple categories.

[0027] As another example, a task can be an agent control task, where the inputs are sequences of observations or other data characterizing an environmental state, and the outputs define actions to be taken by the agent in response to the most recent data in the sequence. The agents can be, for example, real-world or simulated robots, control systems for industrial facilities, or control systems that control different kinds of agents.

[0028] As another example, the task can be a genomics task, where the input is a sequence representing a fragment of a DNA sequence or other molecular sequence, and the output is either an embedding of the fragment for use in a downstream task, e.g., by using unsupervised learning techniques on a dataset of DNA sequence fragments, or the output of a downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting the functional impact of non-coding variants, etc.

[0029] In some cases, the machine learning task is a combination of multiple individual machine learning tasks. That is, the system is configured to perform multiple different individual machine learning tasks, e.g., two or more of the machine learning tasks described above. For example, the system can be configured to perform multiple individual natural language understanding tasks using a network input that includes an identifier of the individual natural language understanding task to be performed on the network input.

[0030] 1 illustrates an example of a neural network system 100. Neural network system 100 is an example of a system implemented as a computer program on one or more computers at one or more locations that may implement the systems, components, and techniques described below.

[0031] The system 100 is a system that processes a network input 102 using an attention neural network 110 to generate a network output 112 for a machine learning task, such as one of the tasks listed above.

[0032] Generally, the neural network layers 120 in the attention neural network 110 include one or more initial neural network layers, e.g., an embedding layer and optionally one or more additional layers, a sequence of attentive layers 130, and one or more output layers that process the output of the final attentive layer 130 in the sequence as part of generating the network output 112.

[0033] As an example, when the network input 102 is an input sequence, the attention neural network 110 can process the network input 102 in a single forward pass and generate the network output 112.

[0034] As another example, when the network input 102 is an input sequence or is mapped to an input sequence by an encoder neural network, and the network output 112 is also a sequence containing multiple elements, the attention neural network 110 may operate autoregressively to generate the network output 112 over multiple time steps. At each time step, the attention neural network 110 processes the network input 102 (or the sequence generated from the network input 102) and already generated elements of the output sequence to generate the next element or elements of the output sequence.

[0035] Each attentional layer 130 includes a respective layer input at each of one or more positions and operates on the respective input sequence with a gated attention unit 132 .

[0036] The gated attention unit 132 receives the input sequence of layer 130 and applies the attention mechanism to the input sequence of the layer to generate an output sequence.

[0037] As an example, as described above, when the network input 102 is an input sequence, the attention neural network 110 can process the network input 102 in a single forward pass and generate the network output 112. In this example, the attentive layer 130 applies non-causal self-attention.

[0038] As another example, as described above, when the network input 102 is an input sequence or is mapped to an input sequence by an encoder neural network, and the network output 112 is also a sequence containing multiple elements, the attention neural network 110 can operate autoregressively to generate the network output 112 over multiple time steps. At each time step, the attention neural network 110 processes the network input 102 (or a sequence generated from the network input) and already generated elements of the output sequence to generate the next element or elements of the output sequence. In this example, the attentive layer 130 applies causal self-attention.

[0039] However, unlike a conventional attention layer in a conventional attention neural network, the attentive layer 130 does not first apply an attention mechanism to the input to generate an attention output, and then apply a further feedforward transformation to the attention output to generate the output of the attentive layer 130.

[0040] Instead, the gated attention unit 132 also applies a set of feedforward transformations to the input sequence, and then "gates" the output of the attention mechanism with the output of the feedforward transformations.

[0041] More specifically, unit 132 generates a respective projection input for each layer input by processing the layer inputs using one or more “first” feedforward neural network layers, and applies an attention mechanism to the input sequence to generate a respective attended layer input for each layer input.

[0042] For example, the first feedforward layer may include a linear layer followed by an activation function layer. The same first feedforward layer is applied to each layer input, e.g., in parallel, to generate a projection layer input for each layer input.

[0043] In another example, the layer inputs can be chunked into multiple chunks and a different set of first feedforward layers can be applied to the layer inputs of each chunk, i.e., each chunk has a corresponding set of first feedforward layers.

[0044] Unit 132 then generates a respective initial output for each layer input by computing an element-wise product between each projected input for the layer input and each attended layer input for the layer input, i.e., by "gating" each projected input with each attended layer input that is determined based at least in part on the other layer inputs.

[0045] Unit 132 then generates unit 132 outputs from the respective initial outputs. For example, unit 132 may generate respective updated outputs, i.e., respective unit outputs, for each layer input by processing the initial outputs using one or more second feedforward neural network layers.

[0046] For example, the second feedforward layer may include a linear layer, optionally followed by an activation function layer, and the same second feedforward layer may be applied for each layer input.

[0047] In another example, the layer inputs can be chunked into multiple chunks and a different set of second feedforward layers can be applied to the layer inputs of each chunk, i.e., each chunk has a corresponding set of second feedforward layers.

[0048] The attention mechanism applied by unit 132 depends on the configuration of the attention neural network 110, as will be explained in more detail below.

[0049] In some cases, the output sequence generated by unit 132 is the output sequence of attentive layer 130, and in other cases, attentive layer 130 applies one or more transformations, e.g., residual connections, normalization operations, or both, to the outputs generated by the units to generate the output sequence of layer 130.

[0050] FIG. 2 illustrates the operation of the gated attention unit 132 in the attentive layer 130.

[0051] As shown in FIG. 2, unit 132 receives an input sequence 202 for layer 130 that includes a respective layer input at each of one or more positions.

[0052] In general, the input sequence 202 is an intermediate representation of the current input to the attention neural network 110. That is, when the attention neural network 110 generates the network outputs 112 from the network inputs 102 in a single forward pass, the input sequence 202 is an intermediate representation of the network inputs 102. When the network outputs 112 are a sequence and the neural network 110 generates the network outputs 112 over multiple time steps, the input sequence 202 is an intermediate representation of the already generated network outputs 112 or an intermediate representation of a combinatorial sequence including the network inputs 112 and the already generated network outputs 112.

[0053] The unit 132 generates a respective projection input for each layer input of the input sequence 202 by processing the layer input using one or more first feedforward neural network layers. For example, the one or more first feedforward neural network layers may include a dense layer 210, optionally followed by an element-wise activation function. A dense layer is a fully connected neural network layer with a dense weight matrix, i.e., not constrained to have a certain percentage of values ​​equal to zero.

[0054] Unit 132 applies the attention mechanism 214 to the input sequence to generate a respective attended layer input for each layer input. In particular, unit 132 generates a query Q, a key K, and a value V from the layer inputs of the input sequence 202, and then applies the attention mechanism using Q, K, and V.

[0055] More specifically, unit 132 processes the layer inputs using one or more third feedforward neural network layers to generate respective queries and respective keys for each layer input, processes the layer inputs using one or more fourth feedforward neural network layers, e.g., dense layer 212, to generate respective values ​​for each layer input, generate a respective set of attention weights for each layer input from the respective queries for the layer inputs and the respective keys for the layer inputs, and apply, for each layer input, the respective set of attention weights for the layer input to the respective values ​​for the layer input.

[0056] Then, for each layer input, unit 132 uses the attended layer input for the layer input to "gate" 220 the projected layer input. In other words, unit 132 generates a respective initial output for each layer input by computing an element-wise product between the respective projected input for the layer input and the respective attended layer input for the layer input.

[0057] Thus, for each layer input, unit 132 element-wise multiplies the output of the attention mechanism for that layer input with the output of the feedforward operation for that layer input.

[0058] Unit 132 then generates, from the respective initial outputs for the layer inputs, a sequence of outputs for layer 130. The sequence of outputs for layer 130 includes the respective layer outputs for each layer input.

[0059] In some implementations, the initial layer output for each layer input is the layer output for the layer input.

[0060] In some other implementations, unit 132 generates a respective updated output for each layer input by processing the initial output using one or more second feedforward neural network layers. For example, the one or more second feedforward neural network layers can be another dense layer 240.

[0061] Layer 130 may then treat each updated output as a layer output for the layer input, or may optionally further modify each updated output. For example, layer 130 may generate a respective layer output for each layer input by applying a normalization operation, such as layer normalization, residual connections, or both, to each updated output for the layer input.

[0062] By utilizing the gating mechanism 220, the system reduces the burden on the self-attention mechanism, i.e., the quality of the resulting network output becomes less dependent on the modeling capabilities of the self-attention mechanism.

[0063] Thus, unit 132 can use a less computationally intensive attention mechanism than the multi-head self-attention (MHSA) mechanism used by conventional attention neural networks to achieve comparable or better results.

[0064] An example of such an attention mechanism is shown in Figure 2. In the example of Figure 2, the system reduces the computational cost of attention in various ways.

[0065] As an example, the system may compute query Q and key K by applying a linear transformation 218, e.g., multiplication with a dense weight matrix, to the layer inputs to generate a shared representation Z, and then applying per-dimensional scale and shift transformations to each dimension of Z to generate Q and K. In some implementations, the scale and shift transformations may be the same learned transformation, and query Q and key K may be equal to each other. In some other implementations, the system may learn separate transformations for each dimension for query Q and key K. Thus, rather than using separate dense matrices for both Q and K, unit 132 may use a shared dense weight matrix.

[0066] In the example of FIG. 2, unit 132 then determines whether a matrix A containing the set of respective attention weights for the layer inputs is A=relu 2 (Q.K. T +b), The attention mechanism is calculated so as to satisfy the following formula: 2 is the ReLU2 element-wise activation function, Q is the matrix of queries over the layer inputs, K is the matrix of keys over the layer inputs, and b is a bias, e.g., the relative position bias.

[0067] As another example, unit 132 may determine that a matrix A containing a set of respective attention weights for layer inputs is A=relu 2 (Qdiag(γ)K T +b), The attention mechanism can be calculated to satisfy 2 is the ReLU2 element-wise activation function, Q is a matrix of queries for each layer input, diag(γ), is a diagonal matrix with vector γ on the diagonal and zeros in all other entries, K is a matrix of keys for each layer input, and b is a bias, e.g., relative position bias.

[0068] In this example, rather than applying per-dimension scale and shift transformations, the system can set the query Q and key K equal to each other by using matrix Z as both Q and K. As an example, the system can use per-dimension transformations during training to determine the parameters of the scale and shift transformations, and then after training, Zdiag(γ)Z T The output of QK T Then, we generate a diagonal matrix diag(γ) equal to Qdiag(γ)K with Q=K=Z. T The attention mechanism can be applied using

[0069] While the attention mechanism in the above example still has quadratic complexity with respect to sequence length, it reduces the computational cost of attention in a number of ways.

[0070] As an example, the attention mechanism described is a single-headed attention mechanism, rather than the MHSA that many conventional approaches require.

[0071] As another example, the described attention mechanism uses an element-wise activation function (ReLU2) instead of the more computationally expensive softmax that many conventional approaches require.

[0072] As yet another example, the dimensionality of the queries and keys can be reduced relative to the layer inputs, making the multiplications required to generate the queries and generate the attention weights computationally less expensive. For example, if each layer input has a dimension of 512, 768, or 1024, the queries and keys can have a dimension of 128, i.e., at least four times lower than the layer inputs.

[0073] Furthermore, although not shown in FIG. 2, the use of a gating mechanism allows unit 132 to implement an approximate attention mechanism with less than quadratic complexity with respect to sequence length.

[0074] For example, the attention unit may implement a partial attention mechanism where the attention matrix A is constrained to be a sparse matrix. Examples of such attention schemes include local window attention, local + sparse attention, axial attention, and attention mechanisms that follow learnable patterns by hashing or clustering.

[0075] As another example, the attention unit may implement a linear attention mechanism that linearizes the attention computation by decomposing the attention matrix and then rearranging the order of the matrix multiplications. When applying a linear attention mechanism, the matrix containing each set of attended inputs for a layer input is Q(K T V). This rearrangement reduces the complexity from quadratic to linear with respect to sequence length.

[0076] As yet another example, the attention unit can approximate the original quadratic attention using mixed chunk attention, which merges benefits from both partial and linear attention.

[0077] When mixed chunk attention is used, the layer input is split into multiple chunks, then secondary attention is applied within each chunk and linear attention is applied across the chunks. Each chunk is a partition of the layer input sequence and includes a disjoint subset of the layer inputs from each chunk. For example, unit 132 or layer 130 may split the input sequence into chunks before processing the input sequence.

[0078] Applying mixed chunk attention is described in more detail below with reference to FIG.

[0079] Further, when mixed chunk attention is used, unit 132 may generate the projection layer inputs as described above, or may generate the projection layer inputs for each chunk independently.

[0080] That is, the one or more first feedforward neural network layers may include a respective set of one or more first feedforward neural network layers for each chunk, and for each chunk, unit 132 may generate a respective projection input for each layer input of the chunk by processing the layer input of the chunk using the respective set of one or more first feedforward neural network layers for the chunk. For example, the one or more first layers for each chunk may have the same architecture but different parameters after training.

[0081] Similarly, when mixed chunk attention is used, unit 132 may generate updated outputs as above, or independently for each chunk.

[0082] That is, the one or more second feedforward neural network layers may include a respective set of one or more second feedforward neural network layers for each chunk, and for each chunk, unit 132 may generate a respective updated output for each layer input of the chunk by processing the initial layer output for the layer input of the chunk using the respective set of one or more second feedforward neural network layers for the chunk. For example, the one or more second layers for each chunk may have the same architecture but different parameters after training.

[0083] 3 is a flow diagram of an exemplary process 300 for processing an input sequence to an attentive layer to generate an output sequence for the attentive layer. For convenience, the process 300 is described as being performed by one or more computer systems at one or more locations. For example, an attentive layer included in a neural network system, such as one of the attentive layers 130 included in the neural network system 100 of FIG. 1 when appropriately programmed, can perform the process 300.

[0084] The system obtains an input sequence for a layer, step 302. As noted above, the input sequence has a respective layer input at each of one or more positions.

[0085] The system generates a respective projection input for each layer input by processing the layer input using one or more first feedforward neural network layers (step 304).

[0086] The system applies the attention mechanism to the input sequence to generate a respective attended layer input for each layer input (step 306). As noted above, in some implementations, the attention mechanism is of quadratic complexity but at reduced computational cost compared to traditional MHSA approaches. In some other implementations, the attention mechanism is an approximation attention mechanism of less than quadratic complexity. For example, as described in more detail below with reference to FIG. 4, the attention mechanism can be a local attention mechanism, a partial attention mechanism, or a mixed chunk attention mechanism.

[0087] Furthermore, as noted above, in some implementations the attention mechanism is a causal attention mechanism, while in other implementations the attention mechanism is acausal.

[0088] The system generates a respective initial output for each layer input by computing the element-wise product between each projected input for the layer input and each attended layer input for the layer input (step 308). That is, the system "gates" each projected input for the layer input with each attended layer input for the layer input.

[0089] The system generates a sequence of outputs for the layer from each initial output for the layer input (step 310).

[0090] 4 is a flow diagram of an exemplary process 400 for performing mixed chunk attention. For convenience, the process 400 is described as being performed by one or more computer systems at one or more locations. For example, an attentive layer included in a neural network system, such as one of the attentive layers 130 included in the neural network system 100 of FIG. 1, when appropriately programmed, can perform the process 400.

[0091] To apply the attention mechanism, the system may perform steps 402-406 for each chunk. For example, the system may perform steps 402-406 in parallel for each chunk.

[0092] The system generates a respective query and a respective key for each layer input in the chunk by processing the layer input using one or more third feedforward neural network layers for the chunk (step 402).

[0093] In some implementations, the third feedforward layer is the same for each chunk, while in other implementations, each chunk has its own set of feedforward layers.

[0094] Furthermore, in some implementations, the query and key for each layer input in a chunk are the same, i.e., the system processes the layer inputs using one or more feedforward layers for the chunk to generate a shared representation Z that is used as both the query and the key for the inputs in the chunk.

[0095] The system generates a respective value for each layer input in the chunk by processing the layer input with one or more fourth feedforward neural network layers for the chunk, step 404. In some implementations, the fourth feedforward layers are the same for each chunk, while in other implementations, each chunk has its own set of feedforward layers.

[0096] The system applies a secondary attention mechanism to the layer inputs in the chunk using the respective queries, keys, and values ​​for the layer inputs in the chunk to generate respective attended secondary inputs for each layer input in the chunk (step 406).

[0097] In some implementations, as part of applying the secondary attention mechanism, the system applies per-dimensional parameters, such as per-dimensional scale and shift parameters, to the query and the key.

[0098] In particular, the system uses a suitable secondary attention mechanism to generate a respective set of attention weights for each layer input in the chunk from each query for the layer input of the chunk and each key for the layer input, and for each layer input, applies the respective set of attention weights for the layer input to each value for the layer input, generating a respective attended secondary input for the layer input.

[0099] The system may apply any suitable secondary attention mechanism within the chunks.

[0100] For example, the system may apply one of the two examples of secondary attention mechanisms described above.

[0101] That is, as an example, let us consider a matrix A that contains a set of attention weights for each layer input of chunk g. g teeth,

number

[0102] As another example, let us consider a matrix A that contains a set of attention weights for each layer input of chunk g. g teeth,

number

[0103] As noted above, in some implementations, the system may use a second exemplary attention mechanism during training while using a first attention mechanism after training.

[0104] The system applies a linear attention mechanism across the chunks to generate respective attended linear inputs for each layer input (step 408).

[0105] In some implementations, as part of applying the quadratic attention mechanism, the system applies per-dimension parameters to the query and the key, such as per-dimension scale and shift parameters. In particular, the linear and quadratic attention mechanisms can have different parameters per dimension as a result of training a neural network.

[0106] For example, the system may generate a matrix of linear attention inputs for the layer inputs of chunk g by applying a linear attention mechanism.

number

number

[0107] As another example, the system can generate a matrix of attended linear inputs for the layer inputs of chunk g by applying a linear attention mechanism.

number

number

[0108] As noted above, in some implementations, the system may use a second exemplary attention mechanism during training while using a first attention mechanism after training.

[0109] As noted above, in some implementations, the attentive layer applies causal self-attention.

[0110] In these cases, both quadratic and linear attention mechanisms are causal.

[0111] In particular, a linear attention mechanism is a causal attention mechanism that only attends to chunks that are before chunk g in the chunk ordering. For example, the system may generate a matrix of attended linear inputs for chunk g's layer inputs generated by applying the linear attention mechanism:

number

number

[0112] As another example, the system can generate a matrix of attended linear inputs for the layer inputs of chunk g by applying a linear attention mechanism.

number

number

[0113] As noted above, in some implementations, the system may use a second exemplary attention mechanism during training while using a first attention mechanism after training.

[0114] Similarly, a secondary attention mechanism is a causal attention mechanism such that each position in a chunk does not attend any positions after it in the chunk. For example, the system can implement this causality by masking the matrix of attention weights A, so that for any position i, the attention weight of all positions after position i in the chunk is zero.

[0115] For each layer input, the system combines each attended linear input for the layer input with each attended secondary input for the layer input to generate an attended input for the layer input (step 410). For example, the system can add each attended linear input for the layer input with each attended secondary input for the layer input to generate an attended input for the layer input.

[0116] By combining the outputs of both linear and quadratic attention, the system integrates the benefits from both partial and linear attention. For example, when attention is causal, the system can achieve significant training speedups compared to using linear attention over the entire input sequence while maintaining constant decoding and memory costs per step for each decoding step during inference. More generally, mixed chunk attention requires significantly less sequence dependency than traditional linear attention while significantly reducing the computations in quadratic attention. Because modern hardware accelerators that perform matrix multiplication in hardware, e.g., GPUs, TPUs, and other ASICs, can perform additions and multiplications in hardware much faster than memory accesses to read new data from memory, the reduction in sequence dependency and the resulting reduction in the number of memories required to read data from previous steps from memory results in the described attention scheme being efficiently implemented in hardware accelerators both during inference and training.

[0117] Thus, the described scheme is theoretically efficient and practically efficient when deployed across one or more hardware accelerators.

[0118] On the other hand, Fig. 4 describes that the mixed chunked attention uses second-order attention within a chunk and chunked linear attention across chunks, and in fact, any suitable partial attention mechanism can be used instead of the second-order attention. That is, any suitable partial attention mechanism can be combined with the chunked linear attention as described above. For example, the second-order attention described above can be replaced with a non-local partial attention mechanism or any other suitable partial attention mechanism.

[0119] Before using the neural network to perform a machine learning task, the training system trains the neural network to perform the task, i.e., determines trained values ​​for the parameters of the neural network, i.e., the layer of the sequence, the output layer(s), and the parameters of the embedding layer used to generate the input to the first layer of the sequence. For example, the training system can use traditional machine learning techniques to train the neural network from scratch on the training data of the task to minimize the loss function of the task, e.g., cross-entropy loss, negative log-likelihood loss, etc. As another example, the training system can first pre-train the neural network on an unsupervised objective and then fine-tune the neural network on the training data of the task. As yet another example, the training system can train the neural network on both unlabeled data and the training data of the task by semi-supervised learning.

[0120] During training, the training system can incorporate any number of techniques to improve the speed, effectiveness, or both of the training process. For example, the system can use dropout, label smoothing, or both to reduce overfitting. As another example, the system can perform training using a distributed architecture that trains multiple instances of the neural network in parallel. Additionally, as described above, the system can first pre-train the neural network on a large unsupervised dataset by unsupervised learning, e.g., to minimize the BERT loss or other unsupervised loss, and then fine-tune the neural network on task-specific training data to optimize the loss function for the task.

[0121] An "embedding" as used herein is a vector of numeric values, eg, floating point or other types of numeric values, having a given dimensionality, eg, having a given number of values.

[0122] The term "configured" is used herein in connection with systems and computer program components. A system of one or more computers configured to perform a particular operation or action means that the system has installed thereon software, firmware, hardware, or a combination thereof that causes the system to perform an operation or action. A computer program or programs configured to perform a particular operation or action means that the program or programs contain instructions that, when executed by a data processing device, cause the device to perform an operation or action.

[0123] Embodiments of the subject matter and functional operations described herein can be implemented in digital electronic circuitry, or tangibly embodied computer software or firmware, or computer hardware, including, for example, the structures disclosed herein and their structural equivalents, or one or more combinations thereof. Embodiments of the subject matter described herein can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions, encoded on a tangible, non-transitory storage medium for execution by or control of the operation of a data processing apparatus. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or one or more combinations thereof. Alternatively, or additionally, the program instructions can be encoded in an artificially generated propagated signal, e.g., a mechanically generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to a receiving device suitable for execution by the data processing apparatus.

[0124] The term "data processing apparatus" refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including, by way of example, a programmable processor, a computer, or multiple processors or computers. An apparatus can also be or further include special purpose logic circuitry, such as an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). In addition to hardware, an apparatus can also optionally include code that creates an execution environment for a computer program, such as code that constitutes a processor firmware, a protocol stack, a database management system, an operating system, or one or more combinations thereof.

[0125] A computer program, which may also be referred to as a program, software, software application, app, module, software module, script, or code, can be written in any form of programming language, including compiled or interpreted languages, declarative or procedural languages, and it can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may correspond to a file in a file system, but this is not necessarily the case. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple associated files, e.g., files that store one or more modules, subprograms, or code portions. A computer program can be deployed to be executed on one computer or on multiple computers located at one site or distributed across multiple sites and interconnected by a data communication network.

[0126] As used herein, the term "database" is used broadly to refer to any collection of data. The data need not be structured in any particular way, or at all, and may be stored on storage devices in one or more locations. Thus, an index database may contain multiple collections of data, each of which may be organized and accessed in different ways.

[0127] Similarly, the term "engine" is used broadly herein to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. In general, an engine may be implemented as one or more software modules or components, or may be installed on one or more computers in one or more locations. In some cases, one or more computers are directed to a particular engine, and in other cases, multiple engines may be installed and executed on the same computer or computers.

[0128] The processes and logic flows described herein may be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data to generate output. The processes and logic flows may also be performed by special purpose logic circuitry, such as an FPGA or an ASIC, or a combination of special purpose logic circuitry and one or more programmed computers.

[0129] A computer suitable for executing a computer program may be based on a general-purpose or special-purpose microprocessor, or on both, or on other types of central processing units. In general, the central processing unit receives instructions and data from a read-only memory, a random access memory, or on both. The basic elements of a computer are a central processing unit for implementing and executing instructions, and one or more memory devices for storing instructions and data. The central processing unit and memory may be supplemented by, or incorporated with, special-purpose logic circuitry. In general, a computer also includes one or more mass storage devices, such as magnetic disks, magneto-optical disks, or optical disks, for storing data, or is operatively connected to receive data from or transmit data to them, or both. However, a computer need not have such devices. Furthermore, a computer may be embedded in other devices, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device, such as a universal serial bus (USB) flash drive, to name a few examples.

[0130] Computer-readable media suitable for storing computer program instructions and data include all types of non-volatile memory, media, and memory devices including, by way of example, semiconductor memory devices, such as EPROM, EEPROM, and flash memory devices, magnetic disks, such as internal hard disks or removable disks, magneto-optical disks, and CD-ROM and DVD-ROM disks.

[0131] To interact with a user, embodiments of the subject matter described herein can be implemented in a computer having a display device, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user, and a keyboard and pointing device, such as a mouse or trackball, by which the user can provide input to the computer. Other types of devices can also be used to provide interaction with the user. For example, feedback provided to the user can be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback, and input from the user can be received in any form, such as acoustic, spoken language, or tactile input. Additionally, a computer can interact with a user by sending and receiving documents to a device used by the user, such as by sending a web page to a web browser on a user's device in response to a request received from the web browser. A computer can also interact with a user by sending text messages or other forms of messages to a personal device, such as a smartphone running a messaging application, and receiving a response message from the user in return.

[0132] A data processing apparatus for implementing machine learning models may also include dedicated hardware accelerator units for handling the general and numerically-intensive parts of, for example, machine learning training or machine learning production, i.e., inference, workload.

[0133] The machine learning model can be implemented and deployed using a machine learning framework, for example, the TensorFlow framework or the Jax framework.

[0134] An embodiment of the subject matter described herein can be implemented in a computing system that includes a back-end component, e.g., a data server, or includes a middleware component, e.g., an application server, or includes a front-end component, e.g., a client computer having a graphical user interface, web browser, or app that allows a user to interact with an embodiment of the subject matter described herein, or any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN), a wide area network (WAN), e.g., the Internet.

[0135] A computing system may include clients and servers. Clients and servers are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and by virtue of having a client-server relationship to each other. In some embodiments, the server sends data, e.g., HTML pages, to a user device, e.g., for the purpose of displaying the data to and receiving user input from a user interacting with the device acting as a client. Data generated at the user device, e.g., results of user interaction, can be received from the device at the server.

[0136] Although the present specification contains many specific implementation details, these should not be construed as limiting the scope or scope of any invention or patentable scope, but as a description of features specific to a particular embodiment of a particular invention. Certain features described in the context of separate embodiments herein can also be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Furthermore, even if features may be described above as functioning in a particular combination and originally claimed as such, one or more features from the claimed combination may in some cases be deleted from the combination, and the claimed combination may be directed to a subcombination or a variation of the subcombination.

[0137] Similarly, although the figures show operations in a particular order, and the claims describe operations in a particular order, this should not be understood as requiring that such operations be performed in the particular order or sequential order shown, or that all of the operations shown be performed, to achieve desirable results. In certain situations, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the program components and systems described may generally be integrated into a single software product or packaged into multiple software products.

[0138] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims may be performed in a different order and still achieve desirable results. By way of example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. 1. A system for performing a machine learning task on a network input to generate a network output, the system comprising: one or more computers; When executed by the one or more computers, An attention neural network configured to perform the machine learning task, the attention neural network including a plurality of attentional layers, each attentional layer comprising: obtaining an input sequence for the layer, the input sequence including a respective layer input at each of one or more locations; generating respective projection inputs for each layer input by processing the layer inputs using one or more first feedforward neural network layers; applying an attention mechanism to the input sequence to generate, for each layer input, a respective attended layer input; generating a respective initial output for each layer input by computing an element-wise product between the respective projected input for the layer input and the respective attended layer input for the layer input; one or more storage devices storing instructions that cause the one or more computers to implement the attention neural network configured to receive the input sequence of the layer and generate the output sequence of the layer by performing operations that include: generating an output sequence of the layer from the respective initial outputs for the layer inputs, the output sequence of the layer including a respective layer output for each layer input; Including, the system.

2. The operation includes:

2. The system of claim 1, further comprising: generating respective updated outputs for each layer input by processing the initial outputs with one or more second feedforward neural network layers.

3. The system of claim 2 , wherein the respective layer output for each layer input is the respective updated output for that layer input.

4. The operation includes:

3. The system of claim 2, further comprising: generating the respective layer outputs for each layer input by applying a normalization operation, a residual connection, or both to the respective updated outputs for the layer inputs.

5. applying an attention mechanism to the input sequence to generate a respective attended layer input for each layer input; generating a respective query and a respective key for each layer input by processing the layer inputs using one or more third feedforward neural network layers; generating a respective value for each layer input by processing the layer inputs using one or more fourth feedforward neural network layers; generating a respective set of attention weights for each layer input from the respective queries for the layer inputs and the respective keys for the layer inputs; and for each layer input, applying the respective set of attention weights for that layer input to the respective value for that layer input.

6. 6. The system of claim 5, wherein the respective keys and the respective values ​​for each layer input are derived from a same shared representation that is generated by processing the layer inputs using the one or more third feedforward neural network layers.

7. The system of claim 5 or claim 6, wherein the dimension of each query and key is less than the dimension of the layer inputs and the layer outputs.

8. The matrix A containing the respective sets of attention weights for the layer inputs is Arelu ２ (Qdiag(γ)K Ｔ +b)、 In the formula, relu ２ 8. The system of claim 5, wherein: R is an element-wise activation function of ReLU2; Q is a matrix of the respective queries on the layer inputs; diag(γ) is a diagonal matrix with vector γ on the diagonal and zeros in all other entries; K is a matrix of the respective keys on the layer inputs; and b is a bias.

9. The system of claim 8 , wherein b is a relative position bias.

10. The system according to any one of claims 1 to 5, wherein the attention mechanism is a linear attention mechanism.

11. The system according to any one of claims 1 to 5, wherein the attention mechanism is a partial attention mechanism.

12. The layer inputs are divided into a plurality of chunks, and applying an attention mechanism to the input sequence to generate a respective attended layer input for each layer input includes: For each chunk, generating a respective query and a respective key for each layer input of the chunk by processing the layer inputs using one or more third feedforward neural network layers for the chunk; generating a respective value for each layer input of the chunk by processing the layer inputs using one or more fourth feedforward neural network layers for the chunk; applying a secondary attention mechanism to the layer inputs of the chunks using the respective queries, keys, and values ​​for the layer inputs of the chunks to generate respective attended secondary inputs for each layer input of the chunks; applying a linear attention mechanism across the chunks to generate a respective attended linear input for each layer input; for each layer input, combining the respective attended linear input for the layer input with the respective attended secondary input for the layer input to generate the attended input for the layer input; The system according to any one of claims 1 to 5, comprising:

13. 13. The system of claim 12, wherein combining the respective attended linear inputs for the layer inputs and the respective attended secondary inputs for the layer inputs to generate the attended input for the layer input comprises adding the respective attended linear inputs for the layer inputs and the respective attended secondary inputs for the layer inputs to generate the attended input for the layer input.

14. applying a secondary attention mechanism to the layer inputs of the chunks using the respective queries, keys, and values ​​for the layer inputs of the chunks to generate respective attended secondary inputs for each layer input of the chunks; generating a respective set of attention weights for each layer input of the chunk from the respective queries for the layer inputs of the chunk and the respective keys for the layer inputs; and for each layer input, applying the respective set of attention weights for the layer input to the respective values ​​for the layer input to generate the respective attended secondary inputs for the layer input, wherein a matrix A includes the respective sets of attention weights for the layer inputs of the chunk g. ｇ teeth, [0010] In the formula, relu ２ is the element-wise activation function of ReLU2, and Q ｇ is the matrix of the respective queries to the layer inputs of the chunk g, diag(γ) is a diagonal matrix with vector γ on the diagonal and zeros in all other entries, and K ｇ 14. The system of claim 12 or claim 13, wherein x is a matrix of the respective keys for the layer inputs of the chunk g, and b is a bias.

15. the matrix of attended linear inputs for the layer inputs of chunk g generated by applying the linear attention mechanism [0025] but, [0030] In the formula, Q ｇ is the matrix of queries for the layer inputs of chunk g, diag(λ) is a diagonal matrix with vector λ on the diagonal and zeros in all other entries, T is the total number of layer inputs in the input sequence, C is the number of layer inputs in each chunk, and K ｈ is the matrix of the respective keys of the layer inputs of the chunk h, and V ｈ 15. The system of claim 12, wherein x is a matrix of the respective values ​​of the layer inputs of the chunk h.

16. The linear attention mechanism is a causal linear attention mechanism, and a matrix of the attended linear inputs to the layer inputs of chunk g generated by applying the causal linear attention mechanism [0045] but, [0050] In the formula, Q ｇ is the matrix of the respective queries to the layer inputs of the chunk g, diag(λ) is a diagonal matrix with vector λ on the diagonal and zeros in all other entries, summed over the chunks before chunk g in the order of the chunks in the causal linear attention mechanism, and K ｈ is the matrix of the respective keys of the layer inputs of the chunk h, and V ｈ 15. The system of claim 12, wherein x is a matrix of the respective values ​​of the layer inputs of the chunk h.

17. the one or more first feedforward neural network layers include a respective set of one or more first feedforward neural network layers for each chunk; generating respective projection inputs for each layer input by processing the layer inputs using one or more first feedforward neural network layers; 17. The system of claim 12, further comprising: for each chunk, generating a respective projection input for each layer input of the chunk by processing the layer inputs of the chunk using the respective set of one or more first feedforward neural network layers for the chunk.

18. the one or more second feedforward neural network layers include a respective set of one or more second feedforward neural network layers for each chunk; Processing the initial outputs using one or more second feedforward neural network layers to generate respective updated outputs for each layer input, The system of claim 17 when also dependent on claim 2, comprising, for each chunk, generating respective updated inputs for each layer input of the chunk by processing the initial outputs for the layer inputs of the chunk using the respective set of one or more second feedforward neural network layers for the chunk.

19. 10. The system of any preceding claim, wherein the neural network includes one or more initial layers, the sequence of attentive layers, and one or more output layers.

20. The system of claim 19 , wherein the initial layer comprises an embedding layer.

21. 21. The system of claim 19 or claim 20, wherein the network output is a sequence, and the neural network autoregressively generates the network output over a number of time steps, and at each time step, the neural network processes the network input and any already generated elements of the network output.

22. 22. The system of claim 21, wherein the attention mechanism in each attentive layer is a causal attention mechanism.

23. The system of claim 22 when also dependent on claim 16, wherein the secondary attention mechanism within each chunk is causal.

24. The system of claims 19-20, wherein the network input is a sequence and the attention mechanism in each attentive layer is non-causal.

25. One or more computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to implement an attention neural network as set forth in any preceding claim.

26. Receiving a network input; processing said network inputs using the attention neural network of any preceding claim to generate network outputs for said network inputs; A method comprising:

Images