Efficient decoding of output sequences using adaptive early exit

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for generating an output sequence using an autoregressive decoder neural network. In particular, during generation, adaptive advance exit is used to reduce the time required to generate the output sequence.

Description

Background Art

[0001] This specification relates to using a neural network to process an input to generate an output sequence.

[0002] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict the output of a received input. In addition to the output layer, some neural networks also include one or more hidden layers. The output of each hidden layer is used as the input of the next layer in the network (i.e., another hidden layer or the output layer). Each layer of the network generates an output from the received input according to the current values ​​of the corresponding set of parameters. Summary of the invention

[0003] This specification describes a system implemented as a computer program on one or more computers at one or more locations that uses a decoder neural network to perform a machine learning task on a network input.

[0004] In particular, the system generates an output sequence that includes a corresponding word-gram from the word-gram vocabulary at each of a plurality of output time steps.

[0005] When generating the output sequence, the system can perform "early exiting". Early exiting refers to using the output of an intermediate layer of the decoder neural network instead of the output of the last layer of the decoder neural network to generate the output token. Performing early exiting can make the generation of tokens more computationally efficient because the computation of any layer after the intermediate layer does not need to be performed.

[0006] "Early exit" is called "adaptive" because the system determines to use the output of an intermediate layer for a given output token only when the termination criterion is met. If any of the intermediate layers do not meet the termination criterion when generating a given output token, the system uses the output of the last layer to generate the given output token.

[0007] The subject matter described in this specification can be implemented in specific embodiments to achieve one or more of the following advantages.

[0008] Recent advances in Transformer-based large language models (LLMs) have brought significant performance improvements across many tasks. These gains are accompanied by a dramatic increase in the size of the model, potentially resulting in slow and expensive use at inference time. However, in practice, a series of generation by the LLM when generating an output sequence consists of different difficulty levels. While some predictions do benefit from the full capacity of the model, other continuations are more trivial and can be solved by reduced computation. This specification describes a framework for dynamically allocating different amounts of computation per input and generation time step using "early exit decoding". In particular, in early exit decoding, the system determines when to stop processing and generate the next output token at each time step. By using the described framework, the system can reduce computation, for example, achieving a potential acceleration of three times or more when generating an output sequence, while provably maintaining high performance. The system can also reduce latency during generation while maintaining high performance, and can use "early exit decoding" to improve the overall throughput of the system. The system can also reduce memory usage when skipping the processing of subsequent layers. The system efficiently and dynamically allocates computing resources as needed when processing each position / at each time step.

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

[0010] Other features, aspects, and advantages of the subject matter will be apparent from the description, drawings, and claims. BRIEF DESCRIPTION OF THE DRAWINGS

[0011] Figure 1 is a diagram of an example sequence generation system.

[0012] Figure 2 is a flow chart of an example process for generating an output sequence.

[0013] Figure 3 is a flow chart of an example process for processing a sequence of hidden states using layers of a decoder neural network.

[0014] Figure 4 is a flow chart of an example process for determining a sharing threshold.

[0015] Figure 5 An example of using early exit to generate an output sequence is shown.

[0016] Like reference numbers and designations in the various drawings indicate like elements. DETAILED DESCRIPTION

[0017] Figure 1An example sequence generation system 100 is shown. Sequence generation system 100 is an example of a system implemented as a computer program on one or more computers in one or more locations in which the systems, components, and techniques described below may be implemented.

[0018] System 100 uses a neural network to perform machine learning tasks on network input 102.

[0019] The machine learning task may be any machine learning task that generates a network output, which is the output sequence 112. That is, the machine learning task requires processing the network input 102 to generate the output sequence 112.

[0020] The following are some examples of machine learning tasks that system 100 may be configured to perform.

[0021] As an example, if the input to a neural network is a sequence of text in one language, the output generated by the neural network may be a sequence of text in another language that is a proper translation of the input text into the other language.

[0022] As another example, the task may be an audio processing task.

[0023] For example, if the input to a neural network is a sequence representing spoken utterances, the output generated by the neural network may be a textual transcription of the utterances.

[0024] As another example, if the input to a neural network is a sequence representing spoken utterances, the output generated by the neural network may indicate whether a particular word or phrase (a "hot word") was spoken in the utterances.

[0025] As another example, the task may 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, waveform, or other data defining the audio of the text spoken in the natural language.

[0026] As another example, the task may be a text generation task where the input is a text sequence and the output is another text sequence, e.g., a completion of the input text sequence, a response to a question posed in the input sequence, or a text sequence about a topic specified by the first text sequence.

[0027] As another example, the input to the text generation task may be an input other than text, such as an image, and the output sequence may be text describing the input. As another example, the input may include both text and other data (e.g., an image).

[0028] As another example, the task may be a computer code generation task, where the input is a text sequence describing the functionality of a piece of computer code, or a computer code sequence to be modified or completed, or both, and the output is a computed code sequence that modifies the computer code, a computed code sequence having the functionality described by the text sequence, or a computed code sequence of both.

[0029] As another example, the task may be an image generation task where the input is a conditional input and the output is an input sequence of intensity values ​​for pixels of an image.

[0030] As another example, the task may be an agent control task, where the input is a sequence of observations or other data characterizing the state of the environment, and the output defines an action to be performed by the agent in response to the latest data in the sequence. The agent may be, for example, a real-world or simulated robot, a control system for an industrial facility, or a control system that controls different kinds of agents.

