Reward-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 the benefit of priority of U.S. Provisional Patent Application No. 63 / 377,532, filed on September 28, 2022, the entire disclosure of which is incorporated herein by reference.

[0002] This specification relates to performing inference tasks using a machine learning model such as a neural network.

Background Art

[0003] A neural network is a machine learning model that utilizes one or more layers of non - linear 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 the received input according to the current values of each set of parameters.

[0004] The parameters of a machine learning model can be determined 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, the weights and bias coefficients of the network layers of the neural network.

Prior Art Documents

Non - Patent Documents

[0005]

Non - Patent Document 1

[0006] This specification describes methods, computer systems, and apparatus, including computer programs encoded on a computer storage medium, 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 comprising a sequence of one or more inference steps to produce the final answer to the query.

[0007] Language models and language-generating neural networks can perform a variety of tasks, including translation tasks (if the training corpus contains words in different languages) and arithmetic tasks, among many other tasks. In relation to reasoning tasks, language models can be used to generate responses to queries (e.g., questions). Responses to queries can be given as text sequences in the form of yes / no binary answers, answers providing one or more numerical values, answers identifying one or more objects or events, or answers identifying a planned path or action. Responses to queries can be used in many ways by humans or computer systems. For example, the answers and reasoning processes may be useful in themselves and may be used to give warnings and / or to control the movement of one or more robotic agents or objects, such as autonomous vehicles.

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

[0009] In one example, the system can be used to diagnose failures in mechanical systems operating in a real-world environment. Input queries may include observations about the mechanical system from one or more sensors (e.g., a camera, microphone, accelerometer, temperature sensor, etc.) and questions about the mechanical system. In some implementations, observations may include detected electronic signals, such as motor current or temperature signals, and / or image or video data from, for example, a camera or LiDAR sensor, data from a sensor on a mechanical agent, or data from a sensor placed separately from the mechanical agent in the environment. The observations and questions can be proposed to generate a linguistic expression of the query. For example, the query may include general questions such as "Given the measurements from sensors A, B, and C, is the system functioning correctly?" or "Given the measurements from A, B, and C, what is wrong with the system?" or specific requests such as "Is component X faulty?". The output response may provide a final answer to the question and the reasoning process leading to that final answer.

[0010] In another example, the environment may be an educational environment, and for instance, the system could be deployed as part of an educational software program that helps users learn or practice one or more corresponding skills. For example, the input query may be a mathematical word problem, and the output response may include the correct answer to the word problem and the intermediate steps to arrive at the correct answer. The output response could be given to the student as proof of how to solve the mathematical problem, or it could be used as a scoring 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 about a task and may include, for example, a request to perform the task. The output response may be used to control, for example, a mechanical system (which may be called a mechanical agent) or a computer system to perform the task. For example, the input query may include a high-level question from, for example, a human, to perform the task, such as, "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 the planned procedure and the reasoning steps for determining the planned procedure. The output response can be used to control one or more mechanical agents, for example, a robotic arm, to perform the planned procedure. For example, the output response may include one or more control signals to control that one or more mechanical agents, such as position, velocity, or force / torque / acceleration data for one or more joints of a part of the robot or another mechanical agent. In some implementations, the input query may specify observations of a real-world environment, often relating to a real-world task, and the output response may specify actions or paths that the user should take in the real-world environment 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. Output responses generated using the methods described herein can greatly improve the usefulness and safety of machine agents, computer systems, or human users in performing real-world tasks in a real-world environment, especially for tasks where generating an inaccurate response can have significant negative consequences. In particular, the output response can be used to perform the task in principle and logically.

[0012] In some implementations, the reasoning process within the output response provides a human-interpretable explanation that can be accessed later when certain criteria are met, such as 1) whether the user should perform an action specified in the final response within the output response, or 2) to evaluate the performance of the language model or to diagnose the cause of an error that occurs as a result of performing an action specified in the 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 a query to a response in a causal chain. Such a human-interpretable explanation can overcome or mitigate some of the drawbacks of using language models (especially neural network-based language models) as a "black box" by allowing human users to understand the reasons behind the responses given by the language model.

