Building a Transformer

David I. Inouye

Dynamic Content · Last Updated: October 30, 2026

A transformer combines context mixing and feature transformation

Transform features → preserve a residual path → normalize → stack blocks.

Keep the same task: use an email prefix to predict its next token.

Attention mixes positions; a shared MLP transforms each row

Three token positions are drawn side by side with inputs below and outputs above. Three token rows enter the same MLP independently. Each expands from d to m features, applies ReLU, and returns to d; there are no cross-position arrows.

\[\operatorname{MLP}(H)=\operatorname{ReLU}(HW_1^{\mathsf T}+\boldsymbol{1}\boldsymbol{b}_1^{\mathsf T})W_2^{\mathsf T}+\boldsymbol{1}\boldsymbol{b}_2^{\mathsf T}\]

\[H\in\mathbb{R}^{L\times d},\quad W_1\in\mathbb{R}^{m\times d},\quad W_2\in\mathbb{R}^{d\times m}\]

Same parameters at every position; separate parameters in each block.

A residual path adds a learned update without changing the shape

Reading upward, an L by d state branches into an unchanged bypass and a learned update F, also L by d. Elementwise addition returns L by d.

\[H_{\mathrm{out}}=H+F(H)\]

The attention output projection and the MLP’s last layer both return \(d\) features.

LayerNorm normalizes one token’s features at a time

Like BatchNorm: subtract a mean, divide by a standard deviation, then learn a scale and shift.

\[\mu_\ell=\frac1d\sum_{j=1}^{d}h_{\ell j},\qquad \sigma_\ell^2=\frac1d\sum_{j=1}^{d}(h_{\ell j}-\mu_\ell)^2\]

\[\operatorname{LN}(\boldsymbol{h}_\ell)=\boldsymbol{\gamma}\odot \frac{\boldsymbol{h}_\ell-\mu_\ell\boldsymbol{1}} {\sqrt{\sigma_\ell^2+\epsilon}}+\boldsymbol{\beta}\]

BatchNorm (CNN): one channel across batch examples and spatial locations.
LayerNorm: one token across its \(d\) features—the last axis of [B,L,d].

At evaluation, LayerNorm still uses that token’s statistics; BatchNorm usually uses running statistics.

Four minutes. Learned gamma and beta each have d entries and are shared across token positions. Contrast with the previously discussed BatchNorm: this LayerNorm uses no other examples or positions and has no running batch statistics. Epsilon is a small positive constant. Population variance divides by d. With column h=[1,3]^T, gamma=[1,1]^T, beta=[0,0]^T, and epsilon approaching zero, the normalized row approaches [-1,1]^T. After learned scaling/shift, the output need not have zero mean or unit variance. Both train and evaluation compute statistics from the current token row. PyTorch nn.LayerNorm(d) matches this choice; normalizing over [L,d] would mix positions and can leak future information. [Sources] source:pytorch-normalization-2-7, LayerNorm documentation and forward implementation; source:pytorch-batchnorm2d-2-7. BatchNorm is shown in the CNN layout already taught; in a fully connected layer it instead aggregates each feature across examples. Both methods use a learned per-feature/channel affine transform, but normalize over different groups. PyTorch BatchNorm defaults to tracking running statistics; training/evaluation details and variance estimators differ, so the modules are not otherwise identical.

A pre-norm block adds two shape-preserving updates

From bottom to top, H passes through LayerNorm, causal MHA and residual addition, then LayerNorm, a row MLP and residual addition. Both stages preserve L by d shape.

\[U=H+\operatorname{MHA}_{\mathrm{causal}}(\operatorname{LN}_1(H))\] \[H_{\mathrm{out}}=U+\operatorname{MLP}(\operatorname{LN}_2(U))\]

Pre-norm is the usual modern LLM default. The original Transformer used post-norm.

Pre-norm	Post-norm
\(H+F(\operatorname{LN}(H))\)	\(\operatorname{LN}(H+F(H))\)
Normalize before the branch.	Normalize after addition.

