Allocating computational resources between model size and training data while training a machine learning model

Abstract

A method, system, and apparatus, including a computer program encoded on a computer storage medium, for training a machine learning model to perform a machine learning task. In one aspect, a method implemented by one or more computers is described. The method includes obtaining data defining a computational budget characterizing an amount of computational resources allocated to train a machine learning model to perform a machine learning task, processing the data defining the computational budget with an allocation mapping according to a set of allocation mapping parameters to generate an allocation tuple defining (i) a target model size for the machine learning model and (ii) a target amount of training data for training the machine learning model, instantiating the machine learning model, the machine learning model having the target model size, and obtaining the target amount of training data for training the machine learning model.

Description

[Technical field]

[0001] This specification relates to processing data using machine learning models. [Background technology]

[0002] A machine learning model receives input and generates an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models and generate an output based on the received input and on the values ​​of the parameters of the model.

[0003] Some machine learning models are deep models that utilize multiple layers of the model 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 of the Invention [Problem to be solved by the invention]

[0004] This specification generally describes a training system implemented as a computer program on one or more computers at one or more locations that can train machine learning models to perform machine learning tasks. As described below, the training system can implement various methods and be realized in various systems to train machine learning models. [Means for solving the problem]

[0005] In one aspect, a method implemented by one or more computers for training a machine learning model is described. The method includes obtaining data defining a computational budget characterizing an amount of computational resources allocated to train a machine learning model to perform a machine learning task, and processing the data defining the computational budget with an allocation mapping according to a set of allocation mapping parameters to generate an allocation tuple defining (i) a target model size for the machine learning model and (ii) a target amount of training data for training the machine learning model. According to the method, selecting a model size of the machine learning model as the target model size and training the machine learning model with the target amount of training data is predicted to optimize the performance of the machine learning model in a machine learning task subject to a constraint that an amount of computational resources used to train the machine learning model satisfies a threshold defined by the computational budget. The method further includes instantiating a machine learning model, the machine learning model having a target model size, obtaining a target amount of training data for training the machine learning model, and training the machine learning model having the target model size with the target amount of training data.

[0006] In some implementations, values ​​for the set of allocation mapping parameters are determined by operations including identifying a plurality of trial allocation tuples, each trial allocation tuple defining (i) a trial model size for the machine learning model and (ii) a trial amount of training data for training the machine learning model; for each of the plurality of trial allocation tuples, determining a performance measure characterizing performance of the trial machine learning model at the machine learning task obtained by selecting the model size of the trial machine learning model as the trial model size and training the trial machine learning model with the trial amount of training data; and determining values ​​for the set of allocation mapping parameters based on the performance measures corresponding to the plurality of trial allocation tuples.

[0007] In some implementations, determining values ​​of the set of allocation mapping parameters based on the performance measures corresponding to the multiple trial allocation tuples includes, for each of a multiple computational budgets, determining an optimal model size and an optimal amount of training data corresponding to the computational budget based on the performance measures corresponding to the multiple trial allocation tuples, and determining values ​​of the set of allocation mapping parameters based on the optimal model size and the optimal amount of training data corresponding to each of the multiple computational budgets.

[0008] In some implementations, determining values ​​of the set of allocation mapping parameters based on an optimal model size and an optimal amount of training data corresponding to each of the multiple computational budgets includes adapting values ​​of the set of allocation mapping parameters based on an optimal model size and an optimal amount of training data corresponding to each of the multiple computational budgets.

[0009] In some implementations, determining, for each of a plurality of computational budgets, an optimal model size and an optimal amount of training data corresponding to the computational budget includes determining a respective performance curve for each of a plurality of trial model sizes based on performance measures corresponding to a plurality of trial allocation tuples, where the performance curves for the trial model sizes define a continuous mapping from possible computational budgets to predicted performance measures, where the predicted performance measures corresponding to the possible computational budgets define predicted performance of a trial machine learning model having the trial model size trained using an amount of computational resources that satisfies a threshold defined by the possible computational budget, and determining an optimal model size and an optimal amount of training data corresponding to each computational budget using the performance curves.

[0010] In some implementations, determining the performance curve for the trial model size includes determining the performance curve for the trial model size by interpolating performance measures of trial allocation tuples that correspond to the trial model size.

[0011] In some implementations, using the performance curves to determine an optimal model size and an optimal amount of training data corresponding to each computational budget includes, for each computational budget of the multiple computational budgets, determining an optimal performance curve from among the performance curves that achieves an optimal performance measure for the computational budget, determining an optimal model size as the trial model size corresponding to the optimal performance curve, and determining an optimal amount of training data based on the computational budget and the optimal model size.

[0012] In some implementations, determining, for each of a plurality of computational budgets, an optimal model size and an optimal amount of training data corresponding to the computational budget includes determining a respective performance curve for each of the plurality of computational budgets based on performance measures corresponding to a plurality of trial allocation tuples, where the performance curves for the computational budgets define a continuous mapping from possible model sizes to predicted performance measures, where the predicted performance measures corresponding to the possible model sizes define predicted performance of trial machine learning models having the possible model sizes trained using an amount of computational resources that satisfies a threshold defined by the computational budget, and using the performance curves to determine an optimal model size and an optimal amount of training data corresponding to each computational budget.

[0013] In some implementations, determining the performance curve for the computational budget includes determining the performance curve for the computational budget by interpolating performance measures of trial allocation tuples that correspond to the computational budget, where a trial allocation tuple corresponds to the computational budget if training a trial machine learning model having a trial model size defined by the trial allocation tuple with a trial amount of training data defined by the trial allocation tuple results in using an amount of computational resources that satisfies a threshold defined by the computational budget.

[0014] In some implementations, using the performance curves to determine an optimal model size and an optimal amount of training data corresponding to each computational budget includes, for each computational budget of the multiple computational budgets, determining an optimal model size as the model size that optimizes the performance curve corresponding to the computational budget, and determining an optimal amount of training data based on the computational budget and the optimal model size.

[0015] In some implementations, determining values ​​of the set of allocation mapping parameters based on performance measures corresponding to the plurality of trial allocation tuples includes determining a set of parameters of a performance estimation function configured to process data defining (i) an input model size and (ii) an input amount of training data to generate a predicted performance measure characterizing a predicted performance of a machine learning model having the input model size trained with the input amount of training data in a machine learning task, including adapting values ​​of the set of parameters of the performance estimation function based on the performance measures corresponding to the plurality of trial allocation tuples, and determining values ​​of the set of allocation mapping parameters using the performance estimation function.

[0016] In some implementations, determining values ​​of the set of allocation mapping parameters using the performance estimation function includes determining values ​​of the set of allocation mapping parameters to map each input computational budget to a target model size and target amount of training data that optimizes the performance estimation function subject to the constraint that training a machine learning model having a target model size with a target amount of training data uses an amount of computational resources given by the input computational budget.

[0017] In some implementations, adapting values ​​of a set of parameters of a performance estimation function based on performance measures corresponding to the plurality of trial allocation tuples includes adapting values ​​of the set of parameters of the performance estimation function to minimize, for each trial allocation tuple, a measure of error between (i) the performance measure corresponding to the trial allocation tuple and (ii) a predicted performance measure generated by processing a trial model size and a trial amount of training data defined by the trial allocation tuple with the performance estimation function.

[0018] In some implementations, the error measure includes Huber loss.

[0019] In some implementations, for each of the plurality of trial allocation tuples, determining a performance measure corresponding to the trial allocation tuple includes training a trial machine learning model having a trial model size with the trial volume of training data using a learning rate schedule selected based on the trial volume of training data.

[0020] In some implementations, the allocation mapping increases the target model size and the target amount of training data at substantially the same rate as the computational budget increases.

[0021] In some implementations, the machine learning task includes a language modeling task.

[0022] In some implementations, the machine learning model includes a neural network model.

[0023] In some implementations, the method further includes receiving a model input to the machine learning model; and processing the model input with the machine learning model according to trained values ​​of a set of model parameters of the machine learning model to generate a model output.

[0024] In some implementations, the machine learning model includes a multimodal model in which one or both of the model input and the model output include images or audio, and the multimodal model is configured to process model inputs including at least one of visual tokens representing pixels of still or moving images, data representing audio waveforms, and text tokens representing sequences of text, to generate model outputs including text tokens, images, or audio representing the model inputs.

[0025] In some implementations, the method is used to adapt a machine learning model to a particular computing hardware, the machine learning model includes a neural network model, and the particular computing hardware includes a plurality of neural network accelerators, in which case the method further includes determining an energy budget for training the machine learning model, the energy budget defining a total number of floating point operations for training the machine learning model, determining a computational budget from the energy budget, determining hardware specifications of the particular computing hardware on which the machine learning model is to be trained, the hardware specifications defining a number of neural network accelerators in the particular computing hardware, using any of the above-mentioned methods to determine a target model size for the machine learning model, the target model size defining a number of trainable parameters of the machine learning model, using any of the above-mentioned methods to determine a target amount of training data for training the machine learning model, the target amount of training data defining a number of training tokens to be used to train the model, and training the machine learning model having the defined number of trainable parameters on the particular computing hardware using the defined number of training tokens.

[0026] In a second aspect, a method implemented by one or more computers is described. The method includes receiving a model input to a machine learning model, and processing the model input with the machine learning model according to trained values ​​of a set of model parameters of the machine learning model to generate a model output. In this case, the machine learning model has been generated by operations including obtaining data defining a computational budget characterizing an amount of computational resources allocated to train the machine learning model to perform a machine learning task, and processing the data defining the computational budget with an allocation mapping according to a set of allocation mapping parameters to generate an allocation tuple defining (i) a target model size for the machine learning model and (ii) a target amount of training data for training the machine learning model. According to the method, selecting a model size of the machine learning model as the target model size and training the machine learning model with the target amount of training data is predicted to optimize the performance of the machine learning model in a machine learning task subject to a constraint that an amount of computational resources used to train the machine learning model satisfies a threshold defined by the computational budget. The method further includes instantiating a machine learning model, where the machine learning model has a target model size; obtaining a target amount of training data for training the machine learning model; and training the machine learning model having the target model size with the target amount of training data.

