Reward aware fine-tuning of generative neural networks

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for generating an output using a generative neural network. One of the methods include: obtaining a prompt input; determining, for each of a plurality objectives that correspond to different aspects of an output to be generated by a generative neural network, a weight to assign to the objective; generating a prefix input that defines the weight assigned to each of the plurality of objectives; and processing, using the generative neural network, the prefix input and the prompt input to generate the output.

Description

REWARD AWARE FINE-TUNING OF GENERATIVE NEURAL NETWORKSBACKGROUND

[0001] This specification relates to training neural networks.

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

[0003] This specification describes a system implemented as computer programs on one or more computers in one or more locations that trains a generative neural network that has parameters to update the parameters of the generative neural network based on optimizing a combined objective. The combined objective includes a weighted combination of reward scores that each correspond to a different objective.

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

[0005] Performance of a generative neural network on many content generation applications, e.g., conditional text, image, or audio generation tasks, can be improved when the generative neural network is able to generate outputs that achieve specific objectives that each correspond to a different aspect of an output. Examples of such outputs include an output sequence of text that is a helpful and concise response to an input sequence of text, a video that have frames that are free from blurring artifacts and that are temporally coherent, and so on. Specifically, in the prior example, one objective measures the helpfulness of the output sequence and another objective measures the conciseness of the output sequence.

[0006] This capability to generate outputs that achieve specific objectives saves computational and memory resources that are otherwise required by, e.g., the execution of multiple runs of the generative neural network to generate multiple outputs that may not achieve all of these objectives, until an output that simultaneously achieve these objectives is generated; or the further processing, e.g., modification, of such outputs to obtain the output that simultaneously achieve these objectives. This also improves user experience with these content generation applications.

[0007] Some existing systems maintain multiple generative neural networks in order to be able to generate outputs that achieve different combinations of objectives, e.g., by generating a first generative neural network through a fine-tuning process for a first combination of two objectives and separately generating a second generative neural network through another fine-tuning process for a second combination of the two objectives where the two objectives are weighted differently than the first combination.

[0008] However, repeatedly executing a fine-tuning process and subsequently storing a finetuned generative neural network for each different combination of obj ectives consumes a significant amount of computational and memory7resources when the space of possible combined objectives is large, e.g., when there are one hundred, one thousand, or more different ways in which the objectives can be weighted in the combination of objectives. In fact, executing a fine-tuning process for each different combination of objectives becomes practically infeasible when the space of possible combined objectives is large.

[0009] By leveraging reinforcement learning to perform reward-aware fine-tuning of a single instance of a generative neural network based on optimizing a combined objective that includes various weighted combinations of reward scores that correspond to different objectives, the system described in the specification reduces the amount of computational and memory resources that are needed to train and store the generative neural networks to generate outputs that can achieve different combinations of objectives because executing a fine-tuning process and storing a fine-tuned generative neural network for each different combination of objectives is no longer required.

[0010] At inference time, a new7prefix input that defines desired (e.g., customized) weights (and, correspondingly, the prioritization) of the individual objectives can be generated and then provided (as part of a context input) to the generative neural network for processing. Thus, the system can enhance the performance of the generative neural network to generate outputs that will achieve specific objectives in accordance with the desired prioritization, with reduced consumption of computational and memory resources relative to existing systems.

[0011] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.BRIEF DESCRIPTION OF THE DRAWINGS

[0012] FIG. 1 is a diagram of an example training system.

[0013] FIG. 2 is an illustration of example operations performed by a training system to finetune a generative neural network based on optimizing a combined objective.

[0014] FIG. 3 is a flow diagram of an example process for fine-tuning a generative neural network based on optimizing a combined objective.

[0015] FIG. 4 is a flow diagram of an example process for using a generative neural network to generate an output.

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

[0017] FIG. 1 is a diagram of an example training system 100. The training system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented.

[0018] The training system 100 trains a generative neural network 110 that is used to perform one or more generative tasks. In particular, the generative neural netw ork 110 is a neural network having parameters 116 and that can be configured through training to process a context input 102 in accordance with the parameters 116 to generate an output 152 for a generative task based on the context input 102. For example, the parameters 116 include weights and, optionally, biases of the layers of the generative neural network 110.

[0019] After training, the training system 100 or another inference system 150 can deploy the generative neural network 110 on one or more computing devices to perform inference for the one or more generative tasks, i.e., to generate new' outputs 152 for the generative tasks based on new' context inputs 102.

[0020] In some implementations, the generative neural network 110 can be configured as, or include, a generative (large) language model or a multi-modal model, e.g., a visual and language model, that generates an output for the generative task in an auto-regressive manner.