Pre-norm leaves the bypass untouched.

Four minutes. Follow each branch and bypass in the diagram. LN1 and LN2 have separate learned scale/shift parameters. MHA includes all Q/K/V projections, causal masking, head concatenation and output projection. Each intermediate is L by d; the MLP expands internally and returns to d. Dropout is omitted for clarity. This is our chosen pre-norm variant, not the original 2017 post-norm arrangement or an exact named production model. The full model needs a final normalization after the stack. F denotes either the attention or MLP branch. The slide calls pre-norm the usual modern LLM default, not a universal standard across all Transformer architectures: recent models also use other placements, including Peri-LN. Pre-norm has favorable initialization-gradient behavior in Xiong et al., which helps explain its use for deep stacks; do not claim that post-norm is obsolete or every model uses pre-norm. Llama 2 uses pre-normalization with RMSNorm, a different normalization operator, not exactly the LayerNorm shown here. The remembered original Add & Norm diagram performs the residual sum first and LayerNorm after it, not LayerNorm immediately before the plus operation. [Sources] source:xiong-layernorm-transformer-2020, Figure 1 and Table 1; source:vaswani-attention-2017; source:touvron-llama2-2023, section 2.2; source:kim-periln-2025, sections 1–2.

Stacked blocks turn token features into next-token logits

Token IDs at the bottom enter learned embeddings plus additive positions, then N distinct pre-norm blocks, final LayerNorm, and a vocabulary head producing L by V logits. An alternative replaces additive positions with RoPE in every block's queries and keys.

\[H^{(0)}=E[w_{0:L},:]+P,\qquad Y=\operatorname{LN}_{\mathrm{final}}(H^{(N)})W_{\mathrm{vocab}}^{\mathsf T}+\boldsymbol{1}\boldsymbol{c}^{\mathsf T}\]

\[H^{(N)}\in\mathbb{R}^{L\times d},\qquad W_{\mathrm{vocab}}\in\mathbb{R}^{v\times d},\qquad Y\in\mathbb{R}^{L\times v}\]

Four minutes. E is the learned vocabulary table, w_0:L is the code-style slice of L input IDs, and P stacks transposed position columns as rows. N denotes block count, preserving L for sequence length. Head output dimension v is vocabulary size (V in code), distinct from the earlier projected value matrix V. Apply softmax to each vocabulary-logit row and compare with the shifted target. Our illustrative model does not tie embedding/head weights. For the alternative RoPE variant use H0=E[w,:] without P, then rotate Q/K within every attention sublayer; retain embeddings, residuals, MLPs, final normalization and vocabulary head. Do not silently combine both positional schemes. Both variants still generate autoregressively. [Sources] source:vaswani-attention-2017; source:xiong-layernorm-transformer-2020; source:su-roformer-2021.

One length-\(L\) input solves \(L\) prefix-prediction problems

Along the bottom, five inputs BOS, please, send, the, report feed causal output rows. Targets at the top are shifted one position: please, send, the, report, EOS. Every output reads only its prefix and contributes one loss.

\[\mathcal{L}(\theta)=-\frac{1}{L}\sum_{\ell=0}^{L-1}\log p_\theta(w_{\ell+1}\mid w_0,\ldots,w_\ell).\]

Known inputs allow parallel work; the mask preserves each allowed prefix.

Four minutes. The same synthetic sequence used for RNN training supplies five inputs and five shifted targets. BOS is the start token and EOS is the final target. Output y_ell is a column of vocabulary logits; the model stacks its transpose in row ell. Softmax and cross-entropy act independently on each output row before averaging the losses. Explain the report prediction: the output at input “the” sees BOS please send the, but cannot see the following input “report,” even though that token is present elsewhere in the training array. Each row equals the last-row prediction from running only its allowed prefix, with matching positions and deterministic settings. Thus we solve L overlapping prefix problems in one shared pass, without separately running L growing prefixes. With dropout, independent runs can differ because of sampled masks; this does not change causal information access. Padding positions are excluded from the loss in variable-length batches. [Sources] source:vaswani-attention-2017, sections 3.1 and 5.1; source:d2l-rnn.