[0027] In a third aspect, a system is described that includes one or more non-transitory computer storage media that store instructions that, when executed by one or more computers, cause the one or more computers to perform operations of any of the methods described above.

[0028] In a fourth aspect, a system is described that includes one or more computers and one or more storage devices communicatively coupled to the one or more computers, the one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of any of the methods described above.

[0029] As described above, the training system can train the machine learning model to perform the machine learning task, for example, by implementing any of the methods described above. The machine learning model can be any suitable type of machine learning model. For example, the machine learning model can be a neural network model, a random forest model, a support vector machine model, or any combination thereof, or any other type of machine learning model.

[0030] The machine learning model may have any suitable machine learning model architecture. For example, if the machine learning model is a neural network model, the neural network model may have an attention-based neural network architecture (e.g., a transformer architecture), a convolutional architecture, a fully connected architecture, or any other suitable neural network architecture. In particular, the neural network model may include any suitable type of neural network layer (e.g., a convolutional layer, an attention layer, a fully connected layer, a regression layer, etc.) in any suitable number (e.g., 10 layers, 100 layers, or 1000 layers) and connected in any suitable configuration (e.g., as a linear sequence of layers or as a directed graph of layers).

[0031] Throughout this specification, the "model size" of a machine learning model may refer to the number of (trainable) parameters required to implement the machine learning model, such as weights, biases, matrix elements, etc.

[0032] Throughout this specification, a "computational budget" characterizes the amount of computational resources allocated to train a machine learning model to perform a machine learning task. A computational budget can be measured in floating-point operations (FLOPs), i.e., the total number of operations available to train a machine learning model. Note that FLOPs, denoted with a lowercase "s", should not be confused with floating-point operations per second (FLOPS), denoted with a capital "S", which corresponds to the rate of operations that can be performed by a particular computing hardware. In some cases, FLOPs can be determined by FLOPS multiplied by the total computation time, such as FLOPs=FLOPS×time.

[0033] In some implementations, the method is used to adapt a machine learning model to a particular computing hardware. In such cases, the machine learning model includes a neural network model, and the particular computing hardware includes multiple neural network accelerators. A neural network accelerator is a dedicated hardware used to accelerate neural network calculations, such as a GPU (graphics processing unit) or a TPU (tensor processing unit). In general, a neural network accelerator is configured to perform hardware matrix multiplication (e.g., may include a set of one or more multiply-accumulate units (MACs)), for example, using parallel computing.

[0034] The method may include determining an energy budget for training the machine learning model, the energy budget defining a total available number of floating-point operations for training the machine learning model. The energy budget may be determined according to a target carbon footprint for training. The computational budget may be determined from the energy budget. Depending on how the energy budget is expressed, e.g., the units in which the budget is expressed, both the energy budget and the computational budget may be defined as a total available number of floating-point operations (FLOPs). Alternatively, the computational budget may be determined from an energy budget expressed in terms of electrical energy based on known (e.g., average) energy usage of the computing hardware. In some implementations, the majority (e.g., nearly all) of the floating-point operations (FLOPs) performed during training are performed by the neural network accelerator, and thus the energy budget may be approximated under this assumption, e.g., using the energy consumption of the floating-point operations in one of the neural network accelerators to determine the computational budget.

[0035] The method can be used to determine the hardware specifications of a particular computing hardware on which a machine learning model is to be trained, i.e., the hardware specifications that define the number of neural network accelerators included in the particular computing hardware. The method is then used to determine a target model size for the machine learning model, the target model size defining the number of trainable parameters of the machine learning model. The method is also used to determine a target amount of training data for training the machine learning model, the target amount of training data defining the number of training data items, in particular training tokens, to be used to train the model. As described later, such training tokens represent training data items such as text tokens representing words or word pieces, or visual tokens representing intensity values ​​for pixels of still or moving images, e.g., for regions of an image. The method trains the machine learning model, i.e., a neural network with a defined number of trainable parameters, on the particular computing hardware using the defined number of training tokens. For a given energy and computation budget, many neural network models have been found to be too large and trained with too few tokens, indicating that some previously made assumptions about machine learning models are incorrect. That is, it has been found that the general trend of increasing model size effectively results in models with poor performance, and it has been surprisingly found that constraining the model to fit within certain hardware constraints and applying the techniques described herein can result in substantially better performance than previously expected. Some results in support of the techniques described are presented below.

[0036] The machine learning models used for the techniques described herein may be configured to perform any suitable machine learning task.

[0037] In particular, the machine learning models can be configured to process any suitable model inputs including, for example, one or more of images, audio waveforms, point clouds (e.g., generated by a lidar or radar sensor), protein representations, molecular representations, sequences of words (e.g., forming one or more sentences or paragraphs), video (e.g., represented as a sequence of video frames), or combinations thereof.

[0038] The machine learning model can be configured to generate any model output that characterizes the model input. For example, the model output can be a classification output, a regression output, a sequence output (i.e., including a sequence of output elements), a segmentation output, or a combination thereof.

[0039] A few examples of machine learning tasks that can be performed by the machine learning model are described in more detail below.

[0040] In some implementations, the machine learning model is configured to process model inputs representing pixels of an image to generate a classification output that includes a respective score for each object category (e.g., vehicle, pedestrian, bicyclist, etc.) in a set of possible object categories. The score for an object category may define the likelihood that the image shows an object that belongs to the object category.

[0041] In some implementations, the machine learning model is configured to process model inputs representing audio samples in an audio waveform to perform speech recognition, i.e., to generate output defining a sequence of phonemes, graphemes, characters, or words that correspond to the audio waveform.

[0042] In some implementations, the machine learning model is configured to process model inputs representing words in a sequence of words to perform natural language processing tasks, such as topic classification or summarization. To perform topic classification, the machine learning model generates an output including a respective score for each topic category in a set of possible categories (e.g., sports, business, science, etc.). The scores for the topic categories can define the likelihood that the sequence of words is related to the topic category. To perform summarization, the machine learning model generates an output including an output sequence of words that has a shorter length than the input sequence of words and captures important or relevant information from the input sequence of words.

[0043] In some implementations, the machine learning model performs a machine translation task by processing a model input representing a sequence of text, e.g., a sequence of words, phrases, characters, or word pieces in one language, to generate an output that may be a translation of the sequence of text into another language, i.e., a sequence of text in said other language that is a translation of the input sequence of text. As a specific example, the task may be a multilingual machine translation task, and the machine learning model is configured to translate between multiple different source and target language pairs. In this example, the source language text may be supplemented with an identifier indicating the target language into which the machine learning model should translate the source language text.

[0044] In some implementations, the machine learning model is configured to perform an audio processing task. For example, if the model input represents a speech utterance, the output generated by the machine learning model may be a score for each of a set of text fragments, each score representing an estimated likelihood that the text fragment is a correct transcription for the utterance. As another example, if the model input represents a speech utterance, the output generated by the machine learning model may indicate whether a particular word or phrase (a "hot word") was spoken. As another example, if the model input represents a speech utterance, the output generated by the machine learning model may identify the natural language of the speech uttered.

[0045] In some implementations, the machine learning model is configured to perform natural language processing or understanding tasks, such as entailment tasks, paraphrasing tasks, text similarity tasks, sentiment tasks, sentence completion tasks, grammaticality tasks, etc., operating on a set of model inputs representing text in some natural language.

[0046] In some implementations, the machine learning model is configured 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, waveform, or other data that defines the audio of the text spoken in the natural language.

[0047] In some implementations, the machine learning model is configured to perform a text generation task, where the model input represents a sequence of text and the 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 on a topic specified by the initial sequence of text. As another example, the model input can represent data other than text, e.g., images, and the output sequence can be text that describes the data represented by the model input.

[0048] In some implementations, the machine learning model is configured to perform an image generation task, where the model inputs represent conditioning inputs and the output is a sequence of intensity value inputs for pixels of the image.

[0049] In some implementations, the machine learning model is configured to perform an agent control task, where the model input represents a sequence of one or more observations or other data characterizing the state of the environment, and the output defines an action to be performed by the agent in response to the most recent data in the sequence. The agent may be a machine agent that performs an action in a real-world environment to perform a task, where the observations may include any type of observation, e.g., image observations, and the model output may include a control signal for controlling the agent to perform the task. Optionally, the model input may include other information, e.g., text tokens for text that defines the task to be performed. The agent may be a real-world or simulated robot, a control system for an industrial facility, or a control system that controls different types of agents.

[0050] In some implementations, the machine learning models are configured to perform genomics tasks, where the model inputs represent fragments of DNA sequences or other molecular sequences, and the outputs include, for example, promoter site predictions, methylation analysis, predictions about functional effects of non-coding variants, and the like.

[0051] In some implementations, the machine learning model is configured to perform, for example, a protein modeling task, where the model input represents a protein and the model output characterizes the protein For example, the model output can characterize the predicted stability of the protein or the predicted structure of the protein.

[0052] In some implementations, the machine learning model is configured to perform, for example, a point cloud processing task, where the model input represents a point cloud (e.g., generated by a lidar or radar sensor) and the model output, for example, characterizes a type of object represented by the point cloud.

[0053] In some implementations, the machine learning model is configured to perform a language modeling task, for example, by autoregressively generating an output sequence of text data. More specifically, the machine learning model can be configured to generate a sequence of output text tokens (text tokens may include, for example, characters, word pieces, words, n-grams, etc.). The machine learning model can generate an output sequence of text tokens over a sequence of time steps. At each time step, the machine learning model can generate an output token at a respective position in the sequence of output text tokens. The machine learning model can condition the generation of a text token at a position in the output sequence on the text tokens generated for each of one or more preceding positions in the output sequence. For example, to generate a text token at a position in the output sequence of text tokens, the machine learning model can process data including the text tokens generated for one or more preceding positions in the output sequence of text tokens to generate a score distribution over a set of possible text tokens. The machine learning model can then use the score distribution to select a token for the position, for example, by selecting a token having a highest score in the score distribution.

