Attention neural networks with mixture of depths attention layer blocks

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium for processing an input sequence using a neural network that includes one or more mixture of depths attention layer blocks that can make dynamic token-level routing decisions across the depth of the network.

Description

[0001] ATTENTION NEURAL NETWORKS WITH MIXTURE OF DEPTHS ATTENTION LAYER BLOCKS

[0002] CROSS REFERENCE TO RELATED APPLICATION

[0003] This application claims priority to U.S. Provisional Application No. 63 / 572,112 filed on March 29, 2024. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application.

[0004] BACKGROUND

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

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

[0007] SUMMARY

[0008] 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. More specifically, this specification describes the use and training of a neural network that includes one or more mixture of depths attention layer blocks that can make dynamic token-level routing decisions across the depth of the network. The mixture of depths attention layer block can select specific input positions to route through a layer of the block for more processing and input positions to route around the block to save compute.

[0009] A mixture of depths attention layer block, as used in the specification, is a layer block including attention layers that does not route each token of an input sequence through the block. As a result, individual tokens can pass through different numbers of layer blocks through the depth of the model.

[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. Generally, neural network models spread floating point operations (FLOPs) uniformly across input sequences. In a neural network model. FLOPs can be a measure of the amount of compute required to generate an output for a given input. Thus, neural networks generally expend the same amount of compute per token in a forward pass.

[0011] However, not all tokens require the same amount compute to process in order to generate accurate output or representation of the token. That is, for some or all of the tokens in any given sequence, not all of the FLOPs performed by the neural network may be necessary to generate a high quality output. This can be due to many factors, including token complexity, ambiguity, and importance.

[0012] For example, in natural language processing, a word-level tokenization sequence may include tokens for words such as "the" or "and” that are less important for the semantic context of the sequence and thus, the neural network may not need to perform as many FLOPs when processing these tokens, i.e., relative to the other tokens in the sequence, to generate a high quality output for the sequence.

[0013] As another example, in image processing, tokens that represents pixel values for pixels in regions or patches of the image that are entirely background, e.g.. not the focal point of the image, are not as important to the semantic context of the input sequence of the image and thus, the neural network may not need to perform as many FLOPs when processing these tokens, i.e., relative to the other tokens in the sequence, to generate a high quality output for the sequence.

[0014] As another example, ambiguity can influence how many FLOPs are necessary for certain tokens in a sequence to ensure a high quality output. For example, in natural language processing, certain tokens can represent ambiguous words, such as homographs, e.g., words that are spelled the same but have different meanings. As a particular example, the word '‘lead” can mean "to guide’’ or "‘a type of metal” and may require more FLOPs to generate a high quality output, i.e., to allow the neural network to distinguish the meaning of the word based on the context of the rest of the sequence.

[0015] As another example, complexity can influence the FLOPs given to certain tokens in the sequence. For example, in audio processing, a token sequence representing audio data being spoken in a natural language can include a segment of the audio that is entirely silent. While potentially being important to the semantic context of the audio data, a silent segment of audio data is not as complex as segments of audio that include a spoken utterance and thus, the neural network may not need to perform as many FLOPs when processing these tokens, i.e., relative to the other tokens in the sequence, to generate a high quality output for the sequence.

[0016] To address this issue, the neural network system as described in this specification, by virtue of including one or more mixture of depths attention layer blocks, can leam to dynamically allocate compute, e.g., FLOPs, to specific positions in a sequence, optimizing the allocation along the sequence for different layers across the model depth. Thus, compute expenditure is dynamic and context sensitive at the token-level, e.g. tokens that need more processing can be given more compute and tokens that need less processing are given less compute. This results in the neural network system (i) generating higher-quality outputs given the same number of FLOPs as conventional systems, (ii) generating the same or higher quality outputs in significantly fewer FLOPs as conventional systems, or (iii) both.

[0017] In some cases, the compute expenditure by the described neural network system, while being dynamic and context sensitive at the token-level, can nonetheless be entirely predictable in sum total. By defining compute prior to training, hardware efficiency gains can be anticipated and exploited ahead of time without sacrificing the performance of the model. That is, by specifying the sum total of the compute that will be performed on any given input while dynamically varying only how that compute is allocated across the tokens in the sequence, the system can optimize the deployment of the neural network, e.g., on hardware accelerators or other parallel processing hardware.

[0018] In some cases, the described neural network system can require less memory and / or fewer hardware accelerators (e.g., Tensor Processing Units) than similarly performing systems that do not comprise mixture of depths attention layer blocks. For example, mixture of depths attention layer blocks can allow for a smaller KV cache to be used during autoregressive sampling in some cases.

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

[0020] According to a first aspect there is provided a method performed by one or more computers. The method includes receiving an input sequence including 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. The neural network includes a plurality' of layer blocks that include one or more mixture of depths attention layer blocks, and processing the input sequence includes, for each mixture of depths attention layer block selecting, using a first routing model for the mixture of depths attention layer block and from respective first embeddings for each of the input positions, a first proper subset of the input positions for the mixture of depths attention layer block. Processing the input sequence further includes processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block to generate respective second embeddings for each of the input positions in the first proper subset, including applying an attention mechanism for the mixture of depths attention layer block to the respective first embeddings of each of the input positions in the first proper subset of the input positions and generating a respective second embedding for each of the input positions not in the first proper subset without applying the attention mechanism for the mixture of depths attention layer block to the input positions not in the first proper subset.

[0021] In some implementations, selecting, using a first routing model and from the respective first embeddings for each of the input positions, a first proper subset of the input positions to be processed by the attention layer block includes processing the respective first embedding for each of the input positions using the first routing model for the mixture of depths attention layer block to generate a respective routing score for each of the input positions and selecting the first proper subset of the input positions based on the respective routing scores for the input positions.

[0022] In some implementations, selecting the first proper subset of the input positions based on the respective routing scores includes selecting the k input positions having the highest routing scores, wherein k is an integer that is less than the total number of input positions in the sequence.

[0023] In some implementations, processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block to generate the respective second embeddings for each of the input positions in the first proper subset includes processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block to generate respective initial second embeddings for each of the input positions in the first proper subset. Processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block further includes, for each input position in the first proper subset, scaling the initial second embedding for the input position by the routing score for the input position to generate the second embedding for the input position.

[0024] In some implementations, each mixture of depths attention layer block includes (i) an attention layer that applies the attention mechanism for the mixture of depths attention block and (ii) a feedforward subnetwork. Processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block to generate respective second embeddings for each of the input positions in the first proper subset includes processing the respective first embeddings for each of the input positions in the first proper subset of the input positions using the attention layer to generate a respective attended embedding for each of the input positions in the first proper subset and processing the respective attended embeddings for each of the input positions for the first proper subset of the input positions using the feedforward subnetwork to generate a respective third embedding for each of the input positions in the first proper subset.

[0025] In some implementations, processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block further includes applying a residual connection to the respective third embeddings for each of the input positions in the first proper subset to generate the respective initial second embeddings.

[0026] In some implementations, generating the respective second embeddings for each of the input positions not in the first proper subset without applying the attention mechanism for the mixture of depths attention layer block to the input positions not in the first proper subset includes generating the respective second embeddings for each of the input positions not in the first proper subset without processing any embeddings for the input positions not in the first proper subset using the attention layer or the feedforward subnetwork.

[0027] In some implementations, generating the respective second embeddings for each of the input positions not in the first proper subset includes, for each of the input positions not in the first proper subset, setting the second embedding for the input position equal to the first embedding for the input position.

[0028] 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. The method further includes training the neural network and the respective first routing models for the mixture of depths attention layer blocks on a loss function based on the respective probability distributions.

[0029] In some implementations, the loss function includes a first term that measures, for each input position, 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.

[0030] In some implementations, the loss function includes a respective auxiliary loss term for each mixture of depths attention layer block that measures a difference between (i) the respective routing scores for the input positions and (ii) a target set of routing scores that assigns each input position in the first proper subset a score of 1 and each input position not in the first proper subset a score of 0.

[0031] In some implementations, the method further includes training a respective inference-time routing model for each mixture of depths attention layer, where the respective inference-time routing model is configured to receive, for each input position, the respective first embedding for the input position and to generate as output a score that indicates a likelihood that the input position will be in the first proper subset, and where training the respective inference-time routing model includes training the respective inference-time routing model on an objective that measures, for each input position, an error between the score for the input position and a target score for the input position is equal to a first score if the input position is in the first proper subset and a second score if the input position is not in the first proper subset.

[0032] In some implementations, the first routing model for the mixture of depths attention layer is a linear neural network layer or a multi-layer perceptron (MLP).

[0033] In some implementations, the one or more layer blocks further includes one or more attention layer blocks that are not mixture of depths attention layer blocks.

[0034] In some implementations, the one or more layer blocks includes one or more mixture of depths attention layer blocks.

[0035] In some implementations, for one or more of the mixture of depths attention layer blocks, the feedforward subnetwork is a mixture of experts (MoE) subnetwork that includes a routing function and one or more expert feedforward neural networks.

[0036] According to the first aspect there is provided a second method performed by one or more computers and for auto-regressively generating an output sequence including a respective output token at each of plurality of output positions. The method includes, for each of the plurality of output positions receiving a preceding output token generated for the preceding output position and processing the preceding output token using a neural network to generate the output token at the output position. The neural netw ork includes one or more layer blocks that include one or more mixture of depths attention layer blocks. Processing the preceding output token includes, for each mixture of depths atention layer block, determining, using a routing model for the mixture of depths atention layer block and from a respective first embedding for the preceding output position, whether to include the preceding output position in a first proper subset of the output positions for the mixture of depths atention layer block and in response to determining to include the preceding output position in the first proper subset of the output positions, processing the respective first embedding for the preceding output position using the mixture of depths atention layer block to generate a respective second embedding for the preceding output position, comprising applying an atention mechanism for the mixture of depths atention layer block to the respective first embedding for the preceding output position. The method further includes, in response to determining not to include the preceding output position in the first proper subset of the output positions, generating the respective second embedding for the preceding output position without applying the attention mechanism for the mixture of depths atention layer block to the respective first embedding for the preceding output position.

[0037] In some implementations, determining whether to include the preceding output position in the first proper subset includes processing the respective first embedding for the preceding output position using the routing model to generate a score for the preceding output position and determining whether to include the preceding output position in the first proper subset using the score.

[0038] In some implementations, determining whether to include the preceding output position in the first proper subset using the score includes including the preceding output position in the first proper subset when the score exceeds a threshold score.

[0039] In some implementations, processing the respective first embedding for the preceding output position using the mixture of depths atention layer block to generate a respective second embedding for the preceding output position includes processing the respective first embedding for the preceding output position using the mixture of depths atention layer block to generate a respective initial second embedding for the preceding output position and scaling the initial second embedding for the input position by a routing score for the input position to generate the second embedding for the input position.

[0040] In some implementations, the routing score for the input position is the score for the input position.

[0041] In some implementations, the routing model is the first routing model of any one of the above implementations of the first method. In some implementations, the routing model is the inference-time routing model of the first method.

[0042] In some implementations, processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate a respective second embedding for the preceding output position includes processing the respective first embedding for the preceding output position using a first routing model for the mixture of depths attention layer block to generate the routing score for the preceding output position.

[0043] In some implementations, applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position includes applying the attention mechanism for the mixture of depths attention layer block to (i) the respective first embedding for the preceding output position and (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset.

[0044] In some implementations, the output sequence is being generated in response to an input sequence, and wherein applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position includes applying the attention mechanism for the mixture of depths attention layer block to (i) the respective first embedding for the preceding output position, (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset, and (iii) respective first embeddings in a first proper subset of input positions in the input sequence.

[0045] In some implementations, applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position includes applying the attention mechanism for the mixture of depths attention layer block to (i) the respective first embedding for the preceding output position and (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset but not to (iv) respective first embeddings for output positions in the output sequence that were previously not selected for inclusion in the first proper subset. Further in some implementations, each mixture of depths attention layer block includes (i) an attention layer that applies the attention mechanism for the mixture of depths attention block and (ii) a feedforward subnetwork. Processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate the respective second embedding for the preceding output position includes processing the respective first embedding for the preceding output position using the attention layer to generate a respective attended embedding for the preceding output position and processing the respective attended embedding for the preceding output position using the feedforward subnetwork to generate a respective third embedding for the preceding output position.

[0046] In some implementations, processing the respective first embedding for the preceding output position using the mixture of depths attention layer block further includes applying a residual connection to the respective third for the preceding output position to generate the respective initial second embedding for the preceding output position.

[0047] In some implementations, generating the respective second embedding for the preceding output position without applying the attention mechanism for the mixture of depths attention layer block to the preceding output position includes generating the respective second embedding for the preceding output position without processing any embedding for the preceding output position using the attention layer or the feedforward subnetwork.

[0048] In some implementations, in response to determining not to include the preceding output position in the first proper subset of the output positions, generating the respective second embedding for the preceding output position without applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position includes setting the second embedding for the preceding output position equal to the first embedding for the preceding output position.

[0049] In some implementations, the routing model for the mixture of depths attention layer is a linear neural network layer or a multi-layer perceptron (MLP).

[0050] In some implementations, the one or more layer blocks further includes one or more attention layer blocks that are not mixture of depths attention layer blocks.

[0051] In some implementations, the one or more layer blocks includes one or more mixture of depths attention layer blocks. In some implementations, for one or more of the mixture of depths attention layer blocks, the feedforward subnetwork is a mixture of experts (MoE) subnetwork that includes a routing function and one or more expert feedforward neural networks.

[0052] According to the first aspect there is provided a third method performed by one or more computers. The method includes receiving an input sequence including a respective input token at each of one or more input positions and processing the input sequence using a neural network to generate a network output. Processing the input sequence using a neural network to generate a network output includes selecting, using a first routing model and from respective first embeddings for each of the input positions, a first proper subset of the input positions to be processed by an attention layer block and processing the respective first embeddings for the first proper subset of the input positions using the attention layer block to generate respective second embeddings for each of the input positions in the first proper subset, comprising applying the attention mechanism of the attention layer block for each of the input positions in the first proper subset of the input positions. Processing the input sequence using a neural network to generate a network output further includes generating a respective second embedding for each of the input positions not in the first proper subset without applying the attention mechanism for the attention layer block to the input positions not in the first proper subset.

[0053] According to the first aspect there is provided a fourth method performed by one or more computers and for auto-regressively generating an output sequence including a respective output token at each of one or more output positions, the method including, for each of the one or more output positions receiving a preceding output token generated for the preceding output position and processing the preceding output token using a neural network to generate the output token at the output position. Processing the preceding output token using a neural network to generate the output token at the output position includes determining, using an inference routing model and from a respective first embedding for the preceding output position, whether to include the preceding output position in a first proper subset of the output positions and in response to determining to include the preceding output position in the first proper subset of the output positions, processing the respective first embedding for the preceding output position using the attention layer block to generate a respective second embedding for the preceding output position, including applying an attention mechanism of the attention layer block to the respective first embedding for the preceding output position. Then, in response to determining not to include the preceding output position in the first proper subset of the output positions, generating the respective second embedding for the preceding output position without applying the attention mechanism for the attention layer block to the respective first embedding for the preceding output position.

[0054] According to the first aspect there is provided a fifth method performed by one or more computers. The method including 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. Processing the input sequence using a neural network to generate a network output includes selecting, using a first routing model and from respective first embeddings for each of the input positions, a first proper subset of the input positions to be included as keys for an attention mechanism of an attention layer block and processing the respective first embeddings using the attention layer block to generate respective second embeddings for each of the input positions, including applying the attention mechanism of the attention layer block to the respective first embedding using one or more sets of keys, wherein each set of keys includes a respective key for each of the input positions in the first proper subset of the input positions and does not include a respective key for any of the input positions not in the first proper subset.

[0055] According to the first aspect there is provided a sixth method performed by one or more computers and for auto-regressively generating an output sequence including a respective output token at each of one or more output positions. The method includes, for each of the one or more output positions receiving a preceding output token generated for the preceding output position and processing the preceding output token using a neural network to generate the output token at the output position. The processing of the preceding output token using a neural network to generate the output token at the output position includes determining, using an inference routing model and from a respective first embedding for the preceding output position, whether to include the preceding output position in a first proper subset of the output positions and processing the respective first embedding for the preceding output position using the attention layer block to generate a respective second embedding for the preceding output position, comprising applying an attention mechanism of the attention layer block to the respective first embedding for the preceding output position using one or more sets of keys, wherein each set of keys includes a respective key for each output position that has been selected for inclusion in the first proper subset and does not include a respective key for any of the output positions that have not been selected for inclusion in the first proper subset. According to the first aspect there is provided a seventh method where the neural network is configured to perform a machine learning task including an image processing task, where the neural network is configured to process an input image to generate: (i) a network output including scores for each of a set of object categories, each score representing an estimated likelihood that the input image includes an image of an object belonging to the category; (ii) a network output including identifying locations in the input image at which particular types of objects are depicted; or (iii) a network output assigning each pixel of the input image to a category from a set of categories.

[0056] According to a second aspect there is provided the methods of the first 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 respective operations of the respective methods.

[0057] According to a third aspect there is provided the methods of the first 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 respective operations of the respective methods.

[0058] Other features, aspects and advantages of the subj ect matter will become apparent from the description, the drawings, and the claims.

[0059] BRIEF DESCRIPTION OF THE DRAWINGS

[0060] FIG. 1 shows an example neural network system that includes one or more mixture of depths attention layer blocks within a neural network.

[0061] FIG. 2A shows an example mixture of depths attention layer block during inference time.

[0062] FIG. 2B shows an example mixture of depths attention layer block during auto regressive inference.

[0063] FIG. 3 shows example graphs of the routing decisions of different layer blocks.

[0064] FIG. 4 shows example routing methods for a mixture of depths attention layer block.

[0065] FIG. 5 is a flow' diagram of an example process of the example neural netw ork system.

[0066] FIG. 6 is a flow diagram of sub-steps of one of the steps of the process of FIG. 6.

[0067] FIG. 7 is a flow diagram of an example process of auto-regressively generating an output sequence using the example neural network system. FIG. 8 is a flow diagram of sub-steps of one of the steps of the process of FIG. 8.

[0068] FIG. 9 is a diagram demonstrating an example of the routing decisions of a trained neural network including one or more mixture of depths attention layer blocks.

[0069] FIG. 10 shows an example of the performance of the neural network in autoregressive mode.

[0070] FIG. 11 shows an example mixture of depths modification to a mixture of experts approach.

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

[0072] DETAILED DESCRIPTION

[0073] FIG. 1 shows an example neural network system 100 that includes one or more mixed depths attention layer blocks within a neural network 130. The neural network system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations in which the systems, components and techniques described below are implemented.

[0074] The neural network system 100 receives as input an input sequence 105. The input sequence 105 can include a respective input token at each of one or more input positions.

[0075] The neural network 130 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 105 and to generate any kind of score, classification, or regression output based on the input sequence 105.

[0076] Some examples of machine learning tasks that a neural netw ork 130 when implemented using one of the architectures described below or other known architectures can be configured to perform now follow.

[0077] In any of the implementations below-, the neural network 130 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 105 received by the neural network 130 are (generated from) user inputs and the output sequences generated by the neural netw ork can be used to generate responses to the user inputs.

[0078] In implementations the neural network 130 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. In some cases, the neural network 130 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 105 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 130 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 130 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 130 can identity’ 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 170 generated by the neural network 130 can assign each pixel of the input image to a category from a set of categories. In some other cases, the neural network 130 is a neural network that is configured to perform an image generation task, where the input is a conditioning input, and the output 170 is a sequence of intensity7value inputs for the pixels of an image.

[0079] As one example, the task may be a neural machine translation task. For example, if the input to the neural network 130 is a sequence of text, e.g.. a sequence of words, phrases, characters, or word pieces, in one language, the output 170 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, word pieces, or characters of the first language, and the vocabulary7for the output tokens may be words, word pieces, or characters of the other language. As a particular example, the task may be a multi-hngual machine translation task, where a single neural network is configured to translate betw een 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 130 should translate the source language text.

[0080] Some implementations may be used for automatic code generation. For example, the input tokens may represent words, word pieces, 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 web page. As another example, the task may be an audio processing task. For example, if the input to the neural network 130 is a sequence representing a spoken utterance, the output generated by the neural network may be a score for each of a set of pieces of text, each score representing an estimated likelihood that the piece of text is the correct transcript for the utterance, e.g., a speech to text task. As another example, if the input sequence 105 to the neural network 130 is a sequence representing a spoken utterance, the network output 170 generated by the neural network 130 can indicate whether a particular word or phrase ("hotword") was spoken in the utterance. As another example, if the input sequence 105 to the neural network 130 is a sequence representing a spoken utterance, the network output 170 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.

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

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

[0083] 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 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. As another example, the task can be a text generation task, where the input sequence 105 is a sequence of text, and the network output 170 is another sequence of text, e.g., a completion of the input sequence 105 of text, a response to a question posed in the input sequence 105, 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 network output sequence 170 can be text that describes the input.

[0084] In some implementations the input sequence 105 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 105 represents compressed data, and the network output sequence 170 represents a decompressed version of the data, e.g., image data, text data, audio data, or any other type of data.

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

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

[0087] 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 position, joint 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 observ ations may similarly include one or more of the position, linear or angular velocity, 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.

[0088] 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 / acceleration 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.

[0089] 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 100 implementing the neural network 130 may be used to select actions in the simulated environment during training or evaluation of the system 100 and, after training, or evaluation, or both, are complete, the action selection policy may be deployed for controlling a real-world agent in the 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 100 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.

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

[0091] For example, a system 100 implementing the neural network 130 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 100 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 100 instructed the user to perform. Using the monitoring system, the system 100 can determine whether the task has been completed. The system 100 may identify actions which the user performs incorrectly with more than a certain probability. If so. when the system 100 instructs the user to perform such an identified action, the system 100 may warn the user to be careful. Alternatively, or additionally, the system 100 may leam not to instruct the user to perform the identified actions, i.e., ones which the user is likely to perform incorrectly.

[0092] 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 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 130, 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.

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

[0094] As another example, the task can be a genomics task, where the input sequence 105 is a sequence representing a fragment of a DNA sequence or other molecule sequence and the network output 170 is either an embedding of the fragment for use in a downstream task, e.g., by making use of an unsupervised learning technique on a data set of DNA sequence fragments, or an output for the downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, and so on.

[0095] In some cases, the machine learning task is a combination of multiple individual machine learning tasks, i.e., the system 100 is configured to perform multiple different individual machine learning tasks, e.g., tw o or more of the machine learning tasks mentioned above. For example, the system 100 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.

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

[0097] 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 130 includes both a computer vision neural netw ork and a text processing neural netw ork. That is, the target network output 170 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 open-vocabulary image classification, open-vocabulary object detection, image captioning, text-based image search, image-based retrieval, and so on.

[0098] As some further examples a multi-modal processing task can involve processing a text input including a sequence of text or audio data representing values of an audio w aveform, 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 170. The network output 170 may include any form of output appropriate to the task performed. For example, the network output 170 may include 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 170 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 condition or characteristic of the real world. For example, the video or audio may be captured from the real- world. The network output 170 can then provide an answer, e g., in natural language, to a question asked about the real-world input.

[0099] Also, or instead the network output 170 may include 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 170 may include non-textual action selection data for selecting an action to be performed by an agent controlled by the network output 170, 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 170 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 170 that performs the task.

[0100] 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 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. More generally, the task to be performed by the neural network 130 can be specified by the input sequence 105. As a particular example, the input sequence 105 can include a prompt or an instruction that specifies the task that is to be performed by the neural network 130. Optionally, in this example, the input sequence 105 also includes context for performing the task. 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.

[0101] As described above, the neural network 130 can be configured to receive any kind of input sequence 105 and to generate any kind of network output 170 based on the input sequence 105.

[0102] In some implementations, after training, the neural network 130 can be referred to as an auto-regressive neural network, i.e., because the neural network auto-regressively generates an output sequence of tokens. More specifically, the auto-regressively generated output 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.

[0103] For example, the neural network 130 can be an auto-regressive attention neural network that includes (i) one or more attention layer blocks that each apply a selfattention operation and (ii) an output subnetwork that processes an output of the last attention block to generate the score distribution.

[0104] In this example, the neural network 130 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.

[0105] 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. 11446, 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; and Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell. et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.

[0106] In some other implementations, after training, the neural network 130 can operate in a non-auto-regressive mode. In these implementations, the neural network 130 processes an input sequence of tokens in a single forward pass to generate a network output for the input sequence, e g., any of the classification or regression outputs described above.

[0107] More specifically, the neural network 130 can include one or more layer blocks.

[0108] The one or more layer blocks can include one or more mixture of depths attention layer blocks and, in some cases, multiple mixture of depths attention layer blocks. For example, the neural network 130 can include two mixture of depths attention layer blocks, e.g., a mixture of depths attention layer block B 145 and a mixture of depths attention layer block C 155. In some implementations, the one or more layer blocks can also include other layer blocks in addition to mixture of depths attention layer blocks. For example, the neural network 130 can also include a layer block A 135 along with the mixture of depths attention blocks.

[0109] The one or more layer blocks can be arranged in any appropriate arrangement. In some implementations, the one or more mixture of depths attention layer blocks and the one or more other layer blocks can be interleaved or alternating. In some implementations, the one or more mixture of depths attention layer blocks can be separated by one or more other layer blocks. For example, mixture of depths attention layer block B can be separated from mixtures of depths attention layer block C by one or more other layer blocks, e.g., layer block A. As another example, mixture of depths attention layer block C can be separated from another mixture of depths attention layer block two other layer blocks.

[0110] A layer block, as used in this specification, is a collection of one or more neural network layers, e.g., arranged sequentially such that the output of each neural network layer except the last is used as input to the next neural network layer in the layer block.

[0111] For example, the one or more neural network layers in the layer block can include one or more attention layers that each use an attention mechanism to apply a self-attention operation.

[0112] Generally, to apply the self-attention operation in an attention layer, each attention mechanism uses one or more attention heads.

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

[0114] As a particular example, in an attention head of a self-attention neural network layer, the attention mechanism may be configured to apply each of a query transformation, a key transformation, and a value transformation, to the attention layer input for each embedding of an input sequence X to derive a respective query' vector, keyvector, and value vector which are used to determine the updated embedding. The query, key, and value transformations can be any respective linear transformations or any other appropriate learned transformation. For example, the attention head can generate an updated embedding for each input position computing a weighted sum of the values, weighted by a similarity- function of the query for the input position to the corresponding key. The similarity function may comprise, e.g.. a dot product, cosine similarity-, or other similarity measure.

[0115] When the attention head uses position encoding, the application of the dot product attention function, the computation of the queries, keys, and values, or both depend on the relative or absolute positions of the embeddings corresponding to the queries, keys, and values within the input sequence. For example, an implementation of RoPE can involve determining, for a given query at a respective input position, a query rotation matrix that represents the absolute or relative position of the respective input position of the query, e.g., an index of the input position in the sequence; determining, for a given key at a respective input position, a key rotation matrix that similarly represents the absolute or relative position of the respective input position of the key. e.g., an index of the input position in the sequence, and multiplicatively combining the query rotation matrix, the key rotation matrix, the query (vector), and the key (vector), to determine a weight value between the query' and the key that is dependent on a relative distance between the position corresponding to the key and the position corresponding to the query'.

[0116] As another example, an implementation of ALiBi can involve adding a linear bias matrix to a weight determined from a combination of the key and the query.

[0117] When the attention head does not use position encoding, both the application of the dot product attention function and the computation of the queries, keys, and values, are independent of the relative or absolute positions of the embeddings corresponding to the queries, keys, and values within the input sequence.

[0118] Each query, key, value can be a vector that includes one or more vector elements. When there are multiple attention heads, the attention block 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.

[0119] 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., non-zero attention weights for a given position are computed only for positions that are within the local window of the given position.

[0120] In some cases, because the attention applied by the attention layers is causal, the system 100 can store, for any given attention mechanism and when generating the output for any given input position, the embeddings or the keys and values already computed for earlier input positions steps rather than re-computing the embeddings (or the keys and values) for earlier time steps.

[0121] Thus, in these cases, updating the respective embeddings for each of the input positions by applying an attention mechanism to the respective embeddings refers to updating the respective embedding for the last input position in the current input sequence using keys and values or embeddings for the other input positions that have been retrieved from memoty (e.g., from a “cache'’). Storing keys and values in a memory for later re-use will generally be referred as storing the keys and values in a “KV cache.”

[0122] In some implementations, some, or all of the layer blocks in the neural network can include other types of layers in addition to attention layers, e.g., normalization layers, residual connection layers, feedforward layers, and so on.

[0123] In some cases, some or all of the feedforward subnetworks within the layer blocks in the neural network are implemented as sparse mixture of experts (MoE) layers while in other cases some or all the feedforward layers are dense feedforward layers.

[0124] A mixture of experts (MOE) layer can use conditional logic to route tokens to one or more expert multi-layer perceptrons (MLPs). That is, the input tokens are routed to the most relevant one or more “experts” for processing. The mixture of experts (MOE) layers can include a routing function and a plurality of so-called expert feedforward neural networks (i.e., feedforward neural network layers that are different from one another, e.g., having different values of their respective trainable parameters and / or different architectures). MOE layers are described in, for example, Shazeer et al. arXiv: 1701.06538.

[0125] Differing from the mixture of depths attention layer block, the mixture of expert layer routes every7token to one or more of the expert feedforward neural networks in the layer. As will be described below, in the mixture of depths attention layer block, a token may not be routed through any portion of the layer, including attention mechanisms or feedforward networks of the layer block.

[0126] The neural network 130 can receive an input sequence 105 that includes a respective input token at each of one or more input positions and process the input sequence 105 to generate a network output 170.

[0127] At inference time, the neural network system 100 can operate in one of two modes: anon-autoregressive or an autoregressive mode.

[0128] In the non-autoregressive mode, the system 100 processes the input sequence of tokens in a single forward pass, i.e., at a single time step to generate the network output 170 for the time step. Thus, in this mode, the mixture of depths attention layer block receives a respective first embedding 137 for each of the input positions in the input sequence 105, e.g., generated as output by the preceding component of the neural network 130. That is, the system 100 receives a set of first embeddings 137 for the input positions of the input sequence 105 to process using the mixture of depths attention layer block. In the autoregressive mode, the system 100 generates an output sequence of tokens by, at each of multiple time steps, processing the most recently generated token in the output sequence to generate a new output token to be added to the output sequence. Thus, in this mode, at each time step, the mixture of depths attention layer block receives a first embedding 137 of the preceding output position of the input sequence 105, e.g., a first embedding of the token generated at the preceding time step. That is, the system 100 receives a single first embedding 137 to process using the mixture of depths attention layer block.

[0129] The mixture of depths attention layer block B 145 can then determine for any given token, e.g., first embedding 137. whether to perform the processing of the layers in the block B 145 or route the token around the block B 145.

[0130] The mixture of depths attention layer block B 145 can then process the first embedding(s) 137 to generative respective second embeddings 147 for each received first embedding.

[0131] The second embedding(s) 147 can then be processed by one or more additional layer blocks, including mixture of depths attention layer blocks, in the neural network 130 to generate the network output 170.

[0132] Examples of various network outputs 170 are described in detail above.

[0133] FIG. 2A shows an example mixture of depths attention layer block during non- autoregressive inference. The mixture of depths attention layer block 245 can include a routing model 210, an attention layer 225 and (optionally) a feedforward subnetwork 230.

[0134] When performing non-autoregressive inference, the system receives an input sequence and processes the input sequence in a single forward pass to generate the network output.

[0135] In some implementations, the system 100 can process the input sequence to generate a classification or regression output.

[0136] As another example, and as will be discussed in more detail below; the neural network 130 can operate in non-autoregressive mode during training of the neural network to perform auto-regressive inference. In this example, the network output during training is a respective probability distribution over a vocabulary of tokens for each of multiple positions in the input sequence.

[0137] During the processing of the input sequence, the mixture of depths attention layer block 245 can dynamically allocate the available compute by making decisions per-token, in each layer block, about where to spend compute from an available budget. The mixture of depths attention layer block 245 can achieve this by either applying a computation to a token or passing it around the block, leaving the token unchanged and therefore, saving compute.

[0138] As a first example, the mixture of depths attention layer block 245 can receive as input first embeddings 137 representing a first embedding for each of the input positions, e.g., generated as output by the preceding layer block of the neural network.

[0139] The mixture of depths attention layer block 245 can then select, using a first routing model 210 for the mixture of depths attention layer block 245 and from respective first embeddings for each of the input positions, a first proper subset of input positions 212.

[0140] The first routing model 210 for the mixture of depths attention layer can be a linear neural network layer or a multi-layer perceptron (MLP). The first routing model 210 can determine the tokens to be routed through the layer block using any appropriate routing method.

[0141] The mixture of depths attention layer block 245 can process the respective first embedding 137 for each of the input positions using the routing model 210 to generate a respective routing score for each of the input positions.

[0142] The mixture of depths attention layer block 245 can then select the first proper subset of the input positions 212 (i.e., a subset including one or more but not all of the input positions 212) based on the respective routing scores for the input positions.

[0143] For example, in expert-choice mixture of depths routing, the first proper subset 212 can be selected by selecting the k input positions having the highest routing scores, where k is a positive integer that is less than the total number of input positions in the sequence.

[0144] The expert-choice mixture of depths routing method is described in more detail with reference to FIG. 4.

[0145] The number of input positions in the first proper subset of input positions 212 can be fixed and is generally small relative to the total number of positions in the input sequence, e.g., the input sequence 105 of FIG. 1. For example, the number of positions in the first proper subset 212 can be one half of the total number, one quarter of the total number, or even one eighth of the total number.

[0146] In some implementations, the neural network system 100 can make the determination of the number of input positions to include in the first proper subset 212 based on a static compute budget for the system 100. By setting a static compute budget, the compute expenditure of the neural network system 100 can be predictable in sum total, while maintaining dynamic and context sensitive routing at the token-level. More specifically, the neural network system 100 can dynamically determine the top-k most context-essential tokens, or input positions, to include in the first proper subset 212.

[0147] The system 100 can determine the compute that is associated with the number of tokens routed through the block and determine the number of tokens to route so that the maximum number of tokens does not exceed the compute budget for the block. That is, the neural network system 100 can maintain a static compute budget while enhancing model performance by dynamically routing a static number of tokens through the block. A mixture of depths attention layer block can select different proper subsets of input positions 212 for different input sequences 105. That is. because the selection of input positions for the first proper subset of input position 212 is based on the first embeddings 137 and for different input sequences, a given mixture of depths attention layer block will receive different first embeddings 137, the given mixture of depths attention layer will generally select different proper subsets of the input positions for different input sequences.

[0148] The mixture of depths attention layer block 245 can then process the respective first embeddings for the first proper subset of the input positions 212 to generate respective second embeddings 147 for each of the input positions in the first proper subset 212. This processing can include applying an attention mechanism 225, e.g., selfattention, for the mixture of depths attention layer block to the respective first embeddings of each of the input positions in the first proper subset of input positions 212 to generate respective second embeddings for the input positions in the first proper subset 212.

[0149] To apply the attention mechanism 225, the mixture of depths attention layer block 245 can generate a set of queries, a set of keys, and a set of values for the input positions. A query, key, and value can be generated by applying any respective linear transformations or any other appropriate learned transformation to the first embedding of the input position for each input position.

[0150] The mixture of depths attention layer block 245 can apply any variety of variants of query-key-value attention, e.g., a dot product attention function or a scaled dot product attention function with or without positional encodings, to each of the first embeddings of the input positions in the first proper subset 212 to generate respective second embeddings for the input positions in the first proper subset 212. As the attention mechanism is only applied to the input positions in the first proper subset 212, e.g., input positions routed through the block, the set of queries is generated from each of the first embeddings of the input positions in the first proper subset 212. That is, the set of queries for the attention mechanism 225 can include the respective query7for each of the input positions in the first proper subset of the input positions 212 and does not include a query for any of the input positions not in the first proper subset.

[0151] In some implementations, the set of keys and the set of values can include a key and value for each input position in the input sequence, including those not in the first proper subset.

[0152] In some other implementations, the keys and values can be generated from the first proper subset of input positions 212 for the attention mechanism 225. That is, the set of keys and the set of values can include the respective key and value for each of the input positions in the first proper subset 212 of input position but does not include a key or value for any of the input positions in the input sequence.

[0153] In these implementations, the mixture of depths attention layer block 245 can apply7the attention mechanism using one or more sets of keys and one or more sets of values, where each set of keys and each set of values includes the respective key for each of the input positions in the first proper subset of the input positions and does not include a key for any of the input positions not in the first proper subset.

[0154] The processing of the embeddings for each of the input positions in the first proper subset 212 can also include applying a feedforward subnetwork 230 to each of the input positions 212 after applying the attention mechanism 225.

[0155] In some implementations, the feedforward network 230 of the mixture of depths attention layer block 245 can be implemented as sparse mixture of experts (MoE) layers while in other cases the feedforward network 230 can be a dense feedforward layer.

[0156] In some implementations, processing the respective first embeddings 137 for the first proper subset to generate the respective second embeddings for each input position in the first proper subset 212 can include processing the respective first embeddings 137 for the first proper subset 212 of the input position using the mixture of depths attention layer block 245 to generate initial second embeddings for each of the input positions in the first proper subset 212, e.g., as described above. Then, for each input position in the first proper subset 212, the mixture of depths attention layer block 245 can scale the initial second embedding by the routing score for the input position to generate the second embedding 235 for the input position.

[0157] The mixture of depths attention layer block 245 can also generate a respective second embedding for each of the input positions 214 not in the first proper subset without applying the attention mechanism for the mixture of depths attention layer block 245 to the other input positions 214 (or e.g., any other attention mechanism) not in the first proper subset. The processing for the input positions 214 not in the first proper subset can be generally much more computationally efficient than the processing for the input positions in the first proper subset 212. For example, the mixture of depths attention layer block 245 can, for each of the input positions 214 not in the first proper subset 212, set the second embedding for the input position equal to the first embedding for the input position. Thus, the mixture of depths attention layer block 245 can dynamically allocate FLOPs (or compute) to specific positions in a sequence by only selecting a fraction of the input positions for inclusion in the first proper subset 212. Because different mixture of depths attention layer blocks can select different first proper subsets, the neural network, e.g., the neural network 130 of FIG. 1. optimizes the allocation along the input sequence, e.g., the input sequence 105 of FIG. 1, for different layer blocks across the model depth. Thus, the described techniques expand FLOPs non-uniformly across the time and model depth dimensions. When the size of the proper subsets is fixed, compute expenditure is entirely predictable in sum total, but dynamic and context-sensitive at the token-level. This results in a high performing neural network that uses less compute relative to other approaches.

[0158] Prior to using the neural network, e.g., the neural network 130 of FIG. 1, to generate network outputs 170, the system 100 or another system trains the neural network 130 on training data.

[0159] As described above, in some implementations, during training, the input sequence is a training input sequence, and the network output can be a respective probability distribution over a set of input tokens for each of the input positions.

[0160] In some implementations, training the neural network can further include training the neural network and one or more routing models for autoregressive inference. That is, the one or more routing models can be trained to route tokens based on the probability, or the prediction, that they will be in the top-k tokens for the mixture of depths attention layer block. In some implementations, training the neural network further includes training the neural network and the respective first routing models, e.g.. the first routing model 210, for autoregressive inference for the mixture of depths attention layer blocks on a loss function based on the respective probability distributions.

[0161] The loss function can include a first term that measures, for each input position, 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.

[0162] The loss function can further include a respective auxiliary loss term for each mixture of depths attention layer block that measures a difference between (i) the respective routing scores for the input positions and (ii) a target set of routing scores that assigns each input position in the first proper subset a first score (e.g., 1) and each input position not in the first proper subset a second score (e.g., 0) different to the first score.

[0163] In some implementations, the training can include training a respective inferencetime routing model for the neural network 130 when in autoregressive mode. The inference time routing model will be described in further detail below with reference to FIG. 2B.

[0164] The inference time routing model can be configured to receive, for each input position, the respective first embedding for the input position and to generate as output a score that indicates a likelihood that the input position will be in the first proper subset. Training the respective inference-time routing model can include training the respective inference-time routing model on an objective (i.e., an objective function) that measures, for each input position, an error between the score for the input position and a target score for the input position that is equal to a first score if the input position is in the first proper subset and a second score if the input position is not in the first proper subset.

[0165] FIG. 2B shows an example mixture of depths attention layer block during autoregressive inference.

[0166] When performing autoregressive inference, the system receives an input sequence and processes the input sequence one token at a time to generate a network output.

[0167] As one example, the network output can be for a text, image, or other data item generation task. As another example, the network output can be temporal prediction, e.g., prediction of future values based on temporal data.

[0168] The mixture of depths attention layer block 245 can include a routing model 210. an attention layer 225 and optionally, a feedforward subnetwork 230. The mixture of depths attention layer block 245 can first receive as input a first embedding 137 for the preceding output position in the output sequence, e.g., the token generated at the immediately preceding time step.

[0169] The mixture of depths attention layer block 245 can then determine, using a routing model 210 for the mixture of depths attention layer block 245 and from the respective first embedding of the preceding output position, whether to include the preceding output position in a first proper subset 212 of the output positions for the mixture of depths attention layer block 245.

[0170] In some implementations, the routing model 210 can be the same first routing model described above with reference to FIG 2A, e.g., the routing model used during training / non-autoregressive inference.

[0171] In some implementations, the routing model 210 can be a different, inference-time routing model, e.g., one that has been trained as described with reference to FIG. 2A.

[0172] In some implementations, the inference-time routing model can be a linear neural network layer or an MLP.

[0173] The inference time routing model can use any appropriate routing method to determine the tokens to be routed through or around the mixture of depths attention layer block 245.

[0174] To determine whether to include the preceding output position in the first proper subset 212. the mixture of depths attention layer block 245 can process the respective first embedding for the preceding output position using the routing model 210 to generate a score for the proceeding output position and determine whether to include the preceding output position in the first proper subset 212 using the score.

[0175] In some implementations, the mixture of depths attention layer block can include the preceding output position in the first proper subset 212 when the score exceeds a threshold score.

[0176] In expert-choice mixture of depths routing, the mixture of depths attention layer block can include the preceding output position in the first proper subset 212 when the score is a top-k score, e.g.. the routing model should prioritize routing of the preceding output position through the mixture of depths attention layer block 245. However, to determine whether a given token’s routing score is among the top-k for the input sequence, e.g., that the token should be routed through the block, depends on the values of the routing score for tokens that come after it, which cannot be accessed when autoregressively sampling. In some implementations, the routing model for the mixture of depths attention layer block 245 can predict whether the preceding output position will be in a first proper subset of output positions when generating the output position during auto-regressive sequence generation. That is, the model can predict the likelihood that the given output position will be in the top-k of output positions in the mixture of depths attention layer block among all of the output positions in the output sequence (or in some cases, in a combined sequence that includes the input sequence and output sequence.

[0177] In response to determining to include the preceding output positions in the first proper subset 212 of the output positions, the mixture of depths attention layer block 245 processes the respective first embedding 137 for the preceding output position to generate a respective second embedding 235 for the preceding output position, which includes applying an attention mechanism 225 for the mixture of depths attention layer block 245 to the respective first embedding 217 for the preceding output position.

[0178] The attention mechanism can be the same attention mechanism described above with reference to FIG. 2A or a different attention mechanism.

[0179] In some implementations, the attention mechanism for the mixture of depth attention layer block can be applied to (i) the respective first embedding for the preceding output position and (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset.

[0180] In some implementations where the output sequence is being generated in response to the input sequence 105, the attention mechanism for the mixture of depth attention layer block can be applied to (i) the respective first embedding for the preceding output position, (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset and (iii) respective first embeddings in a first proper subset of input positions in the input sequence.

[0181] In some implementations, the attention mechanism for the mixture of depth attention layer block can be applied to (i) the respective first embedding for the preceding output position and (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset but not to (iii) respective first embeddings for output positions in the output sequence that were previously not selected for inclusion in the first proper subset.

[0182] In some implementations, to process the respective first embeddings to generate respective second embeddings for the preceding output position, the attention layer block can apply the attention mechanism to the respective first embedding of the preceding output position using one or more sets of keys, where each set of key includes a respective key for each of the output positions that has been selected for inclusion in the first proper subset of the input positions and does not include a key for any of the output positions that have not been selected for inclusion in the first proper subset.

[0183] The processing of the first embedding 137 for the preceding output position if in the first proper subset 212 can also include applying a feedforward subnetwork 230 to the first embedding 137 after applying the attention mechanism 225.

[0184] In some implementations, the feedforward network 230 of the mixture of depths attention layer block can be implemented as sparse mixture of experts (MoE) layers while in other cases the feedforward network 230 can be a dense feedforward layer.

[0185] In some implementations, to process the respective first embedding 137 for the preceding output position using the mixture of depths attention layer block 245 to generate a respective second embedding 235 for the preceding output position, the mixture of depths attention layer block 245 can first process the respective first embedding 137 for the preceding output position using the mixture of depths attention layer block 245 to generate a respective initial second embedding for the preceding output position. Then, the mixture of depths attention layer block 245 can scale the initial second embedding for the input position by a routing score for the input position to generate the second embedding 235 for the input position. The routing score can be the score generated by the routing model for the input position.

[0186] In response to determining not to include the preceding output position in the first proper subset 212 of the output positions, the mixture of depths attention layer block 245 can generate the respective second embedding 235 for the preceding output position without applying the attention mechanism 225 for the mixture of depths attention layer block 245 to the respective first embedding 137 for the preceding output position. For example, in some implementations, the mixture of depths attention layer block 245 can set the second embedding 235 for the preceding output position equal to the first embedding 137 for the preceding output position. That is. the mixture of depths attention layer block 245 can "route" the first embedding 137 around the block 264 by not applying the layers of the block 245, e.g., the attention mechanism 225 and feedforward network 230. FIG. 3 shows example graphs representing routing decisions of different approaches in contrast to the routing decisions of a neural network including one or more mixture of depths attention layer blocks.

[0187] The three approaches compared in FIG. 3 are a traditional approach, an “early exit” approach and the mixture of depths approach described in this specification.

[0188] To begin, neural network 310 showcases a traditional approach that routes every input position in an input sequence through every layer block in the neural network 310. For example, for both layer block A 315 and layer block B 317, as well as every other layer block in the neural network 310, input position C 311 and input position D 313 are routed through both layer blocks. As traditional layer blocks route every' input position in the input sequence through the block, neural networks including traditional layer blocks are not computationally efficient.

[0189] The neural networks 330 and 350 showcase two different routing methods for conditional computation or learned mechanisms to determine when and how to expend computation: an early exit approach and a trained mixture of depths approach.

[0190] The early exit approach focuses on learning to decide when to end computation on a given token, allowing the token to skip any remaining layer blocks after the exit decision is made. The neural network 330 demonstrates the routing decisions of “early exit” layer blocks that engages blocks serially, e.g., input positions cannot route through a later block if not routed through the previous block. That is. the input positions exit the neural network 330 “early” by not routing through later layer blocks as demonstrated in the graph of neural network 330. For example, input position C 331 routes through layer block A 335 but takes an “early exit” and does not route through layer block B 337 or any of the other layer blocks in the neural network 330 whereas for input position D 333 routes through both layer block A 335 and layer block B 337 as well as every other layer block in the neural network 330.

[0191] Similar to an early exit approach, mixture of depths can choose whether to route a token through a layer block. However, in a mixture of depths approach, a token can skip middle layer blocks, then be routed through later layer blocks. The neural network 350 is an example of a neural network that includes one or more trained mixture of depths attention layer blocks as described in this specification. The neural network 350 showcases the routing decisions of the layer blocks for an input sequence of 64 input positions. As depicted in the graph, input positions can be processed by later layer blocks, despite passing through relatively few total blocks throughout the model’s depth. For example, for input position A 351, the input position 351 is routed through layer block A 355, not layer block B 357, and then routed through layer block C 359. As another example, input position B 353 is not routed through layer block A 355, is routed through layer block B 357, is not routed through layer block C 359, and then is later routed through layer block D 361. As another example (not depicted in FIG. 3), an input position can be routed through every other layer block in the neural network.

[0192] As shown in the routing decisions of neural network 350, one or more mixture of depths attention blocks can increase computational efficiency by routing tokens through particular layer block’s attention mechanisms and not every' layer block’s attention mechanism. Since some tokens take a second route around the block, the neural network 350 has a smaller total FLOP footprint compared to the traditional neural network 310. Unlike the early exit approach seen in neural network 330, because tokens can skip middle layer blocks and then be routed through a later layer block in the neural network, the tokens of input positions in neural network 350 can be updated via self-attention with tokens that have gone through all the middle layer blocks, benefitting from the attention mechanisms of layer blocks without necessarily having to be processed by that layer block.

[0193] FIG. 4 shows different example routing methods for a mixed depth attention layer block.

[0194] The three routing methods compared in FIG. 4 are token-choice routing, expertchoice routing, and expert-choice mixture of depths routing.

[0195] For a mixture of depths attention layer block, learned routing is preferable. The layer block should be able to learn which tokens require more or less processing than others, expending fewer FLOPS and being more computationally efficient.

[0196] The routing method 410 showcases token-choice routing, where tokens are routed to the computational path of their choice. A routing model produces per-token probability distributions across computational paths and then the tokens are shuttled to the path they prefer, e.g.. the path with the highest probability. As an example, in the mixture of depths attention layer block the computational paths can be a path through the layer block, and a path around the layer block. In the specific example depicted in FIG. 4, there are three computational paths. For example, token A 411 and token B 413 choose computational path A 421 while token C 415 choose computational path B 423. If a given path exceeds its capacity, then surplus tokens must be dropped. The exact token that is ultimately dropped depends on the precise implementation in the underlying code for the routing model. For example, priority is often given to those tokens that come earlier in the sequence or batch order. As a specific example, token D 417 has also chosen computational path A 421 and is dropped because too many tokens have chosen computational path A 421 and token D 417 is later in the sequence than token A 411 or token B 413. One of the problems with token choice routing is that the routing can have load balancing problems since there is no guarantee that tokens divide themselves appropriately between the possible paths, e.g.. even- token prefers one computational path.

[0197] In expert-choice routing, as seen in routing method 430, rather than having tokens choose the path they prefer, each computational path instead chooses the top-k tokens based on the tokens’ preference, e.g., routing weight. For example, computational path A 441 can choose the top-2 tokens according to their routing weight: token A 431 and token F 437. Here, tokens can be dropped if they are not among the top-k with respect to any of the computational paths and some tokens can be routed to multiple paths. For example, token C 435 is not in the top-2 tokens for any of the computational paths and is dropped while token F 437 is routed to both computational path A 441 and computational path C 445. As opposed to token choice routing where not every computational path has the same number of tokens, this routing method ensures a perfect load balance since k tokens are guaranteed to be shuttled to each path.

[0198] To determine the routing weights, each token can be processed by a routing model to produce a scalar weight. For a set of token embeddings in a sequence of length S for a given layer / , represented as X1= x | i is an integer, 1 < i < S, the router weight for a given token embedding is a scalar produced as a result of a linear projection: ri = w^xt

[0199] For a mixture of depths attention layer block, expert-choice mixture of depths routing can be used. Essentially, expert choice routing is used, but there is only a single “expert,” splitting the tokens into two mutually exclusive sets: a first set in which tokens are routed through the layer block, and a second set of input positions in which tokens are routed around the layer block. That is, the expert-choice mixture of depths routing can leverage the knowledge that tokens will be dropped if k is less than the sequence length to split the tokens into two sets: the top-k tokens that will be routed through the computational path, e.g.. through the layer block, and the dropped tokens that will be routed around the computational path, e.g., around the layer block. For example, as seen in routing method 450, computational path A 461, the only computational path, token A 451 and token F 457 are both chosen as the top-2 tokens while the four other tokens are dropped. In this example, token A 451 and token F 457 are tokens part of the first proper set that is routed through the mixture of depths attention layer block while the rest of the tokes are part of the second set that is routed around mixture of depths attention layer block. This routing model efficiently splits the tokens into two sets with one set of tokens to be routed around the computation path, allowing the model to expend fewer FLOPs in a given forward pass. Also, because the top-k operation depends on the router weights, the expert-choice mixture of depths routing allows for relative routing weights to help determine which tokens most need the block’s computations. That is, the routing method can try to ensure that the most critical tokens are among the top-k by setting their weights appropriately.

[0200] In some implementations, an inference time routing model, e.g., a small auxiliary’ MLP predictor, that receives the same inputs as the routing model, but whose output is a prediction whether that token will be among the top-k or not in the sequence, can be introduced to utilize the expert route mixture of depths approach in an autoregressive model as described with reference to FIG. 2B.

[0201] FIG. 5 is a flow diagram of an example process 500 of the example neural network system. The below example process 500 can be used during training or non-auto- regressive inference.

[0202] For convenience, the process 500 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g.. the neural network system 100 depicted in FIG 1, appropriately- programmed in accordance with this specification, can perform the process 500.

[0203] The system can receive an input sequence including a respective input token at each of one or more input positions (step 502).

[0204] The system can then process the input sequence using a neural network to generate a network output (step 504).

[0205] In some implementations, the system 100 can process the input sequence to generate a classification or regression output.

[0206] As another example, and as will be discussed in more detail below, the neural network can operate in non-autoregressive mode during training of the neural network to perform auto-regressive inference. In this example, the network output during training is a respective probability distribution over a vocabulary of tokens for each of multiple positions in the input sequence.

[0207] The below steps are sub steps of the step 504, further describing how the neural network processes the input sequence to generate a network output.

[0208] The neural network includes one or more layer blocks that include one or more mixture of depths attention layer blocks.

[0209] In some implementations, the one or more layer blocks can further include one or more attention layer blocks that are not mixture of depths attention layer blocks.

[0210] For each mixture of depths attention layer block, the system can first select, using a first routing model for the mixture of depths attention layer block and from respective first embeddings for each of the input positions, a first proper subset of input positions (step 506). The system can select the first proper subset of the input positions by processing the respective first embedding for each of the input positions using the first routing model for the mixture of depths attention layer block to generate a respective routing score for each of the input positions. The system can then select the first proper subset of the input positions based on the respective routing scores for the input positions.

[0211] In some implementations, the system can select the k input positions having the highest routing scores to be in the first proper subset, where k is an integer that is less than the total number of input positions in the sequence.

[0212] In some implementations, the first routing model is a linear neural network layer or a multi-layer perceptron (MLP).

[0213] The system can then process the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block to generate respective second embeddings (step 508). This processing can include applying an attention mechanism, e.g., self-attention, for the mixture of depths attention layer block to the respective first embeddings of each of the input positions in the first proper subset of input positions to generate respective second embeddings for the input positions in the first proper subset. This processing can further include applying a feedforward subnetwork to each of the input positions after applying the attention mechanism.

[0214] In some implementations, the feedforward network can be implemented as a sparse mixture of experts (MoE) layer while in some other cases the feedforward network can be a dense feedforward layer.

[0215] The system can generate a respective second embedding for each of the input positions not in the first proper subset (step 510). In some implementations, the system can generate the respective second embedding for the preceding output position without using the attention layer or the feedforward subnetwork. For example, the respective second embedding for the preceding output position can be generated without applying the attention mechanism by setting the second embedding for the preceding output position equal to the first embedding for the preceding output positions.

[0216] FIG. 6 is a flow diagram of sub-steps of step 508 of the process 500 of FIG. 6.

[0217] The system can include one or more mixture of depths attention layer blocks that include (i) an attention layer that applies the attention mechanism for the mixture of depths attention block and (ii) a feedforward subnetwork (step 602).

[0218] In some implementations, the feedforward subnetwork is a mixture of experts (MoE) subnetwork that includes a routing function, e.g., another routing model, and one or more expert feedforward neural networks.

[0219] In some implementations, the feedforward subnetwork is a dense feedforward layer.

[0220] The system can process the respective first embedding for the first proper subset of the input positions using the mixture of depths attention layer block to generate respective second embeddings for each of the input positions in the first proper subset (step 604).

[0221] In some implementations, processing the respective first embeddings for the first proper subset to generate the respective second embeddings for each input position in the first proper subset can further include, for each input position in the first proper subset, the mixture of depths attention layer block scaling the initial second embedding by the routing score for the input position to generate the second embedding for the input position.

[0222] The below steps (step 606-608) are sub steps of step 604 and further detail the generation of the respective second embeddings for each of the input positions in the first proper subset.

[0223] The system can process the respective first embeddings for each of the input positions in the first proper subset of the input positions using the attention layer to generate a respective attended embedding (i.e., respective embeddings generated as outputs from the attention layer) for each of the input positions in the first proper subset (step 606).

[0224] To apply the attention mechanism, the system can generate a set of queries, a set of keys, and a set of values for the input positions. A query, key, and value can be generated by applying any respective linear transformations or any other appropriate learned transformation to the first embedding of the input position for each input position.

[0225] The system can apply any variety of variants of query-key-value attention, e.g., a dot product attention function or a scaled dot product attention function with or without positional encodings, to each of the first embeddings of the input positions in the first proper subset to generate respective attended embeddings for the input positions in the first proper subset.

[0226] As the attention mechanism is only applied to the input positions in the first proper subset, e.g., input positions routed through the block, the set of queries is generated from each of the first embeddings of the input positions in the first proper subset. That is. the set of queries for the attention mechanism can include the respective query for each of the input positions in the first proper subset of the input positions and does not include a query for any of the input positions not in the first proper subset.

[0227] In some implementations, the set of keys and the set of values can include a key and value for each input position in the input sequence, including those not in the first proper subset.

[0228] In some other implementations, the keys and values can be generated from the first proper subset of input positions for the attention mechanism. That is, the set of keys and the set of values can include the respective key and value for each of the input positions in the first proper subset of input position but does not include a key or value for any of the input positions in the input sequence.

[0229] In these implementations, the mixture of depths attention layer block can apply the attention mechanism using one or more sets of keys and one or more sets of values, where each set of keys and each set of values includes the respective key for each of the input positions in the first proper subset of the input positions and does not include a key for any of the input positions not in the first proper subset.

[0230] The system can process the respective attended embeddings for each of the input positions for the first proper subset of the input positions using the feedforward subnetwork to generate a respective third embedding for each of the input positions in the first proper subset (step 608).

[0231] The system can process the respective attended embeddings by applying any appropriate transformation and / or activation functions to the attended embeddings to generate respective third embeddings. In some implementations, the transformations can include linear transformations, e.g., applying weights and biases.

[0232] In some implementations, the activation functions can include a rectified linear unit (ReLU) activation function, a sigmoid activation function, a softmax activation function, or any other appropriate activation function.

[0233] In some implementations, to generate the respective second embeddings for each of the input positions in the first proper subset of input positions from the respective third embedding, the system can apply a residual connection to the respective third for each of the inputs positions in the first proper subset to generate the respective initial second embeddings.

[0234] FIG. 7 is a flow diagram of an example process 700 of autoregressively generating an output sequence using the example neural network system. The output sequence of the system can include a respective output token at each of one or more input positions.

[0235] For convenience, the process 700 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g.. the neural network system 100 depicted in FIG 1. appropriately programmed in accordance with this specification, can perform the process 700.

[0236] For each of the one or more output positions, the system can first receive the preceding output token generated for the preceding output position (step 702).

[0237] The system can then process the preceding output token using a neural network to generate the output token at the output position (step 704).

[0238] The neural network can include one or more layer blocks that include one or more mixture of depths attention layer blocks.

[0239] In some implementations, the one or more layer blocks further includes one or more attention layer blocks that are not mixture of depths attention layer blocks.

[0240] The below steps describe the steps of processing the preceding output token for each mixture of attention layer block.

[0241] The system can determine, using a routing model for the mixture of depths attention layer block and from a respective first embedding for the preceding output position, whether to include the preceding output position in a first proper subset of the output positions for the mixture of depths attention layer block (step 706).

[0242] In some implementations, the routing model for the mixture of depths attention layer is a linear neural network layer or a multi-layer perceptron (MLP). In some implementations, the routing model is the first routing model, e.g., the routing model used during training or non-auto-regressive inference.

[0243] In some implementations, the routing model can be an inference time-routing model.

[0244] In some implementations, determining whether to include the preceding output position in the first proper subset can include processing the respective first embedding for the preceding output position using the routing model to generate a score for the preceding output position. The generated score can then be used to determine whether to include the preceding output position.

[0245] In some implementations, the preceding output position can be included in the first proper subset when the score exceeds a threshold score.

[0246] In some implementations, the system can include the preceding output position in the first proper subset when the score is a top-k score, e.g., the routing model should prioritize routing of the preceding output position through the mixture of depths attention layer block. However, to determine whether a given token’s routing score is among the top-k for the input sequence, e.g., that the token should be routed through the block, depends on the values of the routing score for tokens that come after it, which cannot be accessed when autoregressively sampling.

[0247] In some implementations, the routing model can predict whether the preceding output position will be in a first proper subset of output positions when generating the output position during auto-regressive sequence generation. That is, the model can predict the likelihood that the given output position will be in the top-k of output positions in the mixture of depths attention layer block among all of the output positions in the output sequence (or in some cases, in a combined sequence that includes the input sequence and output sequence.

[0248] In this specification, the score, generated by the routing model, for the input position can also be referred to as the routing score for the input position.

[0249] In response to determining to include the preceding output position in the first proper subset of the output positions, the system can process the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate a respective second embedding for the preceding output position (step 708).

[0250] In some implementations, processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate a respective second embedding can further include scaling the initial second embedding for the input position by a routing score for the input position to generate the second embedding for the input position.

[0251] In response to determining not to include the preceding output position in the first proper subset of the output positions, the system can generate the respective second embedding for the preceding output position without applying the attention mechanism for the mixture of depths attention layer block (step 710).

[0252] In some implementations, the system can generate the respective second embedding for the preceding output position without using the attention layer or the feedforward subnetwork. For example, the respective second embedding for the preceding output position can be generated without applying the attention mechanism by setting the second embedding for the preceding output position equal to the first embedding for the preceding output positions.

[0253] FIG. 8 is a flow' diagram of sub-steps of step 708 of the process 700 of FIG. 8. The below steps further describe the process 700 of using the mixture of depths attention block to process the respective first embedding of the preceding output position.

[0254] As described above, the system can include one or more mixture of depths attention layer blocks. The mixture of depths attention layer blocks can include (i) an attention layer and (ii) a feedforward subnetwork (step 802).

[0255] In some implementations, the feedforward subnetwork is a mixture of experts (MoE) subnetw ork that includes a routing function, e.g., another routing model, and one or more expert feedforward neural netw orks.

[0256] The system can process the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate the respective second embedding for the preceding output position (step 804).

[0257] The below steps (step 806 and 808) are sub-steps of step 804 and further detail the generation of the respective second embedding for the preceding output position.

[0258] The system can first process the respective first embedding for the preceding output position using the attention layer to generate a respective attended embedding for the preceding output position (step 806).

[0259] To generate the respective second embedding using the mixture of depths attention layer block, the system can apply the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position. In some implementations, applying the attention mechanism can include applying the attention mechanism for the mixture of depths attention layer block to (i) the respective first embedding for the preceding output position and (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset.

[0260] In some implementations where the output sequence is being generated in response to the input sequence, the attention mechanism for the mixture of depth attention layer block can be applied to (i) the respective first embedding for the preceding output position, (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset and (iii) respective first embeddings in a first proper subset of input positions in the input sequence.

[0261] In some implementations, the attention mechanism for the mixture of depth attention layer block can be applied to (i) the respective first embedding for the preceding output position and (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset but not to (iv) respective first embeddings for output positions in the output sequence that were previously not selected for inclusion in the first proper subset.

[0262] The system can then process the respective attended embedding for the preceding output position using the feedforward subnetwork to generate a respective third embedding for the preceding output position (step 808).

[0263] The system can process the respective attended embeddings by applying any appropriate transformation and / or activation functions to the attended embeddings to generate respective third embeddings.

[0264] In some implementations, the transformations can include linear transformations, e.g., applying weights and biases.

[0265] In some implementations, the activation functions can include a rectified linear unit (ReLU) activation function, a sigmoid activation function, a softmax activation function, or any other appropriate activation function.

[0266] To generate the respective second embedding for the preceding output position from the respective third embedding, the system can apply a residual connection to the respective third for the preceding output position to generate the respective initial second embedding for the preceding output position. FIG. 9 illustrates an example of the routing decisions of a neural network, including one or more mixture of depths attention layer blocks, that is trained according to the techniques described in this specification.

[0267] The diagram 903 shows the routing decision of a trained neural network 950 that includes interleaving routing blocks, e g., blocks that alternate between 100% routing blocks and mixture of depths attention layer blocks that route at a lower capacity. As a specific example, the diagram 103 showcases mixture of depths attention layer blocks that route 12.5% of tokens through the block.

[0268] As seen in diagram 903, the network can preferentially route certain tokens through more layer blocks along the model’s depth relative to other tokens.

[0269] For example, input token 921 is not routed through any of the mixture of depths attention layer blocks and is only routed through the layer blocks that routes every token through the block. In comparison, input token 923 is routed through almost every mixture of depths attention layer block in the neural network 950.

[0270] That is, as seen in the diagram 903, the neural network 950 is able to allocate more compute to certain tokens that need more processing than others. In doing so. the neural network 950 can limit compute to tokens that do not need further processing.

[0271] The neural network 950 can determine the tokens to process through the mixture of depths attention layer blocks using the routing weights of the tokens. The router weights of the tokens routed through the mixture of depth attention layer block or around the layer block are shown in graph 907. As demonstrated in the graph, tokens with a router weight above 0.5 were routed through the mixture of depths attention layer block.

[0272] Despite aggressive routing (e.g., not routing 87.5% of tokens through a block), the trained neural network 950 can achieve performance improvements relative to baselines.

[0273] Thus, the trained neural network 950 can dynamically allocate compute to tokens that need more processing and limit compute to tokens that need less processing while generating the same or higher quality' outputs than conventional networks in significantly fewer FLOPs.

[0274] FIG. 10 shows an example of the performance of a trained neural network when operating in auto-regressive mode.

[0275] Graph 1005 compares the performance of the trained neural network in non- autoregressive mode (using top-k routing) and in autoregressive mode using the predictor-based approach that uses the separate inference-time routing model for autoregressive inference . Graph 1005 showcases a minimal performance degradation during autoregressive inference when using the inference-time routing model, because of the accuracy of training the routing model on the predictor based approach. As seen in graph 1009, the top-k prediction accuracy, i.e., the accuracy of the routing model in predicting whether any given token should be part of the top-k, can reach upwards of 97% accuracy during training.

[0276] Similar to described above with FIG. 9, there are neural network variants (e.g., different number of FLOPs), that are better performing than the baseline, while requiring fewer FLOPs per forward pass.

[0277] Thus, during autoregressive inference, the trained neural network can allocate more compute to certain tokens that need more processing than others, while generating the same or higher quality outputs that conventional networks with fewer FLOPs.

[0278] FIG. 11 shows an example modification of a mixture of experts model to include a mixture of depths technique.

[0279] While a mixture of depths model can be used on its own, as an alternative, the mixture of depths technique described in this specification can be used to modify a mixture of experts model. That is, the mixture of depths technique can be naturally integrated with mixture of experts (MoE) models to make MoE models more robust.

[0280] As shown in graph 1110, the improvements of mixture of depths described in this specification can be compounded with those of mixture of experts to create a mixture of depths and experts (MoDE) approach that reduces normalized loss over both a baseline, traditional approach, or the regular mixture of experts approach.

[0281] The mixture of experts model can be modified by the mixture of depths approach in one or more ways, including a staged MoDE approach and an integrated MoDE approach.

[0282] A staged MoDE is a mixture of depths and experts technique that first implements mixture of depth machinery prior to mixture of experts machinery. That is, tokens can be routed around or towards the block prior to the self-attention step. For example, in layer 1130. the staged MoDE model can have a routing model 1135 before the self-attention step that routes tokens either through the layer or route around the layer 1139. For the tokens routed through the layer, the tokens can go through the self-attention step and then be routed by a second routing model 1145 that routes the tokens to the mixture of experts feedforward subnetworks 1147. Essentially, the mixture of depths technique splits the tokens first and then the chosen routed tokens are later sent to the best expert using MoE. This approach is advantageous as it allows for tokens to skip the self-attention step using mixture of depth (MoD).

[0283] An integrated MoDE is a mixture of depths and experts technique which uses one routing operation to funnel tokens to either experts or no-op operations. That is, a “no-op” expert is integrated into the conventional MLP experts as the route around option for tokens not being processed by the layer. For example, in layer 1150, all tokens are processed by the self-attention step of the layer and then sent to a routing model 1155. The routing model 1155 can receive all tokens for the input sequence and can funnel tokens to the chosen expert in the mixture of experts 1157, including the no-op expert 1159. The no-op expert 1159 can represent the route around option in the mixture of depths portion of the MoDE approach. This approach is advantageous as it simplifies the routing machinery, e.g., only one routing model is needed, instead of two. It also allows tokens to explicitly learn to choose the residual, “no-op” path around the experts, instead of preferring an expert but being dropped in a convention MoE model.

[0284] 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 its 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.

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

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

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

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

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

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

[0291] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

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

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

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

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

[0296] 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 sen’ er, 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.

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

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

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

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

[0301] What is claimed is:

Claims

CLAIMS1. 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 mixture of depths attention layer blocks, and wherein processing the input sequence comprises, for each mixture of depths attention layer block: selecting, using a first routing model for the mixture of depths attention layer block and from respective first embeddings for each of the input positions, a first proper subset of the input positions for the mixture of depths attention layer block; processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block to generate respective second embeddings for each of the input positions in the first proper subset, comprising applying an attention mechanism for the mixture of depths attention layer block to the respective first embeddings of each of the input positions in the first proper subset of the input positions; and generating a respective second embedding for each of the input positions not in the first proper subset without applying the attention mechanism for the mixture of depths attention layer block to the input positions not in the first proper subset.

2. The method of claim 1, wherein selecting, using a first routing model and from the respective first embeddings for each of the input positions, a first proper subset of the input positions to be processed by the mixture of depths attention layer block comprises: processing the respective first embedding for each of the input positions using the first routing model for the mixture of depths attention layer block to generate a respective routing score for each of the input positions; and selecting the first proper subset of the input positions based on the respective routing scores for the input positions.

3. The method of claim 2, wherein selecting the first proper subset of the input positions based on the respective routing scores comprises: selecting the k input positions having the highest routing scores, wherein k is an integer that is less than the total number of input positions in the sequence.

4. The method of any one of claim 2 or claim 3, wherein processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block to generate the respective second embeddings for each of the input positions in the first proper subset comprises: processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block to generate respective initial second embeddings for each of the input positions in the first proper subset; and for each input position in the first proper subset, scaling the initial second embedding for the input position by the routing score for the input position to generate the second embedding for the input position.

5. The method of any preceding claim, wherein: each mixture of depths attention layer block comprises (i) an attention layer that applies the attention mechanism for the mixture of depths attention block and (ii) a feedforward subnetw ork, and processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block to generate respective second embeddings for each of the input positions in the first proper subset comprises: processing the respective first embeddings for each of the input positions in the first proper subset of the input positions using the attention layer to generate a respective attended embedding for each of the input positions in the first proper subset; and processing the respective attended embeddings for each of the input positions for the first proper subset of the input positions using the feedforward subnetwork to generate a respective third embedding for each of the input positions in the first proper subset.

6. The method of claim 5. when dependent on claim 4, wherein processing the respective first embeddings for the first proper subset of the input positions using the mixture of depths attention layer block further comprises: applying a residual connection to the respective third embeddings for each of the input positions in the first proper subset to generate the respective initial second embeddings.

7. The method of claim 5 or claim 6, wherein generating the respective second embeddings for each of the input positions not in the first proper subset without applying the attention mechanism for the mixture of depths attention layer block to the input positions not in the first proper subset comprises: generating the respective second embeddings for each of the input positions not in the first proper subset without processing any embeddings for the input positions not in the first proper subset using the attention layer or the feedforward subnetwork.

8. The method of any preceding claim, wherein generating the respective second embeddings for each of the input positions not in the first proper subset comprises: for each of the input positions not in the first proper subset, setting the second embedding for the input position equal to the first embedding for the input position.

9. 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 and the respective first routing models for the mixture of depths attention layer blocks on a loss function based on the respective probability distributions.

10. The method of claim 9, wherein: the loss function includes a first term that measures, for each input position, 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.

11. The method of claim 9 or claim 10, when dependent on claim 2, wherein: the loss function includes a respective auxiliary loss term for each mixture of depths attention layer block that measures a difference between (i) the respective routing scores for the input positions and (ii) a target set of routing scores that assigns each input position in the first proper subset a score of 1 and each input position not in the first proper subset a score of 0.

12. The method of claim 9 or claim 10, further comprising: training a respective inference-time routing model for each mixture of depths attention layer, wherein the respective inference-time routing model is configured to receive, for each input position, the respective first embedding for the input position and to generate as output a score that indicates a likelihood that the input position will be in the first proper subset, and wherein training the respective inference-time routing model comprises training the respective inference-time routing model on an objective that measures, for each input position, an error between the score for the input position and a target score for the input position is equal to a first score if the input position is in the first proper subset and a second score if the input position is not in the first proper subset.

13. The method of any preceding claim, wherein the first routing model for the mixture of depths attention layer is a linear neural network layer or a multi-layer perceptron (MLP).

14. The method of any preceding claim, wherein the plurality of layer blocks further comprises one or more attention layer blocks that are not mixture of depths attention layer blocks.

15. The method of any preceding claim, wherein the plurality of layer blocks comprises a plurality of mixture of depths attention layer blocks.

16. The method of any preceding claim, when dependent on claim 5, wherein, for one or more of the mixture of depths attention layer blocks, the feedforward subnetwork is a mixture of experts (MoE) subnetwork that includes a routing function and a plurality of expert feedforward neural networks.

17. A method performed by one or more computers and for auto-regressively generating an output sequence comprising a respective output token at each of plurality’ of output positions, the method comprising, for each of the plurality of output positions: receiving a preceding output token generated for the preceding output position; and processing the preceding output token using a neural network to generate the output token at the output position, wherein: the neural network comprises a plurality of layer blocks that include one or more mixture of depths attention layer blocks, and wherein processing the preceding output token comprises, for each mixture of depths attention layer block: determining, using a routing model for the mixture of depths attention layer block and from a respective first embedding for the preceding output position, whether to include the preceding output position in a first proper subset of the output positions for the mixture of depths attention layer block; and in response to determining to include the preceding output position in the first proper subset of the output positions, processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate a respective second embedding for the preceding output position, comprising applying an attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position; and in response to determining not to include the preceding output position in the first proper subset of the output positions, generating the respective second embedding for the preceding output position without applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position.

18. The method of claim 17, wherein determining whether to include the preceding output position in the first proper subset comprises: processing the respective first embedding for the preceding output position using the routing model to generate a score for the preceding output position; and determining whether to include the preceding output position in the first proper subset using the score.

19. The method of claim 18, wherein determining whether to include the preceding output position in the first proper subset using the score comprises including the preceding output position in the first proper subset when the score exceeds a threshold score.

20. The method of any one of claim 18 or claim 19, wherein processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate a respective second embedding for the preceding output position comprises: processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate a respective initial second embedding for the preceding output position; and scaling the initial second embedding for the input position by a routing score for the input position to generate the second embedding for the input position.

21. The method of claim 20, wherein the routing score for the input position is the score for the input position.

22. The method of claim 21, wherein the routing model is the first routing model of any one of claims 1-16.

23. The method of claim 20 or claim 21, wherein the routing model is the inferencetime routing model of claim 12.

24. The method of claim 20 or claim 23, wherein processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate a respective second embedding for the preceding output position comprises: processing the respective first embedding for the preceding output position using a first routing model for the mixture of depths attention layer block to generate the routing score for the preceding output position.

25. The method of any one of claims 17-24, wherein applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position comprises: applying the attention mechanism for the mixture of depths attention layer block to (i) the respective first embedding for the preceding output position and (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset.

26. The method of claim 25, wherein the output sequence is being generated in response to an input sequence, and wherein applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position comprises: applying the attention mechanism for the mixture of depths attention layer block to (i) the respective first embedding for the preceding output position, (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset, and (iii) respective first embeddings in a first proper subset of input positions in the input sequence.

27. The method of claim 25 or claim 26. wherein applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position comprises: applying the attention mechanism for the mixture of depths attention layer block to (i) the respective first embedding for the preceding output position and (ii) respective first embeddings for output positions in the output sequence that were previously selected for inclusion in the first proper subset but not to (iv) respective first embeddings for output positions in the output sequence that were previously not selected for inclusion in the first proper subset.

28. The method of any one of claims 17-27, wherein: each mixture of depths attention layer block comprises (i) an attention layer that applies the attention mechanism for the mixture of depths attention block and (ii) a feedforward subnetwork, and processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate the respective second embedding for the preceding output position comprises: processing the respective first embedding for the preceding output position using the attention layer to generate a respective attended embedding for the preceding output position; and processing the respective attended embedding for the preceding output position using the feedforward subnetwork to generate a respective third embedding for the preceding output position.

29. The method of claim 28, when dependent on claim 20, wherein processing the respective first embedding for the preceding output position using the mixture of depths attention layer block further comprises: applying a residual connection to the respective third for the preceding output position to generate the respective initial second embedding for the preceding output position.

30. The method of claim 28 or claim 29, wherein generating the respective second embedding for the preceding output position without applying the attention mechanism for the mixture of depths attention layer block to the preceding output position comprises: generating the respective second embedding for the preceding output position without processing any embedding for the preceding output position using the attention layer or the feedforward subnetwork.31 . The method of any one of claims 17-30, wherein in response to determining not to include the preceding output position in the first proper subset of the output positions, generating the respective second embedding for the preceding output position without applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position comprises: setting the second embedding for the preceding output position equal to the first embedding for the preceding output position.

32. The method of any one of claims 17-31, wherein the routing model for the mixture of depths attention layer is a linear neural network layer or a multi-layer perceptron (MLP).

33. The method of any one of claims 17-32, wherein the plurality of layer blocks further comprises one or more attention layer blocks that are not mixture of depths attention layer blocks.

34. The method of any one of claims 17-33, wherein the plurality of layer blocks comprises a plurality of mixture of depths attention layer blocks.

35. The method of any one of claims 17-34, when dependent on claim 28, wherein, for one or more of the mixture of depths attention layer blocks, the feedforward subnetwork is a mixture of experts (MoE) subnetwork that includes a routing function and a plurality of expert feedforward neural networks.

36. 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, comprising: selecting, using a first routing model and from respective first embeddings for each of the input positions, a first proper subset of the input positions to be processed by an attention layer block; processing the respective first embeddings for the first proper subset of the input positions using the attention layer block to generate respective second embeddings for each of the input positions in the first proper subset, comprising applying the attention mechanism of the attention layer block for each of the input positions in the first proper subset of the input positions; and generating a respective second embedding for each of the input positions not in the first proper subset without applying the attention mechanism for the attention layer block to the input positions not in the first proper subset.

37. A method performed by one or more computers and for auto-regressively generating an output sequence comprising a respective output token at each of plurality’ of output positions, the method comprising, for each of the plurality of output positions: receiving a preceding output token generated for the preceding output position; and processing the preceding output token using a neural network to generate the output token at the output position, comprising: determining, using an inference routing model and from a respective first embedding for the preceding output position, whether to include the preceding output position in a first proper subset of the output positions; in response to determining to include the preceding output position in the first proper subset of the output positions, processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate a respective second embedding for the preceding output position, comprising applying an attention mechanism of the mixture of depths attention layer block to the respective first embedding for the preceding output position; and in response to determining not to include the preceding output position in the first proper subset of the output positions, generating the respective second embedding for the preceding output position without applying the attention mechanism for the mixture of depths attention layer block to the respective first embedding for the preceding output position.

38. 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, comprising: selecting, using a first routing model and from respective first embeddings for each of the input positions, a first proper subset of the input positions to be included as keys for an attention mechanism of an attention layer block; and processing the respective first embeddings using the attention layer block to generate respective second embeddings for each of the input positions, comprising applying the attention mechanism of the attention layer block to the respective first embedding using one or more sets of keys, wherein each set of keys includes a respective key for each of the input positions in the first proper subset of the input positions and does not include a respective key for any of the input positions not in the first proper subset.

39. A method performed by one or more computers and for auto-regressively generating an output sequence comprising a respective output token at each of plurality of output positions, the method comprising, for each of the plurality of output positions: receiving a preceding output token generated for the preceding output position; and processing the preceding output token using a neural network to generate the output token at the output position, comprising: determining, using an inference routing model and from a respective first embedding for the preceding output position, whether to include the preceding output position in a first proper subset of the output positions; and processing the respective first embedding for the preceding output position using the mixture of depths attention layer block to generate a respective second embedding for the preceding output position, comprising applying an attention mechanism of the mixture of depths attention layer block to the respective first embedding for the preceding output position using one or more sets of keys, wherein each set of keys includes a respective key for each output position that has been selected for inclusion in the first proper subset and does not include a respective key for any of the output positions that have not been selected for inclusion in the first proper subset.

40. The method of any one of the preceding claims, wherein the neural network is configured to perform a machine learning task comprising an image processing task, wherein the neural network is configured to process an input image to generate: (i) a network output comprising scores for each of a set of object categories, each score representing an estimated likelihood that the input image comprises an image of an object belonging to the category; (ii) a network output comprising identifying locations in the input image at which particular types of objects are depicted; or (hi) a network output assigning each pixel of the input image to a category from a set of categories.

41. 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-40.

42. 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-40.

Images