Reward model-based reinforcement learning for performing inference tasks

Abstract

A method, system, and apparatus, including a computer program encoded on a computer storage medium, for training a language model to perform an inference task. The system acquires a plurality of training examples, each training example including a respective sample query text sequence characterizing a respective sample query and a respective reference response text sequence including a reference final answer for the respective sample query. The system trains a reward model with the plurality of training examples. The reward model is configured to receive input including the query text sequence characterizing the query and one or more inference steps generated in response to the query, process the input, and calculate a reward score indicating how successful the one or more inference steps are in producing a correct final answer to the query. The system trains the language model using the trained reward model.

Description

[Technical Field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims priority to U.S. Provisional Patent Application No. 63 / 377,532, filed September 28, 2022, the entire disclosure of which is incorporated herein by reference.

[0002] This specification relates to using machine learning models, such as neural networks, to perform inference tasks. [Background technology]

[0003] A neural network is a machine learning model that uses 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 the 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 the received input according to the current values ​​of its respective set of parameters.

[0004] The parameters of the machine learning model can be determined during the training process based on training data including one or more training examples. For example, a neural network can be trained by updating network parameters including, for example, weights and bias coefficients of network layers of the neural network. [Prior art documents] [Non-patent literature]

[0005] [Non-Patent Document 1] 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

Non-Patent Document 2

[0006] This specification describes methods, computer systems, and apparatuses, including computer programs encoded on computer storage media, for performing inference tasks using language models. In particular, the inference task includes providing a response to an input query. More specifically, the response may include (1) a final answer to the query and (2) an inference process that includes a sequence of one or more inference steps to produce the final answer to the query.

[0007] Language models and language generation neural networks can perform a variety of tasks, such as translation tasks (if the training corpus includes words in different languages) and arithmetic tasks, among many other tasks. In connection with inference tasks, language models can be used to generate responses to queries (e.g., questions). Responses to queries can be provided as text sequences in the form of, for example, a yes / no binary answer, an answer providing one or more numerical values, an answer identifying one or more objects or events, an answer identifying a planned path or action, etc. Responses to queries can be used in many ways by humans or computer systems. For example, the answers and inference processes can be useful themselves, or they can be used to provide warnings and / or to control the movement of one or more robotic agents or objects, e.g., autonomous vehicles.

[0008] In some implementations, the input query may be a natural language query about an environment, particularly a real-world environment. The output response may be a natural language reply or natural language output statement also about the environment. For example, the output response may provide information about the environment, in some implementations, information about or specifying an action to be taken in the environment.

[0009] In one example, the system can be used to diagnose faults in a mechanical system operating in a real-world environment. An input query can include, for example, observations about the mechanical system from one or more sensors (e.g., a camera, a microphone, an accelerometer, a temperature sensor, etc.) and a question about the mechanical system. In some implementations, the observations can include sensed electronic signals, such as motor currents or temperature signals, and / or image or video data, for example, from a camera or LIDAR sensor, data from a sensor on the mechanical agent or from a sensor located separately from the mechanical agent in the environment. The observations and question can be suggested to generate a linguistic expression of the query. For example, the query can include a general question such as, "Given the measurements from sensors A, B, and C, is the system operating normally?" or "Given the measurements from A, B, and C, what is wrong with the system?" or a specific request such as, "Is component X faulty?" The output response can provide a final answer to the question and the reasoning process that led to the final answer.

[0010] In another example, the environment may be an educational environment; for example, the system may be deployed as part of an educational software program that assists users in learning or practicing one or more corresponding skills. For example, the input query may be a mathematical word problem, and the output response includes a correct answer to the word problem and intermediate steps to arrive at the correct answer. The output response may be given to the student as a proof of how to solve the mathematical problem or may be used as a grading tool to provide feedback on the student's practice.

[0011] In another example, the system can be used for natural language control of tasks in a real-world environment. That is, the input query may be task-related and may include, for example, a request to perform the task. The output response can be used to control, for example, a mechanical system (which may be referred to as a machine agent) or a computer system to perform the task. As an example, the input query may include a high-level question, for example, from a human, to perform the task, such as, for example, "What is the most cost-effective way to manufacture this component with our current equipment?" (The real-world environment may be a manufacturing facility in such a case.) The output response includes a final answer specifying a planned procedure and reasoning steps for determining the planned procedure. The output response can be used to control one or more machine agents, for example, a robotic arm, to perform the planned procedure. For example, the output response may include one or more control signals for controlling the one or more machine agents, such as position, velocity, or force / torque / acceleration data for one or more joints of a part of the robot or another machine agent. In some implementations, the input query may relate to a real-world task and may specify observations of the real-world environment, and the output response may specify actions or paths to be taken in the real-world environment by the user to perform the task. The system may provide the output response to the user, for example, via a display or audio output device, to guide the user to perform the task in the real-world environment. The output response generated using the methods described herein can greatly improve the usability and safety with which a machine agent, computer system, or human user can perform real-world tasks in a real-world environment, especially for tasks where generating an incorrect response may have significant negative consequences. In particular, the output response may be used to principledly and logically perform a task.

[0012] In some implementations, the reasoning process in the output response provides a human-interpretable explanation that can be used to 1) determine, for example, by a user, whether to implement an action specified in a final answer in the output response, or 2) be accessed later when certain criteria are satisfied to evaluate the performance of the language model or diagnose the cause of an error that occurs as a result of implementing an action specified in a response generated by the language model. The human-interpretable explanation may include, for example, a sequence of logical steps, presented as natural language statements, leading from the query to the response in a causal chain. Such human-interpretable explanations can overcome or mitigate some of the drawbacks of typically using language models (especially neural network-based language models) as a "black box" by allowing a human user to understand the reason for the response given by the language model.