[0054] Optionally, a machine learning model configured to perform a language modeling task may be conditioned using one or more conditioning inputs. For example, a machine learning model may be conditioned on a question, and the machine learning model may autoregressively generate an output sequence of textual data that provides an answer to the question. As another example, a machine learning model may be conditioned on a task and a programming language, and the machine learning model may autoregressively generate an output sequence of textual data that defines instructions in the programming language to accomplish the task. As another example, a machine learning model may be conditioned on input instructions, e.g., a set of textual instructions, and the machine learning model may autoregressively generate an output sequence of textual data responsive to the set of input instructions.

[0055] In some cases, the machine learning model configured to perform a language modeling task may be implemented as a neural network model, which may include an attention neural network layer, such as a self-attention neural network layer, a mutual attention neural network layer, or both.

[0056] In some implementations, the machine learning model is configured to perform a combination of a plurality of individual machine learning tasks, e.g., two or more of the machine learning tasks mentioned above. For example, the machine learning model can be configured to perform a plurality of individual natural language understanding tasks, and a model input is processed by the machine learning model that includes an identifier for the individual natural language understanding task to be performed on the model input.

[0057] As a specific example, the machine learning model may include a multimodal model in which one or both of the model input and the model output include images or audio. For example, a multimodal machine learning model may be configured to process model inputs including visual tokens representing pixels of still or moving images (e.g., point cloud images) and / or data representing audio waveforms (e.g., values ​​or features of audio waveforms such as audio tokens and / or text tokens representing sequences of text) to generate model outputs (e.g., text tokens representing still or moving images or audio waveforms and / or sequences of intensity value inputs for pixels of the image or sequences of values ​​defining the audio waveform). A visual token may represent multiple pixels in a region of an image, for example, as a feature of the region. Such a multimodal model may perform any of the aforementioned tasks using multimodal inputs or by providing multimodal outputs or by converting between different input and output modes (e.g., text / image / audio). For example, the multimodal model may generate text that represents, describes (e.g., subtitling), or otherwise characterizes an image or audio input, e.g., by answering questions related to the image or audio input, such as a physical prediction of a state of an object represented by the image or audio. As another example, the multimodal model may generate images or audio that are represented, described, or otherwise characterized by the text input, or that are otherwise responsive to the text input, e.g., that represent an image or audio answer to a text question.

[0058] Throughout this specification, the term "optimize" may refer to predictive optimization or approximate optimization (i.e., not exact optimization).

[0059] Throughout this specification, a "performance measure" of a machine learning model in a machine learning task can refer to a measure of how effectively the machine learning model performs the machine learning task. For example, the system described herein can measure the performance of the machine learning model using, for example, a loss function or objective function that characterizes the predictive accuracy of the machine learning model. That is, the performance measure can be expressed as the value of the loss function or objective function used to train the machine learning model. The system can measure the performance of the machine learning model, for example, on a set of training data used to train the machine learning model, or on a set of validation data that is excluded from training the machine learning model, i.e., the machine learning model is not trained on the validation data set. Examples of loss functions / objective functions can include, for example, a cross-entropy objective function, a squared error objective function, etc.

[0060] For example, determining an optimal model size and an optimal amount of training data for a given computational budget may involve determining an optimization that optimizes a performance measure, e.g., optimizing the value of an objective function used to train the machine learning model, or minimizing the value of a loss function used to train the machine learning model, respectively.

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

[0062] The training system described herein trains a machine learning model to perform a machine learning task that is subject to a computational budget that allocates a limited amount of computational resources (e.g., FLOPs) for training the machine learning model. Training a machine learning model that is subject to a computational budget involves a tradeoff between (i) the model size of the machine learning model and (ii) the amount of training data used to train the machine learning model. For example, increasing the model size of the machine learning model may require a reduction in the amount of training data that can be used to train the machine learning model, i.e., to ensure that the computational resources consumed during training do not exceed the computational budget. To be precise, the amount of training data used to train the machine learning model generally refers to the amount of training data that the machine learning model sees during training, and not necessarily the total amount of training data at the disposal of the training system. For example, if there is a limited amount of training data available for a particular machine learning model and / or a particular machine learning task, the training system can sample from the available training data multiple times (if necessary). In this case, the amount of training data used to train the machine learning model may include multiple instances of the same data (e.g., tokens).

[0063] The training system uses an allocation mapping to determine a trade-off between a model size of a machine learning model and an amount of training data used to train the machine learning model. In particular, the allocation mapping processes data defining a computational budget to generate data defining a target model size and a target amount of training data that is predicted to optimize performance of the machine learning model subject to constraints defined by the computational budget. That is, the training system uses the allocation mapping to determine an allocation of computational resources between model size and training data to optimize performance (e.g., predictive accuracy) of the machine learning model.

[0064] The trial system can determine the mapping parameters of the allocation mapping by training a range of trial machine learning models to vary the model size, the amount of training data used for training, and the available computational resources specified by the computational budget, and determine empirical data for various trial machine learning models. The optimization system can then use the resulting data characterizing the performance of the machine learning models across the model size, the amount of training data, and the computational budget to adapt the parameters of the allocation mapping. For example, the optimization system can interpolate (and extrapolate) the data to different computational budgets, while simultaneously determining the optimal model size and the optimal amount of training data for the computational budget. After determining the parameters of the allocation mapping, the training system can then use the allocation mapping to determine the optimal trade-off between model size and the amount of training data each time the training system trains a machine learning model to perform a machine learning task.

[0065] By optimally allocating computational resources between model size and amount of training data, the training system can significantly increase the efficiency of computational resource usage during training. For example, the training system can reduce the likelihood of wasting computational resources by overtraining the machine learning model, i.e., training the machine learning model with a set of training data beyond the point where no further gains in the machine learning model's performance are achieved. As another example, the training system can select a model size for the machine learning model that is significantly smaller than would otherwise be selected (for the same set of training data), thereby reducing the use of computational resources during fine-tuning and downstream use of the trained machine learning model.

[0066] The details of one or more embodiments of the subject matter herein 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, drawings, and claims. [Brief description of the drawings]

[0067] [Figure 1] FIG. 1 is a block diagram of an example training system capable of training a machine learning model having a target model size with a target amount of training data for performing a machine learning task. [Diagram 2] 1 is a flow diagram of an example process for training a machine learning model having a target model size with a target amount of training data for performing a machine learning task. [Diagram 3] FIG. 1 is a block diagram of an example trial system that can determine values ​​for a set of allocation mapping parameters based on performance measures of a trial machine learning model. [Figure 4] 1 is a flow diagram of an example process for determining values ​​for a set of allocation mapping parameters based on performance measures of a trial machine learning model. [Diagram 5] 1 is a block diagram of two example optimization systems that can determine values ​​for a set of allocation mapping parameters based on a performance curve. [Figure 6] 1 is a flow diagram of an example process for determining values ​​for a set of allocation mapping parameters based on an optimal model size and an optimal amount of training data for a given computational budget. [Figure 7A] 1 is a flow diagram of an example process for determining an optimal model size and optimal amount of training data for a given computational budget based on a performance curve. [Figure 7B] FIG. 13 illustrates an example of generating a set of allocation mapping parameters using a performance curve that defines a continuous mapping from a possible computational budget to a predicted performance measure. [Figure 8A]4 is a flow diagram of another example process for determining an optimal model size and an optimal amount of training data for a given computational budget based on a performance curve. [Figure 8B] FIG. 13 illustrates an example of generating a set of allocation mapping parameters using respective performance curves for each of multiple possible computational budgets. [Figure 9A] 4 is a block diagram of another example optimization system that can use a performance estimation function to determine values ​​for a set of allocation mapping parameters. [Figure 9B] 4 is a block diagram of another example optimization system that can use a performance estimation function to determine values ​​for a set of allocation mapping parameters. [Figure 10] 4 is a flow diagram of an example process for determining values ​​for a set of allocation mapping parameters using a performance estimation function. [Figure 11A] FIG. 1 illustrates example experimental results comparing (i) the performance of an “optimal computational” machine learning model generated by the training system described herein and (ii) the performance of alternative machine learning models. [Figure 11B] FIG. 1 illustrates example experimental results comparing (i) the performance of an “optimal computational” machine learning model generated by the training system described herein and (ii) the performance of alternative machine learning models. DETAILED DESCRIPTION OF THE PREFERRED EMBODIMENTS

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

[0069] Large-scale machine learning models, such as large-scale language models (e.g., machine learning models including neural networks performing language modeling tasks as described above), deep learning models, generative models, discrimination and classification models, regression models, etc., have been implemented with a large number of parameters, e.g., more than 10 billion parameters, or more than 50 billion parameters, or more than 100 billion parameters, or more than 250 billion parameters, or more than 500 billion parameters. In particular, large-scale language models (LLMs) have demonstrated impressive performance in many machine learning tasks (e.g., language modeling tasks) using a variety of training and evaluation protocols, including zero-shot, few-shot, and fine-tuning.

[0070] However, the computational and energy costs of training a large-scale machine learning model (e.g., LLM) can be substantial and rise with increasing model size. Indeed, the allocated training computation (i.e., computational budget), e.g., how many accelerators (e.g., high-performance computing units) are available and for how long the accelerators are available, may be known in advance. In some situations, a machine learning model may be trained only once (or a few times), and therefore, accurately estimating the best model hyperparameters for a given computational budget may be of considerable benefit. For example, reducing the model size of a machine learning model can significantly reduce inference costs and facilitate downstream implementation in resource-constrained environments. The energy costs of a large-scale machine learning model are amortized by using the model for inference and fine-tuning. The benefits of a more optimally trained, smaller model therefore go beyond the immediate benefits of improved performance.

[0071] In this regard, the training system described herein can predict a target model size and a target amount of training data that are predicted to (approximately) optimize the performance of the machine learning model for a given computational budget, i.e., such that the training of the machine learning model is computationally optimal. In some cases, the training system can determine that training an optimal computational machine learning model for a given computational budget may require, for example, substantially increasing the amount of training data as opposed to increasing the model size. For example, for some optimal computational machine learning models (e.g., LLM), the training system can determine that the model size and the training data size scale at (approximately) equal rates with respect to the computational budget.

