Input Token IDs and Batch Construction

A custom collate function that only groups input token IDs into batches gets the model partway to being trainable. 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 answers one 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 applied to each sequence. The most useful way to understand it is to watch that operation happen to concrete token IDs, so the next section does exactly that. Keep the question in mind as you read.

Take a small input sequence of four token IDs: 15, 42, 7, and 90. The target sequence is built in three moves. First, drop the first input ID (15). Second, shift every remaining ID one slot toward the front, so 42 moves into the first slot, 7 into the second, and 90 into the third. Third, append an end-of-text token to fill the last slot, because nothing in the input is left to move there.

The diagram shows the pattern. Each target position holds the token that comes next in the input. Target position 0 holds the input ID at position 1, target position 1 holds the input ID at position 2, and so on. The window that reads from the input advances by one position, so the target lines up with the input one step ahead. Written with a symbol, the rule for every position i that has a successor is $y_{i} = x_{i+1}$, where $x$ is the input and $y$ is the target. The last target position has no successor, so it is the special case the end-of-text token handles.

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, which is why the batching process must deliver targets alongside inputs.

The same alignment appeared earlier in pretraining. There, the model learns to predict the next token in a sequence, so 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. The shift-and-append procedure is therefore not a new idea. It is the next-token relationship written down for a whole batch.

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.

Notice what the pipeline keeps constant. The batch still groups sequences the same way the earlier collate function did, and the target batch is produced from the same sequences, so row i of the target batch is always the shifted version of row i of the input batch. The training loop can then read both outputs in step.

Batches often contain sequences of different lengths, so shorter ones are padded. Consider an input that already ends in end-of-text tokens: [5, 6, 50256, 50256, 50256]. Dropping the first ID and shifting left brings the 6 into the first slot. The padding entries of the input shift along with everything else, so the remaining slots hold end-of-text tokens. The final slot then receives a fresh end-of-text token, giving the target [6, 50256, 50256, 50256, 50256].

The final target position always receives an end-of-text token, even when the input already contains padding. That is why the target keeps the same length as the input. The rule is applied the same way whether the sequence ends in real text or in padding, so the collate function does not need a special case for padded sequences.