[0013] In one particular embodiment, this specification describes a system implemented as a computer program on one or more computers in one or more locations for training a language model for performing inference tasks.

[0014] The system includes a reward model configured to receive an input consisting of 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 score indicating how successful one or more inference steps are in producing the correct final answer to the query.

[0015] The system can train its reward model using either output-monitoring or process-monitoring training with multiple training examples. Generally, each training example includes a sample query text sequence characterizing each sample query (i.e., training query) and a reference response text sequence containing the final reference ("ground truth") answer to each sample query. For each training example, the system can generate one or more candidate inference processes in response to the training example's sample query. Each candidate inference process includes one or more inference steps. Based on the reference response text sequences in the training example, the system can assign a target score to each inference step in each candidate inference process, indicating whether the inference step is correct or incorrect. The system then trains its 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.

[0016] In output-monitoring training, the system determines whether a candidate reasoning process generates a reference final answer within the training examples. If a candidate reasoning process generates a reference final answer within the training examples, the system assigns a reward score to each reasoning step in the candidate reasoning process to indicate that the reasoning step is correct. If a candidate reasoning process does not generate a final answer within the training examples, the system assigns a reward score to each reasoning step in the candidate reasoning process to indicate that the reasoning step is incorrect.

[0017] In process monitoring training, a reference response text sequence includes a reference inference process which includes a sequence of inference steps, with the final answer being the last inference step. The system assigns a target score to the current inference step based on whether the inference steps generated up to the current inference step 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., 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.

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

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

[0020] In some implementations, to generate an expert process, after one or more inference steps are generated for an expert inference process, the system uses a policy language model to generate a plurality of next candidate inference steps. The system uses a 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 inference 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 is reached. The system can then use the resulting inference process as the expert inference process.

[0021] In another aspect, the present specification describes a system implemented as a computer program on one or more computers in one or more locations that uses a language model to perform an inference task.

[0022] The system can obtain an input text sequence characterizing an input query. The system further obtains 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 for the sample query, the reward model being configured to process an input including a query text sequence characterizing a query and one or more inference steps generated in response to the query to calculate a reward score indicating how successful the one or more inference steps are in producing a correct final answer for the query.

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

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

[0025] The subject matter described in this specification may be implemented in a particular implementation form to realize one or more of the following advantages.

[0026] This specification provides techniques for using a reward model to induce reinforcement learning of a language model to perform an inference task. The reward model can be trained and applied using either purely output-result-based monitoring or process-based monitoring, but the reward model is capable of generating a reward signal not only for the output result of performing the task (e.g., the final answer to a query), but also for individual inference steps. By using the reward model to give a reward signal to the inference steps, the language model is improved when performing the inference task, for example, by reducing output result errors and / or inference process errors.

Brief Description of the Drawings

[0027] [Figure 1A] It is a diagram showing an exemplary training system for training a language model to perform an inference task. [Figure 1B] It is a diagram showing a reinforcement learning process implemented by a training system. [Figure 2A] This figure shows an example of using a reward model to select an expert reasoning process. [Figure 2B] This figure shows an example of using a reward model to generate an expert reasoning process. [Figure 3] This is a flowchart illustrating an exemplary process for training a language model. [Figure 4] This is an illustrative flowchart of the process for performing an inference task using a pre-trained language model. [Modes for carrying out the invention]

[0028] Similar reference numbers and names in various drawings refer to the same elements.

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

[0030] For example, a trained language model can be used to diagnose malfunctions in a machine system operating in a real-world environment. Input queries might include observations of the machine system from one or more sensors, and questions about the machine system, such as, "Given the measurements from sensors A, B, and C, is the system functioning correctly?" Output responses could provide a final answer to the question and the reasoning process leading to that final answer.

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

