Cascade-aware training for language model neural networks

Abstract

Methods, systems, and apparatuses, including computer programs encoded on computer storage media, for training a student language model neural network for deployment in a cascade with a teacher neural network. That is, by training a student neural network using techniques that incorporate the difficulty of accurately predicting the target token for each output position of a target output of a training example for each training example for both the student and the teacher language model neural networks, the described techniques result in a student teacher cascade with higher overall task performance per unit of computational cost.

Description

[0001] CASCADE-AWARE TRAINING FOR LANGUAGE MODEL NEURAL NETWORKS

[0002] CROSS-REFERENCE TO RELATED APPLICATION

[0003] This application claims priority to U.S. Application No. 63 / 650,854, filed May 22, 2024. the disclosure of which is incorporated herein by reference.

[0004] BACKGROUND

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

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

[0007] SUMMARY

[0008] This specification describes a system implemented as computer programs on one or more computers in one or more locations that trains a student neural network for deployment in a cascade with a teacher neural network.

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

[0010] There can exist a tradeoff between task performance and computational cost when employing neural networks to carry out a machine learning task. That is, a neural network can have a greater performance for a higher computational cost (e.g., a language model neural network can process complex network inputs to generate reliable network outputs with high latency) or a lower performance for a lower computational cost (e.g.. a language model neural network only generates reliable network outputs for simple network inputs but with low latency).

[0011] In order to mitigate the undesirable need to tradeoff performance for computational cost, one can employ a cascade, in which two or more neural networks of various performance capabilities with various respective computational costs process a network input to generate a network output in order to perform a machine learning task. For example, a cascade of a single student neural network (with a lower performance and a lower latency) and a single teacher neural network (with a higher performance and a higher latency) can process a network input according to deferral rules (i.e. , a method of routing the network input to one neural network or the other to generate a network output) to maximize total performance over many network inputs for a given computational cost or latency. Put simply, this example cascade can perform a machine learning task by processing network inputs using the student neural network when its performance is sufficient to generate an acceptable network output and can use the teacher neural network when greater performance is needed to generate the network output. In this way, the higher cost / latency of using the teacher neural network is only paid when necessary.

[0012] In some prior art methods, the neural networks of a cascade are trained independently. Two challenges with the use of the cascade technique are how to route network inputs to the appropriate neural network and how to ensure that the neural networks complement each other to maximize the overall task performance of the cascade per unit of computational cost. These challenges are especially difficult when the cascade includes language model neural networks, because language model neural networks process network inputs at the token level (i.e.. process network inputs that include a sequence of tokens one token at a time) and generate network outputs at the token level (i.e., generate network outputs that include a sequence of tokens one token at a time) but the network inputs are routed at the sequence level (i.e., routing considers the sequence of tokens included in the network input in its entirety).

[0013] Since the training of a language model neural network and the processing and generating of network inputs and outputs occur at the token level but the use of the cascade to generate network outputs to perform a machine learning task is at the sequence level (i.e., the routing of network inputs occurs at the sequence level), the "token level to sequence level gap’’ prevents the language model neural networks of the cascade from efficiently complementing each other, leading to a suboptimal performance of the cascade. That is, the overall task performance of the cascade per unit of computational cost will not be maximized unless the language model neural networks of a cascade account for the presence of each other during training. But training neural networks of a cascade to account for the presence of each other is a challenge due to network input processing and sequence output generation occurring at the token level but network input routing occurring at the sequence level.

[0014] This specification describes techniques that can address the aforementioned challenges. That is, this specification describes techniques for training a student language model neural network to be deployed in a cascade with a teacher language model neural network. The techniques include receiving a set of one or more training examples, where each training example includes (i) a network input and (ii) a target output for the network input. The target output includes a respective target token at each of a plurality of output positions, and each target token is selected from a vocabulary of tokens. For each training example, the techniques include processing the network input in the training example using the student language model neural network to generate, for each output position, a respective student score distribution that assigns a respective student score to each token in the vocabulary' of tokens. Additionally for each training example, the techniques include processing the network input in the training example using the teacher language model neural network to generate, for each output position in the target output in the training example, a respective teacher score distribution that assigns a respective teacher score to each token in the vocabulary of tokens. Lastly, for each training example, the techniques include determining, for each of the output positions and from the respective teacher score distribution and the respective student score distribution for the output position, a respective weight for the output position that represents a difficulty of accurately predicting the target token at the output position. After determining the weights for each training example, the techniques include training the student neural network on a loss function that depends on, for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position and (ii) the student score assigned to the target token at the output position by the student score distribution for the output position.

[0015] By training the student language model neural network cascade using the abovedescribed loss function, the described techniques result in a student teacher cascade that includes a student trained ‘‘aware of its role in the cascade”. That is, the cascade is more efficient (i.e., the overall task performance of the cascade per unit of computational cost is higher or average latency can be reduced considering a particular workload of inputs) because the student language model neural network is trained considering the difficulty’ of accurately predicting the target token at the output position for both the student and the teacher language model neural networks.

[0016] After training using the described techniques, the resulting student teacher cascade has improved task performance relative to a conventional student teacher cascade for the same compute budget. That is. because the described techniques improve the performance of the student language model network in a way that enhances its performance for a scope of appropriate network inputs, the described techniques lower the rate that a network input needs to be deferred to the teacher language model neural network. Yet, the overall performance of the cascade is not sacrificed (because the improved student language model neural network can appropriately generate network outputs for a greater proportion of received network inputs) and the compute cost (e.g., number of floating point operations required to generate a network output for the network input) or latency of using the teacher language model neural network is avoided.

[0017] In some implementations, the student language model neural network and teacher language model neural network are deployed on different devices and operate as a distributed system. For example, the student language model neural network can be deployed on a resource constrained user device that a teacher language model neural network cannot be deployed on, e.g. the teacher language model neural network has too large a memory' footprint or has unacceptable latency to be deployed on the user device. The teacher language model neural network can be deployed on a remote server having the necessary computational resources and the user device can be configured to invoke the teacher language model neural network over a network connection when it is determined the student language model neural network is unlikely to be able to provide satisfactory' output. It follows that the described techniques not only improve the overall performance of the cascade for a given compute budget relative to a conventional cascade, but also reduce the rate of network input deferral and therefore the amount of data transmitted over a network, e.g., sending a network input to the teacher language model neural netw ork. In addition, the overall latency, which can include a network latency cost, for a given workload of inputs can be reduced. It will be appreciated that the student language model neural network and teacher language model neural network can also be implemented on different devices when training the student language model neural network.

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

[0019] According to a first aspect there is provided a method performed by one or more computers and for training a student language model neural network to be deployed in a cascade with a teacher language model neural network. The method includes receiving a set of one or more training examples, where each training example includes (i) a network input and (ii) a target output for the network input. The target output includes a respective target token at each of a plurality of output positions and each target token is selected from a vocabulary of tokens. The method further includes, for each training example, processing the network input in the training example using the student language model neural network to generate, for each output position, a respective student score distribution that assigns a respective student score to each token in the vocabulary of tokens. Then the method includes, for each training example, processing the network input in the training example using the teacher language model neural network to generate, for each output position in the target output in the training example, a respective teacher score distribution that assigns a respective teacher score to each token in the vocabulary7of tokens. After processing the network input in the training example using the student language model neural network and the teacher language model neural network, the method includes, for each training example, determining, for each of the output positions and from the respective teacher score distribution and the respective student score distribution for the output position, a respective weight for the output position that represents a difficulty of accurately predicting the target token at the output position. After performing the above for each training example, method includes training the student neural network on a loss function that depends on. for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position and (ii) the student score assigned to the target token at the output position by the student score distribution for the output position.

