Carbon Note 核心学堂 Carbon Note Academy

Practice: Mathematical Foundations and Complete Architecture Derivation of Transformer Models

Since Vaswani et al. introduced the Transformer in 2017, the architecture has served as the absolute mathematical foundation for Large Language Models (LLMs) and Generative AI.

To truly master the Transformer, one must grasp the rigorous mathematical machinery powering its forward activations and backward pass gradients. This laboratory report presents a systematic, step-by-step mathematical derivation of the Transformer architecture, ranging from numerical stability analysis to complete backward gradient flows.

1. Scaled Dot-Product Attention

The Self-Attention layer computes semantic correlations across sequence tokens. It accepts three main inputs: Query $\mathbf{Q}$, Key $\mathbf{K}$, and Value $\mathbf{V}$.

A. Core Forward Formulations

The attention output matrix $\mathbf{O}$ is formulated as (split column view of LaTeX code and math representation):

$\mathrm{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \mathrm{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}\right)\mathbf{V}$

Where $\mathbf{Q} \in \mathbb{R}^{n \times d_k}$, $\mathbf{K} \in \mathbb{R}^{n \times d_k}$, and $\mathbf{V} \in \mathbb{R}^{n \times d_v}$. $n$ represents the sequence token length, and $d_k, d_v$ are feature projection dimensions.

To break this down, we define the raw similarity logits as $\mathbf{S}$:

$\mathbf{S} = \frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}} \quad \Rightarrow \quad S_{ij} = \frac{1}{\sqrt{d_k}} \sum_{m=1}^{d_k} Q_{im} K_{jm}$

Applying a row-wise Softmax normalization yields the attention weight probability matrix $\mathbf{A}$:

$A_{ij} = \frac{\exp(S_{ij})}{\sum_{p=1}^{n} \exp(S_{ip})}$

The final representation $\mathbf{O} \in \mathbb{R}^{n \times d_v}$ is the weighted summation of values:

$\mathbf{O} = \mathbf{A}\mathbf{V} \quad \Rightarrow \quad O_{ij} = \sum_{q=1}^{n} A_{iq} V_{qj}$

B. Numerical Stability: Why divide by $\sqrt{d_k}$?

Without the scaling factor $\sqrt{d_k}$, extreme feature dimensions push dot-product outputs to high-magnitude regions, causing the Softmax function's gradient to saturate and disappear.

Proof: Assume query vector components $\mathbf{q} \in \mathbb{R}^{d_k}$ and key vector components $\mathbf{k} \in \mathbb{R}^{d_k}$ are independent, identically distributed (i.i.d.) random variables with mean $0$ and unit variance $1$:

$\mathbb{E}[q_i] = \mathbb{E}[k_i] = 0, \quad \mathrm{Var}(q_i) = \mathrm{Var}(k_i) = 1 \quad (\forall i = 1, \dots, d_k)$

Let the unscaled dot product be $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$. Due to independent components:

$\mathbb{E}[q_i k_i] = \mathbb{E}[q_i] \mathbb{E}[k_i] = 0$

$\mathrm{Var}(q_i k_i) = \mathbb{E}[(q_i k_i)^2] - (\mathbb{E}[q_i k_i])^2 = \mathbb{E}[q_i^2] \mathbb{E}[k_i^2] - 0 = (\mathrm{Var}(q_i) + \mathbb{E}[q_i]^2)(\mathrm{Var}(k_i) + \mathbb{E}[k_i]^2) = 1 \cdot 1 = 1$

Since the summations are mutually independent, variance additive properties apply:

$\mathbb{E}[q \cdot k] = \sum_{i=1}^{d_k} \mathbb{E}[q_i k_i] = 0$

$\mathrm{Var}(q \cdot k) = \sum_{i=1}^{d_k} \mathrm{Var}(q_i k_i) = d_k$

To scale the dot product back to a stable distribution with mean $0$ and unit variance $1$, we divide by $\sqrt{d_k}$:

$\mathrm{Var}\left(\frac{q \cdot k}{\sqrt{d_k}}\right) = \frac{1}{d_k} \mathrm{Var}(q \cdot k) = \frac{d_k}{d_k} = 1$

This proves why the scaling parameter $\sqrt{d_k}$ is mathematically necessary. $\text{Q.E.D.}$

2. Multi-Head Attention Mechanism

Multi-Head Attention projects representation tokens into multiple subspaces, allowing parallelized representation pathways (split column view):

$\begin{aligned}
\mathrm{MultiHead}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) &= \mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_h)\mathbf{W}^O \\
\text{where } \mathrm{head}_i &= \mathrm{Attention}(\mathbf{Q}\mathbf{W}_i^Q, \mathbf{K}\mathbf{W}_i^K, \mathbf{V}\mathbf{W}_i^V)
\end{aligned}$

Dimensional constraints:

• $\mathbf{W}i^Q \in \mathbb{R}^{d{\text{model}} \times d_k}$
• $\mathbf{W}i^K \in \mathbb{R}^{d{\text{model}} \times d_k}$
• $\mathbf{W}i^V \in \mathbb{R}^{d{\text{model}} \times d_v}$
• $\mathbf{W}^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$

3. Residuals & Layer Normalization (LayerNorm)

LayerNorm normalizes activation flows across the hidden dimension of each token independently, stabilizing forward paths.

A. Layer Statistics

For layer representation output $\mathbf{x} \in \mathbb{R}^d$, we calculate internal mean $\mu$ and variance $\sigma^2$:

$\mu = \frac{1}{d} \sum_{j=1}^{d} x_j, \quad \sigma^2 = \frac{1}{d} \sum_{j=1}^{d} (x_j - \mu)^2$

B. Normalization and Scale-Shift Mapping

Using computed statistics, we scale and shift via learnable scalar weights $\gamma$ and bias $\beta$:

$\hat{x}_j = \frac{x_j - \mu}{\sqrt{\sigma^2 + \epsilon}}$

$y_j = \gamma_j \hat{x}_j + \beta_j$

Where $\epsilon$ is a tiny floating-point value to prevent division by zero.

4. Feed-Forward Network (FFN)

Each attention layer is followed by a Position-wise Feed-Forward network, containing two linear projections and a ReLU (or GELU) activation:

$\mathrm{FFN}(\mathbf{x}) = \max(0, \mathbf{x}\mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2$

Where $\mathbf{W}1 \in \mathbb{R}^{d{\text{model}} \times d_{\text{ff}}}$ and $\mathbf{W}2 \in \mathbb{R}^{d{\text{ff}} \times d_{\text{model}}}$.

5. Sine-Cosine Positional Encodings

Since Attention matrices are permutation-invariant, positional features must be added to input embeddings:

$ \begin{aligned} PE_{(pos, 2i)} &= \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) \ PE_{(pos, 2i+1)} &= \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) \end{aligned} $

6. Attention Layer Backward Pass Gradient Derivation

To detail the backpropagation flow, we calculate complete derivatives for the Scaled Dot-Product Attention module.

Let the incoming upstream loss gradient relative to the attention output $\mathbf{O}$ be $\frac{\partial \mathcal{L}}{\partial \mathbf{O}} \in \mathbb{R}^{n \times d_v}$.

A. Gradient with respect to Value Matrix $\mathbf{V}$

Since $O_{ij} = \sum_{q=1}^{n} A_{iq} V_{qj}$, we take partial derivatives:

$\frac{\partial \mathcal{L}}{\partial V_{kj}} = \sum_{i=1}^{n} \sum_{p=1}^{d_v} \frac{\partial \mathcal{L}}{\partial O_{ip}} \frac{\partial O_{ip}}{\partial V_{kj}} = \sum_{i=1}^{n} \frac{\partial \mathcal{L}}{\partial O_{ij}} A_{ik}$

Expressed as a clean matrix product:

$\frac{\partial \mathcal{L}}{\partial \mathbf{V}} = \mathbf{A}^T \frac{\partial \mathcal{L}}{\partial \mathbf{O}}$

B. Gradient with respect to Attention Matrix $\mathbf{A}$