[0032] In another example, a trained language model can be used for natural language control of tasks in a real-world environment. Input queries can be about the task and may include, for example, requests for performing the task. Output responses can be used to control, for example, a machine system (which may be called a machine agent) or a computer system to perform the task. As an example, an input query may include a high-level question from, for example, a human, about performing the task, such as, for example, "What is the most cost-effective way to manufacture this component with our current equipment?" 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. In some other examples, the reasoning process in the output response can be used to provide a human-interpretable explanation that can be accessed later when certain criteria are met, such as 1) whether, for example, a user should perform the action specified in the final answer in the output response, or 2) 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 may specify the actions or paths that a human user should take in the real-world environment to perform the task, and the system may provide the human user with the output response to guide them to perform the task in the real-world environment, for example, via a display or audio output device.

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

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

[0035] Referring again to Figure 1A, the language model 110 has a set of language model parameters 115. In some implementations, in each of several steps, the language model 110 may process an input containing a query 120 to generate an output 150 that characterizes an inference step 130 in the inference process in response to the query 120. Each of the query 120 and the inference step 130 may contain 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 an 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 compute a reward score 160 which indicates how successful one or more inference steps are in producing the correct final answer to the query.

[0037] The language model 110 may also be a neural network configured to process an input and generate an output containing a probability distribution across 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 appropriate 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 can be fixed, or it can be generated by applying an appropriate tokenizer, such as a byte-pair encoded tokenizer or a SentencePiece tokenizer, to a corpus of text.

[0039] For example, the language model 110 may be an autoregressive neural network. The language model 110 is called an autoregressive neural network because it autoregressively generates an output sequence of tokens by generating each particular token in an output sequence, given a current input sequence containing tokens already generated for any previous position in the output sequence that precedes any (e.g., all) tokens in the output sequence that precede any particular position in the output sequence that precede any particular token, and a context input that provides a context for the output sequence ("context sequence"), for example, a sequence of text tokens representing query 120.

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

[0041] More specifically, in order to generate a specific 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, which assigns a score, e.g., a probability, to each text token in the vocabulary of text tokens. The neural network can then use the score distribution to select a text token from the vocabulary as the specific token. For example, the neural network could greedily select the highest-scoring token, or it could sample tokens from the distribution using, for example, nuclear sampling or another sampling technique.

[0042] As a specific example, the language model 110 may be or include an autoregressive transformer-based neural network comprising (i) a sequence of attentional blocks, each of which applies an autoattentional action, and (ii) an output subnetwork that processes the output of the last attentional 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] However, generally speaking, transformer-based neural networks consist of a sequence of attention blocks, and during the processing of a given input sequence, each attention block in the sequence receives its respective input hidden state for each input token in the given input sequence. The attention block then updates, at least, its hidden state for the last token in the given input sequence, at least in part, by applying self-attention to generate its respective output hidden state for the last token. The input hidden state for the first attention block is the embedding of the input tokens 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 subnetwork processes the output hidden state generated by the last attention block in the sequence for the last input token in the input sequence to generate a score distribution.

[0046] In some implementations, prior to reinforcement learning, the language model 110 can be initialized using a base language model pre-trained with unsupervised learning. That is, the language model 110 can start with the same architecture and weights as the 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 a masked portion of an input text, Next Sentence Prediction (NSP) task to predict whether a second sentence is a successor to the first given two sentences, and / or Autoregressive Pre-training (ARP) task to predict the next token in a sequence given a sequence of tokens. Specifically, the language model 110 may be pre-trained with a maximum likelihood goal on a large dataset of publicly available text, for example, 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 the SFT, the system 100 performs supervised learning to update the parameters of the language model 110, for example, to maximize the log-likelihood at which the language model 110 outputs 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 the 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 can further perform contextual learning of the language model 110 by providing one or more prompts in the input to the 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] System 100 trains the reward model 140 using a set of training examples 148 (i.e., updates the parameters 148 of the reward model 140). Each training example 145 includes a query text sequence that characterizes the respective sample query, and a reference response text sequence. The reference response text sequence includes at least the final reference response for the respective sample query.

