Neural networks with selective attention layers

Abstract

Methods, systems, and apparatuses, including computer programs encoded on computer storage media, for processing an input sequence (or new input token) using a neural network that includes one or more attention layer blocks that each include a selective attention layer. By generating and using aggregated mask scores to modify attention logits used by an attention mechanism, the neural network's performance is improved. Further, by using the aggregated mask scores to also remove one or more context tokens from a context buffer; and by pruning the context buffer based on the aggregated mask scores, the memory cost and computation processing cost of using the neural network to process an input sequence (or new input token) is simultaneously improved.

Description

[0001] Atorney Docket No. 56113-0812WO1

[0002] NEURAL NETWORKS WITH SELECTIVE ATTENTION LAYERS

[0003] CROSS-REFERENCE TO RELATED APPLICATION

[0004] This application claims priority to U.S. Application No. 63 / 686.686, filed August 23, 2024. The disclosure of the foregoing application is hereby incorporated by reference in its entirety.

[0005] BACKGROUND

[0006] This specification relates to processing inputs using neural networks.

[0007] 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 value inputs of a respective set of parameters.

[0008] SUMMARY

[0009] This specification describes a system implemented as computer programs on one or more computers in one or more locations that processes input sequences to perform one or more machine learning tasks.

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

[0011] Some techniques for processing input sequences using neural networks include the use of attention-based layers that operate on a context buffer having a fixed maximum token capacity7. These layers are configured to retain and process all tokens from an input sequence up to this fixed limit. This requirement to maintain and process all tokens within the context buffer hinders the performance of the neural networks for multiple reasons. First, the inclusion of irrelevant tokens (i.e., tokens that do not provide relevant information for the generation of the neural networks’ outputs) in the attention mechanism can introduce noise that degrades the performance of the neural network by interfering with the neural network’s ability to attend to the most salient tokens in the input sequence. Second, maintaining this complete context buffer results in substantial memory requirements for the key -value cache, which is an inefficient use of the fixed sized context buffer’s capacity and creates a memory7bottleneck that limits the speed and Atorney Docket No. 56113-0812WO1 efficiency of the computer-implemented inference process. Third, the attention-based layers’ processing of irrelevant tokens requires use of limited computational processing and memory bandwidth, leading to an inefficient implementation on computer hardware.

[0012] This specification describes techniques that can address the aforementioned challenges by modifying the architecture of the neural network to include one or more selective attention layers. That is, this specification describes techniques that simultaneously increase the performance of a neural network while also reducing the memory and compute requirements for processing an input sequence (or new input token) using the neural network through the use of one or more attention layer blocks that each include a selective attention layer. The described techniques generate and use aggregated mask scores to modify attention logits used by an attention mechanism to improve the neural network’s performance. Further, the described techniques use the aggregated mask scores to also remove one or more context tokens from a context buffer; and by pruning the context buffer based on the aggregated mask scores, the memory7cost and computation processing cost of using the neural network process an input sequence (or new input token) is simultaneously improved.

[0013] Therefore, the described techniques improve the technical field of machine learning. In particular, the described techniques' inclusion of generating a respective set of adjusted attention logits by combining aggregated mask scores with a respective set of attention logits provides a direct means for controlling the internal data processing of the attention mechanism. That is, by dynamically reducing the influence of irrelevant tokens (through the adjusted attention logits), the described techniques can effectively “filter contextual noise” (i.e., avoid attending to irrelevant tokens), which addresses the first issue described above and leads to a more accurate and reliable network output.

[0014] Furthermore, the described techniques provide a further technical effect that improves the functioning of the computer system itself. In particular, the improved control of the attention mechanism described above, in turn, also allows the determined set of aggregated mask scores to serve as a reliable, signal of token relevance. This signal enables a further improvement which is the ability to remove one or more context tokens from an inference context based on the aggregated mask scores. By pruning the least important tokens from the key-value cache, the described techniques directly address the second issue described above by substantially reducing the amount of memory required to store the context. Atorney Docket No. 56113-0812WO1

[0015] This reduction in memory7footprint provides another advantageous effect that addresses the third issue described above. By pruning the key-value cache, the described techniques perform fewer computations during the application of an attention mechanism (e.g., during the query-key dot product and value-weighting steps). This reduces the load on the computer's processing hardware and decreases the required memory bandwidth, resulting in a faster and more efficient computer-implemented process.

[0016] The advantage of a reduced memory footprint is particularly pronounced for hardware accelerators like Graphics Processing Units (GPUs), Tensor Processing Units (TPUs), and ASICs (application-specific integrated circuits). That is, by pruning the keyvalue cache, the described techniques reduce the amount of data that must be maintained by hardware accelerators (e.g.. loaded from off-chip sources), directly decreasing the required memory bandwidth bottleneck for inference on such accelerators. Furthermore, with fewer tokens in the cache of these hardware accelerators, the hardware accelerators perform fewer computations, e.g., during the query -key dot product and value-weighting steps, reducing the load on the accelerator's processing cores and resulting in a faster and more efficient computer-implemented process. Furthermore, in systems / devices with a constrained memory space, such as mobile devices or hardware accelerator units, the techniques described herein can effectively allow a larger context window to be stored within the limited memory space available to the system / device.

[0017] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below7.

[0018] According to a first aspect there is provided a method performed by one or more computers. The method includes receiving an input sequence that includes a respective input token at each of a plurality of input positions. Then, processing the input sequence using a neural network to generate a network output. The neural netw ork includes a plurality of layer blocks that include one or more attention layer blocks that each include a selective attention layer having one or more attention heads. Additionally, processing the input sequence includes, for each selective attention layer, obtaining a respective set of mask scores for each particular input token of the plurality of input tokens that includes a respective mask score for each of the input tokens that measures a degree to which the particular input token is masked by the input token. Then, processing the input sequence includes determining, for each particular input token and from the respective set of mask scores for the particular input token, a respective set of aggregated mask scores. Additionally, processing the input sequence includes, for each attention head of the Atorney Docket No. 56113-0812WO1 selective atention layer, receiving a respective input embedding for each of the input tokens. Then, determining a respective query, key, and value for each of the input tokens from the respective input embeddings for each of the input tokens. Then, determining, for each particular input token and from the query for the particular input token and the respective keys for the input tokens, a respective set of attention logits that includes a respective attention logit for each of the input tokens. Then, combining, for each particular input token, the respective set of attention logits for the particular input token and the respective set of aggregated mask scores for the particular input token to generate a respective set of adjusted attention logits for the particular input token that includes a respective adjusted attention logit for each of the input tokens. Then, generating a respective output embedding for each of the input tokens for the attention head from the respective sets of adjusted attention logits and the respective values for the input tokens.

[0019] In some cases, obtaining a respective set of mask scores for each particular input token of the plurality of input tokens that includes a respective mask score for each of the input tokens that measures a degree to which the particular input token is masked by the input token includes particular steps. These steps include obtaining a respective set of initial mask scores for each particular input token of the plurality of input tokens that includes a respective initial mask score for each of the input tokens. Then, generating the respective sets of mask scores from the respective sets of initial mask scores, which includes, for one or more of the particular input tokens, generating the respective set of mask scores for the particular input token by setting to zero one or more of the initial mask scores in the respective set of initial mask scores for the particular input token.

[0020] In some cases, obtaining a respective set of initial mask scores for each particular input token of the plurality of input tokens that includes a respective initial mask score for each of the input tokens includes determining the respective sets of initial mask scores for the particular input tokens from the respective sets of attention logits determined by one or more of the attention heads.

[0021] In some cases, determining the respective sets of initial mask scores for the particular input tokens from the respective sets of attention logits determined by one or more of the attention heads includes using, as the respective sets of initial mask scores for the particular input tokens, the respective sets of attention logits determined by a designated one of the attention heads.

[0022] In some cases, determining the respective sets of initial mask scores for the particular input tokens from the respective sets of attention logits determined by one or Atorney Docket No. 56113-0812WO1 more of the attention heads includes combining the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the respective sets of initial mask scores for the particular input tokens.

[0023] In some cases, combining the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the respective sets of initial mask scores for the particular input tokens includes averaging the respective sets of attention logits determined by a designated subset of two or more of the attention heads.

[0024] In some cases, generating the respective set of mask scores for the particular input token by setting to zero one or more of the initial mask scores in the respective set of initial mask scores for the particular input token includes setting to zero each initial mask score that has a negative value.

[0025] In some implementations, setting to zero each initial mask score that has a negative value includes applying a rectified linear unit (ReLU) function to the initial mask scores.

[0026] In some cases, generating the respective set of mask scores for the particular input token by setting to zero one or more of the initial mask scores in the respective set of initial mask scores for the particular input token includes setting to zero the initial mask score for the particular token.

[0027] In some cases, a first input token at a first input position in the input sequence is a beginning of sequence token. Additionally, generating the respective set of mask scores for the particular input token by setting to zero one or more of the initial mask scores in the respective set of initial mask scores for the particular input token includes setting to zero the initial mask score for the beginning of sequence token.

[0028] In some implementations, determining, for each particular input token and from the respective set of mask scores for the particular input token, a respective set of aggregated mask scores includes, for each particular input token and for each input token, generating the aggregated mask score by summing the mask scores, in the set of mask scores for the particular input token, for any input tokens that precede the input token in the input sequence.

[0029] In some implementations, combining, for each particular input token, the respective set of attention logits for the particular input token and the respective set of aggregated mask scores for the particular input token to generate a respective set of adjusted attention logits for the particular input token that includes a respective adjusted attention logit for each of the input tokens includes, for each particular input token and for Atorney Docket No. 56113-0812WO1 each input token, subtracting the aggregated mask score for the input token from the attention logit for the input token.

[0030] In some implementations, the processing the input sequence of the method further includes, for each selective attention layer, combining the respective output embeddings for each of the input tokens generated by the attention heads to generate an output of the selective attention layer.

[0031] In some implementations, generating a respective output embedding for each of the input tokens for the attention head from the respective sets of adjusted attention logits and the respective values for the input tokens includes particular steps. These steps include generating, for each particular input token, a respective set of attention weights for the particular input token that includes a respective attention weight for each of the input tokens. Then, for each particular input token, generating the respective output embedding for the particular output token as a weighted sum of the values for the input tokens with each value being weighted by the attention weight for the corresponding input token in the respective set of attention weights for the particular input token.

[0032] In some implementations, the input sequence is a training input sequence for training the neural network, and the network output includes a respective probability distribution over a set of input tokens for each of the input positions. Additionally, the method further includes training the neural network on a loss function based on the respective probability distributions.

[0033] In some implementations, the loss function includes a first term that measures, for one or more of the input positions, a probability assigned to a next input token by the probability distribution for the input position, where the next input token is the input token at an immediately following input position in the input sequence.

[0034] In some implementations, the loss function includes an auxiliary loss term that measures, for each selective attention layer, the sets of aggregated mask scores generated by the selective attention layer for the input tokens.

