Weight averaged rewarded policy training for machine learning models

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for fine-tuning a target machine learning model to perform a target machine learning task. In one aspect, a method comprises: obtaining initial parameters for a target machine learning model; at each interpolation iteration of a sequence of interpolation iterations: training a plurality of auxiliary machine learning models to perform the target machine learning task using training data for the target machine learning task, interpolating the trained parameters for the plurality of auxiliary machine learning models for the interpolation iteration, and updating the current parameters for the target machine learning model using the interpolated parameters for the interpolation iteration; and, after the final interpolation iteration, determining a trained set of parameters for the target machine learning model based on the current parameters for the target machine learning model.

Description

WEIGHT AVERAGED REWARDED POLICY TRAINING EOR MACHINELEARNING MODELSCROSS-REFERENCE TO RELATED APPLICATIONS

[0001] This application claims priority to U.S. Provisional Application No. 63 / 662,986, filed on June 21, 2024. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application.BACKGROUND

[0002] This specification relates to processing data using machine learning models.

[0003] Machine learning models receive an input and generate an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models and generate the output based on the received input and on values of the parameters of the model.

[0004] Some machine learning models are deep models that employ multiple layers of models to generate an output for a received input. For example, a deep neural network is a deep machine learning model that includes an output layer and one or more hidden layers that each apply a nonlinear transformation to a received input to generate an output.SUMMARY

[0005] This specification generally describes a method for training (e g., fine-tuning or aligning) a target machine learning model to perform a target machine learning task. In particular, the described methods can incrementally fine-tune a machine learning model over a sequence of training iterations by, at each training iteration, training multiple sets of parameters (e.g., model weights) for the machine learning model starting from an initial shared set of parameters for the current training iteration and determining an initial shared set of parameters for the next training iteration as an interpolation of the sets of parameters for the current training iteration.

[0006] In general, the target machine learning task can be any of a variety of machine learning tasks. For example, the target machine learning task can include, e.g., a classification task, a regression task, a segmentation task, an auto-encoding task, a data generation task, or any other appropriate task or a combination thereof.

[0007] The target machine learning task can include a variety of specific tasks. The described methods can fine-tune or align the target machine learning model based on, e.g., preferences orrewards for outputs of each of the specific tasks included within the target machine learning task. For example, the target machine learning task can be a multi-modal data generation task that can include generating output data in a variety of data modalities (e.g., text data, image data, audio data, video data, etc.). As a further example, the target machine learning task can include processing a model input to, e.g., generate a multi-modal output that includes data for a plurality of data modalities, generate a model output that includes data for one or more data modalities as specified by the model input, and so on. The described methods can fine-tune or align the target machine learning model to perform such a multi-modal target machine learning task as a whole based on preferences or rewards for the specific tasks of generating each of the output data modalities.

[0008] The target machine learning model can be any appropriate machine learning model configured to perform the target machine learning task. For example, the machine learning model can be a random forest model, a support vector machine model, a regression model, a neural network, and so on. As a particular example, the target machine learning model can be a generative model, such as a Large Language Model (LLM) or a Vision-Language Model (VLM).

[0009] The target machine learning model can be pre-trained to perform one or more pretraining tasks. The pre-training tasks can be more general data processing tasks compared to the target machine learning task and the target machine learning task can be, e.g., a domain specific task, an application specific task, and so on.

[0010] When the target machine learning model is pre-trained, the described methods can more efficiently train the target machine learning model to perform the target machine learning task. For example, when the target machine learning model is pre-trained, the described methods can train the target machine learning model to perform the more specific target machine learning task using, e.g., fewer training examples for the target machine learning task, fewer training iterations, and so on, which can reduce the computational resources (e.g., computational time, memory usage, power consumption, etc.) required to train the target machine learning model to perform the target machine learning task. Additionally, such pre-training of the target machine learning model can enable the target machine learning model to still effectively perform the more general pre-training tasks after being trained to perform the more specific target machine learning task.

[0011] According to one implementation, there is provided a method performed by one or more computers and for training a target machine learning model to perform a target machine learning task, the method comprising: obtaining an initial set of machine learning model parameters for the target machine learning model; initializing, using the initial set of machine learning model parameters, a current set of machine learning model parameters for the target machine learning model; at each interpolation iteration of a sequence of interpolation iterations: for each of a plurality of auxiliary machine learning models for the interpolation iteration: initializing a set of machine learning parameters for the auxiliary machine learning model using the current set of machine learning parameters for the target machine learning model; and training the auxiliary machine learning model to perform the target machine learning task using training data for the target machine learning task to determine a trained set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration; and determining an interpolated set of machine learning model parameters for the interpolation iteration by interpolating the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration; and updating the current set of machine learning model parameters for the target machine learning model based on the interpolated set of machine learning model parameters for the interpolation iteration; and after the final interpolation iteration of the sequence of interpolation iterations, determining a trained set of machine learning model parameters for the target machine learning model based on the current set of machine learning model parameters for the target machine learning model.

[0012] In some implementations, interpolating the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration comprises performing a linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration.

[0013] In some implementations, interpolating the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration comprises performing a spherical linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration.

[0014] In some implementations, performing the spherical linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration comprises performing a layer-wise spherical linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration.

[0015] In some implementations, updating the current set of machine learning model parameters for the target machine learning model based on the interpolated set of machine learning model parameters for the interpolation iteration comprises: updating, based on the interpolated set of machine learning model parameters for the interpolation iteration, the current set of machine learning model parameters for the target machine learning model as an exponential moving average of the interpolated set of machine learning model parameters for each of the interpolation iterations.

[0016] In some implementations, training the auxiliary machine learning model to perform the target machine learning task to determine the trained set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration comprises: for each of a sequence of training steps for the auxiliary machine learning model for the interpolation iteration: updating the set of machine learning parameters for the auxiliary machine learning model for the interpolation using a machine learning technique to optimize an objective function for the target machine learning task.

[0017] In some implementations, updating the set of machine learning parameters for the auxiliary machine learning model for the interpolation using the machine learning technique to optimize the objective function for the target machine learning task comprises: updating the set of machine learning parameters for the auxiliary machine learning model for the interpolation using a gradient of the objective function for the target machine learning task.

[0018] In some implementations, training the auxiliary machine learning model to perform the target machine learning task to determine the trained set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration further comprises: initializing an anchor set of machine learning parameters using the current set of machine learning parameters for the target machine learning model.

[0019] In some implementations, training the auxiliary machine learning model to perform the target machine learning task to determine the trained set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration further comprises, at each of the sequence of training steps for the auxiliary machine learning model for the interpolation iteration: after updating the set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration: updating, based on the set of machine learning model parameters for the auxiliary machine learning model interpolation iteration, the anchor set of machine learning model parameters for the target machine learning model as an exponential moving average of the set of machine learning model parameters over each of the training steps for the auxiliary machine for the interpolation iteration.

[0020] In some implementations, the objective function for the target machine learning task includes a term measuring a Kullback-Leibler divergence between (i) a distribution determined using the auxiliary model for the interpolation iteration in accordance with the set of machine learning model parameters for the auxiliary machine learning model and (ii) a distribution determined using the auxiliary model for the interpolation iteration in accordance with the anchor set of machine learning model parameters.

[0021] In some implementations, determining the trained set of machine learning model parameters for the target machine learning model based on the current set of machine learning model parameters for the target machine learning model comprises: determining the trained set of machine learning model parameters for the target machine learning model as a weighted combination of the current set of machine learning model parameters for the target machine learning model and the initial set of machine learning model parameters for the target machine learning model.

[0022] In some implementations, obtaining the initial set of machine learning model parameters for the target machine learning model comprises: obtaining the initial set of machine learning model parameters for the target machine learning model by pre-training the target machine learning model to perform one or more pretraining tasks.

[0023] In some implementations, the target machine learning model comprises a generative model.

[0024] In some implementations, the target machine learning model comprises a language model.

[0025] In some implementations, the target machine learning model comprises an image generation neural network.

[0026] In some implementations, the target machine learning model is configured to process input token sequences to generate corresponding output token sequences, wherein the input token sequence and the output token sequence comprise tokens from a vocabulary of tokens for the target machine learning model.

[0027] In some implementations, the vocabulary of tokens for the target machine learning model comprises one or more tokens representing text data.

[0028] In some implementations, the vocabulary of tokens for the target machine learning model comprises one or more tokens representing image data.

[0029] In some implementations, the vocabulary of tokens for the target machine learning model comprises one or more tokens representing video data.

[0030] In some implementations, the vocabulary of tokens for the target machine learning model comprises one or more tokens representing audio data.

[0031] In some implementations, the vocabulary of tokens for the target machine learning model comprises one or more tokens representing structured numerical data.

[0032] In some implementations, the vocabulary of tokens for the target machine learning model comprises one or more tokens representing an instruction for a system or agent to perform a task.

[0033] In some implementations, the target machine learning model is configured to interact with a user; and the training data for the target machine learning task comprises a plurality of training examples, wherein each training example comprises (i) an example of a query from anexample user for the training example and (ii) an example response to the respective example of a query from the example user for the training example.

[0034] In some implementations, the target machine learning model is configured to select actions for an agent interacting with an environment to perform a task in the environment; and the training data for the target machine learning task comprises a plurality of training examples, wherein each training example comprises an example selected action for an example agent to perform the task in an example environment for the training example.

[0035] In some implementations, each training example comprises a respective example observation of the example environment for the training example.

[0036] According to another implementation, there is provided a system comprising: one or more computers; and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform the operations of the respective method.

[0037] According to another implementation, there is provided a one or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform the operations of the respective method.

[0038] According to another implementation, there is provided method performed by one or more computers for generating a model output using a target machine learning model, the method comprising: obtaining a trained set of model parameters for the target machine learning model by, at each of a sequence of interpolation iterations: (i) training a plurality of auxiliary machine learning models to perform a target machine learning task and (ii) updating the set of model parameters for the target machine learning model based on an interpolation of the trained model parameters for the plurality of the auxiliary machine learning models; receiving a model input for the target machine learning model; and processing the model input for the target machine learning model using the trained model parameters for the target machine learning model to generate the model output.

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

[0040] Conventional methods for fine-tuning or aligning a machine learning model can erroneously update the target machine learning model away from pre-trained initializations andcan therefore lose the computational benefits of pre-training the target machine learning model. In particular, when fine-tuning or aligning the target machine learning model using relatively few training examples for the target machine learning task, conventional training methods can encounter training errors, such as catastrophic forgetting and reward hacking, caused by the small amount of training data for the target machine learning task.

[0041] By performing a sequence of interpolation iterations between multiple sets of model parameters trained during each interpolation iteration and by using an interpolation (e.g., a linear interpolation, a spherical interpolation, etc.) of the trained sets of model parameters at each interpolation iteration as an initialization for the next interpolation iteration, the described methods can incrementally fine-tune the target machine learning model while avoiding training errors such as catastrophic forgetting and reward hacking. In particular, the described methods can achieve better generalization, combine the strengths of individual trained auxiliary models at each interpolation iteration, and reduce memorization of a small training set for the target machine learning task. By iteratively interpolating multiple sets of model parameters, the described methods can therefore more reliably fine-tune the target machine learning model using less training data for the target machine learning task and using fewer computational resources as compared to conventional methods.

[0042] As an example, fine-tuning large generative models, such as large language models (LLMs) and vision-language models (VLMs), with human feedback is particularly useful in many applications, as these large models can have the computational capability of accurately modeling human preferences and of generating high quality outputs. Fine-tuning generative models using human feedback can enable the generative models to learn human preferences regarding model outputs and to generate more preferable outputs. However, obtaining human feedback can be resource intensive, which can significantly limit the amount of training data available for fine-tuning based on human feedback. Conventional fine-tuning methods, such as reinforcement learning from human feedback (RLHF), can suffer from training errors such as catastrophic forgetting and reward hacking when only a limited amount of human feedback is available. By more reliably fine-tuning generative models using less training data, the described methods can therefore provide an improved method for fine-tuning generative models based on human feedback.

[0043] Additionally, in some implementations, the described methods can, at each interpolation iteration, initialize a respective anchor set of model parameters for each of the auxiliary machine learning models. While training the auxiliary machine learning models, each anchor set of machine learning model parameters can be updated as an exponential moving average of the set of machine learning model parameters over the training steps for the corresponding auxiliary machine learning model. The described methods can train the auxiliary machine learning models using a Kullback-Leibler regularization that encourages the model parameters of the auxiliary machine learning models to remain close to the corresponding anchor sets of model parameters. The exponential moving average anchor sets of model parameters can better maintain essential information from the initial parameters for each interpolation iteration and can act as slow weights that are progressively updated during training of the auxiliary machine learning models.

