EP4652544A1 - Attention neural networks with conditional computation attention layers - Google Patents

Abstract

Systems and methods for processing inputs using attention neural networks with conditional computation attention layers.

Description

ATTENTION NEURAL NETWORKS WITH CONDITIONAL COMPUTATION ATTENTION LAYERS

CROSS-REFERENCE TO RELATED APPLICATION

This application claims priority to U.S. Provisional Application No. 63/452,659, filed on March 16, 2023. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application.

BACKGROUND

This specification relates to performing a machine learning task on a network input using neural networks.

Neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of parameters.

SUMMARY

This specification describes a system implemented as computer programs on one or more computers in one or more locations that performs a machine learning task on a network input.

The machine learning task can be any machine learning task that (i) operates on a network input that is an input sequence, (ii) generates a network output that is an output sequence, or (iii) both.

In particular, the system performs the task using an attention neural network that includes one or more conditional computation attention layers.

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

The techniques described in this specification allow a neural network system to process input sequences, generate output sequences, or both more efficiently than existing attention-based networks both during training and at run-time, i.e., in terms of computational resources (e.g.. memory, computing power, or both), by making use of a conditional attention layer that includes (i) both light and heavy attention heads, (ii) light and heavy branches of feedforward layers, or (iii) both.

The attention layers within some existing attention neural networks employ a dotproduct attention mechanism which involves computing, for every given query vector, respective dot products of the query vector with all of the key vectors. The netw orks typically derive such key vectors or query vector vectors from network inputs that may be sequential. Thus the computational cost could be substantial when applying a dot-product attention mechanism over sequential data that is of significant length. In particular, conventional self-attention layers have a quadratic dependency on the sequence length, resulting in the model consuming a large amount of computational resources when operating on or generating longer sequences. Additionally, processing sequences of tokens requires applying feedforward and projection layers to every token, which can be quite computationally expensive for long sequences, i.e., sequences with large numbers of tokens.

The described techniques, however, address these problems by using conditional computation within the attention layers to reduce the computation cost while maintaining performance. In particular, the described techniques leverage the fact that not all tokens are equally important for generating a high quality output, especially for longer inputs, e.g., input documents or other text sequences, multi-modal inputs, or other inputs that have large numbers of tokens.

Thus, by employing conditional computation within attention layers as described in this specification, the attention neural network can devote more resources to important tokens in feedforward layers, attention layers, or both. As a result, the described attention neural network can achieve stronger performance than state of the art on a variety of long-input tasks with much faster training and inference. Notably, the described neural network can effectively and tractably make use of extremely long inputs, showing strong gains up to 64k token input length and potentially even longer.

Additionally, the described attention neural network is considerably more computationally efficient than other approaches to improving performance on long-input tasks. For example, as seen in Table 1 below, the described techniques (COLT5) require significantly fewer FLOPs per layer than a current state-of-the-art approach (L0NGT5) to improving long-input task performance and a baseline approach (T5), despite (as will be described below) the COLT5 model achieving significantly better performance on long-input tasks than either L0NGT5 or T5.

Table 1

In Table 1, n is the length of the input sequence in terms of tokens and d is the model dimensionality, i.e., the dimensions of the vectors operated on by the attention layers of the neural network. As can be seen from Table 1, given the same model dimensionality⁷, an encoder layer of the described architecture requires significantly fewer FLOPS than either T5 or LongT5 at any input length n.

Training large models that can be used for real-world tasks and then using the models to perform inference can result in significant carbon dioxide (CO2) emissions and a significant amount of electricity usage, e.g., because the data sets on which the training is done are extremely large and the models have significant numbers of parameters. Bydecreasing the number of FLOPs required during training and inference, the described techniques significantly reduce the CO2 footprint of the training and inference processes while also significantly reducing the amount of electricity consumed by the training and inference processes.

Moreover, the system can leverage parallelism in order to mitigate the impact of the additional FLOPs required by the described neural network relative to an initial architecture that does not include the heavy branches of the attention layers as described below. In particular, because the heavy and light attention heads can operate independently, the system can execute the light attention heads in parallel with the heavy attention heads, e.g., by deploying each head on a respective set of one or more hardware devices, e.g., GPUs, TPUS, or other ASICs. Thus, although the described neural network requires additional FLOPs relative to the initial architecture, the additional components (the heavy attention heads and the heavy branch of the feed-forward layers) can be executed in parallel with the corresponding existing components that are shared with the baseline architecture (the light attention heads and the light branch of the feed-forward layers), mitigating the impact on execution latency despite the additional FLOPs.

Additionally, because the heavy branches of each attention layer process only selected, proper subsets of tokens, the light branches can be more computationally efficient than the corresponding components of the initial neural network. For example, the feedforward layers can have a smaller hidden dimension, the light attention heads can use local instead of global attention, there can be fewer light attention heads, or some combination of the above. This can result in the described attention neural network architecture having better performance while being more computationally efficient, e.g., requiring fewer FLOPs, than a comparably sized initial neural network, e.g., a T5 neural network.

The details of one or more embodiments of the subject matter of this specification 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, the drawings, and the claims.

