Attention Mechanism

Summary: Attention computes a weighted sum of values where weights are determined by query-key compatibility. The Transformer uses scaled dot-product attention: softmax(QK^T/√d_k)V, enabling parallelizable, constant-path-length dependency modeling.

Visual Overview

Diagram: (Mermaid diagram - view source for diagram code)

flowchart LR
    subgraph QKV["Q, K, V Projections"]
        X[Input X: n×d_model] --> Q_PROJ[W_Q: d_model×d_k]
        X --> K_PROJ[W_K: d_model×d_k]
        X --> V_PROJ[W_V: d_model×d_v]
        Q_PROJ --> Q[Q: n×d_k]
        K_PROJ --> K[K: n×d_k]
        V_PROJ --> V[V: n×d_v]
    end

subgraph ATTENTION["Scaled Dot-Product Attention"]
    Q --> QT[Q K^T]
    K --> QT
    QT --> SCALE[Scale: /√d_k]
    SCALE --> MASK{Mask?}
    MASK -- Causal --> MASK_OP[Set -∞ for j>i]
    MASK -- None --> SOFTMAX
    MASK_OP --> SOFTMAX[Softmax]
    SOFTMAX --> P[Attention Weights: n×n]
    P --> PV[P V]
    V --> PV
    PV --> OUT[Output: n×d_v]
end

classDef proj fill:#e3f2fd,stroke:#90caf9;
classDef attn fill:#e8f5e9,stroke:#c8e6c9;
classDef data fill:#fff3e0,stroke:#ffcc02;
class Q_PROJ,K_PROJ,V_PROJ,Q,K,V proj;
class QT,SCALE,MASK,MASK_OP,SOFTMAX,P,PV attn;
class X,OUT data;

Core Idea

In plain English: Attention lets each position in a sequence "look at" all other positions and decide how much to incorporate their information. The query asks "what do I need?", keys represent "what I offer", and values are "what I actually contribute".

Instead of sequential processing (RNN) or fixed receptive fields (CNN), attention computes all pairwise relationships simultaneously:

Attention(Q, K, V) = softmax((QK^T)/(sqrt(d_k)))V

Scaled Dot-Product Attention

Why Scaling?

For large d_k, dot products grow large in magnitude, pushing softmax into regions with extremely small gradients. Scaling by 1/sqrt(d_k) counteracts this.

Computational Complexity

• Time: O(n^2 d) for sequence length n, dimension d
• Memory: O(n^2) for attention matrix
• Bottleneck for long sequences → motivates FlashAttention, sparse attention

Attention Variants in Transformer

1. Encoder Self-Attention

• Q, K, V all from encoder layer output
• Each position attends to all positions in the same layer
• Bidirectional context

2. Decoder Self-Attention (Causal)

• Q, K, V from decoder layer output
• Masking: Future positions masked (set to -∞ before softmax)
• Ensures auto-regressive property: position i only attends to ≤ i

3. Encoder-Decoder Attention (Cross-Attention)

• Q from decoder layer, K,V from encoder output
• Allows decoder positions to attend over full input sequence
• Similar to traditional seq2seq attention

Mathematical Properties

Property	Self-Attention	Recurrent	Convolutional
Max path length	O(1)	O(n)	O(log_k n)
Parallelizable	Yes (O(1) seq ops)	No (O(n) seq ops)	Yes
Complexity/layer	O(n^2 d)	O(n d^2)	O(k n d^2)
Learn long-range deps	Easy	Hard	Medium

Extensions & Variants

Variant	Key Idea	Reference
Multi-Head	Parallel attention heads with different projections	multi-head-attention
FlashAttention	IO-aware tiling, recomputation for exact attention	flash-attention
Sparse Attention	Fixed/local/strided patterns (Longformer, BigBird)	—
DSA	Learned sparse pattern via indexer+selector	dsa-attention
Linear Attention	Kernel approximation O(n d) complexity	Katharopoulos et al. (2020)
Retention	Recurrent formulation of attention	Sun et al. (2023)

Attention Visualization Insights

From the original paper (Appendix):

• Heads learn distinct linguistic roles (syntax, anaphora resolution, etc.)
• Head 5 in layer 5: attends to verb-object dependencies
• Head 6 in layer 5: resolves coreference (e.g., "its" → antecedent)
• Some heads attend mostly to delimiter tokens ([CLS], [SEP])

Limitations

• Quadratic complexity in sequence length → limits context window
• Memory intensive O(n^2) attention matrix
• No inductive bias for locality (unlike CNNs)
• Position information only via added encodings

Related Concepts

• transformer-architecture — Full architecture using attention
• multi-head-attention — Parallel attention heads
• flash-attention — Memory-efficient exact attention
• dsa-attention — Dynamic sparse attention

Sources

• arxiv-1706.03762: Section 3.2 — Attention formulation and variants