[0044] Using the exponential moving average anchor sets of model parameters as part of the Kullback-Leibler regularization enables a gradual relaxation of the regularization. The auxiliary model parameters can initially be strongly tied to the initial parameters for each interpolation iteration and can then be updated more aggressively later in training, leading to better performing auxiliary machine learning models. Kullback-Leibler regularization with exponential moving average anchor sets of model parameters can therefore enable more efficient training and can further reduce the computational cost of training the target machine learning model.

[0045] Additionally, in some implementations, the current set of machine learning model parameters for the target machine learning model at each interpolation iteration can be determined as an exponential moving average of the interpolated set of machine learning model parameters for each of the interpolation iterations, e.g., as parameterized by a moving average weighting for the model parameters of the target machine learning model. In some implementations, the final trained set of machine learning model parameters for the target machine learning model can be determined as a weighted combination of the current set of machine learning model parameters for the target machine learning model as of the final interpolation iteration and the initial set of machine learning model parameters (e.g., the pretrained set of model parameters) for the target machine learning model, e.g., as parameterized by a final weighting for the model parameters of the target machine learning model. By adjusting the moving average weighting and the final weighting for the model parameters of the target machine learning model, the described methods can adjust a trade-off between features learnedduring fine-tuning for the target machine learning task and more general knowledge from pretraining, enabling the described methods to perform efficient balancing between general and task-specific features.

[0046] 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

[0047] FIG. l is a block diagram of an example training system.

[0048] FIG. 2 illustrates updating model parameters for a target machine learning model by iteratively interpolating trained model parameters for a set of auxiliary machine learning models.

[0049] FIG. 3 is a flow diagram of an example process for training a target machine learning model to perform a target machine learning task over a sequence of interpolation iterations.

[0050] FIG. 4 is a flow diagram of an example process for training an auxiliary machine learning model to perform a target machine learning task.

[0051] FIG. 5 illustrates an example algorithm that a training system can use to train a target machine learning model.

[0052] FIG. 6 illustrates experimental results demonstrating a performance of example target machine learning models that have been trained using the described methods.

[0053] Like reference numbers and designations in the various drawings indicate like elements.DETAILED DESCRIPTION

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

[0055] The training system 100 can train a target machine learning model 102 to perform a target machine learning task using a set of training data 104 for the target machine learning task.

[0056] In general, the target machine learning task can be any of a variety of machine learning tasks. For example, the target machine learning task can include, e.g., a classification task, aregression task, a segmentation task, an auto-encoding task, a data generation task, or any other appropriate task or a combination thereof.

[0057] The target machine learning model 102 can be any appropriate machine learning model configured to perform the target machine learning task by processing model inputs for the target machine learning task to generate model outputs for the target machine learning task. For example, the machine learning model can be a random forest model, a support vector machine model, a regression model, a neural network, and so on.

[0058] The model inputs for the target machine learning task can include any appropriate type of data for the target machine learning task. In particular, the target machine learning model 102 can be configured to process inputs of any of multiple modalities, e.g., image, video, audio, text, or any other appropriate modality for the target machine learning task, or any combination thereof. The target machine learning model 102 can have a single set of model parameters while, at the same time, being substantially flexible such that it is able to process any one, or a combination, of the multiple modalities.

[0059] In some implementations, the model inputs can characterize entities, including images, audio waveforms, point clouds (e.g., generated by a lidar or radar sensor), proteins, sequences of words (e.g., that form one or more sentences or paragraphs), videos (e.g., represented as sequences of video frames), or any other appropriate type of data or a combination thereof. As used herein an image, i.e. a static or moving image (video), may include a point cloud.

[0060] As a particular example, an entity can include multiple units arranged in a spatial or temporal structure, e.g., an entity can be an image and each unit can be a pixel in the image. Each unit in an entity can have an associated data element embedding that can characterize, e.g., a position of the unit in the spatial or temporal structure and / or features associated with the unit in the spatial structure. In another particular example, an entity can be a text, e.g., an entity can be a sequence of words, phrases, characters, or word pieces, in one or more languages. In yet another particular example, an entity can include a combination of different modalities, e.g., an entity can include a combination of audio, video, and label data. As another example, an entity can include sensor data from one or more sensors that are configured to perceive one or more characteristics of a real-world environment such as image data, audio data, or so-called “fine touch” sensor data (which permits localization). Although a number of examples of model inputs are described above, generally, the model input can have any appropriate dimensionality and structure for thetarget machine learning task, and can include any appropriate type(s) of data for the machine learning task.

[0061] The model outputs for the target machine learning task can be, e.g., classification outputs, regression outputs, sequence outputs (i.e., that includes a sequence of output elements), segmentation outputs, auto-encoding outputs, or any other appropriate model output or a combination thereof.

[0062] As a particular example, the target machine learning task can be a generative machine learning task that includes processing input prompts to generate corresponding output data items. For example, the target machine learning task can include receiving an input query from a user and processing the received query to generate an output as a response to the received query. The target machine learning task can include, e.g., generating output text, an output image, output audio, an output video, and so on in response to a user query. As another example, the target machine learning task can include selecting actions for an agent interacting with an environment to perform a task in the environment. As a further example, the target machine learning task can include processing an input prompt that includes data characterizing the environment (e.g., data characterizing an observation of the environment) to generate a selected action for the agent as the model output.

[0063] The target machine learning model 102 can be a neural network with any appropriate architecture for processing model inputs (e.g., input prompts) for the target machine learning task to generate model outputs (e.g., output data items) for the target machine learning task. In particular, the target machine learning model 102 can be a neural network that includes any of a variety of processing layers (e.g., feedforward layers, convolutional layers, recurrent layers, attention layers, graph processing layers, etc.) in any appropriate combination for performing the target machine learning task.

[0064] For example, the target machine learning model 102 can be a sequence processing neural network configured to generate output sequences (e.g., output token sequences) representing output data items for the target machine learning task by processing input sequences (e.g., input token sequences) representing input prompts for the target machine learning task. As a further example, the target machine learning model 102 can be an auto-regressive generative model (e.g., a Transformer, a recurrent neural network, etc.) that can auto-regressively generate output sequences for the target machine learning task. A transformer neural network is a neural networkthat includes a stack of transformer blocks, each typically including an attention or self-attention neural network layer, generally followed by a feedforward neural network layer (where a selfattention neural network layer applies a self-attention operation, e.g., QKV self-attention, to elements of an embedding, to update each element of the embedding).

[0065] The target machine learning model 102 can, for example, be a large language model (LLM) that can generate tokenized representations of text data; a vision-language model (VLM) that can generate tokenized representations of image or video data, e.g., in response to a text input or that can generate tokenized representations of text, e.g., in response to an image input; an audio model that can input or generate tokenized representations of audio data; or a multimodal model that can generate output token sequences representing text data, image data or audio data, e.g., in response to inputs characterizing input text, input images input audio; and so on.

[0066] Generally, prior to the training of the target machine learning model 102 by the system 100, the target machine learning model 102 can have already been trained across one or more previous training stages.

[0067] For example, the one or more previous training stages can include a pre-training stage. During the pre-training stage, the target machine learning model 102 can have been trained by the system 100 or a separate system to perform one or more pre-training tasks.

[0068] As a particular example, when the target machine learning model 102 is a sequence processing neural network, the one or more pre-training tasks can be next token prediction tasks, e g., tasks that requires predicting, given a current sequence of tokens, a next token that follows the current sequence in training data for the pre-training tasks. For example, as part of pre-training, the target machine learning model 102 can have been trained to optimize a maximum-likelihood objective 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.

[0069] As another example, the one or more previous training stages can include one or more additional training stages, e.g., that occur after the pre-training stage. For example, the one or more previous training stages can include any one or more of: a supervised fine-tuning stage, a reinforcement learning stage, a preference learning stage, an instruction tuning stage, and so on.

[0070] Such training of the target machine learning model 102 over the one or more previous training stages can enable the training system 100 to more efficiently train the model 102 to perform the target machine learning task (e.g., using less training data, fewer training iterations, etc.) and can enable the target machine learning model 102 to perform the one or more pretraining tasks after being trained by the system 100 to perform the target machine learning task.

[0071] In particular, the training system 100 can efficiently fine-tune or align the target machine learning model 102 to generate more preferable outputs for the target machine learning task using the training data 104. When the target machine learning model 102 is a large generative model, such as a large language model or a vision-language model with hundreds of millions or billions of parameters, the target machine learning model 102 can have a computational capability to accurately model human preferences for the target machine learning task and to generate high quality (e.g., more preferable) outputs for the target machine learning task. While the one or more previous training stages of the target machine learning model 102 can enable the model to process inputs and generate outputs for the target machine learning task, the model 102 often requires additional fine-tuning to correctly model preferences regarding the target machine learning task. By fine-tuning the target machine learning model 102 using training data 104 that includes feedback (e.g., human feedback) for example outputs for the target machine learning task, the training system 100 can specifically fine-tune the model 102 to produce more preferable outputs for the target machine learning task.

[0072] The target machine learning task can include a variety of specific tasks. The training system 100 can fine-tune or align the target machine learning model based on, e.g., preferences or rewards for outputs of each of the specific tasks included within the target machine learning task. For example, the target machine learning task can be a multi-modal data generation task that can include generating output data in a variety of data modalities (e.g., text data, image data, audio data, video data, etc.). As a further example, the target machine learning task can include processing a model input to, e.g., generate a multi-modal output that includes data for a plurality of data modalities, generate a model output that includes data for one or more data modalities asspecified by the model input, and so on. The training system 100 can fine-tune or align the target machine learning model to perform such a multi-modal target machine learning task as a whole based on preferences or rewards for the specific tasks of generating each of the output data modalities.

[0073] Example target machine learning tasks and example architectures for the target machine learning model 102 are described in more detail later in this specification.

[0074] The training data 104 for the target machine learning task can include a plurality of training examples 106 for the target machine learning task. Each of the training examples 106 can include a training input (e.g., an example prompt) for the training example and a target output (e.g., an example model output) for the training example.

[0075] In some implementations, each training example can include a quality score for the training example that measures a quality of the target output for the training example. As an example, the quality score for each training example can characterize a human preference for the target output for the training example. As another example, the quality score for each training example can be a task reward for the target machine learning task that characterizes, e.g., a degree of success, an accuracy, and so on associated with the target machine learning model 102 performing the target machine learning task by processing the training input for the training example to generate the target output for the training example.

[0076] In general, the system 100 can receive the quality scores for the training examples from any of a variety of sources, e g., from a user, from another system, as an output from a trained model (e.g., a reward model).

[0077] The quality score for each training example can be determined based on the training input and the target output for the training example. The quality score for each training example can be determined based on, for example, rating or preference information corresponding to the target output for the training example, as provided by a user or a trained machine model. The quality score for each training example can be determined based on execution of one or more processes based on the target output for the training example. For example, a target output can include computer language such as programming code which, when executed by a computer, causes the computer to carry out a process and the quality score for the target output can be determined based on the execution of the process, for example based on whether the process was successfully executed to completion, or based on metrics relating to the process such as memory usage forexample. In another example, a target output can include data representing actions to be taken by an agent (e g., a mechanical agent such as a robot) and the quality score for the target output can be determined based on the execution of the action, for example in a simulated environment or a real world environment.

[0078] The quality scores for the training examples can be determined using a trained reward model. The reward model can have been trained using a dataset comprising example inputs, example outputs, and target quality scores based on rating or preference information or based on execution of a process, for example.

[0079] The training system 100 can train the target machine learning model 102 by, at each of a sequence of interpolation iterations, training each of a plurality of auxiliary machine learning models 106 to perform the target machine learning task and interpolating the trained model parameters of the auxiliary machine learning models 106 to update the model parameters of the target machine learning model 102. Each of the auxiliary machine learning models 106 can have the same architecture as the target machine learning model 102 while including different machine learning model parameters than those of the target machine learning model 102. For example, when the target machine learning model 102 is a neural network, the auxiliary machine learning models 106 can be neural networks that each include a same set of processing layers (e.g., feedforward layers, convolutional layers, recurrent layers, attention layers, graph processing layers, etc.) as the target machine learning model 102 in a same combination as the model 102, while each processing layer of the auxiliary machine learning models 106 can be parameterized by different machine learning model parameters as compared to the corresponding processing layer of the model 102.

[0080] At each interpolation iteration, the training system 100 can determine initial model parameters 108 for each of the auxiliary machine learning models 106 for the interpolation using the current model parameters of the target machine learning model 102 as of the interpolation iteration. In particular, the initial model parameters 108 for each of the auxiliary machine learning models 106 can be the current model parameters of the target machine learning model 102 as of the interpolation iteration.

