From Self-Attention to GQA

What Is the Attention Mechanism

Attention is one of the most important innovations in modern deep learning. It allows a model to dynamically focus on information at different positions when processing sequence data.

In the Transformer architecture, the core idea of attention is: let each position's token perform a weighted sum over all other tokens according to its relevance to them.

The Standard Self-Attention Computation

Given an input sequence, Self-Attention computes its output through the following steps:

1. Linear projection: project the input $X$ into Query ($Q$), Key ($K$), and Value ($V$)

   $$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

2. Compute attention scores: compute similarity through the dot product of Query and Key

   $$S = \frac{QK^T}{\sqrt{D}}$$

   where $D$ is the head dimension, and dividing by $\sqrt{D}$ is for numerical stability (scaled dot-product)

3. Apply Softmax: convert scores into a probability distribution

   $$A = \text{softmax}(S)$$

4. Weighted sum: weight the Value with the attention weights

   $$O = AV$$

Intuition:

• $Q$ (Query): what information do I want?
• $K$ (Key): what information do I provide?
• $V$ (Value): what information do I actually contain?
• Attention score: measures how well "what I need" matches "what you provide"

PyTorch Reference Implementation

import torch
import torch.nn.functional as F

def self_attention(X, W_q, W_k, W_v):
    """
    X: (batch, seq_len, d_model)
    W_q, W_k, W_v: (d_model, d_head)
    """
    Q = X @ W_q  # (batch, seq_len, d_head)
    K = X @ W_k
    V = X @ W_v

    # Compute attention scores
    scores = Q @ K.transpose(-2, -1)  # (batch, seq_len, seq_len)
    scores = scores / (K.shape[-1] ** 0.5)  # scale

    # Softmax normalization
    attn_weights = F.softmax(scores, dim=-1)

    # Weighted sum
    output = attn_weights @ V  # (batch, seq_len, d_head)

    return output, attn_weights