LM3-T6 — Tensor Shapes and Matrix Operations in Attention

Learning objectives

After this lesson, you should be able to:

• Name the four dimensions B, n, H and d_h and give the shape of every tensor in multi-head attention.
• Explain matrix products of batched tensors, including broadcasting of the mask.
• Estimate memory and arithmetic cost for the projections and the attention matrices, and find the sequence length at which the n² term dominates.
• Recognise and debug common shape bugs, and use scaled_dot_product_attention.

The idea: dimensions with jobs

Four letters will carry you through any attention code:

• B: the batch, independent sequences processed together. No information crosses it.
• n (or T): the sequence length, the tokens in a sequence. Attention mixes information along this axis.
• H: the number of heads. Heads are independent of one another until the output projection.
• d and d_h: the model feature width and the width per head, d_h = d/H. Linear layers mix information along this axis.

Which axis does what?

A good way to read code is to ask, for each operation, which axis it acts on. A linear layer acts on the last axis (features). The softmax acts on the last axis of the score tensor, which is the key axis. The matrix product q @ k.transpose(2, 3) contracts the feature axis. The product weights @ v contracts the key axis. Anything that touches the batch axis is a bug in a model whose sequences are independent.

Batched matrix products

PyTorch’s @ on tensors with more than two dimensions treats all leading dimensions as batch dimensions and multiplies the last two as matrices. So (B, H, n, d_h) @ (B, H, d_h, n) gives (B, H, n, n): B · H independent matrix products of (n × d_h) by (d_h × n). This is why the head dimension is moved next to the batch dimension: it turns “do this for each head” into a single call.

Broadcasting

The causal mask has shape (n, n), but the scores are (B, H, n, n). Operations broadcast by aligning shapes from the right and stretching dimensions of size 1 (or missing leading dimensions), so the mask applies to every sequence and head without copies. This is convenient and also dangerous: a mask of shape (B, n) broadcast against (B, H, n, n) would silently align B with the wrong axis.

The mathematics: shapes, memory and arithmetic

With input x of shape (B, n, d) and H heads of width d_h = d/H, the shapes are:

Memory of the score tensor. The tensor in step 4 (and again after the softmax) has

B, H, n

batch size, number of heads, sequence length.

bytes

4 for float32, 2 for float16 or bfloat16.

M_scores

memory of one score tensor; training keeps the scores and the weights for the backward pass.

For B = 1, H = 12, n = 1,024 and float32: 12 · 1,048,576 · 4 = 50,331,648 bytes, about 50 MB per tensor per layer. At n = 4,096 it is 16 times larger, about 805 MB per tensor per layer.

Arithmetic (multiply-adds counted as 2 FLOPs). The four projections (Q, K, V, output) each multiply (B · n) × d by d × d: 2 · B · n · d² FLOPs each. The two attention products (QK^T and WV) each cost 2 · B · H · n² · d_h = 2 · B · n² · d. Totals:

Their ratio is n / (2d). So the attention-matrix products dominate once n > 2d. For GPT-2 small (d = 768) at n = 1,024 (B = 1): projections 8 · 1,024 · 768² = 4,831,838,208 FLOPs, attention products 4 · 1,024² · 768 = 3,221,225,472 FLOPs; the crossover is at n = 1,536. At 8,192 tokens the attention products would be over five times the projections.

A worked example: a shape trace you can run

The code below runs one layer for B = 2, n = 6, d = 8, H = 2 and checks every shape from the table with an assertion. It also compares the result with the fused function.

import torch
import torch.nn.functional as F

B, n, d, H = 2, 6, 8, 2
dh = d // H
torch.manual_seed(0)
x = torch.randn(B, n, d)
Wq, Wk, Wv = (torch.randn(d, d) / d ** 0.5 for _ in range(3))

q, k, v = x @ Wq, x @ Wk, x @ Wv
assert q.shape == (B, n, d)

split = lambda t: t.view(B, n, H, dh).transpose(1, 2)
qh, kh, vh = split(q), split(k), split(v)
assert qh.shape == (B, H, n, dh)

scores = qh @ kh.transpose(2, 3) / dh ** 0.5
assert scores.shape == (B, H, n, n)