[0050] As described above, the reward model 140 is configured to receive an input 150 which includes a query text sequence characterizing a query and one or more inference steps generated in response to the query, and to process the input 150 to compute a reward score 160 indicating how successful one or more inference steps are in producing the correct final answer to the query. In some implementations, for each inference step in the inference process generated in response to the query, the reward model 140 outputs a score for the inference step. In one particular example, the reward model 140 can be trained to generate scores for inference steps 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 to indicate whether the inference process, including that step, leads to the 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 the reference final answer.

[0052] In some implementations of training the ORM, system 100 obtains multiple candidate inference processes for each training example in response to sample queries of the training example. As described below with reference to Figure 1B, system 100 can generate candidate inference processes using a policy language model. 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 example. 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 the reference final answer, system 100 assigns a target score to each inference step in the candidate inference process indicating that the inference step is correct. Otherwise, system 100 assigns a target score to each inference step in the candidate inference process indicating that the inference step is incorrect. System 100 then trains a reward model 140 to generate a reward score 160 for the inference steps in the candidate inference process that matches the corresponding target score for the inference step. For example, system 100 can 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 scores for each specific inference step in the candidate inference process, indicating 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 the intermediate inference steps of the inference process.

[0054] In some implementations for training PRM, the reference response text sequence for each training example includes a reference inference process, which is a sequence of inference steps with the final answer as the last step. 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., 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 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 a human annotator with (i) a sample query and reference response text sequence for a training example, and (ii) a candidate inference process obtained for the training example. System 100 can then receive input from the annotator identifying the first inference step in the candidate inference process as a critical error, if any. In one example, a critical error is defined as a step where the information represented is incorrect, or where it is no longer possible to reach the correct answer without undoing that step. After receiving 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, the 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 to generate output tokens indicating the correctness of the inference steps, for example, "correct" or "incorrect". The 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 the language model 110, the reward model 140 can be (i) initialized with a pre-trained base language model, (ii) fine-tuned using SFT, and / or (iii) undergo contextual learning through fusion shot prompting.

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

[0058] Figure 1B shows an example of an RM-based reinforcement learning process performed by training system 100. In the exemplary process described below, system 100 uses an expert iteration method 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: (i) policy improvement and (ii) distillation.

[0059] In the example shown in Figure 1B, during policy refinement, system 100 uses the policy language model 112 and the trained reward model 140 to perform exploration steps to produce the expert inference process 190. During distillation, system 100 uses the expert inference process 190 to train the 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 can be used to generate the expert inference process 190 (e.g., after undergoing contextual learning using pre-training, SFT, and / or fusion shot prompting). The system 100 can perform expert iterations for multiple iterations. After the language model 110 has been trained using the expert inference 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 inference process 190 for the next iteration. In some implementations, the system can also use the policy language model 112 to generate candidate inference processes for training the reward model 140.

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

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

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