BRIEF DESCRIPTION OF THE DRAWINGS

FIG. 1 shows an example neural network system.

FIG. 2 shows an example architecture of a conditional computation attention layer.

FIG. 3A is a flow diagram of an example process for processing an input using a conditional computation attention layer.

FIG. 3B shows an example of an attention pattern of a conditional computation attention layer.

FIG. 4 is a flow diagram of an example process for processing an input using a conditional feedforward layer.

FIG. 5 is a flow diagram of an example process for training an attention neural network that includes conditional computation attention layers.

FIG. 6 shows an example of the performance of the described techniques.

FIG. 7 shows another example of the performance of the described techniques. Like reference numbers and designations in the various drawings indicate like elements.

DETAILED DESCRIPTION

This specification describes a system implemented as computer programs on one or more computers in one or more locations that performs a machine learning task on a network input to generate a network output for the machine learning task.

The machine learning task can be any machine learning task that (i) operates on a network input that is an input sequence of tokens, (ii) generates a network output that is an output sequence of tokens, or (iii) both. A token, as used in this specification, is an ordered collection of numerical values, e.g., a vector of floating point or other numerical values, that represents a given input or output.

For example, tokens can be generated by processing the corresponding inputs or outputs through an appropriate embedding layer or embedding neural network or can be the outputs of intermediate layers of the neural network.

Some examples of machine learning tasks that the system can be configured to perform follow.

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

As another example, the task may be an audio processing task. For example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network may be a score for each of a set of pieces of text, each score representing an estimated likelihood that the piece of text is the correct transcript for the utterance. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can indicate whether a particular word or phrase (“hotword”) was spoken in the utterance. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can identify the natural language in which the utterance was spoken.

As another example, the task can be a natural language processing or understanding task, e.g., an entailment task, a paraphrase task, a textual similarity task, a sentiment task, a sentence completion task, a grammaticality task, and so on, that operates on a sequence of text in some natural language.

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

As another example, the task can be a health prediction task, where the input is a sequence derived from electronic health record data for a patient and the output is a prediction that is relevant to the future health of the patient, e.g., a predicted treatment that should be prescribed to the patient, the likelihood that an adverse health event will occur to the patient, or a predicted diagnosis for the patient.

As another example, the task can be a text generation task, where the input is a sequence of text, and the output is another sequence of text, e.g., a completion of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text that is about a topic specified by the first sequence of text. As another example, the input to the text generation task can be an input other than text, e.g., an image, and the output sequence can be text that describes the input.

As another example, the task can be an image generation task, where the input is a conditioning input, e.g., text, a lower-resolution image, or a partial image, and the output is a sequence of intensity value inputs for the pixels of an image.

As another example, the task can be an image processing task. For example, the input can be the intensity values of the pixels of the image or an encoded representation of the intensity values of the pixels generated by an encoder neural netw ork, and the netw ork output can be (i) an image classification output that classifies the input image into one of a plurality of object categories (ii) an object detection output, i.e., a sequence that specifies the coordinates of one or more bounding boxes in the image that are predicted to encompass objects or (iii) a segmentation output that classifies each pixel in the input image into one of a plurality of categories. As another example, the input can include the intensity values of the pixels of the image or an encoded representation of the intensity values of the pixels generated by an encoder neural netw ork and optionally text, and the netw ork output can be text that characterizes the image, e.g., captions the image or answers a question posed by the text in the input about the image.

As another example, the task can be an audio generation task, where the input is a conditioning input, e.g., text, an image, or context audio, and the output is a sequence of tokens that represents audio.

As another example, the task can be an audio processing task. For example, the input include audio or an encoded representation of the audio generated by an encoder neural network, and the network output can be text or an image that characterizes the audio. As another example, the task can be an agent control task, where the input is a sequence of observations or other data characterizing states of an environment and the output defines an action to be performed by the agent in response to the most recent data in the sequence. The agent can be, e.g., a real-world or simulated robot, a control system for an industrial facility, or a control system that controls a different kind of agent.

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 molecule sequence and the output is either an embedding of the fragment for use in a downstream task, e g., by making use of an unsupervised learning technique on a data set of DNA sequence fragments, or an output for the downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, and so on.

As another example, the task can be a multi-modal task that requires generating a multi-modal output, processing a multi-modal input, or both.

For example, a multi-modal output or multi-modal input can include tokens representing two or more of audio, visual inputs (images, video, or both) or text.

In some cases, the machine learning task is a combination of multiple individual machine learning tasks, i.e., the system is configured to perform multiple different individual machine learning tasks, e.g., two or more of the machine learning tasks mentioned above. For example, the system can be configured to perform multiple individual natural language understanding tasks, with the network input including an identifier for the individual natural language understanding task to be performed on the network input or a prompt that describes the natural language understanding task.

FIG. 1 shows an example neural network system 100. The neural network system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented.

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, e.g., one of the tasks described above.

Generally, the neural network layers 120 within 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 attention layers 130, and one or more output layers that process the output of the last attention layer 130 in the sequence as part of generating the network output 1 12. As one 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 to generate the network output 1 12.

