Batching Input Data for Language Models

Up to this point, attention has been applied to a single input text. A language model rarely works that way. Training pipelines group many texts together so the hardware can process them in one pass, and the attention code has to accept that grouped input without modification. The class in this article is built with that requirement in mind. It extends the self-attention class from the earlier section by adding two things: a causal mask that hides future tokens, and a dropout layer applied to the attention weights.

To test batch handling without a real data loader, the source duplicates one tokenized input text and stacks the two copies along a new first dimension. The result is a three-dimensional tensor: two inputs, six tokens each, and every token is a three-dimensional embedding vector. Those three numbers are the key to every shape question later in this article, so it helps to know what each one indexes.

Read the shape from left to right. The leading 2 says how many sequences travel together. The 6 says how many token positions each sequence has. The trailing 3 is the embedding size, meaning the number of values in each token's vector. The attention code must keep the first two dimensions meaningful throughout, because the scores are computed for every pair of tokens within each sequence.

Read the class before running anything. Most lines look familiar from the self-attention version: three linear projections produce queries, keys, and values, and the forward method computes scores, normalizes them, and mixes the values. The two additions are the dropout layer created in the constructor and the mask registered with register_buffer. Each addition appears in a different place in the forward method, so the code itself is the best map of where things happen.

The transpose deserves a closer look. Keys are transposed across dimensions 1 and 2 while the batch dimension stays in position 0. That single arrangement lets one matrix multiplication compute scores for every input in the batch at once, rather than looping over the inputs one at a time.

The shape prediction from the opening question works because every tensor can be followed through the forward method. Start with the input of shape 2 by 6 by 3. The three linear projections each map the embedding dimension from 3 to 2, so queries, keys, and values all have shape 2 by 6 by 2. Multiplying queries by the transposed keys produces a score for every pair of tokens within each sequence, giving a 2 by 6 by 6 matrix. The mask, the scaling, and the softmax all operate on that square matrix without changing its shape. Finally, multiplying the weights by the values produces context vectors of shape 2 by 6 by 2.

Tracing shapes is also the fastest way to locate a bug. If the printed shape differs from the one you traced, the error is in a step where the last dimension or the token dimension changed unexpectedly. A wrong output dimension usually points to the projection layers, while a wrong square score matrix points to the transpose or to a mismatch between the query and key shapes.

A causal mask enforces a simple rule: a token may attend to itself and to earlier tokens, but never to tokens that come later in the sequence. This restriction prevents data leakage in autoregressive models, which generate text one token at a time and must not peek at the words they have not yet produced.

The class builds the mask with torch.triu applied to a square matrix of ones. Passing diagonal=1 zeroes the main diagonal and everything below it, so the surviving ones sit strictly above the diagonal and mark exactly the future positions. Read the mask row for query token 2 above: keys 0, 1, and 2 are earlier or current tokens and remain visible, while keys 3, 4, and 5 lie in the future and are blocked.

Before softmax, blocked positions are filled with -inf. Because the exponential of negative infinity is zero, those positions receive zero weight after normalization, and the remaining weights in each row still sum to one. Masking happens on the scores, so the mask does not change the shape of the score matrix.

The mask is a fixed tensor. It holds no learnable weights, so it is not a model parameter. It still has to live on the same device as the model's parameters, because the masked fill operates on scores that are computed on that device. The register_buffer call exists for this case. A buffer is stored with the model but is not treated as a trainable parameter, and the source notes that buffers move automatically to the appropriate device, CPU or GPU, together with the model.

Dropout is a regularization technique that randomly sets a fraction of tensor elements to zero, which helps prevent overfitting. In this class it is created as nn.Dropout(dropout) in the constructor and applied to the attention weights after the softmax. The source applies dropout to the normalized weights rather than to the raw scores or the input embeddings, so each entry in the weight matrix can be zeroed independently during training.

The order of operations is the important part. Masking comes first, so blocked future positions are already -inf and get zero weight from the softmax. Softmax then normalizes each row. Dropout comes last and only perturbs the finished weights. The context vectors are produced from those dropped-out weights and the values, and the shape is still (2, 6, 2) at the end.