[0064] Figure 2B shows an example of using a policy language model 112 and a reward model 140 to generate an expert inference process 190. In this case, after one or more inference steps have been generated for the expert inference process, the system uses the policy language model 112 to generate multiple next candidate inference steps 182. The system uses a 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 inference process and repeats the process until the next step matches the final answer indicator (e.g., until the next inference step contains the text corresponding to the final answer) or until the maximum number of steps is reached. In this case, the system can use a reward model 140 trained with process monitoring training, and thus the reward model (process monitoring reward model, i.e., 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 are correct. A policy that maximizes the PRM score generally selects each step in a way that maximizes the RM estimated probability that the steps so far are correct. If the steps taken so far have been correct, this usually means that such a policy minimizes the probability of making a mistake in the current step.

[0065] Figure 3 is a flowchart of an exemplary process 300 for training a language model to perform an inference task. For convenience, the process 300 is described as being performed by a system consisting of one or more computers in one or more locations. For example, a training system appropriately programmed according to this specification, such as the training system 100 shown in Figure 1, can perform the process 300.

[0066] In 310, the system acquires multiple training examples. Each training example includes a sample query text sequence characterizing each sample query, and a reference response text sequence containing the reference final answer for each sample query. The system can acquire training examples from one or more of various data sources, for example, from a database or a labeled dataset. As an example for explanation, the publicly available GSM8K dataset includes mathematical word problems and natural language solutions to these word-mathematical problems.

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

[0068] As described above, the system can train a reward model based on candidate inference processes and target scores assigned to the inference steps within those processes. The target scores can be assigned based on (i) output-monitoring training, where the training of the reward model is monitored by the reference final answer of a training example, or (ii) process-monitoring training, where the training of the reward model is monitored by the process represented by the inference process.

[0069] In 330, the system trains a language model using a pre-trained reward model. As described above, the system updates the parameters of the language model by performing RM-based reinforcement learning. 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 containing a query and one or more inference steps, and to generate output tokens that indicate the correctness of the inference steps, for example, "correct" or "incorrect". The reward model is used to provide reward signals to individual inference steps during reinforcement learning. This is in contrast to conventional techniques where the language model is rewarded only based on the correctness of the final answer during reinforcement learning. As discussed with reference to Figures 5A-6, by using the reward model to provide reward signals to inference steps, the language model is improved when performing inference tasks, for example, by reducing output result errors and / or inference process errors.

[0070] Figure 4 is a flowchart 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 consisting of one or more computers in one or more locations.

[0071] In step 410, the system receives an input query. The input query can be represented by a text token sequence.

[0072] In 420, the system processes an input query using a trained language model to generate multiple candidate output inference processes in response to the input query. Each candidate output inference process is a text token sequence representing a sequence of inference steps, where the final inference step in the process is either the final answer generated in response to the input query, or the final answer. As described with reference to Figure 1, the language model is autoregressive and configured to process the current input sequence to generate a probability distribution of tokens in the vocabulary for the next token in the output, so 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] In step 430, the system uses a ranking strategy to select the best inference process from the candidate output inference processes. In some implementations, the system can use a reward model trained using the training process described with reference to Figure 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 with its corresponding performance score and identify the final answer with the maximum 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 maximum total weight. The system can then select from the candidate output inference processes that lead to the identified "correct" final answer based on the performance score.

[0074] In step 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 output a response if the performance score estimated by the reward model for the best inference process falls below a threshold. That is, before outputting the best inference process as an output response, the system determines whether the score estimated by the reward model for the best inference process falls below a threshold, and can only output the best inference process as an output response if it determines that the score does not fall below the threshold.

[0075] This specification uses the term “configured” in relation to system and computer program components. A system consisting of one or more computers being configured to perform a particular operation or action means that the system has installed software, firmware, hardware, or a combination thereof that causes the system to perform the operation or action when it is running. A computer program being configured to perform a particular operation or action means that the program contains instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0076] The subject matter and functional operating embodiments described herein can be implemented in digital electronic circuit mechanisms, tangibly implemented computer software or firmware, computer hardware including the structures disclosed herein and their structural equivalents, or in combination of one or more of these. Embodiments of the subject matter described herein can be implemented as one or more computer programs, for example, as one or more modules of computer program instructions encoded on a tangible non-temporary storage medium for execution by a data processing device or for controlling the operation of a data processing device. The computer storage medium may be a machine-readable storage device, a machine-readable storage board, 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 artificially generated propagating signals, such as mechanically generated electrical, optical, or electromagnetic signals, which are generated to encode information for transmission to a suitable receiver device for execution by a data processing device.

[0077] The term "data processing device" refers to data processing hardware and encompasses all kinds of devices, machines, and equipment for processing data, including, for example, a programmable processor, a single computer, or multiple processors or computers. A device may be a special-purpose logic circuit mechanism, such as an FPGA (Field-Programmable Gate Array) or ASIC (Application-Specific Integrated Circuit), or may further include these. Optionally, in addition to hardware, a device may include code that creates an execution environment for computer programs, such as processor firmware, a protocol stack, a database management system, an operating system, or code constituting one or more of these, or a combination thereof.