[0081] The training system 100 can then, over a sequence of training steps for the interpolation iteration, train each of the auxiliary machine learning models 106 to optimize an objective function for the target machine learning task. In particular, at each training step, the system 100can process training inputs 110 for a plurality of training examples for the training step using the auxiliary machine learning models 106 to generate corresponding model outputs 112 for the training examples for the training step. The system 100 can process the model outputs 112 and corresponding target outputs 114 for the training examples for the training step using an update system 116 to determine parameter updates 118 for the auxiliary machine learning models 106.

[0082] The update system 116 can determine the parameter updates 118 for the auxiliary machine learning models 106 using any appropriate machine learning technique to optimize the objective function for the target machine learning task. The objective function can be any appropriate objective function for the target machine learning task. For example, the objective function can measure a cross-entropy loss between the model outputs 112 and the corresponding target outputs 114. As another example, the objective function can measure an error (e.g., a root mean-squared error, an LI loss, an L2 loss, etc.) between the model outputs 112 and the corresponding target outputs 114. As another example, the objective function can measure a likelihood of the auxiliary machine learning models 106 generating the target outputs 114 by processing the corresponding training inputs 110. As another example, the objective function can measure a reward for the target machine learning task based on the model outputs 112 and the quality scores for the corresponding target outputs 114.

[0083] At each training step, the training system 100 can use different training examples from the training data 104 to train each of the auxiliary machine learning models 106. Therefore, once trained, the auxiliary machine learning models 106 will each have different sets of trained model parameters.

[0084] At each interpolation iteration, after training the auxiliary machine learning models 106, the training system 100 can update the target machine learning model 102 by interpolating the trained model parameters of the auxiliary machine learning models 106. In particular, the system 100 can update the model parameters of the target machine learning model 102 to be the interpolated model parameters 120 of the auxiliary machine learning models 106.

[0085] Updating the model parameters of the target machine learning model 102 as the interpolated model parameters 120 of the auxiliary machine learning models 106 enables the system 100 to incrementally fine-tune the target machine learning model 102. As compared to directly training the target machine learning model 102 to optimize the objective function for the target machine learning task training the target machine learning model 102 using theinterpolated parameters of the auxiliary machine learning models 106 can achieve better generalization, combine the strengths of the trained auxiliary machine learning models 106 at each interpolation iteration, and reduce memorization of a small training set for the target machine learning task. By iteratively interpolating the model parameters of the auxiliary machine learning models 106, the system 100 can therefore more reliably fine-tune the target machine learning model using less training data for the target machine learning task and using fewer computational resources as compared to conventional methods.

[0086] In some implementations, at each interpolation iteration, the training system 100 can initialize a respective anchor set of model parameters for each of the auxiliary machine learning models 106. The objective function for the target machine learning task can include a Kullback- Leibler regularization that encourages the model parameters of the auxiliary machine learning models 106 to remain close to the corresponding anchor sets of model parameters. Such regularization based on the anchor sets of model parameters for the auxiliary machine learning models 106 can enable the system 100 to better maintain essential information from the initial parameters for each interpolation iteration while training the auxiliary machine learning models 106.

[0087] While training the auxiliary machine learning models 106, the system 100 can update the anchor set of model parameters for each of the auxiliary machine learning models 106 as an exponential moving average of the model parameters for the auxiliary machine learning model over the training steps. The exponential moving average anchor sets of model parameters can act as slow weights that are progressively updated during training of the auxiliary machine learning models 106.

[0088] Using the exponential moving average anchor sets of model parameters as part of the Kullback-Leibler regularization can enable a gradual relaxation of the regularization. The model parameters for the auxiliary machine learning models 106 can initially be strongly tied to the initial parameters for each interpolation iteration and can then be updated more aggressively later in training, leading to better performing auxiliary machine learning models 106.

[0089] In some implementations, the training system 100 can update the model parameters for the target machine learning model 102 as an exponential moving average of the interpolated set of machine learning model parameters for each of the interpolation iterations. In some implementations, after the final interpolation iteration, the system 100 can determine the finaltrained set of model parameters for the target machine learning model 102 as a weighted combination of the current set of model parameters for the target machine learning model 102 as of the final interpolation iteration and the initial set of machine learning model parameters (e.g., the pre-trained set of model parameters) for the target machine learning model 102. By adjusting weights for the exponential moving average and the final weighted combination of the model parameters of the target machine learning model 102, the system 100 can adjust a trade-off between features learned during fine-tuning for the target machine learning task and more general knowledge from pre-training, enabling the system 100 to balance between task specific performance for the target machine learning task and retaining general knowledge from pretraining.

[0090] An example process by which the training system 100 can train the auxiliary machine learning models 106 and update the model parameters for the target machine learning model 102 is described in more detail below with reference to FIG. 2 and FIG. 3.

[0091] As illustrated by experimental results provided in FIG. 6 below, the described methods can enable the training system 100 to train the target machine learning model 102 to attain a more favorable trade-off between performance on the target machine learning task and the strength of the regularization used to retain general knowledge learned during pre-training as compared to conventional training methods. In particular, for a same strength of the regularization used to retain general knowledge learned during pre-training, the system 100 can train the target machine learning model 102 to attain better performance on the target machine learning task as compared to conventional training methods.

[0092] After training by the training system 100, the target machine learning model 102 can be used to perform the target machine learning task by receiving and processing model inputs (e.g., input prompts) for the task (e.g., from a user, another system, etc.) to generate model outputs (e.g., output data items) for the task.

[0093] Example target machine learning tasks and example architectures for the target machine learning model 102 are described below.

[0094] In some implementations, the target machine learning task can be to perform an image or audio segmentation task. For example, the target machine learning model 102 can process a model input that includes an image or a series of audio samples that represent an audio waveform to generate a model output that defines, for each pixel in the input image or for each audiosample, a respective score distribution over a set of possible classes. The score for a class can define a likelihood that a corresponding pixel or sample is included in the class. For example, for an image the possible classes can include, e.g., water, building, vehicle, pedestrian, etc. The image can be a moving image, i.e. a video. Similarly, the target machine learning task can be to perform an image depth prediction task where the model output defines a predicted depth value for a corresponding pixel to obtain a (spatial 3D) depth map for the image.

[0095] In some implementations, the target machine learning task can be to perform a protein modeling task, e.g., where the model input characterizes a protein (e.g., a multiple sequence alignment for the protein), and the model output characterizes a predicted structure of the protein (e.g., a respective three-dimensional (3-D) spatial position and orientation of each amino acid in the protein structure).

[0096] In some implementations, the target machine learning task can be to perform an agent control task, where the model input represents a sequence of one or more observations or other data characterizing states of an environment and the model output defines an action to be performed by the agent in response to the most recent data in the sequence. The environment can be a real-world or a simulated environment, and the agent can be, e.g., a robot, an autonomous land, sea, or air vehicle, or a control system for an industrial facility.

[0097] In some implementations, the target machine learning task can be to process a model input that represents audio samples in an audio waveform to perform speech recognition, e.g., to generate a model output that characterizes a sequence of phonemes, graphemes, characters, or words corresponding to the audio waveform. The audio samples define a sequence of samples and optionally the model input may also include an encoding of the position of a sample in the sequence. Each dimension of the model output can, for example, correspond to a respective time interval in the audio waveform and can define a respective score distribution over a set of possible phonemes, graphemes, characters, or words.

[0098] In some implementations, the target machine learning task can be to perform a reconstruction task such as an auto-encoding task, e.g., by processing a model input to generate a model output that defines a predicted reconstruction of the model input. The model input can be, e.g., an image or video, an audio waveform, a point cloud, or a sequence of text.

[0099] In some implementations, the target machine learning task can be a neural machine translation task, e g., to process a model input that represents a sequence of text, e.g., a sequenceof words, phrases, characters, or word pieces, in one language, to generate a model output that is a translation of the sequence of text into another language, i.e., a sequence of text in the other language that is a translation of the input sequence of text.

[0100] In some implementations, the target machine learning task can be to perform an audio or audio-visual processing task. For example, if the model input represents a spoken utterance, then the model output generated by the target machine learning model 102 can be a score for each of a set of pieces of text, each score representing an estimated likelihood that the piece of text is the correct transcript for the utterance. In this example, each query embedding can be a representation of, or correspond to, a respective piece of text. As another example, if the model input represents a spoken utterance, the model output generated by the target machine learning model 102 can indicate whether a particular word or phrase (“hotword”) was spoken in the utterance.

[0101] In some implementations, the target machine learning task can be to perform a text to speech task, where the model input represents text in a natural language or features of text in a natural language and the model output is a spectrogram, a waveform, or other data defining audio of the text being spoken in the natural language.

[0102] In some implementations, the target machine learning task can be to perform a health prediction task, where the model input represents data derived from electronic health record data for a patient and the model output is a prediction that is relevant to the future health of the patient, e.g., a predicted treatment that should be prescribed to the patient, the likelihood that an adverse health event will occur to the patient, or a predicted diagnosis for the patient.

[0103] In some implementations, the target machine learning task can be to perform a text generation task, where the model input represents a sequence of text, and the model output is another sequence of text, e.g., a completion of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text that is about a topic specified by the first sequence of text. As another example, the model input can represent data other than text, e.g., an image, and the output sequence can be text that describes the data represented by the network input.

[0104] In some implementations, the target machine learning task can be to perform an image generation task, where the model input represents a conditioning input and the model output is a sequence of intensity values for the pixels of an image.

[0105] In some implementations, the target machine learning task can be to perform a genomics task, where the model input represents a fragment of a DNA sequence or other molecule sequence and the model output is either an embedding of the fragment for use in a downstream task, e.g., by making use of an unsupervised learning technique on a data set of DNA sequence fragments, or an output for the downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, and so on

[0106] In some implementations, the target machine learning task can be to perform an image, video or audio classification task. The model input can then represent the pixels of an image or video, or samples of an audio waveform as previously described. Optionally the model input can include a position encoding for the pixels of the image or video or the samples of the audio; for a video the position encoding may be a spatial and / or temporal position. The model output defines a classification of the image, video or audio. The classification can include a respective score for each object category in a set of possible object categories (e.g., for an image, vehicle, pedestrian, bicyclist, etc.). The score for an object category can define a likelihood that the model input comprises an object that belongs to the object category, e.g. that the image or video depicts the object. The classification for a video may comprise a classification of an action depicted in the video, e.g. for gesture recognition.

[0107] As previously mentioned, in general in the above examples references to an image or video include a point cloud. For example, in some implementations, the target machine learning task can be to perform a point cloud processing task, e.g., where the model input represents a point cloud (e.g., generated by a lidar or radar sensor) and the model output characterizes, e.g., a type of object represented by the point cloud.

[0108] In some implementations, the target machine learning task can be to perform an image or video captioning task. For example, the target machine learning model 102 can process a model input that includes an image or a video comprising a series of images to generate a model output that defines, for the image or for one or more of the series of images, a set of tokens that describe the image or video, where the tokens may represent words, parts of words, or sentences. The tokens may be defined deterministically, or stochastically, e.g. by sampling from one or more defined distributions.

[0109] In some implementations, the model input is a multimodal input and the target machine learning task can be to perform a multimodal task. In general such a multimodal input is a combination of two or more different types or modalities of data, where the different types of data can, but need not, represent the same or overlapping objects in the model input using the different modalities. Such multimodal data can be audio-visual data, including a combination of pixels of an image or of video and audio data representing samples of an audio waveform. As another example such multimodal data can include a combination of text data, e.g. tokens, representing text in a natural language and pixels of an image or of video or samples of an audio waveform. The multimodal task may be any of the tasks described above, but with the inclusion of an additional mode of data in the network input. As some examples the multimodal task can include a classification task, a segmentation task, a speech recognition task, a reconstruction task, a recognition task, a captioning task, and an agent control task.

[0110] In general in the above examples the model input can include raw data e.g. raw pixel data or audio samples, or feature embeddings, e.g., defining spatial, temporal or spatio-temporal features obtained by pre-processing with a feature encoder neural network.

[0111] In some implementations, the target machine learning task can be to generate data items as model outputs by processing prompts as model input. When the target machine learning task includes generating data items by processing prompts, the training examples for the target machine learning task can include example prompts and example data items.

[0112] As an example, in some implementations, the target machine learning task can include processing an input prompt to generate an output data item. The input prompt and the output data item can include any of a variety of modalities of data, e.g., text data, image data, audio data, structured numerical data, and so on. In some implementations, the input prompt and / or the output data item can include multi-modal data, e.g., data for multiple different modalities. The quality scores for the output data items can characterize a quality or a perceived quality of the output data items. For example, the quality scores for the data items can characterize, e.g., perceptual scores for the data items, human feedback regarding the data items, and so on. As another example, the output data items can be used as part of performing a downstream task and the quality scores for the data items can be performance metrics for the downstream task as attained using the output data items.