[0035] In some cases, the auxiliary’ loss term measures an approximate amount of memory required to perform auto-regressive inference to generate at least a portion of the input sequence given that tokens are removed from an inference context for the autoregressive inference based on the sets of aggregated mask scores.

[0036] According to a second aspect there is provided a method performed by one or more computers. The method includes receiving a new input token and processing the new input token using a neural network to generate a network output. The neural network Atorney Docket No. 56113-0812WO1 includes a plurality of layer blocks that include one or more attention layer blocks that each include a selective attention layer having one or more attention heads. Additionally, processing the new input token includes, for each selective attention layer, obtaining a set of mask scores that includes a respective mask score for each token in a set of tokens that includes the new input token and a set of context tokens. The respective mask score measures a degree to which the new token is masked by the token. After obtaining the set of mask scores, the method includes determining, from the mask scores, a set of aggregated mask scores that includes a respective aggregated mask score for each token in the set of tokens. Afterwards, the method includes, for each attention head of the selective attention layer, receiving an input embedding of the new input token. Then, determining a respective query, key. and value for the new input token. Then, obtaining a respective key and value for each of a set of context tokens for the selective attention layer. Then, determining, from the query for the new input token, the key for the new input token, and the respective keys for the context tokens, a set of attention logits that includes a respective attention logit for each of the tokens in the set. Then, combining the set of attention logits and the respective set of aggregated mask scores to generate a set of adjusted attention logits that includes a respective adjusted attention logit for each of the tokens in the set of tokens. Then, generating an output embedding for the new token from the set of adjusted attention logits and the respective values for the new token and the context tokens.

[0037] In some cases of the second aspect, obtaining a set of mask scores that includes a respective mask score for each token in a set of tokens that includes the new input token and a set of context tokens includes particular steps. These particular steps include obtaining a set of initial mask scores that includes a respective initial mask score for each of the tokens in the set. Then, generating the set of mask scores from the set of initial mask scores, which includes setting to zero one or more of the initial mask scores in the set of initial mask scores.

[0038] In some cases of the second aspect, obtaining a set of initial mask scores that includes a respective initial mask score for each of the tokens in the set includes determining the set of initial mask scores from the respective sets of attention logits determined by one or more of the attention heads.

[0039] In some cases of the second aspect, determining the set of initial mask scores from the respective sets of attention logits determined by one or more of the attention heads Atorney Docket No. 56113-0812WO1 includes using, as the set of initial mask scores, the respective set of attention logits determined by a designated one of the attention heads.

[0040] In some cases of the second aspect, determining the set of initial mask scores from the respective sets of attention logits determined by one or more of the attention heads includes combining the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the set of initial mask scores.

[0041] In some cases of the second aspect, combining the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the set of initial mask scores includes averaging the respective sets of attention logits determined by a designated subset of two or more of the attention heads.

[0042] In some cases of the second aspect, setting to zero one or more of the initial mask scores in the set of initial mask scores includes setting to zero each initial mask score that has a negative value.

[0043] In some cases of the second aspect, setting to zero each initial mask score that has a negative value includes applying a rectified linear unit (ReLU) function to the initial mask scores.

[0044] In some cases of the second aspect, setting to zero one or more of the initial mask scores in the set of initial mask scores includes setting to zero the initial mask score for the new token.

[0045] In some cases of the second aspect, a first input token of the context tokens is a beginning of sequence token. Additionally, setting to zero one or more of the initial mask scores in the set of initial mask scores includes setting to zero the initial mask score for the beginning of sequence token.

[0046] In some cases of the second aspect, determining a set of aggregated mask scores includes, for each token in the set, generating the aggregated mask score by summing the mask scores for any tokens that precede the token in an ordering of the tokens in the set.

[0047] In some cases of the second aspect, combining the set of attention logits and the set of aggregated mask scores includes, for each token in the set, subtracting the aggregated mask score for the token from the attention logit for the token.

[0048] In some cases of the second aspect, the method further includes updating the set of context tokens by adding the new token to the set of context tokens.

[0049] In some cases of the second aspect, the method further includes, determining that a number of tokens in the set of context tokens exceeds a maximum memory size for the Atorney Docket No. 56113-0812WO1 atention layer and. in response, removing one or more of the context tokens from the set based on the aggregated mask scores for the context tokens.

[0050] In some cases of the second aspect, removing one or more of the context tokens from the set based on the aggregated mask scores for the context tokens includes removing one or more of the context tokens having the highest aggregated mask scores.

[0051] In some cases of the second aspect, obtaining a respective key and value for each of a set of context tokens for the selective atention layer includes obtaining the respective keys and values for the context tokens from a memory of the one or more computers. Additionally, removing one or more of the context tokens from the set based on the aggregated mask scores for the context tokens includes removing the respective keys and values for the one or more context tokens from the memory.

[0052] In some cases of the second aspect, the plurality of layer blocks include a plurality of atention layer blocks that each include a respective selective atention layer.

[0053] In some cases of the second aspect, two or more of the selective atention layers have different maximum memory sizes.

[0054] In some cases of the second aspect, the processing the input sequence of the method further includes, for each selective atention layer, combining the respective output embeddings for the new token generated by the attention heads to generate an output of the selective attention layer.

[0055] In some cases of the second aspect, generating the output embedding for the new token for the atention head includes generating, from the adjusted atention logits, a set of atention weights that includes a respective atention weight for each of the tokens in the set. Then, generating the output embedding as a weighted sum of the values for the tokens in the set with each value being weighted by the atention weight for the corresponding token in the set of atention weights.

[0056] In some cases of the second aspect, the network output includes a respective score for each of a vocabulary of tokens. Additionally, the method further includes selecting a token to follow the new token in an output sequence using the network output.

[0057] According to a third aspect, there is provided the methods of the first aspect or second aspect performed by a system that includes one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one more computers to perform the operations of the respective method. Atorney Docket No. 56113-0812WO1

[0058] According to a fourth aspect, there is provided the methods of the first aspect or second aspect performed by one or more computer storage media storing instructions that when executed by one or more computers cause the one more computers to perform the operations of the respective method.

[0059] Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.

[0060] BRIEF DESCRIPTION OF THE DRAWINGS

[0061] FIG. 1 shows a selective attention system.

[0062] FIG. 2A is a flow diagram of an example process for processing an input sequence with a neural network.

[0063] FIG. 2B is a flow' diagram of an example process for applying a selective attention mechanism within a selective attention layer.

[0064] FIG. 3A is a flow diagram of an example process for processing a new input token during autoregressive inference.

[0065] FIG. 3B is a flow diagram of an example process for applying a selective attention mechanism during a single step of autoregressive inference.

[0066] FIG. 4 is a flow' diagram of an example process for training a neural network.

[0067] FIG. 5 is an example of the performance of the described techniques.

[0068] FIG. 6 is an example of the performance of the described techniques.

[0069] FIG. 7 is an example of the performance of the described techniques.

[0070] FIG. 8 is an example of the performance of the described techniques.

[0071] FIG. 9 is an example of the performance of the described techniques.

[0072] FIG. 10 is an example visualization of the selective attention mechanism.

[0073] Like reference numbers and designations in the various drawings indicate like elements.

[0074] DETAILED DESCRIPTION

[0075] FIG. 1 shows a selective attention system 100. The 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. Atorney Docket No. 56113-0812WO1

[0076] The system 100 processes input sequences 102 to perform one or more machine learning tasks. In particular, the system 100 processes an input sequence 102 using a neural network 104 to generate a network output 110.

[0077] The neural network 104 can be configured through training to perform any kind of machine learning task, i.e., can be configured to receive any kind of input sequence 102 and to generate any kind of score, classification, or regression output 110 based on the input sequence 102.

[0078] In some situations, after training, the neural network 104 can be referred to as an auto-regressive neural network, i.e., because the neural network 104 auto-regressively generates an output 110 that is an output sequence of tokens. More specifically, the auto- regressively generated output 110 is created by generating each particular token in the output sequence conditioned on a current input sequence that includes any tokens that precede the particular token in the output sequence, i.e., the tokens that have already been generated for any previous positions in the output sequence that precede the particular position of the particular token.

[0079] For example, the neural network 104 can be an auto-regressive attention neural network that includes (i) a plurality of attention layer blocks that each apply a selfattention operation and (ii) an output subnetw ork that processes an output of the last attention block to generate the score distribution.

[0080] In this example, the neural network can have any of a variety of Transformerbased neural network architectures. Examples of such architectures include those described in J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas. L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute- optimal large language models, arXiv preprint arXiv:2203. 15556, 2022; J.W. Rae, S. Borgeaud, T. Cai, K. Millican, J. Hoffmann, H. F. Song, J. Aslanides, S. Henderson, R. Ring, S. Young, E. Rutherford, T. Hennigan, J. Menick, A. Cassirer, R. Powell, G. van den Driessche, L. A. Hendricks, M. Rauh, P. Huang, A. Glaese, J. Welbl, S. Dathathri, S. Huang, J. Uesato, J. Mellor, I. Higgins, A. Creswell, N. McAleese, A.Wu, E. Eisen, S. M. Jayakumar. E. Buchatskaya, D. Budden. E. Sutherland, K. Simonyan. M. Paganini. L. Sifre, L. Martens, X. L. Li, A. Kuncoro, A. Nematzadeh, E. Gribovskaya, D. Donato, A. Lazaridou, A. Mensch, J. Lespiau, M. Tsimpoukelli, N. Grigorev, D. Fritz, T. Sottiaux, M. Pajarskas, T. Pohlen, Z. Gong, D. Toyama, C. de Masson d'Autume, Y. Li, T. Terzi, V. Mikulik, I. Babuschkin. A. Clark, D. de Las Casas. A. Guy, C. Jones, J. Bradbury, M. Johnson, B. A. Hechtman, L. Weidinger, I. Gabriel, W. S. Isaac, E. Lockhart, S. Atorney Docket No. 56113-0812WO1

[0081] Osindero, L. Rimell, C. Dyer, O. Vinyals, K. Ayoub, J. Stanway, L. Bennett, D. Hassabis, K. Kavukcuoglu. and G. Irving. Scaling language models: Methods, analysis & insights from training gopher. CoRR, abs / 2112.1 1446, 2021; Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 1910. 10683, 2019; Daniel Adiwardana, Minh-Thang Luong, David R. So, Jamie Hall, Noah Fiedel, Romal Thoppilan, Zi Yang, Apoorv Kulshreshtha. Gaurav Nemade, Yifeng Lu, and Quoc V. Le. Towards a human-like open-domain chatbot. CoRR, abs / 2001.09977, 2020; Team, G., Mesnard, T., Hardin, C., Dadashi, R., Bhupatiraju, S., Pathak, S., Sifre, L., Riviere, M., Kale, M.S., Love. J. and Tafti, P., 2024. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08295; and Team, G., Anil, R., Borgeaud, S., Alayrac, J.B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A.M., Hauth, A., Millican, K. and Silver, D., 2023. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312. 11805.

