Thursday, April 17, 2025

The Mathematical Foundation of Transformers

Mathematical Foundation of the Transformer Architecture (Detailed)

Introduction: The Transformer Paradigm Shift

The Transformer model, presented in the paper "Attention Is All You Need" (Vaswani et al., 2017), marked a significant departure from prevalent sequence-to-sequence architectures like Recurrent Neural Networks (RNNs) and Long Short-Term Memory networks (LSTMs). Instead of processing sequences step-by-step (recurrently), which inherently limits parallelization and struggles with long-range dependencies, the Transformer relies entirely on attention mechanisms. These mechanisms allow the model to weigh the importance of different parts of the input sequence when processing a specific part, regardless of their distance. This design enables massive parallelization and has proven highly effective for various Natural Language Processing (NLP) tasks and beyond.

We will dissect the mathematical underpinnings of its key components: Input Processing, Attention Mechanisms, Encoder Layers, Decoder Layers, Output Generation, and the crucial aspects of Training and Inference.

Core Concepts and Notation (Glossary)

Before diving into the formulas, let's define the key symbols and concepts used throughout:

  • Sequences:
    • x = (x_1, ..., x_n): The input sequence of tokens (e.g., words, subwords).
    • y = (y_1, ..., y_m): The target output sequence of tokens.
    • n: Length of the input sequence.
    • m: Length of the output sequence.
  • Dimensionality:
    • d_model: The core dimensionality used for embeddings and throughout most layers of the Transformer (e.g., 512, 768). This represents the "width" of the information pathway.
    • d_k: Dimensionality of keys and queries in the attention mechanism.
    • d_v: Dimensionality of values in the attention mechanism.
    • d_{ff}: Dimensionality of the inner layer in the Feed-Forward Networks (often 4 * d_model).
    • V_{size}: The size of the vocabulary (total number of unique tokens the model knows).
  • Matrices and Vectors:
    • W (e.g., W^Q, W^K, W^V, W^O, W_1, W_2, W_e, W_{final}): Learnable weight matrices used in linear transformations. These matrices contain the parameters the model learns during training. The superscripts or subscripts indicate their specific role.
    • b (e.g., b_1, b_2, b_{final}): Learnable bias vectors, added after matrix multiplications in linear transformations.
    • X, Y, Z, Q, K, V: Matrices representing collections of vectors (e.g., X often represents the matrix of input embeddings, where each row is a token's vector).
    • x_i, y_j, z, v: Individual vectors (often representing a single token or position).
  • Parameters:
    • θ: Represents the entire set of learnable parameters in the model ({W..., b..., γ..., β...}). Training aims to find the optimal θ.
    • γ, β: Learnable scale and shift parameters used in Layer Normalization.
  • Hyperparameters: Parameters set before training, defining the architecture and training process.
    • h: Number of parallel attention heads in Multi-Head Attention.
    • N: Number of identical layers stacked in the encoder and decoder.
    • p_{drop}: Dropout probability (rate at which activations are set to zero during training).
    • ε (epsilon): A small constant used for numerical stability (e.g., in Layer Normalization, Adam optimizer).
    • η: Learning rate for the optimizer.
    • β_1, β_2: Exponential decay rates for moment estimates in the Adam optimizer.
    • λ: Weight decay coefficient (for regularization).
    • k: Beam width (for beam search decoding).
  • Operators and Functions:
    • Matrix Multiplication: Standard matrix multiplication (e.g., QK^T). Dimensions must align.
    • T: Transpose operator (swaps rows and columns of a matrix, e.g., K^T).
    • +: Element-wise addition (matrices/vectors must have the same shape).
    • ⊙: Element-wise multiplication (Hadamard product).
    • softmax(): A function that converts a vector of arbitrary real numbers into a probability distribution (non-negative numbers that sum to 1).
    • LayerNorm(): Layer Normalization function.
    • Concat(): Concatenation operation (joining matrices/vectors along a specified dimension).
    • Lookup(): Embedding lookup operation (retrieving the vector corresponding to a token index).
    • max(0, ...): Rectified Linear Unit (ReLU) activation function. Sets negative values to zero.
    • ∇: Gradient operator (represents the vector of partial derivatives).
    • Σ: Summation operator.
    • arg max: Returns the argument (index) that maximizes the function.

The Anatomy: Forward Pass Mathematics

This describes how input data flows through the network to produce an output, assuming the model parameters θ are fixed.

1. Input Processing: From Tokens to Vectors

Purpose: To convert the symbolic input tokens (like words) into numerical vectors that capture semantic meaning and include positional information, as the core Transformer architecture doesn't inherently process sequences sequentially.

a. Token Embeddings:

  • Concept: Each unique token in the vocabulary is associated with a dense vector of size d_model. This vector is learned during training and aims to capture the token's semantic properties.
  • Mechanism: An embedding matrix W_e of shape (V_{size}, d_model) stores these vectors. For an input sequence x = (x_1, ..., x_n), where each x_i is a token index, we perform a lookup.
  • Formula:
    \[ X_{emb} = \text{Lookup}(x, W_e) \]
  • Explanation:
    • x: The input sequence of token indices.
    • W_e: The learnable embedding matrix. W_{e,k} is the vector for the k-th token in the vocabulary.
    • Lookup: Operation that retrieves the row vector from W_e corresponding to each token index in x.
    • X_{emb}: The resulting matrix of shape (n, d_model), where the i-th row is the embedding vector for token x_i.

b. Positional Encoding (PE):

  • Concept: Since self-attention is permutation-invariant (shuffling the input would yield the same weighted sums if not for PE), we need to explicitly inject information about the position of each token in the sequence.
  • Rationale: The chosen sine and cosine functions have properties that allow the model to easily learn relative positioning. Specifically, PE_{pos+k} can be represented as a linear transformation of PE_{pos}, making it easier for the model to understand relative distances. Using alternating sine and cosine across the d_model dimensions with varying frequencies allows encoding position uniquely.
  • Mechanism: A fixed (non-learned, usually) matrix PE of shape (max_seq_len, d_model) is pre-calculated.
  • Formulas: For position pos (0 to n-1) and dimension index j (0 to d_model-1):
    \[ PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i / d_{model}}}\right) \quad \text{for } j=2i \text{ (even dimensions)} \]
    \[ PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i / d_{model}}}\right) \quad \text{for } j=2i+1 \text{ (odd dimensions)} \]
  • Explanation:
    • pos: The index (position) of the token in the sequence (0-based).
    • i: Index iterating through pairs of dimensions (i goes from 0 to d_model/2 - 1).
    • d_model: The embedding dimension.
    • 10000: A large constant chosen by the authors; changing it affects the range of frequencies used.
    • The formulas generate values between -1 and 1 based on the position and the dimension index, creating a unique positional signature for each position.