[0113] As another example, in some implementations, the target machine learning task can be a reinforcement learning task that involves controlling an agent to perform one or more agent tasks while interacting with an environment. In the context of reinforcement learning, the target machine learning model 102 can be considered to be a policy for the agent, the prompts for the target machine learning task can include observations of an environment of an agent and the output data items for the target machine learning task can characterize actions for the agent to perform the agent’s tasks. The quality scores for the output data items can be rewards associated with performance of the agent tasks by the agent.

[0114] As described above, the target machine learning model 102 can be a language model or vision language model neural network. In general, a (vision) language model neural network can be a neural network that has been trained so that, given a text prompt that includes a sequence of tokens in a natural language, the neural network can generate the next token in the sequence. This process can be repeated to extend the text prompt one token at a time to generate a natural language output, i.e., to generate the natural language output auto-regressively token by token. At each “time step,” the language model neural network processes the current sequence to generate a probability distribution over a vocabulary of tokens. The next token can then be selected using the probability distribution, e.g., by sampling from the distribution using nucleus sampling or another sampling technique or by selecting the highest-probability token. The tokens in the vocabulary can include any of a variety of tokens, e.g., some combination of words, sub-words, characters, punctuation and other symbols, and numbers. In general, the language model neural network is trained on a corpus of text made up of tokens from the vocabulary (and optionally other tokens that can be mapped to a designated out-of-vocabulary token), to predict the next token in a sequence of tokens from the training data. The (vision) language model neural network can be an autoregressive Transformer neural network.

[0115] A (vision) language model neural network can be made to perform a particular task by providing a natural language description of the desired response as an input or “prompt” (input sequence). In some cases, the prompt can be a few-shot prompt where a few, e.g., 1 to 10, examples of a query and an example output are provided in the text prior to the actual query.

[0116] A (vision) language model neural network can be “fine-tuned” to perform a particular task, by obtaining a pre-trained language model neural network trained on a large corpus of examplesas previously described and then further training part of all of the language model neural network on a relatively small number of examples particular to the type of task that is to be performed.

[0117] The target machine learning model 102 can be a large language model neural network, e.g., one that has greater than 1 billion, 10 billion or 100 billion trained parameters. The target machine learning model 102 can have been trained on greater than 10 billion, 100 billion or 1000 billion words or tokens representing words or other tokens.

[0118] The model inputs and the model outputs can be sequences of elements referred to herein as tokens. A “token” as used in this specification is a vector of numerical values having a specified dimensionality, i.e., the number of numerical values is constant across different tokens. Each token can include a respective predetermined or learned embedding (an ordered collection of numerical values having a pre-determined dimensionality.

[0119] In some implementations, the model inputs and the model outputs can include tokens representing text, e.g., words, wordpieces or characters, in a natural or computer language. For example, text can be received, e.g., as a series of encoded characters, e.g., UTF-8 encoded characters; such “characters” can include Chinese and other similar characters, as well as logograms, syllabograms and the like. A text encoder, i.e., a tokenizer, can process a sequence of text to represent the text as a series of text tokens from a vocabulary of text tokens, e.g., that each represent words, wordpieces or characters in a natural or computer language. The computer language can be any formal language used to communicate with a computer, e.g., a markup language, or a command or configuration language, or a data exchange language such as ISON, or a programming language. The tokenizer can, e.g., implement BPE (Byte Pair Encoding) or Wordpiece tokenization. Optionally the text can be obtained from audio data representing speech; the output tokens can be converted into audio data that represent speech corresponding to the text.

[0120] In some implementations, the model inputs and the model outputs can include image tokens representing images. Each image token can include a block encoding of values of the pixels in a different region of an image that maps a set of values of the pixels to a respective image token. The block encoding can be obtained using a neural network such as a Transformer neural network.

[0121] As used herein an image can be any still or moving image, i.e., the image can be part of a video, in 2D or 3D, and can be a monochrome, color or hyperspectral image, i.e., including monochrome or color pixels. As defined herein an “image” includes a point cloud, e.g., from a LIDAR system, and a “pixel” includes a point of the point cloud. An image can be captured by acamera or other image sensor from the real world; and objects in the image can include physical objects, represented by the image.

[0122] In some implementations, the model inputs and the model outputs can include tokens representing audio waveforms. For example, a set (sequence) of input or output tokens can represent audio data representing a waveform e g., instantaneous audio amplitude values or timefrequency audio data. Each audio token can include a block encoding of the audio waveform in a different time segment of the audio that maps a set of values representing the audio waveform to a respective image token. The block encoding can be obtained using a neural network such as a Transformer neural network.

[0123] In a multimodal system audio data or an image can be flagged by a start-of-audio token or start-of-image token.

[0124] In some implementations the model inputs can include tokens representing text, pixels of an image, or an audio waveform and the target machine learning model 102 can generate the output sequence of tokens to perform tasks represented by the input sequence of tokens.

[0125] In some implementations the target machine learning task can include an image or audio generation task. The input sequences of tokens can then characterize images or audio to be generated, and the output sequences of tokens can include tokens defining images or audio waveforms characterized by the input sequences of tokens, e.g., text tokens.

[0126] In some implementations the target machine learning task can include an image or audio processing task. The input sequences of tokens can define image or audio inputs, and the output sequences of tokens can include tokens defining text that describes the image or audio inputs. As some examples, the target machine learning task can include a speech recognition task, an object or action detection task, a classification task, a captioning task, a question-answering task, or a character or word recognition task.

[0127] In some implementations the target machine learning task can include a multimodal processing task in which the input sequences of tokens and / or the output sequences of tokens can include multimodal data. For example, an input sequence of tokens can characterize both an image or audio input and a text input and a corresponding output sequence of tokens can include tokens defining a result of an image or audio processing task defined by the text, such as an open vocabulary classification or object detection task.

[0128] In general, multimodal data includes a combination of two or more different types of data, e.g., two or more of audio data, image data, text data, or graph data. As one example the multimodal data can include audio-visual data, comprising a combination of pixels of an image or of video and audio data representing values of a digitized audio waveform. As another example the multimodal data can include a combination of i) text data representing text in a natural language and ii) pixels of an image or of video or audio data representing values of an audio waveform.

[0129] Some examples of multimodal tasks include: open-vocabulary image classification (the output can classify the image input based on a text input comprising text descriptions of one or more classes in the image); open -vocabulary object detection (the output can detect one or more objects in the image input based on a text input comprising text descriptions of the one or more objects); image captioning (the output can comprise text that describes the image input); text-based image search (the output can identify from amongst multiple images in the image input one or more images that meet a text description of images to be retrieved, the text description being provided in a text input); image-based retrieval (the output can identify from amongst multiple images in the image input one or more images that match a further image in the image input), and so on. The multimodal processing tasks to be performed can be defined by text in the input sequences.

[0130] In some implementations the target machine learning task can include an agent control task in which the agent interacts with an environment to perform the task. The agent can be a mechanical agent such as a robot or (semi-)autonomous vehicle, interacting with a real-world environment to perform the task. The target machine learning model 102 can be trained to control a simulated version of the agent in in a simulated version of the environment and then afterwards used to control the real agent in the real-world environment. The input sequence of tokens can include tokens that represent an observation of the environment, e.g., an image captured by a camera or other imaging device from a real-world environment. The output sequences of tokens comprises tokens that define one or more actions to be performed by the agent in the environment in response to the observation.

[0131] In some implementations the target machine learning model 102 can be stored on a user computing device, i.e., a device local to the user, such as a mobile device, e.g., a mobile phone, or a smart speaker.

[0132] In some implementations the target machine learning model 102 can be implemented on a remove server in communication with a user computing device over a wired or wireless network communications link between the user computing device and the server.

[0133] The user computing device can be provided with an input mechanism, such as a text or voice interface, that enables user input from the user in a natural language. The user computing device can be provided with an output mechanism that provides a system output for the user in the or another natural language, e.g., as speech or text; or in some other way, e.g., by displaying an image. The input and output mechanisms can include, e.g., a keyboard, microphone, speaker, display, and / or camera.

[0134] As an example, the input mechanism can include a system configured to input audio data characterizing a speech waveform of speech representing the input from the user in a natural language, and a system configured to convert the audio data into tokens representing the speech in the natural language, e.g., representing a transcription of the spoken input. The output mechanism can include a system configured to receive tokens representing the output for the user in the or another natural language and a system configured to convert the received tokens into audio data representing a waveform of speech representing the output to the user in the natural language, i.e., representing spoken words.

[0135] As a further example, the target machine learning model 102 can be deployed in an environment that enables a user to provide a request for the system, e.g., to process a multimodal input to generate a corresponding output sequence. Users can provide requests, e.g., by way of a user interface or through an application programming interface (API). The request can be transmitted from a user device, e.g., over a data communications network such as the internet, to one or more computers implementing the system, e.g., in a data center. The system can generate a output sequence and then transmit the output sequence to a user device over a data communications network.

[0136] A user computing device can be provided, as an interface for the target machine learning model 102, with an input mechanism that enables user input from the user in a natural language and an output mechanism that provides a system output to the user in the natural language. The input and output mechanism can include, e.g., a keyboard and display. Also or instead the input and output mechanism can include a speech-based mechanism. For example, the input mechanism can include a system configured to input audio data characterizing a speech waveform of speechrepresenting the input from the user in the natural language and configured to convert the audio data into tokens representing the speech in the natural language, e.g., representing a transcription of the spoken input. The output mechanism can include a system configured to receive tokens representing the output to the user in the natural language and a system configured to convert the received tokens into audio data representing a waveform of speech representing the output to the user in the natural language, i.e., representing spoken words.

[0137] In some implementations the input sequences include one or more natural language statements relating to an environment, in particular a real-world environment, and include natural language requests relating to the environment. Similarly the output sequences can include natural language replies or natural language output statements that also relate to the environment, i.e., providing information relating to the environment, in some implementations relating to or specifying actions to be taken in the environment.

[0138] The target machine learning model 102 can be used for diagnosing faults, or for correcting undesired behavior, in a mechanical or computing system operating in the real world environment. The model inputs can include descriptions and / or images of observations of the mechanical or computing system, e.g., of operation of the system, optionally obtained from one or more sensors sensing a condition or operation of the system. An image observation can be converted into a text description, e.g., using an image captioning system or in other ways. The generated output sequences can include images, audio, or text that identify (describe) likely causes of the faults or undesired behavior. This can be used to repair the faults or correct the behavior. The preference measures for the target machine learning task can define relatively more useful types of output for repairing faults or correcting behavior, and other aspects of the responses as previously described.

[0139] The target machine learning model 102 can be used for controlling a mechanical agent such as a robot or vehicle. For example, the model inputs can include descriptions of tasks to be performed, and the generated output sequences can include lists of sub-tasks to be performed by the mechanical agent (trained to perform such sub-tasks), in order to perform the tasks. The preference measures for the target machine learning task can define relatively more preferable or useful types of sub-task, task safety, efficiency, and so on.

[0140] As another example, the environment can be a computer security monitoring environment, e.g., the system can be deployed as part of a system that monitors the security of one or more computers. For example, the environment can be a computer network security monitoringenvironment, and the system can be deployed as part of a system that monitors the security of one or more computers on a computer network, e.g., a wireless network, a cellular network, a local area network and / or the internet. As another example, the environment can alternatively or additionally be a computer system security monitoring environment and the system can be deployed as part of a system that monitors the system for the presence of computer viruses and / or an unresolved software vulnerability, e.g., a zero-day exploit. A software vulnerability can be resolved by updating the software (e.g., patching) and / or removing (e.g., uninstalling) the software from the computer system. In these examples, the natural language requests can query whether computer security incidents have been resolved (e.g., “has the incident been resolved?”) and the model inputs can include relevant statements from system logs, i.e., that are potentially relevant to the events being queried. A computer security incident can be, e.g., a data breach, an unauthorized log-in or other access of a secured system, a detection of a computer virus or detection of a software vulnerability. An incident can be “resolved” when the underlying incident is no longer a threat to the security of the computer system e.g., the computer virus has been removed, the access to the secured system has been removed, the data breach has been mitigated, or the software having the vulnerability has been updated or removed. The system can use the model inputs 204 to generate replies to the requests that include natural language statements indicating whether the incidents have been resolved, optionally displaying evidence used to determine this.

[0141] The model inputs can include one or more of: code snippets from the software code, system logs, program logs, or other artifacts that should be left on the computer by running the program, or verification rules that represent requirements for the execution of the software program, or natural language statements describing the computer system on which the software executes. In general, the model inputs can include relevant statements, i.e., statements that are potentially relevant to the events being queried.

[0142] In some implementations obtaining input data from the environment can include obtaining, from the system logs, the data characterizing the computer network, or both, or from other data as described above, one or more observations of the computer network (which here includes computers on the network), and processing the one or more observations to generate a natural language representation of the one or more observations. The natural language requests can relate to the computer security incidents or to the secure operation of the computer network. The target machine learning task can include using the natural language representations of the one or moreobservations to provide one or more of the natural language statements describing the computer network, and using the natural language replies or the natural language output statements to identify a security status of the computer network or a security flaw in the computer network.