[0013] In one particular aspect, this specification describes a system implemented as a computer program on one or more computers at one or more locations that trains language models to perform inference tasks.

[0014] The system includes a reward model configured to receive input including a query text sequence characterizing the query and one or more inference steps generated in response to the query, process the input, and calculate a score indicating how successful the one or more inference steps are in producing a correct final answer to the query.

[0015] The system can train the reward model using output-supervised training or process-supervised training using multiple training examples. Generally, each training example includes a respective sample query text sequence characterizing a respective sample query (i.e., a training query) and a respective reference response text sequence containing a reference (“ground truth”) final answer to the respective sample query. For each training example, the system can generate one or more candidate inference processes in response to the sample queries of the training example. Each candidate inference process includes one or more inference steps. The system can assign a target score for each inference step in each candidate inference process based on the reference response text sequences in the training examples, where the target score indicates whether the inference step is correct or incorrect. The system then trains the reward model to generate reward scores for inference steps in the candidate inference processes that match the corresponding target scores for the inference steps.

[0016] In output result monitoring training, the system determines whether a candidate reasoning process produces a reference final answer in the training examples. If the candidate reasoning process produces a reference final answer in the training examples, the system assigns each reasoning step in the candidate reasoning process a reward score indicating that the reasoning step is correct. If the candidate reasoning process does not produce a final answer in the training examples, the system assigns each reasoning step in the candidate reasoning process a reward score indicating that the reasoning step is incorrect.

[0017] In process monitoring training, the reference response text sequence includes a reference reasoning process that includes a sequence of reasoning steps, with the final answer as the last reasoning step. The system assigns a target score to the current reasoning step based on whether the reasoning steps generated up to and including the current reasoning step match the sequence of reasoning steps in the reference reasoning process. That is, the target score for the current reasoning step may be determined based on a comparison (e.g., a similarity score) between each of the reasoning steps generated up to (and including) the current reasoning step and the corresponding reasoning step in the sequence of reasoning steps in the reference reasoning process.

[0018] To train the language model, the system can receive training queries and use the policy language model and the trained reward model to generate one or more expert reasoning processes in response to the training queries, where each respective expert reasoning process includes a respective plurality of reasoning steps. The system can use the one or more expert reasoning processes to train the language model. The policy language model can be a model configured to receive input including a query text sequence characterizing the query and one or more reasoning steps generated in response to the query, and process the input to generate a next reasoning step.

[0019] In some implementations, to generate expert processes, the system uses a policy language model to generate multiple candidate expert reasoning processes in response to a training query. The system can use the reward model to generate a performance score for each of the candidate expert reasoning processes based on reward scores generated for reasoning steps in the candidate expert reasoning processes. The system selects one or more candidate expert reasoning processes with the highest performance scores as the expert reasoning processes.

[0020] In some implementations, to generate an expert process, after one or more inference steps have been generated for the expert reasoning process, the system uses the policy language model to generate multiple next candidate inference steps. The system uses the reward model to generate a reward score for each of the next candidate inference steps, and selects the next candidate inference step with the highest reward score as the next step for the expert reasoning process. The system may repeat the above steps until the next step matches the final answer indicator (e.g., until the next inference step contains text corresponding to the final answer) or until the maximum number of steps has been reached. The system can then use the resulting inference process as the expert reasoning process.

[0021] In another aspect, this specification describes a system implemented as a computer program on one or more computers at one or more locations that uses a language model to perform reasoning tasks.

[0022] The system may obtain an input text sequence characterizing an input query. The system further obtains a reward model trained with a plurality of training examples, each training example including a respective sample query text sequence characterizing a sample query and a respective reference response text sequence, the reference response text sequence including at least a reference final answer to the sample query, and the reward model is configured to process the input including the query text sequence characterizing the query and one or more inference steps generated in response to the query, and calculate a reward score indicating how successful the one or more inference steps are in producing a correct final answer to the query.

[0023] The system uses a language model to generate a plurality of candidate inference processes as responses to an input query, each of which includes a plurality of candidate inference steps including a candidate final answer. The system may select a best inference process from the plurality of candidate inference processes based on one or more reward scores calculated by the trained reward model, and output the best inference process as an output response to the input query.

[0024] This specification also describes computer-implemented methods performed by the above-described systems. This specification further describes one or more computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform the above-described methods.

[0025] The subject matter described herein can be implemented in particular implementations to realize one or more of the following advantages.

[0026] This specification provides techniques for using reward models to guide reinforcement learning of language models to perform inference tasks. Reward models can be trained and applied using either purely output-based or process-based monitoring, but reward models can generate reward signals not only for the output of a task (e.g., a final answer to a query) but also for individual inference steps. By using reward models to provide reward signals for inference steps, the language model improves as it performs an inference task, e.g., by reducing output error and / or inference process error. [Brief explanation of the drawings]

[0027] [Figure 1A] FIG. 1 illustrates an exemplary training system for training a language model to perform an inference task. [Figure 1B] FIG. 1 illustrates the reinforcement learning process performed by the training system. [Figure 2A] FIG. 1 illustrates an example of using a reward model to select an expert reasoning process. [Figure 2B] FIG. 1 illustrates an example of using a reward model to generate an expert reasoning process. [Figure 3] 1 is a flow diagram of an example process for training a language model. [Figure 4] 1 is a flow diagram of an example process for performing an inference task using a trained language model. DETAILED DESCRIPTION OF THE INVENTION

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