As another example, when the network input 102 is an input sequence or has been mapped to an input sequence by an encoder neural network and the network output 112 is also a sequence that includes multiple elements, the attention neural network 110 can operate auto-regressively and generate the network output 1 12 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 102) and the already generated elements of the output sequence to generate the next one or more elements of the output sequence.

As yet another example, when the network input 102 is an input sequence and the network output 1 12 is also a sequence that includes multiple elements, the attention neural network 110 can include an encoder neural network that generates a respective encoded representation of each of the inputs in the input sequence in a single forward pass and a decoder neural network that operates auto-regressively and generates the network output 112 over multiple time steps. At each time step, the decoder neural network processes the encoded representations and the already generated elements of the output sequence to generate the next one or more elements of the output sequence. In these examples, some of the attention layers 130 are in the encoder neural network while others are in the decoder neural network.

Each attention layer 130 operates on a respective input sequence that includes a respective input vector at each of one or more positions.

Moreover, each of the layers 130 includes an attention mechanism layer and, in some implementations, a feed-forward layer.

The attention mechanism layer receives the input sequence for the layer and applies an attention mechanism on the input sequence for the layer to generate an attended input sequence.

The attention mechanism applied by the attention mechanism layer depends on the configuration of the attention neural network.

As one example, as described above, when the netw ork input 102 is an input sequence, the attention neural network 110 can process the network input 102 in a single forward pass to generate the network output 112. In this example, the attention mechanism layers apply non-causal self-attention. As another example, as described above, when the network input 102 is an input sequence or has been mapped to an input sequence by an encoder neural network and the network output 1 12 is also a sequence that includes multiple elements, the attention neural network 110 can operate auto-regressively and 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) and the already generated elements of the output sequence to generate the next one or more elements of the output sequence. In this example, the attention mechanism apply causal self-attention.

As another example, when the network input 102 is an input sequence and some of the attention layers are in the encoder portion of the attention neural network 110 and other layers are in the decoder portion of the attention neural network 110, the attention neural network 110 can process the network input 102 using the encoder portion in a single forward pass to generate an encoded representation of the input. In this example, the attention mechanism layers within the encoder portion apply non-causal self-attention. The decoder portion of the attention neural network 110 can then operate auto-regressively and generate the network output 112 over multiple time steps. At each time step, the attention neural network 110 processes the already generated elements of the output sequence to generate the next one or more elements of the output sequence conditioned on the encoded representation. In this example, some of the attention mechanism layers in the decoder apply causal selfattention while others of the attention mechanism layers in the decoder apply cross-attention between the already generated elements of the output sequence and the encoded representation.

Generally, however, some or all of the attention mechanism layers are conditional computation attention layers 132. These layers are also referred to as "conditional attention layers.’’

Each conditional computation attention layer 132 makes use of conditional computation in order to selectively perform more processing for certain tokens in the input sequence to the layer 132.

In particular, each conditional attention layer includes (i) both light and heavy attention heads, (ii) light and heavy branches of feedforward layers, or (iii) both.

For a given conditional computation attention layer 132, all of the tokens in the input sequence are processed using the light attention heads and the light branch, while only respective proper subsets of the tokens are processed using the heavy attention heads and the heavy branch. Conditional attention layers are described in more detail below with reference to FIGS. 2 and 3.

As used in this specification, the term “learned” means that an operation or a value has been adjusted during the training of the attention neural network.

As a particular example, the attention neural network 110 can include an encoder neural network and a decoder neural network. In the encoder neural network, some or all of the attention layers are conditional attention layers 132 while in the decoder neural network, none of the layers are conditional attention layers. In this example, the system 100 can use any of a variety of optimizations to the attention mechanisms employed by the decoder neural network to further increase the computational efficiency of the neural network 110. For example, the attention layers in the decoder can employ one or more of: multi-query (rather than multi-head) attention, key-value caching, implementing local attention mechanisms, and so on, for one or more of the layers in the decoder.

FIG. 2 is an illustration 200 of an example of a conditional computation attention layer 132 that employs conditional computation for the attention heads within the layer and the feed-forward neural network within the layer. In some other examples, the conditional computation attention layer 132 can only employ conditional computation for the attention heads, and not for the feed-forward neural network. In yet other examples, the conditional computation attention layer 132 can employ conditional computation for the feed-forward neural network, but not for the attention heads.

In particular, as shown in FIG. 2, the conditional attention layer 132 includes a light branch 210 and a heavy branch 250.

The light branch 210 includes light attention heads 212 and a light branch 214 of feedforward layers while the heavy branch 250 includes heavy attention heads 252 and a heavy branch 254 of feedforward layers.

Generally, each token within the input sequence to the layer 132 is processed using all of the components in the light branch 210.

Each component in the heavy branch 250, however, only processes a respective proper subset of the tokens in the input sequence.

To that end, the light branch 210 includes a first learned routing function 216 and third learned routing function 218.

The first learned routing function 216 selects a respective proper subset of query tokens q to be updated by each of the heavy attention heads 252 in the heavy branch 250. In particular, in the example of FIG. 2, each heavy attention head 252 receives the same proper subset of query tokens q.

