Ground-Up VLA: Multi-Head Attention

From GPU basics to robot foundation models

We’ve spent a couple of posts looking at practical implementations of deep learning on the WikiText-2 dataset. This post will bring us back to a key bit of theory that we said we’d revisit when we first covered the concept of attention: multi-head attention.

Intuitively, we can think of the different heads of attention for a deep learning pipeline like the members of a politician’s staff. Each member’s role is to analyze the possible actions that a politician can take from a different angle. The politician makes informed decisions based on the combination of assessments done by each staff member.

For instance, if a politician is trying to take a national leadership role in a democratic system, the staff might have an economics expert and an elections expert. Each expert is given the same query: “Should we raise taxes?” The economics expert might say this would be great for balancing the country’s budget, while the elections expert might say that we are so close to elections that raising taxes would hurt re-election chances.

Much like each expert comes pre-biased with life experiences that inform how that person will see a query, each head of attention looks at a linear projection of the query unique to itself. In other words, the raw input is the same, but the way the question is perceived might be different. Additionally, each head also has its own weights for projecting keys and values.

In a deep learning pipeline, the various heads tend to take on somewhat more abstract concepts than domestic economics and election impact. For instance, one head might examine tone of voice, while another looks at the frequency of complex jargon. In both levels of abstraction, however, the math works the same.

Single-Head Self-Attention

Let’s recap the concept of single-headed self-attention, and then expand the math to multi-head attention. For a more detailed explanation, see the prior self-attention post. Given an input sequence X, the goal of self-attention is to determine how the tokens of X relate to one another. There are three linear projections that are produced: the query, key, and value. Each projection is produced by multiplying the input X by a weight matrix associated with the projection type.

• Query Q = XW_Q: what information is this word looking for from other words?

• Key K = XW_K: what kind of information does each word offer to others?

• Value V = XW_V: what information does a word provide if another word decides to pay particular attention to it?

Once the three linear projections are produced, we can compute the attention results.

• Compute compatibility scores S = QK^T / sqrt(d_model).

• Normalize S via the softmax function: A = Softmax(S). A is the attention matrix.

• Produce the output Y = AV.

A key part to remember before delving into multi-head attention is the shape of these matrices. The input sequence X, consisting of n tokens, has dimensions n x d_model, and each weight matrix (W_Q, W_K, W_V) has dimensions d_model x d_model, which together mean that the projections Q, K, and V each have dimensions n x d_model. This also means that S and A both have dimensions n x n, and that Y has dimensions n x d_model (the same as X).

Multi-Head Self-Attention

In multi-head attention, we evenly divide the features of a model among its heads; for this post, we’ll stick to self-attention, but as we’ll see in the future, there are other types of attention where the multi-head mechanism also works. Suppose d_model = 512 and there are two heads of attention. Each head will be allocated 256 of the total 512 features that the model can reason about. This simple change has many downstream effects. Let’s start by examining the effect on the queries, keys, and values.

The Q, K, and V matrices will be produced once per head. Each head will have its own queries, keys, and values — and corresponding weight matrices. In our case, therefore, we will have the following:

• Head 1: Q⁽¹⁾ = XW_Q⁽¹⁾, K⁽¹⁾ = XW_K⁽¹⁾, and V⁽¹⁾ = XW_V⁽¹⁾

• Head 2: Q⁽²⁾ = XW_Q⁽²⁾, K⁽²⁾ = XW_K⁽²⁾, and V⁽²⁾ = XW_V⁽²⁾

Here, each weight matrix W has dimensions d_model x d_head, where d_head = d_model / 2, and each linear projection has dimensions n x d_head.

The rest of the self-attention math largely remains the same, except that we replace d_model with d_head where d_model had been used previously. So, S⁽¹⁾ and A⁽¹⁾ would still have size n x n, and Y⁽¹⁾ would have size n x d_head. At the end of the self-attention block, the outputs Y⁽¹⁾ and Y⁽²⁾ are concatenated together to produce an output of size n x d_model that can be presented as the output of self-attention.

Returning to our politician’s staff member analogy, the different heads conduct independent analyses of the input, and the reports produced by each are stapled together. This is not quite a synthesized result, however, so multi-head attention introduces one more step. The politician’s chief of staff will examine the stapled reports and produce an executive summary that pulls insights and information from both reports.

Mathematically, the chief of staff’s executive summary consists of a new linear projection that combines the outputs of each head of attention. In other words, suppose the matrix Y is a horizontal concatenation of the outputs of each head: [ Y⁽¹⁾ Y⁽²⁾ ]. The output of self-attention will now be Z = YW_O. Here, the dimensions of Y are n x d_model, the dimensions of W_O are d_model x d_model, and the dimensions of the final output Z are n x d_model.

Recap

Today, we re-covered what single-head attention is, and took a look at what changes about the math when we do multi-head attention. Fortunately, the change is not too dramatic. Each head of attention is assigned a subset of the total feature space d_model to work with, and we just perform attention for each head as we did before — only with a smaller feature space d_head < d_model.

There is one extra step when performing multi-head attention. Once each head has produced its output, a final linear projection will mix the results from each head together. This mixed result synthesizes the different perspectives from each head into a unified understanding of the input sequence. From our politician’s staff analogy, this is the Chief of Staff producing an executive summary encapsulating the perspectives of the various experts on the staff.

Next time, we’ll take a look at another key component of deep learning pipelines: optimizers. We’ve already seen hints of what these are. The simple one that we used initially was stochastic gradient descent, and the more complex one we introduced but did not explain in the previous post is Adam.