Transformer architecture: trace one complete block

Summary

A transformer block lets positions exchange information, then transforms each position’s features. Residual additions and normalization connect those operations without changing the token-by-feature output shape.

Task: trace a two-token, two-feature pre-norm block with input $X = [[1, 3], [3, 1]]$ . Use per-row LayerNorm with unit gain, zero bias, and zero epsilon only for this nonconstant-row calculation. Production normalization uses declared positive epsilon. Set query/key projections to zero and value/output projections to identity, with one causal attention head. Use a row-wise FFN with identity matrices, zero biases, and ReLU. What reaches each residual output, and can the second input row change the first output?

The minimum mental model

A transformer block is:

x = x + Attention(LayerNorm(x))   # mix across positions
x = x + FFN(LayerNorm(x))         # mix across features

The second line uses the first line’s updated tensor, not the original input. Write these states as $X_{1} = X + A (LN (X))$ and $X_{2} = X_{1} + F (LN (X_{1}))$ . Token embeddings and the task’s output projection sit outside this block.

The retained figure shows unmasked encoder connectivity; our numerical trace changes visibility to a causal decoder mask. Attention also contains feature projections: the picture distinguishes primary communication roles, not exclusive ownership of feature mixing.

One tensor, two dependency patterns

Attention mixes token positions; the feed-forward layer transforms each position independently

Read it this way: Attention can exchange information across these token rows. Each feed-forward application transforms only its own row at this layer. Can one feed-forward application exchange information between different token rows?

Reference drawing

One tensor, two dependency patterns

Attention mixes token positions; the feed-forward layer transforms each position independently

Derivation

hold the T × d tensor shape fixed. Here, unmasked encoder attention connects every token row; a causal decoder omits future links. The shared FFN transforms each row independently. Each sublayer contributes an update through its own residual addition. Attention's internal feature projections are not shown. Sources: Vaswani et al. (2017) and the PyTorch encoder-layer documentation.

Check: Can token 1 directly affect token 2 inside the FFN sublayer? Where can that communication occur instead?

Compute the two residuals

Each input row has mean $2$ and variance $1$ , so normalized rows are $[- 1, 1]$ and $[1, - 1]$ . Zero queries and keys give equal logits on allowed positions. The causal first row sees only itself; the second averages both value rows.

State	Row 0	Row 1
$X$	$[1, 3]$	$[3, 1]$
$LN (X)$	$[- 1, 1]$	$[1, - 1]$
Attention contribution	$[- 1, 1]$	$[0, 0]$
First residual $X_{1}$	$[0, 4]$	$[3, 1]$
$LN (X_{1})$	$[- 1, 1]$	$[1, - 1]$
Row-wise FFN contribution	$[0, 1]$	$[1, 0]$
Second residual $X_{2}$	$[0, 5]$	$[4, 1]$

The first residual’s row zero has mean $2$ and variance $4$ ; renormalization still gives $[- 1, 1]$ . ReLU acts within each row, and the same FFN parameters serve both rows. No FFN edge directly joins different positions, although its inputs already contain attention-derived context.

Change normalization placement

New task: use post-norm, $h = LN (X + A (X))$ , then $Y = LN (h + F (h))$ . Picture and change: move normalization outside each addition and recompute branch inputs too. Check: isolate zero branches on row $[1, 3]$ . Pre-norm returns $[1, 3]$ ; post-norm returns $[- 1, 1]$ under the same toy normalization. Use when: the architecture changes residual-normalization order; its direct identity path is no longer the same.

Change position visibility

New task: encode a complete document bidirectionally instead of predicting a next token. Remove the causal restriction where allowed, while retaining padding validity and position semantics. With the original input, both attention rows now average to $[0, 0]$ , giving first final row $[1, 4]$ .

Check: replace the second input row with $[1, 3]$ . Bidirectional attention now changes the first final row to $[0, 5]$ . The causal block’s first final row remains $[0, 5]$ before and after that perturbation. Use when: the entire input is available; the future-perturbation invariant is no longer required.

Shapes and failure boundaries

