Transformers and attention deep dive

For large d_k, the dot products q·k grow with variance proportional to d_k (assuming unit-variance components), pushing the softmax into regions with extremely small gradients (saturation). Dividing by sqrt(d_k) keeps the variance of the scores ~1, preventing softmax saturation and stabilizing gradients during training.

scores = (Q @ K.transpose(-2,-1)) / math.sqrt(d_k)
attn = scores.softmax(dim=-1)
out = attn @ V

Both compute and memory are O(n^2 * d) for the QK^T score matrix and its softmax, since every token attends to every other token. The n^2 scaling dominates for long sequences, and materializing the n x n matrix makes memory the practical bottleneck, which motivates approaches like FlashAttention and sparse/linear attention.

FlashAttention tiles the computation into blocks that fit in fast SRAM and uses an online (streaming) softmax with running max and normalizer statistics, so it never materializes the full n x n score matrix in HBM. It is IO-aware: it trades extra recomputation for far fewer slow memory reads/writes, reducing memory from O(n^2) to O(n) while producing exact (not approximate) outputs.

Self-attention is permutation-equivariant, so without positional information it cannot distinguish token order. Sinusoidal encodings are fixed functions added to embeddings (extrapolate somewhat); learned absolute encodings are trained per-position (limited to trained length); RoPE rotates query/key vectors by position-dependent angles, encoding relative position directly in the dot product and generalizing better to longer contexts.

Multiple heads let the model jointly attend to information from different representation subspaces and positions simultaneously, learning distinct relationships (e.g., syntactic vs. semantic) in parallel. A single large head would average these into one attention pattern, reducing the diversity of relations the layer can capture, even at equal parameter/compute cost.

MQA uses a single shared key/value head for all query heads; GQA groups query heads to share a smaller number of K/V heads (between MQA and full MHA). This shrinks the KV cache size during autoregressive decoding, the dominant memory/bandwidth cost at inference, with minimal quality loss compared to full multi-head attention.

During generation the keys and values of previously processed tokens do not change, so they are cached and reused instead of recomputed at each step. This turns per-step attention into computing only the new token's query against cached K/V, reducing per-token cost and avoiding redundant recomputation of the prefix. Its size grows with sequence length and is a major inference memory consumer.

Post-LN (original) applies LayerNorm after the residual addition; Pre-LN applies it inside the residual branch before the sublayer. Pre-LN keeps a clean identity residual path, yielding better-behaved gradients at initialization and allowing stable training of deep models often without learning-rate warmup, whereas Post-LN tends to need warmup and careful tuning but can reach slightly better final performance.

# Pre-LN block
x = x + attn(layernorm1(x))
x = x + ffn(layernorm2(x))

It enforces autoregression: position i may only attend to positions <= i, preventing information leakage from future tokens that would make next-token prediction trivial during teacher-forced training. It is implemented by adding -inf (a large negative value) to the upper-triangular entries of the score matrix before softmax, driving those attention weights to zero.

mask = torch.triu(torch.ones(n,n), diagonal=1).bool()
scores.masked_fill_(mask, float('-inf'))

If you permute the input tokens, the self-attention outputs are permuted identically (with no positional encoding), because attention treats inputs as an unordered set. This matters because language and time series are order-dependent, so positional information must be injected externally; it also explains why Transformers can be applied to set-structured data naturally.

The FFN is two linear layers with a nonlinearity, applied independently to each position, typically expanding from d_model to ~4*d_model and back (e.g., 512 -> 2048 -> 512). It provides per-token nonlinear feature transformation and is where a large share of parameters and compute reside; modern variants use GLU/SwiGLU gating for better quality.

RoPE multiplies query and key vectors by a rotation matrix whose angle depends on absolute position; because rotations compose, the dot product q_m·k_n depends only on the relative offset (m - n). This makes attention scores a function of relative distance, which generalizes more gracefully to unseen sequence lengths and supports interpolation/extension techniques (e.g., NTK/linear scaling, YaRN).

An attention sink is the tendency of many heads to allocate large attention weight to the first token(s) regardless of content. Because softmax forces weights to sum to 1, heads that have 'nothing relevant' to attend to dump probability onto an always-present anchor (often the BOS/first token), which acts as a no-op. Preserving these sink tokens is crucial for streaming/sliding-window inference (StreamingLLM).

Softmax couples all keys through normalization, giving expressive data-dependent weighting but forcing the O(n^2) score matrix. Linear attention replaces exp(q·k) with a kernel feature map phi: phi(q)·phi(k), so by associativity you compute phi(K)^T V once (O(n) in sequence length) and reuse it, trading some expressiveness/accuracy for linear complexity.

# linear attention: O(n) via associativity
KV = phi(K).transpose(-2,-1) @ V
out = phi(Q) @ KV

Residual (skip) connections add the sublayer input to its output, creating an identity path that lets gradients flow directly to earlier layers and mitigates vanishing gradients in deep stacks. They also let each block learn a perturbation of the identity, which stabilizes optimization and is what makes very deep Transformers trainable when combined with normalization.

In cross-attention, queries come from the decoder's current hidden states, while keys and values come from the encoder's output representations. This lets each decoding position selectively read from the full encoded source sequence, which is the mechanism that aligns target tokens to relevant source content (e.g., in translation).

