Given Q of shape (B, n_heads, T, d_head), the per head attention output before concatenation has shape ___.

Fill in blank·Easy·4.0 · 0·~1 min·Asked atFiddler Ai Mistral AI Tata Digital·Relevant atMicrosoft

Attempt it

Given Q of shape (B, n_heads, T, d_head) and the standard self-attention formula Attn(Q, K, V) = softmax(QK^T / sqrt(d_k)) V, the per head attention output before concatenation has shape .

**Trace the shapes step by step.** Start with Q, K, V each of shape `(B, n_heads, T, d_head)` where B is batch size, T is sequence length, and d_head is the per-head dimension. The first matmul is `Q @ K^T`. K has shape `(B, n_heads, T, d_head)`; transposing the last two dims gives `(B, n_heads, d_head, T)`. Matmul against Q produces `(B, n_heads, T, T)`. This is the per-head attention-score matrix: for each batch element, for each head, a `(T, T)` matrix saying how strongly each query token attends to each key token. Scaling by `1 / sqrt(d_head)` and applying softmax along the last dimension does not change shape; the result is still `(B, n_heads, T, T)`. Now every row of the inner matrix is a probability distribution over the T key positions. The final matmul is `attention_weights @ V`. The weights are `(B, n_heads, T, T)`; V is `(B, n_heads, T, d_head)`. The inner T dimensions collapse, and the result has shape `(B, n_heads, T, d_head)`. So each query token in each head ends up with a single d_head-wide context vector, which is the weighted sum of all V vectors using the attention weights as the mixing coefficients. The full output tensor has shape `(B, n_heads, T, d_head)`, ready to be concatenated across heads back into `(B, T, d_model)`.

The per-head attention output has shape `(B, n_heads, T, d_head)`, which is identical to the shape of Q (and of K and V). The arithmetic is short. Start with Q, K, V all of shape `(B, n_heads, T, d_head)`. Compute scores: $$\text{scores} = \frac{QK^\top}{\sqrt{d_\text{head}}}, \quad \text{shape} = (B, n_\text{heads}, T, T)$$ Apply softmax along the last dimension; shape is unchanged. Multiply by V, where V is `(B, n_heads, T, d_head)`. The last T of the scores collapses with the first T of V, yielding shape `(B, n_heads, T, d_head)`. The useful intuition: each query token receives one context vector per head. The context vector has dimension d_head because it is a weighted sum of V vectors, and each V vector is d_head-wide. After this, you typically `transpose` the n_heads axis with the T axis and `view` to merge n_heads with d_head, producing `(B, T, d_model)` where `d_model = n_heads * d_head`. That reshape is the 'concatenate heads' step. The output projection W_O then operates on the d_model dimension. Memorizing the shape pipeline is foundational: it lets you debug any attention implementation by trace-printing tensor shapes at each step and catching the bug at the first place the shape diverges from the expected sequence.

TL;DR

The per-head output is (B, n_heads, T, d_head). The softmax produces a (T, T) attention matrix per head, then multiplying by V of shape (T, d_head) recovers d_head columns.

Memory aid

Easy to grasp

Picture a librarian handling T books at once. For each book the librarian asks 'how related are you to every other book in the room?' and gets back a list of T relatedness scores, one per other book. The librarian then makes one summary per book by mixing the contents of every book according to those scores. The summary for each book is a single page exactly as wide as one book's original page count. The librarian does this for every book (T books) and ends up with T summaries, each the same width. Now imagine a whole roomful of librarians doing this in parallel, each focused on a different angle, for several different stories at once. The full pile of summaries has the shape (B, n_heads, T, d_head).

Key concepts

Multi-head attention Tensor shape arithmetic

Concept explanation~2 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

Tensor-shape arithmetic is the most foundational debugging skill for any transformer implementation. Most attention bugs manifest as shape mismatches, and being able to predict the shape at every step of the attention pipeline is the difference between a one-minute fix and a one-hour debug session.