[0082] More specifically, the neural network 104 includes plurality of layer blocks. A layer block, as used in this specification, is a collection of one or more neural network layers.

[0083] The plurality of layer blocks includes one or more attention layer blocks 106 that each include a selective attention layer 108 and any number of attention layer blocks that do not include a selective attention layer 108. In some cases, the plurality of layer blocks includes a plurality of attention layer blocks 106 that each include a selective attention layer 108 and any number of attention layer blocks that do not include a selective attention layer 108. In some other cases, within the plurality of layer blocks, all of the attention layer blocks include a selective attention layer 108.

[0084] During training (and during inference when inference is not auto-regressive), the neural network 104 receives an input sequence 102 that includes a respective input token at each of a plurality of input positions.

[0085] The neural network 104 processes the input sequence 102 to generate a network output 110.

[0086] During auto-regressive sequence generation after training, the system 100 auto- regressively generates an output 110 that is an output sequence that includes a respective output token at each of plurality of output positions.

[0087] To do this, for each output position (after the very first output position in the sequence), the system 100 receives a preceding output token ('‘new input token”) Atorney Docket No. 56113-0812WO1 generated for the preceding output position and processes the preceding output token using the neural network 104 to generate the output token at the output position. This processing is conditioned on a set of context tokens, which includes all tokens that precede the new input token in the sequence (e.g., all tokens from an original input prompt as well as previously generated output tokens), to generate the output token for the current output position.

[0088] During this processing, a given attention head of a given selective attention layer 108 receives an input embedding 112 for the new input token and generates as output an output embedding 122 for the new input token, i.e., to update the embedding for the new input token.

[0089] An embedding is an ordered collection of numerical values, e.g., a vector of floating point or other numerical values. For example, an ^-dimensional vector of numbers, where n can be any positive integer.

[0090] Generally, each attention mechanism uses one or more attention heads.

[0091] Each attention head generates a set of queries, a set of keys, and a set of values, and then applies any of a variety of variants of query -key -value (QKV) attention, e.g., a dot product attention function or a scaled dot product attention function, using the queries, keys, and values to generate an output. Each query, key, value can be a vector that includes one or more vector elements.

[0092] When there are multiple attention heads, the attention layer block 106 then combines the outputs of the multiple attention heads, e.g., by concatenating the outputs and, optionally, processing the concatenated outputs through a linear layer. The attention mechanisms can be local or global (or some attention mechanisms can be local while others for other layer blocks are global). For local attention mechanisms, for each position, the positions that are used to generate the queries, keys, and values for the position are defined by the local window size for the local attention mechanism, i.e., nonzero attention weights for a given position are computed only for positions that are within the local window of the given position. For global attention mechanisms, in contrast, there is no local window, and non-zero attention weights for a given position can be computed for all other positions in the input sequence.

[0093] In some cases, because the attention applied by the attention layer blocks 106 is causal, the system 100 can store, in memory and for any given attention head and when generating the output for any given input position, the embeddings or the keys and values already computed for earlier input positions (i.e., for the “context’’ tokens that precede the Atorney Docket No. 56113-0812WO1 current token in the sequence) rather than re-computing the embeddings (or the keys and values) for earlier time steps. Storing the keys and values in a memory is also referred to as maintaining a KV cache.

[0094] Some or all of the layer blocks in the neural network 104 can also include other ty pes of lay ers, e.g., normalization layers, residual connection layers, feedforward layers, and so on. In some cases, some or all of the feedforward subnetworks within the layer blocks in the neural network 104 are implemented as sparse mixture of experts (MoE) layers while in other cases all the feedforward layers are dense feedforward layers.

[0095] Generally, the selective attention layers 108 modify the standard self-attention operation to improve the performance of the neural netw ork 104. Within each selective attention layer 108, each attention head performs this modification by first using a set of mask scores 114 to determine a set of aggregated mask scores 116. The attention head then combines these aggregated mask scores 116 with its attention logits 118 to generate a final set of adjusted attention logits 120. This process reduces attention to unneeded elements, improving the overall performance of the neural network 104.

[0096] In some implementations, the system 100 can use the sets of aggregated mask scores 116 to remove certain tokens from the context of the neural network 104, e. g. , by removing the stored keys and values from memory7(e.g., the KV cache), during the generation of any given output sequence. This makes the inference process of the neural network 104 significantly more memory efficient, i.e., allows the system 100 to generate output sequences using significantly less memory. By using the aggregated mask scores 116 to identify and remove only the most unneeded tokens, the system 100 achieves significant memory savings without degrading the performance of the neural network 104.

[0097] Some examples of machine learning tasks that a neural network when implemented using one of the architectures described above or other known architectures can be configured to perform follow.

[0098] In any of the implementations below, the neural network may be deployed as part of a chat bot. dialogue agent, or other software tool that receives inputs from users and provides outputs in response to the received input, e.g., as part of a conversation or dialogue. In these implementations, the input sequences received by the neural network are (generated from) user inputs and the output sequences generated by the neural network can be used to generate responses to the user inputs. Atorney Docket No. 56113-0812WO1

[0099] In implementations the neural network may be configured as, or include, a generative (large) language model or a multi-modal model, e.g., a visual and language model, to perform these example machine learning tasks.

[0100] In some cases, the neural network is a neural network that is configured to perform an image processing task, i.e., receive an input image and to process the input image to generate a network output for the input image. For example, the input sequence may comprise tokens representing pixel values for pixels in regions or patches of the image. For example, the task may be image classification and the output generated by the neural network for a given image may be scores for each of a set of object categories, with each score representing an estimated likelihood that the image contains an image of an object belonging to the category. As another example, the task can be image embedding generation and the output generated by the neural network can be a numeric embedding of the input image. As yet another example, the task can be object detection and the output generated by the neural network can identify locations in the input image at which particular types of objects are depicted. As yet another example, the task can be image segmentation and the output generated by the neural network can assign each pixel of the input image to a category from a set of categories. In some other cases, the neural netw ork is a neural network that is configured to perform an image generation task, where the input is a conditioning input and the output is a sequence of intensity value inputs for the pixels of an image that has been generated conditioned on the input image, e.g., with the content and / or style of the input image.

[0101] 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 w ords, 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. The vocabulary for the input tokens may be words, w ordpieces or characters of the first language, and the vocabulary for the output tokens may be words, wordpieces or characters of the other language. 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. Atorney Docket No. 56113-0812WO1

[0102] Some implementations may be used for automatic code generation. For example the input tokens may represent words, wordpieces or characters in a first natural language and the output tokens may represent instructions in a computer programming or markup language, or instructions for controlling an application program to perform a task, e.g., build a data item such as an image or w eb page.

[0103] 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 netw ork 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, e g. a speech to text task. 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 be a classification of the spoken utterance into one of a plurality of categories, for example an identity of the natural language in which the utterance was spoken.

[0104] 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.

[0105] 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.

[0106] 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. Such electronic health data may, for example, comprise one or more sequences of physiological data taken from a patient, with the output being a corresponding prediction that relates to those sequences of data. Examples of physiological data and a corresponding prediction include: blood glucose measurements, with the prediction being a predicted future blood glucose measurement or the prediction of a hyper- or hypo-glycemic event; a heart rate, with the Atorney Docket No. 56113-0812WO1 prediction being the presence or absence of a heart condition, or a future cardiac event; blood pressure measurements, with the prediction being the risk of a future heart condition; or the like.

[0107] 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.

[0108] In some implementations the input sequence represents data to be compressed, e.g., image data, text data, audio data, or any other type of data; and the output sequence a compressed version of the data. The input and output tokens may each comprise any representation of the data to be compressed / compressed data, e.g., symbols or embeddings generated / decoded by a respective neural network. In some complementary implementations the input sequence represents compressed data and the output sequence represents a decompressed version of the data, e.g., image data, text data, audio data, or any other type of data.

[0109] 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. The observations may comprise sensor data captured by sensors associated with (e.g.. part of) the agent, for example visual data, LIDAR data, sonar data, agent configuration data (e.g.. joint angles), agent orientation data, or the like.

[0110] In some implementations, the environment is a real-world environment, the agent is a mechanical (or electro-mechanical) agent interacting with the real-world environment, e g., a robot or an autonomous or semi-autonomous land, air, or sea vehicle operating in or navigating through the environment, and the actions are actions taken by the mechanical agent in the real-world environment to perform the task. For example, the agent may be a robot interacting with the environment to accomplish a specific task, e.g., to locate or manipulate an object of interest in the environment or to move an object of interest to a specified location in the environment or to navigate to a specified destination in the environment. Atorney Docket No. 56113-0812WO1

[0111] In these implementations, the observations may include, e.g., one or more of: images, object position data, and sensor data to capture observations as the agent interacts with the environment, for example sensor data from an image, distance, or position sensor or from an actuator. For example in the case of a robot, the observations may include data characterizing the current state of the robot, e.g., one or more of: joint positionjoint velocity, joint force, torque or acceleration, e.g.. gravity-compensated torque feedback, and global or relative pose of an item held by the robot. In the case of a robot or other mechanical agent or vehicle the observations may similarly include one or more of the position, linear or angular velocity7, force, torque or acceleration, and global or relative pose of one or more parts of the agent. The observations may be defined in 1 , 2 or 3 dimensions, and may be absolute and / or relative observations. The observations may also include, for example, sensed electronic signals such as motor current or a temperature signal; and / or image or video data for example captured by a camera or a LIDAR sensor, e.g., data from sensors of the agent or data from sensors that are located separately from the agent in the environment.

[0112] In these implementations, the actions may be control signals to control the robot or other mechanical agent, e.g., torques for the joints of the robot or higher-level control commands, or the autonomous or semi-autonomous land, air, sea vehicle, e.g., torques to the control surface or other control elements, e.g., steering control elements of the vehicle, or higher-level control commands. The control signals can include for example, position, velocity, or force / torque / accel eration data for one or more joints of a robot or parts of another mechanical agent. The control signals may also or instead include electronic control data such as motor control data, or more generally data for controlling one or more electronic devices within the environment the control of which has an effect on the observed state of the environment. For example in the case of an autonomous or semi-autonomous land or air or sea vehicle the control signals may define actions to control navigation, e.g., steering, and movement e.g., braking and / or acceleration of the vehicle.

[0113] In some implementations the environment is a simulation of the above-described real-world environment, and the agent is implemented as one or more computers interacting with the simulated environment. For example, a system implementing the neural network may be used to select actions in the simulated environment during training or evaluation of the system and, after training, or evaluation, or both, are complete, the action selection policy may be deployed for controlling a real-world agent in the Atorney Docket No. 56113-0812WO1 particular real-world environment that was the subject of the simulation. This can avoid unnecessary wear and tear on and damage to the real-world environment or real-world agent and can allow the control neural network to be trained and evaluated on situations that occur rarely or are difficult or unsafe to re-create in the real-world environment. For example the system may be partly trained using a simulation of a mechanical agent in a simulation of a particular real-world environment, and afterwards deployed to control the real mechanical agent in the particular real-world environment. Thus in such cases the observations of the simulated environment relate to the real-world environment, and the selected actions in the simulated environment relate to actions to be performed by the mechanical agent in the real-world environment.

[0114] In some implementations, as described above, the agent may not include a human being (e.g., it is a robot). Conversely, in some implementations the agent comprises a human user of a digital assistant such as a smart speaker, smart display, or other device. Then the information defining the task can be obtained from the digital assistant, and the digital assistant can be used to instruct the user based on the task.

[0115] For example, a system implementing the neural network may output to the human user, via the digital assistant, instructions for actions for the user to perform at each of a plurality of time steps. The instructions may for example be generated in the form of natural language (transmitted as sound and / or text on a screen) based on actions chosen by the system. The system chooses the actions such that they contribute to performing a task. A monitoring system (e.g., a video camera system) may be provided for monitoring the action (if any) which the user actually performs at each time step, in case (e.g., due to human error) it is different from the action which the system instructed the user to perform. Using the monitoring system the system can determine whether the task has been completed. The system may identify actions which the user performs incorrectly with more than a certain probability. If so, when the system instructs the user to perform such an identified action, the system may warn the user to be careful. Alternatively or additionally, the system may learn not to instruct the user to perform the identified actions, i.e., ones which the user is likely to perform incorrectly.

[0116] More generally, the digital assistant instructing the user may comprise receiving, at the digital assistant, a request from the user for assistance and determining, in response to the request, a series of tasks for the user to perform, e.g., steps or sub-tasks of an overall task. Then for one or more tasks of the series of tasks, e.g., for each task, e.g., until a final task of the series the digital assistant can be used to output to the user an Atorney Docket No. 56113-0812WO1 indication of the task, e.g., step or sub-task, to be performed. This may be done using natural language, e.g., on a display and / or using a speech synthesis subsystem of the digital assistant. Visual, e.g., video, and / or audio observations of the user performing the task may be captured, e.g., using the digital assistant. A system as described above may then be used to determine whether the user has successfully achieved the task, e.g., step or sub-task, i.e., from the answer as previously described. If there are further tasks to be completed the digital assistant may then, in response, progress to the next task (if any) of the series of tasks, e.g., by outputting an indication of the next task to be performed. In this way the user may be led step-by-step through a series of tasks to perform an overall task. During the training of the neural network, training rewards may be generated, e.g., from video data representing examples of the overall task (if corpuses of such data are available) or from a simulation of the overall task.

[0117] In a further aspect there is provided a digital assistant device including a system as described above. The digital assistant can also include a user interface to enable a user to request assistance and to output information. In implementations this is a natural language user interface and may comprise a keyboard, voice input-output subsystem, and / or a display. The digital assistant can further include an assistance subsystem configured to determine, in response to the request, a series of tasks for the user to perform. In implementations this may comprise a generative (large) language model, in particular for dialog, e.g., a conversation agent such as Sparrow (Glaese et al. arXiv:2209. 14375) or Chinchilla (Hoffmann et al. arXiv:2203. 15556). The digital assistant can have an observation capture subsystem to capture visual and / or audio observations of the user performing a task; and an interface for the above-described language model neural network (which may be implemented locally or remotely). The digital assistant can also have an assistance control subsystem configured to assist the user. The assistance control subsystem can be configured to perform the steps described above, for one or more tasks, e.g., of a series of tasks, e.g., until a final task of the series. More particularly the assistance control subsystem and output to the user an indication of the task to be performed, capture, using the observation capture subsystem, visual or audio observations of the user performing the task, determine from the above-described answer whether the user has successfully achieved the task. In response the digital assistant can progress to a next task of the series of tasks and / or control the digital assistant, e.g., to stop capturing observations. Atorney Docket No. 56113-0812WO1

[0118] 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 dow nstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, and so on.

[0119] 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.

[0120] In some cases, the machine learning task is a multi-modal processing task that requires processing multi-modal data. In general, multi-modal data is a combination of two or more different types of data, e.g., two or more of audio data, image data, text data, or graph data. As one example the multi-modal data may comprise audio-visual data, comprising a combination of pixels of an image or of video and audio data representing values of a digitized audio waveform. As another example the multi-modal data may comprise a combination of i) text data representing text in a natural language and ii) pixels of an image or of video or audio data representing values of an audio waveform. Optionally, but not necessarily, the different types of data may represent the same or overlapping objects using the different modalities (types), and when processing multimodal data the data may be mapped into a common embedding space.

