NN LAB 🧠

In the seminal paper Attention Is All You Need (Vaswani et al. 2017), the authors note that for large values of $d_k$, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. Here is the complete variance stabilization proof.

1. Statistical Formulation of the Raw Dot Product

Let $\mathbf{q}, \mathbf{k} \in \mathbb{R}^{d_k}$ be independent random query and key vectors whose components are independent and identically distributed (i.i.d.) random variables with zero mean and unit variance:

The unscaled dot product $S$ is given by the sum of pairwise products:

2. Expectation and Variance of Each Component Product

For each term $X_i = q_i k_i$, using the independence of $q_i$ and $k_i$:

The variance of each individual product term is:

3. Variance of the Total Sum

Because all $d_k$ terms $X_1, \dots, X_{d_k}$ are mutually independent, the variance of the sum is the sum of the variances:

Consequently, the standard deviation is $\sigma_S = \sqrt{d_k}$. In modern models where $d_k = 128$, the standard deviation of raw scores is $\sqrt{128} \approx 11.31$.

4. Vanishing Softmax Gradients Under High Variance

Consider the softmax function $p_i = \text{softmax}(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j} e^{z_j}}$. Its partial derivative with respect to any input logit $z_j$ is:

When logits have large variance ($\sigma \approx 11.3$), the maximum logit $z_{\max}$ severely outstrips all others, causing the softmax probability distribution to saturate into a one-hot distribution:$p_{\max} \approx 1.0$ and $p_k \approx 0.0$ for all $k \neq \max$. Substituting these saturated probabilities back into the derivative:

Gradients vanish entirely, and backpropagation fails to update earlier layers.

5. Temperature Rescaling to Unit Variance

By dividing the raw dot product by $\sqrt{d_k}$, the variance of the scaled attention scores becomes:

Key Takeaway: Scaling by $1/\sqrt{d_k}$ guarantees that the variance of the attention logits remains precisely $1.0$ regardless of projection dimension, maintaining softmax sensitivity and preventing gradient collapse across arbitrarily deep networks.

The placement of normalization layers represents one of the most critical architectural evolutions between the original 2017 Transformer and modern foundation LLMs (LLaMA, Mistral, Gemma, DeepSeek).

Why Modern LLMs Use Pre-RMSNorm: Pre-LN eliminates the need for delicate warmups and enables training networks with hundreds of layers (e.g. LLaMA 70B with 80 layers). Furthermore, replacing LayerNorm with RMSNorm ($x / \text{RMS}(x)$) discards mean centering, saving $\approx 10\%$ to $15\%$ of memory bandwidth without any degradation in perplexity.

On modern GPU hardware (NVIDIA A100, H100, B200), compute throughput (Tensor Cores) has scaled dramatically faster than memory bandwidth (HBM). Consequently, standard self-attention is not compute-bound—it is strictly memory bandwidth (IO) bound.

1. The Memory Hierarchy Bottleneck in Standard Attention

Modern accelerators feature a multi-tier memory hierarchy:

• High-Bandwidth Memory (HBM): Large capacity (80 GB), but relatively slow bandwidth ($\approx 3.35\text{ TB/s}$ on H100 SXM).
• On-Chip Static RAM (SRAM): Ultra-fast ($\approx 33\text{ TB/s}$, $10\times$ faster), but small capacity ($\approx 228\text{ KB}$ per Streaming Multiprocessor, $\approx 50\text{ MB}$ total).

Standard attention materializes intermediate matrices in HBM at each step:

1. Read $Q, K$ from HBM $\to$ compute $S = Q K^T \in \mathbb{R}^{N \times N}$ in SRAM $\to$ write $S$ back to HBM ($O(N^2)$ memory access).
2. Read $S$ from HBM $\to$ compute $A = \text{softmax}(S)$ in SRAM $\to$ write $A$ back to HBM ($O(N^2)$ memory access).
3. Read $A, V$ from HBM $\to$ compute $O = A V \in \mathbb{R}^{N \times d}$ $\to$ write $O$ to HBM.

2. FlashAttention Tiling & Online Softmax Algorithm

FlashAttention (Dao et al. 2022, 2023) completely avoids materializing the quadratic $N \times N$ attention matrix in HBM. It partitions inputs into blocks $Q_i \in \mathbb{R}^{B_r \times d}$ and $K_j, V_j \in \mathbb{R}^{B_c \times d}$ that fit entirely within fast on-chip SRAM. To compute softmax without seeing all tokens at once, it utilizes the Online Softmax recurrence (Milakov & Gimelshein 2018):

Theoretical IO Complexity: FlashAttention reduces total HBM data transfers from $O(N^2)$ down to $O(N)$ while computing mathematically exact (non-approximate) attention. This yields a $2\times$ to $4\times$ wall-clock speedup and enables context windows of $128\text{k}$ to $1\text{M}$ tokens.

In production LLM serving, generation takes place token-by-token in an autoregressive loop. While the initial prompt processing (prefill) is compute-heavy, autoregressive generation (decoding) is strictly limited by memory bandwidth.

1. Arithmetic Intensity of Autoregressive Decoding

To generate token $t+1$, the query is a single vector $\mathbf{q}_{t+1} \in \mathbb{R}^{1 \times d_{\text{model}}}$. However, to attend to all preceding context tokens, the GPU must fetch all previously computed Key and Value vectors for all layers from HBM:

Modern Tensor Cores operate at hundreds of teraflops, yet memory bus speeds are orders of magnitude lower. During generation, GPU compute cores idle for over 95% of execution time waiting for KV cache memory loads!

2. Memory Scaling Comparison: MHA vs MQA vs GQA

Consider serving a 70B parameter model ($L = 80$ layers, $H = 64$ heads, $d_{\text{head}} = 128$) with sequence length $S = 4096$ and batch size $B = 16$ in FP16 precision (2 bytes):

The Architectural Verdict: While MQA achieves maximum bandwidth reduction, empirical research demonstrates quality degradation on complex reasoning and long-context retrieval tasks. GQA strikes the optimal Pareto frontier, achieving $95\%$ of MQA's speed with $100\%$ of MHA's quality.