[0078] Computer programs may be called, or described as, programs, software, software applications, apps, modules, software modules, scripts, 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 standalone programs or as modules, components, subroutines, or other units suitable for use in a computing environment. A program may, but may not, correspond to a file in a file system. A program may be stored in a file portion that holds one or more scripts stored in other programs or data, such as a markup language document, in a single file dedicated to the program in question, or in multiple collaborative files, such as files that store one or more modules, subprograms, or parts of code. Computer programs may be deployed to run on one computer or on multiple computers located in one place, or they may be distributed across multiple locations and interconnected by data communication networks.

[0079] In this specification, the term “database” is used broadly to refer to any collection of data, but data does not need to be structured in any particular way, or does not need to be structured 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, in this specification, the term “engine” is used broadly to refer to a software-based system, subsystem, or process 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 in one or more locations. In some cases, one or more computers will be dedicated to a particular engine, while in other cases, multiple engines may be installed and running on the same computer or on multiple computers.

[0081] The processes and logic flows described herein may be implemented by one or more programmable computers running one or more computer programs to perform their functions by acting on input data and generating outputs. The processes and logic flows may be implemented by special-purpose logic circuit mechanisms, such as FPGAs or ASICs, or by a combination of special-purpose logic circuit mechanisms and one or more programmed computers.

[0082] A computer suitable for running computer programs 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 read-only memory or random-access memory, or both. Essential elements of a computer are a central processing unit for executing or running instructions, and one or more memory devices for storing instructions and data. The central processing unit and memory can be complemented by or incorporated into special-purpose logic circuit mechanisms. Generally, a computer may include one or more mass storage devices for storing data, such as magnetic, magneto-optical disks, or optical disks, or be operably coupled to mass storage devices to receive data from or transfer data, or both. However, a computer is not required to have such devices. Moreover, a computer can be incorporated into another device, for example, a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, such as 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, for example, semiconductor memory devices such as EPROM, EEPROM, and flash memory devices, magnetic disks such as internal hard disks or removable disks, magneto-optical disks, and CD-ROM and DVD-ROM disks, and all forms of non-volatile memory, media, and memory devices.

[0084] To enable interaction with the user, embodiments of the subject matter described herein can be implemented on a computer having a display device for displaying information to the user, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, and a keyboard and pointing device for the user to provide input to the computer, such as a mouse or trackball. Interaction with the user can also be provided using other types of devices, for example, the feedback provided to the user can be any form of sensory feedback, such as visual feedback, auditory feedback, or haptic feedback, and input from the user can be received in any form, including acoustic, voice, or haptic input. Furthermore, the computer can interact with the user by sending documents to and receiving documents from a device used by the user, for example, by sending a web page to a web browser on the user's device in response to a request received from a web browser. The computer can also interact with the user by sending text messages or other forms of messages to a personal device, such as a smartphone running a messaging application, and receiving response messages from the user in return.

[0085] Data processing devices for implementing machine learning models may also include, for example, special-purpose hardware accelerator units for processing the common and numerical computation parts of machine learning training or production, such as inference and workload.

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

[0087] Embodiments of the subject matter described herein can be implemented in a computing system that includes, for example, a data server as a backend component, or a middleware component, for example, an application server, or a client computer having a frontend component, for example, a graphical user interface, a web browser, or an application for a user to interact with the implementation of the subject matter described herein, or any combination of one or more such backend, middleware, or frontend components. The components of the system can be interconnected by digital data communications of any form or medium, such as communication networks. Examples of communication networks include local area networks (LANs) and wide area networks (WANs), such as the Internet.

[0088] A computing system may include a client and a server. Clients and servers are generally separate from each other and typically interact through a communication network. The client-server relationship arises from computer programs running on each computer and having a client-server relationship with each other. In some embodiments, the server sends data, such as an HTML page, to a user device to display data to a user interacting with a device acting as a client and to receive user input from the user. Data generated on the user device, such as the results of user interaction, may be received from the device to the server.

[0089] This specification includes many specific implementation details, but these should not be interpreted as limitations on the scope of any invention or the scope of what can be claimed, but rather as descriptions of features that may be specific to a particular embodiment of a particular invention. Furthermore, some features described herein 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 separately or in any suitable partial combination in multiple embodiments. Moreover, features are described above as acting in several combinations, and may be initially claimed as such, but one or more features from a claimed combination may, in some cases, be removed from that combination, and the claimed combination may be a partial combination or a variation of a partial combination.