The first learned routing function (or, equivalently, a second learned routing function), also separately selects a respective proper subset of key-value tokens v to be used by each of the heavy attention heads 252 in the heavy branch 250 to update the query' tokens q. In the example of FIG. 2, each heavy attention head 252 receives the same proper subset of key-value tokens v.

Generally, the query tokens q can be different from the key-value tokens v. Separately selecting query' and key-value tokens allows the model to differentiate between tokens that require additional information and those that possess such information.

In some implementations, the set of heavy attention heads is more computationally expensive than the set of light attention heads. For example, the set of heavy attention heads can include a larger number of attention heads than the set of light attention heads. As a particular example, the set of heavy attention heads can include 2x, 3x, or 4x, the total number of heads as the light attention heads.

Similarly, the third learned routing function 216 selects a respective proper subset of heavy tokens m to be updated by the heavy branch 254 of feedforward layers in the heavy branch 250.

In some implementations, the heavy branch of feedforward layers is more computationally expensive than the light branch of feedforward layers.

For example, the heavy branch of feedforward layers can have a larger hidden dimension than the light branch of feedforward layers. As a particular example, the heavy branch of feedforward layers can have a hidden dimension that is 2x, 3x, or 4x larger than the light branch of feedforward layers.

That is, each branch of feedforward layers can include a sequence of feedforward layers, e.g., fully-connected layers. Optionally, one or more of these layers can be followed by a non-linear activation function. The size (dimensionality) of the input vectors received by each subsequent layer in the sequence after the first layer is referred to as the "‘hidden dimension’⁷ of the branch of feedforward layers.

As another example, the heavy branch of feedforward layers can have more feedforward layers than the light branch.

The processing of each of the components of the layer 132 is described in more detail below with reference to FIGS. 3 and 4. Thus, by performing the computation of the layer “conditionally’' as described in FIG. 2, the layer leverages the fact that some token representations may benefit from more processing than others and routes only those tokens to the heavy branch.

Additionally, the layer leverages the fact that most tokens have simple, local interactions, but some tokens benefit from heavier processing and long-range interactions and routes only those tokens for processing by the heavy attention heads.

Thus, including the layer within the neural network 110 results in improved performance, especially for longer input tasks. However, because only a proper subset of tokens are routed to the heavy branch, this improved performance comes without a corresponding significant increase in computational complexity. In fact, by decreasing the computational complexity of the light branch, e.g., by (i) decreasing the hidden dimensions employed by the light feed-forward layers, the light attention heads, or both, (ii) by using fewer light attention heads, or (iii) both, the layer can actually have reduced computational complexity and require fewer FLOPs than existing approaches despite this improved performance.

Furthermore, the system can execute the light attention heads 212 in parallel with the heavy attention heads 252, e g., by deploying each head on a respective set of one or more hardware devices, e.g., GPUs, TPUS, or other ASICs. This can mitigate latency caused by the use of the heavy branch in addition to the light branch.

FIG. 3A is a flow diagram of an example process 300 for processing an input using a conditional attention layer. For convenience, the process 300 will be described as being performed by a system of one or more computers located in one or more locations. For example, aneural network system, e.g., neural network system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 300.

The process 300 can be performed by each conditional attention layer to generate a respective attended input sequence for the conditional attention layer. As described above, each conditional attention layer includes multiple attention heads and the multiple attention heads include one or more light attention heads and one or more heavy attention heads.

The system receives an input sequence of tokens to the conditional attention layer (step 302). The input sequence of tokens generally includes a plurality of tokens. For example, for the first attention layer in the neural network, the tokens in the sequence can be embeddings generated by an embedding layer of the neural network. For each other attention layer in the neural network, the tokens in the sequence can be the outputs of the preceding layer in the neural network. The system processes the input sequence using each of the one or more light attention heads to generate, for each light attention head, a respective initial attended token for each of the tokens in the input sequence (step 304).

Generally, each light attention head applies a respective light attention mechanism over the tokens in the input sequence to generate the respective initial attended tokens for the light attention head.

As described above, the “light” attention mechanism can be less computationally expensive than the “heavy” attention mechanism.

For example, the light attention mechanisms can be local attention mechanisms. In this example, for each light attention head, applying the respective light attention mechanism over the tokens in the input sequence includes, for each token in the input sequence, attending over tokens within a local context window of the token within the input sequence to generate the respective initial attended token. The local context window is a fixed length window of tokens that is smaller than the total number of tokens in the input sequence. Thus, applying a local attention head is less computationally expensive than standard “dense” attention where each token attends across the entire input because fewer non-zero attention weights need to computed.

For each of the one or more heavy attention heads, the system selects, as query tokens for the heavy attention head, a first proper subset of the tokens in the input sequence and selects, as key-value tokens for the heavy attention head, a second proper subset of the tokens in the input sequence (step 306).

The subsets are referred to as “proper” because each proper subset includes less than all of the tokens in the input sequence.

Generally, the system selects the query’ tokens in a learned manner.