[0029] This specification describes techniques for training language models to perform inference tasks. In particular, the language models are trained to respond to queries (e.g., questions) and generate responses that include (e.g., as the final inference step) a sequence of inference steps that lead to a final answer to the query. The described techniques can be used to train language models to perform a variety of inference tasks.

[0030] In one example, the trained language model can be used to diagnose faults in a mechanical system operating in a real-world environment. An input query can include, for example, observations about the mechanical system from one or more sensors and a question about the mechanical system, such as, "Given the measurements from sensors A, B, and C, is the system operating correctly?" The output response can provide a final answer to the question and the reasoning process that led to the final answer.

[0031] In another example, the trained language model can be deployed as part of an educational software program to help users learn or practice one or more corresponding skills. For example, the input query can be a mathematical word problem, and the output response includes a correct answer to the word problem and intermediate steps to arrive at the correct answer. The output response can be given to the student as a proof of how to solve the mathematical problem or can be used as a grading tool to provide feedback on the student's practice.

[0032] In another example, the trained language model can be used for natural language control of tasks in a real-world environment. The input query can be related to the task and include, for example, a request to perform the task. The output response can be used to control, for example, a mechanical system (which may be referred to as a machine agent) or a computer system to perform the task. As an example, the input query can include a high-level question, for example, from a human, to perform the task, such as, "What is the most cost-effective way to manufacture this component with our current equipment?" The output response includes a final answer that specifies a planned procedure and reasoning steps for determining the planned procedure. The output response can be used to control one or more machine agents, for example, a robotic arm, to perform the planned procedure. In some other examples, the reasoning process in the output response can be used to 1) determine, for example, by a user, whether to perform an action specified in the final answer in the output response, or 2) provide a human-interpretable explanation that can be used to later access, when certain criteria are satisfied, to evaluate the performance of the language model or to diagnose the cause of an error that occurs as a result of performing the action specified in the response generated by the language model. In some implementations, the output response can specify an action or path to be taken in the real-world environment by the human user to perform the task, and the system can provide the output response to the human user, for example via a display or audio output device, to guide the user to perform the task in the real-world environment.

[0033] 1A shows an exemplary training system 100 for training a language model 110 to perform an inference task. Training system 100 is an example of a system, implemented as a computer program on one or more computers at one or more locations, in which the systems, components, and techniques described below may be implemented.

[0034] Generally, the language model 110 is trained to respond to a query 120 by generating a response 130 that includes a sequence of reasoning steps that lead to a final answer to the query.

[0035] 1A , language model 110 has a set of language model parameters 115. In some implementations, at each of a plurality of steps, language model 110 can process input including query 120 to generate output 150 that characterizes an inference step 130 in the inference process in response to query 120. Each of query 120 and inference step 130 can include a text token sequence.

[0036] The training system 100 includes a reinforcement learning engine 135 that performs reinforcement learning to update the language model parameters 115 of the language model 110 using a reward model 140. The reward model 140 is configured to process input 150, which includes a query text sequence characterizing a query and one or more inference steps generated in response to the query 120, and to calculate a reward score 160 that indicates how successful the one or more inference steps are in producing a correct final answer to the query.

[0037] The language model 110 may be a neural network configured to process an input and generate an output that includes a probability distribution over a set of text tokens in a vocabulary of text tokens, where the probability of each token represents the likelihood that the text token will immediately follow the input.

[0038] The vocabulary of text tokens may include any suitable tokens that appear in natural language text, such as ASCII characters, words, word fragments, or variously distributed n-grams. For example, the vocabulary of text tokens may be fixed or may have been generated by applying a suitable tokenizer, such as a byte-pair encoding tokenizer or a SentencePiece tokenizer, to a corpus of text.

[0039] For example, language model 110 may be an autoregressive neural network. Language model 110 is referred to as an autoregressive neural network because the neural network autoregressively generates the output sequence of tokens by generating each particular token in the output sequence conditioned on a current input sequence that includes any (e.g., all) tokens that precede the particular text token in the output sequence, i.e., tokens already generated for any previous position in the output sequence that precedes the particular position of the particular token, and a context input that provides a context for the output sequence (a “context sequence”), e.g., a text token sequence representing query 120.

[0040] For example, the current input sequence when generating a token at any given position in the output sequence may include a context sequence and a token at any preceding position that precedes the given position in the output sequence. As a particular example, the current input sequence may include a context sequence followed by a token at any (e.g., all) preceding positions that precede the given position in the output sequence. Optionally, the context and the current output sequence may be separated by one or more predetermined tokens in the current input sequence.

[0041] More specifically, to generate a particular token at a particular position in a candidate output sequence, the neural network 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 text token in a vocabulary of text tokens. The neural network can then use the score distribution to select a text token from the vocabulary as the particular token. For example, the neural network can greedily select the highest-scoring token or can sample the token from the distribution using, for example, kernel sampling or another sampling technique.

[0042] As a particular example, the language model 110 may be or include an autoregressive transformer-based neural network that includes (i) a sequence including multiple attention blocks, each applying a self-attention operation, and (ii) an output subnetwork that processes the output of the last attention block to generate a score distribution.