[0020] In some cases, the teacher language model neural network is an already -trained neural network that is held fixed during the training of the student language model neural network.

[0021] In some cases, the student language model neural network has been pre-trained prior to the training of the student language model neural network to be deployed in the cascade with the teacher language model neural network.

[0022] In some implementations, determining, for each of the output positions and from the respective teacher score distribution and the respective student score distribution for the output position, a respective weight for the output position that represents a difficulty of accurately predicting the target token at the output position includes determining whether the target token at the output position has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position. Then, determining the respective weight for the output position based on whether the target token at the output position (i) has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position or (ii) has been assigned a highest teacher score of any vocabulary token in the respective teacher score distribution for the output position.

[0023] Further in some implementations, determining the respective weight based on whether the target token at the output position (i) has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position or (ii) has been assigned a highest teacher score of any vocabulary token in the respective teacher score distribution for the output position includes assigning a first weight to the output position and assigning a second weight to the output position. In particular, it includes assigning a first weight to the output position when the target token at the output position (i) has been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position, (ii) has been assigned the highest teacher score of any vocabulary token in the respective teacher score distribution for the output position, or (iii) has been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position and has been assigned the highest teacher score of any vocabulary token in the respective teacher score distribution for the output position. Then it includes assigning a second weight to the output position when the target token at the output position (i) has not been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position and (ii) has not been assigned the highest teacher score of any vocabulary token in the respective teacher score distribution for the output position.

[0024] Further in some implementations, the first weight is one and the second weight is zero.

[0025] In some cases, the loss function measures, for each training example, a sum of, for each output position, a product of the respective weight for the output position and a negative log likelihood computed from the student score assigned to the target token at the output position by the student score distribution for the output position.

[0026] In some cases, the loss function is a distillation loss that depends on, for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position, (ii) the student score assigned to the target token at the output position by the student score distribution for the output position, (iii) the student score distribution for the output position, and (iv) the teacher score distribution for the output position. Further in some cases, the loss function measures, for each training example, a sum of. for each output position, a product of the respective weight for the output position and a token-level distillation loss. The token level distillation loss is computed based on a negative log likelihood computed from the student score assigned to the target token at the output position by the student score distribution for the output position and a difference measure between the student score distribution for the output position, and the teacher score distribution for the output position.

[0027] Further in some cases, the token-level distillation loss is a weighted sum of the negative log likelihood and the difference measure.

[0028] In some cases, the teacher language model neural network has more parameters than the student language model neural network.

[0029] According to a second aspect, there is provided the methods of the first aspect performed by a system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one more computers to perform the operations of the respective method.

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

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

[0032] BRIEF DESCRIPTION OF THE DRAWINGS

[0033] FIG. 1 shows a cascade training system.

[0034] FIG. 2 is a flow diagram of an example process for training a student language model neural network to be deployed in a cascade with a teacher language model neural network.

[0035] FIG. 3 is a flow diagram of an example process for generating an output for a task using a student language model neural network deployed in a cascade with a teacher language model neural network.

[0036] FIG. 4 is a flow diagram of an example process for updating trainable parameters of a student neural network.

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

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

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

[0040] DETAILED DESCRIPTION

[0041] This specification descnbes a system implemented as computer programs on one or more computers in one or more locations that trains a student neural network for deployment in a cascade with a teacher neural network. In particular, this specification describes techniques for training the student neural network using a paradigm that accounts for this deployment scenario and therefore improves the performance of the overall cascade after the student neural network has been trained. The training is therefore referred to as “cascade-aware training.”

[0042] In particular, after training, the system or another inference system receives a network input for performing a task using the student language model neural network. For example, the student language model neural network can be deployed on a user device and the system can receive the network input at the user device, e.g., from a user of the device.

[0043] The system processes the network input using the student language model neural network to generate an output sequence for the network input and to determine whether to defer the network input to the teacher neural network (i.e. whether processing of the network input by the teacher neural network is required).

[0044] In response to determining not to defer the network input, the system uses the output sequence generated by the student as the output for the task. In response to determining to defer the network input, the system provides the network input as an input to the teacher language model neural network. As one example, the teacher language model neural network can be deployed on one or more computers that are remote from the user device, and the system can provide the network input as input by providing the network input from the user device to the one or more remote computers over a data communication network.

[0045] As a particular example, the teacher neural network may be a larger neural network than the student neural network and that therefore cannot be deployed on the user device, e.g., because the larger neural network has an excessive memory footprint or because performing inference using the larger neural network given the limited processing capacity’ of the user device results in unacceptable latency.

[0046] The system can then obtain, as output from the teacher language model neural network, a "teacher" output sequence and use the teacher output sequence as output for the task.

[0047] Thus, the system can selectively defer difficult network inputs to the more computationally expensive and / or remotely located teacher neural network while processing other inputs using the student neural network.

[0048] Further details for generating an output for a task using a student language model neural network deployed in a cascade with a teacher language model neural network are described below.

[0049] FIG. 1 shows an example cascade training system 100. The cascade training 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.

[0050] The system 100 can perform cascade-aware training over multiple training steps.

[0051] At each of the multiple training steps, the system 100 receives a set of one or more training examples. For convenience, FIG. 1 displays the system 100 receiving a single training example 102, but it should be understood that the system 100 can receive and process more than one training example in parallel.

[0052] Each training example 102 includes (i) a network input 104 and (ii) a target output 106 for the network input 104, and the target output 106 includes a respective target token at each of a plurality of output positions, with each target token being selected from a vocabulary of tokens. Generally, the network input 104 can include any type of input data (e.g., natural language text data, audio data, image data, video data, any combination of these data, and so on) and can be represented as an input sequence, e.g., a sequence of natural language text, image pixels or patches, video frames, video frame patches, audio waveform time windows, spectrogram amplitude frequency -time windows, any combination of these elements, and so on. The system 100 can represent the input sequence as a sequence of tokens, e.g., sequence of text tokens, e.g., words, word pieces, bytes, characters, numbers, punctuation, or other text symbols and tokens representing other types of data, e.g., image data, video data, audio data, and so on. Likewise, the tokens of the target output 104 are defined in the same manner as those of the input sequence of tokens above. Thus, whilst the student and teacher neural networks are referred to as language model neural networks, neither are limited to processing of text only.

[0053] The system 100 can generate a sequence of tokens for the network input 104 by mapping the input sequence representation of the network input 104 to a sequence of tokens.

[0054] For example, if the network input 104 includes natural language text data, then the system 100 can, e.g., map each character, word, or sub-word of the natural language text representation to a corresponding token by applying a text tokenizer to the input text.

[0055] As another example, if the network input 104 includes audio data, then the system 100 can, e.g.. convert the audio into a spectrogram and map small segments (i.e.. frequency, time patches of the spectrogram) to corresponding tokens, e.g., by applying an audio encoder neural network.

[0056] As another example, if the network input 104 includes image data, then the system 100 can, e.g.. divide each image into patches or pixels and map each patch or pixel to a corresponding token, e.g., by applying an image encoder neural network to the patch embeddings.