[0021] Examples of such neural networks include those described in Cohn Raffel, et al.Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 1910.10683. 2019; Daniel Adiwardana. et al. Towards a human-hke open-domain chatbot. CoRR, abs / 2001.09977, 2020; Tom B Brown, et al. Language models are few-shotlearners. arXiv preprint arXiv:2005. 14165, 2020; Aakanksha Chowdhery, et al. PaLM: Scaling Language Modeling with Pathways. arXiv preprint arXiv: 2204.02311; and Rohan Anil, et al. Palm 2 technical report. arXiv preprint arXiv:2305.10403, 2023.

[0022] In some implementations, the generative neural network 110 can be configured as diffusion model neural network that generates an output for the generative task across multiple updating iterations by performing a reverse diffusion process.

[0023] Examples of such neural networks include those described in Chitwan Saharia. et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36494, 2022; Aditya Ramesh, et al.Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125; Robin Rombach. et al. High-resolution image synthesis with latent diffusion model, Proceedings of the IEEE / CVF conference on computer vision and pattern recognition. 2022; and Huiwen Chang, et al. Muse: Text-to-image generation via masked generative transformers. arXiv preprint arXiv:2301.00704, 2023.

[0024] The generative neural network 110 can be configured to perform any of a variety of generative tasks, i.e., to generate outputs that include data of any type (modality) based on context inputs that include data of any type (modality).

[0025] In some implementations, the generative neural network 110 can be configured to generate an output that is an output sequence made up of tokens from a vocabulary. The vocabulary of tokens can include any of a variety of tokens that represent text symbols or other symbols. For example, the vocabulary of tokens can include one or more of characters, sub-words, words, punctuation marks, numbers, or other symbols that appear in a corpus of natural language text and / or computer code.

[0026] Additionally, or alternatively, the vocabulary of tokens can include tokens that can represent data other than text. For example, the vocabulary of tokens can include image tokens that represent a discrete set of image patch embeddings of an image that can be generated by an image encoder neural network based on processing the image patches of the image. As another example, the vocabulary of tokens can include audio tokens that represent code vectors in a codebook of a quantizer, e.g., a residual vector quantizer.

[0027] As an example, the generative neural network 110 can generate text sequences, i.e., each output generated by the generative neural network 110 is a sequence of text tokens from a vocabulary of text tokens that includes, e.g., one or more of characters, sub-words, words, punctuation marks, numbers, or other symbols that appear in natural language text. Forexample, the inference system 150 can use the generative neural network 110 to generate text sequences and provide the text sequences for presentation to users.

[0028] As another example, the generative neural network 110 can generate images or videos that have multiple frames (where each frame is an image) by generating images, e.g., either as sequences of pixels or through an iterative denoising process. For example, the output 152 generated by the generative neural network 110 is a plurality of color values for pixels in an image arranged according to a specified order. As another example, the output 152 generated by the generative neural network 110 is a plurality of tokens that represent image patch embeddings of an image which can then be processed by a decoder neural network to generate the image. For example, the inference system 150 can use the generative neural network 110 to generate an image or a video conditioned on a context input 102 that includes a text description of the content of the image or the video.

[0029] As another example, the context input 102 can include an input sequence of text and the output 152 is another sequence of text, e.g., a completion of the input sequence of text, a paraphrase of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text that is about a topic specified by the input sequence of text. As another example, the context input 102 can be an input other than text, e.g., an image, and the output 152 can be text that describes the input.

[0030] As another example, the context input 102 represents data to be compressed, e.g., image data, text data, audio data, or any other type of data; and the output 152 is a compressed version of the data. The tokens included in the output 152 can include any representation of compressed data, e.g., symbols or embeddings to be decoded by a respective neural network.

[0031] As a particular example, the inference system 150 can be part of a dialog system and the context input 102 can include audio or text from the most recent conversational turn submitted by a user of the dialog system during the dialog while the output 152 is the next turn in the conversation, e.g., either text or audio that is a response to the most recent conversational turn. Optionally, the context input 102 can also include one or more historical conversational turns that occurred earlier in the conversation.

[0032] As another particular example, the inference system 150 can be part of a machine translation system and the context input 102 can include text in a source language while the output 152 is text in a target language that is a translation of the source text into the target language.

[0033] As another particular example, the inference system 150 can be part of a computer code generation system and the context input 102 can be a text description of a desired piece of code or a snippet of computer code in a programming language and the output 152 can be computer code, e.g., a snippet of code that is described by the context input or a snippet of code that follows the context input in a computer program.

[0034] As another particular example, the inference system 150 can be part of an image editing system and the context input 102 can be an initial image and the output 152 is an edited version of the initial image, e.g., where the output 152 is an enhanced (e.g., optimized) version of the initial image that has higher quality (e.g., because it includes fewer blurring artifacts) than the initial image, where the output 152 is an image that has a higher resolution than the initial image, or where the output 152 is an inpainted image that includes a reconstruction of a portion of the initial image that was damaged, deteriorated, or missing.

