Combined Projections Introduce a Head Axis

One combined projection can reproduce separate head projections when its final feature axis is split and transposed correctly. Trace named axes through reshape and transpose, and diagnose operations that preserve element count but change meaning.

Computing each head with a separate matrix is clear notation, but an implementation can combine those matrices into one larger projection. The mathematics does not change; the head axis becomes visible through reshape and transpose operations.

Concatenate Projection Matrices by Columns

Let

W_Q=\left[W_Q^{(1)}\;W_Q^{(2)}\;\cdots\;W_Q^{(h)}\right].

If every head has query width $d_k$ , then

W_Q:d_{model}\times(hd_k).

The combined product

Q_{flat}=XW_Q

has shape $(B,T,hd_k)$ . Its final axis contains the head-1 query features, followed by head 2, and so on. Reshaping exposes that organization:

(B,T,hd_k)\rightarrow(B,T,h,d_k).

This operation does not reorder values. It gives names to two factors of the same final axis.

Transpose for Parallel Attention

Attention is easiest to express with the head axis before the two position axes:

(B,T,h,d_k)\rightarrow(B,h,T,d_k).

Now the score product is

(B,h,T,d_k)(B,h,d_k,T)\rightarrow(B,h,T,T).

The causal mask may broadcast from $(1,1,T,T)$ across batch and head axes. Softmax still normalizes only the final key-position axis. Heads and queries do not compete in one shared softmax.

The same transforms apply to values:

V:(B,T,hd_v)\rightarrow(B,T,h,d_v)\rightarrow(B,h,T,d_v).

After attention, the head results reverse the process:

(B,h,T,d_v)\rightarrow(B,T,h,d_v)\rightarrow(B,T,hd_v).

The last step joins the head and per-head value axes. It must not merge the position axis into the features.

Reshape and Transpose Do Different Jobs

Reshape factors or joins axes without choosing a new semantic order.
Transpose changes the order of named axes.
Concatenate joins separate tensors along a stated axis.

Some array libraries require a contiguous copy after transposition before a view-like reshape can join axes. That is a storage-layout constraint, not a new attention equation. The semantic requirement is that the resulting feature axis contain each position's head outputs in the declared order.

The Separate and Combined Forms Must Agree

For the numerical trace, combining $P_1$ and $P_2$ by columns gives a $4\times4$ identity matrix. Projecting $X$ once therefore returns its four features. Reshaping those features into two width-2 heads reconstructs exactly the separate $P_1$ and $P_2$ results.

This is a strong implementation test: calculate heads separately, calculate them with combined projections, and compare every head output before $W_O$ .

Q1. Repair the head transform

A combined query projection has shape $(B,T,hd_k)=(2,5,12)$ with $h=3$ and $d_k=4$ . Which sequence gives the shape needed for parallel attention?

Choose one

Reshape to (2, 5, 3, 4), then transpose to (2, 3, 5, 4)

Reshape to (2, 3, 4, 5), then keep that order

Reshape to (3, 2, 5, 4), moving heads before the batch

Select one choice, then check.

Hint

Begin with axis names

(B,T,hd_k)

Solution

Reshape to

(B,T,h,d_k)=(2,5,3,4)

, then transpose the middle axes to obtain

(B,h,T,d_k)=(2,3,5,4)

Trace Axis Names, Not Only Sizes

Two axes may both have length 8 and still mean different things. Record $(B,h,T,d_k)$ beside the tensor during each transformation. Numeric dimensions alone cannot reveal that a position axis and head axis were accidentally swapped.