[0057] As another example, if the network input 104 includes video, then the system 100 can, e.g., divide each video into a sequence of images and divide each image into patches or pixels and map each patch or pixel to a corresponding token. Alternatively, a token can represent a spatio-temporal portion of the video, e.g., by applying a video encoder neural network.

[0058] In some cases, the tokens of the input sequence representation of the network input 104 belong to a vocabulary (i.e., a finite set of tokens). Additionally, in some cases, the vocabulary is the same vocabulary as that of the tokens of the target output but, in other cases, it is a different vocabulary than that of the tokens of the target output.

[0059] For each training example 102, the system 100 processes the network input 104 in the training example 102 using the student language model neural network 108 to generate, for each output position, a respective student score distribution that assigns a respective student score to each token in the vocabulary of tokens.

[0060] The system also processes the network input 104 in the training example using the teacher language model neural network 110 to generate, for each output position in the target output 106 in the training example 102, a respective teacher score distribution that assigns a respective teacher score to each token in the vocabulary of tokens.

[0061] The system 100 determines, for each of the output positions and from the respective teacher score distribution and the respective student score distribution for the output position, respective weights 116 for the output positions that each represents a difficulty of accurately predicting the target token at the output position. For example, the difficulty can be based upon whether at least one of the student language model neural network or the teacher language model neural network has correctly predicted the target token. If neither neural network can correctly predict the target token, it may be considered to be too difficult for the student language model neural netw ork to leam given that is likely to have limited capacity.

[0062] The system 100 then trains the student neural network 108 on a loss function that depends on, for each training example 102 and for each output position in the target output 106 in the training example 102, (i) the respective weight for the output position and (ii) the student score assigned to the target token at the output position by the student score distribution for the output position. Generally, the loss function places less weight on output positions at which the target token is more difficult to predict. For example, the loss function can exclude, i.e., assign a zero weight to, output positions at which the target token is difficult to accurately predict, and only include, i.e., assign anon-zero weight, e.g., one, to, output positions at which the target token is not difficult to accurately predict. The student language model neural network can therefore focus its limited capacity on learning "easier" token predictions. This can increase the overall accuracy and confidence of predictions provided by the student language model neural network, particularly on the types of queries that the student language model neural network will be expected to handle in the cascade and reduce the number of times the teacher language model neural network will have to be invoked. An example of architectures and uses for the student neural network now follows. It will be understood that the teacher neural network can have similar architectures and uses as the student neural network.

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

[0064] In particular, to generate a particular token at a particular position within an output sequence, the language model neural network can process the current input sequence to generate a score distribution (e.g., a probability7distribution) that assigns a respective score, e.g., a respective probability, to each token in the vocabulary of tokens. The language model neural network can then select, as the particular token, a token from the vocabulary using the score distribution. For example, the neural network can greedily select the highest-scoring token or can sample, e.g., using nucleus sampling or another sampling technique, a token from the distribution.

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

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

[0067] Osindero, L. Rimell, C. Dyer, O. Vinyals, K. Ayoub, J. Stanway, L. Bennett, D. Hassabis, K. Kavukcuoglu. and G. Irving. Scaling language models: Methods, analysis & insights from training gopher. CoRR. abs / 2112. 11446. 2021; Cohn Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning wi th a unified text-to-text transformer. arXiv preprint arXiv: 1910. 10683, 2019; Daniel Adiwardana, Minh-Thang Luong, David R. So, Jamie Hall, Noah Fiedel, Romal Thoppilan, Zi Yang, Apoorv Kulshreshtha. Gaurav Nemade, Yifeng Lu, and Quoc V. Le. Towards a human-like open-domain chatbot. CoRR, abs / 2001.09977, 2020; and Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda AskelL et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020.

[0068] More specifically, the neural network includes plurality of layers that include a plurality of attention layers.

[0069] Each attention layer receives a respective hidden state for each of the input positions and updates the respective hidden states for each of the input positions by applying an attention mechanism to the respective hidden states.

[0070] Generally, the task of the network input can be any task that requires generating an output sequence that includes a respective output token at each of multiple output positions. Examples of such tasks include computer code generation or editing tasks, text generation or editing tasks, image, video, or audio understanding tasks, and so on.

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

[0072] In any of the implementations below, the neural network may be deployed as part of a chat bot, dialogue agent, or other software tool that receives inputs from users and provides outputs in response to the received input, e.g., as part of a conversation or dialogue. In these implementations, the input sequences received by the neural network are (generated from) user inputs and the output sequences generated by the neural network can be used to generate responses to the user inputs. In implementations the neural network may be configured as, or include, a generative (large) language model or a multi-modal model, e.g., a visual and language model, to perform these example machine learning tasks.

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

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

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

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

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

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

[0079] As another example, the task can be a health prediction task, where the input is a sequence derived from electronic health record data for a patient and the output is a prediction that is relevant to the future health of the patient, e.g., a predicted treatment that should be prescribed to the patient, the likelihood that an adverse health event will occur to the patient, or a predicted diagnosis for the patient. Such electronic health data may, for example, comprise one or more sequences of physiological data taken from a patient, with the output being a corresponding prediction that relates to those sequences of data. Examples of physiological data and a corresponding prediction include: blood glucose measurements, with the prediction being a predicted future blood glucose measurement or the prediction of a hyper- or hypo-glycemic event; a heart rate, with the prediction being the presence or absence of a heart condition, or a future cardiac event; blood pressure measurements, with the prediction being the risk of a future heart condition; or the like.

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

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

[0082] As another example, the task can be an agent control task, where the input is a sequence of observations or other data characterizing states of an environment and the output defines an action to be performed by the agent in response to the most recent data in the sequence. The agent can be, e.g., a real-world or simulated robot, a control system for an industrial facility', or a control system that controls a different kind of agent. The observations may comprise sensor data captured by sensors associated with (e.g. part ol the agent, for example visual data. LIDAR data, sonar data, agent configuration data (e.g. joint angles), agent orientation data, or the like.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

[0097] In some cases, the teacher language model neural network differs from the student language model neural network. For example, as described above, in some cases the teacher language model neural network is larger than the student language model neural network. For example, the teacher language model neural network can be larger relative to the student language model neural network in that it can have a larger architecture (i.e., can have more parameters, e.g., more neural network layers, more neurons per layer, more attention blocks, more attention heads per layer, and so on), can require using more computational resources (e.g., has a larger memory footprint or requires using many hardware accelerators), or both.

[0098] FIG. 2 is a flow diagram of an example process 200 for training a student language model neural network to be deployed in a cascade with a teacher language model neural network. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a cascade training system, e.g., the cascade training system 100 of FIG. 1 , appropriately programmed in accordance with this specification, can perform the process 200.

[0099] The student language model neural network and teacher language model neural network can each have any of a variety of neural network architectures. That is, the student and teacher neural networks can have any appropriate architecture in any appropriate configuration that can process a network input and generate an output sequence, including fully connected layers, convolutional layers, recurrent layers, attention-based layers, and so on, as is appropriate. For example, as described above, the student and teacher language model neural networks can be auto-regressive neural network models. As particular examples, the student and teacher language model neural networks can belong to the Gemini, Gemma, or PaLM 2 family of neural networks.