[0072] Moreover, as defined herein, a large-scale machine learning model may not need to be trained to the lowest possible loss to be computationally optimal. That is, the described technique describes how to optimize the loss for a given computational budget, taking into account that the machine learning model may not need to be trained to convergence. For example, as described later, some implementations of the system use a performance estimation function that takes this into account, for example, the function includes a term that represents the residual part of the loss due to the machine learning model not being trained to convergence.

[0073] For reference, some LLMs include Transformer Neural Networks, i.e., neural network models with a Transformer architecture. In general, a Transformer Neural Network may be a neural network model characterized by having a series of self-attention neural network layers. The self-attention neural network layers have an attention layer input for each element of the input and are configured to apply an attention mechanism to the attention layer input to generate an attention layer output for each element of the input. There are many different attention mechanisms that can be used. Some of these LLMs may use a Transformer Neural Network as an encoder, some may use a Transformer Neural Network as a decoder, while some may use one Transformer Neural Network as an encoder and another Transformer Neural Network as a decoder coupled to the encoder. Just by way of example, some LLMs are decoder-only models.

[0074] These and other features are described in more detail below.

[0075] 1 illustrates an example training system 100 that can train a machine learning model 102 having a target model size 132 with a target amount of training data 134 for performing a machine learning task 104. Training system 100 is an example of a system implemented as a computer program on one or more computers at one or more locations in which the systems, components, and techniques described below are implemented.

[0076] In general, for any particular machine learning model 102 configured to perform any particular machine learning task 104, the training system 100 determines a target model size N that is predicted to be computationally optimal. t 132 and the target amount of training data D t134. In other words, the target sizes 132 and 134 are predicted to optimize the (predicted) performance of the model 102 in the task 104, subject to the constraint that the amount of computational resources used for training (F) satisfies a threshold defined by the computational budget C 112, e.g., F=C or F≦C. The computational budget 112 defines the amount of computational resources allocated for training. For example, the allocated computational resources may be fixed by the available computational architecture (e.g., several accelerators, servers, GPU clusters, supercomputers, combinations thereof, etc.) and may not (or should not) be exceeded. Alternatively or additionally, the amount of allocated resources may be fixed to limit the energy consumption associated with training the machine learning model 102, such as, for example, so that multiple machine learning models may be in parallel training to reduce the environmental impact. In either case, the training system 100 can enable a reduction in the amount of both computational and energy resources expended on training the machine learning model 102 while enabling the machine learning model to achieve acceptable performance in the machine learning task 104.

[0077] For reference, the model size N may refer to, for example, the number of parameters that can be utilized by the machine learning model 102 that are required to implement the machine learning model 102. The amount of training data D, or training data size, may refer to a particular size of a particular training dataset 144 that can be used to train the machine learning model 102. For example, the training data size may refer to the number of tokens contained in the training dataset 144. More precisely, the amount of training data D used to train the machine learning model 102 refers to the amount of training data that the machine learning model 102 sees during training. Thus, the training dataset 144 may include multiple instances of the same token if the total training data available to the training system 100 is limited. As mentioned above, the computational budget 112 may refer to the amount of computational resources allocated to training the machine learning model 102 and may be measured in terms of the total number of floating-point operations (FLOPs). In some cases, the computational budget 112 may also be measured in terms of the total number of instructions, the total computation time, memory space, or a combination thereof (e.g., as a weighted sum). The amount of computational resources used during training, F (also called total computation), can be measured in the same units as the computational budget 112.

[0078] To determine the target sizes 132 and 134 for the machine learning model 102, the training system 100 first obtains (e.g., receives) data defining the computational budget 112. For example, the data may be provided to the training system 100 by a user or an automated process that seeks to implement an optimal computational training regime for the machine learning model 102 under the computational budget 112. For ease of explanation, the data defining the computational budget 102 may be described as being provided by a server 110, such as a cloud server, a local server, or a remote server.

[0079] The training system 100 is αβ120 to process the data defining the computational budget 112 to generate an allocation tuple [N t ,D t ] 130. The allocation tuple 130 is generated based on the target model size N t 132 and target data size D t 134. In general, the allocation mapping A αβ 120 is a function parameterized by a set of allocation mapping parameters 126 {α, β} that define how the allocation mapping 120 determines an optimal computational allocation of the computational budget C between possible model sizes N and possible data sizes D. As discussed above, the optimal computational allocation is determined by the target size N as the model size and data size, as follows: t and D. t Corresponds to the selection of [N t (C), D t (C)]=A αβ (C)

[0080] For clarity, let α = {α0,α1,...,α n} and β={β0,β1,...,β n} is the allocation mapping A αβ 120, how to calculate the budget C for each target model size N t and target data size D t The subsets α and β may share common parameters and do not necessarily have the same number of parameters. In general, the allocation mapping 120 can assume any functional form based on the particular set of mapping parameters 126. Some examples are described below.

[0081] In some implementations, the allocation mapping 120 may be expressed as a linear function with the mapping parameters 126 being the slope and intercept, for example, as follows: [N t (C), Dt (C)] = [α0,β0] + [α1,β1]C

[0082] In some implementations, the allocation mapping 120 may be expressed as a power law where the mapping parameters 126 are the coefficients and exponents, for example, as follows:

[0083]

number

[0084] In this case, when the machine learning system 102 is an LLM, the training system 100 may determine that in some scenarios, α≈β≈0.5 characterizes optimal computational scaling of model size and data size with the computational budget, i.e., in these cases, the target model size 132 and the target data size 134 should scale at a rate substantially equal to the computational budget 112.

[0085] In some implementations, the allocation mapping 120 may be expressed as a polynomial or Taylor series of a particular degree n, where the mapping parameters 126 are the coefficients of the polynomial, for example, as follows:

[0086]

number

[0087] More generally, in some implementations, the allocation mapping 120 may be expressed as a set of basis functions (eg, of order n) where the mapping parameters 126 are the coefficients of the basis functions, for example, as follows:

[0088]

number

[0089] Basis function f n,q(C) may be polynomial basis functions, Lagrangian basis functions, B-spline basis functions, Fourier basis functions, exponential basis functions, or any suitable set of basis functions of a desired order. In some cases, the basis functions themselves may depend on the allocation mapping parameters 126.

[0090] The value of the mapping parameter 126 determines the exact functional dependence of the allocation mapping 120 on the computational budget 112. In particular, the training system 100 determines the allocation t and D. t optimizes the performance L(N,D) of the machine learning model 102 on the machine learning task 104, subject to the constraint that the total computation F(N,D) is equal to the computational budget C. In other words,

[0091]

number

[0092] It is.

[0093] The above formula is the allocation tuple [N t ,D t ] 130, the machine learning model 102 associated with the training F(N t ,D t)=C, while optimizing its performance in the machine learning task 104 after training. For reference, the computation function F(N,D) represents the total computation used to train a machine learning model 102 having a particular model size N with a particular amount of training data D. The performance function L(N,D) represents a performance measure (e.g., pre-training loss) of the machine learning model 102 in the machine learning task 104 given the particular sizes N and D of the model 102. Note that the exact functional dependency between the computation function F(N,D) and the performance function L(N,D) is generally not known a priori, since it depends on the sizes N and D (e.g., "global" properties) of the particular machine learning model 102 that characterize its overall architecture. Thus, determining an appropriate allocation mapping 120 that satisfies the above constraints is a challenging problem. Various systems and methods for determining (e.g., empirically estimating) the allocation mapping 120 are described in detail with respect to FIGS. 3-10.

[0094] After generating the allocation tuples 130, the training system 100 instantiates 142 the machine learning model 102 with the target model size 132. The training system 100 then trains the machine learning model 102 with a training dataset 144 having the target amount of training data 134. For example, the training system 100 may obtain the training dataset 144 from the server 110 or other means. As described above, training begins with instantiating the allocation tuples [N t ,D t ] 130. In other words, training consumes the allocated computational resources defined by the computational budget 112, and the performance of the machine learning model 102 may be optimized for the machine learning task 104 given the computational budget 112.

[0095] After being trained, the machine learning model 102 can be deployed for use in performing a machine learning task 104. For example, the machine learning model 102 can be deployed in an environment that can enable a user to provide a request for the machine learning model 102 to process specified model inputs and generate corresponding model outputs. The user can provide the request, for example, by a user interface or through an application programming interface (API). The request can be transmitted from a user device (e.g., via a data communication network such as the Internet) to one or more computers implementing the machine learning model 102, for example, in a data center. The machine learning model 102 can process the model inputs specified by the user request to generate corresponding model outputs and then transmit the model outputs to the user device (e.g., via a data communication network).

[0096] 2 is a flow diagram of an exemplary process for training a machine learning model having a target model size with a target amount of training data for performing a machine learning task. For convenience, the process 200 is described as being performed by a system of one or more computers at one or more locations. For example, a training system suitably programmed in accordance with this specification, such as the training system 100 of FIG. 1, can perform the process 200.

[0097] The training system obtains data defining a computational budget characterizing an amount of computational resources allocated to train the machine learning model to perform the machine learning task (210). The training system can obtain the data defining the computational budget from a user, e.g., via a user interface or application programming interface (API), or from an external resource management system, e.g., that manages computational resources in one or more data centers.

[0098] The training system processes the data defining the computational budget with the allocation mapping according to the set of allocation mapping parameters to generate an allocation tuple defining (i) a target model size for the machine learning model and (ii) a target amount of training data for training the machine learning model (220). The training system generates the allocation tuple such that selecting the model size of the machine learning model as the target model size and training the machine learning model with the target amount of training data is predicted to optimize performance of the machine learning model at a machine learning task subject to the constraint that the amount of computational resources used to train the machine learning model satisfies a threshold defined by the computational budget.

[0099] The training system instantiates the machine learning model, which has a target model size (230). For example, the training system can generate an instance of the machine learning model, which includes determining an architecture of the machine learning model and initializing values ​​of a set of model parameters of the machine learning model. The training system can determine the architecture of the machine learning model, for example, by mapping a target model size of the machine learning model to a corresponding machine learning model architecture (e.g., according to a predefined architecture mapping). The architecture of the machine learning model can be defined, for example, by a set of architecture hyperparameters, and the system can generate values ​​for each architecture hyperparameter according to the target model size. For example, in an implementation in which the machine learning model is implemented as a neural network, the set of architecture hyperparameters can include hyperparameters that specify the number of layers in the neural network, the configuration of each layer in the neural network, and a directed graph that defines the connectivity between the layers of the neural network. The training system can initialize values ​​of the set of model parameters of the machine learning model using any suitable initialization technique, for example, random initialization or Glorot initialization.

