Efficient decoding of output sequences using early exit adaptors

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for performing a machine learning task. One of the methods includes generating an output sequence by, at each of a plurality of output time steps: generating a current input sequence from at least the tokens at output time steps that precede the output time step in the output sequence; generating a respective embedding for each input in the current input sequence; and processing the respective embeddings for the inputs in the current input sequence through the layers in the sequence of layers until a termination criterion is satisfied.

Description

[0001]Attorney Docket No.: 45288-0544WO1 EFFICIENT DECODING OF OUTPUT SEQUENCES USING EARLY EXIT ADAPTORS CROSS-REFERENCE TO RELATED APPLICATION This application claims priority to U.S. Provisional Application No.63 / 702,152, filed on October 1, 2024. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application. BACKGROUND This specification relates to processing inputs using neural networks to generate output sequences. 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., another hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of parameters. SUMMARY This specification describes a system implemented as computer programs on one or more computers in one or more locations that performs a machine learning task on a network input using a decoder neural network. According to one aspect there is provided a method performed by one or more computers and for auto-regressively generating, using a decoder neural network, an output sequence that includes a respective token from a vocabulary of tokens at each of a plurality of output time steps, wherein the decoder neural network comprises a sequence of layers, an output subnetwork and, for each layer of a subset of the layers in the sequence, a respective early-exit adaptor, and wherein the method comprises: generating the output sequence by, at each of the plurality of output time steps: generating a current input sequence from at least the tokens at output time steps that precede the output time step in the output sequence; generating a respective embedding for each input in the current input sequence; processing the respective embeddings for the inputs in the current input sequence through the layers in the sequence of layers until a termination criterion is satisfied, comprising, for each layer starting from the first layer in the sequence and until the termination criterion is satisfied: obtaining data derived from a respective input embedding for the layer for each input in the Attorney Docket No.: 45288-0544WO1 current input sequence, and generating a respective updated embedding for at least a last input in the current input sequence at least in part by applying an attention mechanism using the data derived from the respective input embeddings; when the layer is in the subset, determining whether the termination criterion is satisfied at the layer; when the termination criterion is satisfied at one of the layers in the subset: processing the updated embedding for the last input in the current input sequence generated by the layer at which the termination criterion is satisfied using the respective early-exit adaptor for the layer to generate an adapted embedding; processing the adapted embedding using the output subnetwork to generate a probability distribution over the tokens in the vocabulary; and selecting the token at the output time step using the probability distribution. In some implementations, the output sequence is conditioned on a network input. In some implementations, the current input sequence further comprises one or more inputs corresponding to the network input. In some implementations, processing the respective embeddings for the inputs in the input sequence through the layers in the sequence of layers until a termination criterion is satisfied comprises: determining that the termination criterion is satisfied at the last layer in the sequence when the termination criterion was not satisfied at any preceding layer. In some implementations, the method further comprises when the termination criterion is satisfied at the last layer in the sequence, processing the updated embedding for the last input in the current input sequence generated by the last layer using the output subnetwork to generate a last probability distribution over the tokens in the vocabulary; and selecting the token at the output time step using the last probability distribution. In some implementations, processing the respective embeddings for the inputs through the layers comprises: once the termination criterion is satisfied at a given layer in the sequence, refraining from performing any processing for any additional layers that are after the given layer in the sequence at the output time step. In some implementations, processing the respective embeddings for the inputs through the layers comprises: once the termination criterion is satisfied at a given layer in the sequence, refraining from performing any processing using any additional respective early- exit adaptors for any additional layers that are after the given layer in the sequence at the output time step. In some implementations, each respective early-exit adaptor comprises a multilayer perceptron, and wherein the early-exit adaptor is configured to process the updated embedding by: applying a non-linear transformation to the updated embedding to generate a Attorney Docket No.: 45288-0544WO1 transformed embedding; and applying a projection to the transformed embedding to generate the adapted embedding. In some implementations, each respective early-exit adaptor comprises a linear projection layer, and wherein the early-exit adaptor is configured to process the updated embedding by applying a linear transformation to the updated embedding to generate the adapted embedding. In some implementations, each respective early-exit adaptor comprises a set of low- rank adaptation (LoRA) weights, and wherein processing the updated embedding for the last input in the current input sequence generated by the layer at which the termination criterion is satisfied comprises processing the updated embedding in accordance with the low-rank adaptation weights. In some implementations, each respective early-exit adaptor comprises a LoRA projection layer, and wherein the early-exit adaptor is configured to process the updated embedding by applying a transformation to the updated embedding to generate the adapted embedding. In some implementations, determining whether the termination criterion is satisfied at the layer comprises: generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input generated by the layer; and determining that the termination criterion is satisfied when the confidence score for the layer is greater than or equal to a threshold value. In some implementations, generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: generating a probability distribution over the tokens in the vocabulary from the updated respective input embedding for the last input in the current input sequence; and determining the confidence score based on a highest probability in the probability distribution. In some implementations, generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: generating a probability distribution over the tokens in the vocabulary from the updated respective input embedding for the last input in the current input sequence; and determining the confidence score based on a difference between a highest probability in the probability distribution and a second highest probability in the probability distribution. Attorney Docket No.: 45288-0544WO1 In some implementations, generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: determining the confidence score based on a similarity measure between (i) the updated respective input embedding for the last input in the current input sequence generated by the layer and (ii) the updated respective input embedding for the last input in the current input sequence generated by the preceding layer in the sequence. In some implementations, generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: processing the updated respective input embedding for the last input in the current input sequence using an early exit classifier to generate the confidence score. In some implementations, the respective early-exit adaptor for each layer has been trained independently from each other respective early-exit adaptor. In some implementations, the decoder neural network has been trained on a layer- wise loss function that measures, for each layer, a difference between an output probability distribution and a ground-truth probability distribution. In some implementations, the output probability distribution comprises a first probability distribution over the tokens in the vocabulary and a second probability distribution over the tokens in the vocabulary, and wherein the first probability distribution is weighted by a probability that the termination criterion is satisfied for the layer, and wherein the second probability distribution is weighted by a probability that the termination criterion is not satisfied for the layer. In some implementations, the probability that the termination criterion is satisfied for the layer comprises a highest probability of the first probability distribution. In some implementations, the first probability distribution over the tokens in the vocabulary is generated by: processing the updated embedding for the last input in the current input sequence generated by the layer using the respective early-exit adaptor for the layer to generate the adapted embedding for the layer; and applying a softmax function to the adapted embedding for the layer. In some implementations, the second probability distribution over the tokens in the vocabulary is generated by: processing the updated embedding for the last input in the current input sequence generated by the layer through a subsequent layer and the respective early- exit adaptor for the subsequent layer to generate the adapted embedding for the subsequent layer; and applying a softmax function to the adapted embedding for the subsequent layer. Attorney Docket No.: 45288-0544WO1 In some implementations, the respective early-exit adaptor for each layer has been trained while holding parameters of the layers fixed. In some implementations, the output sequence is conditioned on a network input for a machine learning task, and wherein the decoder neural network is configured to generate the output sequence in accordance with a set of low-rank approximation weights for the machine learning task. In some implementations, the respective early-exit adaptor for each layer has been trained while holding parameters of the layers and the set of low-rank approximation weights fixed. In some implementations, (a) the network input comprises at least one of one or more of image data, audio data, video data; and the output sequence comprises data describing content of the image data, audio data or video data; (b) the network input comprises data defining a context, and the output sequence comprises image data, audio data or video data having content defined by the context; (c) the network input comprises sensor data output by a sensor and describing a real-world environment, and the output sequence comprises control data for controlling an agent in the real-world environment described by the sensor data; or (d) the network input comprises electronic health record data for a patient and the output sequence comprises a treatment suggestion, diagnosis or prognosis for the patient. According to another aspect there is provided a computer-implemented method comprising: obtaining data specifying a pre-trained neural network, wherein the pre-trained neural network comprises a sequence of layers; and training one or more early-exit adaptors for the pre-trained neural network, wherein each early-exit adaptor is associated with a respective layer from a subset of the sequence of layers and is configured to generate an adapted embedding from an updated embedding generated by the associated layer. In some implementations, training the one or more early-exit adaptors comprises training each of the one or more early-exit adaptors on a plurality of training examples that each comprise (a) a training input that includes one or more tokens, and (b) a target output that includes a ground-truth sequence of output tokens following the training input. In some implementations, parameters of the sequence of layers are held fixed during the training of the one or more early-exit adaptors. In some implementations, the pre-trained neural network comprises a set of low-rank approximation weights, and wherein the parameters of the sequence of layers and the set of low-rank approximation weights are held fixed during the training of the one or more early- exit adaptors. Attorney Docket No.: 45288-0544WO1 In some implementations, each early-exit adaptor is trained independently from each other early-exit adaptor. In some implementations, training the one or more early-exit adaptors comprises training the one or more early-exit adaptors on a layer-wise loss function that measures, for each layer, a difference between an output probability distribution and a ground-truth probability distribution. In some implementations, the output probability distribution is based on a weighted combination of a first probability distribution given a determination to exit at the respective layer and a second probability distribution given a determination not to exit at the respective layer. In some implementations, the first probability distribution is weighted by a probability that a termination criterion is satisfied for the layer, and wherein the second probability distribution is weighted by a probability that the termination criterion is not satisfied for the layer. According to another aspect there is provided 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 or more computers to perform the operations of the methods described herein. According to another aspect there is provided one or more computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform the operations of the methods described herein. According to another aspect there is provided 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 or more computers to implement a decoder neural network comprising: a sequence of layers, each layer configured to process an input embedding to generate an updated embedding; and a respective early-exit adaptor for each layer of a subset of the layers in the sequence, wherein each respective early-exit adaptor is configured to process the updated embedding from the corresponding layer to generate an adapted embedding that is aligned with a final updated embedding of a last layer of the sequence of layers. In some implementations, each respective early-exit adaptor comprises a multilayer perceptron (MLP). Attorney Docket No.: 45288-0544WO1 In some implementations, each respective early-exit adaptor comprises one of: a linear projection layer; a set of low-rank adaptation (LoRA) weights; or a LoRA projection layer. In some implementations, the decoder neural network is configured to process the updated embedding from the corresponding layer using the respective early-exit adaptor for the corresponding layer when a termination criterion is satisfied. The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages. Recent advances in machine learning models comprising decoder neural networks, such as Transformer-based large language models (LLMs) have led to significant performance improvements across many tasks. In particular, for complex tasks such as a plan generation task, generating a chain of thought alongside the final answer can enhance reasoning. However, this improvement in performance increases inference time and computational cost, creating challenges for practical applications. This specification describes a system that integrates early exit models and early-exit adaptors. For example, the system includes a decoder neural network that includes a sequence of layers and one or more early-exit adaptors, allowing the decoder neural network to dynamically decide whether to exit at intermediate layers based on the complexity of the chain of thought generation process. The system can dynamically adjust the number of layers used to process each token in the chain of thought, so that some tokens are generated by shallower or earlier intermediate layers, while others are generated by deeper or later intermediate layers, or the last layer. For example, the system can include an exit router (or “routing mechanism”) for each layer that determines, e.g., for each token, whether to exit at the layer for an optimal trade-off between quality and cost. Thus the system can reduce compute at inference while maintaining high performance, balancing output quality and computational cost. For example, during generation of an output sequence, layers of the decoder neural network that are not expected to improve the quality of the output sequence appreciably can be bypassed such that processing operations that would otherwise have been carried out for these layers are not performed. To enable high-quality token generation with fewer layers, the system includes early exit adaptors. The early exit-adaptor for each layer is configured to align intermediate representations generated by the layer with the final output of the last layer. Because the intermediate representations generated by intermediate layers are aligned with the final output of the last layer, the system can address potential performance degradation due to Attorney Docket No.: 45288-0544WO1 exiting early (i.e., not performing the processing operations of the later layers), and thus enable high-quality exits at earlier layers (i.e., representations or embeddings that can be decoded to determine high-quality tokens), which reduces the computational resources that would otherwise be required for processing using all of the layers of the neural network. Furthermore, the system can minimize or otherwise reduce the computational overhead for the early-exit adaptors. For example, the early-exit adaptors can have any of a variety of lightweight architectures. Some conventional early exit models are trained with a language modeling loss that only optimizes the final layer’s predictions, resulting in intermediate layers that are not well- suited for decoding tokens directly. The system described in this specification includes early- exit adaptors that enable the neural network to produce high-quality tokens from earlier layers. For example, the early-exit adaptors and a routing mechanism that determines whether to early-exit at each layer can be jointly optimized. For example, the neural network can be trained on a loss function that encourages token alignment only when the routing mechanism predicts a token should be early-exited, allowing the neural network to focus its computational power where necessary. As used herein, a routing mechanism (or exit router) can be a component (e.g., comprising a neural network) associated with a layer of the sequence of layers of the decoder neural network that determines whether a termination criterion is satisfied at the layer. In some implementations, the system can train the early exit-adaptors independently from one another using a layer-wise language modeling loss. Thus each early-exit adaptor can be trained to specialize in producing accurate token predictions without interference from the training of higher-layer adaptors. Some conventional early exit models are fine-tuned on certain machine learning tasks for better performance, but do not perform well on reasoning tasks. In some examples, the system described in this specification uses softmax routing to dynamically adjust exit points based on token-level confidence, which allows the model to balance computational depth with the complexity of the generated chain of thought. For example, the system can determine whether a termination criterion is satisfied for a layer based on a confidence score for the layer. For example, the confidence score can be indicative of a lack of uncertainty in the output of the layer, i.e., a lack of uncertainty when selecting a token for the output sequence using the output of the layer. The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, Attorney Docket No.: 45288-0544WO1 and advantages of the subject matter will become apparent from the description, the drawings, and the claims. BRIEF DESCRIPTION OF THE DRAWINGS FIG.1 shows an example task performance system. FIGS.2A-2B are a flow diagram of an example process for generating an output sequence. FIG.3 shows an example process for training a decoder neural network. FIG.4 shows an example decoder neural network. FIG.5 shows the performance of an example task performance system. Like reference numbers and designations in the various drawings indicate like elements. DETAILED DESCRIPTION FIG.1 is a diagram of an example task performance system 100. The task performance system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented. The task performance system 100 is a system that performs a machine learning task using a decoder neural network 110. In particular, the system 100 generates an output sequence 130 that includes a respective token from a vocabulary of tokens at each of multiple output time steps. In some examples, the vocabulary of tokens can include any of a variety of tokens that represent text symbols or other symbols. For example, the vocabulary of tokens can include one or more of characters, sub-words, words, punctuation marks, numbers, or other symbols that appear in a corpus of natural language text and / or computer code. Additionally, or alternatively, the vocabulary of tokens can include tokens that can represent data other than text, such as images, videos, or audio. For example, the vocabulary of tokens can include image tokens that represent a discrete set of image patch embeddings of an image that can be generated by an image encoder neural network based on processing the image patches of the image. As another example, the vocabulary of tokens can include audio tokens that represent code vectors in a codebook of a quantizer, e.g., a residual vector quantizer. Attorney Docket No.: 45288-0544WO1 The machine learning task can be any machine learning task that generates a network output that includes an output sequence. For example, the machine learning task can require processing a network input to generate an output sequence. As an example, the machine learning task can be a chain-of-thought (COT) reasoning task. For example, the output sequence can outline the reasoning process for responding to a network input. As a particular example, the network input can include a plan generation task, e.g., planning to carry out a multi-step task, a trip planning task, a robotic planning task, etc. To perform the machine learning task, the system uses a neural network 110 to generate the output sequence 130. The system 100 can generate the output sequence 130 autoregressively. That is, the system 100 can generate each particular token in the output sequence conditioned on a current input sequence that includes any (e.g., one or more or all) tokens that precede the particular token in the output sequence (in some implementations, the system can also be conditioned on the network input, e.g., tokens from an input sequence of the network input, when generating each token). The neural network 110 can have any appropriate architecture for generating an output sequence. For example, the neural network 110 can include any appropriate types of neural network layers (e.g., self-attention layers, fully connected layers, convolutional layer, and so forth) in any appropriate numbers (e.g., 5 layers, 10 layers, or 50 layers) and connected in any appropriate configuration (e.g., as a directed graph of layers). As an example, the neural network 110 can include an auto-regressive decoder neural network. As other examples, the neural network 110 can include a recurrent neural network (RNN)-based model, a state space model (SSM), or a hybrid model that includes a RNN- based model and a self-attention based model, e.g., as described in De, Soham et al., “Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models”, arXiv preprint arXiv:2402.19427 (2024). Some example suitable architectures for the decoder neural network 110 are described in further detail in Gemini Team, et al., “Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805”, 2023; Gemma Team, et al., “Gemma: Open Models Based on Gemini Research and Technology”. arXiv preprint arXiv:2403.08295; or Chowdhery, et al., “PaLM: Scaling Language Modeling with Pathways”, arXiv preprint arXiv:2204.02311. The decoder neural network 110 is configured to receive a respective embedding for each input (e.g., each token) in an input sequence. The input sequence includes a respective Attorney Docket No.: 45288-0544WO1 input at each of one or more input positions. An embedding can refer to an ordered collection of numerical values, e.g., a vector, matrix, or other tensor of numerical values. The decoder neural network 110 includes a sequence of multiple layers 120a-n. The decoder neural network 110 is configured to process the respective embeddings for the inputs in the input sequence through the sequence of layers to generate a respective updated embedding for each input in the input sequence. Each layer operates on a respective input sequence that includes a respective input vector at each of one or more positions. The input vectors will also be referred to as “embeddings” or “hidden states” in this specification. Each layer is configured to obtain data derived from a respective input embedding for each input (e.g., each token) in the input sequence and update the respective input embedding for at least a last input (e.g., each token) in the input sequence at least in part by applying an attention mechanism using the data derived from the respective input embeddings. In some implementations, each of the layers can include a self-attention neural network layer followed by a feedforward neural network. A self-attention neural network layer is a neural network layer that applies an attention mechanism, e.g., to update the input embedding for at least the last input (e.g., token) based on the input embeddings of the other inputs (e.g., tokens) in the input sequence. Some examples of self-attention layers are described in Vaswani et al. “Attention is all you need”, 31st Conference on Neural Information Processing Systems (NIPS 2017), Long Beach, CA, USA. In some examples, each of the self-attention layers can include one or more normalization layers, e.g., an RMSNorm layer. While generating the output sequence 130, to generate a token at an output time step, the neural network 110 is further configured to process an updated embedding for the last input (e.g., token) in the input sequence, e.g., using an output subnetwork 126, to generate a probability distribution over the tokens in the vocabulary. The system 100 selects the token at the output time step using the probability distribution. The output subnetwork 126 is configured to process an embedding (such as the updated embedding for the last token in the input sequence) to generate a probability distribution over the tokens in the vocabulary. The output subnetwork 126 can have any appropriate architecture for generating a probability distribution from an embedding. In implementations where the output subnetwork 126 is implemented as a neural network, the output subnetwork 126 can include any appropriate types of neural network layers (e.g., fully connected layers, message passing layers, convolutional layers, attention layers, recurrent layers, pooling layers, and so forth), in Attorney Docket No.: 45288-0544WO1 any appropriate number (e.g., 5 layers, or 10 layers, or 50 layers), and connected in any appropriate configuration (e.g., as a directed graph of layers). As an example, the output subnetwork 126 can include a linear layer and a softmax layer or can be a multi-layer perceptron (MLP). In some examples, while generating the output sequence 130, the system can perform “early exiting.” Early exiting refers to generating an output token using the output of an intermediate layer of the decoder neural network rather than the output of a last layer of the decoder neural network. That is, when early exiting occurs at a particular intermediate layer of the sequence of layers, the output embedding (e.g., an adapted embedding as described below) generated by the intermediate layer may be processed by the output subnetwork without having been processed by the layers that are after the layer in the sequence of layers. Performing early exiting can make the generation of the token more computationally efficient because the computation of any layers after the intermediate layer at which the system early exits does not need to be performed. To perform early exiting, the system processes the respective embeddings for the inputs in the input sequence through the sequence of layers until a termination criterion 106 is satisfied. The system determines whether the termination criterion is satisfied, e.g., using a routing mechanism, at one or more layers of a subset 122 of layers in the sequence 120a-n. The subset 122 can include one or more layers of the layers 120a-n that are before the last layer of the sequence. As an example, the subset 122 can include every layer of the layers 120a-n except for the last layer. As another example, the subset 122 can include intermediate layers that are at a depth that is before a particular layer in the sequence, intermediate layers starting at a particular layer in the sequence, intermediate layers at a predetermined interval, e.g., every other intermediate layer, etc. The term “depth” is used here to count layers; the first layer has a depth of one, and each successive layer has a “depth” one higher than the one before. When the termination criterion is satisfied, the system 100 can use an early-exit adaptor 124 for the layer to process the updated embedding for the last input in the current input sequence generated by the layer to generate an adapted embedding that can be used for generating an output token. In cases where the system 100 determines that the termination criterion is satisfied at an intermediate layer, the system selects a token at the output time step at the intermediate layer based on the adapted embedding generated for the intermediate layer at which the termination criterion is satisfied, rather than selecting a token at the last layer of the decoder neural network 110. Attorney Docket No.: 45288-0544WO1 Thus when the termination criterion is satisfied at an intermediate layer at a particular output time step, the system refrains from performing any processing (i.e., performs no processing) for any additional layers in the sequence of layers that are after the intermediate layer, and refrains from performing any processing (i.e., performs no processing) using any additional respective early-exit adaptors for any additional layers that are after the intermediate layer at the output time step. Thus the system adds capacity or predictive power to the neural network through the additional parameters of the early-exit adaptors, without significantly increasing the computational burden or latency at inference. In general, the termination criterion will be satisfied at different intermediate layers for different output time steps. The decoder neural network 110 includes, for each layer of the subset 122 of layers in the sequence 120a-n, a respective early-exit adaptor 124. Each early-exit adaptor 124 can have any appropriate architecture for generating an adapted embedding from an updated embedding. In implementations where early-exit adaptor 124 is implemented as a neural network, the early-exit adaptor 124 can include any appropriate types of neural network layers (e.g., fully connected layers, message passing layers, convolutional layers, attention layers, recurrent layers, pooling layers, and so forth), in any appropriate number (e.g., 5 layers, or 10 layers, or 50 layers), and connected in any appropriate configuration (e.g., as a directed graph of layers). As particular examples, the early-exit adaptor 124 can include a multilayer perceptron (MLP), e.g., that processes the updated embedding to generate the adapted embedding, a linear projection layer, e.g., that processes the updated embedding to generate the adapted embedding, a set of low-rank adaptation, also referred to as low-rank approximation (LoRA) weights, e.g., that is used to generate an output that is added to the updated embedding to generate the adapted embedding, or a LoRA projection layer, e.g., that is applied to the updated embedding to generate the adapted embedding. Examples of LoRA weights / projection layers are described in in Hu, E. et al., LoRA: Low-rank adaptation of large language models, arXiv:2106.09685 (2021) and Houlsby, N. et al., Parameter-efficient transfer learning for NLP, arXiv:1902.00751 (2019). For example, the low-rank adaptation weights can comprise respective numerical values of a first weight matrix and a second weight matrix. The system may then process an input using the first weight matrix to generate an output that is then processed using the second weight matrix to generate the adapted embedding, e.g., the output of the second weight matrix can be combined with (added to) the updated embedding to generate the adapted embedding. The Attorney Docket No.: 45288-0544WO1 first and second weight matrices may be used together in this way to approximate updates to the embedding that could otherwise have been obtained by fine-tuning the layer of the decoder neural network, but advantageously requiring fewer weight values to be trained compared to fine-tuning of the layer. Generating the adapted embedding using the early-exit adaptor 124 is described in further detail below with reference to FIG.2. The system determines to use the output of the intermediate layer for a given output token only when the termination criterion 106 is satisfied. If the termination criterion is not satisfied for any of the intermediate layers when generating the given output token (i.e., the output token for a given output time step), the system uses the output of the last layer to generate the given output token, i.e., the output of each intermediate layer is provided as input to a next intermediate layer or, in the case of the last intermediate layer, to the last layer, i.e., processing is performed by each of the layers in the sequence in turn. Thus, in cases where the system performs early exiting, the system can generate tokens of the output sequence 130 using fewer computational resources than would otherwise be required to process input sequences through all layers of the neural network 110. An example process for generating the output sequence 130 is described in further detail below with reference to FIGS.2A-2B. In some implementations, the decoder neural network 110 can have been pre-trained to perform the machine learning task. In some examples, the pre-trained neural network 310 can have been pre-trained to perform a different machine learning task or a more general machine learning task. For example, the decoder neural network can have been trained using a machine learning training technique, e.g., a gradient descent with backpropagation training technique that uses a suitable optimizer, e.g., stochastic gradient descent, RMSprop, Adam optimizer, or Adafactor optimizer, to optimize an objective function for a next token prediction task. For example, the system can train the neural network to optimize an objective function that, for each training example, measures an error (e.g., cross-entropy error) between (i) the ground-truth sequence of output tokens specified by the training example and (ii) the output tokens generated by the neural network for the training input specified by the training example. In some implementations, the decoder neural network 110 can use task-specific low- rank adaptation (LoRA) weights for a specific machine learning task. For example, one or more of the layers of the decoder neural network 110 can each include task-specific LoRA weights. Attorney Docket No.: 45288-0544WO1 For example, the output sequence can be conditioned on a network input for a specific machine learning task. The decoder neural network can be configured to generate the output sequence in accordance with the set of low-rank approximation weights for the specific machine learning task. For example, the set of low-rank approximation weights for the machine learning task can be representative of updates to the decoder neural network that would have been applied during fine-tuning of the decoder neural network for the specific machine learning task e.g., by optimizing a product of two smaller matrices in order to reduce the dimensionality of the calculation required to compute the change in weights required by the update. For example, performing a low-rank approximation refers to breaking up a matrix containing target update weights, e.g., the network weights identified for updating, into a product of two smaller matrices that when multiplied together can recover the values of the network weights with high fidelity. In particular, the low-rank decomposition can represent^^ + ^^ ≈ ^^ + ^^, where ^^ is the update to the base network weights ^^ and theproduct BA approximates ^^. “Low-rank” can refer to matrices that have a lower matrix rank than the update that would have been applied during fine-tuning. ,By using task-specific LoRA weights, the system can achieve better performance on the specific task while reducing the number of parameters of the neural network that are trained compared to fine-tuning for the specific machine learning task. The set of low-rank approximation weights for each machine learning task can have been determined by training the decoder neural network on a training dataset for the machine learning task. For example, the system can train the set of low-rank approximation weights while holding parameters of layers of the decoder neural network fixed. For example, the system can perform gradient updates in accordance with an objective function, e.g., based on minimizing a loss function applicable to the decoder neural network, in order to generate an approximation of the target update weights using the low-rank approximation. Training the decoder neural network is described below in further detail with reference to FIG.3. FIGS.2A-2B are a flow diagram of an example process 200 for generating an output sequence. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a task performance system, e.g., the task performance system 100 of FIG.1, appropriately programmed, can perform the process 200. The system auto-regressively generates an output sequence that includes a respective token from a vocabulary of tokens at each of multiple output time steps using a decoder Attorney Docket No.: 45288-0544WO1 neural network. As described above with reference to FIG.1, the decoder neural network includes a sequence of layers, an output subnetwork and, for each layer of a subset of the layers in the sequence, a respective early-exit adaptor. To generate the output sequence, the system performs the steps 202-206 at each of multiple output time steps. The system generates a current input sequence (step 202). For example, the system can generate the current input sequence from at least the tokens at output time steps that precede the output time step in the output sequence. For example, at output time step t, the system generates the current input sequence to include one or more tokens generated for output time steps 1, ..., t-1. In some examples, the current input sequence can include any (e.g., one or more or all) tokens generated at output time steps that precede the output time step in the output sequence. In some other examples, the system can include one or more tokens generated at output time steps that precede the output time step in the output sequence. As a particular example, when the system uses KV-caching, the current input sequence can include the token generated at the immediately preceding output time step. In some examples, when the system uses KV-caching, the system can include, in the output sequence, the token generated at the immediately preceding output time step, and key and value vectors from a cache for each self-attention layer for inputs at output time steps that precede the immediately preceding output time step. Implementations of KV-caching are described in more detail below. In some examples, the output sequence is conditioned on a network input. In these examples, the system receives the network input. The network input can include, for example, data for performing a machine learning task. The current input sequence can include one or more inputs corresponding to the network input. In some examples, the system can obtain the one or more inputs corresponding to the network input. As an example, one or more of the inputs can include a token from a vocabulary. For example, the system can generate the one or more inputs by processing the network input using an appropriate tokenizer. For example, the system can tokenize a given sequence of words by applying a tokenizer, e.g., the SentencePiece tokenizer (Kudo et al., arXiv:1808.06226) or another tokenizer, to divide the sequence into tokens from the vocabulary. Attorney Docket No.: 45288-0544WO1 As another example, one or more of the inputs can include a soft token. For example, the system can obtain the one or more inputs by generating an embedded representation of at least a portion of the network input. The system generates a respective embedding for each input in the current input sequence (step 204). For example, the system can map each of one or more tokens to corresponding numerical values in accordance with a predefined mapping to generate each respective embedding. As an example, the predefined mapping can be represented as an embedding matrix that has learned values. In other examples, the predefined mapping can be represented as a different learned module such as a neural network. In some examples where one or more of the inputs includes a soft token, the system can use the input as the respective embedding. The system processes the respective embeddings for the inputs in the current input sequence through the layers in the sequence of layers until a termination criterion is satisfied (step 206). An example process for processing the respective embeddings through the layers is described in further detail with reference to FIG.2B. Referring to FIG.2B, the system performs the steps 212-222 for each layer starting from the first layer in the sequence and until the termination criterion is satisfied. The system obtains data derived from a respective input embedding for the layer for each input in the current input sequence (step 212). At the first layer, the respective input embedding for each input in the current input sequence includes the respective embedding for each input in the current input sequence generated in step 204. For each subsequent layer, the respective input embedding for each input in the current input sequence includes the respective output embedding for the input generated using the preceding layer. The data derived from a respective input embedding for the respective layer for each input in the current input sequence can include query (Q), key (K), and value (V) vectors for the input. In some examples where the current input sequence includes each input of the input sequence, the system can generate query, key, and value matrices. For example, to generate the query matrix, the system can apply a query transformation, e.g., parameterized by a learned weight query matrix, to the respective input embeddings for the current input sequence. To generate the key matrix, the system can apply a key transformation, e.g., Attorney Docket No.: 45288-0544WO1 parameterized by a learned key weight matrix, to the respective input embeddings for the current input sequence. To generate the value matrix, the system can apply a value transformation, e.g., parameterized by a learned value weight matrix, to the respective input embeddings for the current input sequence. In examples where the current input sequence includes one input, e.g., where the system uses KV-caching, the system can generate query, key, and value vectors for the input. For example, to generate the query vector, the system can apply a query transformation, e.g., parameterized by a learned query weight matrix, to the respective input embedding for the input. To generate the key vector, the system can apply a key transformation, e.g., parameterized by a learned key weight matrix, to the respective input embedding for the input. To generate the value vector, the system can apply a value transformation, e.g., parameterized by a learned value weight matrix, to the respective input embedding for the input. The system can obtain a key matrix by combining, e.g., appending, the key vector for the input to a retrieved key matrix. The retrieved key matrix can include a key vector for each input at output time steps that precede the output time step for the input. The retrieved key matrix can have been retrieved from a KV-cache for the layer. The system can obtain a value matrix by combining, e.g., appending, the value vector for the input to a retrieved value matrix. The retrieved value matrix can include a value vector for each input at output time steps that precede the output time step for the input. The retrieved value matrix can have been retrieved from a KV-cache for the layer. The system can retrieve cached key and value matrices for inputs at output time steps that precede the output time step for the input from a KV-cache for the layer. The system generates a respective updated embedding for at least a last input in the current input sequence (step 214). For example, the system can generate the respective updated embedding at least in part by applying an attention mechanism using the data derived from the respective input embeddings. For example, the system can generate a respective output embedding for each input in the current input sequence. As an example, when the current input sequence includes each input of the input sequence, the system can generate a respective output embedding for each input. As another example, when the current input sequence includes one input, e.g., where the system uses KV-caching, the system can generate a respective output embedding for the input. Attorney Docket No.: 45288-0544WO1 The system can generate each respective output embedding by applying the attention mechanism over the query, key, and value vectors for at least the input to generate a respective attention output for the input. The system can process the respective attention output, e.g., using a feedforward network, to generate the respective output embedding. The attention mechanism is configured to determine relationships within the input sequence. For example, when the current input sequence includes each input of the input sequence, the attention mechanism may, for example, be a dot product attention operation applied by applying each query vector to each key vector to determine respective weights for each value vector, then combining the value vectors using the respective weights to determine the corresponding attention output for each input of the input sequence. As another example, when the current input sequence includes one input, e.g., where the system uses KV-caching, the attention mechanism may, for example, be a dot product attention operation applied by applying the query vector to each key vector in the key matrix to determine respective weights for each value vector in the value matrix, then combining the value vectors using the respective weights to determine the attention output for the input. In some implementations, the system can scale the attention output by a scaling factor e.g. by the square root of the dimensions of the queries and keys, to implement scaled dot product attention. Thus, for example, an output of the attention operation may be determined as where dk is a dimension of the key (and value) vector. the attention mechanism can implement multi-head attention, that is, it may apply multiple different attention mechanisms in parallel. The attention outputs of these can then be combined, e.g. concatenated, to generate the attention output. In some implementations, the attention mechanism can apply a learned linear transformation defined by a learned output projection matrix to the combination of attention outputs to reduce to the original dimensionality. The system can generate the output embedding for the input by processing the attention output of the attention mechanism, e.g., using a feed-forward network. The feed-forward network is configured to process an attention output to generate an output embedding in accordance with a learned feed-forward weight matrix for the feed- forward network. The feed-forward network can have any appropriate architecture for processing an attention output to generate an output embedding. As a particular example, the feed-forward network can include one or more linear layers and an activation function, e.g., a Attorney Docket No.: 45288-0544WO1 Rectified Linear Unit (ReLU) or Gaussian Error Linear Unit Gated Linear Unit (GeGLU) activation function. When the layer is in the subset, the system determines whether the termination criterion is satisfied at the layer (step 216). In some examples, the system can determine whether the termination criterion is satisfied at a layer by generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence generated by (or provided to) the layer. The system can determine that the termination criterion is satisfied when the confidence score for the layer is greater than or equal to a threshold value. By determining whether the termination criterion is satisfied based on the confidence score for the layer, the system can dynamically determine the number of layers used to process each token, balancing computational depth with the complexity of the generated output sequence, and allowing for more efficient use of computational resources. As an example, the confidence score can be based on a softmax distribution generated by the layer. For example, the system can generate the confidence score by generating a probability distribution over the tokens in the vocabulary from the updated respective input embedding for the last input in the current input sequence. For example, the probability distribution can be a softmax distribution generated by projecting the updated respective embedding using a linear layer to generate a vector that includes a score for each token in the vocabulary, and processing the vector of scores using a softmax layer to generate the softmax distribution. In some examples, the system can determine the confidence score based on a highest probability in the probability distribution. As another example, the system can generate the confidence score by generating a probability distribution over the tokens in the vocabulary from the updated respective input embedding for the last input in the current input sequence as described above. The system can determine the confidence score based on a difference between a highest probability in the probability distribution and a second highest probability in the probability distribution. As another example, the system can generate the confidence score based on a similarity measure e.g., cosine similarity, between (i) the updated respective input embedding for the last input in the current input sequence generated by the layer and (ii) the updated respective input embedding for the last input in the current input sequence generated by the preceding layer in the sequence. Attorney Docket No.: 45288-0544WO1 As another example, the system can generate the confidence score by processing the updated respective input embedding for the last input in the current input sequence using an early exit classifier to generate the confidence score. The early exit classifier is configured to predict whether the updated respective input embedding for the last input in the current input sequence is reliable, e.g., whether the updated input embedding is expected to lead to the same token being selected for the output time step as when the updated input embedding is processed further by the one or more layers after the layer in the sequence of layers. The early exit classifier can have any appropriate architecture for predicting whether the updated respective input embedding for the last input in the current input sequence is reliable. In implementations where the early exit classifier is implemented as a neural network, the early exit classifier can include any appropriate types of neural network layers (e.g., fully connected layers, message passing layers, convolutional layers, attention layers, recurrent layers, pooling layers, and so forth), in any appropriate number (e.g., 5 layers, or 10 layers, or 50 layers), and connected in any appropriate configuration (e.g., as a directed graph of layers). In some examples, the early exit classifier is trained on top of the hidden states using a cross-entropy loss against an oracle that compares each layer’s prediction to the final model output. For example, the classifier machine learning model can have been trained using a per-layer independent cross-entropy loss against a target label that indicates whether the token with the highest probability in the probability distribution at the layer is the same as the token with the highest probability in the probability distribution of the last layer. The target label can have been generated by comparing each layer’s token with the highest probability to the token with the highest probability at the last layer. When the termination criterion is satisfied at one of the layers in the subset, the system performs steps 218-222. The system processes the updated embedding for the last input in the current input sequence generated by the layer at which the termination criterion is satisfied using the respective early-exit adaptor for the layer to generate an adapted embedding (step 218). The adapted embedding is aligned to the output space of the output subnetwork. That is, the early-exit adapter can be trained to generate adapted embeddings that can be processed using the output subnetwork to generate a probability distribution over the tokens in the vocabulary from which an appropriate (e.g., high quality) token can be selected. As an example, a respective early-exit adaptor can include a multilayer perceptron. In these examples, the early-exit adaptor is configured to process the updated embedding by applying a non-linear transformation to the updated embedding to generate a transformed Attorney Docket No.: 45288-0544WO1 embedding, and applying a projection to the transformed embedding to generate the adapted embedding. For example, the projection can project the transformed embedding to the output space of the output subnetwork. Because the multilayer perceptron is similar to other layers of the decoder neural network, e.g., a feed-forward network, the system can integrate the multilayer perceptron into the decoder neural network. By leveraging the same structure as existing layers, the system can perform the computations of the adaptors using the same optimized pathways, improving computational efficiency. As another example, a respective early-exit adaptor can include a linear projection layer. In these examples, the early-exit adaptor is configured to process the updated embedding by applying a linear transformation to the updated embedding to generate the adapted embedding. Because the linear projection layer requires a relatively small amount of computational resources for storing and inference, the linear projection layer can be advantageous in scenarios where computational resources are constrained. As another example, a respective early-exit adaptor can include a set of low-rank adaptation (LoRA) weights. In these examples, processing the updated embedding for the last input in the current input sequence generated by the layer at which the termination criterion is satisfied can include processing the updated embedding in accordance with the low-rank adaptation weights. For example, the system can add the updated embedding to the output of a LoRA projection defined by the LoRA weights that is applied to the last input in the current input sequence. For example, the system can apply LoRA weights to the key, query, value, and output matrices of the attention mechanism to generate the output of the LoRA projection, to reduce the number of trainable parameters. The system can thus adapt the attention weights in intermediate layers, allowing for effective token prediction while minimizing computational cost. As another example, a respective early-exit adaptor can include a LoRA projection layer. In these examples, the early-exit adaptor is configured to process the updated embedding by applying a transformation to the updated embedding to generate the adapted embedding. The system can apply the transformation using a LoRA projection layer that includes a frozen base matrix, e.g., that can be randomly initialized, and low-rank trainable matrices, e.g., that are updated during training of the early-exit adaptors. The system can thus reduce the rank of the hidden weight matrices, lowering the number of parameters, and increasing the computational efficiency. Attorney Docket No.: 45288-0544WO1 The system processes the adapted embedding to generate a probability distribution over the tokens in the vocabulary (step 220). For example, the system can process the adapted embedding using the output subnetwork to generate the probability distribution. For example, the system can process the adapted embedding using a linear layer to generate a vector that includes a score for each token in the vocabulary. The system can, for example, process the vector of scores using a softmax layer to generate the probability distribution. The probability distribution assigns a respective probability to each token in the vocabulary. The system selects the token at the output time step using the probability distribution (step 222). For example, the neural network can greedily select the highest-scoring token or can sample, e.g., using top-k sampling, nucleus sampling or another sampling technique, a token from the probability distribution. Alternatively, in some implementations, instead of processing the adapted embedding using the output subnetwork to generate a probability distribution from which to select the token at the output time step, in response to determining that the termination criterion is satisfied, the system can generate the token at the output time step as the adapted embedding for the last input in the current input sequence. The system can use the adapted embedding as a soft token, e.g., for purposes of image, audio, or video generation. The system can decode the token at the output time step using an appropriate decoder neural network, e.g., for image generation, audio generation, or video generation. Once the termination criterion is satisfied at a given layer in the sequence, the system refrains from performing (i.e., does not perform) any processing for any additional layers that are after the given layer in the sequence at the output time step. The system can proceed to step 202 for a next output time step, without performing any processing for any additional layers. Once the termination criterion is satisfied at a given layer in the sequence, the system refrains from performing (i.e., does not perform) any processing using any additional respective early-exit adaptors for any additional layers that are after the given layer in the sequence at the output time step. The system can proceed to step 202 for a next output time step, without performing any processing for any additional layers or using any additional respective early-exit adaptors. That is, for any given input, the system uses at most one of the respective early-exit adaptors. In examples where the termination criterion is satisfied at an intermediate layer in the sequence, the system uses only the early-exit adaptor for the intermediate layer at which the termination criterion was satisfied. Attorney Docket No.: 45288-0544WO1 In examples where the system uses KV-caching, once the termination criterion is satisfied at a given layer in the sequence, the system updates the KV-cache for each subsequent layer to include the key vector and the value vector for the last input in the current input sequence. For example, the system can copy the updated embedding for the last input in the current input sequence generated by the layer at which the termination criterion is satisfied. The system can apply a key transformation to the updated embedding to generate the key vector. The system can apply a value transformation to the updated embedding to generate the value vector. Thus, in cases where the system early exits for a token at an output time step at a particular layer, the system can access the key and value vector for the token at subsequent time steps at subsequent layers, allowing for approximation of missing hidden states and ensuring the smooth operation of the attention mechanism during the generation of subsequent tokens. When the layer is in the subset and the system determines that the termination criterion is not satisfied, the system proceeds to perform steps 212-222 for subsequent layers in the sequence of layers. In some cases, the system determines that the termination criterion is satisfied at the last layer in the sequence when the termination criterion was not satisfied at any preceding layer. That is, the system determines to exit at the last layer in the sequence when the termination criterion was not satisfied at any preceding layer. When the termination criterion is satisfied at the last layer in the sequence, the system processes the updated embedding for the last input in the current input sequence generated by the last layer using the output subnetwork to generate a last probability distribution over the tokens in the vocabulary. For example, the system can process the updated embedding using a linear layer to generate a vector that includes a score for each token in the vocabulary. The system can process the vector of scores using a softmax layer to generate the probability distribution. The probability distribution assigns a respective probability to each token in the vocabulary. The system selects the token at the output time step using the last probability distribution. For example, the neural network can greedily select the highest-scoring token or can sample, e.g., using top-k sampling, nucleus sampling or another sampling technique, a token from the probability distribution. Alternatively, in some implementations, in response to determining that the termination criterion is satisfied, the system generates the token at the output time step as the updated embedding for the last input in the current input sequence. The system can decode Attorney Docket No.: 45288-0544WO1 the token at the output time step using an appropriate decoder neural network, e.g., for image generation or audio generation. In examples where the output time step is the last output time step, the system can output the output sequence. In some examples, the system can process the output sequence using an appropriate detokenizer for the machine learning task. FIG.3 shows an example process 300 for training a decoder neural network. For convenience, the process 300 will be described as being performed by a system of one or more computers located in one or more locations. For example, a task performance system, e.g., the task performance system 100 of FIG.1, appropriately programmed in accordance with this specification, can perform the process 300. The system can perform the process 300 to train the decoder neural network, e.g., the early-exit adaptors. For example, the system can train the decoder neural network by training the respective early-exit adaptor for each layer while holding parameters of the layers of the decoder neural network fixed. In examples where the decoder neural network includes a set of LoRA weights, the system can train the respective early-exit adaptor for each layer while holding parameters of the layers and the set of low-rank approximation weights fixed. The system can train the decoder neural network over a sequence of training iterations. At each training iteration, the system can receive one or more training examples for the training iteration (step 302). For example, each training example can include (a) a training input that includes one or more tokens, and (b) a target output that includes a ground-truth sequence of output tokens following the training input. As a particular example, the training input can include an instruction to perform a task and the target output can include a response for the instruction. The system can generate output probability distributions for the training iteration (step 304). For example, for each training example for the training iteration, the system can process the training input using the decoder neural network to generate an output probability distribution for each token of a training sequence of tokens. The system can update parameters of the decoder neural network (step 306). For example, the system can determine gradients of a loss function that depends on the output probability distributions and can use the gradients to update parameter values, e.g., of the early exit adaptors, of the decoder neural network. As an example, the loss function can be a classification objective function, e.g., a log likelihood loss function, that measures a difference between the output probability distributions generated by the decoder neural network and ground truth probability distributions for the target outputs. Attorney Docket No.: 45288-0544WO1 As a particular example, the loss function can be a layer-wise loss function that measures, for each layer, e.g., in the subset, a difference between an output probability distribution and a ground-truth probability distribution. For example, in cases where the subset includes every layer except the last layer, the loss function can be described as , where is the negative log-likelihood loss for layer i. The weights Thus the respective early-exit adaptor for each layer i can be trained independently from each other respective early-exit adaptor. Thus, early-exit adaptors for shallower layers may in that case have no impact on early-exit adaptors at higher-layers, and each early-exit adaptor can be trained to specialize in producing accurate token predictions without interference from the training of higher-layer adaptors. In some implementations, the system trains the decoder neural network to generate output probability distributions conditioned on early exit probabilities. For example, the system can train the early exit adaptors in conjunction with the decision of the routing mechanism about whether a token should exit early. Thus, the training of each early-exit adaptor is influenced by the real-time, probabilistic decision to exit of the routing mechanism. Rather than training every early-exit adaptor to generate accurate outputs for every token, the system can train each early-exit adaptor to specialize only on the tokens that the routing mechanism is confident should exit at the corresponding layer, improving specialization of each early-exit adaptor, and allowing early-exit adaptors to optimize accuracy and computational efficiency more effectively. For example, the system can train the decoder neural network to generate output probability distributions that include a first probability distribution over the tokens in the vocabulary and a second probability distribution over the tokens in the vocabulary. For example, the output probability distribution can be a combination, e.g., a weighted sum, of the first probability distribution given a determination to exit, and the second probability distribution given a determination of not to exit. The first probability distribution can be weighted by a probability that the termination criterion is satisfied for the layer, and the second probability distribution can be weighted by a probability that the termination is not satisfied for the layer. The probability that the Attorney Docket No.: 45288-0544WO1 termination criterion is satisfied for the layer can include a highest probability of the first probability distribution. Thus the system can use the probability that the termination criterion is satisfied at each layer, instead of a binary decision of whether to exit, to determine the value of the loss function for propagating back the gradients from the loss function, allowing the system to focus the training of the early exit adaptor for the layer on tokens that the routing mechanism predicts should exit at the layer. As an example, for the next token at layer i, the output probability distribution can be computed as , where softmax confidence score of the most probable token prediction at layer i, computed as , where Y is the set of possible tokens for the next step that the token does not exit at layer i, which can be . is the first probability distribution, or the given the token exits at layer i, and is the second probability or the approximated probability over the next token given the token does not exit at layer i. As described above with reference to FIG.2, the first probability distribution and the second probability distribution can be generated using a softmax layer. For example, the system can generate the first probability distribution by processing the updated embedding for the last input in the current input sequence generated by the layer using the respective early-exit adaptor for the layer to generate the adapted embedding for the layer, and applying a softmax function to the adapted embedding for the layer. For example, the system can generate the first probability distribution as for the layer i, where is the updated embedding for the last input and is the early exit adaptor for the The system can generate the second probability distribution by processing the updated embedding for the last input in the current input sequence generated by the layer through a subsequent layer and the respective early-exit adaptor for the subsequent layer to generate the Attorney Docket No.: 45288-0544WO1 adapted embedding for the subsequent layer, and applying a softmax function to the adapted embedding for the subsequent layer. For example, the system can approximate the second probability distribution as , where if thetoken does not exit at is passed to thenext layer, where the prediction for continues at layer , is the early exit adaptor for the layer . Thus the train the decoder neural network to account for the decision of the routing mechanism on whether to exit or continue processing at the next layer so that an early exit adaptor for a layer focuses on aligning predictions only when the routing mechanism indicates an early exit for a token at the layer. For example, by jointly optimizing the routing mechanism and early-exit adaptors, the system allows layers to specialized in predicting subsets of tokens, and the routing mechanism determines when to leverage each layer’s strength. The system can thus enable a more efficient allocation of computational resources, particularly for lightweight early exit adaptors, which can rival the performance of more complex models while being significantly cheaper to serve. The system can determine whether training is complete (step 308). If the system determines that training is not complete, the system can continue to a next training iteration (e.g., return to step 302). The system can determine whether training is complete using any of a variety of criteria. For example, the system can determine that training is complete after a pre-determined number of training iterations. As another example, the system can determine that training is complete when a value of the loss function for the training iteration falls below a predetermined threshold. As another example, the system can determine that training is complete when a difference between the value of the loss function for the current training iteration and a value of the loss function for a previous training iteration falls below a predetermined threshold. When the system determines that training is complete, the system can return the trained decoder neural network (step 310). FIG.4 shows an example decoder neural network 400. The decoder neural network 400 is an example of the decoder neural network 110 described above with reference to FIG. 1. Attorney Docket No.: 45288-0544WO1 The decoder neural network 400 includes an embedding layer 410. The embedding layer 410 is configured to generate a respective embedding for each input in a current input sequence as described above with reference to step 204 of FIG.2. The decoder neural network 400 includes multiple layers (“Transformer Layer 1..i..N”). The decoder neural network 400 is configured to process the respective embeddings for each input in the current input sequence through the multiple layers until a termination criterion is satisfied as described above with reference to FIG.2. Given an input prompt , the decoder neural network 400 is configured to autoregressively generate output using hidden states computed across the layers. For generating the next token , the hidden state at layer i is computed as an output subnetwork 430 (“Decoder Projection”). The output subnetwork 430 is configured to map hidden states, i.e., the updated embeddings generated by the Transformer Layer N, or the adapted embeddings generated by early-exit adaptors at layers 1...N-1, to the vocabulary. The output subnetwork 430 generates the logits for the token distribution where is the shared output projection matrix. . The system can thus ensure that decoder neural network memory overhead from having a separate output projection matrix at each layer. Early exiting allows the model to dynamically decide whether to terminate the generation at an intermediate layer, based on a confidence score, rather than always computing the final output at the last layer N. For a given token , the system checks the confidence at each layer i as where computes the confidence from the If where is a predefined threshold, the system determines to exit early. The system performs early exiting using a respective early-exit adaptor for the layer i. FIG.4 shows that an early exit adaptor (“EE Adaptor 1...i...N”) is added to each intermediate layer 1...i...N. The input to the early exit adaptor Aiis the hidden state hiof layer i. The hidden state includes the intermediate token representations computed by the model up to that layer. The early exit adaptor processes hiand produces an adapted hidden state , which is aligned to the output space. Attorney Docket No.: 45288-0544WO1 To generate the token for the intermediate layer i, the output of the early exit adaptor is passed through the shared output projection matrix Wi to generate the logits for token prediction. For example, the token can be selected using the probability distribution . 1...i...N, and the output subnetwork 430, are frozen. network 400 includes a set of task-specific LoRA weights that was trained prior to the training of the early exit adaptors, the set of LoRA weights is also frozen. FIG.5 shows the performance of an example task performance system. The example task performance systems can include a decoder neural network 110 as described above with reference to FIGS.1-4. FIG.5 shows the performance of example task performance systems in terms of accuracy and average depth per token on a COT reasoning task, compared to conventional techniques. The example task performance systems include, e.g., a decoder neural network with early exit adaptors that include MLPs (“Backbone w / MLP adaptor”), a decoder neural network with early exit adaptors that include a linear projection layer (“Backbone w / LP adaptor”), a decoder neural network with early exit adaptors that include LoRA weights (“Backbone w / LoRA adaptor” and “Backbone w / Single LoRA adaptor”). The conventional techniques include a neural network without early exit training and early exit adaptors (“Backbone w / o EE adaptation”) and a neural network conventionally trained for early exit (“Baseline EE adaptation”). FIG.5 shows that the example task performance systems described in this specification have higher accuracy when exiting across the various layers e.g., at layers before the final layer 28, than the conventional techniques. In some examples, the system can generate the decoder neural network from a pre- trained neural network. The pre-trained neural network can include a sequence of layers. For example, the system can obtain data specifying a pre-trained neural network. The system can train one or more early-exit adaptors for the pre-trained neural network. Each early-exit adaptor is associated with a respective layer from a subset of the sequence of layers and is configured to generate an adapted embedding from an updated embedding generated by the associated layer. Further details for training the one or more early-exit adaptors are described above with reference to FIG.3. Attorney Docket No.: 45288-0544WO1 For example, the system can train each of the one or more early-exit adaptors on multiple training examples that each comprise (a) a training input that includes one or more tokens, and (b) a target output that includes a ground-truth sequence of output tokens following the training input. Parameters of the sequence of layers are held fixed during the training of the one or more early-exit adaptors. In some examples, the pre-trained neural network includes a set of low-rank approximation weights. In these examples, the parameters of the sequence of layers and the set of low-rank approximation weights are held fixed during the training of the one or more early-exit adaptors. In some examples, each early-exit adaptor is trained independently from each other early-exit adaptor. In some examples, the system can train the one or more early-exit adaptors by training the new neural network on a layer-wise loss function that measures, for each layer, a difference between an output probability distribution and a ground-truth probability distribution. The output probability distribution can be based on a weighted combination of a first probability distribution given a determination to exit at the respective layer and a second probability distribution given a determination not to exit at the respective layer. The first probability distribution is weighted by a probability that the termination criterion is satisfied for the layer, and the second probability distribution is weighted by a probability that the termination criterion is not satisfied for the layer. In some examples, the system can output data representing the trained decoder neural network. For example, the system can output data representing the trained decoder neural network for storage. As an example, the system can implement the decoder neural network that includes the sequence of layers. Each layer is configured to process an input embedding to generate an updated embedding. The decoder neural network can include the respective early-exit adaptor for each layer of a subset of the layers in the sequence. Each respective early-exit adaptor can be configured to process the updated embedding from the corresponding layer to generate an adapted embedding that is aligned with a final updated embedding of a last layer of the sequence of layers. The decoder neural network can be configured to process the updated embedding from the corresponding layer using the respective early-exit adaptor for the corresponding layer when a termination criterion is satisfied. In examples where the output sequence includes one or more soft tokens at each of multiple output time steps, the system can generate the output sequence by, at each of the multiple output time steps, obtaining a current input sequence that includes one or more soft tokens from at least the tokens at output time steps that precede the output time step in the Attorney Docket No.: 45288-0544WO1 output sequence. The system can process the respective embeddings for each input in the current input sequence through the layers in the sequence of layers until a termination criterion is satisfied. The respective embeddings for each input in the current input sequence can include the one or more soft tokens. Some examples of machine learning tasks that the system can be configured to perform follow. As one example, if the input to the neural network is a sequence of text in one language, the output generated by the neural network may a sequence of text in another language that is proper translation of the input text into the other language. As another example, the task can be a text generation task, where the input is a sequence of text, and the output is another sequence of text, e.g., a completion of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text that is about a topic specified by the first sequence of text. As another example, the input to the text generation task can be, or comprise, an input other than text, e.g., an image (a term used here to mean pixel-level data defining at least one intensity value for at least one multi-dimensional array of pixels; the image may for example be the output of a still or video camera arranged to capture the (still or moving) image by imaging part of the real world), and the output sequence can be text that describes the input. As another example, the input can include both text and other data, e.g., images. As another example, the task can be a text generation task, where the input is a sequence of text, and the output is another sequence of text, e.g., a completion of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text that is about a topic specified by the first sequence of text. As another example, the task can be a computer code generation task, where the input is a sequence of text describing the functionality of a piece of computer code, or a sequence of computer code to be modified or completed, or both and the output is a sequence of computed code that modifies the computer code, that has the functionality that is described by the sequence of text, or both. As another example, the task can be an image generation task, where the input is a conditioning input and the output is a sequence of intensity value inputs for the pixels of an image. As another example, the task can be a health prediction task, where the input is electronic health record data for a patient (e.g. comprising one or more numerical measurements which have been made of the patient and / or one or more medical images of the Attorney Docket No.: 45288-0544WO1 patient) and the output is a prediction that is relevant to the future health of the patient, e.g., a predicted treatment that should be prescribed to the patient, the likelihood that an adverse health event will occur to the patient, or a predicted diagnosis for the patient. As another example, the task can be an agent control task, where the input is a sequence of observations or other data characterizing states of an environment (e.g., a real- world 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 (that is, an electromechanical apparatus), a control system for an industrial facility, or a control system that controls a different kind of agent. As another example, the task can be a robotic planning task. For example, the input can include data, e.g., text, image, etc., specifying a current state and a goal for an agent (e.g., electromechanical robot). The output can include data representing a sequence of steps for accomplishing the goal. As another example, the task can be a trip or event planning task. For example, the input can include data, e.g., text, image, etc., specifying a goal. The output can include data representing a plan, e.g., an itinerary, for the goal. The sequence of actions for generating the output can include determining objectives for the goal, e.g., using a language model neural network or VLM, determining sub-events, determining deadlines, determining a timeline, and generating data representing the plan. As another example, the task can be a resource allocation task, e.g., for allocating computational resources. For example, the input can include data specifying a goal for the allocation of resources and information about the computational resources, specifications, and components. The output can include data representing a resource management plan. As another example, the task can include a mathematical reasoning task, e.g., a mathematical problem covering arithmetic operations and requiring step by step reasoning to arrive at the correct solution. As another example, the task may be an audio processing task, e.g., of processing an input which comprises an audio signal (such as sound samples captured from the real world by a microphone), e.g. to generate an outcome descriptive of content of the audio signal. For example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network may be a text transcript for the utterance. Attorney Docket No.: 45288-0544WO1 As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can indicate whether a particular word or phrase (“hotword”) was spoken in the utterance. As another example, the task can be a text to speech task, where the input is text in a natural language or features of text in a natural language and the network output is a spectrogram, a waveform, or other data defining audio of the text being spoken in the natural language. As a particular example of this, the system can perform speaker-conditioned text- to-speech, where the input is a text transcript and audio of a speaker and the output is a verbalization of the text transcript being spoken by the speaker. As another example, the task may be an audio generation task for generating an output audio signal (e.g. a set of sound samples (e.g. amplitude or Fourier components) at each time point of a plurality of time points within a time period). For example, the output audio signal can be a speech signal and the system can generate the speech signal unconditionally, e.g., resulting in a speech signal being generated that is drawn from a distribution represented by the training data set(s) on which the neural network was trained. As another example, the output audio signal can be a different type of audio signal, e.g., music, ambient sounds, sound events, and so on, and the system can generate the audio signal unconditionally, e.g., resulting in an audio signal being generated that is drawn from a distribution represented by the training data set(s) on which the neural network was trained. As another example, the system can receive a context to generate an audio signal and generate the output audio signal conditioned on the received context. For example, the output audio signal can be a speech or other audio signal, i.e., so that the system generates the output audio signal conditioned on a context that is an input audio signal. For example, the context can include an audio input as an input audio signal. In some examples, the audio input can include words spoken by a particular speaker, and the output audio signal can be a continuation of the words spoken by the particular speaker. In some examples, the audio input can include music, and the output audio signal can be a continuation of the music in the audio input. In some implementations, the output audio signal that is generated can be a prediction of an audio signal that follows the input audio signal. For example, the context can be an input speech signal that is a question that is asked by one speaker and the output audio signal can be an output speech signal that is an answer to the question spoken by the same speaker or another speaker. As another example, the context can be an input speech signal that is a Attorney Docket No.: 45288-0544WO1 first portion of an utterance spoken by one speaker and the input audio signal can be an output speech signal that is a completion of the utterance spoken by the speaker or by another speaker or a response to the input utterance. In some implementations, the input audio signal can include a melody, and the output audio signal can be music that follows the melody. The system can generate an output audio signal that is music that follows the melody. As another example, the system can perform sound separation on the input audio signal to generate the output audio signal. For example, the input audio signal can include both speech and music or other background noise and the output audio signal can represent only the speech. As another example, the input audio signal can include speech from multiple speakers (and optionally background noise) and the output audio signal can include only the speech of one of the speakers. In some examples, the system can perform audio- conditional separation, i.e., when the input audio signal can include an extra audio input acoustically similar to one of the sources. As another example, the system can perform speech-to-speech translation, where the input speech and the output speech represent the same semantic content but spoken differently. For example, the input audio signal can include speech in one natural language and the output audio signal can represent speech in a target, different natural language that is a translation of the input speech into the target language. As another example, the input audio signal can include speech spoken by a first speaker and the output audio signal can represent the same semantic content as the input speech but spoken by a different speaker. As another example, the input audio signal can include a first speech segment and the output audio signal can include a second, shorter speech segment that summarizes the semantic content of the first speech. As another example, the context can include both audio data and text data. In this example, the output audio signal can include speech that reflects the text data. In some implementations, the output audio signal can include music that reflects the text data. In these implementations, the system generates an output audio signal that reflects the text data. For example, the system can perform transcript conditioned speech enhancement, where the context is a text transcript and noisy audio corresponding to the text transcript and the output audio signal is clean audio corresponding to the text transcript. As another example, the system can perform transcript-based audio in-filling, where the context is a text transcript and audio corresponding to a portion of the text transcript, and the output audio signal corresponds to a different portion of the text transcript. Attorney Docket No.: 45288-0544WO1 As another example, the context can include both audio data and visual, e.g., image or video, data. In this example, the output audio signal can include speech that describes the visual data or music that reflects the visual data. For example, the system can perform audio-video continuation, where the system receives a partial audio track with corresponding video, and the output audio signal is a continuation of the partial audio track. As another example, the system can perform cross-modal in-filling, where the system receives video and an audio track corresponding to a portion of the video and the output audio signal is an audio track corresponding to a different portion of the video. As another example, the context input can include only visual data. For example, the system can perform image-conditioned audio generation, where the system receives an input image and generates an output audio signal that describes the image. As another example, the context input can include only text data. For example, the system can perform sound synthesis based on text descriptions, e.g., where the input is text that describes an audio signal and the output is an output audio signal that is characterized by the text. Additional description of tasks that the neural network (the “generative” neural network) can perform now follow. Generally, the generative neural network is configured to process a conditioning input (“input prompt”) to generate a data item. Generally the data item represents a response to the conditioning input which may be, e.g. a “prompt” for the generative neural network. For example the conditioning input can characterize one or more desired properties for the generated data item. In some implementations the method is adapted to run on parallel processing computer system comprising a plurality of hardware computing devices configured to operate in parallel. Each hardware computing device may comprise a neural network accelerator, i.e. specialized hardware that is used to accelerate neural network computations, such as a GPU (Graphics Processing Unit) or TPU (Tensor Processing Unit). In general a neural network accelerator is configured to perform hardware matrix multiplications; it can include a set of one or more multiply accumulate units (MACs). For example, the system can deploy multiple instances of the neural network across multiple devices and use the instances to parallelize the generation of the outputs. In some implementations the generative neural network generates an output token sequence from an input token sequence including the conditioning input. The generative neural Attorney Docket No.: 45288-0544WO1 network may then be configured to process the input token sequence to generate for each position in the output token sequence, a respective score for each token in a vocabulary of output tokens, that is used to select an output token for the output token sequence. In some implementations the tokens can represent text, e.g., words, wordpieces or characters, in a natural or computer language. For example text may be received, e.g., as a series of encoded characters, e.g. UTF-8 encoded characters; such “characters” can include Chinese and other similar characters, as well as logograms, syllabograms and the like. A text encoder, i.e. a tokenizer, can process a sequence of text to represent the text as a series of text tokens from a vocabulary of text tokens, e.g. that each represent words, wordpieces or characters in a natural or computer language. The computer language may be any formal language used to communicate with a computer, e.g. a markup language, or a command or configuration language, or a data exchange language such as JSON, or a programming language. The tokenizer can, e.g., implement BPE (Byte Pair Encoding) or Wordpiece tokenization. Optionally the text can be obtained from audio data representing speech; the output tokens may be converted into audio data that represent speech corresponding to the text. Also or instead the tokens may represent an image. For example a set (sequence) of input or output tokens can represent an image. Each image token may comprise a block encoding of values of the pixels in a different region of an image that maps a set of values of the pixels to a respective image token. The block encoder may comprise a neural network, e.g. having one or more (self-)attention layers, such as a Transformer neural network. Also or instead the tokens may represent an audio waveform. For example a set (sequence) of input or output tokens can represent audio data representing an waveform e.g. instantaneous audio amplitude values or time-frequency audio data. Each image token may comprise a block encoding of the audio waveform in a different time segment of the audio that maps a set of values representing the audio waveform to a respective image token. The block encoder may comprise a neural network, e.g. having one or more (self-)attention layers, such as a Transformer neural network. In a multimodal system audio data or an image may be flagged by a start-of-audio token or start-of-image token. In some implementations the generative neural network can be a multimodal network that is configured to process a conditioning input comprising one or more of text data, audio data defining an audio signal (e.g. as amplitude values of the audio signal or as a time-frequency representation of the audio signal), or a still or moving image (e.g. as image pixel values), to generate a data item that can similarly comprise text data, audio data, or a still or moving image. Attorney Docket No.: 45288-0544WO1 For example the conditioning input may comprise text and the data item may comprise an image or an audio signal that represents speech an image generated in response to the text, e.g. described by the text. Also or instead the conditioning input may comprise an audio signal that represents speech, or an image, and the data item may comprise text, e.g. that describes the conditioning input. As another example the conditioning input may comprise an observation, e.g. of a real world environment, e.g. from sensor such as a camera or other image sensor; and optionally additional information such as information defining a particular task to be deformed. The output data item may comprise agent control data that defines one or more actions to be performed by an agent, e.g. by a mechanical agent such as a robot or autonomous vehicle, to perform a task. The reward model(s) may, e.g., define a preferred trajectory of motion of the mechanical agent in the (real-world) environment. In some implementations the generative neural network may comprise a language and / or image generation neural network, that may have been trained before being fine-tuned by the above described method. The conditioning input may comprise a prompt, e.g. a natural or computer language prompt for the generative neural network. The generated data item may comprises a natural or computer language and / or image response to the prompt. In general the generative neural network can have any appropriate architecture for processing the conditioning input to generate the data item. As one example, the generative neural network may comprise an auto-regressive generative model (e.g., a Transformer, a recurrent neural network, etc.) that can auto- regressively generate an output sequence as the data item based on the conditioning input. The generative model can, for example, comprise a large language model (LLM) that can auto- regressively generate tokenized representations of text data, a vision-language model (VLM) that can auto-regressively generate tokenized representations of image or video data, e.g. in response to a text conditioning input or that can auto-regressively generate tokenized representations of text, e.g. in response to an image conditioning input, an audio language model that can auto-regressively generate tokenized representations of text data, or a multimodal model that can that can generate tokens representing any of text, image or audio, e.g. in response to a conditioning input comprising any of text, image or audio, and so forth. As another example, the generative neural network may comprise a diffusion model (e.g., a denoising diffusion model, a score-based diffusion model, a latent diffusion model, etc.) that can generate the data item by repeatedly transforming samples from a noise distribution (e.g., a Gaussian distribution) based on the conditioning input over a sequence of iterations. Attorney Docket No.: 45288-0544WO1 For example, the generative neural network may comprise a diffusion model that transforms samples from the noise distribution using a denoising neural network with any appropriate architecture (e.g., a convolutional neural network, a recurrent neural network, etc.). Such a diffusion model may be used to generate, e.g., a still or moving (video) image. As another example, the generative neural network may comprise a neural network that can generate the data item by transforming samples from a noise distribution (e.g., a Gaussian distribution). The generative neural network may comprise, e.g., a generator network of a generative adversarial network, a decoder of a variational auto-encoder, a normalizing flow, and so on. As used herein an image may be any still or moving image, i.e. the image may be part of a video, in 2D or 3D, and may be a monochrome, color or hyperspectral image, i.e. comprising monochrome or color pixels. As defined herein an “image” includes a point cloud e.g. from a LIDAR system, and a “pixel” includes a point of the point cloud. An image may have been captured by a camera or other image sensor from the real world; and objects in the image may comprise physical objects, represented by the image. According to another aspect, there is provided a system that includes one or more computers and one or more storage devices communicatively coupled to the one or more computers and storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the previously described method. According to another aspect, there is provided one or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations of the previously described method. Example hardware implementations In some implementations the generative neural network, e.g. a language model or a visual language model, is stored on a user computing device, i.e. a device local to the user, such as a mobile device e.g. a mobile phone, or a smart speaker. In some implementations the generative neural network is implemented on a remove server in communication with a user computing device over a wired or wireless network communications link between the user computing device and the server. The user computing device may be provided with an input mechanism, such as a text or voice interface, that enables user input from the user in a natural language. The user computing device may be provided with an output mechanism that provides a system output for the user in the or another natural language e.g. as speech or text; or in some other way, e.g. Attorney Docket No.: 45288-0544WO1 by displaying an image. The input and output mechanism may comprise, e.g., a keyboard, microphone, speaker, display, and / or camera. As an example the input mechanism may comprise a system configured to input audio data characterizing a speech waveform of speech representing the input from the user in a natural language, and configured to convert the audio data into tokens representing the speech in the natural language, e.g. representing a transcription of the spoken input. The output mechanism may comprise a system configured to receive tokens representing the output for the user in the or another natural language and a system configured to convert the received tokens into audio data representing a waveform of speech representing the output to the user in the natural language, i.e. representing spoken words. As a further example, the trained system can be deployed in an environment that enables a user to provide a request for the system, e.g. to process a multimodal conditioning input to generate a corresponding data item output. A users can provide the request, e.g., by way of a user interface or through an application programming interface (API). The request can be transmitted from a user device, e.g., over a data communications network such as the internet, to one or more computers implementing the system, e.g., in a data center. The system can generate a data item and then transmit the data item to a user device over a data communications network. Example multimodal applications The generative neural network may comprise a multimodal machine learning system such as a visual language model (VLM). That is implementations of the generative neural network can perform a multimodal task in which the conditioning input and data item, collectively, comprise data of multiple different types. As used herein text can include numbers, punctuation, special symbols, and so forth. In some implementations, after training, a particular task that is to be performed by the generative neural network can be described by part or all of a sequence of text in the conditioning input to the system. For example in a conditioning input that includes an image such a prompt might specify “Generate a caption”, “Generate a description”, “Answer the following question: [about the image or video]”, or “Detect a person”. Where the system is used for an agent control task a prompt may define “Take the knife out of the drawer”, or “Q: What action should the robot take to take the knife out of the drawer?”. Also or instead such a prompt may give one or more examples of a task to be performed. The generative neural network can be trained on multiple natural and / or computer languages and the prompt may then specify a language to use. Attorney Docket No.: 45288-0544WO1 This specification uses the term “configured” in connection with, or in relation to, systems and environments, as well as computer program components. For a system of one or more computers or environment to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For instance, configuring a system might involve installing a software library with specific algorithms, updating firmware with new instructions for handling data, or adding a hardware component for enhanced processing capabilities. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions. Embodiments of the subject matter and the functional operations described in this specification can be implemented in various forms, including 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 storage device such as a machine-readable storage device, a hard drive or solid-state drive (SSD), a storage medium, 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, e.g., carry, information for transmission to suitable receiver apparatus, e.g., a receiving device or system, for execution by a data processing apparatus. Furthermore, implementations may leverage emerging technologies like quantum computing or neuromorphic computing for specific applications, and may be deployed in distributed or cloud-based environments where components reside on different machines or within a cloud infrastructure. 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 or processing unit, a computer, multiple processors or computers, e.g., working together, graphics processing units (GPUs), or tensor Attorney Docket No.: 45288-0544WO1 processing units (TPUs). 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. Embodiments may particularly benefit from utilizing the parallel processing capabilities of GPUs, in a General-Purpose computing on Graphics Processing Units (GPGPU) context, where code specifically designed for GPU execution, often called kernels or shaders, is employed. Similarly, TPUs excel at running optimized tensor operations crucial for many machine learning algorithms. By leveraging these accelerators and their specialized programming models, the system can achieve significant speedups and efficiency gains for tasks involving artificial intelligence and machine learning, particularly in areas such as computer vision, natural language processing, and robotics. 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. The specific implementation of the computer programs may involve a combination of traditional programming languages and specialized languages or libraries designed for GPGPU programming or TPU utilization, depending on the chosen hardware platform and desired performance characteristics. 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. Attorney Docket No.: 45288-0544WO1 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, e.g., located at a single site or distributed across multiple 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. Examples of engine functions within the context of AI and machine learning could include data pre-processing and cleaning, feature engineering and extraction, model training and optimization, inference and prediction generation, and post-processing of results. The specific design and implementation of engines will depend on the overall architecture and the distribution of computational tasks across various hardware components, including CPUs, GPUs, TPUs, and other specialized processors. 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. Additionally, graphics processing units (GPUs) and tensor processing units (TPUs) can be utilized to enable concurrent execution of aspects of these processes and logic flows, significantly accelerating performance. This approach offers significant advantages for computationally intensive tasks often found in AI and machine learning applications, such as matrix multiplications, convolutions, and other operations that exhibit a high degree of parallelism. By leveraging the parallel processing capabilities of GPUs and TPUs, significant speedups and efficiency gains compared to relying solely on CPUs can be achieved. 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, for even greater performance or energy efficiency in specific use cases. 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 (CPU). Additionally, graphics processing units (GPUs), tensor processing units (TPUs), and other machine learning accelerators can be employed to enhance performance, particularly for tasks involving artificial intelligence and machine learning. These accelerators often work in conjunction with CPUs, handling specialized computations while the CPU manages overall system operations and other tasks. Generally, a central processing unit will receive instructions and data from a read only memory or a random access memory or both. The Attorney Docket No.: 45288-0544WO1 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 specific configuration of processing units and memory will depend on factors like the complexity of the AI model, the volume of data being processed, and the desired performance and latency requirements. 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., hard drives, SSDs, flash memory for persistent data storage, 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. Embodiments can be implemented on a wide range of computing platforms, from small embedded devices with limited resources to large- scale data center systems with high-performance computing capabilities. 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, read-only memory (ROM), solid-state drives (SSDs), and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks or hard disk drives (HDDs); magneto optical disks; and optical discs such as CD ROM and DVD-ROM disks and Blu-ray discs. The specific type of computer-readable media used will depend on factors such as the size of the data, access speed requirements, cost considerations, and the desired level of portability or permanence. 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 or an organic light-emitting diode (OLED) display, for displaying information to the user and a key vectorboard and a pointing device, e.g., a mouse or a trackball, keyboard, touchscreens, voice commands, gesture recognition, or other input modalities depending on the specific device and application, 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 Attorney Docket No.: 45288-0544WO1 documents from a device or application 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. The selection of input and output modalities will depend on the specific application and the desired form of user interaction. Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and compute-intensive parts of machine learning training or production, i.e., inference, workloads. Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework, or a Jax framework. These frameworks offer comprehensive tools and libraries that facilitate the development, training, and deployment of machine learning models. 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, a back-end server, or cloud-based infrastructure, or that includes a middleware component, e.g., an application server, a middleware server, or application programming interface (API), to facilitate communication and data exchange, 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. For instance, the described functionality could be implemented solely on a client device (e.g., for on-device machine learning) or deployed as a combination of front-end and back-end components for more complex applications. 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. The specific system architecture and choice of components will depend on factors such as the scale of the application, the need for real-time processing, data security requirements, and the desired user experience. The computing system can include clients and servers. A client and server are generally remote from each other, e.g., geographically separated, and typically interact through a communication network. The specific type of network, such as a local area network (LAN), a wide area network (WAN), or the Internet, will depend on the reach and Attorney Docket No.: 45288-0544WO1 scale of the application. 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, e.g., designed to communicate with each other using appropriate protocols. These protocols may include HTTP, TCP / IP, or other specialized protocols depending on the nature of the data being exchanged and the security requirements of the system. In some embodiments, a server transmits data, e.g., an HTML page, to a user device such as a computer, smartphone, or tablet, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. The client device can then process the received information and display results to the user, and potentially send data or feedback back to the server. Data generated at the user device, e.g., a result of the user interaction, can be received at the server from the device, e.g., for further processing or storage. This allows for dynamic interactions between the user and the system, enabling a wide range of applications and functionalities. While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination. Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products. Attorney Docket No.: 45288-0544WO1 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. This specification also includes the following clauses: Clause 1. A method performed by one or more computers and for auto-regressively generating, using a decoder neural network, an output sequence that includes a respective token from a vocabulary of tokens at each of a plurality of output time steps, wherein the decoder neural network comprises a sequence of layers, an output subnetwork and, for each layer of a subset of the layers in the sequence, a respective early- exit adaptor, and wherein the method comprises: generating the output sequence by, at each of the plurality of output time steps: generating a current input sequence from at least the tokens at output time steps that precede the output time step in the output sequence; generating a respective embedding for each input in the current input sequence; processing the respective embeddings for the inputs in the current input sequence through the layers in the sequence of layers until a termination criterion is satisfied, comprising, for each layer starting from the first layer in the sequence and until the termination criterion is satisfied: obtaining data derived from a respective input embedding for the layer for each input in the current input sequence, and generating a respective updated embedding for at least a last input in the current input sequence at least in part by applying an attention mechanism using the data derived from the respective input embeddings; when the layer is in the subset, determining whether the termination criterion is satisfied at the layer; when the termination criterion is satisfied at one of the layers in the subset: processing the updated embedding for the last input in the current input sequence generated by the layer at which the termination criterion is satisfied using the respective early-exit adaptor for the layer to generate an adapted embedding; processing the adapted embedding using the output subnetwork to Attorney Docket No.: 45288-0544WO1 generate a probability distribution over the tokens in the vocabulary; and selecting the token at the output time step using the probability distribution. Clause 2. The method of clause 1, wherein the output sequence is conditioned on a network input. Clause 3. The method of clause 2, wherein the current input sequence further comprises one or more inputs corresponding to the network input. Clause 4. The method of any preceding clause, wherein processing the respective embeddings for the inputs in the input sequence through the layers in the sequence of layers until a termination criterion is satisfied comprises: determining that the termination criterion is satisfied at the last layer in the sequence when the termination criterion was not satisfied at any preceding layer. Clause 5. The method of clause 4, further comprising when the termination criterion is satisfied at the last layer in the sequence, processing the updated embedding for the last input in the current input sequence generated by the last layer using the output subnetwork to generate a last probability distribution over the tokens in the vocabulary; and selecting the token at the output time step using the last probability distribution. Clause 6. The method of any preceding clause, wherein processing the respective embeddings for the inputs through the layers comprises: once the termination criterion is satisfied at a given layer in the sequence, refraining from performing any processing for any additional layers that are after the given layer in the sequence at the output time step. Clause 7. The method of any preceding clause, wherein processing the respective embeddings for the inputs through the layers comprises: once the termination criterion is satisfied at a given layer in the sequence, refraining from performing any processing using any additional respective early-exit adaptors for any additional layers that are after the given layer in the sequence at the output time step. Clause 8. The method of any preceding clause, wherein each respective early-exit adaptor comprises a multilayer perceptron, and wherein the early-exit adaptor is configured to process the updated embedding by: applying a non-linear transformation to the updated embedding to generate a transformed embedding; and applying a projection to the transformed embedding to generate the adapted embedding. Attorney Docket No.: 45288-0544WO1 Clause 9. The method of any preceding clause, wherein each respective early-exit adaptor comprises a linear projection layer, and wherein the early-exit adaptor is configured to process the updated embedding by applying a linear transformation to the updated embedding to generate the adapted embedding. Clause 10. The method of any of clauses 1-8, wherein each respective early-exit adaptor comprises a set of low-rank adaptation (LoRA) weights, and wherein processing the updated embedding for the last input in the current input sequence generated by the layer at which the termination criterion is satisfied comprises processing the updated embedding in accordance with the low-rank adaptation weights. Clause 11. The method of any of clauses 1-8, wherein each respective early-exit adaptor comprises a LoRA projection layer, and wherein the early-exit adaptor is configured to process the updated embedding by applying a transformation to the updated embedding to generate the adapted embedding. Clause 12. The method of any preceding clause, wherein determining whether the termination criterion is satisfied comprises: generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input generated by the layer; and determining that the termination criterion is satisfied when the confidence score for the layer is greater than or equal to a threshold value. Clause 13. The method of clause 12, wherein generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: generating a probability distribution over the tokens in the vocabulary from the updated respective input embedding for the last input in the current input sequence; and determining the confidence score based on a highest probability in the probability distribution. Clause 14. The method of clause 12, wherein generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: generating a probability distribution over the tokens in the vocabulary from the updated respective input embedding for the last input in the current input sequence; and determining the confidence score based on a difference between a highest probability in the probability distribution and a second highest probability in the probability distribution. Attorney Docket No.: 45288-0544WO1 Clause 15. The method of clause 12, wherein generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: determining the confidence score based on a similarity measure between (i) the updated respective input embedding for the last input in the current input sequence generated by the layer and (ii) the updated respective input embedding for the last input in the current input sequence generated by the preceding layer in the sequence. Clause 16. The method of clause 12, wherein generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: processing the updated respective input embedding for the last input in the current input sequence using an early exit classifier to generate the confidence score. Clause 17. The method of any preceding clause, wherein the respective early-exit adaptor for each layer has been trained independently from each other respective early-exit adaptor. Clause 18. The method of any preceding clause, wherein the decoder neural network has been trained on a layer-wise loss function that measures, for each layer, a difference between an output probability distribution and a ground-truth probability distribution. Clause 19. The method of clause 18, wherein the output probability distribution comprises a first probability distribution over the tokens in the vocabulary and a second probability distribution over the tokens in the vocabulary, and wherein the first probability distribution is weighted by a probability that the termination criterion is satisfied for the layer, and wherein the second probability distribution is weighted by a probability that the termination criterion is not satisfied for the layer. Clause 20. The method of clause 19, wherein the probability that the termination criterion is satisfied for the layer comprises a highest probability of the first probability distribution. Clause 21. The method of any of clauses 19-20, wherein the first probability distribution over the tokens in the vocabulary is generated by: processing the updated embedding for the last input in the current input sequence generated by the layer using the respective early-exit adaptor for the layer to generate the adapted embedding for the layer; and applying a softmax function to the adapted embedding for the layer. Clause 22. The method of any of clauses 19-21, wherein the second probability distribution over the tokens in the vocabulary is generated by: processing the updated embedding for the last input in the current input sequence Attorney Docket No.: 45288-0544WO1 generated by the layer through a subsequent layer and the respective early-exit adaptor for the subsequent layer to generate the adapted embedding for the subsequent layer; and applying a softmax function to the adapted embedding for the subsequent layer. Clause 23. The method of any preceding clause, wherein the respective early-exit adaptor for each layer has been trained while holding parameters of the layers fixed. Clause 24. The method of any preceding clause, wherein the output sequence is conditioned on a network input for a machine learning task, and wherein the decoder neural network is configured to generate the output sequence in accordance with a set of low-rank approximation weights for the machine learning task. Clause 25. The method of clause 24, wherein the respective early-exit adaptor for each layer has been trained while holding parameters of the layers and the set of low-rank approximation weights fixed. Clause 26. The method of any of clauses 2-25 when dependent on clause 2, wherein: (a) the network input comprises at least one of one or more of image data, audio data, video data; and the output sequence comprises data describing content of the image data, audio data or video data; (b) the network input comprises data defining a context, and the output sequence comprises image data, audio data or video data having content defined by the context; (c) the network input comprises sensor data output by a sensor and describing a real- world environment, and the output sequence comprises control data for controlling an agent in the real-world environment described by the sensor data; or (d) the network input comprises electronic health record data for a patient and the output sequence comprises a treatment suggestion, diagnosis or prognosis for the patient. Clause 27. 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 or more computers to perform the respective operations of any one of clauses 1-26. Clause 28. One or more computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform the respective operations of the method of any one of clauses 1-26. Clause 29. A computer-implemented method comprising: obtaining data specifying a pre-trained neural network, wherein the pre-trained neural Attorney Docket No.: 45288-0544WO1 network comprises a sequence of layers; and training one or more early-exit adaptors for the pre-trained neural network, wherein each early-exit adaptor is associated with a respective layer from a subset of the sequence of layers and is configured to generate an adapted embedding from an updated embedding generated by the associated layer. Clause 30. The method of clause 29, wherein training the one or more early-exit adaptors comprises training each of the one or more early-exit adaptors on a plurality of training examples that each comprise (a) a training input that includes one or more tokens, and (b) a target output that includes a ground-truth sequence of output tokens following the training input. Clause 31. The method of any one of clauses 29-30, wherein parameters of the sequence of layers are held fixed during the training of the one or more early-exit adaptors. Clause 32. The method of any one of clauses 29-31, wherein the pre-trained neural network comprises a set of low-rank approximation weights, and wherein the parameters of the sequence of layers and the set of low-rank approximation weights are held fixed during the training of the one or more early-exit adaptors. Clause 33. The method of any one of clauses 29-32, wherein each early-exit adaptor is trained independently from each other early-exit adaptor. Clause 34. The method of any one of clauses 29-33, wherein training the one or more early-exit adaptors comprises training the one or more early-exit adaptors on a layer-wise loss function that measures, for each layer, a difference between an output probability distribution and a ground-truth probability distribution. Clause 35. The method of clause 34, wherein the output probability distribution is based on a weighted combination of a first probability distribution given a determination to exit at the respective layer and a second probability distribution given a determination not to exit at the respective layer. Clause 36. The method of clause 35, wherein the first probability distribution is weighted by a probability that a termination criterion is satisfied for the layer, and wherein the second probability distribution is weighted by a probability that the termination criterion is not satisfied for the layer. Clause 37. A system comprising: one or more computers; and one or more storage devices storing instructions that, when executed by the one or Attorney Docket No.: 45288-0544WO1 more computers, cause the one or more computers to perform the respective operations of any one of clauses 29-36. Clause 38. One or more computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform the respective operations of the method of any one of clauses 29-36. Clause 39. 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 or more computers to implement a decoder neural network comprising: a sequence of layers, each layer configured to process an input embedding to generate an updated embedding; and a respective early-exit adaptor for each layer of a subset of the layers in the sequence, wherein each respective early-exit adaptor is configured to process the updated embedding from the corresponding layer to generate an adapted embedding that is aligned with a final updated embedding of a last layer of the sequence of layers. Clause 40. The system of clause 39, wherein each respective early-exit adaptor comprises a multilayer perceptron (MLP). Clause 41. The system of clause 39, wherein each respective early-exit adaptor comprises one of: a linear projection layer; a set of low-rank adaptation (LoRA) weights; or a LoRA projection layer. Clause 42. The system of any one of clauses 39-40, wherein the decoder neural network is configured to process the updated embedding from the corresponding layer using the respective early-exit adaptor for the corresponding layer when a termination criterion is satisfied. What is claimed is:

Claims

Attorney Docket No.: 45288-0544WO1 CLAIMS 1. A method performed by one or more computers and for auto-regressively generating, using a decoder neural network, an output sequence that includes a respective token from a vocabulary of tokens at each of a plurality of output time steps, wherein the decoder neural network comprises a sequence of layers, an output subnetwork and, for each layer of a subset of the layers in the sequence, a respective early- exit adaptor, and wherein the method comprises: generating the output sequence by, at each of the plurality of output time steps: generating a current input sequence from at least the tokens at output time steps that precede the output time step in the output sequence; generating a respective embedding for each input in the current input sequence; processing the respective embeddings for the inputs in the current input sequence through the layers in the sequence of layers until a termination criterion is satisfied, comprising, for each layer starting from the first layer in the sequence and until the termination criterion is satisfied: obtaining data derived from a respective input embedding for the layer for each input in the current input sequence, and generating a respective updated embedding for at least a last input in the current input sequence at least in part by applying an attention mechanism using the data derived from the respective input embeddings; when the layer is in the subset, determining whether the termination criterion is satisfied at the layer; when the termination criterion is satisfied at one of the layers in the subset: processing the updated embedding for the last input in the current input sequence generated by the layer at which the termination criterion is satisfied using the respective early-exit adaptor for the layer to generate an adapted embedding; processing the adapted embedding using the output subnetwork to generate a probability distribution over the tokens in the vocabulary; and selecting the token at the output time step using the probability distribution.Attorney Docket No.: 45288-0544WO1 2. The method of claim 1, wherein the output sequence is conditioned on a network input.

3. The method of claim 2, wherein the current input sequence further comprises one or more inputs corresponding to the network input.

4. The method of any preceding claim, wherein processing the respective embeddings for the inputs in the input sequence through the layers in the sequence of layers until a termination criterion is satisfied comprises: determining that the termination criterion is satisfied at the last layer in the sequence when the termination criterion was not satisfied at any preceding layer.

5. The method of claim 4, further comprising when the termination criterion is satisfied at the last layer in the sequence, processing the updated embedding for the last input in the current input sequence generated by the last layer using the output subnetwork to generate a last probability distribution over the tokens in the vocabulary; and selecting the token at the output time step using the last probability distribution.

6. The method of any preceding claim, wherein processing the respective embeddings for the inputs through the layers comprises: once the termination criterion is satisfied at a given layer in the sequence, refraining from performing any processing for any additional layers that are after the given layer in the sequence at the output time step.

7. The method of any preceding claim, wherein processing the respective embeddings for the inputs through the layers comprises: once the termination criterion is satisfied at a given layer in the sequence, refraining from performing any processing using any additional respective early-exit adaptors for any additional layers that are after the given layer in the sequence at the output time step.

8. The method of any preceding claim, wherein each respective early-exit adaptor comprises a multilayer perceptron, and wherein the early-exit adaptor is configured to process the updated embedding by: applying a non-linear transformation to the updated embedding to generate a transformed embedding; and applying a projection to the transformed embedding to generate the adapted embedding.Attorney Docket No.: 45288-0544WO1 9. The method of any preceding claim, wherein each respective early-exit adaptor comprises a linear projection layer, and wherein the early-exit adaptor is configured to process the updated embedding by applying a linear transformation to the updated embedding to generate the adapted embedding.

10. The method of any of claims 1-8, wherein each respective early-exit adaptor comprises a set of low-rank adaptation (LoRA) weights, and wherein processing the updated embedding for the last input in the current input sequence generated by the layer at which the termination criterion is satisfied comprises processing the updated embedding in accordance with the low-rank adaptation weights.

11. The method of any of claims 1-8, wherein each respective early-exit adaptor comprises a LoRA projection layer, and wherein the early-exit adaptor is configured to process the updated embedding by applying a transformation to the updated embedding to generate the adapted embedding.

12. The method of any preceding claim, wherein determining whether the termination criterion is satisfied at the layer comprises: generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input generated by the layer; and determining that the termination criterion is satisfied when the confidence score for the layer is greater than or equal to a threshold value.

13. The method of claim 12, wherein generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: generating a probability distribution over the tokens in the vocabulary from the updated respective input embedding for the last input in the current input sequence; and determining the confidence score based on a highest probability in the probability distribution.Attorney Docket No.: 45288-0544WO1 14. The method of claim 12, wherein generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: generating a probability distribution over the tokens in the vocabulary from the updated respective input embedding for the last input in the current input sequence; and determining the confidence score based on a difference between a highest probability in the probability distribution and a second highest probability in the probability distribution.

15. The method of claim 12, wherein generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: determining the confidence score based on a similarity measure between (i) the updated respective input embedding for the last input in the current input sequence generated by the layer and (ii) the updated respective input embedding for the last input in the current input sequence generated by the preceding layer in the sequence.

16. The method of claim 12, wherein generating a confidence score for the layer from at least the updated respective input embedding for the last input in the current input sequence comprises: processing the updated respective input embedding for the last input in the current input sequence using an early exit classifier to generate the confidence score.

17. The method of any preceding claim, wherein the respective early-exit adaptor for each layer has been trained independently from each other respective early-exit adaptor.

18. The method of any preceding claim, wherein the decoder neural network has been trained on a layer-wise loss function that measures, for each layer, a difference between an output probability distribution and a ground-truth probability distribution.

19. The method of claim 18, wherein the output probability distribution comprises a first probability distribution over the tokens in the vocabulary and a second probability distribution over the tokens in the vocabulary, and wherein the first probability distribution is weighted by a probability that the termination criterion is satisfied for the layer, and wherein the second probability distribution is weighted by a probability that the termination criterion is not satisfied for the layer.Attorney Docket No.: 45288-0544WO1 20. The method of claim 19, wherein the probability that the termination criterion is satisfied for the layer comprises a highest probability of the first probability distribution.

21. The method of any of claims 19-20, wherein the first probability distribution over the tokens in the vocabulary is generated by: processing the updated embedding for the last input in the current input sequence generated by the layer using the respective early-exit adaptor for the layer to generate the adapted embedding for the layer; and applying a softmax function to the adapted embedding for the layer.

22. The method of any of claims 19-21, wherein the second probability distribution over the tokens in the vocabulary is generated by: processing the updated embedding for the last input in the current input sequence generated by the layer through a subsequent layer and the respective early-exit adaptor for the subsequent layer to generate the adapted embedding for the subsequent layer; and applying a softmax function to the adapted embedding for the subsequent layer.

23. The method of any preceding claim, wherein the respective early-exit adaptor for each layer has been trained while holding parameters of the layers fixed.

24. The method of any preceding claim, wherein the output sequence is conditioned on a network input for a machine learning task, and wherein the decoder neural network is configured to generate the output sequence in accordance with a set of low-rank approximation weights for the machine learning task.

25. The method of claim 24, wherein the respective early-exit adaptor for each layer has been trained while holding parameters of the layers and the set of low-rank approximation weights fixed.

26. The method of any of claims 2-25 when dependent on claim 2, wherein: (a) the network input comprises at least one of one or more of image data, audio data, video data; and the output sequence comprises data describing content of the image data, audio data or video data; (b) the network input comprises data defining a context, and the output sequence comprises image data, audio data or video data having content defined by the context;Attorney Docket No.: 45288-0544WO1 (c) the network input comprises sensor data output by a sensor and describing a real- world environment, and the output sequence comprises control data for controlling an agent in the real-world environment described by the sensor data; or (d) the network input comprises electronic health record data for a patient and the output sequence comprises a treatment suggestion, diagnosis or prognosis for the patient.

27. 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 or more computers to perform the respective operations of any one of claims 1-26.

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

29. A computer-implemented method comprising: obtaining data specifying a pre-trained neural network, wherein the pre-trained neural network comprises a sequence of layers; and training one or more early-exit adaptors for the pre-trained neural network, wherein each early-exit adaptor is associated with a respective layer from a subset of the sequence of layers and is configured to generate an adapted embedding from an updated embedding generated by the associated layer.

30. The method of claim 29, wherein training the one or more early-exit adaptors comprises training each of the one or more early-exit adaptors on a plurality of training examples that each comprise (a) a training input that includes one or more tokens, and (b) a target output that includes a ground-truth sequence of output tokens following the training input.

31. The method of any one of claims 29-30, wherein parameters of the sequence of layers are held fixed during the training of the one or more early-exit adaptors.Attorney Docket No.: 45288-0544WO1 32. The method of any one of claims 29-31, wherein the pre-trained neural network comprises a set of low-rank approximation weights, and wherein the parameters of the sequence of layers and the set of low-rank approximation weights are held fixed during the training of the one or more early-exit adaptors.

33. The method of any one of claims 29-32, wherein each early-exit adaptor is trained independently from each other early-exit adaptor.

34. The method of any one of claims 29-33, wherein training the one or more early-exit adaptors comprises training the one or more early-exit adaptors on a layer-wise loss function that measures, for each layer, a difference between an output probability distribution and a ground-truth probability distribution.

35. The method of claim 34, wherein the output probability distribution is based on a weighted combination of a first probability distribution given a determination to exit at the respective layer and a second probability distribution given a determination not to exit at the respective layer.

36. The method of claim 35, wherein the first probability distribution is weighted by a probability that a termination criterion is satisfied for the layer, and wherein the second probability distribution is weighted by a probability that the termination criterion is not satisfied for the layer.

37. 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 or more computers to perform the respective operations of any one of claims 29-36.

38. One or more computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform the respective operations of the method of any one of claims 29-36.Attorney Docket No.: 45288-0544WO1 39. 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 or more computers to implement a decoder neural network comprising: a sequence of layers, each layer configured to process an input embedding to generate an updated embedding; and a respective early-exit adaptor for each layer of a subset of the layers in the sequence, wherein each respective early-exit adaptor is configured to process the updated embedding from the corresponding layer to generate an adapted embedding that is aligned with a final updated embedding of a last layer of the sequence of layers.

40. The system of claim 39, wherein each respective early-exit adaptor comprises a multilayer perceptron (MLP).

41. The system of claim 39, wherein each respective early-exit adaptor comprises one of: a linear projection layer; a set of low-rank adaptation (LoRA) weights; or a LoRA projection layer.

42. The system of any one of claims 39-40, wherein the decoder neural network is configured to process the updated embedding from the corresponding layer using the respective early-exit adaptor for the corresponding layer when a termination criterion is satisfied.

Images