[0121] As a particular example, the task is a multi-modal processing task that requires processing both text and image inputs, so that the neural network includes both a computer vision neural network and a text processing neural network. That is, the target output to be generated by the computer vision neural network for a given image depends on one or more outputs generated by the text processing neural network for one or more corresponding text inputs (and vice versa). Examples of such tasks include openvocabulary image classification, open-vocabulary object detection, image captioning, text-based image search, image-based retrieval, and so on. Atorney Docket No. 56113-0812WO1

[0122] As some further examples a multi-modal processing task can involve processing a text input comprising a sequence of text or audio data representing values of an audio waveform, e.g., instantaneous amplitude data or time-frequency domain data, or an image or video (or encoded versions of these inputs) to generate the network output. The network output may comprise any form of output appropriate to the task performed. For example the network output may comprise text in a natural or computer language that defines a result of the task, e.g., for tasks such as image captioning, video or audio question answering (answering a natural language question about a visual or audio input), or object detection or instance segmentation. For example in a video or audio question answering task the question can define an information content extraction task, to extract information from the content of the video or audio, or the question can define a reasoning task such as a predictive reasoning task (e.g. “what would happen next?”), a counterfactual reasoning task (e.g. “what would happen if..?”), or a causal reasoning task (e.g. “why did X happen?”). The network output can provide an answer in any convenient form. e g. tokens representing natural language. An input to the system may be obtained from a sensor sensing the real world, e.g. a a condition or characteristic of the real world. For example the video or audio may be captured from the real-world. The network output can then provide an answ er, e.g. in natural language, to a question asked about the real- world input.

[0123] Also or instead the network output may comprise data defining an image, video or audio object, e.g., as specified by the input (e.g. by a natural language description of one or more characteristics of the object), e.g., in a generative task. As a further alternative the network output may comprise non-textual action selection data for selecting an action to be performed by an agent controlled by the network output, e g. as described above, e.g. in response to an input that includes a natural language description of a physical or other task to be performed by the agent. As another example the network output may also or instead define an intermediate step to be performed during the task, e.g., a call to a software API for a software tool that is used when performing the task; the input may then receive an output from the software tool that is used to generate a final network output that performs the task.

[0124] More generally, the multi-modal processing task may correspond to any of the tasks previously described for any of the types of data making up the multi-modal combination. For example, an accuracy of the previously described tasks may be increased when the task is applied to multi-modal data combining the data for which the Atorney Docket No. 56113-0812WO1 task has been previously described and another type of data. For example detection or classification of an object or event may be improved when data of multiple different types (modalities) is processed.

[0125] More generally, the task to be performed by the neural network can be specified by the input sequence. As a particular example, the input sequence can include a prompt or an instruction that specifies the task that is to be performed by the neural network. Optionally, in this example, the input sequence also includes context for performing the task.

[0126] In general, in implementations of the described techniques the input data, e.g., text, audio, and / or an image or video, may be encoded into a sequence of input tokens in any convenient manner; and output tokens may be similarly decoded into text, audio, and / or image or video data according to the particular task or tasks to be performed.

[0127] FIG. 2A is a flow diagram of an example process 200 for processing an input sequence with a neural network. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a selective attention system, e.g., the selective attention system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 200.

[0128] The system receives an input sequence that includes a respective input token at each of a plurality of input positions (step 202).

[0129] Generally, the input sequence can represent any type of data (e g., text, image, audio, video, or any combination of these). For example, the input sequence can be the result of the system tokenizing text, image, audio, video, or any combination of these.

[0130] As a particular example, for text, the system can use a text tokenizer to partition the text into a sequence of word or sub-word tokens. For example, the system can apply the Byte-Pair Encoding (BPE), WordPiece, or SentencePiece tokenizers to divide the natural language text data into tokens from a vocabulary.

[0131] As another particular example, for images, the system can partition the image into a grid of fixed-size patches, e.g.. 16x16 pixel patches. Each patch is then treated as a single token and ordered into a sequence, for example, according to a raster scan order.

[0132] As another particular example, for audio, the system can first convert a raw audio waveform into a time-frequency representation, e.g., a spectrogram. The system can then partition the spectrogram into a sequence of frames, where each frame is treated as a token. Atorney Docket No. 56113-0812WO1

[0133] As another particular example, for video, the system can sample frames from the video at a particular rate. Each sampled frame can then be processed as a token. For instance, the system can partition each sampled frame into a grid of patches, with the final input sequence being a flattened sequence of all patches from all sampled frames, ordered temporally.

[0134] After the system tokenizes the input data, the system generates a sequence of input embeddings by mapping each token to, e.g.. a high-dimensional numerical vector. The system generally performs this mapping using a trainable embedding layer(s) of the neural network. In some implementations, the system also adds a positional encoding to each input embedding to provide the neural network with information about the position of the corresponding token in the sequence.

[0135] The system can receive the input sequence from a user, another system, or system maintained data.

[0136] For example, the system can receive an input sequence from a user as part of an automatic code generation task, where the input sequence is a natural language text prompt specifying a desired computer programming function, e.g., "write a Python function to sort a list."

[0137] As another example, the system can receive an input sequence from another system, e.g., a sensor system, as part of an agent control task, where the input sequence is a sequence of observations characterizing a state of an environment as part of an agent control task.

[0138] As another example, the system can receive an input sequence from system- maintained data, e.g., by retrieving a sequence derived from historical operation data for a machine from a database as part of a machine failure prediction task.

[0139] The system processes the input sequence using a neural network to generate a network output (step 204).

[0140] As described above, the neural network can have any of a variety of neural network architectures. That is, the neural network can have any appropriate architecture in any appropriate configuration that can process an input sequence to generate a network output, including fully connected layers, convolutional layers, recurrent layers, attentionbased layers, and so on, as is appropriate.

[0141] Generally, however, the neural network includes one or more attention layer blocks that each include a selective attention layer. Atorney Docket No. 56113-0812WO1

[0142] In some implementations, the neural network is configured to operate autoregressively. where it processes a cunent input sequence to generate a single next token for an output sequence. In other implementations, the neural network is configured to process the entire input sequence in a single forward pass to generate a corresponding network output.

