BATCH SELECTION FOR TRAINING MACHINE-LEARNED LARGE LANGUAGE MODELS
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.
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.