[0031] To perform machine learning tasks, system 100 uses a neural network that includes an autoregressive decoder neural network 110 that includes a sequence of multiple layers 120.

[0032] Each layer 120 operates on a respective input sequence comprising, at each of one or more positions, a respective input vector.In this specification, the input vectors will also be referred to as "hidden states".

[0033] Furthermore, each of layers 120 includes an attention layer and, in some implementations, a feed-forward layer.

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

[0035] The attention mechanism applied by the attention layer depends on the configuration of the attention neural network, as will be described in more detail below.

[0036] The feed-forward layer (when included) then operates on the attended input sequence to generate an output sequence of the layer.

[0037] More specifically, the decoder neural network 110 operates autoregressively to generate outputs in an output sequence, and during training, attention sublayers within some or all layers of the decoder apply masked self-attention (also known as “causal” self-attention) to the partially generated output sequences.

[0038] In masked self-attention, the attention mechanism is masked so that any given position in the input sequence does not pay attention to any positions after the given position in the input sequence.

[0039] In general, the layers 120 within a neural network may be arranged in any of a variety of configurations.

[0040] As an example, the neural network may include an encoder neural network that encodes the network input 102 to generate an encoded representation of the input, e.g., a sequence of encoded vectors representing the network input 102. For example, when the input is a sequence, the encoded representation may include a corresponding embedding for each input in the sequence. The encoder neural network may also be an attention-based neural network or another suitable type of neural network that maps the network input to an encoded representation.

[0041] When the neural network includes both an encoder and a decoder 110, some of the layers in the decoder 110 apply cross attention to the encoded representations, while other layers apply self attention to the output sequence.

[0042] In some implementations, the decoder neural network 110 operates directly on the input sequence, i.e., the neural network includes only the decoder 110 and not the encoder. For example, the input sequence may be the network input 102, or may have been generated by another neural network (e.g., by an embedding neural network that transforms the network input 102 into a sequence of tokens) through the network input 102.

[0043] In these implementations, at each time step, the decoder 110 operates on a combined sequence comprising the input sequence and the currently generated output sequence as of that time step, and the layers 120 in the sequence each perform masked self-attention on the combined sequence.

[0044] In other words, the decoder neural network 110 is referred to as an autoregressive neural network because the neural network 110 autoregressively generates an output sequence of tokens by generating each particular token in the output sequence conditioned on the current input sequence, including any tokens preceding the particular text token in the output sequence (i.e., tokens that have been generated for any previous positions in the output sequence preceding the particular position of the particular token) and, in some cases, sequences derived from the network input 102.

[0045] More specifically, to generate a particular word-gram at a particular position within the output sequence, the neural network 110 may process the current input sequence to generate a probability distribution that assigns a corresponding score (eg, a corresponding probability) to each word-gram in the word-gram vocabulary.

[0046] The tokens in the vocabulary may be any suitable text tokens, such as words, word fragments, characters, punctuation marks, etc. representing text elements of one or more natural languages, and optionally, numbers and other text symbols found in a text corpus.

[0047] The neural network 110 may then use the probability distribution to select a word-gram from the vocabulary as the particular word-gram. For example, the neural network 114 may greedily select the word-gram with the highest score, or may sample word-grams from the distribution, for example, using kernel sampling or another sampling technique.

[0048] The decoder neural network 110 includes (i) multiple layers that each apply an attention mechanism and (ii) an output subnetwork that processes the output of the last attention block to generate a score distribution.

[0049] In other words, the decoder neural network 110 includes a sequence of layers, and during processing of a given input sequence, each layer in the sequence receives a corresponding input hidden state for each input token in the given input sequence. The layer then updates each of the hidden states at least in part by applying an attention mechanism to generate a corresponding output hidden state for each of the input tokens.

[0050] The input hidden state of the first attention block is the embedding of the input tokens in the input sequence, and the input hidden state of each subsequent attention block is the output hidden state generated by the previous attention block.

[0051] In this example, the output subnetwork processes the output hidden state generated by the last attention block in the sequence for the last input token in the input sequence to generate a probability distribution. For example, the output subnetwork may include a linear layer followed by a softmax layer. As another example, the output subnetwork may include more layers, for example, a multi-layer perceptron (MLP).

[0052] Generally, unlike conventional systems, to autoregressively generate ("decode") an output sequence using decoder neural network 110, system 100 uses an adaptive early exit scheme to decode the output sequence in a computationally efficient manner and with reduced latency while still maintaining high output quality.

[0053] To employ early exit, the system 100 maintains a respective threshold at each of the output time steps for each layer in the subset of layers 120. Generally, within a given time step, each layer 120 in the subset has the same threshold, but different output time steps have different thresholds.

[0054] In some implementations, the subset of layers 120 includes all layers 120 in the sequence except the last layer in the sequence, ie, all intermediate layers 120 in the sequence.

[0055] In some other implementations, the subset of layers 120 includes less than all of the intermediate layers 120. For example, the subset may include all layers 120 in the sequence except the first and last layers. As another example, the subset may include every other intermediate layer in the sequence. As yet another example, the subset may include all layers to which self-attention is applied, but not the layers to which cross-attention is applied.

