A Short Course for Modern Large Language ModelsGlossary

Lesson 15 of 23 · Transformer block

Read the tensor before the formula.

Your win: explain every axis in \(X^{(\ell)}\in\mathbb R^{B\times T\times d_{\text{model}}}\) and follow one token row through the block.

12 minutesNeeds: vectors and matricesOutcome: keep a shape ledger
Mission link: nearly every confusion in a Transformer block becomes easier once you know which axis is batch, token position, or feature.

Three axes, three jobs

\[X^{(\ell)}\in\mathbb R^{B\times T\times d_{\text{model}}}.\]

Batch \(B\)

How many independent sequences are processed together.

Tokens \(T\)

How many positions are present in each sequence.

Features \(d_{\text{model}}\)

How many numbers represent each token at layer \(\ell\).

Numbers substituted

Let \(B=2\), \(T=4\), and \(d_{\text{model}}=6\). Then:

\[X^{(\ell)}\in\mathbb R^{2\times4\times6},\qquad 2\cdot4\cdot6=48\text{ scalars}.\]

The slice \(X^{(\ell)}[1,2,:]\) is one six-number vector: batch item 1, token position 2, every feature.

\[X^{(\ell)}[1,2,:]=\begin{bmatrix}0.4&-0.2&1.1&0.0&0.7&-0.5\end{bmatrix}.\]

Explore a \(2\times4\times6\) hidden-state tensor

Token ↓Feature coordinates →

The block preserves the outer shape

\(B\times T\times d_{\text{model}}\) → attention workspace → \(B\times T\times d_{\text{model}}\) → FFN workspace → \(B\times T\times d_{\text{model}}\)

Internal widths change temporarily, but each residual addition requires the sublayer output to return to \(d_{\text{model}}\).

Index discipline: use \(i\) for a query position and \(j\) for a key/value position. Both range over the token axis \(T\); neither is a feature dimension.

Code checkpoint · tokens become hidden states

The implementation starts with integer token IDs and looks up one learned feature vector per token. The batch and token axes stay in place.

Show the PyTorch shape transition
token_ids = torch.tensor([[4, 7, 2], [9, 1, 5]])  # [B=2, T=3]
embedding = nn.Embedding(vocab_size=32, embedding_dim=8)
x = embedding(token_ids)                           # [B=2, T=3, D=8]

assert x.shape == (2, 3, 8)

Trace it: Which axis is introduced by the embedding lookup, and which two axes are preserved?

Retrieval check

What is the shape of \(X[b,t,:]\)?

Practice before moving on

  1. For \(B=3,T=5,d_{\text{model}}=8\), how many scalars are in \(X\)?
  2. What does \(X[2,4,:]\) select?
  3. What is the shape of \(X[:,3,:]\)?
  4. If the block input is \(4\times128\times512\), what must its final output shape be for residual addition?
  5. Explain why \(T\) counts token positions but \(d_{\text{model}}\) does not.
Check solutions
  1. \(3\cdot5\cdot8=120\).
  2. The eight-feature vector for batch item 2 at token position 4.
  3. \(B\times d_{\text{model}}=3\times8\) in this example.
  4. \(4\times128\times512\).
  5. The \(T\) axis indexes sequence locations; the feature axis stores the numerical representation at each location.

Primary source: Vaswani et al., “Attention Is All You Need,” Sections 3.1–3.2. Track every occurrence of \(d_{\text{model}}\), \(d_k\), and \(d_v\).

Ask the teaching agent to give you three tensor shapes and check your verbal interpretation.