As a particular example, the system can apply a first learned routing function to the tokens in the input sequence to generate a respective first routing score each of the tokens in the input sequence and then select the q tokens having the highest first routing scores as the query tokens, where q is an integer that is less than the total number of tokens in the input sequence, thereby ensuring that only a proper subset of the tokens are selected.

For example, to apply the first learned routing function, the system can, for each token, compute a dot product between the token and a first learned embedding and then apply a generalized softmax to the dot products between the tokens and the first learned embedding to generate the routing scores.

Generally, the system also selects the key-value tokens in a learned manner. As a particular example, the system can apply a second learned routing function to the tokens in the input sequence to generate a respective second routing score each of the tokens in the input sequence and then select the v tokens having the highest second routing scores as the uery tokens, where v is an integer that is less than the total number of tokens in the input sequence, thereby ensuring that only a proper subset of the tokens are selected.

For example, to apply the second learned routing function, the system can, for each token, compute a dot product between the token and a second learned embedding and then apply a generalized softmax to the dot products between the tokens and the second learned embedding to generate the routing scores.

In some cases, the number of key-value tokens v is greater than the number of query tokens q. As a particular example, v can be equal to 2q. This allows the heavy attention heads to use context from a larger number of other tokens in the sequence without increasing the number of tokens for which attention needs to be performed.

For each of the one or more heavy attention heads, the system processes the query⁷ tokens and the key -value tokens using the heavy attention head to generate, for each query token, a respective initial attended token for the heavy attention head (step 308).

Steps 306 and 308 can, in some implementation, be performed in parallel to step 304.

Generally, for each query⁷ token, the heavy attention head attends over all of the keyvalue tokens to generate the respective initial attended token for the heavy attention head. That is. the heavy attention head applies a dense attention mechanism that is global over all of the selected key-value tokens.

Thus, as a result of performing steps 304-308, the system computes (i) for all of the tokens in the sequence, a respective initial attended token for each light attention head and (ii) for each heavy attention head, only a respective initial attended token for the query tokens selected for the heavy attention head.

The system then generates an attended input sequence of tokens.

As part of generating the attended input sequence, the system combines, for each token, the respective initial attended tokens generated for the token (step 310) to generate a sequence of combined tokens.

In particular, for any token that was not selected as a query token for any heavy attention head, the system can generate the combined token for the token from the initial attended tokens generated for the token by the light attention head(s). That is, when there is more than one light attention head, the system can combine, e.g., sum or average, the initial attended tokens generated for the token by the light attention heads. For the remaining tokens, the system can generate the combined token for the token from the initial attended tokens generated for the token by the light attention head(s) and for the heavy attention heads.

For example, when the query tokens are the same for all heavy attention heads, for each query token, the system can combine the respective initial attended tokens generated by the one or more heavy attention heads to generate a heavy attended token for the token and then scale the heavy attended token by the first routing score for the query token to generate a scaled heavy attended token.

In general, scaling tokens by the corresponding routing score provides a learning signal to the corresponding learned embedding during training while making use of the generalized softmax normalizes the scaling to provide a better distributed learning signal to all tokens.

The system can then combine the respective initial attended tokens generated by the one or more light attention heads to generate a light attended token for the token and combine the light attended token and the scaled heavy attended token, e.g., by summing the light attended token and the scaled heavy attended token.

The system can optionally perform one or more additional operations on the combined tokens to generate the attended input sequence. Examples of additional operations include residual connection operations, normalization operations (e.g., layer normalization or group normalization), and so on.

FIG. 3B shows an example 350 of an attention pattern of a conditional computation attention layer when the input sequence includes fourteen tokens.

In the example of FIG. 3B, the light attention heads apply a light attention mechanism that is a local attention mechanism with a local window of size five. That is, the local window includes, for each token, the token and the two tokens before and after the token in the sequence.

Additionally, in the example of FIG. 3B, the heavy attention heads process two query tokens and four key-value tokens.

Thus, the light branch performs local attention for each token while, in the higher capacity heavy branch, 2 selected query tokens q attend to 4 separately selected key and value tokens v. That is, the heavy branch performs “dense” attention, but on a much smaller set of selected key and value tokens v.

Because the selection of query tokens q and the key value tokens v is performed using respective learned routing functions, the model can effectively differentiate between tokens that require additional information (and select those tokens as query tokens) and those that possess such information (and select those tokens as key-value tokens).

FIG. 4 is a flow diagram of an example process 400 for processing an input using a conditional feedforward layer. For convenience, the process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, aneural network system, e.g., neural network system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 400.

The system can perform the process 400 for each conditional feedforward layer in the neural network.

The system receives the attended input sequence of tokens generated by the conditional attention layer (step 402) within the same attention layer as the conditional feedforward layer.

The system selects, as heavy tokens, a proper subset of the tokens in the attended input sequence (step 404).

Generally, the system selects the heavy’ tokens in a learned manner.

As a particular example, the system can apply a third learned routing function to the tokens in the attended input sequence to generate a respective third routing score each of the tokens in the attended input sequence and then select the k tokens having the highest third routing scores as the heavy tokens, where k is an integer that is less than the total number of tokens in the attended input sequence, thereby ensuring that only a proper subset of the tokens are selected.