c. Input Representation:

  • Concept: Combine the semantic information (token embeddings) with the positional information.
  • Formula:
    \[ X_{input} = X_{emb} + PE_{[0:n, :]} \]
  • Explanation:
    • X_{emb}: Matrix of token embeddings (n, d_model).
    • PE_{[0:n, :]}: The first n rows of the pre-calculated Positional Encoding matrix, corresponding to the positions in the input sequence.
    • +: Element-wise addition. The positional encoding vector is added to the token embedding vector for each position.
    • X_{input}: The final input matrix (n, d_model) fed into the first encoder/decoder layer. (Note: Dropout is typically applied to this sum during training).

2. Scaled Dot-Product Attention: The Core Mechanism

Purpose: To allow each position in a sequence to attend to (draw information from) all positions (including itself) in another sequence (or the same sequence in self-attention), based on the similarity between their representations.

  • Inputs: Three matrices derived from the input sequence(s):
    • Q (Query): Matrix of query vectors (n_q, d_k). Represents the perspective of the position asking for information.
    • K (Key): Matrix of key vectors (n_k, d_k). Represents the perspective of the positions providing information, used for matching against queries.
    • V (Value): Matrix of value vectors (n_k, d_v). Represents the actual information content to be aggregated from the positions providing information.
    • Note: In self-attention (within encoder or decoder), Q, K, V are derived from the same input sequence (n_q = n_k = n or m). In encoder-decoder attention, Q comes from the decoder sequence, K and V from the encoder output (n_q = m, n_k = n). d_k and d_v are typically d_model / h.
  • Formula:
    \[ \text{Attention}(Q, K, V) = \underbrace{\text{softmax}\left(\frac{\overbrace{QK^T}^{\text{Scores}}}{\underbrace{\sqrt{d_k}}_{\text{Scale}}}\right)}_{\text{Attention Weights } A} \underbrace{V}_{\text{Values}} \]
  • Step-by-Step Explanation:
    1. Score Calculation: Scores = QK^T
      • K^T: Transpose of the Key matrix, shape (d_k, n_k).
      • Q @ K^T: Matrix multiplication of Queries (n_q, d_k) and transposed Keys (d_k, n_k). The result is the Score matrix S of shape (n_q, n_k).
      • S_{ij}: The dot product between the i-th query vector (Q_i) and the j-th key vector (K_j). This scalar value represents the raw compatibility or similarity between query i and key j. Higher dot product means higher similarity (assuming vectors point in similar directions).
    2. Scaling: Scaled Scores = Scores / sqrt(d_k)
      • d_k: Dimensionality of the key/query vectors.
      • sqrt(d_k): Square root of the key dimension.
      • Rationale: Dot products can grow large in magnitude as the dimension d_k increases. Large inputs to the softmax function can lead to extremely small gradients (vanishing gradients), making learning difficult. Scaling by sqrt(d_k) counteracts this effect, keeping the variance of the scores approximately constant regardless of d_k.
    3. Attention Weights (Softmax): A = softmax(Scaled Scores)
      • softmax(): Applied row-wise to the Scaled Scores matrix (n_q, n_k). For each row i (representing query i):
        \[ A_{ij} = \frac{\exp(S_{ij} / \sqrt{d_k})}{\sum_{l=1}^{n_k} \exp(S_{il} / \sqrt{d_k})} \]
      • Purpose: Converts the scaled similarity scores into a probability distribution for each query. A_{ij} represents the proportion of attention query i should pay to value j. All weights A_{ij} for a fixed i are non-negative and sum to 1.
    4. Output Calculation: Output = A V
      • A: Attention weights matrix (n_q, n_k).
      • V: Value matrix (n_k, d_v).
      • A @ V: Matrix multiplication. The result is the Output matrix Z of shape (n_q, d_v).
      • Interpretation: Each output vector Z_i (the i-th row of Z) is a weighted sum of all value vectors in V. The weights are given by the i-th row of the attention matrix A. Positions j with higher attention weights A_{ij} contribute more of their value vector V_j to the output Z_i. This aggregates information from the source sequence (V) based on query-key similarity.
  • Masking (Optional):
    • Purpose: In specific contexts, like decoder self-attention, we need to prevent a position from attending to subsequent positions (to maintain the autoregressive property – prediction depends only on past outputs).
    • Mechanism: Before the softmax step, a mask matrix (containing 0s for allowed positions and a large negative number like -∞ or -1e9 for forbidden positions) is added to the Scaled Scores.
    • Masked Scaled Scores = Scaled Scores + Mask
    • When softmax is applied, exp(-∞) becomes 0, ensuring that forbidden positions receive zero attention weight.