[0100] In some cases, the student language model neural network and teacher language model neural network are configured to be remote relative to each other (i.e., the system requires communication over a network, e.g., cloud-based network, the internet, a local area network, wireless network, etc.) to facilitate communication between the student and teacher.

[0101] In some cases, the teacher language model neural network is an already-trained neural network that is held fixed during the training of the student language model neural network.

[0102] For example, the teacher language model neural network can be a pre-trained language model neural network (e.g., pre-trained on a next-token prediction task or masked-token prediction task for a variety of token sequences associated with a variety7of domains, e.g., sequences representing bodies of texts found in websites, articles, books, essays, etc., or sequences representing question and answering, multilanguage question and answering, language translation, reasoning, mathematical reasoning, computer programs, and so on). As described above, examples of teacher language model neural networks can include those belonging to the Gemini, Gemma, or PaLM 2 families of neural networks.

[0103] In some cases, the student language model neural network has been pre-trained prior to the training of the student language model neural network to be deployed in the cascade with the teacher language model neural network.

[0104] For example, the student language model neural network can be pre-trained as described above for the teacher language model neural network, and specific examples of the student language model neural network also include neural networks belonging to the Gemini, Gemma, or PaLM 2 family of neural networks.

[0105] In some implementations, the teacher language model neural network has more parameters than the student language model neural network.

[0106] For example, the teacher neural network can be Gemini 1.0 Ultra with trillions of parameters and the student neural network can be Gemini 1 .0 Nano-2 or Pro with billions or hundreds of billion parameters respectively, all of which are available from Google and are described in arXiv:2312. 11805.

[0107] As another example, the teacher neural network can be Gemma-7B with 7 billion parameters and the student neural network can be Gemma-2B with 2 billion parameters respectively as are described in arXiv: 2403.08295.

[0108] As another example, the teacher and student language model neural networks can both belong to the PALM-2 family as described in arXiv:2305. 10403, where the teacher neural network can be PALM-2 Otter and the student neural network can be PALM-2 Gecko.

[0109] The system receives a set of one or more training examples, where each training example includes (i) a network input and (ii) a target output for the network input (step 202). The target output includes a respective target token at each of a plurality of output positions, each target token is selected from a vocabulary of tokens.

[0110] As described above, the network input can include any type of input data and can be mapped to a sequence of tokens.

[0111] For example, the system can map a network input x to a sequence of M tokens as x ■■= (xlt... , xM), where each token xtcomes from a vocabulary V that includes V tokens. Similarly, the target output y includes a target token at each of plurality of output positions (e.g., A target tokens, where N>0) such that the target output can be defined as y:= where the target output tokens can also come from the same vocabulary

[0112] V.

[0113] For each training example the system receives, the system performs steps 204-208 below.

[0114] The system processes the network input in the training example using the student language model neural network to generate, for each output position, a respective student score distribution that assigns a respective student score to each token in the vocabulary of tokens (step 204).

[0115] For example, to process the network input using the student language model neural network to generate student score distributions for each output position, the system can process the sequence of tokens representing the network input x using a student language model neural network p using '‘teacher forcing”. That is, after training, the student language model neural network can process a network input in an auto-regressive manner to generate a student score for each output position sequentially. But, during training (e.g., during example process 200), the student language model neural network can process the network input in a single forward pass to generate a student score distribution for each output position all at once. For example, the system can use the student language model to compute the probability of an output sequence y given the input sequence x as p(y ) represents the portion of the output sequence before position i, y<T= 0 represents no output sequence, and p(-| x, y < i) represents a per-token probability7mass function, i.e., a student score distribution that assigns a respective student score to each token in the vocabulary of tokens for output position i. From this point forward. pswill denote the student language model neural network p, as described above. Teacher forcing refers to processing the network input and the respective target output to generate student score distributions for each output position all at once in one forward pass using the appropriate values of x and y<tof the target output for each output position yt. While processing the network input in an auto-regressive manner refers to generating each student score distribution for output ytsequentially, conditioned on previously generated y<;and x.

[0116] The system processes the network input in the training example using the teacher language model neural network to generate, for each output position in the target output in the training example, a respective teacher score distribution that assigns a respective teacher score to each token in the vocabulary of tokens (step 206).

[0117] For example, to process the network input using the teacher language model neural network to generate teacher score distributions for each target output position, the system can process the sequence of tokens representing the network input x using a teacher language model neural network p using teacher forcing that computes the teacher score distributions of each output position ytof an output sequence y given appropriate values of x and y<;in a single forward pass (and, after training, computes the output sequence of y in an auto-regressive manner)as was described for the student language model neural network psabove. From this point forward, pTwill denote the teacher language model neural network p, as described above. So the term pT(-| x,y < i) represent represents a per-token probability mass function, i.e., a teacher score distribution that assigns a respective teacher score to each token in the vocabulary of tokens for output position i.

[0118] The system determines, for each of the output positions and from the respective teacher score distribution and the respective student score distribution for the output position, a respective weight for the output position that represents a difficulty of accurately predicting the target token at the output position (step 208).

[0119] For example, the term atcan represent the respective weight for the output position i that represents a difficulty of accurately predicting the target token at the output position which the system determines using the respective teacher score distribution and student score distribution for the output position i. In some cases, when the system performs step 208, the system determines whether the target token at the output position has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position. Then, the system determines the respective weight for the output position based on whether the target token at the output position (i) has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position or (ii) has been assigned a highest teacher score of any vocabulary token in the respective teacher score distribution for the output position.