mask = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)       # shape (n, n) broadcasts over B and H
weights = torch.softmax(scores.masked_fill(mask, float("-inf")), dim=-1)
assert torch.allclose(weights.sum(-1), torch.ones(B, H, n))

ctx = weights @ vh
assert ctx.shape == (B, H, n, dh)
merged = ctx.transpose(1, 2).contiguous().view(B, n, d)
assert merged.shape == (B, n, d)

fused = F.scaled_dot_product_attention(qh, kh, vh, is_causal=True)
print(torch.allclose(ctx, fused, atol=1e-5))                             # True
print(scores.numel(), "scores per layer for this tiny example")           # 144

# the same bug-prone step done wrongly: view without the transpose mixes tokens and features
wrong = q.view(B, H, n, dh)                                              # right shape, wrong meaning
print(wrong.shape == qh.shape, torch.allclose(wrong, qh))                # True False

The last two lines show the most instructive bug. q.view(B, H, n, dh) has exactly the shape of the correct tensor, so nothing raises an error, but it cuts the memory in the wrong places and mixes tokens with features. Only a comparison with a reference, or a leakage test, reveals it.

A memory estimate you can do on paper

GPT-2 small, n = 1,024, B = 8, float32, twelve layers. One score tensor is 8 × 50 MB = 403 MB. The weights after the softmax are the same size, and both are kept for the backward pass: about 0.8 GB per layer, 9.7 GB for twelve layers, from the attention matrices alone. This is why training scripts use smaller batches or fused kernels, and why the book’s code uses a context of 256 for demonstrations on a laptop.

Engineering implications and limitations

Print shapes while you write. Add a one-line shape check at each step (assert q.shape == (B, H, n, dh)). Remove or keep them as a cheap guard.

Make shapes distinct in tests. If B, n, H and d_h all equal 4, a transposed axis is invisible. Use different small numbers (2, 6, 2, 4) as the notebook does.

View needs contiguous memory. After transpose or permute call .contiguous() before view, or use reshape.

Softmax over the key axis. The last dimension of the scores. Normalising over the query axis gives distributions that do not mean what the model needs.

Dimension order in transpose. k.transpose(2, 3) or k.transpose(-2, -1) swaps the last two axes only; a bare k.T reverses all axes and breaks batched tensors.

Use the fused function for speed. torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) takes (B, H, n, d_h) tensors and returns the context with the same shape. On supported hardware it uses kernels that avoid building the n × n matrix, so memory grows linearly in n instead of quadratically. The cost is that you cannot inspect the weights; use the explicit version for debugging and teaching.

Cost at long context. Beyond n ≈ 2d the quadratic term dominates both time and memory. Long-context work (sparse, sliding-window and linear attention, Module 10) targets this term.

Mixed precision. Half-precision scores can overflow before the softmax; implementations compute the softmax in higher precision or use the stable form. Rely on library functions rather than rewriting them.

Common misconceptions

Engineering takeaways

1. Four dimensions carry attention: batch B, tokens n, heads H and feature width d (with d_h = d/H).
2. The score and weight tensors have shape (B, H, n, n); attention mixes tokens along the n axes, linear layers mix features.
3. Memory for scores is B·H·n²·bytes: about 50 MB per tensor per layer for GPT-2 small at n = 1,024 in float32.
4. FLOPs: projections 8·B·n·d², attention products 4·B·n²·d; the second dominates once n > 2d.
5. Check shapes with distinct numbers, compare against scaled_dot_product_attention, and run a leakage test.

References and further reading

See Chapter 3 of the book for the shapes in the multi-head implementation; the attention cost analysis here is standard arithmetic that you can verify with the code.

• PyTorch documentation: torch.nn.functional.scaled_dot_product_attention.
• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.
• PyTorch documentation: scaled_dot_product_attention (official documentation)
• Build a Large Language Model (From Scratch), Chapter 3: Coding Attention Mechanisms (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

That completes Module 3. Next, Module 4 assembles these pieces into a complete GPT model.

Related content

Nothing below is required to finish this lesson.

• Before
  LM3-T5: Multi-Head Attention

  Split attention into heads, compute them in parallel and combine them with an output projection.

• Lab
  LM3-L3: Implement Causal Multi-Head Attention in PyTorch

  Implement causal self-attention and multi-head attention in PyTorch and check them against a reference.