Backpropagation follows different dependency graphs

The same five inputs and targets are used for both models. The RNN has a left-to-right hidden-state chain with reverse gradient arrows. A simplified transformer layer has causal connections between two rows of token states; the last row's reverse paths fan back to all allowed earlier positions. All losses contribute gradients.

Transformer benefit: parallel work across positions in forward and backward computation.

Five minutes. The loss boxes abbreviate the vocabulary head and cross-entropy; embeddings, normalization, residuals and MLPs are omitted for readability. In the RNN, h_ell depends on h_(ell-1), so hidden states require a sequential recurrence even when all inputs are known. Backpropagation through time follows that chain in reverse and adds contributions from losses at all positions. Input projections or output heads can still be batched; “sequential” refers to the recurrent state dependency, not every operation.

In the transformer, every upper row reads only allowed rows from the PREVIOUS layer, not already-updated rows of the same layer. All upper rows can therefore be computed together with masked matrix operations. Backpropagation reverses these edges and accumulates contributions to earlier representations and shared parameters; orange attention edges highlight only the final row to avoid clutter. All other rows also contribute. It does not require an RNN-style token-by-token reverse recurrence. Layers remain dependent in depth; gradient reductions, attention work, memory traffic and hardware determine actual speed. This capacity to exploit parallel hardware is a major computational advantage for training, not a claim of a measured speedup for our untrained model. Generation still depends on unknown future input choices. The numerical probe checks both logits and parameter gradients against the sum of separate prefix losses. [Sources] source:vaswani-attention-2017, section 4; source:d2l-bptt.

Can one forward pass score all shifted targets?

The training batch contains each full email.

Which model change would let an early prediction read its future target?

A. Compute all causal output rows together

B. Normalize each token over its own features

C. Remove the causal attention mask

C. Parallel computation is valid when every path respects the allowed prefix.

A causal decoder is the common decoder-only LLM architecture

Three upward architecture diagrams. Encoder-only uses Q/K/V from all input embeddings and a task head, with ViT as an example. Our causal decoder uses Q/K/V from the prefix with causal self-attention. The original 2017 encoder–decoder has source self-attention, target causal self-attention, and cross-attention with Q from target and K/V from source. Residual and normalization paths are omitted.

Our email model is a causal decoder. Other tasks can use different information paths and heads.

Three minutes. Wrap up training before broadening the architecture family. The diagrams show attention sources, with residual paths, normalization and repeated depth omitted. Q/K/V boxes denote learned projections, not raw token IDs; embeddings also include the chosen position mechanism. An encoder lets every supplied position attend to every other supplied position. A Vision Transformer uses patch embeddings, unrestricted encoder attention, and an image-classification head (typically from a class-token representation). The causal decoder-only form used in many generative LLMs predicts the next token from an allowed prefix. Not every language model is decoder-only. The original 2017 Transformer was an encoder–decoder for sequence-to-sequence prediction, demonstrated on translation. An encoder–decoder contextualizes a source sequence, then generates a target using causal target self-attention plus cross-attention to the full encoded source. Source and target lengths can differ. In cross-attention the target supplies Q and the source supplies K/V. This brief inset introduces information flow, not a second block derivation or a translation experiment.

Shared parameters and a different dependency structure provide different context access: fixed windows omit distant inputs, RNNs compress them through a sequential state chain, and attention reads allowed earlier representations directly. Access is a capability, not evidence of better predictions or length generalization. Full dense attention still has quadratic pairwise work. The RNN/bigram comparison remains the measured evidence; this transformer requires evaluation under the same frozen protocol. Next, return to our causal decoder and turn its next-token distributions into a generated continuation. [Sources] source:vaswani-attention-2017; source:touvron-llama2-2023; source:d2l-rnn; source:dosovitskiy-vit-2020, section 3 and Figure 1.