[0035] The generative neural network 110 is typically trained using a two-stage approach: a pre-training stage followed by a fine-tuning stage, where at least the fine-tuning stage takes place at the training system 100. For example, the training system 100 can receive data specifying a pre-trained generative neural network 110 from another system, and then perform the fine-tuning of the pre-trained generative neural network 110.

[0036] In the pre-training stage, the generative neural network 110 is pre-trained by the training system 100 or another system based on optimizing one or more unsupervised or selfsupervised objective functions, e.g., a maximum-likelihood objective function, on a large dataset of text in one or more natural languages, e.g., text that is publicly available from the Internet or another text corpus, a large dataset of computer code in one or more programming languages, e.g., Python, C++, C#, Java, Ruby, PHP, and so on, e.g., computer code that is publicly available from the Internet or another code repository, a large dataset of audio samples, e.g., audio recordings or waveforms that represent the audio recordings, a large dataset of images where each image includes an array of pixels, a large dataset of videos where each video includes a temporal sequence of frames, or a large multi-modal dataset that includes a combination of two or more of these datasets. Optionally the pre-training stage also involves supervised fine-tuning (SFT) to tram the generative neural network 1 10 on specific generative tasks.

[0037] In the fine-tuning stage, the pre-trained generative neural network 110 is then adjusted to the generative tasks through fine-tuning adaptation. The generative tasks can include any combination of one or more of the generative tasks mentioned below and possibly other tasks. Examples of fine-tuning adaptation technique include reinforcement learning fromhuman feedback (RLHF), reinforcement learning from Al feedback (RLAIF), prompt tuning, instruction tuning, and the like, that use different training objectives, different data, or both.

[0038] As part of this fine-tuning, the training system 100 trains the generative neural network 110 through reinforcement learning based on optimizing a combined objective to repeatedly update the values of the parameters 116 of the generative neural network 110, i.e., to generate fine-tuned values of the parameters 116 from their pre-trained values that have been determined as a result of the pre-training. The combined objective includes a weighted combination of reward scores that are generated by multiple reward models and that correspond respectively to multiple different objectives.

[0039] FIG. 1 thus illustrates that the training system 100 includes multiple reward models 130A-N that are configured to generate multiple reward scores 132A-N that correspond respectively to different objectives. For example, reward model 130A corresponds to a first objective and is configured to generate a reward score 132A that corresponds to the first objective, reward model 130B corresponds to a second objective and is configured to generate a reward score 132B that corresponds to the second objective, and so on.

[0040] Generally, the reward scores measure objective-specific qualities of a training output 142 generated by the generative neural network 110 for a generative task. That is, the same training output 142 for the same generative task will be measured differently for different objectives, and hence will have different reward scores generated by different reward models.

[0041] By training the generative neural network 110 based on optimizing a combined objective that includes a weighted combination of the reward scores 132A-N that correspond to multiple objectives, the training system 100 can determine fine-tuned values of the parameters 116 that will cause the generative neural network 110 to generate outputs for the generative task that simultaneously achieve the multiple objectives.

[0042] More specifically, the training system 100 performs the fine-tuning over multiple fine-tuning training steps. At each fine-tuning training step, the training system 100 updates the parameters of the generative neural network 110 using multiple context inputs 122. Each context input 122 includes a prompt input 123 and a prefix input 124.

[0043] The prompt input 123 can include any type of data that provides context for the training output 142 to be generated by the generative neural network 110 the generative task on which it is being trained. For example, the prompt input 123 can include image data, text data, audio data, or another type of data.

[0044] The prompt inputs 123 can be obtained in many different ways. In some implementations, the training system 100 can access a prompt dataset that includes multipleprompt inputs, and then obtain the prompt inputs 123 from the prompt dataset, e.g., by sampling a batch of prompt inputs from the prompt dataset. In some implementations, the training system 100 can generate the prompt inputs 123 on an as-needed basis, e.g., by using another generative neural network.

[0045] The prefix input 124 includes text data that defines the weight to be assigned to each of the multiple objectives, i.e., defines how the multiple objectives should be combined to provide the combined objective that is being optimized. Each weight can be defined as a numeric value. For example, each weight can be defined as an integer or decimal number within a predetermined range of values, e.g., an integer number between 0 and 10, or between -10 and 10. or a decimal number between 0 and 1. or betw een -1 and 1.

[0046] The weights can be determined in many different ways. In some implementations, the training system 100 can receive a user input specifying the weights to be assigned to the multiple objectives. In some implementations, the training system 100 can determine the weight to be assigned to each objective by selecting the weight from a discrete set of possible weights for the objective, e.g., by iterating through a sequence of predetermined weights. In some implementations, the training system 100 can determine the weight to be assigned to each objective by selecting the weight from a continuous distribution of possible w eights for the objective, e.g., by randomly sampling a weight from a predetermined range of values.