[0043] The neural network may have any of a variety of transformer-based neural network architectures. Examples of such architectures are 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. Elsen, 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, BA Hechtman, L. Weidinger、I. Gabriel、WS Isaac、E. Lockhart、S. Osindero、L. Rimell、C. Dyer、O. Vinyals、K. Ayoub、J. Stanway、L. Bennett、D. Hassabis、K. Kavukcuoglu、およびG. Irving, Scaling Language Models: Methods, Analysis & Insights from Training Gopher, CoRR, abs / 2112.11446, 2021, Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, 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 chatbot, CoRR, abs / 2001.09977, 2020 Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, and Amanda Askell, “Language Models Are Few-Shot Learners,” arXiv preprint arXiv:2005.14165.

[0044] Generally, however, a Transformer-based neural network includes a sequence of attention blocks, and during 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 the hidden state for at least the last token in the given input sequence, at least in part, by applying self-attention to generate a respective output hidden state for the last token. The input hidden state for the first attention block is the embedding of the input token in the input sequence, and the input hidden state for each subsequent attention block is the output hidden state generated by the preceding attention block.

[0045] In this example, the output sub-network processes the output hidden state produced by the last attention block in the sequence for the last input token in the input sequence to produce a score distribution.

[0046] In some implementations, prior to reinforcement learning, language model 110 can be initialized with a base language model pre-trained using unsupervised learning. That is, language model 110 can start with the same architecture and weights as a base language model trained on a corpus of text and perform one or more language modeling tasks that do not require labeled training examples. For example, system 100 or another training system can pre-train the base language model on a masked language modeling (MLM) task to predict masked portions of input text, a next sentence prediction (NSP) task to predict whether, given two sentences, the second sentence is a successor to the first, and / or an autoregressive pre-training (ARP) task to predict the next token in a sequence of tokens. As a specific example, language model 110 may be pre-trained with a maximum likelihood objective on text, e.g., a large dataset of publicly available text from the Internet or another text corpus.

[0047] In some implementations, prior to reinforcement learning and after initializing the language model 110, the system 100 can fine-tune the language model 110 using supervised fine-tuning (SFT) based on a set of supervised training examples. Each supervised example includes a training input token sequence and a corresponding target output token sequence. During SFT, the system 100 performs supervised learning to update parameters of the language model 110, for example, to maximize the log-likelihood that the language model 110 will output the target output token sequence for the corresponding training input sequence. In some implementations, the supervised training examples include an input query and an example of a corresponding target inference process for the input query. That is, the training input token sequence is a token sequence representing the input query, and the target output token sequence is a token sequence representing the target inference process.

[0048] In some of these implementations, system 100 may further perform contextual learning of language model 110 by providing one or more prompts among the inputs to language model 110. Each prompt may be an example of an input-output pair, where the input is an example of an input query and the output is an example of an output inference process that should be generated in response to the input query.

[0049] The system 100 trains the reward model 140 (i.e., updates the parameters 148 of the reward model 140) using a set of training examples 148. Each training example 145 includes a respective query text sequence characterizing a respective sample query and a respective reference response text sequence. The reference response text sequence includes at least the reference final answer to the respective sample query.

[0050] As described above, reward model 140 is configured to receive input 150 including a query text sequence characterizing a query and one or more inference steps generated in response to the query, process input 150, and calculate a reward score 160 indicating how successful the one or more inference steps are in producing a correct final answer to the query. In some implementations, for each inference step in the inference process generated in response to a query, reward model 140 outputs a score for the inference step. In one particular example, reward model 140 can be trained to generate a score for the inference step to indicate the likelihood that the inference step is “correct” or “incorrect.”

[0051] In some implementations, the reward model 140 is trained to generate a score for each inference step in the inference process that indicates whether the inference process that includes that step led to a correct final answer. The resulting reward model trained using this process is called an output-monitored reward model (ORM) because its training is monitored by a reference final answer.

[0052] In some implementations of training an ORM, the system 100 obtains multiple candidate inference processes for each training example in response to sample queries in the training examples. As described below with reference to FIG. 1B , the system 100 can generate the candidate inference processes using a policy language model. The system 100 assigns a target score to each inference step in each candidate inference process based on the reference response text sequence in the training examples. The target score indicates whether the inference step is correct or incorrect based on whether the corresponding final answer in the candidate inference process matches the reference final answer in the training example. That is, if the candidate inference process produces a reference final answer, the system 100 assigns each inference step in the candidate inference process a target score indicating that the inference step is correct. Otherwise, the system 100 assigns each inference step in the candidate inference process a target score indicating that the inference step is incorrect. The system 100 then trains a reward model 140 to generate a reward score 160 for the inference step in the candidate inference process that matches the corresponding target score for the inference step. For example, the system 100 may use a backpropagation-based supervised learning technique to train the reward model 140.

[0053] In some implementations, the reward model 140 is trained to generate a score for each particular inference step in a candidate inference process to indicate whether the inference steps up to and including that step are correct. The resulting reward model trained using this process is called a process-monitored reward model (PRM) because its training is monitored by a process that includes intermediate inference steps in the inference process.

[0054] In some implementations of training a PRM, the reference response text sequence for each training example includes a reference inference process, which is a sequence of inference steps with a final answer as the last step. The system 100 assigns a target score to the current inference step based on whether the inference steps generated up to that point match the sequence of inference steps in the reference inference process. That is, the target score for the current inference step may be determined based on a comparison (e.g., a similarity score) between each of the inference steps generated up to (and including) the current inference step and the corresponding inference step in the sequence of inference steps in the reference inference process.