[0100] The training system obtains 240 a target amount of training data for training the machine learning model. For example, to obtain the target amount of training data, the training system can access one or more data storage devices that store a corpus of training data. The system can identify a subset of the corpus of training data that includes the target amount of training data, for example, by randomly sampling training data from the corpus of training data, and then retrieve the selected training data for use in training the machine learning model.

[0101] Training data for training a machine learning model may be generated in any of a variety of possible ways. For example, the training data may include, for example, text sequences scraped (e.g., extracted using systematic and automated techniques) from one or more data sources, for example, one or more databases, or the Internet. A training system may use the text sequences for training a machine learning model to perform a language modeling task, as described in more detail below. As another example, the training data may include a set of training examples, where each training example includes (i) a model input (e.g., an image) to the machine learning model, and (ii) a target output (e.g., an image label), i.e., an output to be generated by the machine learning model by processing the model input. The target output may be generated, for example, through manual annotation, or in any other suitable manner.

[0102] The training system trains the machine learning model having the target model size with the target amount of training data (250). The training system may train the machine learning model with the training data using any suitable machine learning training technique. Some exemplary techniques for training the machine learning model with a set of training data are described below.

[0103] In some implementations, the machine learning model is a neural network model, the set of training data includes a set of text sequences, and the training system trains the neural network to perform a language modeling task. In these implementations, for each text sequence, the training system can use the neural network to process (at least a portion of) the text sequence to generate a score distribution over a set of possible tokens (e.g., text tokens including characters, word pieces, words, n-grams, etc.) for each of one or more positions in the text sequence. The neural network can be configured to generate a score distribution for a position in the text sequence by processing tokens from a preceding position in the text sequence, but not based on the tokens at that position or on tokens at a subsequent position in the text sequence. The training system can train the neural network for each of one or more positions in the text sequence based on an objective function that measures the error (e.g., cross-entropy error) between (i) the tokens at the position in the text sequence and (ii) the score distribution over the set of possible tokens generated by the neural network for that position. Training a neural network based on an objective function may include, for example, determining the gradient of the objective function with respect to parameters of the neural network (e.g., using backpropagation) and using the gradient to adjust the values ​​of the parameters of the neural network (e.g., using the update rule of a suitable gradient descent optimization technique such as RMSprop or Adam).

[0104] In some implementations, the training system trains the machine learning model to perform a supervised machine learning task. For example, the training system can train the machine learning model with a set of training examples, each of which includes (i) a model input and (ii) a target output. Training the machine learning model with the training examples can include training the machine learning model to process the model inputs of the training examples to generate predicted outputs that match the target outputs of the training examples. In particular, the training system can train the machine learning model to optimize, for each training example, an objective function that measures the error (e.g., cross-entropy error or squared error) between (i) the target output of the training example and (ii) the predicted output generated by the machine learning model for the training example.

[0105] 3 illustrates an example trial system 300 that can determine values ​​for a set of allocation mapping parameters 126 based on performance measures 350 of a trial machine learning model 302. Trial system 300 is an example of a system implemented as a computer program on one or more computers at one or more locations in which the systems, components, and techniques described below are implemented.

[0106] The trial system 300, in combination with the optimization system 500, can determine the allocation mapping 120 along with the value of its mapping parameter 126. That is, given a particular machine learning model 102 and a particular machine learning task 104, the trial system 300 can determine the corresponding allocation mapping 120 that provides optimal computational training of the model 102 for the task 104. The trial system 300 can accomplish this by empirically evaluating the performance of multiple trial machine learning models 302 with different trial model sizes 332 and trial data sizes 334. The optimization system 500 can then interpolate (and / or extrapolate) the performance of the trial sizes 332 / 334 to different possible sizes to determine the optimal size. From these results, the optimization system 500 can determine the value of the mapping parameter 126. Three variations of the optimization system 500 that use novel methods of specifying the value of the mapping parameter 126 are described with respect to FIGS. 5-10.

[0107] The trial system 300 can begin by identifying a number of trial allocation tuples 330. Each trial allocation tuple [N i ,D j ]330.ij is the trial model size N of machine learning model 102 i 332.i and training data D for training the machine learning model 102 j 334.j and the amount of trials in j. The trial allocation tuple 330 can be obtained by the trial system 300 in a variety of ways. For example, the trial system 300 may obtain a trial allocation tuple 330 by i and trial data size D j Let us consider the joint probability distribution [N i ,D j]~p(N,D) may be randomly sampled or sampled separately to generate trial allocation tuples 330.ij from various pairs of trial sizes 332.i / 334.j. In other cases, trial allocation tuples 330 may be specified by a user. Moreover, trial system 300 may choose a range of trial sizes and granularity based on a desired level of accuracy of the resulting mapping parameters 126. A larger range and finer granularity may increase accuracy. For example, trial system 300 may use more than 400 trial allocation tuples 330 with trial model sizes 332 ranging from 70 million to 16 billion parameters and trial data sizes 334 ranging from 5 billion to over 400 billion tokens. Note that a single trial model size 332.i may be associated with multiple different trial data sizes 334.j (and vice versa). This allows the trial system 300 to measure the performance of a trial machine learning model 302.ij with a particular trial model size 332.i on multiple different sized training sets 344.j. Along similar lines, a single trial model size 332.i is not necessarily associated with all trial data sizes 334.j (and vice versa). Thus, in some implementations, the trial system 300 may i and D j All combinations of may or may not be used.

[0108] For each trial allocation tuple 330.ij, the trial system 300 instantiates at 142 a trial machine learning model 302.ij with a respective trial model size 332.i. The trial system 300 then trains the trial machine learning model 302.ij with a training data set 344.j having a respective trial amount of training data 334.j. As mentioned above, the trial system 300 may obtain the training data 344.j from the server 110 or other means. The trial system 300 calculates a total computation F of each trial machine learning model 302.ij that characterizes the amount of computational resources used during training of the trial machine learning model 302.ij. ij=F(N i ,D j ) can also be determined. Therefore, for each trial allocation tuple [N i ,D j ] 330.ij provides data points for the computational function F(N,D). In some implementations, the trial system 300 trains the trial machine learning models 302.ij using learning rates that correspond to their trial data sizes 334.j. For example, the trial system 300 may decay (decrease) the learning rate for larger trial data sizes 334.j.

[0109] The trial system 300 evaluates each performance measure L ij =L(N i ,D j ) 350.ij, the performance of each trial machine learning model 302.ij in the machine learning task 104 is measured. Thus, each trial allocation tuple [N i ,D j ] 330.ij also provides data points for the performance function L(N,D). The trial system 300 then uses the optimization system 500 to optimize the performance measure L ij to determine the value of the allocation mapping parameter 126. As mentioned above, three variations of the optimization system 500 are described with respect to Figures 5-11, each capable of processing the performance measures 350 in a different manner to determine the value of the mapping parameter 126.

[0110] 4 is a flow diagram of an exemplary process 400 for determining values ​​for a set of allocation mapping parameters based on performance measures of a trial machine learning model. For convenience, process 400 is described as being performed by a system of one or more computers at one or more locations. For example, a trial system suitably programmed in accordance with this specification, such as trial system 300 of FIG. 3, may perform process 400.

[0111] The trial system identifies a number of trial allocation tuples, where each trial allocation tuple defines (i) a trial model size for a machine learning model and (ii) a trial amount of training data for training the machine learning model (410).

[0112] The trial system determines (420), for each of the plurality of trial allocation tuples, a performance measure that characterizes the performance of the trial machine learning model in the machine learning task obtained by selecting a model size of the trial machine learning model as the trial model size and training the trial machine learning model with the trial amount of training data.

[0113] The trial system determines (430) values ​​for a set of allocation mapping parameters based on performance measures corresponding to the multiple trial allocation tuples.

[0114] 5 illustrates two example optimization systems 500-1 / 500-2 that can determine values ​​for a set of allocation mapping parameters 126 based on a performance curve 520. Optimization systems 500-1 / 500-2 are examples of systems implemented as computer programs on one or more computers at one or more locations in which the systems, components, and techniques described below are implemented.

[0115] Both the first optimization system 500-1 and the second optimization system 500-2 determine the values ​​of the mapping parameters 126 by first determining the respective optimal model sizes 532 and optimal amounts of training data 534 for a given number of computational budgets 312. The optimal sizes 532 / 534 are the optimal computational type for the respective computational budgets 312. The optimization systems 500-1 / 500-2 then interpolate (and / or extrapolate) these data points to fit the mapping parameters 126 of the allocation mapping 120, which establishes a continuous mapping from the computational budgets 120 to the allocation tuples 130. However, the two optimization systems 500-1 / 500-2 may differ in how they determine the optimal sizes 532 / 534 themselves. The first optimization system 500-1 generates the curves by fixing the trial model size 332 and varying the trial data size 334. Conversely, the second optimization system 500-2 generates curves that vary the trial model size 332 while keeping the total computations fixed at the computation budget 312 (i.e., "iso-compute-budget" curves). The first optimization system 500-1 and the second optimization system 500-2 may work separately or jointly to determine the value of the mapping parameter 126. For example, the results of the two optimization systems 500-1 / 500-2 may be averaged and used for different types of machine learning models 102, for different ranges of trial sizes, etc. Details of the first optimization system 500-1 are outlined below, followed by the second optimization system 500-2.