[0143] The network output can be any kind of score, classification, or regression output based on the input sequence.

[0144] For example, as described above, the network output can be a sequence of output tokens, and this sequence of output tokens can represent data of any modality' (e.g., text, image, audio, video, or any combination of these) and any type. The output tokens could represent, e.g., text in a natural or computer language, a set of scores for a classification task, or data defining an action to be performed by an agent.

[0145] As a particular example, for the automatic code generation task described above, the network output can be a sequence of output tokens representing instructions in a computer programming language, e.g., the Python code for the requested sorting function.

[0146] As another particular example, for the agent control task described above, the network output can be a sequence of output tokens that defines an action to be performed by the agent in response to the sequence of observations, e.g., control signals defining torques for the joints of a robot.

[0147] As another particular example, for the machine failure prediction task described above, the network output can be a sequence of output tokens the represents a prediction relevant to the operability of the machine, e.g., a predicted likelihood that the machine will cease to function.

[0148] The system processes the input sequence by passing it sequentially through the plurality of layer blocks of the neural network. The system generates a sequence of input embeddings for each of the input tokens from the input sequence (e.g., by mapping each input token from the input sequence to a corresponding numerical vector, as described above) and provides this sequence as input to the first layer block. For each subsequent layer block, the system provides, as the input embeddings, the sequence of output embeddings generated by the immediately preceding layer block. In this way, the system uses each layer block to process a sequence of input embeddings to generate a sequence of output embeddings, and then uses the sequence of output embeddings from the final layer block to generate the network output. Atorney Docket No. 56113-0812WO1

[0149] Prior to using the neural network to perform machine learning tasks, the system (or another system) pre-trains the neural network, pre-trains and fine-tunes the neural network or fine-tunes a neural network that has already been pre-trained by another system.

[0150] For example, in some implementations, the input sequence is a training input sequence for training the neural network and the network output includes a respective probability distribution over a set of input tokens for one or more of the input positions. That is, the network output is a sequence of probability distributions, where for one or more positions in the input sequence, the corresponding distribution represents the neural network's prediction for the token at that position given all of the preceding tokens. For such cases, the system can train the neural network on a loss function based on the respective probability distributions.

[0151] Further details of training the neural network are described below with reference to FIG. 4.

[0152] FIG. 2A provides an overview of how the system processes an input sequence using a neural network that contains one or more selective attention layers. FIG. 2B illustrates, in detail, the process of step 204 for each selective attention layer and, within each of those layers, for each respective attention head of the selective attention layer.

[0153] In particular, FIG. 2B is a flow diagram of an example process 250 for applying a selective attention mechanism within a selective attention layer. For convenience, the process 250 will be described as being performed by a system of one or more computers located in one or more locations. For example, a selective attention system, e.g., the selective attention system 100 of FIG. 1. appropriately programmed in accordance with this specification, can perform the process 250.

[0154] The system, for each selective attention layer, performs steps 252-254.

[0155] The system obtains a respective set of mask scores for each particular input token of the plurality of input tokens (step 252). The respective set of mask scores includes a respective mask score for each of the input tokens that measures a degree to which the particular input token is masked by the input token. For the input token to “mask7’ the particular token means that the input token reduces the attention the particular token has in the attention mechanism it participates in.

[0156] In some implementations, as part of step 252, the system obtains a respective set of initial mask scores for each particular input token of the plurality of input tokens that includes a respective initial mask score for each of the input tokens that precedes the Atorney Docket No. 56113-0812WO1 particular input token. For example, the system can determine the respective sets of initial mask scores from respective sets of attention logits generated by one or more attention heads of the selective attention layer (as will be described in further detail below). Then, the system generates the respective sets of mask scores from the respective sets of initial mask scores.

[0157] For example, the system, for one or more of the particular input tokens, generates the respective set of mask scores for the particular input token by setting to zero one or more of the initial mask scores in the respective set of initial mask scores for the particular input token. In other words, the system first obtains a set of initial scores and then refines the set by zeroing out certain initial mask scores to produce the final set of mask scores.

[0158] In some cases, the system sets to zero each initial mask score that has a negative value. For example, the system can set to zero each initial mask score that has a negative value by applying a rectified linear unit (ReLU) function to the initial mask scores. The RELU function is a function that outputs the input value if the input value is positive, and outputs zero otherwise.

[0159] In some cases, the system sets to zero the initial mask score for the particular token. In other words, the system sets to zero the initial mask score that a particular input token would apply to itself so that the system can ensure that a particular input token does not mask itself.

[0160] In some implementations, the input sequence includes a “beginning of sequence token” (e.g., <BOS>) that is a first input token at a first input position in the input sequence, and the system sets to zero the initial mask score for the “beginning of sequence token”. In other words, the system can ensure that no input token masks “the beginning of sequence token”. It is advantageous not to mask the “beginning of sequence token” because the neural network can learn to use that token for special operations. For example, the neural network can leam to use the “beginning of sequence token” as a global context, allowing the neural network to store or access sequence-wide information that should always be available.

[0161] The system determines, for each particular input token and from the respective set of mask scores for the particular input token, a respective set of aggregated mask scores (step 254). Because each mask score represents a reduction of attention for the particular token, the aggregated mask score for the particular input token represents the cumulative Atorney Docket No. 56113-0812WO1 degree to which that token has been marked for reduced attention by tokens in the sequence.

[0162] In some implementations, as part of step 254, for each input token, the system generates the aggregated mask score by summing the mask scores, in the set of mask scores for the particular input token, for any input tokens that precede the input token in the input sequence. In other words, the aggregated mask score for the particular token is the running total of all the individual mask scores of the preceding tokens to the input token.

[0163] The system then performs, for each attention head of the selective attention layer, steps 256-264.

[0164] The system receives a respective input embedding for each of the input tokens (step 256). As described above, an input embedding is an ordered collection of numerical values, e.g., an n-dimensional vector of numbers, where n can be any positive integer, that serves as a numerical representation of a corresponding token.

[0165] As an example, for a first layer in the neural network, the system generates the input embedding from the initial tokenized input sequence (e.g., by mapping each token to a corresponding token embedding and adding a positional encoding to each token embedding). For any subsequent layer, the input embedding the system receives is the output embedding generated by the immediately preceding layer in the neural network.

[0166] The system determines a respective query, key, and value for each of the input tokens from the respective input embeddings for each of the input tokens (step 258).

[0167] The system determines, for each particular input token and from the query for the particular input token and the respective keys for the input tokens, a respective set of attention logits that includes a respective attention logit for each of the input tokens (step 260).

[0168] For example, the system can determine the set of attention logits using a scaled dot-product attention mechanism. This mechanism involves computing a dot product between the query of the particular input token and the key of each respective input token in the sequence. The resulting dot product scores represent the similarity between the tokens and are the attention logits.

[0169] In some implementations, the system uses the respective sets of attention logits to determine the respective sets of initial mask scores described above. That is, to obtain a respective set of initial mask scores for each particular input token of the plurality of input tokens that includes a respective initial mask score for each of the input tokens, the Atorney Docket No. 56113-0812WO1 system can determine the respective sets of initial mask scores for the particular input tokens from the respective sets of attention logits determined by one or more of the attention heads.

[0170] For example, in some cases, the system uses, as the respective sets of initial mask scores for the particular input tokens, the respective sets of attention logits determined by a designated one of the attention heads. For example, the system, a user, or another system designate the attention head, e.g., a user input received from a user can designate the attention head or the system can randomly select the attention head from the set of attention heads.

[0171] As another example, in some cases, the system can combine the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the respective sets of initial mask scores for the particular input tokens. For example, the system can average the respective sets of attention logits determined by a designated subset of two or more of the attention heads. The subset of attention heads to be combined can be designated in various ways, such as through random selection or through a pre-configured system setting or a user input.

[0172] The system combines, for each particular input token, the respective set of attention logits for the particular input token and the respective set of aggregated mask scores for the particular input token to generate a respective set of adjusted attention logits for the particular input token that includes a respective adjusted attention logit for each of the input tokens (step 262). For example, the system can, for each input token, subtract the aggregated mask score for the input token from the attention logit for the input token.

[0173] The system generates a respective output embedding for each of the input tokens for the attention head from the respective sets of adjusted attention logits and the respective values for the input tokens (step 264).

[0174] In some implementations, as part of step 264, the system generates, for each particular input token, a respective set of attention weights for the particular input token that includes a respective attention weight for each of the input tokens. Then the system, for each particular input token, generates the respective output embedding for the particular output token as a weighted sum of the values for the input tokens with each value being weighted by the attention weight for the corresponding input token in the respective set of attention weights for the particular input token. Atorney Docket No. 56113-0812WO1

[0175] In some implementations, after the system performs steps 252-264, the system, for each selective attention layer, combines the respective output embeddings for each of the input tokens generated by the attention heads to generate an output of the selective attention layer. For example, the system, for each input token, can combine the output embeddings for the input token by first concatenating the output embeddings from each head for the input token into a single, larger vector. The system can then process this larger vector through a final linear projection layer to produce the final output embedding for the entire selective attention layer for each input token.

[0176] FIG. 3A is a flow diagram of an example process 300 for processing a new7input token during autoregressive inference. 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, a selective attention system, e.g., the selective attention system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 300.

[0177] The system receives a new input token (step 302). The new input token can be either the final token of an initial input sequence, e.g.. a user-provided prompt, or the output token that was generated by the neural network in the immediately preceding step of an ongoing autoregressive generation process.

[0178] In some cases, in addition to the system receiving a new input token, the system maintains a set of context tokens, which includes all tokens that appeared earlier in the sequence to the new input token.

[0179] The system processes the new input token using a neural network to generate a network output (step 304). The neural network includes a plurality7of layer blocks that include one or more attention layer blocks that each include a selective attention layer having one or more attention heads. In some cases, the plurality of layer blocks include a plurality of attention layer blocks that each include a respective selective attention layer

[0180] In some cases, the network output for this step 304 includes a respective score for each of a vocabulary of tokens.

[0181] For example, the system can process a final output embedding for the new input token, e.g., using a final output layer of the neural network, to generate a probability distribution over the tokens in a vocabulary (i.e., scores over the vocabulary of tokens). This probability7distribution can represent a prediction for the next token to be generated in the output sequence, which is distinct from the process described in FIG. 2A where the Atorney Docket No. 56113-0812WO1 network output is a full sequence of output embeddings for all tokens in the input sequence.

[0182] In some cases, when the network output includes a respective score for each of a vocabulary of tokens, the system selects a token to follow the new token in an output sequence using the network output.

[0183] For example, the system can use a sampling strategy based on the scores for the vocabulary of tokens to select the token to follow the new token in an output sequence using the network output.

[0184] As a particular example, the system can select the token with the highest score to follow the new token in an output sequence using the network output (e.g., greedy decoding strategy).