[0120] For example, the equation can represent the system determining whether the target token at the output position i has been assigned a highest score of any vocabulary token in the respective score distribution for the output position for language model neural network p, where the ytterm is the target token for the output position i, the term arg max p(y'\x, y<t) returns the token y'ev within the vocabulary' V with the highest probability mass according to the language model neural network p. yt= arg returns the Boolean value of True if the target token ythas been assigned a highest score of any vocabulary token, and other terms and notation of this equation are defined as described above. So the equations yt= arg max m(y' |x, y<() and yi = arg max pT(y' lx, y<i) are the student language model y'EV y'EV neural network and teacher language model neural network versions respectively of the above.

[0121] Further in some cases, to determine the respective weight for the output position, the system can assign a first weight (value) or a second weight (value) to the output position. For example, the system assigns a first weight to the output position when (i) the target token at the output position has been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position; (ii) the target token at the output position has been assigned the highest teacher score of any vocabulary token in the respective teacher score distribution for the output position; or (iii) the target token at the output position has been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position and has been assigned the highest teacher score of any vocabulary’ token in the respective teacher score distribution for the output position. In other words, the first weight can be assigned when the target token is correctly predicted by at least one of the student or teacher language model neural networks. The system can assign a second weight to the output position when the target token at the output position (i) has not been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position and (ii) has not been assigned the highest teacher score of any vocabulary token in the respective teacher score distribution for the output position. In other words, the second weight can be assigned when neither the student or teacher language model neural networks correctly predicts the target token.

[0122] Generally, the first and second weight values can be any scalar value. For example, the first weight can be one and the second weight can be zero.

[0123] As an example, the equation can represent how the system can determine the respective weight for the output position i by assigning a first weight of 1 and a second weight 0 to the output position as described above, where 1 denotes an indicator function that returns 1 when the argument is a Boolean value of True and 0 otherwise, the term V of the argument of the indicator function denotes a logical OR operation, and other terms and notation of this equation are defined as described above.

[0124] After the system performs steps 204-208, the system trains the student neural network on a loss function (step 210). In particular, the system trains the student neural network on a loss function that depends on, for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position and (ii) the student score assigned to the target token at the output position by the student score distribution for the output position.

[0125] In some cases, the loss function measures, for each training example, a sum of. for each output position, a product of the respective weight for the output position and a negative log likelihood computed from the student score assigned to the target token at the output position by the student score distribution for the output position.

[0126] For example, the equation can represent how the system computes the loss function (denoted as Lcat-xentx, y)) for a single training example with a network input x and target output y with N tokens as described above, where other terms and notation of this equation are defined as described above. In other words, the loss function can be a measure of the prediction accuracy of the student language model neural network weighted by how difficult each target token is considered to be. As discussed above, target tokens considered too difficult to leam for the student language model neural network can be discounted from the loss function.

[0127] In some cases, the loss function is a distillation loss that depends on, for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position, (ii) the student score assigned to the target token at the output position by the student score distribution for the output position, (iii) the student score distribution for the output position, and (iv) the teacher score distribution for the output position.

[0128] Further in some cases, the loss function measures, for each training example, a sum of, for each output position, a product of the respective weight for the output position and a token-level distillation loss. The token-level distillation loss is computed based on a negative log likelihood computed from the student score assigned to the target token at the output position by the student score distribution for the output position and a difference measure between the student score distribution for the output position, and the teacher score distribution for the output position.

[0129] In some implementations, the token-level distillation loss is a weighted sum of the negative log likelihood and the difference measure.

[0130] As an example of the above described possible loss functions, the equation can represent how the system computes the loss function (denoted as Lcat-distx, y')') for a single training example with a network input x and target output y with N tokens as described above. In particular, w denotes a distillation proportion, and

[0131] PrCy' l*, y< Zo,g Ps(y' \x>y<d denotes an example difference measure between the student score distribution and teacher score distribution for the output position, and other terms and notation of this equation are defined as described above. In other words, the loss function can include a combination of the difficulty weighted prediction accuracy measure for the student language model neural network discussed above and a token-level distillation term. The distillation term attempts to ‘‘distill’' the knowledge of the teacher language model neural network into the smaller student language model neural network such that the student language model neural network can provide the same output as the teacher language model neural network for as many inputs as possible. This reduces the number of times that the more computationally expensive teacher language model neural network needs to be invoked for processing inputs when deployed.

[0132] Because the loss function includes respective weights for the output positions that represent a difficulty of accurately predicting the target token at the output position, the weights for the output positions guide the training of the student neural network by focusing only on tokens that fall within the scope of the student teacher cascade and therefore improve the performance of the student teacher cascade. For example, when the weights for the output positions are 1 only if either the student language model neural networks or the teacher language model neural network can predict the token with high confidence, then the training process using the loss function enables the student language model to optimize its capacity by focusing on learnable tokens, which improves the student language model’s overall performance. In other words, the weights for the output positions described above enable the training using a loss function that includes the weights to concentrate the parameter updates for the student language model neural network on learnable tokens, thereby enhancing the student model’s confidence and accuracy and improving the performance of the student teacher cascade.

[0133] The system training the student language model neural network on a loss function generally includes repeatedly updating the trainable parameters of the student language model neural network using a training dataset. Further details on updating trainable parameters of a student language model neural network on a loss function are described below with reference to FIG. 4.

[0134] FIG. 3 is a flow diagram of an example process 300 for generating an output for a task using a student language model neural network deployed in a cascade with a teacher language model 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 cascade training system, e.g., the cascade training system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 300.

[0135] The system receives a network input for performing a task using the student language model neural network (step 302). As described above, generally, the network input can include any type of input data.

[0136] For example, for the task of intent recognition of an utterance, the network input could be the natural language text, “place this song into my newest playlist.” Alternatively, the input may be spoken by a user and captured by a microphone for example. In some cases, the spoken utterance can be transcribed to provide a text input if the model is not configured to process audio input.

[0137] As another example, for the task of language translation, the network input could be natural language text in a first language to be translated to a second language, e.g., the network input could be “Where is the library?” with the task of translating this English sentence to a Spanish sentence.

[0138] As another example, for the task of image captioning, the network input could include image data (e.g., a sequence of image pixels that represent a cat napping) to be labeled using natural language text. That is, the network output can comprise text describing the objects depicted in the image data.

[0139] As another example, for the task of summarizing speech, the network input could include audio data (e.g., audio spectrogram of a classroom lecture or presentation) to be summarized into natural language text. The network output can comprise text that identifies or describes the speakers present in the audio data. The network output can provide a transcription of some of the audio data. Alternatively, a full transcription of the audio data can be provided in a full speech / audio transcription task.

[0140] As another example, for the task of video summarization, the network input could include video data (e.g., a sequence of video frames of a cat climbing a tree) to be summarized as a concise body of natural language text. The network output can comprise text describing the objects depicted in the video and the actions being carried out in the video.

[0141] As another example, for the task of image object detection, the network input could include an image and, optionally, text (e.g., an image and a natural language text description of an object that should be identified in the image). The network output can be a location of the image at which one or more objects have been detected.

[0142] In some cases, the student language model neural network is deployed on a user device and the teacher language model neural network is deployed on one or more remote computers that are remote from the user device. For example, the student language model neural network can be deployed on a user smart phone, tablet or laptop. The teacher language model neural network can be deployed on one or more remote central server computers, cloud server computers or workstation computers for example.

[0143] In some cases, the system receives the network input from a user of the user device.

[0144] For example, a user of a smart phone, tablet, or laptop, can provide the network input to the end user device through the end user device interface (i.e., the touch screen for the smart phone and tablet and the keyboard for the laptop).

[0145] The system processes the network input using the student language model neural network to generate an output sequence for the network input (step 304).

[0146] For example, as described above, the system can process the network input using a student language model neural network that is an auto-regressive neural network that generates an output sequence of tokens by generating each token in the output sequence conditioned on a current input sequence that includes any tokens that precede the particular token in the output sequence.

[0147] For this example auto-regressive neural network, in some cases, the tokens for the output sequence are selected in a “greedy manner'’, i.e., the neural network can greedily select the highest-scoring token. For example, the neural network selects tokens for the output sequence according to their corresponding probabilities as described above. But, in some cases, the probabilities can be modified using a temperature parameter T. For example, the temperature T can modify the probability of selecting token vkas p' (k) = eP Vk) / T

[0148] - —P[r.) / Twhere p (vfe) represents the temperature modified probability of selecting token vk, p(vk)represent the original probability of selecting vk, the index i runs over all eligible tokens for selection and the variable T is the temperature parameter that can be set. The higher the value of T, the more equal the modified probabilities for the tokens become among each other. While the lower the value of T, the more polarizing the modified probabilities for the tokens become relative to the original probabilities, with higher original probabilities becoming higher modified probabilities and lower original probabilities becoming lower modified probabilities. Therefore, the various values of T in the context of temperature decoding control the probabilistic variability of sampled output sequences, with a value of r=1.0 not modifying the original token selection probabilities, lower values of T (e.g., 0.1, 0.2, 0.5, and so on) resulting in sampled output sequences that more often closely align with a ‘highest probability selection procedure’ (i.e.. the system selects each token of the output sequence according to the highest probability over the tokens of the vocabulary) and higher values of T (e.g., 1.1, 1.2, 1.5, 2.0, and so on) resulting in output sequences that more often closely align with a ‘random selection procedure’ (i.e., the system selects each token of the output sequence randomly from among the tokens of the vocabulary).