It restricts each token to attend only to a fixed local window of w neighbors, reducing complexity from O(n^2) to O(n*w) for long sequences. Stacking layers with windowed attention grows the effective receptive field linearly with depth (like a CNN), and it is often combined with a few global tokens to preserve long-range information.

Embeddings are often initialized with small variance (and weights may be tied with the output projection); multiplying by sqrt(d_model) increases their magnitude so they are not dominated by the sinusoidal positional encodings, which have unit-scale components. This balances the relative contributions of content and position at the input.

BERT uses bidirectional self-attention (no causal mask) trained with masked language modeling, giving rich context-aware representations ideal for understanding/classification/retrieval. GPT uses causal (unidirectional) attention trained to predict the next token, making it naturally generative. Bidirectional models can't generate autoregressively without modification; causal models see only the left context.

Dividing logits by temperature T before softmax controls sharpness: T<1 sharpens toward a near-argmax (more peaked, lower entropy) distribution, while T>1 flattens it toward uniform (higher entropy). As T->0 softmax approaches a hard one-hot; as T->inf it approaches uniform. In attention this modulates how concentrated the weighting is; in sampling it controls output diversity.

attn = (scores / T).softmax(dim=-1)

Induction heads are attention heads that implement a copy/completion pattern: having seen '...A B... A', they attend back to the earlier occurrence of A and predict B (the token that followed it before). Their emergence during training correlates strongly with the onset of in-context learning, suggesting they are a core mechanistic ingredient for pattern continuation and few-shot prompting.

LayerNorm normalizes across the feature dimension per token, independent of other examples and sequence positions, so it works with variable-length sequences and small/variable batch sizes and behaves identically in training and inference. BatchNorm depends on batch statistics, which are noisy/ill-defined with variable sequence lengths and autoregressive decoding (batch size effectively 1 per step).

RMSNorm rescales activations by their root-mean-square (no mean subtraction and typically no bias), using only a learned gain. It drops the re-centering step of LayerNorm, which empirically contributes little, making it cheaper to compute while maintaining or improving stability and quality, hence its adoption in models like LLaMA.

def rmsnorm(x, g, eps=1e-6):
    return x * g / x.pow(2).mean(-1, keepdim=True).add(eps).sqrt()

Each head gets d_k = d_v = d_model/h = 512/8 = 64. The h heads compute attention independently in their 64-dim subspaces, their outputs are concatenated back to 512 dimensions, and a final learned linear projection W_O mixes them into the block output.

ALiBi adds no positional embeddings to inputs; instead it adds a static, head-specific linear penalty proportional to the query-key distance directly to the attention scores (closer tokens favored). This biases attention toward recency, requires no learned position parameters, and extrapolates well to sequences far longer than those seen in training.

scores = QK - slope_h * distance_matrix

High attention weight does not imply causal influence on the output: information also flows through residual streams, value vectors, and downstream layers, and different attention distributions can yield the same prediction. Studies show attention is often not faithful as an explanation, so attribution methods (e.g., gradient- or ablation-based) are more reliable for importance claims.

Online softmax computes the softmax over blocks incrementally without seeing all logits at once, maintaining a running maximum m and a running normalizer (sum of exp) l, plus the partially accumulated output. When a new block arrives, it rescales the previous accumulator by exp(m_old - m_new) to correct for an updated max, ensuring numerical stability and exactness.

MoE replaces a single FFN with many expert FFNs and a router that sends each token to only the top-k experts (sparse activation). This decouples total parameter count from per-token compute, giving large model capacity at roughly constant FLOPs per token, but adds routing complexity, load-balancing challenges, and higher memory/communication costs.

W_O is a learned linear map from the concatenated head outputs (h*d_v = d_model) back to d_model, mixing information across heads so subsequent layers see an integrated representation rather than independent per-head slices. Without it, head outputs would remain in separate subspaces and the residual stream would just be a concatenation, limiting cross-head interaction.

Naive extrapolation pushes RoPE rotation angles beyond the range seen in training, producing out-of-distribution high-frequency phases that the model never learned to interpret, degrading performance. Positional interpolation rescales position indices so they map into the trained range (compressing frequencies), keeping the per-dimension phases in-distribution; NTK-aware/YaRN scaling refine this by treating frequency bands differently.

QK-Norm applies normalization (e.g., L2 or LayerNorm/RMSNorm) to the query and key vectors before computing their dot product. It bounds the magnitude of attention logits, preventing the logit explosion / attention entropy collapse that can destabilize training of very large or long-context models, improving stability without changing the attention mechanism's structure.

Prefill processes the whole prompt in parallel as large matrix-matrix multiplies with high arithmetic intensity, saturating compute. Decode generates one token at a time: each step is a matrix-vector operation that must stream the entire model weights and KV cache from memory for very few FLOPs, so throughput is limited by memory bandwidth rather than compute.

Dropout regularizes the model to reduce overfitting; in the original Transformer it is applied to the output of each sublayer before the residual addition, to the attention weights (after softmax), and to the summed embeddings + positional encodings. Many large modern LLMs reduce or remove dropout because massive data and short training schedules make overfitting less of a concern.

With fixed d_model, more heads means smaller per-head dimension d_k = d_model/h, so each head operates in a lower-dimensional subspace with reduced representational capacity and noisier dot-product similarities. Beyond a point the loss of per-head expressiveness outweighs the benefit of more diverse attention patterns, and many heads become redundant or prunable.