The Transformer architecture relies on self-attention. For an input sequence \(X \in \mathbb{R}^{n \times d_{\text{model}}}\), three learned projections produce queries \(Q = XW_Q\), keys \(K = XW_K\), and values \(V = XW_V\), each representing what information a token is looking for, advertising what it contains, and what content to output, respectively. Scaled dot-product attention \(\text{Attention}(Q,K,V) = \text{softmax}(QK^\top/\sqrt{d_k})V\) computes compatibility scores between queries and keys (scaled to prevent softmax saturation), converts them to probabilities via softmax applied row-wise, then weight-averages the values to produce each token's context-aware output. The scaling factor \(1/\sqrt{d_k}\) is needed because \(\mathrm{Var}(q \cdot k) = d_k\)—without scaling, large \(d_k\) pushes softmax toward one-hot, zeroing gradients. The attention weight matrix \(A = \text{softmax}(QK^\top/\sqrt{d_k}) \in \mathbb{R}^{n \times m}\) has row \(i\) as a probability distribution over keys. Multi-head attention runs \(h\) heads in parallel: \(\text{head}_i = \text{Attention}(QW_Q^i, KW_K^i, VW_V^i)\), concatenates results, and projects via output matrix \(W^O \in \mathbb{R}^{hd_v \times d_{\text{model}}}\). With \(d_k = d_v = d_{\text{model}}/h\), total parameters are exactly \(4d_{\text{model}}^2\)—the same as four square matrices regardless of \(h\). Each head operates in its own subspace, capturing diverse patterns simultaneously.
Causal (masked) self-attention enforces autoregression for language models. The mask sets future positions to \(-\infty\) before softmax: \(S_{ij}^{\text{masked}} = S_{ij}\) if \(j \le i\), else \(-\infty\). Implemented as matrix addition \(S^{\text{masked}} = S + M\) with upper-triangular \(-\infty\) mask \(M\), this makes the resulting attention matrix strictly lower-triangular—each token attends only to itself and past tokens. Cross-attention differs by sourcing Q from the decoder and K, V from the encoder: \(Q = H_{\text{dec}}W_Q\), \(K = H_{\text{enc}}W_K\), \(V = H_{\text{enc}}W_V\), allowing decoder queries over encoder representations for source-target alignment. Self-attention differs from a fully connected layer by computing input-dependent weights \(A_{ij} = \text{softmax}(q_i \cdot k_j/\sqrt{d_k})\) that adapt dynamically to each input, rather than using fixed learned parameters. Computational complexity of self-attention is \(\mathcal{O}(n^2 d)\)—quadratic in sequence length is the main bottleneck for long contexts. Flash Attention tiles Q, K, V into SRAM blocks and fuses score computation, softmax, and value aggregation into one GPU kernel, avoiding materializing the full \(n \times n\) score matrix in HBM—reducing memory traffic from \(\mathcal{O}(n^2)\) to \(\mathcal{O}(n)\) and providing significant wall-clock speedups.
Transformers need explicit positional information since attention is permutation-equivariant. Sinusoidal positional encoding assigns \(\text{PE}(\text{pos}, 2i) = \sin(\text{pos}/10000^{2i/d_{\text{model}}})\) and \(\text{PE}(\text{pos}, 2i+1) = \cos(\text{pos}/10000^{2i/d_{\text{model}}})\); pairing sin and cos at each frequency forms 2D rotations enabling relative offset encoding via linear transformations. Position vectors are added to token embeddings: \(\tilde X_i = E_i + \text{PE}(i)\). Learned positional embeddings (BERT) provide flexibility but limit max training length. Relative positional encodings condition attention on offsets: \(S_{ij}^{\text{rel}} = (x_iW_Q)(x_jW_K + r_{i-j}W_K^R)^\top/\sqrt{d_k}\), generalizing better to longer sequences by capturing distance rather than absolute position. The attention masking operation \(A_{ij} = -\infty\) if \(j > i\), else \(z_{ij}\) implements causal masking cleanly—\(-\infty\) becomes 0 after softmax. Weight sharing in self-attention (same \(W_Q, W_K, W_V\) across positions) reduces parameters and enforces equivariance. Knowledge distillation transfers behavior from teacher to student using temperature-softened outputs. The neural tangent kernel \(K(\mathbf{x},\mathbf{x}') = \nabla_\theta f_\theta(\mathbf{x})^\top \nabla_\theta f_\theta(\mathbf{x}')\) describes the infinite-width regime where networks behave like kernel methods.
Each transformer block combines attention with a feed-forward sublayer and residual connections. The FFN \(\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2\) applies position-wise (independent per token, same matrices reused), typically with \(d_{ff} = 4d_{\text{model}}\) expansion. Pre-LN design normalizes before the sublayer: \(y = x + \text{Sublayer}(\text{LayerNorm}(x))\); the original Post-LN normalized after—Pre-LN improves training stability and is now standard (GPT-2, modern LLMs). Layer normalization \(\text{LayerNorm}(x) = \gamma \odot (x-\mu)/(\sigma+\epsilon) + \beta\) standardizes per token across features (with \(\mu = \frac{1}{d}\sum_i x_i\), \(\sigma\) standard deviation), then rescales with learned \(\gamma, \beta\)—preferred over batch norm in transformers because it's batch-independent. The residual formula \(y = x + \text{Sublayer}(x)\) allows gradients to flow directly to earlier layers without passing through the sublayer. The encoder-decoder Transformer has N encoder layers (bidirectional self-attention + FFN) feeding K, V to M decoder layers (causal self-attention + cross-attention + FFN). The linear classifier head projects the final representation \(\hat{\mathbf{y}} = W_{\text{cls}}\mathbf{h} + \mathbf{b}_{\text{cls}}\) with \(W_{\text{cls}} \in \mathbb{R}^{K \times d}\). During autoregressive decoding, the KV cache \(K_{1:t} = [k_1, \ldots, k_t]\), \(V_{1:t} = [v_1, \ldots, v_t]\) stores past projections to avoid recomputation, reducing per-step complexity to \(\mathcal{O}(td)\) and making inference practical. Gradient checkpointing trades computation for memory—recomputing activations during backward pass reduces memory from \(O(N)\) to \(O(\sqrt{N})\) with ~33% extra compute. Mixed precision training scales the loss \(s \cdot \mathcal{L}\) before backward pass and rescales gradients \(s^{-1}\mathbf{g}\) afterward to prevent FP16 underflow.