[0143] As another example, the environment can be a software testing or evaluation environment, e.g., the system can be deployed as part of a system that tests software before deployment or that evaluates already-deployed software to identify bugs. In these examples, when the system tests software before deployment, the natural language requests can ask whether the software will execute as intended, and the model inputs can include code snippets from the software code and, optionally, natural language statements describing the computer system on which the software will execute. The target machine learning model 102 can process the model inputs to generate replies that indicate whether the code will execute as intended, optionally displaying evidence used to determine this. When the system monitors the execution of code after deployment, the natural language requests can ask whether a software program, or a portion of a software program, has executed as intended, and the model inputs can include one or more of: code snippets from the software code, system logs, program logs, or other artifacts that should be left on the computer by running the program, or verification rules that represent requirements for the execution of the software program, or natural language statements describing the computer system on which the software executes. The model 102 can then process the model inputs to generate replies that indicate whether the code has executed as intended, optionally displaying evidence used to determine this. As a particular example, the software program can be part of the boot up of a computer, and the model 102 can generate a reply each time that the computer starts up to verify whether the computer will function correctly after start up.

[0144] As another example, the environment can be an educational environment, e.g., the system can be deployed as part of an education software program that assists a user in learning or practicing one or more corresponding skills. In these examples, the model inputs can include natural language statements describing or referencing a scenario or scene in a real-world or imagined environment, and the requests can be questions about the scenario or scene.

[0145] As another example, the environment can be an information retrieval environment, e.g., the system can be deployed as part of a search engine or other software that allows a user to search for information in a corpus of documents, e.g., the Internet or another electronic document corpus. In these examples, the requests can be any appropriate natural language question, and the replies canoptionally include evidence such as include relevant statements from the corpus of documents, e.g., as identified by searching the corpus using conventional information retrieval techniques.

[0146] In some implementations, the target machine learning model 102 is a visual language model (VLM). In general, the VLM can process input sequences that include tokens that each represent natural language or (a part of) an image or video to generate output tokens that each represent natural language or (a part of) an image or video. For example, the VLM can be configured to describe an image or video using natural language, e.g., to perform an image or video captioning task. As another example, the VLM can be configured to process input tokens representing an image and text tokens representing a query about the image or a request to modifying the image, and to generate output tokens representing an answer to the query or representing a version of the image that has been modified in accordance with the request. The VLM can generate output tokens representing an image or video that is generated in response to input tokens providing a visual and / or audio and / or textual description of a desired image or video.

[0147] In some implementations, the “language” of the language model is not a natural language such (e.g., English), but can instead be a text-based encoding describing an entity or class of entities, e.g., a chemical or biological entity, such as a chemical structure or molecule. For example, the text-based encoding can be a sequence of tokens that defines a molecule or protein, e.g., a sequence specifying an arrangement of atoms or chemical functional groups in a molecule, or the amino acid residues of a protein. The language model can be referred to as a chemical and / or biological language model in such cases. The model inputs therefore be input strings defining chemical (e.g., protein) structures and the model outputs can include output strings defining different chemical structures from the input strings. The strings can be in the Simplified Molecular Input Line Entry System, SMILES, format, for example.

[0148] In another example of a computer language text generation task, a model input can include an image or video and a sequence of text in a computer language for performing a task in relation to the image or video, e.g., a data processing task that involves analyzing the content of the image or video to provide a result of the analysis or, e.g., a search to search for information relating to the content of the image or video. The computer language in the model output can include computer language for invoking a function or calling one or more external APIs. Merely as one example, such an output can be formatted as a JSON object. As previously, the sequence of text in a multimodal input can define the task to be performed and the second modality input caninclude, e.g., an image or video in relation to which the task is to be performed, e.g., a task that involves manipulation of particular types of data that can benefit from access to an API such as mathematical data, date / time related data, scientific data, recent data that can post-date training of the model (that can be accessed by a search function or API), and so on. After training, when the model is used in inference, the model output can include text in the or another computer language for performing a task, e.g., as described above, in relation to an image or video in the second modality input. The target machine learning task can then include using the text in the computer language to perform the task.

[0149] In some implementations, the target machine learning model 102 can be used to interact with a human user of a digital assistant such as a smart speaker, smart display, or other device. For example, information defining a task can be obtained from the digital assistant, and the digital assistant can be used to instruct the user to perform the task. For example, this can include receiving, at the digital assistant, a request from the user for assistance and determining, in response to the request, a series of tasks for the user to perform, e.g., steps or sub-tasks of an overall task. Then for one or more tasks of the series of tasks, e.g., for each task, e.g., until a final task of the series the digital assistant can be used to output to the user an indication of the task, e.g., step or sub-task, to be performed. This can be done using natural language, e.g., on a display and / or using a speech synthesis subsystem of the digital assistant. Visual, e.g., video, and / or audio observations of the user performing the task can be captured, e.g., using the digital assistant. A system can then be used to determine whether the user has successfully achieved the task, e.g., step or sub-task, i.e., from the answer as previously described. If there are further tasks to be completed the digital assistant can then, in response, progress to the next task (if any) of the series of tasks, e.g., by outputting an indication of the next task to be performed. In this way the user can be led step-by-step through a series of tasks to perform an overall task.

[0150] As an illustrative example, a user can be interacting with a digital assistant and ask for help performing an overall task consisting of multiple steps, e.g., cooking a pasta dish. While the user performs the task, the digital assistant receives audio and / or video inputs representative of the user’s progress on the task, e.g., images or video or sound clips of the user cooking. The digital assistant uses model 102 as described above, in particular by providing it with the captured audio and / or video and a question that asks whether the user has completed a particular step, e.g., ‘Has the user finished chopping the peppers?’, to determine whether the user has successfully completedthe step. If the answer confirms that the user has successfully completed the step then the digital assistant progresses to telling the user to perform the next step or, if at the end of the task, or if the overall task is a single-step task, then the digital assistant can indicate this to the user. The digital assistant can then stop receiving or processing audio and / or video inputs to ensure privacy and / or reduce power use.

[0151] The digital assistant can also include a user interface to enable a user to request assistance and to output information. In implementations this is a natural language user interface and can include a keyboard, voice input-output subsystem, and / or a display. The digital assistant can further include an assistance subsystem configured to determine, in response to the request, a series of tasks for the user to perform. In implementations this can include a generative (large) language model, in particular for dialog. The digital assistant can have an observation capture subsystem to capture visual and / or audio observations of the user performing a task; and an interface for the above-described language model neural network (which can be implemented locally or remotely). The digital assistant can also have an assistance control subsystem configured to assist the user. The assistance control subsystem can be configured to perform the steps described above, for one or more tasks e.g., of a series of tasks, e.g., until a final task of the series. More particularly, the assistance control subsystem can output to the user an indication of the task to be performed, capture, using the observation capture subsystem, visual or audio observations of the user performing the task, determine from the above-described answer whether the user has successfully achieved the task. In response, the digital assistant can progress to a next task of the series of tasks and / or control the digital assistant, e.g., to stop capturing observations.

[0152] In some implementations, a particular task that is to be performed by the target machine learning model 102 can be described by part or all of a sequence of text in an input to the model 102. For example, in a model input that includes an image such a prompt can specify, e.g., “Generate a caption”, “Generate a description”, “Answer the following question: [about the image or video]”, or “Detect a person”. Where the model 102 is used for an agent control task a prompt can define, e g., “Take the knife out of the drawer”, or “Q: What action should the robot take to take the knife out of the drawer?”. Also or instead such a prompt can give one or more examples of a task to be performed. The model 102 can be trained on multiple natural and / or computer languages and the prompt can then specify a language to use.

[0153] A few further examples of some machine learning tasks that can be performed by a target machine learning model 102 trained as described herein follow. The tasks described below can include tasks that require spatial awareness or other context from input images or video. For example, a prompt can ask “What is the object in the top left corner?”.

[0154] In general, for the tasks below the model 102 can have been trained or fine-tuned on examples of the input and output for the task. For example, the model 102 can have been trained using still or moving images containing one or more objects or actions, and corresponding sequences of text or other data e.g., describing or classifying the images. However large, “foundation” models can, in general, perform some tasks zero-shot, i.e., without having been specifically trained on those tasks.

[0155] As one example the task can include an object or action detection task. For example, a generated output sequence can include or represent text that describes or otherwise labels detected object(s) or action(s) in an input that includes an image or audio, and can include coordinates such as bounding-box coordinates for the detected object(s) or action(s), e.g., “10 20 90 100 cat 20 30 100 100 dog”.

[0156] As another example the task can include a classification task, e.g., an object or action classification task. A generated output sequence can include data, e.g., text, that classifies the object(s) or action(s) in represented in conditioning data, e.g., in an image or audio, into one of a plurality of classes, or that otherwise classify object(s) or action(s) represented in the conditioning data.

[0157] As another example the task can include a still or moving image describing task, e.g., a captioning task (which, as used here, includes an audio description task to explain what is happening in an image). A generated output sequence can include data, e.g., text, describing an input image or video. For example, a generated output sequence can provide a caption or description or it can count objects in the image or video, or it can provide some other form of description.

[0158] As another example the task can include a still or moving image question-answering task. A generated output sequence can include data, e.g., text, that answers a question about an input, e.g., an input image or audio, where the question is also specified in the input, e.g., as sequence of text. This can be used, e.g., to answer questions about visual plots and charts or about sounds.

[0159] As another example the task can include a character or word recognition task, e.g., an OCR (optical character recognition) task. An input can include a still or moving image and a generated output sequence can include text that represents characters or words in the input, e.g., in a natural language.

[0160] As another example the task can include a still or moving image generation task. A generated output sequence can include image data defining values for pixels of a still or moving image, and an input, e.g., a sequence of text, can describe or characterize the image to be generated. Merely as an example, an image of a plot or chart can be generated to represent the input, e.g., comprising text.

[0161] As another example the task can include a computer language text generation task. An input can include a natural language description of a task to be performed, and optionally an image (if the task is to be performed on or in relation to an image), and a generated output sequence can include text in a computer language to perform the task, e.g., a task of analyzing the content of the image to provide a result of the analysis or to search for information relating to the content of the image.

[0162] As a particular example the computer language in a generated output sequence can include computer language for invoking a function or calling one or more external APIs. Merely as one example, such an output sequence can include data formatted as a JSON object. As previously, an input can define a task to be performed and can also include an image in relation to which the task is to be performed. In general the task can involves manipulation of particular types of data that can benefit from access to an API such as mathematical data, date / time related data, scientific data, recent data that can post-date training of the model 102 (that can be accessed by a search function or API), and so on; and the generated output sequence can include text in a computer language for performing the task. The target machine learning task can include using the text in the computer language to perform the task.

[0163] In general where a generated output sequence includes text, such text can be converted to speech representing the text, and an audio (speech) output provided.

[0164] In some implementations the task can include an agent control task in which the agent interacts with an environment to perform the agent control task. In these implementations an input can include an observation characterizing the environment. For example, an input can include a sequence of text that defines a task to be performed by the agent and an image representing anobservation of the environment, e.g., captured by a camera or other imaging device from a real- world environment. A generated output sequence can include an action selection output, e.g., including text, that is used to select one or more actions to be performed by the agent in the environment in response to the observation. As an illustration the generated output sequence can define an action as text such as “A: 132 114 128 5 25 156”, that can be converted into a control signal for a mechanical agent, such as a robot, e.g., “AT = [0.1, —0.2,0] A / ? = [10°, 25°, —7°]”. The action selection output can also or instead define one or more low-level skills, e.g., from a vocabulary of previously learnt skills. As before, a sequence of text in a model input can describe the task to be performed, e.g., “What action should the robot take to [perform task]”. Examples of systems for controlling an agent that can be fine tuned as described herein can include PaLM-E (Driess, et al., arXiv:2303.03378), RT-1 (Brohan, et al., arXiv:2212.06817), and RT-2 (Brohan, et al., arXiv:2307.15818).

[0165] In some agent control implementations, the environment is a real-world environment and the agent is a mechanical agent interacting with the real-world environment, e.g., a robot or an autonomous or semi-autonomous land, air, or sea vehicle operating in or navigating through the environment, and the actions are actions taken by the mechanical agent in the real-world environment to perform the task. For example, the agent can be a robot or other mechanical agent interacting with the environment to accomplish a specific task, e.g., to locate or manipulate an object of interest in the environment or to move an object of interest to a specified location in the environment or to navigate to a specified destination in the environment. In these implementations, the observations can include, e.g., one or more of: images, object position data, and sensor data to capture observations as the agent interacts with the environment. The actions can define control signals to control the robot or other mechanical agent, e.g., positions, torques, or other control signals for the parts of the mechanical agent, or higher-level control commands.

