Single-Head Attention Mechanism

This article follows a single practical task: setting up an attention module for the smallest GPT-2 model and then checking that the result matches what you expected. Attention modules are easy to build and hard to debug, because a wrong configuration rarely fails with a loud error. A module can be built from the wrong settings, or from an incomplete set of weights, and still run. The only way to catch this is to know which values were supplied, in which order, and what each stage should produce.

The lesson is built around a before-you-theorize approach. Read the initialization for the smallest GPT-2 model, set up its values, construct the module, and watch what the trace reports at each step. The goal is to understand what you are looking at, so that when a result does not match your expectation you know which line to question first.

The multi-head attention call supplies five values by position: the input dimension d_in, the output dimension d_out, the block size, a dropout value of 0.0, and the number of heads num_heads. The lesson does not give a full definition of each argument beyond these values. Treat the call as a contract: the input dimension, output dimension, block size, dropout rate, and head count must all be supplied together, and none of them can be left out or swapped silently.

Each slot should be checked one at a time. When you copy a configuration from one example to another, do not assume the positions line up. Read the call from left to right, name each value as you go, and confirm that it matches the parameter it is meant to fill. Two of these values, d_in and d_out, are the ones most often confused, because both are widths and both are plain integers.

A common assumption is that a module's output dimension always equals the d_out value you pass in. The lesson shows a different situation. It uses a wrapper that takes num_heads=2 and asks for an output dimension of 2, which matches the result that single-head attention produced earlier. To get that output width with two heads, the lesson changes the projection dimension d_out to 1. The value you pass for d_out is therefore not the same thing as the final width you see, and d_out must be reasoned about together with the head count.

The rule to keep in mind is that d_out is a total width that the heads share, not a per-head width you can set independently. If you change num_heads without revisiting d_out, the combined output width will change too. This is why the lesson treats the head count and d_out as a pair, and why a setup that passes the right d_out value for one head count can produce the wrong output width after the head count is changed.

The second setup task is copying learned weights from one attention module into another. The lesson states that the correct assignment transposes the weights of the query, key, and value matrices taken from another attention mechanism. Each target attribute is wrapped in torch.nn.Parameter, so that the copied tensor is registered as a parameter of the new module rather than as a plain tensor.

Read the three assignments as one repeated operation applied to three different matrices. The same pattern covers the query, the key, and the value weights. Because the lesson presents them as a set, a missing line or a line written differently produces a module built from an incomplete set of weights. Nothing in that situation is a syntax error, so the problem stays hidden until the outputs are compared.

On the transpose, the lesson states that the copy is transposed before it is wrapped and assigned. This passage does not spell out the reasoning behind the transpose, so treat it as a required step of the transfer. The transposed matrix is what gets placed into the target attribute, so skipping the transpose or applying it to only one of the three matrices changes what the new module holds.