Causal Attention and Language Modeling

When you write a self-attention module from scratch, you often end up with more than one version of it. An early version might store its trainable projection weights in one layout, and a later, cleaner version might organize the same mathematical operation differently. Both versions can compute the same thing while looking quite different in code.

This raises a practical question. If you have a trained or carefully set-up instance of one version, how do you move its weights into an instance of the other version so that both produce the same outputs? The answer depends on matching the right weight matrices to each other and on handling one detail of how PyTorch stores linear layer weights. This article works through that transfer, then moves on to the two additions that turn basic self-attention into causal attention for language modeling: a causal mask and multiple heads.

The goal of the transfer is not to make two classes look alike. The goal is to make them behave alike. Two classes can have different internal structure, different attribute names, and different ways of storing their parameters, and still be equivalent if they compute the same function of their input.

Behavior is what gets tested. Once the weights are assigned correctly, you feed the same input into each instance and compare the outputs. If the outputs match, the transfer worked. That shared output is the check that confirms the two instances are equivalent for that input.

The transfer runs from an instance of the second version, called SelfAttention_v2, into an instance of the first version, called SelfAttention_v1. Each class holds trainable weight matrices for its query, key, and value projections. Copying weights therefore comes down to two steps: identify which matrix in one class corresponds to which matrix in the other, and then assign values pairwise, one matrix at a time.

Notice that the assignment is pairwise. The query matrix of one class is matched to the query matrix of the other, and the same holds for keys and values. Mixing them up would still produce tensors of compatible shape in some cases, but the instances would then compute different attention scores and disagree on their outputs.

After the three assignments, run the same input through both instances and compare the outputs. Identical results mean the parameters now define the same function, which is what makes the two implementations equivalent.

The step that most often causes trouble is the storage layout. nn.Linear stores its weight matrix in a transposed form. If you write a weight matrix by hand, you would normally give it rows for input features and columns for output features. Inside nn.Linear the orientation is the other way around, with rows for output features and columns for input features.

This is why the transfer must account for the transpose. The weight of the second instance has to be transposed before it is assigned into the first instance's matrix, so that each value lands in the orientation the first class expects. The code in the previous section applies this with .T.

Most failed transfers come from a small set of errors. Each one is easy to make and easy to miss until the outputs are compared.

The next enhancement adds a causal aspect to self-attention. Causal attention modifies the mechanism so that the model cannot access future information in the sequence. This matters for language modeling, where each prediction should depend only on the words that came before it.

Consider why this restriction is needed. If a model could peek at upcoming words while it predicts the current one, it would learn a shortcut that does not exist at generation time. During training it would appear to predict well because it could see the answer, and then fail once it had to produce text one token at a time.

In GPT-style models each token is processed in turn. For the current token, the tokens that come after it in the input are masked out. The row for a given token therefore only distributes its attention over itself and the tokens before it.

Concretely, the attention weights above the diagonal of the attention matrix are masked out, and the remaining weights in each row are normalized so that they sum to 1. Each row is a valid distribution over the tokens that row is allowed to see. The diagram below shows which entries survive for a three-token sequence.

The masked attention weight matrix can be built in three stages, in this order. First, apply the softmax function to the attention scores. Second, zero out the elements above the diagonal so that future positions contribute nothing. Third, normalize the resulting matrix row by row so that each row again sums to 1.

The ordering of the three stages matters. Zeroing the upper triangle after the softmax leaves rows that no longer sum to 1. That is why the normalization stage follows the masking stage. Right after the zeroing, the masked matrix is unnormalized, and it becomes a valid set of distributions only at the end.

The third enhancement is the multi-head component. Here the attention mechanism is split into several heads. Each head learns different aspects of the data, and the heads together let the model attend to information from different representation subspaces at different positions at the same time. The source presents this as a way to improve performance on complex tasks.

A practical GPT-like model uses both enhancements together. The mask controls which positions a token may see, while the heads let it look at those positions in several different ways. Causal masking and multi-head splitting are therefore presented as complementary parts of the same model rather than as alternatives.