[0047] For each context input 122 in the multiple context inputs, the training system 100 processes, using the generative neural network 110 and in accordance with the parameters 116, the context input 122 to generate a training output 142 for the generative task based on the context input 122. The training output 142 can include data of any type, e.g., image data, text data, audio data, or another type of data.

[0048] For each context input 122 in the multiple context inputs, the training system 100 processes, using each of the multiple reward models 130A-N, a reward model input that includes (i) the prompt input 123 included in the context input 122, (ii) the training output 142 that has been generated by the generative neural network 110 based on the context input 122, or both (i) and (ii) to generate a corresponding reward score 132A-N.

[0049] That is. in some cases, a reward score is generated by a reward model from processing a prompt input. In some other cases, a rew ard score is generated by a rew ard model from processing a training output generated by the generative neural netw ork 110. In some other cases, a reward score is generated by a reward model from processing both a prompt input and a training output generated by the generative neural network 110. In fact, in some further cases, a reward score is generated without using a rew ard model. Examples of such a rewardscore include a reward score that depends on an image quality, a reward score that depends on a sequence length for text, and so on.

[0050] Each reward model 130A-N can have any appropriate architecture, e.g., a convolutional architecture, a fully-connected architecture, or an attention architecture, that allows the reward model to process a reward model input to generate a reward score.

[0051] In some implementations, each reward model 130A-N can have a similar architecture as the generative neural network 1 10, but with a different output layer that allows the reward model to generate an output that is a single numeric value that defines the reward score, rather than an output that specifies an output for a generative task, e.g., an output that includes a score distribution that assigns a respective score to each token in a vocabulary of tokens.

[0052] In some implementations, each reward model 130A-N can have been trained by the training system 100, or another training system, on a corresponding reward model training dataset. For each reward model, the corresponding reward model training dataset includes multiple training tuples.

[0053] Each training tuple includes (i) a sample output, (ii) a ground truth reward score for the sample output, and, optionally, (iii) a sample input that provides context for the sample output. The sample input includes the same type(s) of data as the prompt input. The sample output includes the same type(s) of data as the output that is being generated by the generative neural network 110 for the generative task. The ground truth reward score can be a numeric value that represents the objective-specific quality of the sample output, i.e., the target reward value with respect to the corresponding objective that should be determined by the reward model by processing the sample output.

[0054] In these implementations, for each reward model 130A-N, the training system 100 trains the reward model using the training tuples sampled from the corresponding reward model training dataset by updating the values of parameters of the reward model to minimize a reward model loss function, e.g., by using a gradient-based optimization technique. The reward model loss function can include a loss term that measures, for each sampled training tuple, a difference between the ground truth reward score and a predicted reward score generated by the reward model from processing the sample output included in the sampled training tuple in accordance with the values of the parameters of the reward model.

[0055] As mentioned above, because the reward values measure objective-specific qualities of the training outputs 142 generated by the generative neural network 110 for the generative task, and because the same training output 142 for the same generative task may be measureddifferently for different objectives, multiple reward models 130A-N that correspond respectively to the different objectives will generally generate different reward scores than each other, even though they may process the same reward model input.

[0056] At each fine-tuning training step, the training system 100 then uses an optimization engine 140 to train the generative neural network 110 to update the values of the parameters 116 of the generative neural network 110 based on a combined objective.

[0057] The combined objective is generated based on determining a weighted combination, e.g., a weighted sum, of multiple reward scores 132A-N that have been generated by the multiple reward models 130A-N. In the weighted combination, each reward score is weighted by a weight that is defined in the prefix input 124 for the objective that corresponds to the reward score.

[0058] In some implementations, before the reward scores 132A-N are combined, e.g., summed, to generate the weighted combination, the training system 100 can apply normalization to each reward score 132A-N to generate a normalized reward score. In these implementations, the combined objective can thus be a weighted sum of the normalized reward scores.

[0059] The optimization engine 140 can use any appropriate reinforcement learning technique that uses a reinforcement learning objective function that depends on, e.g., includes, the combined objective.

[0060] That is. the optimization engine 140 can do this based on determining gradients of any reinforcement learning objective function that depends on the combined objective with respect to the parameters 116 of the generative neural network 110, and then using the gradients to adjust the values of the parameters 116 to optimize the combined objective, i.e., to increase each of some or all of multiple reward scores 132A-N included in the weighted combination.

[0061] FIG. 2 is an illustration 200 of example operations performed by the training system 100 of FIG. 1 to fine-tune the generative neural network 110 based on optimizing a combined objective. The combined objective includes a weighted combination of reward scores that correspond respectively to multiple different objectives which, in the example of FIG. 2. includes a helpfulness objective and a conciseness objective.