[0090] Similarly, while operations are shown in the drawings and described in the claims in a specific order, this should not be understood as requiring such operations to be performed in a specific illustrated order or sequentially, or that all described operations must be performed to achieve a desired result. In some situations, multitasking and parallel processing may be advantageous. Furthermore, 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 in multiple software products.

[0091] This document describes specific embodiments of the subject matter. Other embodiments are within the scope of the following claims. For example, the actions described in the claims may be performed in a different order and still achieve the desired results. As an example, the process shown in the accompanying drawings does not necessarily require a specific illustrated order or sequence to achieve the desired results. In some cases, multitasking and parallel processing may be advantageous. [Explanation of Symbols]

[0092] 100 training systems, systems 110 Language Models 112 Policy Language Models 115 Language Model Parameters 120 queries 130 inference steps 135 Reinforcement Learning Engine 140 Reward Models 145 Reward Model Parameters 148 Training Examples 150 queries + inference steps 160 Reward Score 170 training queries 180 Candidate Expert Inference Processes 185 Performance Score 185a, b, c Reward score 190 Expert Reasoning Process 290a, b Select 300 processes 400 processes

Claims

1. A computer implementation method for training a language model to perform an inference task, A step of obtaining multiple training examples, each training example comprising a sample query text sequence characterizing each sample query, and a reference response text sequence containing the final reference answer for each sample query, A step of training a reward model with the aforementioned plurality of training examples, wherein the reward model is configured to receive an input comprising a query text sequence characterizing a query and one or more inference steps generated in response to the query, and to process the input to calculate a reward score indicating how successful the one or more inference steps are in producing the correct final answer to the query, The steps include training the language model using the pre-trained reward model and Methods that include...

2. The step of training the language model using the pre-trained reward model is: Steps to receive training queries, A step of using the policy language model and the trained reward model to generate one or more expert inference processes in response to the training query, wherein each expert inference process includes each of the multiple inference steps, The steps include using one or more expert reasoning processes to train the language model. The method according to claim 1, including the method described in claim 1.

3. The step of training the reward model with the aforementioned multiple training examples is performed for each of one or more of the aforementioned training examples. A step of obtaining a plurality of candidate inference processes in response to the sample query of the training example, wherein each candidate inference process includes one or more inference steps, A step of assigning a target score to each inference step in each candidate inference process based on the reference response text sequence in the training example, wherein the target score indicates whether the inference step is correct or incorrect. A step of training the reward model to generate a reward score for the inference step in the candidate inference process that matches the corresponding target score for the inference step; The method according to claim 1, including the method described in claim 1.

4. The step of assigning the target score to each inference step in the candidate inference process is: A step of determining whether the candidate reasoning process generates the reference final answer in the training example, A step in which, in response to the judgment that the candidate reasoning process produces the reference final answer in the training example, a target score is assigned to each reasoning step in the candidate reasoning process to indicate that the reasoning step is correct, In response to the determination that the candidate reasoning process does not produce the final answer in the training example, a step of assigning a target score to each reasoning step in the candidate reasoning process to indicate that the reasoning step is incorrect. The method according to claim 3, including the method described in claim 3.

5. For each training example, the reference response text sequence includes a reference inference process which includes a sequence of inference steps which include the final response as the last inference step, and the step which assigns the target score to each inference step in the candidate inference process is: The method according to claim 3, further comprising the step of assigning the target score for the current inference step based on whether the inference step generated up to the current inference step matches the sequence of inference steps in the reference inference process.

6. The method according to claim 2, wherein the policy language model is configured to receive an input comprising a query text sequence characterizing a query and one or more inference steps generated in response to the query, and to process the input to generate the next inference step.