[0056] The system 100 also enhances the neural network with a corresponding output sub-network for each of the layers 120 in the subset. That is, the decoder neural network 110 includes not only the output sub-network of the last layer 120 in the sequence, but also the corresponding output sub-network of each layer 120 in the subset.

[0057] The corresponding output subnetwork of each of the layers 120 is configured to, at any given output time step, process the updated hidden state of the last word-gram in the input sequence at the given output time step, after being updated by the layer 120, to generate a probability distribution over the word-grams in the vocabulary.

[0058] For example, each output sub-network may have the same architecture (but different parameter values) as the output sub-network that processes the output of the last layer 120 in the sequence.

[0059] When generating an output token at a given output time step, the early exit engine 150 within the system 100 uses the maintained thresholds to determine whether a termination criterion is met for each layer in the subset.

[0060] If the termination criteria are met at a given layer in the subset, the early exit engine 150 determines to process the output of the given layer using the output subnetwork of the layer to select the token at the output time step. Thus, the system 100 can perform an "early exit" and avoid performing computations for any layer after the given layer in the sequence.

[0061] If the termination criterion is not met at any of the layers 120 in the subset, the early exit engine 150 uses the output of the last layer to generate an output token at the output time step, i.e., by processing the outputs of the layers in the sequence using the output subnetwork of the last layer in the sequence as described above.

[0062] References below Figures 2 to 5 The use of thresholds to determine whether termination criteria are met is described in more detail.

[0063] Prior to using the neural networks, system 100 or another training system trains the neural networks, i.e., decoder neural network 110 and (when included) encoder neural network.

[0064] For example, a neural network can be pre-trained on a language modeling task, e.g., a task requiring, given a sequence of current text tokens, to predict the next token following the current sequence in the training data. As a specific example, a neural network can be pre-trained on a large text dataset (e.g., text publicly available from the Internet or another text corpus) according to a maximum likelihood objective.

[0065] In some implementations, after pre-training, the system 100 fine-tunes the neural network according to the goal of the machine learning task, for example, through supervised learning, reinforcement learning, or both.

[0066] As another example, system 100 or another training system may train a neural network according to the objectives of a machine learning task without pre-training.

[0067] As a specific example, during training according to a given objective (i.e., a pre-training objective or an objective for a task), the system can train a neural network to minimize an overall loss that is a weighted sum of the losses of each of the layers in the subset and the last layer.

[0068] The loss for each layer in the subset is equal to the given objective evaluated using the output of the output subnetwork corresponding to the layer.

[0069] In some implementations, the system can set the weights in the weighted sum to favor later layers in the sequence, i.e., so that each layer has a higher weight than any layer that precedes it in the sequence. As a specific example, the weight of any given layer can be equal to or proportional to the index of the layer within the sequence divided by the sum of the indices of the layers in the subset and the last layer. Setting the weights to favor later layers can ensure that the accuracy of predictions generated from later layers is not reduced by attempting to improve the accuracy of predictions generated from earlier layers.

