Next-Token Prediction in LLM Pretraining

A custom collate function that only groups input token IDs into batches gets the model partway to being trainable. The inputs tell the model what it is reading. Training also needs a second piece of information: what the model is supposed to produce at each position in that input. This article asks one specific question. Once you have a batch of input token IDs, how do you build the matching batch of target token IDs that the training loop will compare against?

The answer is a small, mechanical operation on each sequence. The best way to understand it is to watch that operation happen to concrete token IDs, so the sections below trace one sequence step by step before stating any general rule.

Take the input token ID sequence [7, 12, 3, 9]. Each position in this input has a token ID. The target sequence must have the same length, one entry per input position. Ask yourself what the model should produce at position 0. The useful observation is that the token that comes next in the input, at position 1, is the token the model should be producing at position 0.

Read the diagram from top to bottom. The first input ID, 7, is dropped and does not appear in the target at all. The remaining IDs 12, 3, and 9 each move one slot to the left, so target position 0 holds 12, target position 1 holds 3, and target position 2 holds 9. The final target slot, position 3, has no input ID left to move into it, so an end-of-text token fills it. The target is therefore [12, 3, 9, 50256], which has the same length as the input.

Targets matter because training works by comparison. The model produces a prediction for each position, and that prediction is scored against the target at the same position. The result of that comparison is the loss, and the loss drives the weight updates that make the model better. Without a target sequence there is nothing to compare against, so the batching process must deliver targets alongside the inputs.

The same alignment already appeared during pretraining. There, the model learns to predict the next token in a sequence, and the targets line up with the inputs but sit one position ahead. Instruction fine-tuning reuses that alignment: the target at each position is the token that should follow the input token at that position. Seeing the same logic in both settings is a good sign that the shift-and-append procedure is the standard way targets are formed here.

Batches often contain sequences of different lengths, so shorter sequences are padded. The source's second example shows what happens when the input already ends in end-of-text tokens. The input is [5, 6, 50256, 50256, 50256] and the target is [6, 50256, 50256, 50256, 50256]. Dropping the first ID and shifting left brings the 6 into the first slot. Every remaining slot is filled by an end-of-text token, because the padding entries of the input shift along with everything else, and the final slot receives a fresh end-of-text token.

The collate function built earlier produced one batch of input token IDs. The update does not change that job. It adds a second output, so the function hands back the input token IDs and the target token IDs together. Each sequence in the batch gets its target built by the shift-and-append procedure, and the two batches stay aligned position by position, so the training loop can compare each prediction with the right target.