Differentiating with respect to $A_{ik}$:

$\frac{\partial \mathcal{L}}{\partial A_{ik}} = \sum_{j=1}^{d_v} \frac{\partial \mathcal{L}}{\partial O_{ij}} V_{kj}$

Represented in matrix form:

$\frac{\partial \mathcal{L}}{\partial \mathbf{A}} = \frac{\partial \mathcal{L}}{\partial \mathbf{O}} \mathbf{V}^T$

C. Gradient with respect to Logits Matrix $\mathbf{S}$

Softmax maps lines dependently. Using classic Softmax derivatives:

$\frac{\partial A_{ij}}{\partial S_{ik}} = A_{ij}(\delta_{jk} - A_{ik}) \quad (\text{where } \delta_{jk} \text{ is Kronecker delta})$

Applying multivariate chain rule:

$\frac{\partial \mathcal{L}}{\partial S_{ik}} = \sum_{j=1}^{n} \frac{\partial \mathcal{L}}{\partial A_{ij}} \frac{\partial A_{ij}}{\partial S_{ik}} = \sum_{j=1}^{n} \frac{\partial \mathcal{L}}{\partial A_{ij}} A_{ij} (\delta_{jk} - A_{ik}) = \frac{\partial \mathcal{L}}{\partial A_{ik}} A_{ik} - A_{ik} \sum_{j=1}^{n} \frac{\partial \mathcal{L}}{\partial A_{ij}} A_{ij}$

Vectorizing using Hadamard products ($\odot$) and row-wise summation operators ($\mathrm{rowsums}$):

$\frac{\partial \mathcal{L}}{\partial \mathbf{S}} = \mathbf{A} \odot \left( \frac{\partial \mathcal{L}}{\partial \mathbf{A}} - \mathrm{rowsums}\left(\frac{\partial \mathcal{L}}{\partial \mathbf{A}} \odot \mathbf{A}\right) \mathbf{1}^T \right)$

Where $\mathbf{1} \in \mathbb{R}^n$ is an all-ones column vector.

D. Gradients with respect to Query $\mathbf{Q}$ and Key $\mathbf{K}$

Differentiating $\mathbf{S} = \frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}$:

$\mathrm{d}\mathbf{S} = \frac{1}{\sqrt{d_k}} \left( (\mathrm{d}\mathbf{Q})\mathbf{K}^T + \mathbf{Q}(\mathrm{d}\mathbf{K})^T \right)$

Converting to gradients:

$\frac{\partial \mathcal{L}}{\partial \mathbf{Q}} = \frac{1}{\sqrt{d_k}} \frac{\partial \mathcal{L}}{\partial \mathbf{S}} \mathbf{K}$

$\frac{\partial \mathcal{L}}{\partial \mathbf{K}} = \frac{1}{\sqrt{d_k}} \left( \frac{\partial \mathcal{L}}{\partial \mathbf{S}} \right)^T \mathbf{Q}$

We have derived all forward and backward equations representing self-attention loops inside the Transformer block.

7. Common Problems and Pitfalls (FAQ)

1. Why does my multi-line equation or matrix fail to break lines, collapsing everything into a single row?
   • Solution: In LaTeX, line breaks are achieved using double backslashes \\. In a Markdown text editor, ensure there are no trailing whitespace characters or spaces immediately following the \\ command. Also, make sure to leave an empty line before and after the math blocks for optimal rendering.
2. How do I write standard words (like "softmax" or "Attention") inside an equation without them rendering in ugly squished italics?
   • Solution: By default, LaTeX treats sequential letters as a product of individual variables, rendering them in tight italics. To display words properly, use the roman font operator \mathrm{softmax} or the text utility \text{Attention} to restore normal font weights and natural spacings.
3. Why does my subscript look wrong? For example, W_attention only lowers the 'a', leaving the other letters on the main line.
   • Solution: The subscript (_) and superscript (^) operators in LaTeX only apply to the single character directly following them. For multiple characters, you must wrap them in curly braces {}, such as W_{attention} or d_{k}.