3. Multi-Head Attention (MHA): Attending in Parallel Subspaces

Purpose: Instead of performing a single attention calculation with d_model-sized keys/queries/values, MHA projects them into lower-dimensional subspaces (d_k, d_v) multiple times (h heads) in parallel. This allows the model to jointly attend to information from different representational subspaces at different positions, potentially capturing different aspects of relationships.

  • Inputs: Typically Q=K=V=X, where X ∈ R^{seq_len × d_{model}} is the input from the previous layer. For cross-attention, Q comes from the decoder, K, V from the encoder.
  • Hyperparameters:
    • h: Number of attention heads (e.g., 8).
    • d_k = d_v = d_{model} / h. The dimensions are split evenly across heads.
  • Mechanism:
    1. Linear Projections (Per Head i): Project the input X into h different query, key, and value spaces using learned weight matrices for each head.
      \[ Q_i = X W_i^Q, \quad K_i = X W_i^K, \quad V_i = X W_i^V \quad \text{for } i=1, ..., h \]
      • W_i^Q ∈ R^{d_{model} × d_k}, W_i^K ∈ R^{d_{model} × d_k}, W_i^V ∈ R^{d_{model} × d_v}: Learned weight matrices for head i. These projections capture different aspects of the input X.
      • Q_i, K_i ∈ R^{seq_len × d_k}, V_i ∈ R^{seq_len × d_v}: Query, Key, Value matrices for head i.
    2. Parallel Scaled Dot-Product Attention: Apply the attention mechanism independently for each head.
      \[ \text{head}_i = \text{Attention}(Q_i, K_i, V_i) = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i \]
      • head_i ∈ R^{seq_len × d_v}: The output of the attention mechanism for head i.
    3. Concatenation: Concatenate the outputs of all heads along the feature dimension.
      \[ \text{Concat}(\text{head}_1, ..., \text{head}_h) \in \mathbb{R}^{seq\_len \times (h \cdot d_v)} \]
      • Since h * d_v = d_model, the concatenated matrix has shape (seq_len, d_model).
    4. Final Linear Projection: Apply a final linear transformation using another learned weight matrix W^O.
      \[ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h) W^O \]
      • W^O ∈ R^{d_{model} × d_{model}}: Learned output weight matrix.
      • Purpose: Mixes the information captured by the different heads, allowing them to interact and producing a final output representation of size (seq_len, d_model).