This deep dive walks the shape transformations through scaled dot-product attention step by step, explains why each step has the shape it does, surveys the common axis-ordering variants that appear in production code, and connects the shape pipeline to optimizations like FlashAttention that preserve the output shape while reorganizing the computation.

Mental model: the attention output keeps the batch, head, and query-position axes intact, and the trailing dimension matches V's trailing dimension (d_head). Everything else is bookkeeping.

Walking the shapes step by step

Start with the canonical inputs to scaled dot-product attention in the multi-head layout: Q, K, V all of shape (B, n_heads, T, d_head) where B is batch, n_heads is the number of attention heads, T is the sequence length, and d_head is the per-head feature dimension.

Step 1: compute scores Q @ K^T

K has shape (B, n_heads, T, d_head). Transposing the last two dimensions gives (B, n_heads, d_head, T). The matmul Q @ K^T contracts the d_head axis on Q with the d_head axis on K^T, leaving:

\text{scores} \in \mathbb{R}^{B \times n_\text{heads} \times T \times T}

For each batch element and each head, you have a (T, T) matrix where entry (i, j) is the dot product between query token i's representation and key token j's representation.

Step 2: scale and softmax

Divide by sqrt(d_head) for variance control; this is an element-wise operation that does not change shape. Apply softmax along the last dimension (the key axis). Softmax normalizes each row to sum to 1, producing a per-query probability distribution over keys. Shape is unchanged: still (B, n_heads, T, T).

Step 3: matmul against V

V has shape (B, n_heads, T, d_head). The matmul weights @ V contracts the T-axis on weights with the T-axis on V (the second to last axis), producing:

\text{output} \in \mathbb{R}^{B \times n_\text{heads} \times T \times d_\text{head}}

This is the per-head attention output, with the same shape as Q. Each query token has received one d_head-wide context vector per head, computed as a weighted average of all V vectors with weights from the softmax output.

Why the output shape matches Q (and V)

Axis-ordering variants in the wild

How FlashAttention preserves the output shape

Situations where this technique stops working.

2–4 min · Everything important, quickly.

python

import torch
import torch.nn.functional as F

B, n_heads, T, d_head = 2, 8, 16, 64
Q = torch.randn(B, n_heads, T, d_head)
K = torch.randn(B, n_heads, T, d_head)
V = torch.randn(B, n_heads, T, d_head)

scores = Q @ K.transpose(-2, -1) / (d_head ** 0.5)
print(scores.shape)  # (B, n_heads, T, T)

weights = F.softmax(scores, dim=-1)
print(weights.shape)  # (B, n_heads, T, T) -- unchanged by softmax

per_head_output = weights @ V
print(per_head_output.shape)  # (B, n_heads, T, d_head)

# Concatenate heads: (B, n_heads, T, d_head) -> (B, T, d_model)
d_model = n_heads * d_head
concatenated = per_head_output.transpose(1, 2).contiguous().view(B, T, d_model)
print(concatenated.shape)  # (B, T, d_model)

Real products, models, and research that use this idea.

PyTorch's torch.nn.functional.scaled_dot_product_attention returns output of shape (B, n_heads, T, d_head) by default.
FlashAttention v2's API takes (B, T, n_heads, d_head) and returns the same shape; internally it tiles over T to avoid materializing the (T, T) attention matrix.

What an interviewer would ask next. Try answering before peeking at the approach.

QHow do the shapes change in cross-attention versus self-attention?

In cross-attention, Q comes from the decoder (length T_q) while K and V come from the encoder (length T_kv). The score matrix becomes (B, n_heads, T_q, T_kv). The output is (B, n_heads, T_q, d_head): one context vector per decoder query.

2 more follow-ups an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

Confusing the shape of the attention-weight matrix (B, n_heads, T, T) with the shape of the attention output (B, n_heads, T, d_head). The matmul against V brings the trailing dimension back to d_head.

60 second bullets to scan on the way to the call.

Shape of Q, K, V in the standard multi-head layout
Shape of QK^T after the first matmul

Primary sources. Browse if you want the original framing.