Batch selection for training machine-learned large language models

Abstract

A system performs a batch selection process on a small model and uses the results of batch selection to train a large language model (LLM). The system receives training examples and splits the training examples into a holdout set and an evaluation set. Each training example corresponds to a label. The system uses trains a small model using the training examples of the holdout set. The system evaluates the small model on the training examples of the evaluation set, generating a prediction for each training example and computing a loss between the prediction and the training example's label. The system generates an LLM training set by selecting a set of training examples from the evaluation set with the highest loss. The system trains the LLM using the LLM training set.

What is claimed is: 
     
         1 . A method for training a first machine-learned language model, comprising:
 obtaining a holdout set and an evaluation set, the holdout set including a first set of training examples and the evaluation set including a second set of training examples;   dividing the holdout set into a plurality of batches, each batch including a respective subset of training examples;   training a surrogate model using the holdout set by iterating through the plurality of batches, wherein the surrogate model is a second machine-learned language model;   evaluating the trained surrogate model on the evaluation set by:
 generating a set of predictions by applying the trained surrogate model to the evaluation set, the set of predictions including a prediction for a training example in the evaluation set, and 
 determining, for the training example in the evaluation set, a loss based on the set of predictions for the training example; 
   generating a training set for the first machine-learned language model by selecting, from the evaluation set, a set of training examples with the loss within a predetermined range; and   training parameters of the first machine-learned language model for a first epoch using the selected training set.   
     
     
         2 . The method of  claim 1 , wherein the second machine-learned language model has a smaller number of parameters than a number of parameters of the first machine-learned language model. 
     
     
         3 . The method of  claim 1 , further comprising for at least a second epoch, training the parameters of the first machine-learned language model using the selected training set. 
     
     
         4 . The method of  claim 1 , wherein training the surrogate model comprises repeatedly iterating between:
 performing a forward pass step to generate a set of estimated predictions for a batch,   computing a loss function for the batch based on the set of estimated predictions, and   performing a backpropagation step to update parameters of the surrogate model to reduce the loss function.   
     
     
         5 . The method of  claim 1 , wherein a training example is a sequence of tokens. 
     
     
         6 . The method of  claim 1 , further comprising:
 dividing the selected training set into another plurality of batches,   wherein training the first machine-learned language model comprises repeatedly iterating between:
 performing a forward pass step to generate another set of estimated predictions for a batch in the another plurality of batches, 
 computing a loss function for the batch based on the another set of estimated predictions, and 
 performing a backpropagation step to update the parameters of the first machine-learned language model to reduce the loss function. 
   
     
     
         7 . The method of  claim 1 , further comprising obtaining a total budget indicating a total number of training examples for training the parameters of the first machine-learned language model, and wherein a number of training examples in the selected training set is based on the total budget. 
     
     
         8 . A non-transitory computer readable storage medium comprising stored program code, the program code comprising instructions, the instructions when executed causes a processor system to:
 obtain a holdout set and an evaluation set, the holdout set including a first set of training examples and the evaluation set including a second set of training examples;   divide the holdout set into a plurality of batches, each batch including a respective subset of training examples;   train a surrogate model using the holdout set by iterating through the plurality of batches, wherein the surrogate model is a second machine-learned language model;   evaluate the trained surrogate model on the evaluation set by:   generate a set of predictions by applying the trained surrogate model to the evaluation set, the set of predictions including a prediction for a training example in the evaluation set, and   determine, for the training example in the evaluation set, a loss based on the set of predictions for the training example;   generate a training set for the first machine-learned language model by selecting, from the evaluation set, a set of training examples with the loss within a predetermined range; and   train parameters of the first machine-learned language model for a first epoch using the selected training set.   
     
     
         9 . The non-transitory computer readable storage medium of  claim 8 , wherein the second machine-learned language model has a smaller number of parameters than a number of parameters of the first machine-learned language model. 
     
     
         10 . The non-transitory computer readable storage medium of  claim 8 , wherein the instructions further comprise instructions to, for at least a second epoch, train the parameters of the first machine-learned language model using the selected training set. 
     
     
         11 . The non-transitory computer readable storage medium of  claim 8 , wherein the instructions to train the surrogate model comprise instructions to repeatedly iterate between:
 performing a forward pass step to generate a set of estimated predictions for a batch,   computing a loss function for the batch based on the set of estimated predictions, and   performing a backpropagation step to update parameters of the surrogate model to reduce the loss function.   
     
     
         12 . The non-transitory computer readable storage medium of  claim 8 , wherein a training example is a sequence of tokens. 
     
     
         13 . The non-transitory computer readable storage medium of  claim 8 , wherein the instructions further comprise instructions to:
 divide the selected training set into another plurality of batches,   wherein the instructions to train the first machine-learned language model comprise instructions to repeatedly iterate between:
 performing a forward pass step to generate another set of estimated predictions for a batch in the another plurality of batches, 
 computing a loss function for the batch based on the another set of estimated predictions, and 
 performing a backpropagation step to update the parameters of the first machine-learned language model to reduce the loss function. 
   
     
     
         14 . The non-transitory computer readable storage medium of  claim 8 , wherein the instructions further comprise instructions to obtain a total budget indicating a total number of training examples for training the parameters of the first machine-learned language model, and wherein a number of training examples in the selected training set is based on the total budget. 
     
     
         15 . A computer system, comprising:
 a computer processor; and   a non-transitory computer readable storage medium comprising stored instructions that when executed by the computer processor, cause the computer system to:
 obtain a holdout set and an evaluation set, the holdout set including a first set of training examples and the evaluation set including a second set of training examples; 
 divide the holdout set into a plurality of batches, each batch including a respective subset of training examples; 
 train a surrogate model using the holdout set by iterating through the plurality of batches, wherein the surrogate model is a second machine-learned language model; 
 evaluate the trained surrogate model on the evaluation set by: 
 generate a set of predictions by applying the trained surrogate model to the evaluation set, the set of predictions including a prediction for a training example in the evaluation set, and 
 determine, for the training example in the evaluation set, a loss based on the set of predictions for the training example; 
 generate a training set for the first machine-learned language model by selecting, from the evaluation set, a set of training examples with the loss within a predetermined range; and 
 train parameters of the first machine-learned language model for a first epoch using the selected training set. 
   
     
     
         16 . The computer system of  claim 15 , wherein the second machine-learned language model has a smaller number of parameters than a number of parameters of the first machine-learned language model. 
     
     
         17 . The computer system of  claim 15 , wherein the instructions further comprise instructions to, for at least a second epoch, train the parameters of the first machine-learned language model using the selected training set. 
     
     
         18 . The computer system of  claim 15 , wherein the instructions to train the surrogate model comprise instructions to repeatedly iterate between:
 performing a forward pass step to generate a set of estimated predictions for a batch,   computing a loss function for the batch based on the set of estimated predictions, and   performing a backpropagation step to update parameters of the surrogate model to reduce the loss function.   
     
     
         19 . The computer system of  claim 15 , wherein a training example is a sequence of tokens. 
     
     
         20 . The computer system of  claim 15 , wherein the instructions further comprise instructions to:
 divide the selected training set into another plurality of batches,   wherein the instructions to train the first machine-learned language model comprise instructions to repeatedly iterate between:
 performing a forward pass step to generate another set of estimated predictions for a batch in the another plurality of batches, 
 computing a loss function for the batch based on the another set of estimated predictions, and 
 performing a backpropagation step to update the parameters of the first machine-learned language model to reduce the loss function.