[0055] In some implementations of training a PRM, system 100 can use feedback from a human annotator to determine a target score for an inference step in a candidate inference process. For example, system 100 or another system can present (i) sample query and reference response text sequences of training examples and (ii) candidate inference processes obtained for the training examples to the human annotator. System 100 can then receive input from the annotator identifying an initial inference step in the candidate inference process as a significant error, if any. In one example, a significant error is defined as a step where the information represented is incorrect or where it is no longer possible to reach a correct answer without undoing the step. After receiving the input from the annotator, system 100 can use the feedback from the annotator to assign a target score for the candidate inference step.

[0056] In some implementations, reward model 140 can be implemented as a language model configured to process an input sequence including a query and one or more inference steps and generate output tokens indicating the correctness of the inference steps, e.g., “correct” or “incorrect.” Reward model 140 can further output the probability assigned to the “correct” token by the language model as a reward score. In some implementations, similar to language model 110, reward model 140 can (i) be initialized using a pre-trained base language model, (ii) be fine-tuned using SFT, and / or (iii) be subjected to contextual learning through few-shot prompting.

[0057] Once reward model 140 is implemented, reinforcement learning engine 135 performs reward model-based (RM-based) reinforcement learning to update language model 110. During RM-based reinforcement learning, reinforcement learning engine 135 updates a policy that maps inputs, including (i) query 120 and (ii) inference steps 130 generated so far, to outputs that characterize the next inference step, based on the reward score predicted by reward model 140.

[0058] 1B illustrates an example of an RM-based reinforcement learning process implemented by training system 100. In the exemplary process described below, system 100 uses an expert iteration approach for reinforcement learning. Examples of implementing expert iteration include those described in D. Silver et al., "Mastering the game of go without human knowledge," Nature, 550(7676):354-359, 2017, and T. Anthony et al., "Thinking fast and slow with deep learning and tree search," Advances in Neural Information Processing Systems, 30, 2017. Generally, expert iteration alternates between two operations, including (i) policy improvement and (ii) distillation.

[0059] 1B, during policy refinement, system 100 uses policy language model 112 and trained reward model 140 to perform a search procedure to produce expert reasoning process 190. During distillation, system 100 uses expert reasoning process 190 to train language model 110, for example, by using supervised learning.

[0060] The policy language model 112 can be implemented using any suitable language model. In some implementations, the policy language model 112 may be the same model as the language model 110. That is, the language model 110 (e.g., after undergoing pre-training, SFT, and / or contextual learning using few-shot prompting) can be used to generate the expert reasoning process 190. The system 100 can perform expert iterations for multiple iterations. After the language model 110 is trained using the expert reasoning process 190 in a particular iteration, the system 100 can use the language model 110 as the policy language model 112 to generate the expert reasoning process 190 for the next iteration. In some implementations, the system can also use the policy language model 112 to generate candidate reasoning processes for training the reward model 140.

[0061] To generate expert reasoning process 190, system 100 receives training query 170 and uses policy language model 112 and trained reward model 140 to generate expert reasoning process 190 in response to training query 170.

[0062] In some implementations, the system can use the policy language model 112 to generate (e.g., by random sampling) multiple candidate expert reasoning processes 180 in response to a training query 170, use the trained reward model 140 to generate a performance score 185 for each candidate expert reasoning process, and select one or more candidate expert reasoning processes with the highest performance score as the expert reasoning processes 190.

[0063] FIG. 2A illustrates an example of using a reward model 140 to select an expert reasoning process. The system can determine a performance score 185 for a candidate expert reasoning process based on the reward score generated by the reward model for each reasoning step in the candidate expert reasoning process, e.g., by summing the reward scores of corresponding reasoning steps in the candidate expert reasoning process. The system selects the candidate expert reasoning process with the highest performance score as the expert reasoning process 190. In this case, the system can use a reward model 140 trained using output outcome monitoring training, and thus the reward model (output outcome monitoring reward model, or ORM) is configured to generate a reward score for each reasoning step to indicate whether the reasoning step could have been in the reasoning process that led to the correct final answer. A policy that maximizes the ORM score at each step generally maximizes the RM estimated probability at each step of ultimately reaching the correct final answer.

[0064] FIG. 2B illustrates an example of using the policy language model 112 and the reward model 140 to generate an expert reasoning process 190. In this case, after one or more inference steps are generated for the expert reasoning process, the system uses the policy language model 112 to generate multiple next candidate inference steps 182. The system uses the trained reward model to generate a reward score (e.g., 185a, 185b, or 185c) for each of the next candidate inference steps. The system selects the next candidate inference step with the highest reward score as the next step for the expert reasoning process and repeats the process until the next step matches the final answer indicator (e.g., the next inference step contains text corresponding to the final answer) or until the maximum number of steps is reached. In this case, the system can use the reward model 140 trained using process monitoring training, and thus the reward model (process monitoring reward model, or PRM) is configured to generate a reward score for the current inference step to indicate whether the inference steps up to the current inference step were correct. A policy that maximizes the PRM score generally selects each step to maximize the RM estimated probability that the previous step is correct. If the previous steps were correct, this usually means that such a policy minimizes the probability of introducing an error into the current step.

[0065] 3 is a flow diagram of an exemplary process 300 for training a language model to perform an inference task. For convenience, process 300 is described as being performed by a system of one or more computers at one or more locations. For example, a training system suitably programmed in accordance with this specification, such as training system 100 shown in FIG. 1, can perform process 300.