[0185] As another particular example, the system can modify the scores before sampling the tokens to determine the token to follow the new token (e.g., temperature sampling using a temperature parameter to scale scores).

[0186] FIG. 3 A provides an overview of how the system processes a new input token using the neural network to generate a network output. FIG. 3B illustrates, in detail, the process of step 304 for each selective attention layer and, within each of those layers, for each respective attention head of the selective attention layer.

[0187] In particular, FIG. 3B is a flow diagram of an example process 350 for applying a selective attention mechanism during a single step of autoregressive inference. For convenience, the process 350 will be described as being performed by a system of one or more computers located in one or more locations. For example, a selective attention system, e.g., the selective attention system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 350.

[0188] The system, for each selective attention layer, performs steps 352-354.

[0189] The system obtains a set of mask scores that includes a respective mask score for each token in a set of tokens that includes the new input token and a set of context tokens (step 352). The respective mask score measures a degree to which the new token is masked by the token.

[0190] Generally, the set of context tokens includes all tokens that precede the new input token, e.g., an input sequence of tokens and previously generated output tokens.

[0191] In some cases, as part of step 352, the system obtains a set of initial mask scores that includes a respective initial mask score for each of the tokens in the set. Then, the system generates the set of mask scores from the set of initial mask scores, which Atorney Docket No. 56113-0812WO1 includes setting to zero one or more of the initial mask scores in the set of initial mask scores.

[0192] For example, the system can set to zero each initial mask score that has a negative value. As a particular example, the system can set to zero each initial mask score that has a negative value by applying a rectified linear unit (ReLU) function to the initial mask scores.

[0193] As another example, the system can set to zero the initial mask score for the new token. When the system sets the initial mask score for the new token to zero, the system ensures the new token cannot mask itself during the attention calculation.

[0194] As another example, the system can set to zero the initial mask score for the “beginning of sequence token’; That is, when a first input token of the context tokens is a “beginning of sequence token”, the system sets to zero the initial mask score that the new input token would apply to the “beginning of sequence token”. This ensures that the special “beginning of sequence token” is preserved and not masked by any subsequent tokens.

[0195] The system determines, from the mask scores, a set of aggregated mask scores that includes a respective aggregated mask score for each token in the set of tokens (step 354).

[0196] In some cases, as part of step 354, the system, for each token in the set of tokens, generates the aggregated mask score by summing the mask scores for any tokens that precede the token in an ordering of the tokens in the set of tokens.

[0197] The system then performs, for each attention head of the selective attention layer, steps 356-366.

[0198] The system receives an input embedding of the new input token (step 356). For example, if the selective attention layer is the first layer in the neural network, this input embedding is generated from the initial tokenized input, for example, by mapping the token to a token embedding and adding a positional encoding. If the layer is a subsequent layer in the network, the input embedding it receives is the output embedding generated by the immediately preceding layer.

[0199] The system determines a respective query, key, and value for the new input token (step 358).

[0200] For example, the system can determine the query, key, and value by projecting the input embedding of the new input token using respective learned transformations, e.g., respective linear projection layers (e.g., by multiplying the embedding by three distinct Atorney Docket No. 56113-0812WO1 weight matrices), respective convolutional layers, or other ty pes of learned linear or nonlinear transformations.

[0201] The system obtains a respective key and value for each of a set of context tokens for the selective attention layer (step 360).

[0202] For example, the key and value for each context token can have been previously determined by projecting the respect ve embeddings for those tokens using the same learned linear projection layers described above and can be retrieved, e.g., from memory, e.g., the KV cache. In other words, in some cases, the system obtains the respective keys and values for the context tokens from a memory' of the one or more computers.

[0203] The system determines, from the query' for the new input token, the key for the new input token, and the respective keys for the context tokens, a set of attention logits that includes a respective attention logit for each of the tokens in the set (step 362).

[0204] For example, the system can determine the set of attention logits using a scaled dot-product attention mechanism. This involves computing a dot product between the query vector of the new input token and the key vector for each token in the set (i.e., the new input token and all the context tokens). The resulting scores, which represent the similarity' between the new input token and every other token in the context, are the attention logits.

[0205] The system combines the set of attention logits and the respective set of aggregated mask scores to generate a set of adjusted attention logits that includes a respective adjusted attention logit for each of the tokens in the set of tokens (step 364).

[0206] In some cases, as part of step 364, the system, for each token in the set, subtracts the aggregated mask score for the token from the attention logit for the token. This set of adjusted attention logits ensures that tokens with corresponding attention logits that have been lowered significantly using the corresponding aggregated mask score for the token receive less attention weight.

[0207] In some implementations, the system uses the set of attention logits to determine the set of initial mask scores that includes a respective initial mask score for each of the tokens in the set of tokens described above. That is. the system determines the set of initial mask scores from the respective sets of attention logits determined by one or more of the attention heads.

[0208] For example, in some cases, the system uses, as the set of initial mask scores, the respective set of attention logits determined by a designated one of the attention heads. In Atorney Docket No. 56113-0812WO1 some cases, the system, a user, or another system designate the attention head, e.g.. a user input received from a user can designate the attention head.

[0209] As another example, in some cases, the system can combine the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the set of initial mask scores. For example, the system can average the respective sets of attention logits determined by a designated subset of two or more of the attention heads.

[0210] The system generates an output embedding for the new token from the set of adjusted attention logits and the respective values for the new token and the context tokens (step 366).

[0211] In some implementations, as part of step 366, the system generates, from the adjusted attention logits, a set of attention weights that includes a respective attention weight for each of the tokens in the set. Then, the system generates the output embedding as a weighted sum of the values for the tokens in the set with each value being weighted by the attention weight for the corresponding token in the set of attention weights.

[0212] In some implementations, after the system performs steps 352-366. the system updates the set of context tokens by adding the new token to the set of context tokens. This allows the model to use the new token as part of the context for generating all subsequent tokens in the output sequence.

[0213] In some cases, the system determines that a number of tokens in the set of context tokens exceeds a maximum memory size for the attention layer. Then, in response, the system removes one or more of the context tokens from the set based on the aggregated mask scores for the context tokens. For example, the system can remove one or more of the context tokens having the highest aggregated mask scores and, as a consequence, remove the corresponding keys and values of the selected tokens from the KV-cache in computer memory to directly free up computer memory. Thus, the system removes the tokens that are least relevant for future attention mechanisms and frees up memory.

[0214] In some cases, two or more of the selective attention layers have different maximum memory sizes. By allocating different maximum memory sizes to different selective attention layers, the system can use its overall memory resources more efficiently. For example, some selective attention layers may need to maintain a large context, while others may only require a small context. So, allocating different maximum memory sizes for each layer maximizes the utility of the system’s memory usage. Atorney Docket No. 56113-0812WO1

[0215] In some cases, the system determines the maximum memory' sizes for the selective attention layers through a memory size optimization process. For example, the system can repeatedly process a group of input sequences to determine the optimal maximum memory size for each of the selective attention layers.

[0216] As a particular example, to perform the memory size optimization process, the system can begin with a fully trained neural netyvork, where each layer has the maximum possible memory budget. Then, the system can repeatedly evaluate and reduce the memory budget of the layer that causes the smallest negative impact on performance. The system can continue this iterative pruning until the neural netyvork’ s overall performance reaches a predefined target, resulting in a fixed, optimized set of per-layer memory sizes that are used for all subsequent inference tasks.

[0217] As described above, in some cases, yvhen the system obtains a respective key and value for each of a set of context tokens for the selective attention layer, the system can obtain the respective keys and values for the context tokens from a memory- of the one or more computers. For these cases, yvhen the system removes one or more of the context tokens from the set based on the aggregated mask scores for the context tokens, the system removes the respective keys and values for the one or more context tokens from the memory'. Thus, the system can efficiently use computer memory- by storing keys and values of context tokens that will be utilized in the future for quick access yvhile removing the keys and values of context tokens that will not be utilized in the future to create more free memory to process the next token in the sequence.

[0218] In some implementations, after the system performs steps 352-366, the system, for each selective attention layer, combines the respective output embeddings for the neyv token generated by the attention heads to generate an output of the selective attention layer. For example, the system can combine the output embeddings by first concatenating the output embeddings from each head into a single, larger vector, which the system then processes through a final linear projection layer to produce the final output embedding for the new token.

[0219] FIG. 4 is a flow diagram of an example process 400 for training the neural network 104. 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, a selective attention system, e.g., the selective attention system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 400. Atorney Docket No. 56113-0812WO1

[0220] The system repeatedly updates the trainable parameters of the neural network using a training data set. That is. the system can repeatedly perform the following described example process 400 using training examples to repeatedly update all or a subset of the trainable parameters of the neural network from previously undetermined values, e.g., randomly initialized values, or from previously determined values, e.g., pretrained values.

[0221] The training task for example process 400 is a next-token prediction task. Many of the machine learning tasks described above that the neural network can perform, including code generation, text generation, agent control, and so on, can be framed as a next-token prediction task. For example, for the appropriate input sequence, the neural network can produce a network output that is a sequence of output tokens that represents code, actions for an agent, or a selected class.

[0222] The system obtains a training data set that includes training examples (step 402), where each training example includes a training input sequence and a corresponding target output sequence. Each training input sequence includes a respective input token at each of a plurality of input positions. Generally, the target output sequence is a shifted input sequence so that the target token for a given input position is the input token from the immediately following position in the sequence.

[0223] The system, for each training example, generates a network output (step 404).

[0224] For example, for each training example, the system processes the training input sequence using the neural network to generate a respective network output that includes a respective probability distribution over a set of input tokens (e.g., a vocabulary7of tokens) for each of the input positions.

[0225] The system evaluates an objective that includes a loss function for each training example (step 406).

[0226] The loss function for each training example can be any appropriate loss function that measures the difference between the generated network output of the neural network and the target output sequence. For example, the loss function can be based on the respective probability distributions over the set of input tokens for each of the input positions included in the network output described above.

[0227] In some cases, the loss function can include a first term that measures, for one or more of the input positions, a probability assigned to a next input token by the probability' distribution for the input position, where the next input token is the input token at an immediately following input position in the input sequence. That is, for one or more input Atorney Docket No. 56113-0812WO1 positions, the loss function is based on the probability of the target token (i.e., the actual, correct token that immediately follows) included in the network output. In some cases, the one or more input positions includes all the input positions, but, in other cases, the one or more input positions include a designated proper subset of all the input positions.

[0228] For example, the first term of the loss function can be the cross-entropy loss (or negative log-likelihood) or log-perplexity over the set of input tokens of the input sequence.

[0229] In some cases, the loss function can include an auxiliary loss term that measures, for each selective attention layer, the sets of aggregated mask scores generated by the selective attention layer for the input tokens.

[0230] In some cases, the auxiliary loss term measures an approximate amount of memory required to perform auto-regressive inference given that tokens are removed from an inference context for the auto-regressive inference based on the sets of aggregated mask scores.