[0166] In some agent control implementations the agent can be a human agent and the environment can be a real-world environment. For example, the agent can be a human user of a digital assistant such as a smart speaker, smart display, or some other device that is used to instruct the user to perform actions. The task can include any real-world task that the user wishes to perform. The observations can be obtained from an observation capture subsystem, e.g., a monitoring system such as a video camera or sound capture system, to capture visual observationsof the user performing the task. The actions can include instructions in the form of, e.g., text, image, video, or audio data such as speech, that guide the user in performing the task.

[0167] The described systems and techniques can be applied to a wide range of different types of input sequences and output sequences. In implementations of the described techniques the tokens can represent, characterize, or encode any type of information in a sequence, e.g., stream of data. The term “represent” is used, below, generally to refer to any way in which a token can encode part of a sequence. The tokens can include marker tokens, such as a start of sequence token, an end of sequence token, and a separator token (indicating a separation or break between two distinct parts of a sequence). The tokens can be, e.g., drawn from a defined vocabulary of tokens.

[0168] Some of these implementations can be used for natural language tasks such as providing a natural language response to a natural language input, e.g., for question answering, or for text completion. In some implementations the input sequence can represent text in a natural language and the output sequence can represent text in the same natural language, e.g., a longer item of text. For example, in some implementations the input sequence can represent text in a natural language and the output sequence can represent the same text with a missing portion of the text added or filled in. For example, the output sequence can represent a predicted completion of text represented by the input sequence. Such an application can be used, e.g., to provide an autocompletion function, e.g., for natural language-based search. In some implementations the input sequence can represent a text in a natural language, e.g., posing a question or defining a topic, and the output sequence can represent a text in a natural language which is a response to the question or about the specified topic.

[0169] As another example the input sequence can represent a first item of text and the output sequence can represent a second, shorter item of text, e.g., the second item of text can be a summary of a passage that is the first item of text. As another example the input sequence can represent a first item of text and the output sequence can represent an aspect of the first item of text, e.g., it can represent an entailment task, a paraphrase task, a textual similarity task, a sentiment analysis task, a sentence completion task, a grammaticality task, and in general any natural language understanding task that operates on a sequence of text in some natural language, e.g., to generate an output that classifies or predicts some property of the text. For example, some implementationscan be used to identify a natural language of the first item of text, or of spoken words where the input is audio (as described below).

[0170] Some implementations can be used to perform neural machine translation. Thus in some implementations the input tokens can represent words, wordpieces, or characters in a first natural language and the output tokens can represent words, wordpieces or characters in a second, different natural language. That is, the input sequence can represent input text in the first language and the output sequence can represent a translation of the input text into the second language.

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

[0172] Some implementations can be used for speech recognition. In such applications the input sequence can represent spoken words and the output sequence can represent a conversion of the spoken words to a machine- written representation, e.g., text. Then the input tokens can include tokens representing an audio data input including the spoken words, e.g., characterizing a waveform of the audio in the time domain or in the time-frequency domain. The output tokens can represent words, wordpieces, characters, or graphemes of a machine-written, e.g., text, representation of the spoken input, that is representing a transcription of the spoken input.

[0173] Some implementations can be used for handwriting recognition. In such applications the input sequence can represent handwritten words, syllabograms or characters and the output sequence can represent a conversion of the input sequence to a machine- written representation, e.g., text. Then the input tokens can include tokens representing portions of the handwriting and the output tokens can represent words, wordpieces, characters or graphemes of a machine-written, e.g., text, representation of the spoken input.

[0174] Some implementations can be used for text-to- speech conversion. In such applications the input sequence can represent text and the output sequence can represent a conversion of the text to spoken words. Then the input tokens can include tokens representing words or wordpieces or graphemes of the text and the output tokens can represent portions of audio data for generating speech corresponding to the text, e.g., tokens characterizing a portion of a waveform of the speech in the time domain or in the time-frequency domain, or phonemes.

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

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

[0177] In some implementations the input sequence and the output sequence represent different modalities of input. For example, the input sequence can represent text in a natural language and the output sequence can represent an image or video corresponding to the text; or vice-versa. In general, the tokens can represent image or video features and a sequence of such tokens can represent an image or video. There are many ways to represent an image (or video) using tokens. As one example an image (or video) can be represented as a sequence of regions of interest (Rols) in the image, optionally including one or more tokens for global image features. For example, an image can be encoded using a neural network to extract Rol features; optionally (but not essentially) a token can also include data, e.g., a position encoding, representing a position of the Rol in the image. As another example, the tokens can encode color or intensity values for pixels of an image. As another example, some image processing neural network systems, e.g., autoregressive systems, naturally represent images as sequences of image features. As another example, a transformer-based sequence processing neural network system as previously described can be used to process images instead of or as well as text (e.g., if trained on images instead of or as well as text).

[0178] Thus in some implementations at least one of the input sequence and the output sequence is a sequence representing an image or video and can include tokens representing the image or video. For example, the input sequence can be a sequence of text, the input tokens can represent words, wordpieces, or characters and the output sequence can include output tokens representingan image or video, e.g., described by the text, or providing a visual answer to a question posed by the text, or providing a visualization of a topic of the text. In another example the input sequence can include a sequence of input tokens representing an image or video, and the output tokens can represent words or wordpieces, or characters representing text, e.g., for a description or characterization of the image or video, or providing an answer to a question posed visually by the image or video, or providing information on a topic of the image or video.

[0179] In some other implementations both the input sequence and the output sequence can represent an image or video, and both the input tokens and the output tokens can represent a respective image or video. In such implementations the method / system can be configured to perform an image or video transformation. For example, the input sequence and the output sequence can represent the same image or video in different styles, e.g., one as an image the other as a sketch of the image; or different styles for the same item of clothing.

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

[0181] In some implementations the input sequence represents a sequence of actions to be performed by an agent, e.g., a mechanical agent in a real-world environment implementing the actions to perform a mechanical task. The output sequence can include a modified sequence of actions, e.g., one in which an operating parameter, such as a speed of motion or power consumption, has a limited value; or one in which or safety or other boundary is less likely to be crossed. Then both the input tokens and the output tokens can represent the actions to be performed.

[0182] In some implementations the input sequence represents a sequence of health data and the output sequence can include a sequence of predicted treatment. Then the input tokens can represent any aspect of the health of a patient, e.g., data from blood and other medical tests on the patient and / or EHR (Electronic Health Record) data; and the output tokens can represent diagnostic information, e.g., relating to a disease status of the patient and / or relating to suggested treatments for the patient, and / or relating to a likelihood of an adverse health event for the patient.

[0183] As a particular example the model 102 can be a multimodal model neural network in which one or both of the model input (i.e., input sequence) and the model output (i.e., output sequence) include an image or audio. For example the multimodal machine learning model can be configured to process an input sequence including visual tokens representing pixels of a still or moving image (which here can include a point cloud image), and / or data representing an audio waveform, e.g., values or features of the audio waveform such as audio tokens, and / or text tokens representing a sequence of text, to generate an output sequence, e.g., including text tokens representing the still or moving image or audio waveform, and / or a sequence of intensity value inputs for the pixels of an image or a sequence of values defining an audio waveform. A visual token can, e.g., represent multiple pixels in a region of the image, e.g., as features of the region. Such a multimodal model 102 can perform any of the previously described tasks, e.g., using a multimodal input, or by providing a multimodal output, or by converting between different input and output modes (e.g., text / image / audio). For example, it can generate text representing, describing (e.g., captioning), or otherwise characterizing an image or audio input, e.g., by answering a question related to the image or audio input, e.g., relating to a future, e.g., physical prediction of a state of objects represented by the image or audio. As another example it can generate an image or audio represented, described, or otherwise characterized by a text input, or otherwise in response to the text input, e.g., representing an image or audio answer to a text question.

[0184] FIG. 2 illustrates updating model parameters for a target machine learning model by iteratively interpolating trained model parameters for a set of auxiliary machine learning models.

[0185] As described above with reference to FIG. 1, a training system (e.g., the training system 100 of FIG. 1) can train the target machine learning model by, at each interpolation iteration, training a plurality of auxiliary machine learning models to perform the target machine learning task and updating the model parameters of the target machine learning model by interpolating the trained model parameters of the auxiliary machine learning models. For illustrative purposes, FIG. 2 depicts training the target machine learning model using two auxiliary machine learning models, however the training system can, in general, utilize any number of auxiliary machine learning models as part of training the target machine learning model.