[0070] Figure 2 is a flow chart of an example process 200 for generating an output sequence. For convenience, 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 sequence generation system (e.g., Figure 1 The sequence generation system 100 depicted in FIG. 1 may perform process 200 .

[0071] As described above, the output sequence includes a corresponding word-gram at each of a plurality of time steps.

[0072] The system may perform process 200 at each output time step to generate a corresponding word-gram at that output time step.

[0073] That is, the system may continue to execute process 200 to add word-grams to the output sequence until a termination criterion is met, e.g., until the number of word-grams in the output sequence reaches a maximum number, or until the system selects a predetermined end-of-sequence word-gram from the vocabulary indicating the end of the output sequence.

[0074] The system generates a current input sequence for the output time step based on at least a word-gram at an output time step before the output time step in the output sequence (step 202).

[0075] That is, when the decoder is part of a decoder-only neural network, the current input sequence includes one or more input tokens corresponding to the network input, the token at the output time step before the output time step in the output sequence, and optionally, one or more specified separator tokens that separate the token corresponding to the network input from the tokens that have been generated. When the decoder is part of an encoder-decoder neural network, the current input sequence includes the token at the output time step before the output time step in the output sequence and, in some cases, a specified sequence start token.

[0076] The system generates a corresponding hidden state for each input in the current input sequence (step 204).

[0077] To generate the corresponding hidden state, the system processes each of the tokens at the output time step preceding that output time step in the output sequence (as well as any additional tokens in the current input sequence) using an embedding layer.

[0078] In some implementations, since only the last token in the output sequence has been added to the current input sequence since the previous output time step, the system reuses the corresponding hidden states (previously generated using the embedding layer at the previous time step) for other tokens in the current input sequence instead of recalculating the corresponding hidden states at each time step.

[0079] The system processes the corresponding hidden states of the inputs in the current input sequence through the layers in the sequence of layers until a termination criterion is met (step 206).

[0080] As described above, the system maintains a respective threshold value for each of a plurality of output time steps for each layer in the subset. The system uses the respective threshold value for the current output time step to determine when the termination criterion is met.

[0081] In particular, the system may maintain different thresholds for two or more of the output time steps.

[0082] More specifically, in order to maintain the quality of the generated output sequence, the corresponding thresholds at later output time steps in the output sequence can be lower than the corresponding thresholds at relatively earlier output time steps in the output sequence. That is, the system maintains the thresholds so that the termination criteria are more likely to be met at later time steps. This can maintain the quality of the generated output sequence because, due to the autoregressive nature of the generation process, inaccurately selected tokens at earlier time steps may have a greater impact on the overall sequence quality.

[0083] As a specific example, the corresponding thresholds for the output time steps may be determined based on the shared thresholds.

[0084] As a specific example, for each layer in the subset, a corresponding threshold for each output time step may be generated by applying a decreasing function to the shared threshold and the time index of the output time step.

[0085] The decreasing function may be any suitable function that assigns smaller thresholds to time steps with larger time indices. As a specific example, the decreasing function is an exponential function.

[0086] More specifically, the function may be a decay function such that at the output time step t, the shared threshold Generated Threshold satisfy:

[0087]

[0088] in is a function that clips its input to the range between zero and one (inclusive), N is the maximum output length, is a temperature value that may be received as input by a user or may be predetermined, and a is a value between zero and one (excluding zero and one). For example, a may be equal to .7, .8, or .9.

[0089] By using a shared threshold for all layers and generating the threshold for each time step based on the shared threshold, the system can greatly reduce the search space of possible effective thresholds, thereby allowing the system to easily determine the shared threshold that leads to a high-performance early exit scheme.

[0090] References below Figure 3 The use of these thresholds to determine when termination criteria are met is described in more detail.

[0091] References below Figure 4 Determining the sharing threshold is described in more detail.

[0092] Once the termination criteria are met, the system processes the updated hidden states generated by the layer at which the termination criteria were met that satisfied the last input in the current input sequence to generate a probability distribution of word-grams in the vocabulary (step 208).

[0093] As described above, for each layer in the subset, the system maintains a corresponding output subnetwork for each layer and the last layer in the subset.

[0094] Therefore, to generate a probability distribution, the system uses the output subnetwork of the layer at which the termination criterion is met to process the updated hidden state of the last input in the current input sequence generated by the layer at which the termination criterion is met to generate a probability distribution.

[0095] The system uses the probability distribution to select a word-gram at the output time step (step 210). For example, the system may select the word-gram with the highest probability, or may sample word-grams according to the probability distribution.

[0096] Thus, if the termination criteria is met at a given layer before the last layer in the sequence, the system avoids performing any processing on any additional layers after the given layer in the sequence at that output time step.

[0097] By not performing computations on these additional layers when the system determines to "exit early" when generating any given word, the system can generate the output sequence with reduced latency and consume fewer processor cycles at the same time. However, by utilizing a threshold to determine whether to "exit early", the system can maintain a high generation quality of the output sequence.

[0098] In some implementations, since only the last token in the output sequence is added to the current input sequence since the previous output time step, and since the causal mask on the self-attention mechanism ensures that the hidden state of the previously generated token is not updated at the time step after the token is generated, when executing processing of any given layer in the layers in the decoder, the system reuses the corresponding hidden states of other tokens that have been generated in the current input sequence, rather than recalculating the corresponding hidden states at each time step when executing step 206.

[0099] In these implementations, due to early exit, for some output time steps, at any given layer there may not be stored the hidden states of one or more layers preceding the given layer in the sequence. That is, if the system does not perform processing of a given layer when generating a given token at a given time step due to early exit, then the updated hidden states of that layer for that token at that time step will not be available for reuse when needed at a later time step.

[0100] For example, if the system determines early exit after the third layer in the sequence for a particular output token, but has not yet exited early for the next token when processing the next token through the fifth layer in the sequence, then the hidden state generated by the fourth layer from the particular output token is not available to be provided as input to the fifth layer in the sequence.

[0101] With this in mind, the system can take advantage of state propagation. In state propagation, if the system has exited early at layer j when generating a token at time step s, then for all time steps after s and for all layers after layer j in the sequence, the system sets the input hidden state of the token at time step s to be equal to the updated hidden state generated by layer j for the token at time step s. Thus, the system "copies" the last updated hidden state to later layers in the sequence as needed when generating later tokens in the sequence.

[0102] Figure 3 is a flow chart of an example process 300 for processing a hidden state sequence using a layer of a decoder neural network. For convenience, 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 sequence generation system (e.g., Figure 1 The sequence generation system 100 depicted in FIG. 1 may perform process 300 .

[0103] The system receives a corresponding input hidden state of the layer for each input in the current input sequence (step 302).

[0104] The system updates the respective input hidden state for each of the inputs in the current input sequence using the respective input hidden state at least in part by applying an attention mechanism (step 304).

[0105] For example, the system may process the input hidden states using an attention sublayer of a layer that applies an attention mechanism to generate an initial updated hidden state for each of the inputs, and then process the initial updated hidden states using a feedforward sublayer of the layer to generate an updated input hidden state for the layer.

[0106] As described above, when the decoder is part of a decoder-only neural network, the attention mechanism used for each layer is a causal (or "masked") self-attention mechanism. When the decoder is part of an encoder-decoder neural network, the attention mechanism used for some of the layers is a causal (or "masked") self-attention mechanism, while the attention mechanism for other layers is a criss-cross attention mechanism that criss-crosses the encoded representations of the network inputs.

[0107] When the layer is not in the subset and is not the last layer in the sequence, the system determines that the termination criteria are not met and provides the updated hidden state as the input hidden state of the next layer in the decoder neural network.

[0108] When the layer is in the subset, the system generates a confidence score for the layer based on at least the updated corresponding input hidden state generated by the layer for the last input in the current input sequence (step 306).

[0109] The system may generate confidence scores for the layers in any of a variety of ways.

[0110] As an example, the system can generate a probability distribution of word-units in the vocabulary based on the updated corresponding input hidden state of the last input in the current input sequence. For example, the system can use the output subnetwork of the layer to process the updated corresponding input hidden state of the last input in the current input sequence to generate a probability distribution.

[0111] The system can then determine a confidence score based on the difference between the highest probability in the probability distribution and the second highest probability in the probability distribution. For example, the confidence score can be equal to the difference or equal to the output of an increasing function applied to the difference.

[0112] As another example, the system may determine a confidence score based on a similarity measure between (i) the updated corresponding input hidden state generated by the layer for the last input in the current input sequence and (ii) the updated corresponding input hidden state generated by the previous layer in the sequence for the last input in the current input sequence. For example, the confidence score may be equal to the cosine similarity between (i) the updated corresponding input hidden state generated by the layer for the last input in the current input sequence and (ii) the updated corresponding input hidden state generated by the previous layer in the sequence for the last input in the current input sequence.

[0113] As yet another example, the system can use an early exit classifier to process the updated corresponding input hidden state of the last input in the current input sequence to generate a confidence score.

[0114] The early exit classifier is a machine learning model configured to map an input vector (i.e., an input hidden state) to a confidence score. For example, the early exit classifier can be a linear model. As another example, the early exit classifier can be a multilayer perceptron (MLP). As yet another example, the early exit classifier can be a Transformer neural network. In this example, the early exit classifier can also receive other data (e.g., other updated hidden states) as input.

[0115] In particular, the early exit classifier has been trained, i.e., trained by the system or another training system, while keeping the parameters of the decoder neural network fixed to the trained values ​​determined by training the decoder neural network. Therefore, adding the early exit classifier does not have any impact on the basic performance of the decoder neural network.

[0116] As a specific example, the system can train an early exit classifier on a cross entropy loss that measures, for each layer in the subset, the error between a confidence score predicted by the early exit classifier by processing the hidden state of the layer and an oracle score (e.g., a cross entropy loss) that is equal to 1 when the highest scoring token in the probability distribution generated from the hidden state of the layer is the same as the highest scoring token in the probability distribution generated from the hidden state of the last layer, and equal to 0 otherwise. For example, the loss for a given output token in a given training sequence can be the average of the errors of the layers in the subset.

[0117] The system determines that the termination criteria are met (step 308) when the confidence score of the layer is greater than or equal to the threshold of the layer for the output time step, and determines that the termination criteria are not met when the confidence score of the layer is less than the threshold of the layer for the output time step.

[0118] Figure 4 is a flow chart of an example process 400 for determining a sharing threshold. For convenience, process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a sequence generation system (e.g., Figure 1 The sequence generation system 100 depicted in FIG. 1 may perform process 400 .

[0119] The system receives input from a user defining values ​​for tolerance constraints (step 402).

[0120] The tolerance constraint is a constraint on the quality of the output sequence generated by using a corresponding threshold of output time steps relative to always using the updated hidden state generated by the last layer in the sequence to generate the output sequence. That is, the tolerance constraint represents an acceptable drop in quality (or accuracy) when the system uses a threshold instead of always using the hidden state generated by the last layer to employ "early exit".

[0121] As an example, textual consistency can be used to measure the quality of an output sequence. In this example, a decoder neural network is textually consistent (i.e., does not violate a constraint) when, at inference time given a network input, the expected textual dissimilarity (as measured by a textual dissimilarity function) between: (i) the output sequence generated by the network input using a corresponding threshold of output time steps, and (ii) the expected textual difference between the output sequence of the network input using an updated hidden state generated by the last layer in the sequence.

[0122] As another example, risk consistency can be used to measure the quality of an output sequence. In this example, a decoder neural network is risk consistent (i.e., does not violate a constraint) when the difference between the following two terms is less than or equal to a tolerance constraint value: (i) the expected risk for a given network input between a) the output sequence generated by the network input using the corresponding threshold of the output time step and b) the true value output sequence of the network input, as measured by the textual risk function, and (ii) the expected risk for a given network input between c) the output sequence of the network input using the updated hidden state generated by the last layer in the sequence and d) the true value output sequence of the network input, as measured by the textual risk function.