[0116] First Optimization System (FOS) FOS 500-1 determines a respective performance curve 522.i for each trial model size 332.i. i Performance curve for L i (C) is the predicted performance measure L from a possible computational budget C. i In this case, the predictive performance measure is a function of the trial model size N i A trial machine learning model 302 with a possible computational budget F(Ni The total computation F(N i ,D). Similarly, the constraint F(N i ,D)=C is the trial model size N i Given ,we define the equation for the curve from a possible computational budget C to the amount of possible training data D (and vice versa).

[0117] FOS500-1 is a trial model size N i The performance measure L of the trial allocation tuple 330.ij corresponding to ij In other words, the FOS 500-1 can determine the performance curve 522.i for the trial model size 332.i by interpolating i The trial data size D associated with j For performance measure L ij FOS 500-1 can use a variety of different curve fitting techniques to interpolate performance measures 350, including power law fitting, linear regression, polynomial regression, polynomial interpolation, among others.

[0118] FOS500-1 then determines the optimal model size.

[0119]

number

[0120] 532.k and the optimal amount of training data

[0121]

number

[0122] 534.k for each given computation budget C k 312.k. To do so, FOS 500-1 determines the k 312. Optimal performance curve for k L k (Ck ) The optimal performance curve achieves the optimal performance measure for a given computational budget 312.k, i.e., C k When evaluated at , it achieves the minimum among all performance curves 522. L k (C k ) <L i≠k (C k )

[0123] FOS500-1 then calculates the associated trial model size

[0124]

number

[0125] , as the optimal model size 532.k for a given computational budget 312.k. The FOS 500-1 then determines from the optimal model size 532.k and the corresponding computational budget 312.k, e.g., the constraint

[0126]

number

[0127] can be used to determine the optimal data size 534.k. In general, F(N,D) may be any suitable function that characterizes the relationship between the model size N, the amount of training data D, and the computations required F to train a machine learning model 102 having that model size with that amount of training data. For example, in some implementations, the function is assumed or approximated as F(N,D)≈cND, where c is a constant, such as c=6. In other implementations, the trial system 300 and / or optimization system 500 may determine the total computations F expended during training of the trial machine learning model 302, e.g., using the interpolation and other data fitting techniques described herein. ij =F(N i ,D j) we can empirically determine F(N,D).

[0128] The FOS 500-1 then uses the optimal model size 532, the optimal data size 534, and the given computational budget 312 to determine the values ​​of the mapping parameters 126, e.g.,

[0129]

number

[0130] ,

[0131]

number

[0132] and C k For each associated triplet, A αβ (C k )=[N t (C k ),D t (C k )]and

[0133]

number

[0134] For example, the FOS 500-1 may use any of the curve fitting techniques described herein to fit the values ​​of the mapping parameters 126.

[0135] Second Optimization System (SOS) The SOS 500-2 determines a respective performance curve 524.k for each given computational budget 312.k. k Performance curve for L k (N) is the number of possible models of size N that predicts performance measure L kIn this case, the predicted performance measure is the degree to which a trial machine learning model 302 with possible model size N performs well for a given computational budget F(N,D)=C k It refers to its predictive performance when trained with a total computation F(N,D) equal to k is the calculation budget C k defines an equation for the curve from possible model sizes N to possible amount of training data D (and vice versa) given . Thus, performance curve 524 corresponds to the "iso-compute-budget" curve when the respective compute budget 312.k is fixed for each curve 524.k.

[0136] SOS500-2 is a computational budget C k The performance measure L of the trial allocation tuple 330.ij corresponding to ij In other words, the SOS 500-2 can determine the performance curve 524.k for a given computational budget 312.k by interpolating i For performance measure L ij The total calculation is calculated based on the calculation budget F ij =C k Trial data size D to be fixed j SOS 500-2 can use a variety of different curve fitting techniques to interpolate performance measures 350, such as power law fitting, linear regression, polynomial regression, polynomial interpolation, among others.

[0137] SOS500-2 then determines the optimal model size.

[0138]

number

[0139] 532.k and the optimal amount of training data

[0140]

number

[0141] 534.k, each calculation budget C k 312.k. To do so, SOS500-2:

[0142]

number

[0143] We select an optimal model size 532.k as the model size that optimizes the respective performance curve 524.k for a given computational budget 312.k such that .function..times ...

[0144] SOS500-2 then, for example,

[0145]

number

[0146] From the optimal model size 532.k and corresponding computational budget 312.k, an optimal data size 534.k can be determined using: As described above with respect to FOS 500-1, SOS 500-2 can assume the functional form F(N,D) or determine it empirically.

[0147] The SOS 500-2 then uses the optimal model size 532, the optimal data size 534, and the given computational budget 312 to determine the values ​​of the mapping parameters 126, e.g.,

[0148]

number

[0149] ,

[0150]

number

[0151] and C k For each associated triplet, A αβ (C k )=[N t (C k ),D t (C k )]and

[0152]

number

[0153] For example, the SOS 500-2 may use any of the curve fitting techniques described herein to fit the values ​​of the mapping parameters 126.

[0154] 6 is a flow diagram of an exemplary process 600 for determining values ​​for a set of allocation mapping parameters based on an optimal model size and an optimal amount of training data for a given computational budget. For convenience, process 600 is described as being performed by a system of one or more computers at one or more locations. For example, an optimization system suitably programmed in accordance with this specification, such as optimization systems 500-1 and 500-2 of FIG. 5, can perform process 600.

[0155] The optimization system determines (610) for each of a plurality of computational budgets an optimal model size and an optimal amount of training data corresponding to the computational budget based on the performance measures corresponding to the plurality of trial allocation tuples.

[0156] The optimization system determines (620) values ​​for a set of allocation mapping parameters based on an optimal model size and an optimal amount of training data corresponding to each of a plurality of computational budgets.

[0157] In some implementations, step 620 is accomplished by step 622, which proceeds as follows.

[0158] The optimization system adapts (622) values ​​of a set of allocation mapping parameters based on an optimal model size and an optimal amount of training data corresponding to each of a plurality of computational budgets.

[0159] 7A is a flow diagram of an exemplary process 700 for determining an optimal model size and an optimal amount of training data for a given computational budget based on a performance curve. For convenience, the process 700 is described as being performed by a system of one or more computers at one or more locations. For example, an optimization system suitably programmed in accordance with this specification, such as the first optimization system 500-1 of FIG. 5, can perform the process 700.

[0160] The optimization system determines 710 a respective performance curve for each of a plurality of trial model sizes based on the performance measures corresponding to the plurality of trial allocation tuples. The performance curves for the trial model sizes define a continuous mapping from possible computational budgets to predicted performance measures, where the predicted performance measures corresponding to the possible computational budgets define the predicted performance of a trial machine learning model having a trial model size trained using an amount of computational resources that satisfies a threshold defined by the possible computational budget.

[0161] In some implementations, step 710 is accomplished by step 712, which proceeds as follows.

[0162] The optimization system determines (712) a performance curve versus trial model size by interpolating performance measures of trial allocation tuples that correspond to the trial model sizes.

[0163] The optimization system uses the performance curves to determine 720 the optimal model size and optimal amount of training data that corresponds to each computational budget.

[0164] In some implementations, step 720 is accomplished by steps 722-726 that proceed as follows: For each computational budget of the plurality of computational budgets:

[0165] The optimization system determines 722 from among the performance curves an optimal performance curve that achieves the optimal performance measure for the computational budget.

[0166] The optimization system determines (724) the optimal model size as the trial model size that corresponds to the optimal performance curve.

[0167] The optimization system determines (726) an optimal amount of training data based on a computational budget and an optimal model size.

[0168] FIG. 7B illustrates an example of generating a set of allocation mapping parameters using performance curves that define a continuous mapping from possible computational budgets to predicted performance measures. In particular, graph 728 illustrates an example of a performance curve mapping possible computational budgets to predicted performance measures that the system would generate by training trial model sizes ranging from 75 million to 10 billion parameters. In graph 728, the horizontal axis represents possible computational budgets and the vertical axis represents predicted performance measures, which in this case are characterized, for example, as training loss, where lower training loss represents better performance. The system determines an optimal performance curve, for example, for each computational budget by determining a performance curve that represents the best performance measure for the computational budget (in this case, the lowest value for the computational budget). The system then uses the optimal performance curves to generate allocation mapping parameters that define a mapping from possible computational budgets to target model sizes (represented by lines in graph 730) and that define a mapping from possible computational budgets to target amounts of training data (represented by lines in graph 732). In particular, the data points in graph 730 are N t (C) is used to adapt

[0169]

number

[0170] and its counterpart C k Similarly, the data points in graph 732 correspond to pairs of D t (C) is used to adapt

[0171]

number

[0172] and its counterpart C kThis match then corresponds to the appropriate allocation mapping A αβ (C)=[N t (C), D t (C)].

[0173] 8A is a flow diagram of another exemplary process 800 for determining an optimal model size and an optimal amount of training data for a given computational budget based on a performance curve. For convenience, process 800 is described as being performed by a system of one or more computers at one or more locations. For example, an optimization system suitably programmed in accordance with this specification, such as the second optimization system 500-2 of FIG. 5, can perform process 800.

[0174] The optimization system determines a respective performance curve for each of a plurality of computational budgets based on the performance measures corresponding to the plurality of trial allocation tuples (810). The performance curves for the computational budgets define a continuous mapping from possible model sizes to predicted performance measures, where the predicted performance measures corresponding to the possible model sizes define the predicted performance of trial machine learning models having possible model sizes trained using an amount of computational resources that satisfies a threshold defined by the computational budget.

[0175] In some implementations, step 810 is accomplished by step 812, which proceeds as follows: The optimization system determines a performance curve for the computational budget by interpolating performance measures of a trial allocation tuple corresponding to the computational budget, and a trial allocation tuple corresponds to the computational budget if training a trial machine learning model having a trial model size defined by the trial allocation tuple with a trial amount of training data defined by the trial allocation tuple would use an amount of computational resources that satisfies a threshold defined by the computational budget.

[0176] The optimization system uses the performance curves to determine the optimal model size and optimal amount of training data that corresponds to each computational budget (820).

[0177] In some implementations, step 820 is accomplished by steps 822 and 824, which proceed as follows: for each computational budget of the plurality of computational budgets:

[0178] The optimization system determines 822 the optimal model size as the model size that optimizes the performance curve corresponding to the computational budget.

[0179] The optimization system determines 824 an optimal amount of training data based on a computational budget and an optimal model size.

[0180] FIG. 8B shows an example of generating a set of allocation mapping parameters with a respective performance curve for each of a number of possible computational budgets. The performance curves for the computational budgets define a continuous mapping from possible model sizes to predicted performance measures, where the amount of training data used during training is selected to match the total computation used during training to the computational budget. In this case, the computational budget is 6×10 18 ~3×10 21 FLOPs. In particular, in graph 826, the horizontal axis represents possible model sizes and the vertical axis represents a predicted performance measure, in this case characterized by training loss (e.g., where lower training loss represents better performance). The system then uses the performance curves to generate allocation mapping parameters that define a mapping from possible computational budgets to target model sizes (represented as a line in graph 828) and that define a mapping from possible computational budgets to target amounts of training data (represented as a line in graph 830). In particular, the data points in graph 828 are selected within N t (C) is used to adapt

[0181]

number

[0182] and its counterpart C k Similarly, the data points in graph 830 correspond to pairs of D t (C) is used to adapt

[0183]

number

[0184] and its counterpart C k This match then corresponds to the appropriate allocation mapping A αβ (C)=[N t (C), D t (C)].

[0185] 9A and 9B illustrate an example optimization system 500-3 that can use a performance estimation function 540 to determine values ​​for a set of allocation mapping parameters 126. The optimization system 500-3 is an example of a system implemented as a computer program on one or more computers at one or more locations in which the systems, components, and techniques described below are implemented.

[0186] Third Optimization System (TOS) The TOS500-3 uses a different approach compared to the FOS500-1 and SOS500-2. Instead of generating a performance curve, the TOS500-3 uses a performance estimation function

[0187]

number

[0188] to directly estimate the performance function L(N,D). The performance estimation function 540 is configured to process data defining the input model size N and the input amount of training data D to generate a predicted performance measure. The predicted performance measure characterizes the predicted performance of the machine learning model 102 on the machine learning task 104, given that the machine learning model 102 has an input model size N and is trained with the input amount of training data D. Similar to the mapping parameters 126 of the allocation mapping 120, the performance estimation function 540 is parameterized by a set of parameters 542 {γ} that define its functional form. In some implementations, for example when the machine learning model 102 is an LLM, the performance estimation function 540 is

[0189]

number

[0190] can be approximated as:

[0191] In this case, {γ}={E,A,B,α,β} is the number of entries to N and D.

[0192]

number

[0193] is a set of parameters 542 of the performance estimation function 540 that determines the functional dependence of . The first term in equation (★) captures the loss for an ideal generative process in the data distribution. The second term takes into account that the machine learning model with model size N falls short of the ideal generative process. The last term takes into account the machine learning model not being trained to convergence.

[0194] Referring to FIG. 9A, the TOS 500-3 first calculates the performance measure L of the trial allocation tuple 330. ij350 to the predicted performance measures generated by the performance estimation function 540. In particular, the TOS 500-3 uses the performance estimation function 540 to process the trial model size 332.i and trial data size 334.j of each trial allocation tuple 330.ij to generate the corresponding predicted performance measures

[0195]

number

[0196] TOS 500-3 then uses the error measure H 550 to compare the difference between the observed and predicted performance measures, as follows:

[0197]

number

[0198] In some implementations, the error measure 550 is:

[0199]

number

[0200] is the Huber loss corresponding to

[0201] Huber loss (δ=10 -3 ) is generally robust to outliers and therefore favorable for predictive performance.

[0202] The TOS 500-3 then optimizes 902 the error measure 550 with respect to the parameters γ 542 of the performance estimation function 540 to determine respective values ​​for the parameters.

[0203] Referring to FIG. 9B, the TOS 500-3 calculates the unknown performance function L(N,D) using a known performance estimation function

[0204]

number

[0205] TOS500-3 then uses the target model size N t 132 and target data size D t 134 is a performance estimation function for each input computation budget 512

[0206]

number

[0207] Allocate mapping [N t (C), D t (C)]=A αβ (C) Determine values ​​of the mapping parameters 126 to optimize 120. In other words, the target sizes 132 / 134 correspond to extreme values ​​of the performance estimation function 550 for each input computation budget 512, as follows:

[0208]

number

[0209] The above formula is the total calculation F(N t ,D t Note that this is subject to the constraint that F(N,D)=C is equal to the input computation budget 512. TOS 500-3 may implement a computation function of the form F(N,D)≈6ND, which allows TOS 500-3 to estimate values ​​for mapping parameters 126. However, as noted with respect to FOS 500-1 and SOS 500-2, ... ijUsing F(N,D)≈6ND, the TOS 500-3 may determine the value of the mapping parameter 126 as:

[0210]

number

[0211] In the above formula, N t (C) denotes the target model size given a computational budget C, and D t (C) denotes the target amount of training data given a computational budget C, where {α,β}={E,A,B,α,β} are the allocation mapping parameters 126 of the allocation mapping 120 described with reference to equation (★) and correspond to the same parameters 542 of the performance estimation function 540, but for N t and D. t Determine the functional dependence of .

[0212] 10 is a flow diagram of an exemplary process 1000 for determining values ​​for a set of allocation mapping parameters using a performance estimation function. For convenience, process 1000 is described as being performed by a system of one or more computers at one or more locations. For example, an optimization system suitably programmed in accordance with this specification, such as the third optimization system 500-3 of FIG. 9, may perform process 1000.

[0213] The optimization system processes data defining (i) an input model size and (ii) an input amount of training data to determine a set of parameters of a performance estimation function configured to generate a predicted performance measure characterizing a predicted performance of a machine learning model having the input model size trained with the input amount of training data in a machine learning task (1010). The optimization system adapts values ​​of the set of parameters of the performance estimation function based on the performance measures corresponding to the multiple trial allocation tuples.

[0214] In some implementations, step 1010 is accomplished by step 1012, which proceeds as follows.

[0215] The optimization system adapts (1012) values ​​of a set of parameters of a performance estimation function to minimize, for each trial allocation tuple, a measure of error between (i) a performance measure corresponding to the trial allocation tuple and (ii) a predicted performance measure generated by processing the trial model size and amount of trials of the training data defined by the trial allocation tuple with the performance estimation function.

[0216] The optimization system uses the performance estimation function to determine values ​​for a set of allocation mapping parameters (1020).

[0217] In some implementations, step 1020 is accomplished by step 1022, which proceeds as follows.

[0218] The optimization system determines (1022) values ​​for a set of allocation mapping parameters to map each input computational budget to a target model size and a target amount of training data that optimizes a performance estimation function subject to the constraint that training a machine learning model having the target model size with a target amount of training data uses an amount of computational resources given by the input computational budget.

[0219] 11A and 11B show examples of experimental results comparing (i) the performance of an "optimal computational" machine learning model produced by the training system 300 described herein with (ii) the performance of an alternative machine learning model ("Gopher"). The optimal computational machine learning model requires the same computational budget during training as the alternative machine learning model, but has 4 times fewer model parameters and is trained with 4 times more training data. FIG. 11A shows the improvement (measured in bits per byte) of the optimal computational machine learning model compared to the alternative machine learning model on a set of language modeling tasks. FIG. 11B shows the relative improvement (expressed in percentage) of the optimal computational machine learning model compared to the alternative machine learning model on a set of language understanding tasks. It can be appreciated that the optimal computational model produced by the training system 300 described herein significantly outperforms the alternative model.

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

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

[0222] The term "data processing apparatus" refers to data processing hardware and encompasses any kind of apparatus, device, and machine for processing data, including, by way of example, a programmable processor, a computer, or multiple processors or computers. An apparatus may be or further include special purpose logic circuitry, such as an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). An apparatus may optionally include, in addition to hardware, code that creates an execution environment for a computer program, such as code that constitutes a processor firmware, a protocol stack, a database management system, an operating system, or one or more combinations thereof.

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

[0224] As used herein, 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 and installed on one or more computers at one or more locations. In some cases, one or more computers will be dedicated to a particular engine, and in other cases, multiple engines can be installed and operated on the same computer or computers.

[0225] The processes and logic flows described herein may be implemented by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows may also be implemented by special purpose logic circuitry, e.g., FPGAs, or ASICs, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0226] A computer suitable for executing a computer program may be based on a general-purpose or special-purpose microprocessor or both, or any other type of central processing unit. In general, the 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 implementing or executing instructions, and one or more memory devices for storing instructions and data. The central processing unit and memory may be supplemented by or incorporated in special-purpose logic circuitry. In general, a computer also includes one or more mass storage devices, e.g., magnetic disks, magneto-optical disks, or optical disks, for storing data, or is operatively coupled to the mass storage devices to receive data from or transfer data to such mass storage devices, or both. However, a computer need not have such devices. Moreover, a computer may be incorporated into another device, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, such as a Universal Serial Bus (USB) flash drive, to name just a few.

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