[0062] For example, the helpfulness objective may be defined with reference to an amount of time a user spent on viewing an output generated by the generative neural network 110, e.g., the greater the amount of time the user spent on viewing the output, the greater the helpfulness of the output. The helpfulness objective may also be defined in other ways.

[0063] For example, the conciseness objective can be defined with reference to the length (or another dimension, e.g., width or height) of an output generated by the generative neural network 110, e.g., the shorter the length of the output, the greater the conciseness of the output. The conciseness objective may also be defined in other ways.

[0064] The training system 100 generates a context input 122 that includes a prompt input 123 and a prefix input 124. In the example of FIG. 2, the prompt input 123 is a text sequence: “Help me write an email about the importance of AIT

[0065] The prefix input 124 includes text data that defines a weight to be assigned to the helpfulness objective; in the example of FIG. 2, this weight is 0.75. The prefix input 124 also includes text data that defines a weight to be assigned to the conciseness objective; in the example of FIG. 2. this weight is 0.25.

[0066] Optionally, the prompt input 123 and the prefix input 124 can be separated by predetermined tokens within the context input 122. Optionally, the prompt input 123 or the prefix input 124 can be wrapped by predetermined tokens, e.g., “[Begin of System Instruction]” and “[End of System Instruction]”, within the context input 122.

[0067] The generative neural network 110 processes the context input 122 in accordance with the parameters of the generative neural network 110 to generate a training output 142 for a generative task.

[0068] In the example of FIG. 2, the generative task is a text generation task, and the training output 142 generated by the generative neural network 110 is another text sequence that is a response to the text sequence included in the prompt input 123: “Dear [Recipient Name], I am writing to you today to discuss the importance of artificial intelligence (Al). . . ”

[0069] A first reward model 130A which, in the example of FIG. 2, is configured as a helpfulness reward model processes a reward model input that includes (i) the prompt input 123 included in the context input 122, (ii) the training output 142 that has been generated by the generative neural network 110, or both (i) and (ii) to generate a first reward score which, in the example of FIG. 2, is a reward score RHthat corresponds to the helpfulness objective, i.e., that represents a measure of helpfulness of the training output 142 given the prompt input 123.

[0070] A second reward model 130B which, in the example of FIG. 2, is configured as a conciseness reward model processes a reward model input that includes (i) the prompt input 123 included in the context input 122, (ii) the training output 142 that has been generated by the generative neural network 110, or both (i) and (ii) to generate a second reward scorewhich, in the example of FIG. 2, is a reward score Rcthat corresponds to the conciseness objective, i.e., that represents a measure of conciseness of the training output 142.

[0071] The training system 100 determines a combined objective which includes a weighted combination, e.g., weighted sum, of the first reward score that corresponds to the helpfulness objective and the second reward score that corresponds to the conciseness objective. In the example of FIG. 2. using weights defined by the prefix input 124, the combined objective can be computed as a weighted sum as: Reward = 0.75RH+ 0.25 / ?c.

[0072] The optimization engine 140 then updates the parameters 116 of the generative neural network 110 through reinforcement learning based on the combined objective. In particular, the optimization engine 140 performs an iteration of a reinforcement learning technique that uses a reinforcement learning objective function that depends on the combined objective to update the current values of the parameters 116 of the generative neural network 110 to optimize the combined objective.

[0073] By iteratively updating the parameters 116 of the generative neural network 110 in this manner, the training system 100 causes the first and second reward scores included in the weighted combination to be increased and the helpfulness and the conciseness of the training outputs generated by the generative neural network 110 to be improved. In other words, the generative neural network 110 is fine-tuned so that it can generate training outputs that are both helpful and concise for the generative task.

[0074] FIG. 3 is a flow diagram of an example process 300 for fine-tuning a generative neural network based on optimizing a combined objective. For convenience, the process 300 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 depicted in FIG. I, appropriately programmed in accordance with this specification, can perform the process 300.

[0075] In general the system can repeatedly perform iterations of the process 300 to fine-tune the generative neural network for each of multiple generative tasks. The multiple generative tasks can include any combination of one or more of the generative tasks mentioned below and possibly other tasks.

[0076] That is, the system can continue performing iterations of the process 300 until termination criteria for the fine-tuning of the generative neural network on a first generative task have been satisfied, e.g., until the parameters have converged, until a threshold amount of wall clock time has elapsed, or until a threshold number of iterations of the process 300have been performed, and then move onto a second generative task by continue performing iterations of the process 300 until termination criteria for the fine-tuning of the generative neural network on the second generative task have been satisfied, and so on.