[0123] The system can compute the text consistency or text risk function using any appropriate quality metric m that compares the relative quality of two output sequences for a given task. For example, the text consistency function or text risk function for a pair of sequences can be equal to 1 - m, where m is the value of the quality metric for the pair of sequences. Examples of such quality metrics m include ROUGE-L, Token-F1, and BLEURT.

[0124] The system then determines the sharing threshold based on the tolerance constraint. In particular, the system determines the sharing threshold so that using the sharing threshold in the early exit scheme does not violate the tolerance constraint.

[0125] As one example, the system can determine a plurality of candidate values ​​(step 404). For example, the system can select a set of candidate values ​​that are evenly spaced on a grid of values ​​or randomly sampled from an evenly spaced grid of values.

[0126] The system may then select as the shared threshold the lowest candidate value that produces a corresponding threshold that satisfies the tolerance constraints (step 406 ).

[0127] For example, the system can obtain a calibration dataset. The calibration dataset can be provided or identified, for example, by a user as representing network inputs to be processed by the neural network at inference time.

[0128] The calibration dataset includes the set of network inputs and, when risk consistency is used, the true value (or "golden") output sequence for each network input.

[0129] The system can then apply a Learn then Test calibration framework on the calibration data set to determine a subset of candidate values ​​that do not violate the tolerance constraints (a subset of "valid" candidate values) and then select the minimum candidate value in the subset as the shared threshold.