[0231] As an example, the loss function can be where Z is the first term and the log-perplexity loss described above; e is a small weight factor (e.g., 0.001, 0.01, or 0.1); L is the number of layers; and ntpadis the number of non-pad tokens in the input sequence. The term M- / r is the auxiliary loss term and approximates the memory requirements at the i th token for layer I and has a range of (0 < M < z). where i is the number of tokens in the context at step z, T is a constant (e.g., T = 1), F}kis the aggregated mask score for Zcth input position evaluated at the zth input position for layer I so that Sk=i min(Fk,T) represents a soft count of the number of tokens removed from the context. Thus, this example auxiliary loss term represents an estimate of the memory used (the number of active tokens) based on the aggregate mask scores. The use of the function min(-) clamps from above by T so as not to reward increasing it indefinitely (F is already clamped from below to 0). Since the memory required for a given layer is the maximum memory’ required for each of the tokens, this loss function considers the maximum among the M- s.

[0232] The system updates the trainable parameters to optimize the objective (step 408). Atorney Docket No. 56113-0812WO1

[0233] The system can update the trainable parameters of the neural network in any variety of ways, e.g.. gradient based method, evolutionary algorithm-based method, Bayesian optimization, etc.

[0234] For example, the system can optimize the objective using any of a variety of gradient descent techniques (e.g., batch gradient descent, stochastic gradient descent, or mini-batch gradient descent) that include the use of a backpropagation technique to estimate the gradient of the loss with respect to trainable parameters of the neural network and to update the learnable parameters accordingly.

[0235] Generally, the system repeats the above steps until one or more criteria are satisfied (e.g., the system performs a pre-determined number of iterations, the updates to the trainable parameters no longer exceed a pre-determined magnitude of change, a metric regarding a validation dataset exceeds a pre-determined value, and so on).

[0236] After training the neural network in the process described above, the neural network can operate with a substantially smaller memory footprint because the neural network can have its context buffer pruned during inference based on the learned aggregated mask scores. Additionally, by not attending to unneeded tokens, the performance of the neural network is also improved. As a result, the neural network is more computationally efficient (in terms of both memory and computational processing) and has higher performance.

[0237] FIG. 5 is an example 500 of the performance of the described techniques.

[0238] In particular, example 500 shows a pair of plots that compare the performance, measured by validation log-perplexity, of a neural network that processes an input sequence using the described techniques (i.e., a selective attention layer) to a neural network that does not. The left plot shows performance as a function of the length of the input sequence, while the right plot shows performance as a function of the number of parameters in the neural network.

[0239] Example 500 shows that a neural network using the described techniques consistently achieves a lower (i.e., better) validation log-perplexity than a conventional network across various input sequence lengths and model sizes. This advantage is a result of determining aggregated mask scores and combining them with the standard attention logits to generate adjusted attention logits, so that the system can reduce the influence of unneeded or irrelevant tokens in the input sequence. This “filtering of contextual noise” allows the neural network to generate a more accurate network output, resulting in a lower loss. The figure further shows that this performance advantage becomes more Atorney Docket No. 56113-0812WO1 pronounced as the context length increases, because the described techniques’ ability to mask irrelevant tokens becomes more critical as the amount of potential noise in the input sequence grows.

[0240] FIG. 6 is an example of the performance of the described techniques.

[0241] In particular, example 600 shows a plot that compares the performance of a neural network that processes an input sequence using a selective attention layer (i.e. , the described techniques) to a series of conventional neural networks having an increasing number of attention heads. The x’s on the plot corresponds to the described techniques while the immediately adjacent grouping of points connected with a line correspond to conventional techniques. The performance is measured by validation log-perplexity, which is a measure of a loss function based on the network output, as a function of the number of parameters in the attention layer.

[0242] Example 600 shows that the neural networks with selective attention layers (i.e., the described techniques) achieves a performance level equivalent to a conventional neural network with approximately twice the number of attention heads and parameters. While the conventional neural networks must use the brute-force method of adding more parameters and attention heads to achieve improved performance, the described techniques' more efficiently achieve equivalent performance with fewer parameters.

[0243] FIG. 7 is an example of the performance of the described techniques.

[0244] In particular, example 700 shows a comparison of accuracy results for a neural network that processes an input sequence using a selective attention layer (i.e., the described techniques, i.e., the columns labeled “selective”) to a baseline neural network that does not for various ML tasks (i.e., arc, hellaswag, commonsenseQA, openbookQA, and PiQA) and for a plurality’ of different model sizes (d= 16-28, with larger values of d corresponding to larger neural networks).

[0245] Example 700 shows that a neural network implementing the claimed selective attention method (i.e., the described techniques) consistently achieves higher accuracy on most ML tasks compared to the baseline neural network.

[0246] FIG. 8 is an example of the performance of the described techniques.

[0247] In particular, example 800 shows a plot illustrating the performances of neural networks that process an input sequence using a selective attention layer, as a function of memory’ savings achieved by pruning the context tokens. In particular, the plot compares the performance of a neural network trained with a standard loss function (Lppi), e.g., Atorney Docket No. 56113-0812WO1 cross-entropy loss, and one trained with an auxiliary loss term that measures the aggregated mask scores (Lmem) to a baseline conventional neural network (dotted line). The performance is measured by validation log-perplexity, a measure of a loss function based on the network output.

[0248] Example 800 shows that the described techniques allow for a substantial reduction in memory requirements during inference without degrading performance below the baseline. The best performance, i.e., Lmem. is a direct result of training the neural network with the auxiliary loss term that improves the process of determining aggregated mask scores and using these scores to remove the least relevant context tokens when a memory budget is exceeded. The plot further illustrates that including an auxiliary loss term that measures the aggregated mask scores during training results in a model capable of achieving significantly greater memory savings (a memory improvement factor of approximately 47 in this example) compared to standard training.

[0249] FIG. 9 is an example of the performance of the described techniques.

[0250] In particular, example 900 shows a plot that compares the performance of a neural network that processes an input sequence using a selective attention layer (i.e., the described techniques without the use of the auxiliary loss term during training labeled as “Selective attn"’ and with the use of the auxiliary loss term during training “Selective attn Lmem”) to neural networks that use alternative techniques. The performance is measured by validation log-perplexity, a measure of a loss function based on the network output, as a function of a memory improvement factor.

[0251] Example 900 shows that the described techniques achieve a substantially better performance-efficiency trade-off than the alternative methods. The alternative methods, require large drops in performance in order to match the equivalent memory improvement factor of the described techniques.

[0252] FIG. 10 is an example 1000 visualization of the selective attention mechanism.

[0253] In particular, example 1000 show s a visualization of the described techniques’ selective attention mechanism's operation on a plurality of input tokens of an input sequence for four different machine learning tasks (i.e., Variable Assignment, Parity*, Copy, and Language Modeling), where each row represents the input sequence as it grows (e.g., through auto-regressive generation of output tokens). Example 1000 illustrates both the standard attention logits (shading) and the effect of the mask scores (strike-through width) that are used to generate the adjusted attention logits. Atorney Docket No. 56113-0812WO1

[0254] For the subplot labeled “Variable Assignment'’, the input sequence represents a series of variable assignments, and the visualization shows that a particular input token corresponding to a new assignment generates high mask scores for preceding input tokens representing outdated assignments to the same variable. This results in a high aggregated mask score for the “outdated” tokens, reducing their influence and allowing the selective attention mechanism to correctly identify the most recent value when generating a network output. Also, because these outdated tokens are associated with high aggregated mask scores, the described techniques can select these tokens, based on their aggregated mask scores, for removal from the KV cache to improve inference efficiency of the system.

[0255] For the subplot labeled “Parity*”, the input sequence represents a progressive binary sequence where bits in the odd positions are random and bits in the even positions contain the parity of all bits in the earlier odd positions. The selective attention mechanism generates high mask scores for all but the most recent input tokens. This demonstrates the system's ability to determine that a large portion of the context is irrelevant, allowing the selective attention layer to generate a more accurate output by focusing only on the necessary tokens. The tokens with high mask scores, representing the irrelevant history', could then be removed from the KV cache, significantly reducing the memory footprint for this task.

[0256] For the subplot labeled “Copy”, the input sequence represents an arbitrary’ sequence delimited by a special token (i.e., “<copy>”) that is then copied. When processing the initial part of the input sequence, the selective attention mechanism generates low mask scores, retaining the full context in memory. As the selective attention mechanism begins to autoregressively generate the copied sequence, each new input token generates high mask scores for the corresponding token in the original sequence that has just been copied. This shows the mechanism being used to track generation progress. As each original token is masked after being copied, its key and value could be removed from the KV cache, progressively freeing up memory as the generation proceeds.

[0257] For the subplot labeled “Language Modeling”, the input sequence represents a sequence of natural language text. For a multi-token entity in the input sequence (e.g., "Barack Obama"), the visualization shows that a later input token (e.g., "Obama") generates high mask scores for preceding tokens that are part of the same entity (e.g., "Barack"). This allows the selective attention mechanism to consolidate the semantic Atorney Docket No. 56113-0812WO1 meaning and reduce the influence of redundant tokens when generating subsequent output tokens. These redundant, masked tokens could then be removed from the KV cache, allowing the system to maintain a more compact and semantically relevant context.

[0258] 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.

[0259] 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.

[0260] 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. Atorney Docket No. 56113-0812WO1

[0261] 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.

[0262] 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.

[0263] 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.

[0264] 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.

[0265] 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 Atorney Docket No. 56113-0812WO1 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.

[0266] 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.

[0267] 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 cry stal 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 bysending 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.

[0268] 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. Atorney Docket No. 56113-0812WO1

[0269] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a J AX framework.

[0270] 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 web browser, 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 network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0271] 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.

[0272] 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.

[0273] 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 Attorney Docket No. 56113-0812WO1 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.

[0274] 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. Attorney Docket No. 56113-0812WO1CLAIMS1. A method performed by one or more computers, the method comprising: receiving an input sequence comprising a respective input token at each of a plurality of input positions; and processing the input sequence using a neural network to generate a network output, wherein: the neural network comprises a plurality of layer blocks that include one or more attention layer blocks that each include a selective attention layer having one or more attention heads, and wherein processing the input sequence comprises, for each selective attention layer: obtaining a respective set of mask scores for each particular input token of the plurality of input tokens that includes a respective mask score for each of the input tokens that measures a degree to which the particular input token is masked by the input token; determining, for each particular input token and from the respective set of mask scores for the particular input token, a respective set of aggregated mask scores; and for each attention head of the selective attention layer: receiving a respective input embedding for each of the input tokens; determining a respective query, key, and value for each of the input tokens from the respective input embeddings for each of the input tokens; determining, for each particular input token and from the query for the particular input token and the respective keys for the input tokens, a respective set of attention logits that includes a respective attention logit for each of the input tokens; combining, for each particular input token, the respective set of attention logits for the particular input token and the respective set of aggregated mask scores for the particular input token to generate a respective set of adjusted attention logits for the particular input token that includes a respective adjusted attention logit for each of the input tokens; and generating a respective output embedding for each of the input tokens for the attention head from the respective sets of adjusted attention logits and the respective values for the input tokens.Atorney Docket No. 56113-0812WO12. The method of claim 1. wherein obtaining a respective set of mask scores for each particular input token of the plurality of input tokens that includes a respective mask score for each of the input tokens that measures a degree to which the particular input token is masked by the input token comprises: obtaining a respective set of initial mask scores for each particular input token of the plurality of input tokens that includes a respective initial mask score for each of the input tokens; and generating the respective sets of mask scores from the respective sets of initial mask scores, comprising, for one or more of the particular input tokens, generating the respective set of mask scores for the particular input token by setting to zero one or more of the initial mask scores in the respective set of initial mask scores for the particular input token.