[0228] To enable interaction with a user, embodiments of the subject matter described herein can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user, and a keyboard and pointing device, e.g., a mouse or trackball, for allowing the user to provide input to the computer. Other types of devices can be used to provide interaction with the user, e.g., 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. Additionally, a computer can interact with a user by sending and receiving documents to and from a device used by the user, e.g., by sending a web page to a web browser on the user's device in response to a request received from the web browser. A computer can also interact with a user by sending text messages or other forms of messages to a personal device, e.g., a smartphone running a messaging application, and receiving a response message in return from the user.

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

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

[0231] An embodiment of the subject matter described herein can be implemented in a computing system that includes back-end components, e.g., as a data server, or includes middleware components, e.g., an application server, or includes front-end components, e.g., a client computer having a graphical user interface or a web browser, or an app, by which a user may interact with an implementation of the subject matter described herein, or any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include local area networks (LANs) and wide area networks (WANs), e.g., the Internet.

[0232] A computing system may include clients and servers. Clients and servers are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship with each other. In some embodiments, a server transmits data, e.g., HTML pages, to a user device, e.g., for purposes of displaying data to a user interacting with the device acting as a client and receiving user input from a user interacting with the device acting as a client. Data generated at the user device, e.g., results of user interaction, may be received at the server from the device.