7. The step of generating an expert process is: The steps include using the policy language model to generate multiple candidate expert inference processes in response to the training query, A step of using the reward model to generate a performance score for each of the candidate expert inference processes based on the respective reward scores generated for each of the inference steps in the candidate expert inference process, The step of selecting one or more candidate expert inference processes having the highest performance score as the expert inference process, The method according to claim 6, including the method described in claim 6.

8. The step of generating one of the expert processes is, After one or more inference steps have been generated for the expert inference process, the policy language model is used to generate a plurality of next candidate inference steps, A step of using the reward model to generate a reward score for each of the next candidate inference steps, A step of selecting the next candidate inference step with the highest reward score as the next step for the expert inference process, The steps described above are repeated until the next step matches the final answer indicator or the maximum number of steps is reached, and the steps used and selected are repeated. The method according to claim 6, including the method described in claim 6.

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

10. The method according to claim 2, further comprising the step of updating a base language model to generate the policy language model.

11. The step of updating the base language model to generate the policy language model is: The method according to claim 10, further comprising the step of performing supervised fine-tuning of the base language model for a supervised training example.

12. The step of training the aforementioned language model is: The method according to claim 9, comprising the step of initializing the language model based on the base language model.

13. The step of training the aforementioned reward model is: The method according to claim 9, comprising the step of initializing the reward model based on the base language model.

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

15. The method according to claim 1, further comprising the step of using the trained language model to generate output responses to input queries about a real-world environment, wherein the output responses provide information about the real-world environment or specify actions or paths to be taken in the real-world environment.

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

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

18. The method according to claim 1, further comprising the step of using the trained language model to generate an output response to an input query, which includes observations of a machine system operating in a real-world environment, wherein the generated output response is for diagnosing a fault in the machine system.

19. A computer implementation method for performing inference tasks using a language model, Steps include obtaining an input text sequence that characterizes the input query, A step of obtaining a reward model trained on multiple training examples, wherein each training example includes a sample query text sequence characterizing a sample query and a 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 an input including a query text sequence characterizing the 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 the correct final answer to the query, A step of using the language model to generate multiple candidate inference processes as a response to the input query, wherein each candidate inference process includes multiple candidate inference steps, each including a final candidate answer; A step of selecting the best inference process from the plurality of candidate inference processes based on one or more reward scores calculated by a trained reward model, A method comprising the step of outputting the best reasoning process as an output response to the input query.

20. The step of selecting the best inference process from the plurality of candidate inference processes based on one or more scores calculated by the trained reward model is: A step of using the reward model to generate weights for each of the plurality of candidate inference processes, wherein each weight measures the estimated exact probability of each of the candidate inference processes. A step of selecting the optimal final answer as the candidate final answer that has the maximum total weight, A step of selecting the candidate inference process having the highest weight estimated by the reward model among the candidate inference processes that produce the optimal final answer as the best inference process. The method according to claim 19, including the method described in claim 19.

21. The method according to claim 19, wherein the language model is trained using the method according to claim 1.

22. The method according to claim 19, wherein the input text sequence specifies a mathematical word problem, and the output response specifies a solution to the mathematical word problem and inference steps for solving the mathematical word problem.

23. Before outputting the best inference process as the output response, the step of determining whether the score estimated by the reward model for the best inference process is below a threshold, The method according to claim 19, further comprising the step of outputting the best reasoning process as the output response in response only to the determination that the score does not fall below the threshold.

24. The method according to claim 19, 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. The method according to claim 24, wherein the input query relates to a task in a real-world environment, and the method further includes the step of using the output response to control one or more machine agents or computer systems in the real-world environment to perform the task.

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

27. The method according to claim 19, wherein the input query includes observations of a mechanical system operating in a real-world environment, and the output response is for diagnosing a fault in the mechanical system.

28. The method according to claim 1, wherein the language model includes a neural network.

29. One or more computers, When executed by the one or more computers, one or more storage devices store instructions that cause the one or more computers to carry out the method described in any one of claims 1 to 28. A system equipped with these features.

30. One or more computer-readable storage media that, when executed by one or more computers, store instructions causing the one or more computers to perform the method described in any one of claims 1 to 28.

Images