[0130] In general, to perform a learn-before-test calibration framework, the system tests candidate values ​​by calculating a p-value for the candidate value using expected text consistency or expected risk consistency, and rejects the candidate value if the p-value for the candidate value is less than or equal to a level value. The level value specifies the probability that the decoder neural network needs to satisfy the tolerance constraint, and can be predetermined or received as input from a user along with the tolerance constraint value. Rejected candidate values ​​are determined to be valid candidate values, and non-rejected candidate values ​​are not included in the subset.

[0131] As an example, the system can use the random variables of the network input i in the calibration dataset The empirical mean of the given candidate value is calculated The p-value of the random variable is calculated using text consistency or risk consistency .

[0132] For example, for text consistency, the random variable can be equal to the output of the text dissimilarity function for the following two terms: (i) the network input i is obtained by using the given candidate value The output sequence generated by the corresponding thresholds of the derived output time steps (i.e., where the shared threshold is set equal to the given candidate value), and (ii) the output sequence generated by the network input i using the updated hidden states generated by the last layer.

[0133] As another example, for risk consistency, the random variable can be equal to the maximum of the difference between zero and the following two terms: (i) for a) network input i by using the given candidate value The output sequence generated by the corresponding threshold at the derived output time step (i.e., where the shared threshold is set equal to the given candidate value) and b) the output of the text risk function for the true-valued output sequence of network input i, and (ii) the output sequence generated for c) the network input i using the updated hidden states generated by the last layer and d) the output of the text risk function for the true-valued output sequence of network input i.

[0134] For example, the system can set the p-value equal to:

[0135] ,

[0136] Where n is the total number of network inputs i, is the tolerance constraint value, and yes of The empirical average value of .

[0137] As a specific example, to perform a learn-then-test calibration, the system may sort the candidate values ​​in descending order starting with the largest candidate value.

[0138] The system may continue testing each candidate value in order and, upon failing to reject a given candidate value, designate the previous candidate value in order as the shared threshold.

[0139] If the first candidate value is rejected, the system may set the shared threshold to one, i.e., so that the system does not perform early exit, since this would not result in the decoder neural network satisfying the constraints.

[0140] Figure 5 An example 500 of using early exit to generate an output sequence is shown.

[0141] like Figure 5 As shown, the decoder neural network includes 8 layers, and all layers except the last one are in the subset.

[0142] like Figure 5 As shown, at each output time step (each "position" in the output sequence), early exit engine 150 may use early exit to determine which layer outputs to use to generate a word-gram at that output time step.

[0143] For example, in Figure 5 In the example of , the system has generated a sequence including the word "Grant Elliott hit a six toput". When generating the word-grams "Grant" and "Elliott", the system uses the output of layer 1. However, when generating the word-gram "hit", the system uses the output of layer 7, and when generating the word-gram "to", the system uses the output of layer 8.

[0144] like Figure 5 As shown, the early exit engine 150 uses the used tolerance constraint value at each time step and the level value ϵ to determine the threshold. Figure 5In the example of , these two values ​​have been specified by the user to indicate the level of accuracy degradation that the user will accept when generating any given output sequence. Thus, these user-defined values ​​for the "global" (i.e., sequence-wide) tolerance for accuracy degradation guide the local per-time-step early exit decisions made by the early exit engine 150.

[0145] Figure 5 Also shown is an example of state propagation 502. In particular, as shown in example 502, when generating a word-gram at position (output time step) 6, the output of layer 4 for the word-grams at output time steps 1, 2, 4, and 5 is not available (to be provided as input hidden states for layer 5) when performing processing of layer 5 because layer 4 was skipped due to early exit during the generation of these word-grams. Therefore, the system uses the updated hidden states generated by the last layer at which processing was performed when generating the word-grams at those positions as the input hidden states for layer 5 for the word-grams at those positions. In example 502, the system (i) uses the hidden state generated by layer 1 as the hidden state at output time step 1, (ii) uses the hidden state generated by layer 1 as the hidden state at output time step 2, (iii) uses the hidden state generated by layer 2 as the hidden state at output time step 4, and (iv) uses the hidden state generated by layer 1 as the hidden state at output time step 5.

[0146] Therefore, by employing state propagation, the system can avoid incurring additional processing resources to recalculate hidden states that were skipped due to early exit at previous time steps.

[0147] This specification uses the term "configured" in connection with system and computer program components. For a system of one or more computers to be configured to perform a particular operation or action, it means that the system has installed thereon software, firmware, hardware, or a combination thereof that, in operation, causes the system to perform the operation or action. For one or more computer programs to be configured to perform a particular operation or action, it means that the one or more programs include instructions that, when executed by a data processing device, cause the device to perform the operation or action.

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