[0149] The system determines whether to defer the network input to the teacher neural network (step 306) (i.e. determines whether the teacher neural network should process the network input).

[0150] The system can determine to defer the network input to the teacher neural network using any of a variety of methods.

[0151] In some cases, the system determines to defer the network input to the teacher language model neural network based on the network input content. As such, generating the output sequence for the network input using the student language model neural network (step 304) can be optional and only occurs if the system determines not to defer the network input to the teacher.

[0152] For example, if the system detects that the network input contains specific keywords (e.g., “solve'’, “math”, “reason”, etc.), the system, in response, can determine to defer the network input. As another example, if the system determines the size of the network input exceeds a set threshold size (e.g.. the network input corresponds to natural language text with a word count above a threshold), the system, in response, can determine to defer the network input.

[0153] As another example, the system can determine to defer the network input to the teacher language model neural network using an auxiliary machine learning model. For example, the system can use a classification machine learning model to process the network input and generate an output that indicates whether to defer the network input to the teacher language model neural network.

[0154] In some cases, the system determines to defer the network input to the teacher language model neural network using the student language model neural network. That is, the system can use the student language model neural network psto process a network input x. and then the system can use the output sequence of psin order to determine to defer the network input or not. For example, the system can measure a deferral score r based on the output sequence of psin order to determine whether or not to defer the network input. So. for example, the system can determine to defer the network input if the deferral score exceeds a pre-determined threshold value T.

[0155] In some cases, the deferral score r can be based on the "confidence" of the output sequence of ps. i.e., based on the scores / probabilities / uncertainties of the tokens that make up the output sequence.

[0156] For example, the system can use the normalized log-likelihood of the output sequence of psto determine the deferral score r. For example, for an output sequence generated by the student language model neural network denoted as ys~ ps(-| x) where the output sequence ysincludes N tokens, the deferral score can be determined as where log ps(ys,i ys,<t) denotes the log-likelihood of the output token at position I for output sequence ysand other terms and notation of this equation are defined as described above.

[0157] If the system determines not to defer the network input to the teacher neural network, then the system uses the output sequence generated by the student as the output for the task (step 308). If the system executes step 308, the example process 300 is completed.

[0158] But if the system determines to defer the network input to the teacher neural network, then the system provides the network input as an input to the teacher language model neural network (step 310).

[0159] The system processes the network input using the teacher language model neural network to generate an output sequence for the network input (step 312).

[0160] For example, as described above for the student language model neural network, the system can process the network input using the teacher language model neural network that is an auto-regressive neural network.

[0161] The system uses the teacher output sequence as output for the task (step 314). If the system executes step 314, the example process 300 is completed.

[0162] The equation

[0163] Pcas(y\x): = l[r(x) < T] • Ps y\x) + l[r(x) > r] • pT(y | x) represents an example of how the system can perform example process 300, where pras(y|x) denotes the system generating an output y by processing a network input x using a cascade that includes a student language model neural network and a teacher language model neural network, r(x) denotes determining a deferral score (e.g., as described above), the sum l[r(%) < T] • ps(y |x) + l[r(%) > T] • pT(y\x) denotes the output of the student language model neural network when the deferral score does not meet or exceed the threshold value T and the output of the teacher language model neural network otherwise, and other terms and notation of this equation are defined as described above.

[0164] In some implementations, the system can cause the student language model neural network to be deployed on a user device (or other appropriate device). For example, the student language model neural network can be transmitted to the user device over an appropriate network connection. Likewise, the system can cause the teacher language model neural network to be deployed on an appropriate data processing apparatus. In some implementations, the student language model neural network can be trained on a user device and the teacher language model neural network is accessed remotely by the user device.

[0165] FIG. 4 is a flow diagram of an example process 400 for updating trainable parameters of a student neural network. For convenience, the process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a cascade training system, e.g., the cascade training system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 400.

[0166] The system obtains a training data set that includes training examples (step 402), which the system obtains from a user or another system.

[0167] In some cases, the training data set is a “mini -batch’7(i.e.. a subset of any size of all available training examples). In other cases, the training data set is a “batch” (i.e., a set of all available training examples).

[0168] The training data set includes training examples, and each training example includes a network input and corresponding target output. For example, the system obtains training examples as described above with reference to FIG. 1 and FIG. 2.

[0169] The system, for each training example, generates an output (step 404), i.e., the system generates student score distributions, teacher score distributions, weights, and output sequences using both the student language model neural network and teacher language model neural network. For example, the system generates these outputs for each training example as described above with reference to FIG. 1 and FIG. 2.

[0170] The system evaluates an objective using all training examples and respective target outputs (step 406).

[0171] In particular, the objective generally includes a loss function for each training example. For example, as described above with reference to FIG. 1 and FIG. 2. a loss function is one that depends on, for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position and (ii) the student score assigned to the target token at the output position by the student score distribution for the output position.

[0172] As an example, for a single training example that includes a network input x and target output y the loss function for the training example can be Lcatdist(x> y) (as described above and reproduced below') w here the terms and notation of this equation are defined as described above.

[0173] As an example objective, the objective can be to minimize the average of the loss (e.g., the loss Lcat-dist(x, y)) over all training examples included in the training data set (a set of x and y pairs).

[0174] In some cases, the objective includes one or more regularization terms that penalizes higher values for the trainable parameters to reduce the risk of the trainable parameters of the student language neural netw ork overfitting the training examples. For example, the regularization terms can include the LP regularization term A||iv||p , where / . is the regularization parameter, w is the vector of trainable parameters, and p is the norm degree (e.g., p = I for LI regularization, p = 2 for L2 regularization).

[0175] The system updates trainable parameters to optimize the objective (step 408). That is, the system modifies the parameter values of the trainable parameters included in the student language model neural network such that the evaluation of the objective using the student language model neural network improves. The system can update the trainable parameters of the student language model neural network to optimize the objective in any variety of ways, e.g., gradient based method, evolutionary algorithm-based method, Bayesian optimization, etc.

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

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

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

[0179] In particular, example 500 shows two plots that summarize performance in terms of ‘"deferral curves” over two data sets (i.e., the superGLUE data set and the FLAN2021- Cls-Tasks data set) for the described techniques (i.e., the curves labeled CAT-Xent, and CAT-Dist, where CAT denotes ‘"cascade aware training”) and baselines or other techniques (i.e., curves labeled Xent and Dist).

[0180] A deferral curve is a plot of a task dependent performance metric vs FLOPs (floating point operations) per token for a student teacher cascade system such as Accuracy-vs. -FLOPs or BLEU score-vs. -FLOPs. In order to generate the deferral curve, one sw eeps different thresholds r for deferral score r(x). The sw eeping of the threshold causes the cascade to utilize outputs of the student and teacher at vary ing proportions, and, because the student and teacher language model neural networks have a lower and higher FLOP cost per token respectively, the sw eeping results in a curve of performance metric as a function of FLOP per token values. As such, it is possible to compare the overall cost-performance trade-off of different techniques. A particular threshold value provides an operating point along the cost-performance curve.

