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}.\]

Using one-based teaching labels, the slice \(X^{(\ell)}[1,2,:]\) is one six-number vector: batch item 1, token position 2, every feature. In PyTorch's zero-based syntax, the same slice is x[0, 1, :].

\[X^{(\ell)}[1,2,:]=\begin{bmatrix}0.1&0.8&-0.3&1.2&0.2&0.4\end{bmatrix}.\]

The interactive labels use one-based human numbering; translate to zero-based indices when writing a tensor slice in code.

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?

Engineering bridge · the embedding contract

MiniMind names this lookup embed_tokens. Its model forward first unpacks (B, T), embeds integer IDs, and applies the configured dropout while preserving [B, T, D] for every Transformer block.

Show the model entry point

batch_size, seq_length = input_ids.shape       # [B, T]
hidden_states = self.dropout(self.embed_tokens(input_ids))
# hidden_states: [B, T, config.hidden_size]

Contract check: IDs must be integer values with 0 <= id < vocab_size; only the final axis is created by the embedding table.

Retrieval check

Answer once before reopening the code; use the feedback to correct your explanation.

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

STAGE 4 · PROVE IT

Practice studio · recall → build → verify

For \(B=3,T=5,d_{\text{model}}=8\), how many scalars are in \(X\)?
What does \(X[2,4,:]\) select?
What is the shape of \(X[:,3,:]\)?
If the block input is \(4\times128\times512\), what must its final output shape be for residual addition?
Explain why \(T\) counts token positions but \(d_{\text{model}}\) does not.

BUILD CHALLENGE · RUN IT LOCALLY

Engineering challenge · write a shape-safe selector

Implement `select_token(x, batch, token)` for `x.shape == [B, T, D]`. Return one `[D]` vector and raise a useful error for an invalid index. Test it with `B=2, T=3, D=8`.

Standard answer · compare after your attempt

def select_token(x, batch, token):
    B, T, D = x.shape
    if not (0 <= batch < B and 0 <= token < T):
        raise IndexError("batch/token outside the tensor")
    return x[batch, token, :]

assert select_token(torch.zeros(2, 3, 8), 1, 2).shape == (8,)

The contract is a slice, not a reshape: only the batch and token indices are fixed; every feature remains.

Check solutions

\(3\cdot5\cdot8=120\).
The eight-feature vector for batch item 2 at token position 4.
\(B\times d_{\text{model}}=3\times8\) in this example.
\(4\times128\times512\).
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\).