For example, to apply the third learned routing function, the system can, for each token, compute a dot product between the token and a third learned embedding and then apply a generalized softmax to the dot products between the tokens and the third learned embedding to generate the learned routing scores.

For each token in the attended input sequence, the system processes the token through a light branch of feedforward neural network layers to generate a light output token for the token (step 406). That is, the system processes all of the tokens in the attended input sequence through the light branch.

For each heavy' token, the system processes the heavy token through a heavy branch of feedforward neural network layers to generate a heavy output token for the token (step 408). That is. the system processes only the proper subset of heavy tokens through the heavy branch. In some implementations, the “heavy’" branch is more computationally expensive than the light branch.

For example, the heavy branch can have more parameters than the light branch.

As one example of this, the heavy branch can have a larger hidden dimension than the light branch, e.g., resulting in the parameter matrices for the layer(s) in the heavy branch being larger.

As another example of this, the heavy branch can have more layers than the light branch.

In some examples, steps 406 and 408 are performed in parallel.

The system then generates a respective output token for each of the tokens in the attended sequence (step 410).

As part of this, for each heavy token, the system generates a combined output token by combining the light and heavy output tokens for the heavy token.

For example, the system can scale the heavy output token by the third routing score for the heavy token to generate a scaled heavy output token and combine the light output token and the scaled heavy output token, e.g.. by summing the light output token and the scaled heavy output token.

For each token not selected as a heavy token, the system generates the output token for the token from the light output token for the token.

During the processing of a given network input, for each conditional computation attention layer in the attention neural network, the system can perform the processes 300 and 400 to update the input sequence to the layer. By repeatedly performing this processing for all of the attention layers in the attention neural network and then by processing at least part of the output sequence generated by the last attention layer in the attention neural network using one or more output layers, e.g., one or more linear layers optionally followed by a softmax layer or, more generally, a multi-layer perceptron (MLP), the system can generate a network output for a received network input.

That is, the processes 300 and 400 can be performed as part of predicting an output for an input for which the desired output, i.e., the output that should be generated by the system for the input sequence, is not known.

The processes 300 and 400 can also be performed as part of processing inputs derived from a set of training data, i.e., inputs derived from a set of inputs for which the output that should be generated by the system is known, in order to train the attention neural network to determine trained values for the parameters of the attention neural network. The system can repeatedly perform the processes 300 and 400 on inputs selected from a set of training data as part of a conventional machine learning training technique to train the attention layers and the output layer(s) of the neural network, e.g., a gradient descent with backpropagation training technique that uses a conventional optimizer, e.g., stochastic gradient descent, RMSprop, or Adam optimizer, to optimize an objective function that is appropriate for the task that the attention neural network is configured to perform.

During training, the system can incorporate any number of techniques to improve the speed, the 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 the training using a distributed architecture that trains multiple instances of the attention neural network in parallel.

Moreover, the system can first pre-train some or all of the neural network on a large unsupervised data set through unsupervised learning, e.g., to minimize a BERT loss, a PEGASUS loss, a UL2 loss or other unsupervised loss, and then fine-tune the neural network on task-specific training data to optimize the objective function for the task.

In some cases, to improve performance while maximizing training efficiency, the system can pre-train the neural network using a two-level bootstrapped training method that allows the system to iteratively incorporate conditional computation attention layers into the neural network during training. For example, the system can begin the training from a pretrained or randomly initialized initial neural network that has the same architecture as the final neural network but that does not have any conditional attention layers.

FIG. 5 is a flow diagram of an example process 500 for training an attention neural network by iteratively adding conditional computation attention layers. For convenience, the process 500 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g., neural network system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 500.

The system obtains data specifying an initial, pre-trained attention neural network that includes a plurality of original attention layers (step 502). For example, the initial attention neural network can have been pre-trained as described above and can have conventional attention layers in place of conditional computation attention layers, but can otherwise have the same architecture as the attention neural network.

The system then further trains the neural network at multiple training stages.

At each of the plurality of training stages, the system updates the attention neural network by replacing one or more of the original attention layers in the attention neural network as of the training stage with a corresponding new conditional computation attention layer (step 504). For example, the system can replace one or more highest original attention layers in the sequence with a corresponding new conditional computation attention layer. In particular the system can replace a fixed number, e.g., 1, 2, or 3, highest original attention layers in the sequence with a new conditional computation attention layer.

The system obtains training data for the training stage (step 506). For example, the training data can include some or all of the training data used to pre-train the original neural network.

The system trains the attention neural network on the training data by training the one or more corresponding new conditional computation attention layers while holding any original attention layers in the attention neural network fixed (step 508), e.g., using a selfsupervised learning objective as described above.

The system can perform training stages until all of the layers or some predetermined subset of original attention layers have been replaced with (trained) conditional computation attention layers.