[0181] The task for the plots of example 500 is classification so the deferral curve is Accuracy-vs. -FLOPs, and the deferral curve on the left is for classification tasks of the superGLUE data set and the deferral curve on the right is for classification tasks of the FLAN2021 data set.

[0182] The described techniques denoted as the CAT-Xent and CAT-Dist curves refer to the described techniques of training a student language model neural network that belongs to a student teacher cascade above when the loss function are Lcat-xentand Lcat-dist respectively.

[0183] The baseline techniques denoted as Xent and Dist train a student language model neural network that belongs to a student teacher cascade using loss functions that do not take into account the weight for each output position of the target output as does the described techniques above. Instead, the baselines train the student language model neural network using standard one-hot cross entropy (for the case of Xent) and KL-divergence based distillation loss (for the case of Dist).

[0184] The plots of example 500 show that the described techniques benefits the cascade on the SuperGLUE dataset, especially when the student language model neural network is trained with CAT-Xent. Compared with the training with plain one-hot cross entropy (Xent), CAT-Xent reduces 13% FLOPs (teacher model calls) given fixed 87% accuracy.

[0185] With fixed 2 billion total FLOPs budget (~20% teacher model calls), CAT-Xent gets 2% accuracy improvement. Meanwhile, CAT does not clearly downgrade the student model’s original capability (left end of the curves). When applying CAT with distillation, the cascade benefits the majority of the time (CAT-Dist vs. Dist). Similar trends can be seen on FLAN2021 datasets. Both superGLUE and FLAN2021 dataset, distillation improves the student language model’s uality (Dist vs. Xent) compared to with one-hot cross entropy, CAT-Xent levels up the cascade performance and out-performs the cascade with student language model fine-tuned with distillation, especially in the low-FLOPS range.

[0186] Example 500 shows that on average the overall task performance of the cascade per unit of computational cost is higher when using the described techniques than when not using the described techniques.

[0187] FIG. 6 is an example 600 of the performance of the described techniques. FIG. 6 follows the same type of metric (i.e., deferral curves that, in this case, use the BLEU score metric, which is a metric that measure how well a generated text matches a reference) for the described techniques (i.e., CAT-Xent and CAT-Dist) and baseline techniques (i.e., Xent and Dist) as FIG. 5, but shows results for a different machine learning task (i.e., “generation tasks”, i.e., text generation tasks such as language translation, summarization of text, reasoning, question answering, etc.) on different data sets (i.e., the WMT22 and FLAN2021 -Gen-Tasks data sets).

[0188] Generation tasks commonly contain dozens to hundreds of tokens in the target output and is therefore a more challenging task than the classification tasks as was shown in FIG. 5 above. Yet. FIG 6 still shows benefits from cascade aware training for the cascade (i.e., the described techniques). One obvious gain from CAT is that the student language model can enjoy better intrinsic quality with cross-entropy loss (as illustrated by the higher BLEU scores of CAT-Xent versus Xent on the left side of the plot) on WMT22 dataset. This improvement even mitigates the benefits gained from distillation (CAT-Xent versus Dist). This is further confirmed on the FLAN2021 dataset. One reason this may be is that the improvement partially comes from the lower ratio of noisy training samples after token-wise weighting of the described techniques. In addition, FIG.5 illustrates benefits from CAT for the cascade (i.e., the described techniques) on the generation tasks (e.g. CAT-Xent versus Xent on FLAN2021 or WMT22 in the low-to-intermediate deferral regime).

[0189] Example 600 also shows that on average the overall task performance of the cascade per unit of computational cost is higher when using the described techniques than when not using the described techniques.

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

[0191] In particular, example 700 shows two plots of a deferral weighted performance metric for entities of a cascade vs deferral rate.

[0192] The curves present in both plots are the following: the curve labeled “Al CAT” refers to a the student language model neural network of a cascade trained using the described techniques with a loss Lcat-distas described above; the curve labeled “A2 CAT” refers to the teacher language model neural network counterpart of “Al CAT” that make up a cascade; the curve labeled “Al XEnt” refers a student language model neural network of a cascade trained using a standard KL-divergence based distillation loss; and the curve labeled “A2 XEnt” refer to the teacher language model neural network counterpart to “Al XEnt” included in the cascade.

[0193] The y-axis of these plots refers to a “deferral weighted performance metric” for a language model neural network. That is, every point on the graph is an average performance metric for a particular language model (i.e., student or teacher model trained using the described technique or not) on received network inputs belonging to specific dataset multiplied by the deferral rate (if the model is a teacher language model neural network) or by “1 -deferral rate” (if the model is a student language model neural network). For example, at a deferral rate of 0.2, the y-axis value of Al CAT is the average performance of the student neural network trained using the describe techniques multiplied by 1-0.2 (or 0.8); the y-axis value of Al XEnt is the average performance of the student neural network trained using the standard techniques multiplied by 1-0.2 (or 0.8); the y-axis value of A2 CAT is the average performance of the teacher neural network that is a part of a cascade with the student neural network Al CAT multiplied by 0.2; and the y-axis value of A2 XEnt is the average performance of the teacher neural network that is a part of a cascade with the student neural network Al XEnt multiplied by 0.2;

[0194] For the plot on the left, the performance metric of the y-axis is accuracy for the task of multi-rc (Multi -Sentence Reading Comprehension), where multi-rc is a question answering task involving a context paragraph, a related question, and multiple potential answers the answers must be identified as true or false.

[0195] For the plot on the right, the performance metric of the y-axis is EM (exact match rate) for the task of “record” (Reading Comprehension with Commonsense Reasoning, i.e., a multiple-choice QA task involving a news article and a question with a masked entity, where the masked entity from provided options must be predicted).

[0196] Example 700 demonstrates two separate phenomena leading to improved performance of CAT-trained models (i.e., student model trained using the described techniques). On superGLUE multi-rc, Al -CAT starts above Al-Xent when deferral rate is 0; this indicates that Al -CAT is simply a better model than Al-Xent - all the queries are routed to the student model. Therefore Al-CAT must have learned to correctly process some examples which Al-Xent cannot. In the case of superGLUE record, on the other hand, the two models start effectively overlapping on the left-hand side of the figure; it is only in the body of the deferral rate graph that Al-CAT achieves separation from Al-Xent.

[0197] Example 700 shows that the described techniques yield a student language model neural network of a cascade system that has equal or greater performance compared to a student language model in cascade not trained using the described techniques.

[0198] This specification uses the term “configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions. Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

[0199] The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry7, 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.

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

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

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

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

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

[0205] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the w eb 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.

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

[0207] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework.

[0208] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application sen- er, or that includes a front end component, e.g., a client computer having a graphical user interface, a w eb brow ser, or an app through w hich a user can interact w ith an implementation of the subj ect matter described in this specification, or any combination of one or more such back end, middleware, or front end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0209] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at the server from the device.

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

[0211] 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 softw are product or packaged into multiple software products.