[0066] At 310, the system obtains a plurality of training examples. Each training example includes a respective sample query text sequence characterizing a respective sample query and a respective reference response text sequence including a reference final answer to the respective sample query. The system can obtain the training examples from one or more of a variety of data sources, for example, from a database or a labeled dataset. As an illustrative example, the publicly available GSM8K dataset includes mathematical word problems and natural language solutions to the word mathematical problems.

[0067] At 320, the system trains a reward model with a plurality of training examples. The reward model is configured to receive input including a query text sequence characterizing the query and one or more inference steps generated in response to the query, process the input, and calculate a reward score indicating how successful the one or more inference steps are in producing a correct final answer to the query.

[0068] As described above, the system can train a reward model based on candidate reasoning processes and target scores assigned to reasoning steps within the candidate reasoning processes. The target scores can be assigned based on either (i) output-monitored training, in which the training of the reward model is monitored by reference final answers of training examples, or (ii) process-monitored training, in which the training of the reward model is monitored by the process represented by the reasoning processes.

[0069] At 330, the system trains a language model using the trained reward model. As described above, the system performs RM-based reinforcement learning to update the parameters of the language model. That is, during reinforcement learning, the system uses the reward model to generate rewards. As described above, the reward model can be implemented as a language model configured to process an input sequence including a query and one or more inference steps and generate output tokens indicating the correctness of the inference step, for example, “correct” or “incorrect.” The reward model is used to provide reward signals for individual inference steps during reinforcement learning. This is in contrast to conventional techniques in which a language model is rewarded solely based on the correctness of a final answer during reinforcement learning. As discussed with reference to FIGS. 5A-6 , by using the reward model to provide reward signals for inference steps, the language model improves upon performing an inference task, for example, by reducing output result errors and / or inference process errors.

[0070] Figure 4 is a flow diagram of an exemplary process 400 for performing an inference task using a language model trained using the process described with reference to Figures 1-3. For convenience, process 400 is described as being performed by a system of one or more computers at one or more locations.

[0071] At 410, the system receives an input query, which may be represented by a sequence of text tokens.

[0072] At 420, in response to an input query, the system processes the input query using the trained language model to generate multiple candidate output inference processes. Each candidate output inference process is a text token sequence representing a sequence of inference steps where the last inference step in the inference process is or includes a final answer generated in response to the input query. As described with reference to FIG. 1, because the language model is autoregressive and configured to process a current input sequence to generate a probability distribution of tokens in the vocabulary for the next token in the output, the system can use the same model to generate multiple different candidate output sequences in response to the same query by sampling from the probability distribution.

[0073] At 430, the system selects a best inference process from the candidate output inference processes using a ranking strategy. In some implementations, the system can use a reward model trained using the training process described with reference to FIG. 1 to select the best inference process to output. Specifically, the system can use the trained reward model to calculate a performance score for each candidate output inference process. In some implementations, the system can select the candidate output inference process with the highest performance score as the output inference process. In some other implementations, the system can weight the final answer in each candidate output inference process by its corresponding performance score and identify the final answer with the greatest total weight as the “correct” final answer. That is, the total weight for each final answer can be determined by summing the performance scores for each output inference process that has that final answer, and the “correct” final answer is then identified as the final answer with the greatest total weight. The system can then select from among the candidate output inference processes that lead to the identified “correct” final answer based on the performance scores.

[0074] At 440, the system outputs the best inference process. In some implementations, to ensure the quality of the generated response, the system may choose not to provide an output response when the performance score estimated by the reward model for the best inference process is below a threshold. That is, before outputting the best inference process as the output response, the system may determine whether the score estimated by the reward model for the best inference process is below a threshold, and may output the best inference process as the output response only in response to a determination that the score is not below the threshold.

[0075] This specification uses the term "configured" in the context of systems and computer program components. A system of one or more computers configured to perform a particular operation or action means that the system has installed software, firmware, hardware, or a combination thereof that, when run, causes the system to perform the operation or action. One or more computer programs configured to perform a particular operation or action means that one or more programs contain instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0076] Embodiments of the subject matter and functional operations described herein can be implemented in digital electronic circuitry, in tangibly embodied computer software or firmware, in computer hardware containing the structures disclosed herein and their structural equivalents, or in a combination of one or more of these. Embodiments of the subject matter described herein can be implemented as one or more computer programs, e.g., as one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium for execution by or controlling the operation of a data processing apparatus. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of these. Alternatively or additionally, the program instructions may be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, which is generated to encode information for transmission to an appropriate receiver apparatus for execution by a data processing apparatus.

[0077] 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. An apparatus may be or further include special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). In addition to hardware, an apparatus may optionally include code that creates an execution environment for a computer program, e.g., code comprising processor firmware, a protocol stack, a database management system, an operating system, or any combination of one or more of these.

[0078] A computer program may also be called or described as a program, software, software application, app, module, software module, script, or code, and may be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages, and may be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program may be stored in portions of files that hold 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 cooperating files, e.g., files that store one or more modules, subprograms, or portions of code. A computer program may be deployed to be executed on one computer located at one site or on multiple computers, or may be distributed across multiple sites and interconnected by a data communications network.

[0079] The term "database" is used broadly herein to refer to any collection of data, although the data need not be structured in any particular way, or even at all, and can be stored on storage devices in one or more locations. Thus, for example, an index database may contain multiple collections of data, each of which can be organized and accessed in different ways.