[0077] The system obtains a batch of prompt inputs (step 302). The prompt input can include any ty pe of data that provides context for a training output to be generated by the generative neural network for the generative task on which it is being trained. For example, the prompt input can include image data, text data, audio data, or another type of data.

[0078] The prompt inputs can be obtained in many different ways. In some implementations, the system can access a prompt dataset that includes multiple prompt inputs, and then obtain the batch of prompt inputs by sampling from the prompt dataset. In some implementations, the system can generate the prompt inputs on an as-needed basis, e.g., by using another generative neural network.

[0079] The system repeatedly performs steps 304-310 for each prompt input in the batch.

[0080] The system determines a weight to be assigned to each of multiple objectives for training the generative neural network on the prompt input (step 304). The weights can be determined in many different ways. In some implementations, the system can receive a user input specifying the weights to be assigned to the multiple objectives. In some implementations, the system can determine the weight to be assigned to each objective by selecting the weight from a discrete set of possible yveights for the objective, e.g., by iterating through a sequence of predetermined weights. In some implementations, the system can determine the weight to be assigned to each objective by selecting the yveight from a continuous distribution of possible yveights for the objective, e.g., by randomly sampling a weight from a predetermined range of values.

[0081] Therefore, in some implementations, the weights assigned to the multiple objectives are the same for all prompt inputs included in the batch. For example, the yveight assigned to a first objective (or, analogously, a second objective) for one prompt input is the same as the weight assigned to the first objective (or, analogously, the second objective) for another prompt input.

[0082] Alternatively, in other implementations, the weights assigned to the multiple objectives for the prompt inputs included in the batch will differ. For example, the weight assigned to the first objective for one prompt input is different from the weight assigned to the first objective for another prompt input.

[0083] The system generates a prefix input that includes text data that defines the weight assigned to each of multiple objectives (step 306). That is, the prefix input defines howmultiple objectives should be combined to provide the combined objective that is being optimized. Each weight can be defined as a numeric value. For example, each weight can be defined as an integer or decimal number within a predetermined range of values, e.g., an integer number between 0 and 10, or between -10 to 10, or a decimal number between 0 and 1, or between -1 to 1.

[0084] In some implementations where the weights assigned to the multiple objectives are the same for all prompt inputs included in the batch, the prefix inputs generated by the system for the prompt inputs included in the batch can also be the same. Alternatively, in other implementations where the weights assigned to the multiple objectives are different for the prompt inputs included in the batch, the prefix inputs generated by the system for the prompt inputs included in the batch will also differ from each other.

[0085] The system processes, using the generative neural network and in accordance with the parameters of the generative neural network, a context input that includes the prefix input and the prompt input by to generate a training output (step 308). The training output can include data of any t pe, e.g., image data, text data, audio data, or another type of data.

[0086] For each of the multiple objectives, the system processes, using a reward model that corresponds to the objective, a reward model input that includes (i) the prompt input included in the context input, (ii) the training output that has been generated by the generative neural network based on the context input, or both (i) and (ii) to generate a reward score for the objective (step 310). The reward score measures an objective-specific quality of the training output, i.e., measures a quality of the training output with respect to the objective.

[0087] In practice there can be any number of objectives and, hence, the system can use any number of reward models that are each configured to generate reward scores for a corresponding objective.

[0088] For example, as mentioned above, when the generative task is a text generation task, the objectives can include a helpfulness objective and a conciseness objective.

[0089] As a further example, when the generative task is a text generation task, the objectives can additionally or alternatively include a safeness objective, a responsiveness objective, or both.

[0090] For example, the safeness objective may be defined with reference to an amount of toxic content (e.g., harmful, offensive, or otherwise inappropriate content) included in an output generated by the generative neural network, e.g., the lower the amount of toxic content, the greater the safeness of the output. The safeness objective may also be defined in other ways.

[0091] For example, the responsiveness objective may be defined with reference to a semantic similarity between an output generated by the generative neural network and a prompt input based on which the generative neural network generates the output, e.g., the greater the semantic similarity, the greater the responsiveness of the output. The responsiveness objective may also be defined in other ways.

[0092] As another further example, when the generative task is an image generation task where the output is an image, the objectives can include a fidelity objective, a resolution objective, or both.

[0093] For example, the fidelity7objective may be defined with reference to a number of perspective distortions (e.g., stretches or skews) included in an output generated by the generative neural network, e.g., the lower the number of perspective distortions, the greater the fidelity of the output. The fidelity objective may also be defined in other ways.

[0094] For example, the resolution objective may be defined with reference to a number of pixels included in an output generated by the generative neural network, e.g., the greater the number of pixels, the greater the resolution of the output. The resolution objective may also be defined in other ways.

[0095] As another further examples, when the generative task is an audio or video generation task wherein the output is a video or an audio, the objectives can include a temporal coherence objective, a temporal concreteness objective, or both.