4. Add & Norm: Stabilizing Layers

Purpose: Each sub-layer (MHA or FFN) in the encoder and decoder is wrapped by these two operations to improve training stability and gradient flow.

  • Input: z, the input to the sub-layer (e.g., MHA or FFN).
  • Sub-layer Function: Sublayer(z) (e.g., MultiHead(z, z, z)).
  • Formula:
    \[ \text{Output} = \text{LayerNorm}(z + \text{Dropout}(\text{Sublayer}(z))) \]
  • Components:
    1. Sub-layer Execution: Compute the output of the main function (Sublayer(z)).
    2. Dropout: (Applied during training only, see Physiology section). Randomly sets some activations in Sublayer(z) to zero.
    3. Residual Connection (Add): Sum = z + Dropout(Sublayer(z))
      • z: The original input to the sub-layer.
      • +: Element-wise addition.
      • Rationale: Allows the network to learn modifications (Sublayer(z)) to the identity function (z). Gradients can flow directly through the z path, bypassing the sub-layer, which helps mitigate the vanishing gradient problem in deep networks.
    4. Layer Normalization (LayerNorm): Output = LayerNorm(Sum)
      • Concept: Normalizes the activations across the feature dimension (d_model) independently for each position in the sequence. This differs from Batch Normalization, which normalizes across the batch dimension.
      • Rationale: Helps stabilize the hidden state dynamics during training, making the model less sensitive to the scale of parameters and the learning rate. It centers and scales the activations for each position.
      • Mechanism: For each vector v (a row in the Sum matrix, representing one position):
        • Calculate mean μ and variance σ² across the d_model features of v:
          \[ \mu = \frac{1}{d_{model}} \sum_{j=1}^{d_{model}} v_j \]
          \[ \sigma^2 = \frac{1}{d_{model}} \sum_{j=1}^{d_{model}} (v_j - \mu)^2 \]
        • Normalize v:
          \[ \hat{v} = \frac{v - \mu}{\sqrt{\sigma^2 + \epsilon}} \]
          (ε is a small constant, e.g., 1e-5, for numerical stability to avoid division by zero).
        • Scale and Shift: Apply learned parameters γ (scale) and β (shift), both vectors of size d_model.
          \[ \text{LayerNorm}(v) = \gamma \odot \hat{v} + \beta \]
          (⊙ denotes element-wise multiplication). γ and β allow the network to learn the optimal scale and mean for the normalized activations.

5. Position-wise Feed-Forward Networks (FFN): Adding Non-linearity

