Transformer Architecture

Core Idea: Attention Over Sequences

Transformers replace recurrent processing with self-attention, which lets every token attend to every other token in parallel. Each position builds a representation by weighted aggregation of all positions, where weights come from learned similarity scores. This removes sequential bottlenecks and scales well on GPUs.

Scaled dot-product attention

Queries (Q), keys (K), and values (V) are linear projections of hidden states. Attention scores are computed as QK^T, scaled by the square root of head dimension, then softmaxed into weights. The output is the weighted sum of V vectors. Multiple heads run in parallel so the model can capture different relationship types (syntax, coreference, long-range dependencies).

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

Refer: The Illustrated Transformer · Attention Is All You Need

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, V)

Transformer Block Anatomy

A standard decoder block (used in GPT-style LLMs) contains:

• Pre-norm self-attention with residual connection
• Pre-norm feed-forward network (FFN) with residual connection

The FFN is typically two linear layers with a non-linearity or gated activation in between. It acts as a per-token memory and transformation module, while attention handles cross-token mixing. Encoder-decoder models (original Transformer, T5) add cross-attention between encoder and decoder; most LLMs today are decoder-only.

Positional Information

Attention is permutation-invariant without position signals. Models inject order through:

• Absolute positional embeddings added to token embeddings (early GPT, BERT)
• RoPE (Rotary Position Embedding) rotates Q and K vectors by position-dependent angles, encoding relative distance in attention scores. Used in LLaMA, Mistral, and most modern LLMs.
• ALiBi adds a linear bias to attention scores based on token distance, extrapolating to longer contexts without retraining position tables.

RoPE has become the default because it composes well with long-context extensions and keeps relative position information inside attention mechanics.

Efficiency Variants

Grouped-query and multi-query attention

Standard multi-head attention uses separate K and V projections per head. Multi-query attention (MQA) shares one K/V head across all query heads, cutting KV memory during inference. Grouped-query attention (GQA) shares K/V across groups of heads, balancing quality and speed. GQA is common in production LLMs (LLaMA 2/3, Mistral).

Mixture of Experts (MoE)

Instead of one dense FFN per layer, MoE routes each token to a small subset of expert FFNs via a learned gate. Total parameters grow large, but only a fraction activate per token, improving compute efficiency at scale. Challenges include load balancing across experts and more complex training infrastructure.

Causal Masking and Autoregressive Decoding

Decoder-only LLMs apply a causal mask so position i can only attend to positions j ≤ i. This enforces left-to-right generation at training time. At inference, the model appends one token at a time, reusing cached key-value tensors from prior steps (KV cache) to avoid recomputing the full sequence each step.

Scaling and Design Tradeoffs

Transformer quality improves predictably with more parameters, data, and compute (scaling laws). Engineering choices trade off memory, latency, and quality:

• More layers vs wider hidden dimensions
• Dense FFN vs MoE routing
• Full MHA vs GQA for inference memory
• Context length vs quadratic attention cost in naive implementations

FlashAttention and other IO-aware kernels reduce memory bandwidth bottlenecks, making long contexts and large batches practical without changing the mathematical definition of attention.