3. The method of claim 2, wherein obtaining a respective set of initial mask scores for each particular input token of the plurality of input tokens that includes a respective initial mask score for each of the input tokens comprises: determining the respective sets of initial mask scores for the particular input tokens from the respective sets of attention logits determined by one or more of the attention heads.

4. The method of claim 3, wherein determining the respective sets of initial mask scores for the particular input tokens from the respective sets of attention logits determined by one or more of the attention heads comprises: using, as the respective sets of initial mask scores for the particular input tokens, the respective sets of attention logits determined by a designated one of the attention heads.

5. The method of claim 3, wherein determining the respective sets of initial mask scores for the particular input tokens from the respective sets of attention logits determined by one or more of the attention heads comprises: combining the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the respective sets of initial mask scores for the particular input tokens.Atorney Docket No. 56113-0812WO16. The method of claim 5. wherein combining the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the respective sets of initial mask scores for the particular input tokens comprises: averaging the respective sets of attention logits determined by a designated subset of two or more of the attention heads.

7. The method of any one of claims 2-6, wherein generating the respective set of mask scores for the particular input token by setting to zero one or more of the initial mask scores in the respective set of initial mask scores for the particular input token comprises: setting to zero each initial mask score that has a negative value.

8. The method of claim 7, wherein setting to zero each initial mask score that has a negative value comprises applying a rectified linear unit (ReLU) function to the initial mask scores.

9. The method of any one of claims 2-8, wherein generating the respective set of mask scores for the particular input token by setting to zero one or more of the initial mask scores in the respective set of initial mask scores for the particular input token comprises: setting to zero the initial mask score for the particular token.

10. The method of any one of claims 2-9, wherein a first input token at a first input position in the input sequence is a beginning of sequence token, and wherein generating the respective set of mask scores for the particular input token by setting to zero one or more of the initial mask scores in the respective set of initial mask scores for the particular input token comprises: setting to zero the initial mask score for the beginning of sequence token.

11. The method of any preceding claim, wherein determining, for each particular input token and from the respective set of mask scores for the particular input token, a respective set of aggregated mask scores comprises, for each particular input token: for each input token, generating the aggregated mask score by summing the mask scores, in the set of mask scores for the particular input token, for any input tokens that precede the input token in the input sequence.Attorney Docket No. 56113-0812WO112. The method of any preceding claim, wherein combining, for each particular input token, the respective set of attention logits for the particular input token and the respective set of aggregated mask scores for the particular input token to generate a respective set of adjusted attention logits for the particular input token that includes a respective adjusted attention logit for each of the input tokens comprises, for each particular input token: for each input token, subtracting the aggregated mask score for the input token from the attention logit for the input token.

13. The method of any preceding claim, wherein the processing further comprises, for each selective attention layer: combining the respective output embeddings for each of the input tokens generated by the attention heads to generate an output of the selective attention layer.

14. The method of any preceding claim, wherein generating a respective output embedding for each of the input tokens for the attention head from the respective sets of adjusted attention logits and the respective values for the input tokens comprises: generating, for each particular input token, a respective set of attention weights for the particular input token that includes a respective attention weight for each of the input tokens; and for each particular input token, generating the respective output embedding for the particular output token as a weighted sum of the values for the input tokens with each value being weighted by the attention weight for the corresponding input token in the respective set of attention weights for the particular input token.

15. The method of any preceding claim, wherein the input sequence is a training input sequence for training the neural network; the network output comprises a respective probability distribution over a set of input tokens for each of the input positions; and the method further comprises: training the neural network on a loss function based on the respective probability distributions.Atorney Docket No. 56113-0812WO116. The method of claim 15, wherein: the loss function includes a first term that measures, for one or more of the input positions, a probability assigned to a next input token by the probability distribution for the input position, wherein the next input token is the input token at an immediately following input position in the input sequence.

17. The method of claim 15 or claim 16, wherein: the loss function includes an auxiliary loss term that measures, for each selective attention layer, the sets of aggregated mask scores generated by the selective attention layer for the input tokens.

18. The method of claim 17, wherein the auxiliary loss term measures an approximate amount of memory' required to perform auto-regressive inference to generate at least a portion of the input sequence given that tokens are removed from an inference context for the auto-regressive inference based on the sets of aggregated mask scores.Attorney Docket No. 56113-0812WO119. A method performed by one or more computers, the method comprising: receiving a new input token; and processing the new input token using a neural network to generate a network output, wherein: the neural network comprises a plurality of layer blocks that include one or more attention layer blocks that each include a selective attention layer having one or more attention heads, and wherein processing the new input token comprises, for each selective attention layer: obtaining a set of mask scores that includes a respective mask score for each token in a set of tokens that includes the new input token and a set of context tokens, wherein the respective mask score measures a degree to which the new token is masked by the token; determining, from the mask scores, a set of aggregated mask scores that includes a respective aggregated mask score for each token in the set of tokens; and for each attention head of the selective attention layer: receiving an input embedding of the new input token; determining a respective query, key, and value for the new input token; obtaining a respective key and value for each of a set of context tokens for the selective attention layer; determining, from the query for the new input token, the key for the new input token, and the respective keys for the context tokens, a set of attention logits that includes a respective attention logit for each of the tokens in the set; combining the set of attention logits and the respective set of aggregated mask scores to generate a set of adjusted attention logits that includes a respective adjusted attention logit for each of the tokens in the set of tokens; and generating an output embedding for the new token from the set of adjusted attention logits and the respective values for the new token and the context tokens.Atorney Docket No. 56113-0812WO120. The method of claim 19, wherein obtaining a set of mask scores that includes a respective mask score for each token in a set of tokens that includes the new input token and a set of context tokens comprises: obtaining a set of initial mask scores that includes a respective initial mask score for each of the tokens in the set; and generating the set of mask scores from the set of initial mask scores, comprising setting to zero one or more of the initial mask scores in the set of initial mask scores.

21. The method of claim 20, wherein obtaining a set of initial mask scores that includes a respective initial mask score for each of the tokens in the set comprises: determining the set of initial mask scores from the respective sets of attention logits determined by one or more of the attention heads.

22. The method of claim 21, wherein determining the set of initial mask scores from the respective sets of attention logits determined by one or more of the attention heads comprises: using, as the set of initial mask scores, the respective set of attention logits determined by a designated one of the attention heads.

23. The method of claim 21, wherein determining the set of initial mask scores from the respective sets of attention logits determined by one or more of the attention heads comprises: combining the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the set of initial mask scores.

24. The method of claim 23, wherein combining the respective sets of attention logits determined by a designated subset of two or more of the attention heads to generate the set of initial mask scores comprises: averaging the respective sets of attention logits determined by a designated subset of two or more of the attention heads.

25. The method of any one of claims 20-24, wherein setting to zero one or more of the initial mask scores in the set of initial mask scores comprises: setting to zero each initial mask score that has a negative value.Attorney Docket No. 56113-0812WO126. The method of claim 25, wherein setting to zero each initial mask score that has a negative value comprises applying a rectified linear unit (ReLU) function to the initial mask scores.

27. The method of any one of claims 20-26, wherein setting to zero one or more of the initial mask scores in the set of initial mask scores comprises: setting to zero the initial mask score for the new token.

28. The method of any one of claims 20-27, wherein a first input token of the context tokens is a beginning of sequence token, and wherein setting to zero one or more of the initial mask scores in the set of initial mask scores comprises: setting to zero the initial mask score for the beginning of sequence token.

29. The method of any one of claims 19-28, wherein determining a set of aggregated mask scores comprises: for each token in the set, generating the aggregated mask score by summing the mask scores for any tokens that precede the token in an ordering of the tokens in the set.

30. The method of any one of claims 19-29, wherein combining the set of attention logits and the set of aggregated mask scores comprises: for each token in the set, subtracting the aggregated mask score for the token from the attention logit for the token.31 . The method of any one of claims 19-30, further comprising: updating the set of context tokens by adding the new token to the set of context tokens.

32. The method of claim 31, further comprising: determining that a number of tokens in the set of context tokens exceeds a maximum memory size for the attention layer; and in response, removing one or more of the context tokens from the set based on the aggregated mask scores for the context tokens.Attorney Docket No. 56113-0812WO133. The method of claim 32, wherein removing one or more of the context tokens from the set based on the aggregated mask scores for the context tokens comprises: removing one or more of the context tokens having the highest aggregated mask scores.

34. The method of claim 32 or claim 33. wherein obtaining a respective key and value for each of a set of context tokens for the selective attention layer comprises: obtaining the respective keys and values for the context tokens from a memory of the one or more computers; and wherein removing one or more of the context tokens from the set based on the aggregated mask scores for the context tokens comprises: removing the respective keys and values for the one or more context tokens from the memory.

35. The method of any one of claims 19-34, wherein the plurality of layer blocks include a plurality of attention layer blocks that each include a respective selective attention layer.

36. The method of claim 35, when dependent on claim 32, wherein two or more of the selective attention layers have different maximum memory sizes.

37. The method of any one of claims 19-36, wherein the processing further comprises, for each selective attention layer, combining the respective output embeddings for the new token generated by the attention heads to generate an output of the selective attention layer.

38. The method of any one of claims 19-37, wherein generating the output embedding for the new token for the attention head comprises: generating, from the adjusted attention logits, a set of attention weights that includes a respective attention weight for each of the tokens in the set; and generating the output embedding as a weighted sum of the values for the tokens in the set with each value being weighted by the attention weight for the corresponding token in the set of attention weights.Atorney Docket No. 56113-0812WO139. The method of any one of claims 19-38, wherein the network output comprises a respective score for each of a vocabulary of tokens, and wherein the method further comprises: selecting a token to follow the new token in an output sequence using the network output.

40. A 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 more computers to perform the operations of the respective method of any one of claims 1-39.

41. One or more computer storage media storing instructions that when executed by one or more computers cause the one more computers to perform the operations of the respective method of any one of claims 1-39.

Images