For $T$ tokens and width $d$ , each residual contribution must return shape $T \times d$ . A head computes $Q = Z W_{Q}$ , $K = Z W_{K}$ , $V = Z W_{V}$ , and $softmax (Q K^{⊤} / d_{k} + M) V$ . The additive mask $M$ is zero on allowed keys and negative infinity on forbidden keys. Require an allowed key per query; concatenate head outputs and project back to residual width. The position-wise FFN expands $d$ to its chosen hidden width, applies a nonlinearity or gate, then returns to $d$ .

A causal mask alone does not validate training if teacher-forced inputs expose the target token. Check shifted inputs, target alignment, padding, and absolute positions as part of the prediction contract.

What an interviewer expects you to articulate

Name the objective and available positions before naming a model family. Encoder-only, decoder-only, and encoder-decoder systems differ in visibility, cross-attention, and prediction targets. See their task-specific routes rather than assuming every transformer is a causal language model.

Reconstruct: does representing an identity with zero residual branches prove a deeper model always trains at least as well?

Answer

No. Representational availability does not make optimization find that identity. Normalization, finite precision, regularization, and changed training dynamics also prevent an automatic performance ordering.

The why behind each piece

Why attention?

Attention permits content-dependent communication among allowed keys. Dense unmasked attention connects the supplied context in one layer; causal, sparse, and window masks restrict that graph. Parallel training across supplied positions does not remove autoregressive generation dependencies or finite context limits.

Why multi-head?

Heads form separate learned query/key/value subspaces and attention distributions. No theorem assigns each head a particular linguistic role. At fixed total projection widths, parameter counts can match a single-head design. Logits and softmax storage still depend on head count.

Why residual connections?

Residual addition supplies a direct route alongside a learned branch. Its derivative includes an identity contribution, but branch derivatives can cancel or amplify it. Neither vanishing-gradient prevention nor successful deeper training follows automatically.

Why LayerNorm (not BatchNorm)?

Feature-axis LayerNorm computes statistics within each token, avoiding cross-sample statistical coupling. BatchNorm is a different reduction and train/evaluation contract, not mathematically impossible for sequence inputs. The configured normalized axes matter; see the shared-statistics trace.

Why pre-norm vs post-norm?

Their equations and identity paths differ, as the core’s zero-branch check shows. Depth, initialization, warmup, residual scaling, optimizer, and final normalization affect empirical stability and quality. The original transformer used post-norm; that does not define every later architecture.

Why FFN?

It supplies a shared row-wise nonlinear transformation alongside attention’s communication. Attention’s projections already mix features; the FFN is not the sole feature-mixing operation. Hidden width is architectural, not necessarily four times residual width.

Why positional encoding?

Without positional information or asymmetric masks, row permutation permutes attention outputs: the operation is equivariant, not invariant. Causal masks already introduce an order-dependent restriction. Position schemes add distinct absolute or relative signals; defined beyond-training-length values do not guarantee useful extrapolation.

The scale story

Kaplan et al. and Hoffmann et al. fit empirical loss/resource relationships under specified training regimes. Compute-optimal parameter/token allocation depends on those fits and constraints. They do not guarantee monotonic downstream improvements from every added parameter or dataset.

What’s changed since 2017

Choice	Actual change
RMSNorm	Scales by a root mean square without subtracting the token mean.
SwiGLU	Adds a multiplicative gate and separate input projections.
RoPE	Rotates queries and keys by position-dependent angles.
ALiBi	Adds a relative-distance term to allowed attention scores.
Learned or sinusoidal embeddings	Add a specified positional vector to the token representation.
GQA / MQA	Share key/value heads among query heads, changing cache dimensions.
Conditional experts	Route tokens among FFNs with additional capacity and dispatch contracts.
Bias or dropout removal	Changes the specified training computation, not an automatic quality-preserving simplification.

Common confusions

Parameter sharing, masks, normalization, and position mechanisms impose structural priors. The dominant compute or memory term depends on width, length, batch, implementation, and training versus generation. Architecture names do not replace a shape, visibility, or resource ledger.