Purpose: Applied after the attention sub-layer in each encoder/decoder layer. It processes each position's representation independently and identically, adding non-linearity and further transforming the features.

  • Input: NormOutput, the output from the preceding Add & Norm layer (seq_len, d_model).
  • Mechanism: A two-layer feed-forward network applied to each position vector z (each row of NormOutput) independently.
  • Formula:

    \[ \text{FFN}(z) = \underbrace{ \underbrace{\max(0, z W_1 + b_1)}_{\text{ReLU Activation}} W_2 + b_2 }_{\text{Linear 2 + Bias}} \]

  • Explanation:
    • z ∈ R^{d_{model}}: Input vector for a single position.
    • Layer 1:
      • W_1 ∈ R^{d_{model} × d_{ff}}: Learned weight matrix for the first linear transformation. Expands the dimension from d_model to d_{ff}.
      • b_1 ∈ R^{d_{ff}}: Learned bias vector for the first layer.
      • z W_1 + b_1: First linear transformation.
    • Activation:
      • max(0, ...): ReLU (Rectified Linear Unit) activation function. Introduces non-linearity by setting negative values to zero. ReLU(x) = x if x > 0, else 0.
      • Rationale: Non-linearity is crucial for neural networks to learn complex patterns beyond simple linear relationships.
    • Layer 2:
      • W_2 ∈ R^{d_{ff} × d_{model}}: Learned weight matrix for the second linear transformation. Projects the dimension back from d_{ff} to d_model.
      • b_2 ∈ R^{d_{model}}: Learned bias vector for the second layer.
      • (...) W_2 + b_2: Second linear transformation, producing the final output vector of size d_model for that position.
    • Position-wise: The same matrices W_1, b_1, W_2, b_2 are used for every position in the sequence, but the computation is done independently for each position's vector.

6. Encoder Stack: Processing the Input Sequence

Purpose: To generate a rich contextual representation of the input sequence x. It consists of N identical layers stacked on top of each other.

  • Input: X^{(0)} = X_{input} (Input Embeddings + Positional Encoding).
  • Structure (Per Layer l = 1 to N): Each layer takes the output X^{(l-1)} from the previous layer and performs two main operations, each followed by Add & Norm:
    1. Multi-Head Self-Attention: Allows each input position to attend to all other positions in the input sequence.
      • AttnOutput = MultiHead(Q=X^{(l-1)}, K=X^{(l-1)}, V=X^{(l-1)})
      • Norm1 = LayerNorm(X^{(l-1)} + Dropout(AttnOutput))
    2. Position-wise Feed-Forward Network: Further processes each position's representation independently.
      • FFNOutput = FFN(Norm1)
      • X^{(l)} = LayerNorm(Norm1 + Dropout(FFNOutput))
  • Output: EncOutput = X^{(N)} ∈ R^{n × d_{model}}. This matrix contains the final contextualized representations of the input tokens, capturing information from the entire input sequence. It serves as the Key (K) and Value (V) input for the cross-attention mechanism in the decoder.

7. Decoder Stack: Generating the Output Sequence

Purpose: To generate the target sequence y token by token, conditioned on the encoded input EncOutput and the previously generated target tokens. It also consists of N identical layers.

  • Inputs:
    • EncOutput ∈ R^{n × d_{model}}: The final output from the encoder stack.
    • Y_{input} ∈ R^{m × d_{model}}: Target sequence embeddings + Positional Encoding. During training, this is the ground-truth target sequence, shifted right (prepended with a start token <SOS> and excluding the final token). During inference, it's the sequence of tokens generated so far.
  • Structure (Per Layer l = 1 to N): Each layer takes Y^{(l-1)} (output from previous decoder layer, with Y^{(0)} = Y_{input}) and EncOutput, and performs three main operations, each followed by Add & Norm:
    1. Masked Multi-Head Self-Attention: Allows each position in the target sequence to attend to previous positions (including itself) in the target sequence.
      • MaskedAttnOutput = MultiHead(Q=Y^{(l-1)}, K=Y^{(l-1)}, V=Y^{(l-1)}, mask=True)
      • Masking Rationale: Crucial for autoregression. When predicting token y_j, the model should only attend to y_1, ..., y_j and not future tokens y_{j+1}, .... The mask enforces this causality.
      • Norm1 = LayerNorm(Y^{(l-1)} + Dropout(MaskedAttnOutput))
    2. Multi-Head Encoder-Decoder Attention (Cross-Attention): Allows each position in the target sequence (represented by Norm1) to attend to all positions in the encoded input sequence (EncOutput). This is where information flows from the input to the output sequence.
      • CrossAttnOutput = MultiHead(Q=Norm1, K=EncOutput, V=EncOutput)
      • Norm2 = LayerNorm(Norm1 + Dropout(CrossAttnOutput))
    3. Position-wise Feed-Forward Network: Similar to the encoder, further processes each target position's representation independently.
      • FFNOutput = FFN(Norm2)
      • Y^{(l)} = LayerNorm(Norm2 + Dropout(FFNOutput))
  • Output: DecOutput = Y^{(N)} ∈ R^{m × d_{model}}. This matrix contains the final representations for the target sequence positions.

