Performing task processing tasks using smaller text processing neural networks

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for performing task processing tasks using smaller task processing neural networks. For example, training the smaller task processing neural network can include training the task processing neural network through reinforcement learning using a combinations of multiple reward signals. As another example, performing inference using the smaller task processing neural network can include using the smaller task processing neural network as part of a cascade.

Description

[0001] PERFORMING TASK PROCESSING TASKS USING SMALLER TEXT PROCESSING NEURAL NETWORKS

[0002] CROSS-REFERENCE TO RELATED APPLICATIONS

[0003] This application claims the benefit of priority to U.S. Provisional Patent Application No. 63 / 520,912, filed on August 21, 2023, the contents of which are hereby incorporated by reference.

[0004] BACKGROUND

[0005] This specification relates to processing text 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 training system implemented as one or more computer programs on one or more computers that trains a text processing neural network. This specification also describes an inference system implemented as one or more computer programs on one or more computers that performs inference using a trained text processing 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] Large Language Models (LLMs) have demonstrated impressive capabilities for performing text rewriting and other text processing tasks. Nonetheless, the large sizes of these models make them impractical for on-device inference. On-device inference refers to generating text outputs using a neural network that is deployed locally on a user device, such as a mobile user device with limited memory.

[0011] On-device inference is beneficial relative to server-side inference (where the model is deployed on a server, e.g., in a data center, and is accessed over a data communication network) because it allows for enhanced privacy and security of user data, generates outputs with lower latency, and does not consume network bandwidth, i.e., reduces the amount of communication over the data communication network that is required between the user device and the server.