[0186] For example, at a first interpolation iteration, the current model parameters Qinitare set to initial model parameters Gnit(e.g., as obtained by pre-training the target machine learning model to perform one or more pre-training tasks). For the first interpolation iteration, the training systemcan initialize the current model parameters Qinitof the auxiliary machine learning models to be the initial model parameters, dinit, of the target machine learning model. The training system can then train each of the auxiliary machine learning models to perform the target machine learning task to obtain trained model parameters for the auxiliary machine learning models. For example, as illustrated in FIG. 2, the training system can use a reinforcement learning technique (e.g., the REINFORCE algorithm, using a policy gradient technique, etc.) to train the auxiliary machine learning models and obtain trained model parameters, O7.(and 0^, for the auxiliary machine learning models.

[0187] The training system can regularize the training of the auxiliary machine learning models using a regularization based on respective anchor sets of model parameters for the auxiliary machine learning models. For example, as illustrated in FIG. 2, the training system can utilize Kullback-Leibler (KL) regularization with the anchor sets of model parameters as part of training the auxiliary machine learning models.

[0188] The training system can determine anchor set of model parameters for each auxiliary machine learning model as an exponential moving average (EMA) of the model parameters for the auxiliary machine learning models over a sequence of training steps for the auxiliary machine learning models. For example, as illustrated in FIG. 2, the training system can utilize the exponential moving averages 9gmaand 9gmaas the anchor sets of model parameters for the auxiliary machine learning models.

[0189] After training the auxiliary machine learning models, the training system can determine the interpolated parameters for the target machine learning model for the interpolation iteration based on an interpolation of the trained model parameters of the auxiliary machine learning models. For example, as illustrated in FIG. 2, the training system can interpolate the trained model parameters, 0,1.;and 072Z, for the auxiliary machine learning models using spherical linear interpolation (SLERP) to determine an interpolated set of model parameters, 9sierp.

[0190] The training system can use the interpolated set of model parameters, 0s;erp, to determine an updated set of model parameters, 0-„7t, for the target machine learning model. For example, the system can determine the updated set of model parameters, 0-n7t, for the target machine learning model as a weighted combination of the interpolated set of model parameters, 9sierp, and the current model parameters, 0init, for the target machine learning model. For example, as illustrated in FIG. 2, the training system can determine the updated set of model parameters, 0^,by performing a linear interpolation of the parameters Qsierptowards the current parameters Qinit(LITI).

[0191] The training system can use the updated model parameters, Q'init, as the current model parameters for training the auxiliary model parameters during a next interpolation iteration.

[0192] After one or more interpolation iterations, the training system can obtain a final set of interpolated model parameters, Qs'ierp. The training system can use the final set of interpolated model parameters, Qs'ierp, to determine the final trained model parameters for the target machine learning model. For example, the system can determine the final trained model parameters for the target machine learning model as a weighted combination of the final set of interpolated model parameters, Qs'ierp, and the initial model parameters, 0*nit, for the target machine learning model.

[0193] An example process by which the training system can train the target machine learning model over the sequence of interpolation iterations is described in more detail below with reference to FIG. 3.

[0194] FIG. 3 is a flow diagram of an example process 300 for training a target machine learning model to perform a target learning task over a sequence of interpolation iterations. 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 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 300.

[0195] The system can obtain training data for the target machine learning task and an initial set of machine learning model parameters for the target machine learning model (step 302). Each training example can include a training input (e.g., an example prompt) for the training example. In some implementations, each training example can include a target model output (e.g., an example output data item) for the training example. In some implementations, each training example can include a quality score for the training example that measures a quality of the target model output for the training example.

[0196] As described above with reference to FIG. 1, the training inputs and the target outputs can include any of a variety of modalities of data, e.g., text data, image data, audio data, structured numerical data, and so on. In some implementations, the training inputs and / or the target model outputs can include multi-modal data, e.g., data for multiple different modalities.

[0197] The quality scores for the target outputs can characterize a quality or a perceived quality of the target outputs. For example, the quality scores for the target outputs can characterize, e.g., perceptual scores for the target outputs, human feedback regarding the target outputs, and so on. As another example, the target outputs can be used as part of performing a downstream task and the quality scores for the target outputs can be performance metrics for the downstream task as attained using the target outputs.

[0198] In some implementations, the target machine learning task can involve interacting with a user to perform the task, e.g., by generating responses to queries received from the user. When the machine learning task involves interacting with a user, the training inputs can include examples of queries from an example user and target outputs can include example responses to the examples of queries from the example user.

[0199] In some implementations, the machine learning task can be to select actions for an agent interacting with an environment to perform a task in the environment. The target output for a training example can include example selected actions for an example agent to perform the task in an example environment for the training example. The training input for the training example can include example observations of the example environment for the training example.

[0200] The quality scores for the training examples can be performance measures for the machine learning task. When the machine learning task includes selecting actions for an agent interacting with an environment to perform a task in the environment, the quality scores for the training examples can be performance measures for the task in the environment.

[0201] In some implementations, the quality scores for the training examples can be generated by a reward model processing the training examples. For example, the quality scores for the training examples can be generated by processing data characterizing the target outputs using a language model along with a prompt requesting the language model to evaluate the target outputs.

[0202] The initial machine learning model parameters for the target machine learning model can be pre-trained parameters for the target machine learning model, e.g., as obtained by pre-training the target machine learning model to perform one or more pre-training tasks. For example, the system can obtain the initial machine learning model parameters for the target machine learning model by pre-training the target machine learning model to perform the one or more pre-training tasks. In one example, the initial machine learning model parameters 0for the target machinelearning model can be set to dSFT, the parameters resulting from a supervised fine tuning stage. As another example, the system can obtain the initial machine learning model parameters for the target machine learning model from another system that has pre-trained the target machine learning model to perform the one or more pre-training tasks.

[0203] The system can train the target machine learning model to perform the target machine learning task by updating the machine learning model parameters of the target machine learning model over a sequence of one or more interpolation iterations. At each interpolation iteration, the system can perform steps 304 through 308 described below. For the first interpolation iteration, the current model parameters Qinitare set to the initial model parameters

[0204] At each interpolation iteration, the system can initialize the machine learning model parameters for each of a plurality of auxiliary machine learning models using the current set of machine learning model parameters for the target machine learning model (step 304). In particular, the system can initialize the machine learning model parameters for each of the auxiliary machine learning models to be the current set of machine learning model parameters for the target machine learning model as of the interpolation iteration.

[0205] The system can train each of the auxiliary machine learning models to optimize an obj ective function for the target machine learning task using the training data for the target machine learning task (step 306). An example process by which the system can train each of the auxiliary machine learning models is described in more detail below with reference to FIG. 4. By training the auxiliary machine learning models, the system can obtain a respective set of trained machine learning models parameters for each of the auxiliary machine learning models. In some implementations, the distributed nature of the training of the auxiliary machine learning models means that parallelization of the fine-tuning can be implemented, for example with training of auxiliary machine learning models on different servers and intermittent weight sharing across servers.

[0206] In some implementations, one or more diversity strategies can be employed to increase diversity in weights across the auxiliary machine learning models. For example, the order in which text prompts x are given can be varied. The hyperparameters can be varied. The objective function can be varied. For example, the auxiliary machine learning models can be trained with diverse rewards. For example, one or more auxiliary machine learning models can be trained with a length penalty integrated into the reward and the remaining one or more auxiliary machine learningmodels can be trained without the length penalty. Increased diversity in weights and predictive mechanisms across the auxiliary machine learning models can be beneficial for generalization.

[0207] The system can update the machine learning model parameters of the target machine learning model by interpolating the trained machine learning model parameters of the auxiliary machine learning models (step 308). In particular, the system can determine an interpolated set of machine learning model parameters for the interpolation iteration by interpolating the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration.

[0208] For example, in some implementations, the system can determine the interpolated set of machine learning model parameters by performing linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration.

[0209] As another example, the system can determine the interpolated set of machine learning model parameters by performing a spherical linear interpolation (SLERP) of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration. The system can perform a spherical linear interpolation between two sets of machine learning model parameters, Qi and 02, based on the current machine learning model parameters for the target machine learning model, Qintt, following:

[0210] Where A is an interpolation weight for the spherical linear interpolation, Sxis a task vector defined following 8X= 0X— 0jnjt, 82is a task vector defined following 62= 02— dinit, and fl is an angle between the task vectors 81and 82.

[0211] For example, fl can be determined following:

[0212] In general, system can perform a spherical linear interpolation of M sets of machine learning model parameters (e.g., with M > 2), {0m}“=i, following the recursively defined interpolation:

[0213] Equivalently, the system can determine the spherical linear interpolation of M > 2 sets of machine learning model parameters {9m}“=1, following the recursively defined interpolation:

[0214] In some implementations, the system can perform a layer-wise spherical linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration by performing a separate spherical linear interpolation of the machine learning model parameters for each layer of the target machine learning model.

[0215] The system can update the machine learning model parameters of the target machine learning model using the interpolated set of model parameters by any of a variety of methods. For example, in some implementations, the system can update the machine learning model parameters of the target machine learning model to be the interpolated set of model parameters.

[0216] As another example, in some implementations, the system can update the current set of machine learning model parameters for the target machine learning model as an exponential moving average of the interpolated set of machine learning model parameters for each of the interpolation iterations. For example, during an i-th interpolation iteration, the system can determine the updated machine learning model parameters, 9-^, for the target machine learning model based on the current machine learning parameters, 9-nit, and the interpolated set of model parameters, Qinterp, following:

[0217] Where r| is a weight for the exponential moving average. In an example, r| E [0, 1], In one example, | = 0.3. As described above, the weight r| can be selected to balance a trade-off between attaining better performance on the target machine learning task and maintaining pre-trained capabilities of the target machine learning model.

[0218] An example algorithm that the system can use to update the parameters of the target machine learning model is illustrated below with reference to FIG. 5.

[0219] The system can determine whether updating of the target machine learning model is complete (step 310). The system can use any of a variety of criteria to determine whether the updating of the target machine learning model is complete. For example, the system can determine that the updating of the target machine learning model is complete after a pre-determined numberof interpolation iterations. As another example, the system can determine that the updating of the target machine learning model is complete when an average value of the objective function falls below a pre-determined threshold when training the auxiliary machine learning models. As another example, the system can determine that the updating of the target machine learning model is complete when a difference between average values of the objective function for the current training iteration and a previous training iteration falls below a pre-determined threshold while training the auxiliary machine learning models.

[0220] If the system determines that updating the target machine learning model is not complete, the system can continue to a next interpolation iteration (e.g., return to step 304)

[0221] When the system determines that the updating of the target machine learning model is complete, the system can return the trained target machine learning model (step 312). In particular, after the final interpolation iteration of the sequence of interpolation iterations, the system can determine the final trained set of machine learning model parameters for the trained target machine learning model based on the current set of machine learning model parameters for the target machine learning model as of the final interpolation iteration by any of a variety of methods. For example, in some implementations, the system can determine the final trained set of machine learning model parameters to be the current set of machine learning model parameters for the target machine learning model as of the final interpolation iteration. As another example, the system can determine the final trained set of machine learning model parameters, trained, for the target machine learning model as a weighted combination of the current set of machine learning model parameters for the target machine learning model as of the final interpolation iteration, OjA*, and the initial set of machine learning model parameters for the target machine learning model, 9^lZt, following:

[0222] Where 0 < q < 1 is a weight between the initial and final machine learning model parameters of the target machine learning model over the sequence interpolation iterations. When the target machine learning model is pre-trained, the initial set of machine learning model parameters for the target machine learning model, S(nit, are the pre-trained model parameters for the target machine learning. The weight, q, can be selected to balance a trade-off between attaining better performance on the target machine learning task and maintaining pre-trained capabilities of the target machine learning model.

[0223] FIG. 4 is a flow diagram of an example process 400 for training an auxiliary machine learning model to perform a target machine learning task. For convenience, the process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 400.

[0224] The system can train the auxiliary machine learning model over a sequence of training steps. As part of training the auxiliary machine learning model, the system can perform steps 402 through 408 at each training iteration.

[0225] For each of a plurality of training examples for the training step, the system can process the training input for the training example using the auxiliary machine learning model to generate a corresponding model output for the training example (step 402).

[0226] For example, in some implementations, the system can process the training input for each training example using the auxiliary machine learning model to generate an output data item for the training example. When the auxiliary machine learning model is a generative model, the system can determine, for each training example, a likelihood,of the auxiliary machine learning model generating the output data item y for the training example by processing the training input x for the training example.

[0227] As another example, when the auxiliary machine learning model is a generative model and when each training example includes a target output, the system can determine, for each training example, an output likelihood,, of the auxiliary machine learning model generating the target output y for the training example by processing the training input x for the training example.

[0228] The system can determine the likelihood,(y|x), of the auxiliary machine learning model generating an output data item y by processing a training input x by any of a variety of methods. For example, in some implementations, the auxiliary machine learning model can be configured to generate a model output that specifies a distribution of output data items (e.g., by specifying a mean and covariance for a distribution of output data items, by specifying logits or probabilities for a set of output data items, etc.) and the system can determine the likelihood iteaux(y 1%), to be the likelihood of the output data item y according to a distribution of data items specified by the model output generated by the auxiliary machine learning model processing the training input x.

[0229] As another example, in some implementations, the auxiliary machine learning model can be configured to generate output data items for a training input x by sampling noise values z from a prior noise distribution, pz(z), (e.g., a multi -variate Gaussian noise distribution, a multi-variate uniform noise distribution, etc.) and by then processing the sampled noise values and the example prompt to generate an output data item following a mapping defined by the auxiliary machine learning model, / baux(z, x). When the target generative model is configured to generate output data items by transforming sampled noise, the system can determine the likelihood TTem(Y(y|x) to be a likelihood of sampling a noise value from the prior noise distribution that the auxiliary machine learning model maps to the data item y (e.g., a likelihood of sampling a noise value z such that y =As a particular example, the auxiliary machine learning model can define an invertible transformation from noise values to data items,and the system can determinefollowing:

[0230] inverse transformation defined by the auxiliary machine learning model and is a determinant of the Jacobian matrix of / eauX(ylx)

[0231] In some implementations, the auxiliary machine learning model can be configured to auto- regressively generate output data items. For example, the auxiliary machine learning model can be configured to generate a sequence of n data items, y1;n, over a sequence of auto-regressive iterations. At each auto-regressive iteration, the auxiliary machine learning model can process training input and some or all of the data items generated at the previous auto-regressive iterations to generate an output data item for the auto-regressive iteration. When the auxiliary machine learning model auto-regressively generates output data items, the system can determine the likelihood, n0aux(y1:n| x), of the auxiliary machine learning model generating sequence of data items y1:nby processing training input x following:

[0232] Where TT0aux(y1|x) is the likelihood of the auxiliary machine learning model processing the training input x to generate the first data item y13TTg^^ly^ x) is the likelihood of theauxiliary machine learning model processing the training input x and the first data item y±to generate the second data item y2, and so on.

[0233] In some implementations, the system can determine the likelihood,, by determining a log-likelihood,or any other appropriate function of the likelihood ^auxtyW-

[0234] The system can update the target generative machine learning function to optimize an objective function that depends on (i) the likelihoods of the target generative machine learning model generating the example data items by processing the example prompts for the training examples and (ii) differences between the quality scores for the training examples and the corresponding expected quality scores for the training examples (step 404). In particular, the system can update parameters of the target generative machine learning model to optimize the objective function.

[0235] The objective function can be any appropriate objective function for the target machine learning task. For example, when the target machine learning task is a classification task, the objective function can measure a cross-entropy loss between the model outputs generated by the auxiliary machine learning model and the corresponding target outputs for the training examples. As another example, when the target machine learning task is a regression task, the objective function can measure an error (e.g., a root mean-squared error, an LI loss, an L2 loss, etc.) between the model outputs generated by the auxiliary machine learning model and the corresponding target outputs for the training examples. As another example, when the auxiliary machine learning model is a generative model, the objective function can measure the likelihood of the auxiliary machine learning model generating the target outputs for the training examples. As another example, the objective function can measure a reward for the target machine learning task based on, e.g., quality scores for the target outputs.

[0236] In some implementations, when the auxiliary machine learning model is a generative model, the objective function can include a regularization term that measures, for each training example, a difference between a distribution of data items determined by processing the example prompt for the training example using the auxiliary machine learning model and a regularization distribution of data items for the training input for the training example. The regularization distribution of data items for the training example can be determined using an anchor set of machine learning model parameters for the auxiliary machine learning model. In particular, theregularization term can measure, for each training example, a difference between (i) the likelihood of the auxiliary machine learning model generating an output data item for the training example by processing the training input for the training example and (ii) a likelihood of the output data item for the training example as determined by the regularization distribution of data items for the training example.

[0237] For example, for each training example, the regularization term can measure a Kullback- Leibler divergence following:

[0238] Where 0anchor is the anchor set of machine learning model parameters for the target machine learning model and P is a regularization weight. In one example, = 0.1.

[0239] As a particular example, in some implementations for each training example, the objective function can measure the regularized reward:

[0240] Where x is the training input for the training example, y is an output data item for the training example, and r(x,y) is the reward for the training example. In some further examples, the objective function can include a length penalty. For example, when the target machine learning model is a language model, which given a text prompt that includes a sequence of tokens in a natural language auto-regressively generates the next tokens in the sequence, the length penalty can be -0.0005 x len(y), where y is the model output and len(y) is the number of tokens in the model output.

[0241] The system can use any of a variety of anchor sets of model parameters, for determining the regularization loss for the auxiliary machine learning model. For example, in some implementations, the anchor set of model parameters, 9anchor, can be pre-trained model parameters for a target machine learning model. As another example, in some implementations, the system can determine the anchor set of model parameters, 9anchor,as anexponential moving average of the model parameters of the auxiliary machine learning model, 9aux, over the sequence of training steps.

[0242] As part of updating the parameters of the auxiliary machine learning model, the system can determine gradients of the objective function (e.g., gradients of the objective function with respect to the parameters of the auxiliary machine learning model).

[0243] The system can update the parameters of the auxiliary machine learning model using the gradients of the objective function following any appropriate machine learning technique (e.g., following stochastic gradient descent, ADAM, etc.).

[0244] When the objective function measures a reward for each training example, the system can update the parameters of the auxiliary machine learning model using any appropriate machine learning model technique. For example, the system can update the parameters of the auxiliary machine learning model by determining a policy gradient of the objective function with respect to the parameters of the auxiliary machine learning model.

[0245] When the system determines the anchor set of model parameters, 0anchor,as anexponential moving average of the model parameters of the auxiliary machine learning model, the system can update the anchor set of model parameters at each training step (step 406). For example, at a t-th training step, the system can determine the updated anchor set of model parameters, danc / ior, following:

[0246] Where 9„nchoris the anchor set of model parameters for the t-th training step, QaUXis the set of model parameters for the auxiliary model as updated during the t-th training step, and p is a weight for the exponential moving average. In one example, p = 0.01. In general, the weight p can be selected to balance a trade-off between task performance and training speed.

[0247] The system can initialize the anchor set of model parameters for the first training step, anchorst0be the current model parameters of the auxiliary model.

[0248] The system can determine whether training is complete (step 408). The system can use any of a variety of criteria to determine whether the training is complete. For example, the system can determine that training is complete after a pre-determined number of training steps. As another example, the system can determine that training is complete when a value of the objective function falls below a pre-determined threshold. As another example, the system can determine that training is complete when a difference between values of the objective function for the current training step and a previous training step falls below a pre-determined threshold.

[0249] If the system determines that training is not complete, the system can continue to a next training step (e.g., return to step 402)

[0250] When the system determines that the training is complete, the system can return the trained auxiliary generative machine learning model (step 410). As described above with reference to FIG. 3, the system can use the trained machine learning model parameters of the trained auxiliary machine learning model as part of updating the model parameters of a target machine learning model.

[0251] FIG. 5 illustrates an example algorithm that a training system can use to train a target machine learning model. In particular, the algorithm depicted in FIG. 5 can train a target machine learning model (e.g., a generative model, such as an LLM, a VLM, etc.) with previously trained parameters QSft(e.g., as obtained by supervised fine tuning of the target machine learning model) to perform a target machine learning task using a training set of training examples, X, for the target machine learning task (e.g., a set of training examples that includes example prompts for the target machine learning task).

[0252] At each of / interpolation iterations, the algorithm can train M auxiliary machine learning models to perform the target machine learning task and can update the model parameters for the target machine learning model by interpolating the trained model parameters for the auxiliary machine learning models. The algorithm can train the auxiliary machine learning models to optimize a reward for each training example, r(y |x), determined (e.g., by a reward model for the target machine learning task) based on a model output y and an input prompt x for the training example.

[0253] The algorithm can use a set of interpolated weights for the target machine learning model ^sierp fromthe final interpolation iteration and the previously -trained parameters 0syt, to define a Pareto front of model parameters:

[0254] Where r| is a linear interpolation parameter that can control a trade-off between regularization strength and performance on the target machine learning task.

[0255] FIG. 6 illustrates experimental results demonstrating a performance of example target machine learning models that have been trained using the described methods. In particular, FIG. 6 illustrates a comparison of task rewards obtained for a target machine learning task with respect to regularization strength using the described methods 602 (e.g., as obtained using 1, 2, 3, and 4interpolation iterations) as compared with conventional methods 604 (e.g., as obtained by directly training the target machine learning models using the REINFORCE training algorithm with KL- regularization with either an initial set of parameters obtained by supervised fine-tuning or an exponential moving average of the trained model parameters as a regularization anchor set of model parameters). For the described methods, the results are plotted when sliding the weight r| from 0 to 1. Larger values of p provide greater task rewards but also greater KL divergence from the pre-trained model, while smaller values of p provide smaller task rewards with smaller KL divergence from the pre-trained model.

[0256] FIG. 6 shows results for fine-tuning a large language model to perform a conversational task that involves engaging in single-turn and multi-turn conversations with a user. In particular, the large language models are fine-tuned using feedback data characterizing human preferences for the conversation outputs of the large language model.

[0257] The experimental results in FIG. 6 are obtained by fine-tuning a language model that includes 7 billion parameters that has a network architecture as described by the Gemma Team in "Gemma: Open Models Based on Gemini Research and Technology", arXiv preprint arXiv:2403.08295 (2024) using a dataset containing conversation prompts.

[0258] As illustrated in FIG. 6, the outputs generated by the models trained using the methods described in this specification can obtain greater task rewards when trained using a given regularization strength as compared to conventional training methods. As compared to conventional training methods, the methods described in this specification can therefore be used to fine-tune target machine learning models to attain better performance on target machine learning tasks while still maintaining a given level of performance on more general pre-training tasks or to maintain a greater level of performance on the more general pre-training tasks while attaining a same level of performance on the target machine learning task.

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

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

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

[0262] 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 locatedat one site or distributed across multiple sites and interconnected by a data communication network.

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

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

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

[0266] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto-optical disks; and CD-ROM and DVD-ROM disks.

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

[0268] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and computeintensive parts of machine learning training or production, i.e., inference, workloads.

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

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

[0271] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmitsdata, 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.

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

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

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

[0275] What is claimed is:

Claims

CLAIMS1. A method performed by one or more computers and for training a target machine learning model to perform a target machine learning task, the method comprising: obtaining an initial set of machine learning model parameters for the target machine learning model; initializing, using the initial set of machine learning model parameters, a current set of machine learning model parameters for the target machine learning model; at each interpolation iteration of a sequence of interpolation iterations: for each of a plurality of auxiliary machine learning models for the interpolation iteration: initializing a set of machine learning parameters for the auxiliary machine learning model using the current set of machine learning parameters for the target machine learning model; and training the auxiliary machine learning model to perform the target machine learning task using training data for the target machine learning task to determine a trained set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration; and determining an interpolated set of machine learning model parameters for the interpolation iteration by interpolating the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration; and updating the current set of machine learning model parameters for the target machine learning model based on the interpolated set of machine learning model parameters for the interpolation iteration; and after the final interpolation iteration of the sequence of interpolation iterations, determining a trained set of machine learning model parameters for the target machine learning model based on the current set of machine learning model parameters for the target machine learning model.

2. The method of claim 1, wherein interpolating the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iterationcomprises performing a linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration.

3. The method of claim 1, wherein interpolating the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration comprises performing a spherical linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration.

4. The method of claim 3, wherein performing the spherical linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration comprises performing a layer-wise spherical linear interpolation of the trained sets of machine learning parameters for the plurality of auxiliary machine learning models for the interpolation iteration.

5. The method of any preceding claim, wherein updating the current set of machine learning model parameters for the target machine learning model based on the interpolated set of machine learning model parameters for the interpolation iteration comprises: updating, based on the interpolated set of machine learning model parameters for the interpolation iteration, the current set of machine learning model parameters for the target machine learning model as an exponential moving average of the interpolated set of machine learning model parameters for each of the interpolation iterations.

6. The method of any preceding claim, wherein training the auxiliary machine learning model to perform the target machine learning task to determine the trained set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration comprises: for each of a sequence of training steps for the auxiliary machine learning model for the interpolation iteration: updating the set of machine learning parameters for the auxiliary machine learning model for the interpolation using a machine learning technique to optimize an objective function for the target machine learning task.

7. The method of claim 6, wherein updating the set of machine learning parameters for the auxiliary machine learning model for the interpolation using the machine learning technique to optimize the objective function for the target machine learning task comprises: updating the set of machine learning parameters for the auxiliary machine learning model for the interpolation using a gradient of the objective function for the target machine learning task.

8. The method of claim 6 or claim 7, wherein training the auxiliary machine learning model to perform the target machine learning task to determine the trained set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration further comprises: initializing an anchor set of machine learning parameters using the current set of machine learning parameters for the target machine learning model.

9. The method of claim 8, wherein training the auxiliary machine learning model to perform the target machine learning task to determine the trained set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration further comprises, at each of the sequence of training steps for the auxiliary machine learning model for the interpolation iteration: after updating the set of machine learning parameters for the auxiliary machine learning model for the interpolation iteration: updating, based on the set of machine learning model parameters for the auxiliary machine learning model interpolation iteration, the anchor set of machine learning model parameters for the target machine learning model as an exponential moving average of the set of machine learning model parameters over each of the training steps for the auxiliary machine for the interpolation iteration.

10. The method of claim 8 or claim 9, wherein the objective function for the target machine learning task includes a term measuring a Kullback-Leibler divergence between (i) a distribution determined using the auxiliary model for the interpolation iteration in accordance with the set of machine learning model parameters for the auxiliary machine learning model and (ii) adistribution determined using the auxiliary model for the interpolation iteration in accordance with the anchor set of machine learning model parameters.

11. The method of any preceding claim, wherein determining the trained set of machine learning model parameters for the target machine learning model based on the current set of machine learning model parameters for the target machine learning model comprises: determining the trained set of machine learning model parameters for the target machine learning model as a weighted combination of the current set of machine learning model parameters for the target machine learning model and the initial set of machine learning model parameters for the target machine learning model.

12. The method of any preceding claim, wherein obtaining the initial set of machine learning model parameters for the target machine learning model comprises: obtaining the initial set of machine learning model parameters for the target machine learning model by pre-training the target machine learning model to perform one or more pretraining tasks.

13. The method of any preceding claim, wherein the target machine learning model comprises a generative model.

14. The method of any preceding claim, wherein the target machine learning model comprises a language model.

15. The method of any preceding claim, wherein the target machine learning model comprises an image generation neural network.

16. The method of any preceding claim, wherein the target machine learning model is configured to process input token sequences to generate corresponding output token sequences, wherein the input token sequence and the output token sequence comprise tokens from a vocabulary of tokens for the target machine learning model.

17. The method of claim 16, wherein the vocabulary of tokens for the target machine learning model comprises one or more tokens representing text data.

18. The method of claim 16 or claim 17, wherein the vocabulary of tokens for the target machine learning model comprises one or more tokens representing image data.

19. The method of any one of claims 16-18, wherein the vocabulary of tokens for the target machine learning model comprises one or more tokens representing video data.

20. The method of any one of claims 16-19, wherein the vocabulary of tokens for the target machine learning model comprises one or more tokens representing audio data.

21. The method of any one of claims 16-20, wherein the vocabulary of tokens for the target machine learning model comprises one or more tokens representing structured numerical data.

22. The method of any one of claims 16-21, wherein the vocabulary of tokens for the target machine learning model comprises one or more tokens representing an instruction for a system or agent to perform a task.

23. The method of any preceding claim, wherein: the target machine learning model is configured to interact with a user; and the training data for the target machine learning task comprises a plurality of training examples, wherein each training example comprises (i) an example of a query from an example user for the training example and (ii) an example response to the respective example of a query from the example user for the training example.

24. The method of any preceding claim, wherein: the target machine learning model is configured to select actions for an agent interacting with an environment to perform a task in the environment; and the training data for the target machine learning task comprises a plurality of training examples, wherein each training example comprises an example selected action for an example agent to perform the task in an example environment for the training example.

25. The method of claim 24, wherein: each training example comprises a respective example observation of the example environment for the training example.

26. A system comprising: one or more computers; and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform the operations of the respective method of any one of claims 1-25.

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

28. A method performed by one or more computers for generating a model output using a target machine learning model, the method comprising: obtaining a trained set of model parameters for the target machine learning model by, at each of a sequence of interpolation iterations: (i) training a plurality of auxiliary machine learning models to perform a target machine learning task and (ii) updating the set of model parameters for the target machine learning model based on an interpolation of the trained model parameters for the plurality of the auxiliary machine learning models; receiving a model input for the target machine learning model; and processing the model input for the target machine learning model using the trained model parameters for the target machine learning model to generate the model output.

Images