8. Final Output Layer: From Vectors to Probabilities

Purpose: To convert the final decoder output vectors into probability distributions over the entire vocabulary for each position in the target sequence.

  • Input: DecOutput ∈ R^{m × d_{model}}.
  • Mechanism:
    1. Linear Layer: A final linear transformation projects the d_model-dimensional vectors to V_{size}-dimensional vectors (logits).
      \[ \text{Logits} = \text{DecOutput} W_{final} + b_{final} \]
      • W_{final} ∈ R^{d_{model} × V_{size}}: Learned weight matrix.
      • b_{final} ∈ R^{V_{size}}: Learned bias vector.
      • Logits ∈ R^{m × V_{size}}: Each row Logits_j contains the raw scores (logits) for each word in the vocabulary being the next token at position j.
    2. Softmax Function: Converts the logits into probabilities. Applied independently to each row (position) j.
      \[ P(y_j | y_{
      • q_j ∈ R^{V_{size}}: The predicted probability distribution for the token at output position j. q_{j,k} is the predicted probability that the token at position j is the k-th token in the vocabulary.
      • Σ_{k=1}^{V_{size}} q_{j,k} = 1.

The Physiology: Training, Optimization, and Operation

This describes how the model learns its parameters (θ) and how it's used to generate outputs.

9. Training Objective (Loss Function): Guiding the Learning

Purpose: To define a measure of how well the model's predictions match the true target sequences. The goal of training is to minimize this measure.

  • Concept: We want to maximize the likelihood of observing the true target sequence y* given the input x and parameters θ. This is equivalent to minimizing the negative log-likelihood. The standard loss function for this in classification/generation tasks is Cross-Entropy Loss.
  • Standard Cross-Entropy:
    • Setup: For a given input x and target sequence y* = (y*_1, ..., y*_m), the model produces probability distributions q_j = P(y_j | y_{ for each position j. Let p_j be the true distribution for position j, which is a one-hot vector (1 at the index corresponding to the true token y*_j, and 0 elsewhere).
    • Formula (per sequence):
      \[ L_{CE}(\theta) = - \sum_{j=1}^{m} \log P(y_j=y^*_j | y_{ This sums the negative log probability of the correct token at each position. Minimizing this loss forces the model to assign higher probabilities to the correct target tokens.
    • Alternative View: Using the one-hot p_j:
      \[ L_{CE}(\theta) = - \sum_{j=1}^{m} \sum_{k=1}^{V_{size}} p_{j,k} \log(q_{j,k}) \]
      Since p_{j,k} is 1 only when k = y*_j and 0 otherwise, this simplifies to the previous formula.
  • Label Smoothing (Regularization):
    • Rationale: Using hard 0/1 targets in cross-entropy can make the model overconfident and less adaptable. Label smoothing encourages the model to be less certain by slightly distributing the probability mass from the target token to other tokens.
    • Mechanism: Replace the one-hot target p_j with a smoothed version p'_j:
      \[ p'_{j,k} = (1 - \epsilon) p_{j,k} + \epsilon u(k) = \begin{cases} 1 - \epsilon & \text{if } k = y^*_j \\ \epsilon / (V_{size} - 1) & \text{if } k \neq y^*_j \end{cases} \]
      Where ε (epsilon) is a small hyperparameter (e.g., 0.1) and u(k) = 1/(V_{size}-1) is a uniform distribution over non-target tokens (or sometimes 1/V_{size} over all tokens).
    • Loss Formula:
      \[ L_{LS}(\theta) = - \sum_{j=1}^{m} \sum_{k=1}^{V_{size}} p'_{j,k} \log(q_{j,k}) \]
    • The total loss for training is typically averaged over all tokens in a batch of sequences.

10. Optimization: Finding the Best Parameters

Purpose: To adjust the learnable parameters θ (all the Ws, bs, γs, βs) iteratively to minimize the chosen loss function L(θ).

  • Gradient Descent: The core idea is to move the parameters in the direction opposite to the gradient of the loss function.
    • Gradient Calculation (Backpropagation): The gradient ∇_θ L(θ) (the vector of partial derivatives of the loss with respect to each parameter in θ) is calculated efficiently using the backpropagation algorithm. This involves applying the chain rule of calculus backward through the network, starting from the loss and propagating gradients layer by layer through the softmax, linear layers, Add & Norm operations, attention mechanisms, etc., all the way back to the input embeddings.
    • Parameter Update (Basic): θ_{t+1} = θ_t - η ∇_θ L(θ_t), where η is the learning rate.
  • Adam Optimizer (Adaptive Moment Estimation): A popular and effective optimization algorithm that adapts the learning rate for each parameter based on estimates of the first and second moments of the gradients.
    • Mechanism: Maintains moving averages of the gradient (m_t, first moment) and the squared gradient (v_t, second moment).
      \[ m_t = \beta_1 m_{t-1} + (1 - \beta_1) \nabla_\theta L(\theta_t) \quad \text{(Momentum term)} \]
      \[ v_t = \beta_2 v_{t-1} + (1 - \beta_2) (\nabla_\theta L(\theta_t))^2 \quad \text{(RMSProp-like term)} \]
      (β_1, β_2 are hyperparameters, typically ~0.9 and ~0.999).
    • Bias Correction: Corrects for the initialization bias of the moments (which start at 0).
      \[ \hat{m}_t = m_t / (1 - \beta_1^t), \quad \hat{v}_t = v_t / (1 - \beta_2^t) \]
      (t is the iteration number).
    • Parameter Update:
      \[ \theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon_{Adam}} \]
      The update uses the bias-corrected momentum estimate ˆm_t and scales it inversely by the square root of the bias-corrected squared gradient estimate ˆv_t. This effectively gives parameters with consistently large gradients smaller updates and parameters with small or noisy gradients larger updates. ε_{Adam} (e.g., 1e-8) prevents division by zero.
  • Learning Rate Scheduling:
    • Rationale: Using a fixed learning rate η might not be optimal. Often, starting with a smaller rate, increasing it ("warm-up"), and then gradually decreasing it ("decay") leads to better convergence and final performance.
    • Transformer Schedule Example:
      \[ \eta(step) = d_{model}^{-0.5} \cdot \min(step^{-0.5}, step \cdot warmup\_steps^{-1.5}) \]
      • step: Current training step number.
      • warmup_steps: Number of initial steps for linear warm-up.
      • The rate increases linearly for warmup_steps and then decreases proportionally to 1/sqrt(step). The d_model^{-0.5} term scales the overall rate based on the model size.

11. Regularization: Preventing Overfitting

Purpose: Techniques used during training to prevent the model from fitting the training data too closely (overfitting) and improve its ability to generalize to unseen data.

  • Dropout:
    • Mechanism: During each training forward pass, randomly set a fraction p_{drop} of the outputs of a layer (specifically, applied to the output of each sub-layer before the residual addition, and to the sum of embeddings and PEs) to zero. The remaining activations are scaled up by 1 / (1 - p_{drop}) to keep the expected sum the same.
    • Formula (Conceptual): For an activation vector a, Dropout(a)_i = mask_i * a_i / (1 - p_{drop}), where mask_i is 1 with probability 1 - p_{drop} and 0 with probability p_{drop}.
    • Rationale: Prevents complex co-adaptations between neurons, forcing the network to learn more robust features that are not overly reliant on any single activation. Acts like training an ensemble of smaller networks.
    • Note: Dropout is only applied during training and turned off during inference (evaluation/prediction).
  • Label Smoothing: (Described in the Loss Function section). Discourages the model from assigning probability 1.0 to the target class.
  • Weight Decay (L2 Regularization):
    • Mechanism: Adds a penalty term to the loss function proportional to the sum of the squares of the model's weight parameters W.
      \[ L_{total}(\theta) = L_{CE/LS}(\theta) + \frac{\lambda}{2} \sum_{W \in \theta} ||W||_F^2 \]
      (||W||_F^2 is the squared Frobenius norm, sum of squared elements; λ is the regularization strength hyperparameter).
    • Rationale: Encourages the model to learn smaller weights, which often leads to simpler and better-generalizing models. Many optimizers like AdamW incorporate weight decay directly into the update rule rather than modifying the loss.

12. Inference Strategy (Decoding): Generating Output

Purpose: To use the trained Transformer model to generate an output sequence y given an input sequence x.

  • Autoregressive Generation: The decoder generates the output sequence one token at a time. The token generated at step j becomes part of the input to the decoder for generating the token at step j+1.
  • Process:
    1. Encode Input: Compute EncOutput from the input sequence x using the encoder stack. This is done only once.
    2. Initialize Decoder Input: Start with a sequence containing only the start-of-sequence token: y_{current} = (<SOS>).
    3. Iterative Generation (Loop): Repeat until an end-of-sequence token <EOS> is generated or a maximum length m_{max} is reached:
      • Prepare Decoder Input: Create the input matrix Y_{input} from the current sequence y_{current} (token embeddings + positional encodings).
      • Decoder Forward Pass: Feed Y_{input} and EncOutput through the decoder stack to get DecOutput.
      • Get Next Token Probabilities: Apply the final linear layer and softmax to the last time step's vector in DecOutput to get the probability distribution q_{next} over the vocabulary for the next token.
      • Select Next Token: Choose the next token y_{next} based on q_{next} using a decoding strategy:
        • Greedy Decoding: Simplest method. Select the token with the highest probability: y_{next} = arg max_k q_{next, k}. Fast but can lead to suboptimal sequences.
        • Beam Search: More sophisticated. Keep track of the k (beam width) most probable partial sequences (hypotheses) at each step.
          • For each of the k current hypotheses, predict the probabilities for the next token.
          • Expand each hypothesis by considering all possible next tokens, calculating the cumulative probability (usually log probability) of the extended sequences.
          • Select the top k sequences overall based on their cumulative probabilities.
          • Repeat until <EOS> is generated for hypotheses or max length is reached. Finally, select the hypothesis with the best overall score (e.g., highest log probability, possibly normalized by length). Beam search explores a wider range of possibilities than greedy decoding.
      • Append Token: Add the selected y_{next} to the current sequence: y_{current} = (y_{current}, y_{next}).
    4. Final Output: The generated sequence y_{current} (excluding <SOS> and possibly <EOS>).

Additional Considerations

13. Weight Sharing: Parameter Efficiency

  • Concept: To reduce the number of parameters and potentially improve generalization, certain weight matrices can be shared (tied).
  • Common Practice: The input embedding matrix (W_e), the output embedding matrix (used for the decoder input Y_{input}), and the transpose of the pre-softmax final linear layer matrix (W_{final}^T) are often shared.
  • Mathematical Implication: W_e = W_{final}^T (or related by a scaling factor). This links the representation used for input tokens directly to the representation used to predict output tokens.

Conclusion

The Transformer's mathematical foundation is a sophisticated interplay of linear algebra (matrix multiplications for transformations and projections), calculus (gradients for optimization via backpropagation), probability (softmax for attention weights and output distributions, cross-entropy loss), and carefully designed architectural components (self-attention, positional encoding, residual connections, layer normalization). Each mathematical operation and structural choice, from scaling attention scores by sqrt(d_k) to using residual connections and layer normalization, plays a critical role in enabling the model to process sequences effectively, learn long-range dependencies, parallelize computation, and achieve state-of-the-art performance on a wide range of tasks. Understanding these details provides insight into why the Transformer works and how its various components contribute to its overall power.

Posted by at

No comments:

Post a Comment

Subscribe to: Post Comments (Atom)