[0212] This specification also provides the subject-matter of the following clauses: Clause 1. A method performed by one or more computers and for training a student language model neural network to be deployed in a cascade with a teacher language model neural network, the method comprising: receiving a set of one or more training examples, each training example comprising (i) a network input and (ii) a target output for the network input, the target output comprising a respective target token at each of a plurality of output positions, each target token being selected from a vocabulary of tokens; for each training example: processing the network input in the training example using the student language model neural network to generate, for each output position, a respective student score distribution that assigns a respective student score to each token in the vocabulary of tokens; processing the network input in the training example using the teacher language model neural network to generate, for each output position in the target output in the training example, a respective teacher score distribution that assigns a respective teacher score to each token in the vocabulary of tokens; and determining, for each of the output positions and from the respective teacher score distribution and the respective student score distribution for the output position, a respective weight for the output position that represents a difficulty of accurately predicting the target token at the output position; training the student neural network on a loss function that depends on, for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position and (ii) the student score assigned to the target token at the output position by the student score distribution for the output position.

[0213] Clause 2. The method of clause 1 , wherein the teacher language model neural network is an already -trained neural network that is held fixed during the training of the student language model neural network.

[0214] Clause 3. The method of clause 1 or 2, wherein the student language model neural network has been pre-trained prior to the training of the student language model neural network to be deployed in the cascade with the teacher language model neural network.

[0215] Clause 4. The method of any preceding clause, wherein determining, for each of the output positions and from the respective teacher score distribution and the respective student score distribution for the output position, a respective weight for the output position that represents a difficulty of accurately predicting the target token at the output position comprises: determining whether the target token at the output position has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position; determining whether the target token at the output position has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position; and determining the respective weight for the output position based on whether the target token at the output position (i) has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position or (ii) has been assigned a highest teacher score of any vocabulary' token in the respective teacher score distribution for the output position.

[0216] Clause 5. The method of clause 4, wherein determining the respective weight based on whether the target token at the output position (i) has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position or (ii) has been assigned a highest teacher score of any vocabulary' token in the respective teacher score distribution for the output position comprises: assigning a first weight to the output position when the target token at the output position (i) has been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position (ii) has been assigned the highest teacher score of any vocabulary token in the respective teacher score distribution for the output position, or (iii) has been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position and has been assigned the highest teacher score of any vocabulary' token in the respective teacher score distribution for the output position; and assigning a second weight to the output position when the target token at the output position (i) has not been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position and (ii) has not been assigned the highest teacher score of any vocabulary' token in the respective teacher score distribution for the output position.

[0217] Clause 6. The method of clause 5, wherein the first weight is one and the second weight is zero. Clause 7. The method of any preceding clause, wherein the loss function measures, for each training example, a sum of, for each output position, a product of the respective weight for the output position and a negative log likelihood computed from the student score assigned to the target token at the output position by the student score distribution for the output position.

[0218] Clause 8. The method of any one of clauses 1-6, wherein the loss function is a distillation loss that depends on, for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position, (ii) the student score assigned to the target token at the output position by the student score distribution for the output position, (iii) the student score distribution for the output position, and (iv) the teacher score distribution for the output position.

[0219] Clause 9. The method of clause 8, wherein the loss function measures, for each training example, a sum of, for each output position, a product of the respective weight for the output position and a token-level distillation loss that is computed based on: a negative log likelihood computed from the student score assigned to the target token at the output position by the student score distribution for the output position; and a difference measure between the student score distribution for the output position, and the teacher score distribution for the output position.

[0220] Clause 10. The method of clause 9, wherein the token-level distillation loss is a weighted sum of the negative log likelihood and the difference measure.

[0221] Clause 11. The method of any preceding clause, wherein the teacher language model neural network has more parameters than the student language model neural network.

[0222] Clause 12. A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one more computers to perform the operations of the respective method of any one of clauses 1-11.

[0223] Clause 13. One or more computer storage media storing instructions that when executed by one or more computers cause the one more computers to perform the operations of the respective method of any one of clauses 1-11.

[0224] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous. What is claimed is:

Claims

CLAIMS1. A method performed by one or more computers and for training a student language model neural network to be deployed in a cascade with a teacher language model neural network, the method comprising: receiving a set of one or more training examples, each training example comprising (i) a network input and (ii) a target output for the network input, the target output comprising a respective target token at each of a plurality' of output positions, each target token being selected from a vocabulary of tokens; for each training example: processing the network input in the training example using the student language model neural network to generate, for each output position, a respective student score distribution that assigns a respective student score to each token in the vocabulary of tokens; processing the network input in the training example using the teacher language model neural network to generate, for each output position in the target output in the training example, a respective teacher score distribution that assigns a respective teacher score to each token in the vocabulary of tokens; and determining, for each of the output positions and from the respective teacher score distribution and the respective student score distribution for the output position, a respective weight for the output position that represents a difficulty of accurately predicting the target token at the output position; training the student neural network on a loss function that depends on, for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position and (ii) the student score assigned to the target token at the output position by the student score distribution for the output position.

2. The method of claim 1, wherein the teacher language model neural network is an already -trained neural network that is held fixed during the training of the student language model neural network.

3. The method of claim 1 or 2, wherein the student language model neural network has been pre-trained prior to the training of the student language model neural network tobe deployed in the cascade with the teacher language model neural network.

4. The method of any preceding claim, wherein determining, for each of the output positions and from the respective teacher score distribution and the respective student score distribution for the output position, a respective weight for the output position that represents a difficulty’ of accurately predicting the target token at the output position comprises: determining whether the target token at the output position has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position; and determining the respective weight for the output position based on whether the target token at the output position (i) has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position or (ii) has been assigned a highest teacher score of any vocabulary token in the respective teacher score distribution for the output position.

5. The method of claim 4, wherein determining the respective weight based on whether the target token at the output position (i) has been assigned a highest student score of any vocabulary token in the respective student score distribution for the output position or (ii) has been assigned a highest teacher score of any vocabulary token in the respective teacher score distribution for the output position comprises: assigning a first weight to the output position when the target token at the output position (i) has been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position, (ii) has been assigned the highest teacher score of any vocabulary token in the respective teacher score distribution for the output position, or (iii) has been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position and has been assigned the highest teacher score of any vocabulary token in the respective teacher score distribution for the output position; and assigning a second weight to the output position when the target token at the output position (i) has not been assigned the highest student score of any vocabulary token in the respective student score distribution for the output position and (ii) has not been assigned the highest teacher score of any vocabulary token in the respective teacherscore distribution for the output position.

6. The method of claim 5. wherein the first weight is one and the second weight is zero.

7. The method of any preceding claim, wherein the loss function measures, for each training example, a sum of, for each output position, a product of the respective weight for the output position and a negative log likelihood computed from the student score assigned to the target token at the output position by the student score distribution for the output position.

8. The method of any one of claims 1-6, wherein the loss function is a distillation loss that depends on, for each training example and for each output position in the target output in the training example, (i) the respective weight for the output position, (ii) the student score assigned to the target token at the output position by the student score distribution for the output position, (iii) the student score distribution for the output position, and (iv) the teacher score distribution for the output position.

9. The method of claim 8, wherein the loss function measures, for each training example, a sum of. for each output position, a product of the respective weight for the output position and a token-level distillation loss that is computed based on: a negative log likelihood computed from the student score assigned to the target token at the output position by the student score distribution for the output position; and a difference measure between the student score distribution for the output position, and the teacher score distribution for the output position.

10. The method of claim 9, wherein the token-level distillation loss is a weighted sum of the negative log likelihood and the difference measure.1 1. The method of any preceding claim, wherein the teacher language model neural network has more parameters than the student language model neural network.

12. 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 morecomputers to perform the operations of the respective method of any one of claims 1-11.

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

Images