[0233] Although the present 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 certain embodiments of a particular invention. Also, some features described in the context of separate embodiments can be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Furthermore, although features may be described above as performing actions in some combinations and may even be initially claimed as such, one or more features from a claim may in some cases be deleted from the combination, and the claimed combination may be directed to a subcombination, or a variation of the subcombination.

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

[0235] Specific embodiments of the present subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. By way of example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequence, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous. [Explanation of symbols]

[0236] 100 Training System 102 Machine Learning Models, Models 110 Server 300 Trial system, training system 302 Trial Machine Learning Models 500 Optimization System 500-1 Optimization System, FOS 500-2 Optimization System, SOS 500-3 Optimization System, TOS

Claims

1. 1. A method implemented by one or more computers, comprising: obtaining data defining a computational budget characterizing an amount of computational resources allocated to training a machine learning model to perform a machine learning task; processing the data defining the computational budget using an allocation mapping according to a set of allocation mapping parameters to generate an allocation tuple defining (i) a target model size for the machine learning model, and (ii) a target amount of training data for training the machine learning model; selecting a model size of the machine learning model as the target model size, and training the machine learning model with the target amount of training data is predicted to optimize performance of the machine learning model at the machine learning task, subject to a constraint that an amount of computational resources used to train the machine learning model satisfies a threshold defined by the computational budget; instantiating the machine learning model, the machine learning model having the target model size; obtaining the target amount of training data for training the machine learning model; training the machine learning model having the target model size with the target amount of training data; A method comprising:

2. The values ​​of the allocation mapping parameter set are: identifying a plurality of trial allocation tuples, each trial allocation tuple defining (i) a trial model size for the machine learning model and (ii) a trial amount of training data for training the machine learning model; For each of the plurality of trial allocation tuples, determining a performance measure characterizing a performance of the trial machine learning model at the machine learning task obtained by selecting a model size of a trial machine learning model as the trial model size and training the trial machine learning model with the trial amount of training data; determining the values ​​of the allocation mapping parameter set based on the performance measures corresponding to the plurality of trial allocation tuples; The method of claim 1 , wherein the determination is made by operations including:

3. Determining the values ​​of the allocation mapping parameter set based on the performance measures corresponding to the plurality of trial allocation tuples includes: for each of a plurality of computational budgets, determining an optimal model size and an optimal amount of training data corresponding to the computational budget based on the performance measures corresponding to the plurality of trial allocation tuples; determining the values ​​of the allocation mapping parameter set based on the optimal model size and the optimal amount of training data corresponding to each of the plurality of computational budgets; 3. The method of claim 2, comprising:

4. Determining the values ​​of the allocation mapping parameter set based on the optimal model size and the optimal amount of training data corresponding to each of the plurality of computational budgets includes:

4. The method of claim 3, comprising adapting the values ​​of the allocation mapping parameter set based on the optimal model size and the optimal amount of training data corresponding to each of the plurality of computational budgets.

5. Determining, for each of the plurality of computational budgets, the optimal model size and the optimal amount of training data corresponding to the computational budget includes: determining a respective performance curve for each of a plurality of trial model sizes based on the performance measures corresponding to the plurality of trial allocation tuples; the performance curve versus trial model size defines a continuous mapping from possible computational budgets to predicted performance measures, the predicted performance measures corresponding to the possible computational budgets defining predicted performance of a trial machine learning model having the trial model size trained using an amount of computational resources that satisfies a threshold defined by the possible computational budget; To judge and using the performance curves to determine the optimal model size and the optimal amount of training data corresponding to each computational budget; 5. The method of claim 3 or 4, comprising:

6. Determining the performance curve versus trial model size includes: The method of claim 5 , comprising determining the performance curve for the trial model size by interpolating the performance measures of trial allocation tuples that correspond to the trial model sizes.

7. Using the performance curves to determine the optimal model size and the optimal amount of training data corresponding to each computational budget includes, for each computational budget of the plurality of computational budgets: determining an optimal performance curve from among the performance curves for the computational budget that achieves an optimal performance measure; determining the optimal model size as the trial model size that corresponds to the optimal performance curve; determining an optimal amount of training data based on the computational budget and the optimal model size; 7. The method of claim 5 or 6, comprising:

8. Determining, for each of the plurality of computational budgets, the optimal model size and the optimal amount of training data corresponding to the computational budget includes: determining a respective performance curve for each of the plurality of computational budgets based on the performance measures corresponding to the plurality of trial allocation tuples; the performance curve versus computational budget defines a continuous mapping from possible model sizes to a predicted performance measure, the predicted performance measure corresponding to a possible model size defining a predicted performance of a trial machine learning model having said possible model size trained using an amount of computational resources that satisfies a threshold defined by said computational budget; To judge and using the performance curves to determine the optimal model size and the optimal amount of training data corresponding to each computational budget; 4. The method of claim 3, comprising:

9. Determining the performance curve against the computational budget is determining the performance curve for the computational budget by interpolating performance measures of trial allocation tuples corresponding to the computational budget; 9. The method of claim 8, wherein a trial allocation tuple corresponds to a computational budget if training a trial machine learning model having the trial model size defined by the trial allocation tuple with the trial amount of training data defined by the trial allocation tuple would use an amount of computational resources that satisfies a threshold defined by the computational budget.

10. Using the performance curves to determine the optimal model size and the optimal amount of training data corresponding to each computational budget includes, for each computational budget of the plurality of computational budgets: determining the optimal model size as the model size that optimizes the performance curve corresponding to the computational budget; determining an optimal amount of training data based on the computational budget and the optimal model size; 10. The method of claim 8 or 9, comprising:

11. Determining the values ​​of the allocation mapping parameter set based on the performance measures corresponding to the plurality of trial allocation tuples includes: determining a set of parameters of a performance estimator function configured to process data defining (i) an input model size and (ii) an input amount of training data to generate a predicted performance measure characterizing a predicted performance of a machine learning model having said input model size, at said machine learning task, trained with said input amount of training data; adapting values ​​of the parameter set of the performance estimation function based on the performance measures corresponding to the plurality of trial allocation tuples. and using the performance estimation function to determine the values ​​for the allocation mapping parameter set; 3. The method of claim 2, comprising:

12. Determining the values ​​of the allocation mapping parameter set using the performance estimation function includes: determining the values ​​of the set of allocation mapping parameters to map each input computational budget to the target model size and target amount of training data that optimizes the performance estimation function subject to the constraint that training a machine learning model having a target model size with the target amount of training data uses an amount of computational resources given by the input computational budget.

12. The method of claim 11, comprising:

13. adapting the values ​​of the parameter set of the performance estimation function based on the performance measures corresponding to the plurality of trial allocation tuples, 13. The method of claim 11 or 12, comprising adapting, for each trial allocation tuple, the values ​​of the set of parameters of the performance estimation function to minimize a measure of error between (i) the performance measure corresponding to the trial allocation tuple and (ii) a predicted performance measure generated by processing the trial model size and the trial amount of training data defined by the trial allocation tuple with the performance estimation function.

14. The method of claim 13 , wherein the error measure comprises a Huber loss.

15. Determining, for each of the plurality of trial allocation tuples, the performance measure corresponding to the trial allocation tuple comprises:

15. The method of claim 2, comprising training a trial machine learning model having the trial model size on the trial volume of training data using a learning rate schedule selected based on the trial volume of training data.

16. 16. The method of claim 1, wherein the allocation mapping increases the target model size and the target training data amount in substantially the same proportion as the computational budget increases.

17. 17. The method of claim 1, wherein the machine learning task comprises a language modeling task.

18. 18. The method of claim 1, wherein the machine learning model comprises a neural network model.

19. receiving a model input for the machine learning model; processing the model inputs using the machine learning model according to trained values ​​of a set of model parameters of the machine learning model to generate model outputs; 19. The method of any one of claims 1 to 18, further comprising:

20. the machine learning model comprises a multi-modal model, in which one or both of the model input and the model output include images or audio; 20. The method of claim 19, wherein the multimodal model is configured to process the model inputs including at least one of visual tokens representing pixels of still or moving images, data representing audio waveforms, and text tokens representing sequences of text, to generate the model outputs including text tokens, images, or audio representing the model inputs.

21. is used to adapt the machine learning model to specific computing hardware, the machine learning model comprising a neural network model, the specific computing hardware comprising a plurality of neural network accelerators, the method comprising: determining an energy budget for training the machine learning model, the energy budget defining a total number of floating point operations for training the machine learning model; determining the computation budget from the energy budget; determining hardware specifications of the particular computing hardware on which the machine learning model is to be trained, the hardware specifications defining a number of the neural network accelerators in the particular computing hardware; using the method of any one of claims 1 to 20 to determine the target model size for the machine learning model, the target model size defining a number of trainable parameters of the machine learning model; Using the method of any one of claims 1 to 20 to determine the target amount of training data for training the machine learning model, the target amount of training data defining a number of training tokens to be used to train the model; training the machine learning model having the defined number of trainable parameters on the particular computing hardware using the defined number of training tokens; 21. The method of any one of claims 1 to 20, comprising:

22. 1. A method implemented by one or more computers, comprising: receiving a model input for a machine learning model; processing the model inputs using the machine learning model according to trained values ​​of a set of model parameters of the machine learning model to generate model outputs; Including, The machine learning model is obtaining data defining a computational budget characterizing an amount of computational resources allocated to training the machine learning model to perform a machine learning task; processing the data defining the computational budget using an allocation mapping according to a set of allocation mapping parameters to generate an allocation tuple defining (i) a target model size for the machine learning model and (ii) a target amount of training data for training the machine learning model; selecting a model size of the machine learning model as the target model size and training the machine learning model with the target amount of training data is predicted to optimize performance of the machine learning model at the machine learning task, subject to a constraint that an amount of computational resources used to train the machine learning model satisfies a threshold defined by the computational budget. Generating and instantiating the machine learning model, the machine learning model having the target model size; obtaining the target amount of training data for training the machine learning model; training the machine learning model having the target model size with the target amount of training data; The method is generated by operations including:

23. 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 recited in any one of claims 1 to 22.

24. One or more computers; one or more storage devices communicatively coupled to the one or more computers; A system comprising: 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 to 22. system.

Images