[0012] While LLMs have performed well for server-side inference, training and deploying a smaller yet potent language model neural network (“text processing neural network’") for text rewriting or other text processing tasks presents a formidable challenge because it requires balancing the need for a small model size with the need to retain the emergent capabilities of the LLM for performing these tasks.

[0013] Generating this balance would generally require costly data collection, e g., to train a reward model for use in reinforcement learning from human feedback. In addition to being costly, accurate human feedback is difficult and, in some cases, even impossible to obtain in the quantity that would generally be required to yield an amount of training data that is sufficient for use in training the neural network to achieve good performance.

[0014] To address the above challenges, this specification describes an approach for building an edge device-centric text processing model, i.e., a neural network that is optimized for being deployed on an edge device with limited memory and processing capabilities. The described techniques, which include a heuristic reinforcement learning framework which substantially enhances performance without requiring preference data for training a reward model, allow a smaller model to be trained to have performance that is similar to or exceeds that of a larger model that can only be deployed in the cloud.

[0015] To further bridge the performance gap with the larger server-side model or to further improve model performance, this specification also describes an effective approach that combines the mobile text processing model with the server model using a cascade. In particular, the approach is computationally efficient because the same mobile text processing model is used to both generate a candidate output and to determine whether to send the corresponding input to the server-side model for processing (instead of using the candidate output as the final output for the corresponding input). Thus, this “cascade"’ minimizes consumption of netw ork bandwidth and latency w hile maintaining high output quality. Moreover, the system is able to perform this cascade and achieve these benefits with minimal additional computational overhead on the edge device because the same mobile text processing model is used to both generate a candidate output and to determine whether to send the corresponding input to the server-side model for processing.

[0016] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below'. Other features, aspects, and advantages of the subj ect matter will become apparent from the description, the drawings, and the claims.

[0017] BRIEF DESCRIPTION OF THE DRAWINGS

[0018] FIG. 1 shows an example training system.

[0019] FIG. 2 is a flow diagram of an example process for training a text processing neural network through reinforcement learning.

[0020] FIG. 3 shows an example of generating data for training the text processing neural network.

[0021] FIG. 4 is a flow diagram of an example process for performing a text processing task using a cascade.

[0022] FIG. 5 is a flow diagram of an example process for training the text processing neural network to evaluate output sequence quality.

[0023] FIG. 6 shows an example of the training of the text processing neural network to perform quality’ evaluation.

[0024] FIG. 7 shows an example of the performance of the described techniques relative to existing techniques.

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

[0026] DETAILED DESCRIPTION

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

[0028] The training system 100 trains a text processing neural network 110.

[0029] As described in this specification, a text processing neural network, e.g., the text processing neural network 110, is a neural network that is configured to perform a plurality of text processing tasks.

[0030] After the training system 100 has trained the text processing neural network 110, an inference system 150 uses the text processing neural network 110 to perform some or all of the text processing tasks. A “text processing"’ task is any task that requires processing an input 102 that includes a sequence of text, i.e., a sequence of text tokens, generating an output 112 that includes a sequence of text tokens, or both.

[0031] The text tokens can be tokens selected from a vocabulary of text tokens that includes, e.g., one or more of characters, word pieces, words, punctuation marks, numerical symbols, or any other text symbols.

[0032] For example, the text processing tasks can be text rewnting tasks that each require processing an input text sequence to generate an output text sequence that is a rewritten version of the input text sequence.

[0033] For example, one text rewriting task can be to generate an output text sequence that is a more formal version of the input text sequence but that conveys the same semantic meaning.

[0034] As another example, one text rewriting task can be to generate an output text sequence that is a shorter version of the input text sequence but that conveys the same semantic meaning.

[0035] As another example, one text rewriting task can be to generate an output text sequence that is a more elaborate version of the input text sequence but that conveys the same semantic meaning.

[0036] As another example, one text rewriting task can be to generate an output text sequence that is a paraphrased version of the input text sequence, i.e.. one that uses different words from the input text sequence but that conveys the same semantic meaning.

[0037] As another example, one text rewriting task can be to generate an output text sequence that is a proofread version of the input text sequence, i.e., one that corrects grammar and spelling mistakes in the input text sequence.

[0038] The text processing tasks can also include other tasks instead of or in addition to the text rewriting tasks.

[0039] For example, the text processing tasks can include a task that requires generating an output text sequence that is a completion of an input text sequence.

[0040] As another example, the text processing tasks can include a task that requires generating an output text sequence that is an answ er to a query posed by the input text sequence.

[0041] In general, the text processing task that should be performed on any given input 102 is defined by an input prompt 104 that is provided as input to the neural network 110. For example, the input prompt 104 can be a natural language instruction that describes the task to be performed on the corresponding input (e.g., “make this text formal,” “make the text more concise,” “rewrite this sentence.” “please proofread this sentence,” and so on).

[0042] That is, the text processing neural network 110 is configured to process a task input 102 and a prompt 104 that defines the task to generate an output 112 for the task that is defined by the prompt 104.

[0043] The text processing neural network 110 can generally have any appropriate architecture that allows the text processing neural network 110 to map a sequence of input tokens to a sequence of output tokens.

[0044] For example, the text processing neural network 110 can be a language model neural network, i.e., an auto-regressive neural network that generates output sequences of tokens from a vocabulary, e.g., conditioned on a context sequence.

[0045] As a particular example, the tokens in the vocabulary can be text tokens, such that the language model neural network maps a context sequence of text tokens to an output sequence of text tokens. As another particular example, the language model neural network can be a multi-modal neural network, with the tokens in the context sequence including tokens representing another data modality, e.g., images, videos, or audio, instead of or in addition to text tokens.

[0046] The neural network 110 is referred to as an auto-regressive neural network because the neural network 110 auto-regressively generates an output sequence of tokens by generating each particular token in the output sequence conditioned on a current input sequence that includes any tokens that precede the particular text token in the output sequence, i.e., the tokens that have for already been generated for any previous positions in the output sequence that precede the particular position of the particular token, and a context input that provides context for the output sequence. For example, the current input sequence when generating a token at any given position in the output sequence can include the context sequence and the tokens at any preceding positions that precede the given position in the output sequence. As a particular example, the current input sequence can include the context sequence followed by the tokens at any preceding positions that precede the given position in the output sequence. Optionally, the context and the current output sequence can be separated by7one or more predetermined tokens within the current input sequence.

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

[0048] As a particular example, the language model neural network 110 can be an autoregressive Transformer-based 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.

[0049] The neural network 110 can have any of a variety of Transformer-based neural network architectures. Generally, however, the Transformer-based neural network includes a sequence of attention blocks (a block that applies an attention mechanism over a block input to generate a block output), and, during the processing of a given input sequence, each attention block in the sequence receives a respective input hidden state for each input token in the given input sequence. The attention block then updates at least the hidden state for the last token in the sequence at least in part by applying self-attention to generate a respective output hidden state for the last token. The input hidden states for the first attention block are embeddings of the input tokens in the input sequence and the input hidden states for each subsequent attention block are the output hidden states generated by the preceding attention block.

[0050] In this example, the output subnetwork processes the output hidden state generated by the last attention block in the sequence for the last input token in the input sequence to generate the score distribution.

[0051] In particular, the architecture of the text processing neural network 110 is designed to be computationally efficient to allow the text processing neural network 110 to, when deployed on an edge device, generate outputs with minimal latency while still generating high-quality outputs.

[0052] For example, the text processing neural network 110 can be used to generate text in real-time in response to user requests when deployed on an edge device. As a particular example, the text processing neural network 110 can be deployed as part of an on-device assistant or other software application that provides responses to user requests without sending user data to a server or other remote system.

[0053] The text processing neural network 110 can be referred to as a "smaller" text processing neural network because the text processing neural network has a relatively small model size, i.e., number of parameters, number of layers, or both, relative to a neural network that is designed for being deployed in a data center or other cloud computing environment. The model size of the text processing neural network 110 can be chosen to be small enough to run efficiently on a user mobile computer device, e.g., to fit within a target memory constraint and / or have less than a target latency at inference time on the user device.

[0054] For example, relative to a neural network that is deployed in a data center, the text processing neural network 110 can have fewer parameters, e.g., because the neural network 110 includes fewer self-attention layer blocks, because the hidden states operated on by the self-attention layer blocks have a smaller dimensionality', or both.

[0055] That is. in some implementations, after training, the text processing neural network 110 is deployed on an edge device, e.g., a mobile device, e.g., a smartphone or a tablet, a smart speaker or another embedded computing device, and so on. In other words, the inference system 150 is implemented as one or more computer programs running on the edge device. In these implementations, once generated, the inference system 150 can display the generated outputs in a user interface of the edge device or otherwise provide the generated outputs to the user, e.g., by converting the outputs to speech using a text-to-speech system and then playing back the generated speech to the user.

[0056] The training system 100 trains the text processing neural network 110 to effectively perform text processing tasks even though the text processing neural network 110 has an architecture that is computationally-efficient enough to allows the text processing neural network to effectively be deployed on an edge device, e.g., is small enough to fit on local memory of the edge device and / or run with a low latency on the edge device.

[0057] As one particular example, the text processing neural network 110 can have 20x, 5 Ox, or even lOOx fewer parameters than the larger neural network that is deployed in the data center. As a particular example, the text processing neural network 110 can have a PaLM 2-XXS architecture, while the larger neural network can have the PaLM 2-L architecture.

[0058] As at least part of the training, the system 100 trains the neural netw ork 110 through reinforcement learning to perform multiple different text processing tasks using, for each task, weights 120 for each of a set of reward signals that each measure a different property of an output sequence. By using different weights 120 for different tasks, the system 100 can effectively train the neural network 110 without needing humangenerated reward values.

[0059] This training will be described in more detail below with reference to FIGS. 2, 3, and 5.

[0060] In some implementations, to allow the text processing neural network 110 to effectively perform inference when deployed on-device without degrading output quality, the inference system 150 can determine, using the text processing neural network 110, the predicted quality of each output generated by the text processing neural network 110 and then use a different neural network 170, e.g., one that is deployed remotely, e.g., serverside in the cloud, to generate a final output if the predicted quality is too low, e.g., below a threshold quality value.

[0061] This type of processing “cascade” will be described in in more detail below with reference to FIG. 5.

[0062] More generally, however, the training and inference techniques that are described in this specification can be used to train any appropriate type of generative neural network. For example, the techniques can be used to train and perform inference for an audio generation neural network that generates sequences of audio data conditioned on input sequences, e.g., text sequences, images, or other audio sequence. As another example, the techniques can be used to train and perform inference for an image generation neural network that generates sequences that represent images conditioned on input sequences, e.g., text sequences, audio sequences, or sequences that represent other images.

[0063] FIG. 2 is a flow diagram of an example process 200 for training the text processing neural network through reinforcement learning. 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 training system, e.g., the training system 100 of FIG. l, appropriately programmed, can perform the process 200.

[0064] In some implementations, the system performs the process 200 to train the text processing neural network from “scratch,” i.e., starting from parameters that have been randomly initialized.

[0065] In some other implementations, however, the system or another training system “pre-trains” the text processing neural network prior to training the text processing neural network using the process 200. For example, prior to training the text processing neural network through reinforcement learning, the system can train the text processing neural network through supervised fine-tuning on a set of training examples that each include (i) a respective training input prompt for a respective task and (ii) a respective target output sequence for the respective task to be generated by the text processing neural network by processing the respective training input prompt.

[0066] When the text processing tasks require generating an output sequence conditioned on an input sequence, each training example can also include (iii) a respective training input sequence.

[0067] However, this labeled data can be difficult to obtain in sufficient quantity to allow the text processing neural network to be trained effectively. For example, because a given text rewriting task can be described using many different natural language prompts, it can be difficult to collect a sufficiently labeled training data for a given task that include a sufficiently diverse set of input prompts. As another example, because there are many different variants of text rewriting tasks, it can be difficult to collect enough high-quality input sequence - output sequence pairs for any given task to train a neural network to perform well on a wide array of text rewriting tasks.

[0068] To account for this, in some implementations, the system generates some or all of the training examples using another neural network, e g., an off-the-shelf text processing neural network. For example the other neural network can be the server-side neural network described above.

[0069] This is described in more detail below with reference to FIG. 3.

[0070] As another example, prior to the training of the text processing neural network through supervised fine-tuning, the system or the other training system can have pretrained the text processing neural network on a language modeling task or a different selfsupervised learning objective.

[0071] To train the text processing neural network through reinforcement learning, the system maintains data specifying, for each of the plurality of text processing tasks, a respective set of weights that includes a respective weight for each of a plurality of reward signals (step 202).

[0072] Each reward signal is a function that, for any given input prompt for any given task, maps an output sequence for the given input prompt and, optionally, a corresponding input sequence for the input prompt to a score for the reward signal. Generally, different reward signals measure different properties of the output sequences and, in some cases, of the input sequences. Additionally, reward signals can be determined directly from an output sequence (and, in some cases, the corresponding input sequence).

[0073] As will be described in more detail below, the system then uses the reward signals and corresponding weights to determine reward values for use in reinforcement learning training. Thus, the system can perform the training of the text processing neural network through reinforcement learning without making use of any reward models or any preference data that measures user preferences between different output sequences.

[0074] Some examples of reward signals now follow.

[0075] As one example, the reward signals can include a natural language inference (NLI) score generated by an NLI predictor that, for each input prompt, uses the output sequence for the input prompt as a hypothesis and the input sequence corresponding to the input prompt as a premise. In particular, given a “premise” and a “hypothesis”, the NLI score is the probability that the “hypothesis” is correct given the “premise”. In the context of generative neural networks, the NLI score estimates whether or not the LLM’s output prediction preserves meaning and factuality given the source text.

[0076] The system can generally use any appropriate NLI predictor to generate the NLI scores. One example of an NLI predictor that can be used by the sy stem to generate the NLI score is described in Or Honovich, Roee Aharoni, Jonathan Herzig, Hagai Taitelbaum, Doron Kukliansy, Vered Cohen, Thomas Scialom, Idan Szpektor, Avinatan Hassidim, and Yossi Matias. 2022. True: Re-evaluating factual consistency evaluation. In Proceedings of the Second DialDoc Workshop on Document-grounded Dialogue and Conversational Question Answering, pages 161-175.

[0077] As another example, the reward signals can include a reversed NLI score generated by an NLI predictor that, for each input prompt, uses the output sequence for the input prompt as the premise and the input sequence corresponding to the input prompt as the hypothesis. Thus, the reversed NLI score is similar to the NLI score but the premise and the hypothesis are reversed.

[0078] As yet another example, the reward signals can include a length ratio score that, for each input prompt, is based on a ratio of (i) the length of the output sequence for the input prompt to (ii) the length of the input sequence corresponding to the input prompt. For example, the length of a given sequence can be measured by the number of tokens in the sequence. As yet another example, the reward signals can include an edit distance score that, for each input prompt, is based on an edit distance between the output sequence for the input prompt and the corresponding input sequence for the input prompt. The edit distance measures the minimum number of token-level edits (insertions, deletions and substitutions) to convert the input sequence into the output sequence. In some cases, the edit distance score is a normalized score that is also based on the length of the corresponding input sequence for the input prompt.

[0079] As yet another example, the reward signals can include an n-gram frequency score that, for each input prompt, is based on a frequency with which each n-gram that is present in the output sequence occurs in the output sequence. As a particular example of this, the n-gram frequency score can be equal to first value when the most frequency occurring n-gram in the output sequence occurs at least a threshold number of times and equal to a second value when the most frequency occurring n-gram in the output sequence occurs less than the threshold number of times.

[0080] Because different tasks require different types of output signals, the system maintains different weights for different ones of the tasks. For example, the system can receive data specifying the weights for the tasks as an input from a user prior to beginning the training or can perform a hyperparameter search to determine the weights for the tasks.

[0081] The system then repeatedly performs steps 204-212 in order to train the text processing neural network. In some examples, training is performed until a threshold performance is achieved, e.g., a threshold performance on a test dataset and / or relative to the server-side model.

[0082] The system obtains a set of one or more input prompts (step 204). As described above, in some cases, the text processing tasks each require generating an output sequence conditioned on an input sequence. In these cases, the system also obtains a corresponding input sequence for each of the input prompts.

[0083] For example, the system can sample the set of one or more input prompts from a larger set of input prompts or can generate the set of one or more input prompts using a language model neural network as will be described below.

[0084] For each input prompt, the system processes an input that includes the input prompt using the text processing neural network to generate an output sequence for the input prompt (step 206). When the text processing tasks require generating an output sequence conditioned on an input sequence, the input to the text processing neural network further includes the respective input sequence that corresponds to the input prompt.

[0085] For each input prompt, the system determines a text processing task from the plurality of text processing tasks that corresponds to the input prompt (step 208).

[0086] That is, the system identifies which of the text processing tasks the input prompt corresponds to.

[0087] In some cases, e.g.. where the input prompts are sampled from a labeled set of input prompts, the system has access to data that identifies which task each input prompt belongs to and uses this data to determine the corresponding task for any given input prompt.

[0088] In some other cases, e.g.. when the input prompts have been generated by the system, the system does not have any data that identifies correspondences between prompts and tasks.

[0089] In these cases, the system can determine the task that corresponds to any given prompt using a pre-trained language model neural network. For example, this can be the server-side neural network described above.

[0090] In particular, the system can process an input that includes the input prompt using the language model neural network to generate an output that identifies the text processing task that corresponds to the input prompt. For example, the input can also include (i) a respective text label for each of the plurality of text processing tasks and (ii) one or more examples that each include (i) an example input prompt and (ii) an example output that identifies one of the text processing tasks. In this case, the output that identifies one of the text processing tasks can be the text label for the text processing task. That is. by virtue of including the additional data in the input to the language model neural network, the system causes the language model neural network to generate, as output, one of the text labels, i.e., without needing to further train or otherwise adapt the pre-trained neural network to perform task classification.

[0091] For each input prompt, the system determines a respective value for each of the plurality of reward signals from at least the output sequence for the input prompt (step 210). That is, for each reward signal, the system determines the value of the reward signal for the output sequence (and, optionally, the input sequence) by applying the reward signal, i.e., the function, to the output sequence (and, optionally, the input sequence), e.g., as described above for the example reward signals. For each input prompt, the system determines a reward value by combining the respective values for the reward signals in accordance with the respective set of weights for the text processing task that corresponds to the input prompt (step 212).

[0092] As described above, different tasks generally have different weights for the reward signals. As a result, the resulting reward value for any given prompt is tailored to emphasize the properties of output sequences that are required by the corresponding task.

[0093] The system then trains the text processing neural network using the reward values for the one or more input prompts through reinforcement learning (step 214).

[0094] In particular, the system can train the text processing neural network to optimize any appropriate reinforcement learning objective that is based on the reward values for the input prompts. As one example, the system can use a policy gradient reinforcement learning technique to perform the training, e.g., the REINFORCE technique or the Proximal policy optimization (PPO) technique.

[0095] By repeatedly performing steps 204-214 on different sets of input prompts, the system trains the text processing neural network to generate outputs for the various text processing tasks that properly weight the multiple different reward signals for the corresponding task.

[0096] FIG. 3 shows an example 300 of generating data for training the text processing neural network.

[0097] For example, the technique shown in FIG. 3 can be used to generate training examples for training the neural network through supervised fine-tuning. In particular, the system can generate the training examples by processing few -shot prompts using a language model neural network, e.g., a pre-trained, off-the-shelf language model neural network. For example, the language model neural network can be a foundation model available to the system.

[0098] In particular, the example 300 shows an example few-shot prompt 310 for a text rewriting task. The prompt 310 includes multiple examples, with each example including an original text, a rewritten text, and a prompt. The example 300 also shows a query 320 that includes an original text and a rew ritten text.

[0099] By processing the prompt 310 and the query 320, the language model neural network can generate a new prompt 330 that describes how7the rewritten text in the query 320 was generated from the original text in the query 320. In addition, by processing the prompt 310, the query 320, and the new prompt 330, the language model neural network can generate a set of “hallucinated” examples 330.

[0100] The system can then use these hallucinations 330 as candidate training examples for the training of the text processing neural network.

[0101] For example, the system can generate a plurality of initial training examples by processing a plurality of few-shot prompts using the language model neural network and then filter the plurality of initial training examples to generate training examples.

[0102] For example, the system can use the language model neural network to filter the candidate examples. As a particular example, the system can leverage few-shot Chain-of- Thoughts (CoT) reasoning of the language model neural network to cause the language model neural network to determine whether, for each candidate, the response is following the instruction of the prompt to rewrite the original sentence in a good manner. As another example, the system can leverage the self-consistency approach to improve the accuracy of filtering, i.e., by removing candidate examples that are not self-consistent when applying the self-consistency approach. Self-consistency and chain of though reasoning are described in more detail in Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. 2022a. Selfconsistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203. 11171.

[0103] Thus, by using the language model neural network, the system can generate training examples for supervised fine-tuning, training prompts for reinforcement learning, or both.

[0104] FIG. 4 is a flow diagram of an example process 400 for performing a text processing task using a cascade. 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, an inference system, e.g., the inference system 150 of FIG. 1, appropriately programmed, can perform the process 400.

[0105] For example, as described above, the process 400 can be performed by an edge device on which a text processing neural network is deployed. The edge device may be able to defer certain inputs for processing by a larger text processing neural network that is deployed remotely from the edge device, e.g., in a data center. That is, the edge device can provide any given input to the remote, larger text processing neural network for processing. However, deferring an input for processing results in increased latency due to the need to transmit data to and from a remote system, decreases the privacy and security of user data, and consumes network bandwidth. Thus, the system performs the process 400 to accurately determine when to defer an input to the text processing neural network to minimize these undesirable consequences while maintaining high performance on the requisite tasks.

[0106] The system receives an input sequence for a text processing task (step 402).

[0107] The system processes a first input that includes the input sequence using the text processing neural network to generate an output sequence for the task processing task (step 404). For example, the text processing neural network can be deployed locally on the same computer(s) as the system, e.g., on the same edge device.

[0108] The system process a second input that includes the input sequence, the output sequence, and an instruction to evaluate a quality of the output sequence for the text processing task using the text processing neural network to generate a quality score for the output sequence (step 406).

[0109] The quality score for the output sequence is an estimate of the quality of the output sequence, i.e., as an output for the input sequence for the text processing task. That is, the system uses the same text processing neural network to both generate the output sequence and to evaluate the quality of the output sequence.

[0110] To increase the computational efficiency of the evaluation, the system can re-use hidden states of the task processing neural network generated while processing the first input when processing the second input. That is, the system can start the processing of the second input using cached hidden states or keys and values of the attention layers in the language processing neural network generated by processing the first input.

[0111] For example, the system can determine the quality score based on a probability assigned to a predetermined sequence of one or more tokens by the text processing neural network by processing the second input. For example, the quality' score can be the log likelihood assigned to the text sequence “good,’' the text sequence “bad,” or another appropriate text sequence as a result of processing the second input.

[0112] The system determines, based on the quality score, whether to (i) use the output sequence as a final output sequence for the text processing task or (ii) use a different text processing neural network to generate the final output sequence for the text processing task (step 408). For example, the different neural network can be the server-side neural network described above. For example, e.g., when the predetermined sequence of tokens is indicative of a high quality output sequence, the system can determine to use the output sequence as the final output for the text processing task when the quality score exceeds a threshold value.

[0113] As another example, e.g., when the predetermined sequence of tokens is indicative of a low quality output sequence, the system can determine to use the output sequence as the final output for the text processing task when the quality’ score does not exceed the threshold value.

[0114] That is, in response to determining to use the output sequence as the final output sequence, the system provides the output sequence as a response to the input sequence without using the different text processing neural network (step 410).

[0115] In response to determining to use the different text processing neural network to generate the final output sequence for the text processing task, the system provides the input sequence for processing by the different text processing neural network (step 412), obtains an output sequence from the different text processing neural network (step 414), and provides the output sequence obtained from the different text processing neural network as the response to the input sequence (416).

[0116] As described above, using the different text processing neural network is disadvantageous unless using the different text processing neural network will result in a higher quality output sequence.

[0117] For example, the different text processing neural network can have more parameters than the text processing neural network, resulting in increased latency and increased memory consumption when the different neural network needs to be used.

[0118] More generally, generating the final output sequence using the different text processing neural network consumes more computational resources, incurs a higher latency, or both, than generating the output sequence using the text processing neural network.

[0119] For example, as described above, the input sequence can be received by a user device, the text processing neural network can be deployed on the user device, and the different text processing neural network can be deployed on one or more computers that are remote from the user device. Thus, processing inputs using the different neural network requires transmiting data and receiving data over a network.

[0120] In some cases, prior to using the text processing neural network as part of the cascade, the system fine-tunes the text processing neural network to improve the performance of the neural network in evaluating output sequence quality. FIG. 5 is a flow diagram of an example process 500 for training the text processing neural network to evaluate output sequence quality. For convenience, the process 500 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 of FIG. 1, appropriately programmed, can perform the process 500.

[0121] The system receives an input sequence for a text processing task and a candidate output sequence for the text processing task (step 502). For example, the system can have generated the candidate output sequence using the text processing neural network, i.e., by processing an input that includes an input prompt for the task and the input sequence using the text processing neural network.

[0122] The system processes a first input that includes the input sequence for the text processing task and the candidate output sequence for the text processing task using a pretrained language model neural network to generate a text label that indicates a quality of the candidate output sequence for the text processing task (step 504). For example, the pre-trained neural network can be the server-side neural network described above.

[0123] For example, the system can cause the language model neural network to generate the text label through few-shot prompting. That is, the first input can include a few-shot prompt that includes one or more examples, with each example including (i) an example input sequence, (ii) an example candidate output sequence, and (iii) an example text label.

[0124] The system processes a second input that includes the input sequence, the candidate output sequence, and an instruction to evaluate a quality of the output sequence for the task processing task using the text processing neural network to generate an output that assigns a likelihood to the text label (step 506). For example, the likelihood can be the log likelihood assigned to the text label by the output.

[0125] The system trains the text processing neural network on an objective that measures the likelihood assigned to the text label by the output generated by the text processing neural network (step 508).

[0126] FIG. 6 shows an example 600 of the training of the text processing neural network to perform quality evaluation.

[0127] As shown in the example 600, an input that includes a prompt 610 ('‘Make it formal”) and an input sequence 620 (“I’m busy”) is processed using the text processing neural network 110 (“On device LLM”) to generate an output sequence 630 (“I am occupied”). The system then processes a corresponding input, i.e., that includes the prompt 610, the input sequence 620. and the output sequence 630 using another pre-trained text processing neural network 640 (“Off-the-shelf LLM”) to generate a critique 650 of the output sequence 630, i.e., that identifies either the text label “good” or the text label “bad” for the output sequence 630.

[0128] The system 100 also processes another input that includes the prompt 610, the input sequence 620, the candidate output sequence 630. and an instruction 660 to evaluate a quality of the output sequence (“# quality is”) using the text processing neural network 110 to generate an output 670 that assigns a likelihood to the text label “good,” i.e., assigns a likelihood to outputting “good” given the prompt 610. the input sequence 620, the candidate output sequence 630, and the instruction 660.

[0129] The system can then train the text processing neural network 110 on an objective that measures the likelihood, as described above.

[0130] FIG. 7 show s an example 700 of the performance of the described techniques relative to existing techniques.

[0131] In particular, FIG. 7 shows a table 710 showing performance of various techniques on an EditEval data set that requires performing text editing tasks like making text more cohesive and paraphrasing. Here, SFT(Ours) refers to a model trained using supervised fine tuning with data generated as described with reference to FIG. 3 and SFT + heuristic RL(Ours) refers to a model trained using supervised fine tuning with data generated as above with reference to FIG. 3 and then further trained through reinforcement learning as described with reference to FIG. 4. Additionally, in the Table Size XXS is less than half the size of TO / Tk models and size XXS is over 20 times smaller than model size and over 5 times sailer than model size XS

[0132] As can be seen from table 710, Despite their reduced sizes, the described techniques achieve even better performance than most of the other larger models. Relative to similar-sized instruct! on-tuned models, the described techniques win by a large margin.

[0133] FIG.7 also shows a table 720 showing performance of various techniques on an EditEval data set that requires performing text rewriting tasks based on natural language instructions.

[0134] In table 720, “SFT"’ under “Our instruction-tuned On-device LLMs” refers to the SFT(Ours) model as described above and “SFT + Huerisic RL” refers to the SFT + heuristic RL(Ours) model as described above. Under "cascades." SFT + heuristic RL + critique distillation” refers to the performance of the model trained as described above and further trained as described with reference to FIGS. 5 or 6, but without cascading. Two variants of the cascade are also shown, one with the larger model (InsGPT in this example) being called 15% of the time and one with the larger model being called 40% of the time. As can be seen from the table 720, the described techniques achieve the best performance compared with all listed Pre-trained LLMs and Instruction-Tuned LLMs, which have either same or larger size then the model used by the described techniques. When cascaded with InsGPT, the performance is further improved, even if only 15% of outputs are routed to the larger model, indicating that the small, on-device model can accurately determine when routing to the larger model would be most beneficial.

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

[0136] 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. The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g.. code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

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

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

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

[0140] 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 circuitry7. 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.

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

[0142] 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 yvith 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 yvith a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

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

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

[0145] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the 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.

[0146] 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 anses 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.

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

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

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

[0150] What is claimed is:

Claims

CLAIMS1. A method performed by one or more computers and for training a text processing neural network to perform a plurality' of text processing tasks, the method comprising: maintaining data specifying, for each of the plurality of text processing tasks, a respective set of weights that includes a respective weight for each of a plurality of reward signals; obtaining a set of one or more input prompts; for each input prompt: processing an input that comprises the input prompt using the text processing neural network to generate an output sequence for the input prompt; determining a text processing task from the plurality' of text processing tasks that corresponds to the input prompt; determining a respective value for each of the plurality’ of reward signals from at least the output sequence for the input prompt; and determining a reward value by combining the respective values for the reward signals in accordance with the respective set of weights for the text processing task that corresponds to the input prompt; and training the text processing neural network using the reward values for the one or more input prompts through reinforcement learning.

2. The method of claim 1, further comprising: prior to training the text processing neural network using the reward values, training the neural network through supervised fine-tuning on a set of training examples that each comprise (i) a respective training input prompt and (ii) a respective target output sequence.

3. The method of claim 2, wherein each of the plurality' of text processing tasks require generating an output sequence conditioned on an input sequence, and wherein each training example comprises (iii) a respective training sequence.

4. The method of any one of claims 2 or 3, wherein the text processing neural network has been pre-trained on a language modeling task prior to the training of the text processing neural network through supervised fine-tuning.

5. The method of any one of claims 2-4, wherein each training example in a first plurality of the training examples for the supervised fine-tuning has been generated by a first language model neural network by processing a respective few-shot prompt.

6. The method of claim 5, further comprising: generating a plurality of initial training examples by processing a plurality of fewshot prompts using the first language model neural network; and filtering the plurality of initial training examples to generate the first plurality of training examples.

7. The method of any preceding claim, wherein each of the plurality of text processing tasks require generating an output sequence conditioned on an input sequence, and wherein, for each of the one or more input prompts, the input to the text processing neural network further comprises a respective input sequence.

8. The method of any preceding claim, wherein determining a text processing task from the plurality of text processing tasks that corresponds to the input prompt comprises: processing an input comprising the input prompt using a second language model neural network to generate an output that identifies the text processing task that corresponds to the input prompt.

9. The method of claim 8, wherein the input that comprises the input prompt further comprises (i) a respective text label for each of the plurality of text processing tasks and (ii) one or more examples that each include (i) an example input prompt and (ii) an example output that identifies one of the text processing tasks.

10. The method of any preceding claim, when dependent on claim 7, wherein the reward signals comprise a natural language inference (NLI) score generated by an NLI predictor that, for each input prompt, uses the output sequence for the input prompt as a hypothesis and the input sequence in the input for the input prompt as a premise.

11. The method of any preceding claim, when dependent on claim 7, wherein the reward signals comprise a reversed NLI score generated by an NLI predictor that, for each input prompt, uses the output sequence for the input prompt as the premise and the input sequence in the input for the input prompt as the hypothesis.

12. The method of any preceding claim, when dependent on claim 7, wherein the reward signals comprise a length ratio score that, for each input prompt, is based on a ratio of a length of the output sequence for the input prompt as to a length of the input sequence in the input for the input prompt.

13. The method of any preceding claim, when dependent on claim 7, wherein the reward signals comprise an edit distance score that, for each input prompt, is based on an edit distance between the output sequence for the input prompt and the input sequence in the input for the input prompt.

14. The method of claim 13. wherein, for each input prompt, the edit distance score is further based on length of the input sequence in the input for the input prompt.

15. The method of any preceding claim, when dependent on claim 7, wherein the reward signals comprise an n-gram frequency score that, for each input prompt, is based on a frequency with which each n-gram that is present in the output sequence occurs in the output sequence.

16. The method of claim 15, wherein the n-gram frequency score is equal to first value when a most frequency occurring n-gram in the output sequence occurs at least a threshold number of times and is equal to a second value when the most frequency occurring n-gram in the output sequence occurs less than the threshold number of times.

17. A method performed by one or more computers, the method comprising: receiving an input sequence for a text processing task; processing a first input comprising the input sequence using a text processing neural network to generate an output sequence for the task processing task; processing a second input comprising the input sequence, the output sequence, and an instruction to evaluate a quality of the output sequence for the text processing task using the text processing neural network to generate a quality score for the output sequence; and determining, based on the quality score, whether to (i) use the output sequence as a final output sequence for the text processing task or (ii) use a different text processing neural network to generate the final output sequence for the text processing task.

18. The method of claim 17, wherein the different text processing neural network has more parameters than the text processing neural network.

19. The method of claim 17 or 18, wherein generating the final output sequence using the different text processing neural network consumes more computational resources, incurs a higher latency, or both, than generating the output sequence using the text processing neural network.

20. The method of any one of claims 17-19, wherein the input sequence is received by a user device, the text processing neural network is deployed on the user device, and the different text processing neural network is deployed on one or more computers that are remote from the user device.

21. The method of any one of claims 17-20, wherein processing a second input comprising the input sequence, the output sequence, and an instruction to evaluate a quality of the output sequence for the text processing task using the text processing neural network to generate a quality score for the output sequence comprises: re-using hidden states of the task processing neural network generated while processing the first input.

22. The method of any one of claims 17-21, wherein processing a second input comprising the input sequence, the output sequence, and an instruction to evaluate a quality of the output sequence for the text processing task using the text processing neural network to generate a uality score for the output sequence comprises: determining the quality score based on a probability assigned to a predetermined sequence of one or more tokens by the text processing neural network by processing the second input.

23. The method of any one of claims 17-22, wherein determining, based on the quality score, whether to (i) use the output sequence as a final output sequence for the text processing task or (ii) use a different text processing neural network to generate the final output sequence for the text processing task comprises: determining to use the output sequence as the final output for the text processing task when the quality score exceeds a threshold value; and determining to use the different text processing neural network to generate the final output sequence for the text processing task when the quality score does not exceed the threshold value.

24. The method of any one of claims 17-22, wherein determining, based on the quality score, whether to (i) use the output sequence as a final output sequence for the text processing task or (ii) use a different text processing neural network to generate the final output sequence for the text processing task comprises: determining to use the output sequence as the final output for the text processing task when the quality score does not exceed a threshold value; and determining to use the different text processing neural network to generate the final output sequence for the text processing task when the quality' score exceeds the threshold value.

25. A method performed by one or more computers and for training a text processing neural network, the method comprising: receiving an input sequence for a text processing task and a candidate output sequence for the text processing task; processing a first input comprising the input sequence for the text processing task and the candidate output sequence for the text processing task using a pre-trained language model neural network to generate a text label that indicates a quality of the candidate output sequence for the text processing task; processing a second input comprising the input sequence, the candidate output sequence, and an instruction to evaluate a quality of the output sequence for the task processing task using the text processing neural network to generate an output that assigns a likelihood to the text label; and training the text processing neural network on an objective that measures the likelihood assigned to the text label by the output generated by the text processing neural network.

26. The method of claim 25, wherein the first input comprises a few-shot prompt that comprises one or more examples, each example comprising (i) an example input sequence, (ii) an example candidate output sequence, and (iii) an example text label.

27. The method of claim 25 or claim 26, wherein the likelihood is a log likelihood assigned to the text label by the output.

28. The method of any one of claims 25-27, wherein the pre-trained language model neural network has more parameters than the text processing neural network.

29. The method of any preceding claim, wherein the text processing neural network is an auto-regressive neural network that includes a plurality of self-attention layers.

30. A system comprising: one or more computers; and one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform the respective operations of the method of any one of claims 1 -29.

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

Images