[0149] The term "data processing equipment" refers to data processing hardware and encompasses all kinds of equipment, devices and machines for processing data, including, for example, a programmable processor, a computer or multiple processors or computers. The equipment may also be or include a dedicated logic circuit system, for example, an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). In addition to hardware, the equipment may optionally include code that creates an execution environment for a computer program, for example, code constituting processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0150] A computer program, which may also be referred to or described as a program, software, software application, app, module, software module, script, or code, may be written in any form of programming language (including compiled or interpreted languages, or declarative or procedural languages); and it may 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 may be stored in a portion of a file that holds other programs or data (e.g., one or more scripts stored in a language document of a word element), in a single file dedicated to the program in question, or in multiple coordinated files (e.g., files storing one or more modules, subroutines, or code portions). A computer program may be deployed to execute on one computer or on multiple computers located at one site or distributed over multiple sites and interconnected by a data communications network.

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

[0152] Similarly, in this specification, the term "engine" is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a specific engine; in other cases, multiple engines can be installed and run on the same one or more computers.

[0153] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by a dedicated logic circuit system (e.g., an FPGA or ASIC) or by a combination of a dedicated logic circuit system and one or more programmed computers.

[0154] A computer suitable for executing a computer program can be based on a general or special microprocessor or both, or any other kind of central processing unit. Generally, the central processing unit will receive instructions and data from a read-only memory or a random access memory or both. The basic elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by a special logic circuit system or incorporated into the special logic circuit system. Generally, the computer will also include one or more large-capacity storage devices (e.g., magnetic disks, magneto-optical disks, or optical disks) for storing data, or be operably coupled to receive data from one or more large-capacity storage devices or transfer data to one or more large-capacity storage devices or both. However, the computer does not need to have such a device. In addition, the computer can be embedded in another device, for example, a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device (e.g., a universal serial bus (USB) flash drive), to name a few.

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

[0156] To provide interaction with a user, embodiments of the subject matter described in this specification may be implemented on a computer having a display device (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) for displaying information to the user, and a keyboard and pointing device (e.g., a mouse or trackball) through which the user can provide input to the computer. Other kinds of devices may also be used to provide interaction with the user; for example, the feedback provided to the user may be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback; and input from the user may be received in any form, including sound, voice, or tactile input. In addition, the computer may interact with the user by sending documents to and receiving documents from the device used by the user; for example, by sending a web page to a web browser on the user's device in response to a request received from the web browser. Moreover, the computer may interact with the user by sending a text message or other form of message to a personal device (e.g., a smart phone running a messaging application), and receiving a response message from the user in return.

[0157] A data processing device used to implement a machine learning model may also include, for example, dedicated hardware accelerator units for processing general-purpose and computationally intensive parts of machine learning training or production (i.e., inference, workloads).

[0158] Machine learning models can be implemented and deployed using a machine learning framework (e.g., the TensorFlow framework or the Jax framework).