The system can perform the process 500 as part of additional pre-training of the neural network, e.g., through self-supervised learning, or as part of fine-tuning the neural network for a particular task, e.g., through supervised learning on training data for the particular task. In the case where the system performs the process 500 as part of fine-tuning, the system can use different training data and a different, supervised loss function for performing the process 500 than was used to pre-train the original attention neural network.

FIG. 6 shows an example 600 of the performance of the described techniques (COLT5) relative to another technique (LONGT5) that also attempts to perform the performance of a Transformer neural network on long context sequences. In particular, the example 600 show s performance on the y axis and time per sample on the x axis. The performance is measured as the average performance of the corresponding model across multiple data sets, with the performance for each data set measured using a corresponding performance metric, e.g., F, EM. or R_gm.

As shown in FIG. 6, three variants (B, L, and XT, where B, L, and XT refer to different sizes of the models in terms of number of parameters) of the described techniques each outperform comparable variants of the LONGT5 technique in terms of time per sample, and for the L and XL variants, achieve significantly smaller times per sample while achieving comparable or even better performance. FIG. 7 shows another example 700 of the performance of the described techniques (COLT5) relative to another technique (LONGT5) that also attempts to perform the performance of a Transformer neural network on long context sequences. The example 700 shows the performance in terms of Fl on a NarrativeQA data set as a function of inference time per sample for LONGT5 and COLT5 Large models using varying input lengths.

As can be seen from the example 700, COLT5 achieves both stronger performance and faster inference speed at all input lengths and is able to effectively make use of extremely long inputs. COLT5 achieves large quality gains by going from 32k to 64k tokens even while keeping the number of routed tokens constant, showing that the performance gains of COLT5 relative to other architectures strengthen with input length.

An “embedding,’⁷ as used in this specification is a vector of numeric values, e.g., floating point or other type of numeric values, that has a predetermined dimensionality, e.g., has a predetermined number of values.

This specification uses the term “configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.

Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification 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 to control the operation of. data processing apparatus. The computer storage medium can be a machine- readable storage device, a machine-readable storage substrate, a random or serial access memory' device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

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. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

A computer program, which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or 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, but need not, correspond to a file in a file system. 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 coordinated files, e.g., files that store one or more modules, sub programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communication network.

In this specification, the term ' database'’ is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, the index database can include multiple collections of data, each of which may be organized and accessed differently.

Similarly, in this specification the term “engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers. The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory' devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g.. a mobile telephone, 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, e.g., a universal serial bus (USB) flash drive, to name just a few.

Computer readable media suitable for storing computer program instructions and data include all forms of non volatile memory, media and memory devices, including by way of example semiconductor memory' devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto optical disks; and CD ROM and DVD-ROM disks.

To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory^ feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and compute-intensive parts of machine learning training or production, i.e., inference, workloads.

Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a Jax framework.

Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front end component, e.g., a client computer having a graphical user interface, a w eb brow ser, or an app through which a user can interact with an implementation of the subject matter described in this specification, 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 netw ork. Examples of communication networks include a local area netw ork (LAN) and a wide area netw ork (WAN), e.g., the Internet.

The computing system can include clients and servers. A client and server 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 having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g.. for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at the server from the device.

While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

Particular 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 can be performed in a different order and still achieve desirable results. As one 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.

What is claimed is:

Claims

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 and one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to implement: an attention neural network configured to perform the machine learning task, the attention neural network comprising one or more conditional attention layers, each conditional attention layer comprising a plurality of attention heads, the plurality of attention heads comprising one or more light attention heads and one or more heavy attention heads, each conditional attention layer configured to: receive an input sequence of tokens to the conditional attention layer that comprises a plurality of tokens; process the input sequence using each of the one or more light attention heads to generate, for each light attention head, a respective initial attended token for each of the tokens in the input sequence, wherein each light attention head applies a respective light attention mechanism over the tokens in the input sequence to generate the respective initial attended tokens for the light attention head; for each of the one or more heavy attention heads: select, as uery tokens for the heavy attention head, a first proper subset of the tokens in the input sequence and select, as key-value tokens for the heavy attention head, a second proper subset of the tokens in the input sequence; and process the query tokens and the key-value tokens using the heavy attention head to generate, for each query token a respective initial attended token for the heavy attention head, wherein, for each query token, the heavy attention head attends over all of the key-value tokens to generate the respective initial attended token for the heavy attention head; and generating an attended input sequence of tokens, comprising, for each token, combining the respective initial attended tokens generated for the token.

2. The system of claim 1, wherein, for each light attention head, applying a respective light attention mechanism over the tokens in the input sequence comprises, for each token in the input sequence, attending over tokens within a local context window of the token within the input sequence to generate the respective initial attended token.

3. The system of claim 1 or claim 2, wherein selecting, as query tokens for the heavyattention head, a first proper subset of the tokens in the input sequence comprises: applying a first learned routing function to the tokens in the input sequence to generate a respective first routing score each of the tokens in the input sequence; and selecting the q tokens having the highest first routing scores as the query- tokens, where q is an integer that is less than a total number of tokens in the input sequence.

4. The system of claim 3, wherein applying a first learned routing function to the tokens in the input sequence to generate a respective first routing score each of the tokens in the input sequence comprises: for each token, computing a dot product between the token and a first learned embedding.