[0080] Similarly, the term "engine" is used broadly herein 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 and installed on one or more computers at one or more locations. In some cases, one or more computers will be dedicated to a particular engine, and in other cases, multiple engines may be installed and running on the same computer or multiple computers.

[0081] The processes and logic flows described herein may be implemented 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 may also be implemented by special purpose logic circuitry, e.g., FPGAs, or ASICs, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0082] A computer suitable for running a computer program may be based on a general-purpose or special-purpose microprocessor, or both, or any other type of central processing unit. Generally, the central processing unit will receive instructions and data from a read-only memory, 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 memory may be supplemented by, or incorporated in, special-purpose logic circuitry. Generally, a computer also includes one or more mass storage devices, e.g., magnetic, magneto-optical, or optical disks, for storing data, or is operatively coupled to the mass storage devices to receive data from, transfer data to, or both. However, a computer need not have such devices. Moreover, a computer may be incorporated in another device, e.g., a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

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

[0084] To enable user interaction, embodiments of the subject matter described herein 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 pointing device, e.g., a mouse or trackball, for allowing the user to provide input to the computer. Other types of devices can also be used to provide user interaction; 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. Furthermore, a computer can interact with a user by sending documents to and receiving documents from a device used by the user, e.g., by sending a web page to a web browser on the user's device in response to a request received from the web browser. A computer can also interact with a user by sending text messages or other forms of messages to a personal device, e.g., a smartphone running a messaging application, and receiving a reply message in return from the user.

[0085] A data processing apparatus for implementing machine learning models may also include special purpose hardware accelerator units, for example for handling common and numerically intensive parts of the machine learning training or production, e.g., inference, workload.

[0086] Machine learning models can be implemented and deployed using machine learning frameworks, for example the TensorFlow framework.

[0087] Embodiments of the subject matter described herein can be implemented in a computing system that includes back-end components, e.g., as data servers, or includes middleware components, e.g., application servers, or includes front-end components, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user may interact with an implementation of the subject matter described herein, 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 local area networks (LANs) and wide area networks (WANs), e.g., the Internet.

[0088] A computing system may include clients and servers. Clients and servers 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 sends data, e.g., HTML pages, to a user device, e.g., to display data to and receive user input from a user interacting with the device acting as a client. Data generated at the user device, e.g., a result of user interaction, may be received from the device at the server.

[0089] While the 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 a particular invention. Also, some features described in this specification in the context of separate embodiments can be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Furthermore, while features may be described above as acting in some combination and initially claimed as such, one or more features from a claimed combination can, in some cases, be deleted from the combination, and the claimed combination may be directed to a subcombination or a variation of a subcombination.

[0090] Similarly, although operations are illustrated in the figures or recited in the claims in a particular order, this should not be understood as requiring such operations to be performed in the particular order illustrated, or sequentially, or that all of the illustrated operations be performed to achieve desirable results. In some situations, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the above-described embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems generally can be integrated together in a single software product or packaged in multiple software products.

[0091] Specific embodiments of the present 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. By way of 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. [Explanation of symbols]

[0092] 100 Training System, System 110 language models 112 Policy Language Model 115 Language Model Parameters 120 queries 130 Inference Steps 135 Reinforcement Learning Engine 140 Reward Model 145 Reward Model Parameters 148 training examples 150 queries + inference steps 160 Reward Score 170 training queries 180 Candidate Expert Reasoning Process 185 performance score 185a, b, c Reward Score 190 Expert Reasoning Process 290a, b selection 300 processes 400 processes

Claims

1. 1. A computer-implemented method for training a language model to perform an inference task, comprising: obtaining a plurality of training examples, each training example including a respective sample query text sequence characterizing a respective sample query and a respective reference response text sequence including a reference final answer to the respective sample query; training a reward model with the plurality of training examples, the reward model being configured to receive input including a query text sequence characterizing a query and one or more inference steps generated in response to the query, process the input, and calculate a reward score indicating how successful the one or more inference steps are in producing a correct final answer to the query; training the language model using the trained reward model; A method comprising:

2. training the language model using the trained reward model, receiving a training query; using the policy language model and the trained reward model to generate one or more expert reasoning processes in response to the training queries, each respective expert reasoning process including a respective plurality of reasoning steps; using the one or more expert reasoning processes to train the language model; 2. The method of claim 1, comprising:

3. Training the reward model on the plurality of training examples comprises, for each of one or more of the training examples: obtaining a plurality of candidate reasoning processes in response to the sample queries of the training examples, each candidate reasoning process including one or more reasoning steps; assigning a target score to each inference step in each candidate inference process based on the reference response text sequences in the training examples, the target score indicating whether the inference step is correct or incorrect; training the reward model to generate reward scores for the inference steps in the candidate inference processes that match the corresponding target scores for the inference steps; 3. The method of claim 1 or claim 2, comprising:

4. The step of assigning the target score to each inference step in the candidate inference process comprises: determining whether the candidate reasoning process produces the reference final answer among the training examples; responsive to determining that the candidate reasoning process produced the reference final answer in the training examples, assigning to each reasoning step in the candidate reasoning process a target score indicating that the reasoning step is correct; in response to determining that the candidate reasoning process does not produce the final answer in the training examples, assigning to each reasoning step in the candidate reasoning process a target score indicating that the reasoning step is incorrect; 4. The method of claim 3, comprising:

5. For each training example, the reference response text sequence includes a reference inference process including a sequence of inference steps including the final answer as a final inference step, and the step of assigning the target score to each inference step in the candidate inference process includes:

4. The method of claim 3, further comprising assigning the target score for the current inference step based on whether the inference steps generated up to the current inference step match the sequence of the inference steps in the reference inference process.

6. 6. A method as claimed in any one of claims 3 to 5 and dependent on claim 2, wherein the policy language model is configured to receive input including a query text sequence characterising a query and one or more inference steps generated in response to the query, and to process the input to generate a next inference step.

7. The step of generating an expert process comprises: using the policy language model to generate a plurality of candidate expert reasoning processes in response to the training queries; using the reward model to generate a performance score for each of the candidate expert reasoning processes based on a respective reward score generated for each of one or more of the reasoning steps in the candidate expert reasoning processes; selecting one or more candidate expert reasoning processes having the highest performance scores as said expert reasoning processes; 7. The method of claim 6, comprising:

8. The step of generating one of the expert processes comprises: after one or more inference steps have been generated for the expert reasoning process, using the policy language model to generate a plurality of next candidate inference steps; using the reward model to generate a reward score for each of the next candidate inference steps; selecting the next candidate inference step with the highest reward score as the next step for the expert inference process; repeating the using and selecting steps until the next step matches a final answer indicator or a maximum number of steps has been reached.

7. The method of claim 6, comprising:

9. The method of claim 1 , further comprising the step of obtaining a base language model.

10. 9. A method as recited in claim 8 and dependent on claim 2, further comprising updating the base language model to generate the policy language model.

11. The step of updating the base language model to generate the policy model includes: The method of claim 10 , comprising performing supervised fine-tuning of the base language model on supervised training examples.

12. The step of training the language model comprises:

12. The method of claim 9, comprising initializing the language model based on the base language model.

13. Training the reward model comprises:

13. The method of claim 9, comprising initializing the reward model based on the base language model.

14. 14. The method of claim 1, further comprising updating the reward model using the trained language model.

15. 15. The method of claim 1, further comprising using the trained language model to generate an output response to an input query about a real-world environment, the output response providing information about the real-world environment or specifying an action or path to be taken in the real-world environment.

16. 16. The method of claim 15, wherein the input query relates to a task in the real-world environment, the method further comprising using the generated output response to control one or more machine agents or computer systems acting in the real-world environment to perform the task.

17. 16. The method of claim 15, wherein the input query relates to a task in the real-world environment and the output response specifies an action or path to be taken in the real-world environment, the method further comprising providing the output response to the user to guide the user to perform the task in the real-world environment.

18. 16. The method of claim 1, further comprising using the trained language model to generate output responses to input queries comprising observations about a mechanical system operating in a real-world environment, the generated output responses being for diagnosing faults in the mechanical system.

19. 1. A computer-implemented method for performing an inference task using a language model, comprising: obtaining an input text sequence characterizing an input query; obtaining a reward model trained on a plurality of training examples, each training example including a respective sample query text sequence characterizing a sample query and a respective reference response text sequence, the reference response text sequence including at least a reference final answer to the sample query, the reward model being configured to process input including the query text sequence characterizing a query and one or more inference steps generated in response to the query, and to calculate a reward score indicating how successful the one or more inference steps are in producing a correct final answer to the query; using the language model to generate a plurality of candidate inference processes as responses to the input query, each candidate inference process including a respective plurality of candidate inference steps including a candidate final answer; selecting a best reasoning process from the plurality of candidate reasoning processes based on one or more reward scores calculated by the trained reward model; and outputting the best inference process as an output response to the input query.

20. selecting a best reasoning process from the plurality of candidate reasoning processes based on one or more scores calculated by the trained reward model, using the reward model to generate a respective weight for each of the plurality of candidate reasoning processes, the respective weight measuring the estimated probability of correctness of the respective candidate reasoning process; selecting an optimal final answer as the candidate final answer having the greatest total weight; selecting the candidate inference process having the highest weight estimated by the reward model as the best inference process among the candidate inference processes that produce the optimal final answer; 20. The method of claim 19, comprising:

21. 21. The method of claim 19 or claim 20, wherein the language model is trained using a method according to any one of claims 1 to 18.

22. 22. The method of any one of claims 19 to 21, wherein the input text sequence specifies a mathematical word problem, and the output response specifies a solution to the mathematical problem and reasoning steps for solving the mathematical word problem.

23. determining whether the score estimated by the reward model for the best inference process is below a threshold before outputting the best inference process as the output response; and outputting the best inference process as the output response only in response to determining that the score is not below the threshold.

24. 24. The method of any one of claims 19 to 23, wherein the input query relates to a real-world environment and the output response provides information about the real-world environment or specifies an action or path to be taken in the real-world environment.

25. 25. The method of claim 24, wherein the input query relates to a task in a real-world environment, the method further comprising using the output response to control one or more machine agents or computer systems in the real-world environment to perform the task.

26. 25. The method of claim 24, wherein the input query relates to a task in the real-world environment and the output response specifies an action or path to be taken in the real-world environment, the method further comprising providing the output response to the user to guide the user to perform the task in the real-world environment.

27. 25. The method of any one of claims 19 to 24, wherein the input query comprises an observation about a mechanical system operating in a real-world environment, and the output response is for diagnosing a fault in the mechanical system.

28. 28. The method of claim 1, wherein the language model comprises a neural network.

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

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

Images