[0159] Embodiments of the subject matter described in this specification may be implemented in a computing system that includes a back-end component (e.g., as a data server), or that includes a middleware component (e.g., an application server), or that includes a front-end component (e.g., a client computer with 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 components, middleware components, or front-end components. The components of the system may be interconnected by any form or medium of digital data communication (e.g., a communication network). Examples of communication networks include local area networks (LANs) and wide area networks (WANs), such as the Internet.

[0160] A computing system may include a client and a server. The client and the server are usually remote from each other and typically interact through a communication network. The relationship between the client and the server is due to computer programs running on respective computers and having a client-server relationship with each other. In some embodiments, the server sends data (e.g., an HTML page) to a user device, for example, to display data to a user interacting with a device acting as a client and to receive user input therefrom. Data generated at a user device can be received from the device at the server, for example, due to user interaction.

[0161] Although this specification contains many specific implementation details, these details should not be interpreted as limitations on the scope of any invention or the scope that may be claimed, but should be interpreted as descriptions of features that may be peculiar to a particular embodiment of a particular invention. Certain features described in the context of a separate embodiment in this specification may also be implemented in combination in a single embodiment. Conversely, the various features described in the context of a single embodiment may also be implemented individually or in any suitable sub-combination in multiple embodiments. In addition, although features may be described above as working in certain combinations and even initially claimed as such, in some cases, one or more features from the claimed combination may be deleted from the combination, and the claimed combination may involve a variant of a sub-combination or a sub-combination.

[0162] Similarly, although operations are depicted in the drawings and described in the claims in a particular order, this should not be construed as requiring that such operations be performed in the particular order shown or in a sequential order, or that all of the operations shown be performed, to achieve the desired result. In some cases, multitasking and parallel processing may be advantageous. In addition, the separation of various system modules and components in the embodiments described above should not be construed 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.

[0163] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the appended claims. For example, the actions recited in the claims can be performed in a different order and still achieve the desired results. As an example, the processes depicted in the accompanying drawings do not necessarily require the specific order shown or the ordered order to achieve the desired results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. A method executed by one or more computers for self-regressively generating an output sequence using a decoder neural network, the output sequence comprising a corresponding word-gram from a word-gram vocabulary at each of a plurality of output time steps, wherein the decoder neural network comprises a sequence of layers, Wherein the decoder neural network is configured as: receiving a respective hidden state for each input in an input sequence, the input sequence including a respective input at each of one or more input positions, processing respective hidden states of the inputs in the input sequence through the sequence of layers to generate respective updated hidden states for each input in the input sequence, wherein each layer is configured to: receiving a corresponding input hidden state for each input in the input sequence, and updating a respective input hidden state for each of the inputs using the respective input hidden state at least in part by applying an attention mechanism; as well as processing the updated hidden state of the last input in the input sequence to generate a probability distribution over the word-grams in the vocabulary, and The method comprises: maintaining, for each layer in the subset of layers in the sequence, a respective threshold for each of the plurality of output time steps, wherein two or more of the output time steps have different respective thresholds; and The output sequence is generated by performing the following operation at each of the plurality of output time steps: Generate a current input sequence based on at least a word-element at an output time step before the output time step in the output sequence; Generate a corresponding hidden state for each input in the current input sequence; processing corresponding hidden states of the input in the current input sequence through the layers in the sequence of layers until a termination criterion is satisfied, including until the termination criterion is satisfied, for each layer starting with a first layer in the sequence: receiving a corresponding input hidden state for the layer for each input in the current input sequence; and updating a respective input hidden state for each of the inputs in the current input sequence using the respective input hidden state at least in part by applying an attention mechanism; When the layer is in the subset, generating a confidence for the layer based on at least an updated corresponding input hidden state generated by the layer for the last input in the current input sequence; and When the confidence score of the layer is greater than or equal to the threshold of the layer for the output time step, determining that the termination criterion is satisfied; Once the termination criterion is satisfied, processing the updated hidden state of the last input in the current input sequence, generated by the layer at which the termination criterion is satisfied, to generate a probability distribution of the word-grams in the vocabulary; and The word-gram at the output time step is selected using the probability distribution.

2. The method of claim 1, wherein the output sequence is conditioned on the network input.

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

4. The method of claim 2, further comprising: The network input is processed using an encoder neural network to generate an encoded representation of the network input, and wherein one or more of the layers in the sequence applies criss-cross attention to the encoded representation of the network input.

5. A method as claimed in any preceding claim, wherein generating a corresponding hidden state for each input in the current input sequence comprises: Each of the tokens at an output time step preceding the output time step in the output sequence is processed using an embedding layer.

6. A method as claimed in any preceding claim, wherein processing the corresponding hidden states of the inputs in the input sequence through the layers in the sequence of layers until a termination criterion is satisfied comprises: The termination criterion is determined to be satisfied after the last layer in the sequence when the confidence score of any of the layers in the subset is not greater than or equal to the threshold for the layer for the output time step.

7. A method as claimed in any preceding claim, wherein processing the corresponding hidden state of the input by the layer comprises: Once the termination criterion is met at a given layer in the sequence, any processing on any additional layers following the given layer in the sequence is avoided at the output time step.

8. A method as claimed in any preceding claim, wherein generating a confidence score for the layer based at least on the updated corresponding input hidden state of the last input in the current input sequence comprises: generating a probability distribution of the word-grams in the vocabulary according to the updated corresponding input hidden state of the last input in the current input sequence; as well as The confidence score is determined based on a difference between a highest probability in the probability distribution and a second highest probability in the probability distribution.

9. The method of any one of claims 1 to 7, wherein generating a confidence score for the layer based at least on the updated corresponding input hidden state of the last input in the current input sequence comprises: The confidence score is determined based on a similarity measure between (i) the updated corresponding input hidden state generated by the layer for the last input in the current input sequence and (ii) the updated corresponding input hidden state generated by a previous layer in the sequence for the last input in the current input sequence.

10. The method of any one of claims 1 to 7, wherein generating a confidence score for the layer based at least on the updated corresponding input hidden state of the last input in the current input sequence comprises: The updated corresponding input hidden state of the last input in the current input sequence is processed using an early exit classifier to generate the confidence score.

11. The method of claim 10, wherein the early exit classifier has been trained while keeping parameters of the decoder neural network fixed to trained values ​​determined by training the decoder neural network.

12. A method as claimed in any preceding claim, Wherein for each layer in the subset, the corresponding threshold at a later output time step in the output sequence is lower than the corresponding threshold at a relatively earlier output time step in the output sequence.

13. The method of claim 12, further comprising: For each layer in the subset, a corresponding threshold for each output time step is generated by applying a decreasing function to the shared threshold and the time index of the output time step. The method of claim 13 , wherein the decreasing function is an exponential function.

15. A method as claimed in any preceding claim, wherein for each layer in the subset, the respective threshold for the output time step is determined based on a shared threshold.

16. The method of any one of claims 13 to 15, further comprising: receiving an input from a user, the input defining a value of a tolerance constraint on a quality of an output sequence generated by using a corresponding threshold relative to always using an updated hidden state generated by the last layer in the sequence to generate the output sequence; as well as The sharing threshold is determined based on the tolerance constraint.

17. The method of claim 16, wherein determining the sharing threshold based on the one or more tolerance constraints comprises: determining a plurality of candidate values; as well as The lowest candidate value that produces a corresponding threshold that satisfies the tolerance constraint is selected as the shared threshold.

18. A system comprising: one or more computers; as well as One or more storage devices storing instructions, which, when executed by the one or more computers, cause the one or more computers to perform corresponding operations as claimed in any one of claims 1 to 17.

19. 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 corresponding operations of the method of any one of claims 1 to 17.

Images