Fill in the tensor shapes as a batch flows through one transformer block
Attention output and block output both ride the d_model residual stream; only the FFN hidden activation widens to 4 d_model before projecting back.
Picture a busy hallway with a fixed width. People (features) move along it. A meeting room (attention) opens off the hallway, but anything coming out has to fit back into the hallway because that is the only path. Later, a workshop (the feed forward) opens up. Inside the workshop people spread out across four big tables, mix things around, then squeeze back through the door into the same narrow hallway. The hallway width never changes. The workshop tables are wider, but only briefly. That hallway width is d_model, and the wider tables are the 4 d_model FFN hidden.
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.
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.
Every transformer block is a shape-preserving function. The input enters as [B, T, d_model], every internal computation eventually folds back into that same width, and the output leaves as [B, T, d_model] so the next block can pick up where this one left off. Once you internalize this contract, the rest of the block's anatomy becomes obvious.
The block has two sublayers: attention and a feed forward network. Each one reads from a normalized copy of the stream, computes a delta, and adds that delta back into the stream. The delta has to be the same shape as the stream, because tensor addition requires matching dimensions. So attention output: [B, T, d_model]. FFN output: [B, T, d_model]. The only place the block briefly leaves this width is inside the FFN sublayer, between its two linear projections.
This deep dive walks through the shape contract end to end, with the specific tensor shapes at each step, and explains why the 4x FFN expansion is the standard ratio.
Attention sublayer: reshape, compute, project back
Three projections. From the input X of shape [B, T, d_model], the block computes Q = X W_Q, K = X W_K, V = X W_V. Each weight matrix is [d_model, d_model] in the original transformer, producing Q, K, V each of shape [B, T, d_model]. For grouped query attention, W_K and W_V are narrower ([d_model, n_kv_heads * d_head]), but for now assume vanilla MHA.
Reshape into heads. The tensor is viewed as [B, T, n_heads, d_head] where d_head = d_model / n_heads. A transpose then gives [B, n_heads, T, d_head], which is the layout attention actually consumes. This is just a view; no data movement is required, only a stride change.
Per-head attention. Each head computes softmax(Q K^T / sqrt(d_head)) V, yielding [B, n_heads, T, d_head]. Heads run in parallel and never interact during the softmax.
Concatenate and project. The output of per-head attention is reshaped back to [B, T, d_model] (heads stitched together along the last dim). The output projection W_O of shape [d_model, d_model] mixes information across heads and gives the residual update. The block adds this update back to the stream: X = X + W_O(\text{concat}(\text{heads})).
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- Llama 3 70B uses d_model 8192 with SwiGLU at hidden width 28672 (close to (8/3) * 8192), matching the classic 4x parameter budget
- Mixtral 8x7B keeps the same d_model 4096 and SwiGLU shape per expert, switching only the FFN routing across 8 experts
What an interviewer would ask next. Try answering before peeking at the approach.
QIf you doubled n_heads while keeping d_model fixed, which shapes change and which stay the same?
The reshape inside attention shifts because d_head halves. Block input and output, attention sublayer output, FFN hidden, and FFN output all remain unchanged.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Writing the attention output as n_heads times d_head, forgetting that the output projection collapses heads back to d_model before the residual add.
60 second bullets to scan on the way to the call.
Why the residual stream width must equal d_model end to end
How attention reshapes into heads and projects back via W_O
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.