5. The system of claim 4, wherein applying a first learned routing function to the tokens in the input sequence to generate a respective first routing score each of the tokens in the input sequence comprises: for each token, applying a generalized softmax to the dot product between the token and the first learned embedding.

6. The system of any one of claims 1-5, wherein selecting, as key-value tokens for the heavy attention head, a second proper subset of the tokens in the input sequence comprises: applying a second learned routing function to the tokens in the input sequence to generate a respective second routing score each of the tokens in the input sequence; and selecting the v tokens having the highest second routing scores as the key-value tokens, where v is an integer that is less than a total number of tokens in the input sequence.

7. The system of claim 6, wherein applying a second learned routing function to the tokens in the input sequence to generate a respective second routing score each of the tokens in the input sequence comprises: for each token, computing a dot product between the token and a second learned embedding.

8. The system of claim 7, wherein applying a second learned routing function to the tokens in the input sequence to generate a respective second routing score each of the tokens in the input sequence comprises: for each token, applying a generalized softmax to the dot product between the token and the second learned embedding.

9. The system of any one of claims 6-8, when dependent on claim 3, wherein v is greater than q.

10. The system of claim 9, wherein v is equal to 2q.

11. The system of any preceding claim, wherein the query tokens and the key-value tokens are the same for all of the heavy attention heads.

12. The system of claim 11, when dependent on claim 3, wherein, for each token, combining the respective initial attended tokens generated for the token comprises: for each query token: combining the respective initial attended tokens generated by the one or more heavy attention heads to generate a heavy attended token for the token; scaling the heavy attended token by the first routing score for the query⁷ token to generate a scaled heavy attended token; combining the respective initial attended tokens generated by the one or more light attention heads to generate a light attended token for the token; and combining the light attended token and the scaled heavy attended token.

13. The system of claim 12, wherein combining the light attended token and the scaled heavy attended token comprises: summing the light attended token and the scaled heavy attended token.

14. The system of any preceding claim, wherein one or more of the conditional attention layers further comprise a conditional feedforward layer, and wherein the conditional feedforward layer is configured to: receive the attended input sequence of tokens generated by the conditional attention layer; select, as heavy tokens, a proper subset of the tokens in the attended input sequence; for each token in the attended input sequence, process the token through a light branch of feedforward neural network layers to generate a light output token for the token; for each heavy token: process the heavy token through a heavy branch of feedforward neural network layers to generate a heavy output token for the token, and generate an output token for the heavy token, comprising combining the light and heavy output tokens for the heavy token; for each token not selected as a heavy token, generate an output token for the token from the light output token for the token.

15. The system of claim 14, wherein the heavy branch has more parameters than the light branch.

16. The system of claim 15, wherein the heavy branch has a larger hidden dimension than the light branch.

17. The system of claim 14 or claim 15, wherein selecting, as heavy tokens, a proper subset of the tokens in the attended input sequence comprises: applying a third learned routing function to the tokens in the attended input sequence to generate a respective third routing score each of the tokens in the attended input sequence; and selecting the k tokens having the highest third routing scores as the heavy tokens, where k is an integer that is less than a total number of tokens in the attended input sequence.

18. The system of claim 17, wherein applying a third learned routing function to the tokens in the attended input sequence to generate a respective third routing score each of the tokens in the input sequence comprises: for each token, computing a dot product between the token and a third learned embedding.

19. The system of claim 18, wherein applying a third learned routing function to the tokens in the attended input sequence to generate a respective third routing score each of the tokens in the input sequence comprises: for each token, applying a generalized softmax to the dot product between the token and the third learned embedding.

20. The system of any one of claims 14-19 when dependent on claim 17, wherein combining the light and heavy output tokens for the heavy token comprises: scaling the heavy output token by the third routing score for the heavy token to generate a scaled heavy output token; and combining the light output token and the scaled heavy output token.

21. The system of claim 20, wherein combining the light output token and the scaled heavy output token comprises: summing the light output token and the scaled heavy output token.

22. One or more computer storage media storing instructions that when executed by one or more computers cause the one or more computers to implement the attention neural network of any preceding claim.

23. A method comprising: receiving a network input; and processing the network input using the attention neural network of any preceding claim to generate a network output for the network input.

24. A method performed by one or more computers, the method comprising: obtaining data specifying an initial, pre-trained attention neural network comprising a plurality of original attention layers; and at each of a plurality of training stages: updating the attention neural network by replacing one or more of the original attention layers in the attention neural network as of the training stage with a corresponding new conditional attention layer; obtaining training data for the training stage; and training the attention neural network on the training data, comprising: training the one or more corresponding new conditional attention layers while holding any original attention layers in the attention neural network fixed.

25. The method of claim 18, wherein the plurality of attention layers are arranged in a sequence, and wherein updating the attention neural network comprises: replacing one or more highest original attention layers in the sequence with a corresponding new conditional attention layer.

26. The method of claim 18 or 19, wherein training the attention neural network on the training data comprises training the attention neural network on the training data through selfsupervised learning.