[0096] For example, the resolution objective may be defined with reference to a number of artifacts (e.g., blurring of video continuations) that are present in video frames included in an output generated by the generative neural network, e.g., the lower the number of artifacts, the greater the temporal coherence of the output. The temporal coherence objective may also be defined in other ways.

[0097] For example, the temporal concreteness objective may be defined with reference to a number of video frames or audio samples included in an output generated by the generative neural network, e.g., the greater the number of video frames or audio samples, the greater the temporal concreteness of the output. The temporal concreteness objective may also be defined in other ways.

[0098] The system trains the generative neural network on the batch of prompt inputs to update values of the parameters of the generative neural network through reinforcement learning to optimize a combined objective (step 312). For each prompt input in the batch, the combined objective can be determined based on a weighted combination of the reward scores generated by the reward models for the multiple objectives for the prompt input.

[0099] For example, the combined objective can be determined as a weighted sum of the reward scores. In this weighted sum. the reward score for each objective is weighted by the weight assigned to the objective. In some implementations where different prompt inputs in the batch have different weights assigned to them, the weights in the combined objective will be different for different prompt inputs.[000100] In general the system can use any reinforcement learning technique that is appropriate for the generative neural network, i.e., that uses any appropriate reinforcement learning objective function that depends on the combined objective.[000101] Examples of reinforcement learning techniques that can be used include a proximal policy optimization (PPO) algorithm (see John Schulman, et al. "Proximal policy optimization algorithms." arXiv preprint arXiv: 1707.06347 (2017)), a direct preference optimization (DPO) algorithm (see Rafael Rafailov, et al. "Direct preference optimization: Your language model is secretly a reward model." Advances in Neural Information Processing Systems 36 (2024)), and an asynchronous actor-critic (A3C) algorithm (see Volodymyr Mnih, et al. "Asynchronous methods for deep reinforcement learning." International conference on machine learning. PMLR, 2016).[000102] By training the generative neural network based on optimizing a combined objective that includes a weighted combination of the reward scores that correspond to multiple objectives, the system can determine fine-tuned values of the parameters that will cause the generative neural network to generate outputs for the generative task that simultaneously achieve the multiple objectives.[000103] Advantageously, the system can fine-tune the generative neural network to achieve the multiple objectives in a holistic manner. That is, when possible, the generative neural network can generate outputs that achieve multiple objectives and, when achieving multiple objectives is unrealistic, the generative neural network can generate outputs that prioritize the objectives in the combined objective in accordance with their assigned weights. [000104] For example, assuming the multiple objectives include a helpfulness objective and a conciseness objective, and the helpfulness objective is assigned a relatively greater weight than the conciseness objective, then the generative neural network can be fine-tuned to generate outputs that are highly helpful but are somewhat verbose.[000105] As another example, assuming the multiple objectives include a safeness objective and a responsiveness objective, and the safeness objective is assigned a relatively greater weight than the responsiveness objective, then the generative neural network can befine-tuned to generate outputs that are free from toxic content but are somewhat unresponsive, in view of the prompt inputs based on which the outputs are generated.[000106] Further advantageously, the generative neural network trained in this way can generalize during inference to combined objectives that include objectives that are weighted differently than during the fine-tuning.[000107] FIG. 4 is a flow diagram of an example process 400 for using a generative neural network to generate an output. For convenience, the process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, an inference system, e.g., the inference system 150 depicted in FIG. 1, appropriately programmed in accordance with this specification, can perform the process 400.[000108] The system obtains a prompt input (step 402). For example, the prompt input can include image data, text data, audio data, or another type of data. For example, the prompt input can be received by the system as a user input.[000109] For each of multiple objectives that correspond to different aspects of an output to be generated by the generative neural network for a generative task, the system determines a weight to assign to the objective (step 404). In various implementations, the system can make such a determination based on the content of the prompt input, the user which provided the prompt input, the environment in which the generative neural network is deployed, and so on.[000110] In particular, the weights to be assigned to the multiple objectives need not be the same as the weights that were assigned to the multiple objectives during fine-tuning (although the individual objectives will generally be the same). For example, the weight assigned to the first objective for the prompt input obtained at step 402 may be different from the weight(s) assigned to the first objective for any prompt input obtained by the system during the fine-tuning.[000111] The system generates a prefix input that includes text data that defines the weight assigned to each of multiple objectives (step 406).[000112] The system processes, using the generative neural network and in accordance with the fine-tuned values of the parameters, a context input that includes the prefix input and the prompt input to generate an output for the generative task (step 408). By virtue of the fine-tuning of the generative neural netw ork, the output can achieve the multiple objectives in a holistic manner.[000113] For example, assuming that during the fine-tuning, the multiple objectives include a helpfulness objective and a conciseness objective where the helpfulness objective isY1assigned a different (e.g., relatively greater or smaller) weight than the conciseness objective. At step 404, the system assigns equal weights to the helpfulness objective and the conciseness objective. In this example, the outputs generated by the generative neural network can be about equally helpful and concise, despite during the fine-tuning the training outputs generated by the generative neural network were either helpful but verbose, or unhelpful but concise.[000114] This specification uses the term ‘‘configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.[000115] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine- readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.[000116] The term "data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an executionenvironment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.[000117] A computer program, which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form, including as a stand alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communication network.[000118] In this specification, the term “database” is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, the index database can include multiple collections of data, each of which may be organized and accessed differently.[000119] Similarly, in this specification the term “engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.[000120] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers. [000121] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read onlymemory or a random access memory7or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry'. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e g , a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few. [000122] Computer readable media suitable for storing computer program instructions and data include all forms of non volatile memory7, media and memory devices, including by way of example semiconductor memory7devices, 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.[000123] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid cry stal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory7feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.[000124] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardw are accelerator units for processing common and compute-intensive parts of machine learning training or production, i.e., inference, workloads.[000125] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a Jax framework.[000126] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back end, middleware, or front end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.[000127] The computing system can include clients and servers. A client and server are generally remote from each other and ty pically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g.. a result of the user interaction, can be received at the server from the device.[000128] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.[000129] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustratedoperations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.[000130] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.[000131] What is claimed is:

Claims

CLAIMS1. A computed-implemented method for training a generative neural network that has parameters, wherein the method comprises: obtaining a batch of prompt inputs: for each prompt input in the batch: determining a weight to be assigned to each of a plurality of objectives for training the generative neural network on the prompt input; generating a prefix input that defines the weight assigned to each of the plurality of objectives; processing, by the generative neural network, the prefix input and the prompt input to generate a training output: for each of the plurality of objectives, processing at least the training output using a reward model that corresponds to the objective to generate a reward score for the objective that measures a quality of the training output with respect to the objective; and training the generative neural network on the batch of prompt inputs to update values of the parameters of the generative neural network through reinforcement learning to optimize a combined objective that, for each prompt input, is based on a weighted combination of the reward scores generated by the reward models for the plurality of objectives, wherein, for each objective, the reward score for the objective is weighted in the weighted combination by the weight assigned to the objective.

2. The method of claim 1, wherein the generative neural network has been pre-trained on an unlabeled training dataset to optimize one or more unsupervised or self-supervised objective functions.

3. The method of any one of claims 1-2, wherein the reward model that corresponds to each of the plurality of objectives has been trained on a reward model training dataset that comprises a plurality of training tuples, each training tuple comprising at least a sample output and a ground truth reward score for the sample output.

4. The method of any one of claims 1-3, wherein training the generative neural network through reinforcement learning to optimize the combined objective comprises, before generating the weighted combination of the reward scores: normalizing the reward scores generated by the reward models for the plurality of objectives.

5. The method of any one of claims 1-4, wherein determining the weight to be assigned to each of the plurality of objectives comprises, for each of the plurality’ of objectives: selecting the weight from a discrete set of possible weights for the objective.

6. The method of any one of claims 1 -4, wherein determining the weight to be assigned to each of the plurality7of objectives comprises, for each of the plurality' of objectives: sampling the weight from a continuous distribution of possible weights for the objective.

7. A computed-implemented method comprising: obtaining a prompt input; determining, for each of a plurality7objectives that correspond to different aspects of an output to be generated by a generative neural network, a weight to assign to the objective; generating a prefix input that defines the weight assigned to each of the plurality of objectives; and processing, using the generative neural network, the prefix input and the prompt input to generate the output.

8. The method of claim 7, wherein the generative neural network comprises a language model neural network, and wherein the different aspects of the generative output to be generated by the generative neural network comprise two or more of: harmlessness, helpfulness, or conciseness.

9. The method of any one of claims 7-8, wherein the generative neural network has been pre-trained on a dataset to optimize one or more unsupervised or self-supervised objective functions, and then fine-tuned through reinforcement learning to optimize, for each training prompt input, a combined objective that is based on a weighted combination of reward scores for the plurality of objectives, wherein, for each objective, the reward score for the objective (i) measures a quality7of a training generative output with respect to the objective and (ii) is weighted in the weighted combination by a weight assigned to the objective.

10. The method of claim 9, wherein, for each objective, the reward score is generated by a reward model that corresponds to the objective.

11. A system comprising one or more computers and one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or morecomputers to perform the respective operations of any one of the methods of any of the preceding claims 1-10.

12. 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 respective operations of any one of the methods of any of the preceding claims 1-10.

Images