How Transformer Architecture Works: A Deep Dive

Scenario context

The Transformer is the architecture behind every major LLM (GPT, BERT, Claude, Gemini). Its core innovation is Self-Attention — a mechanism that lets the model weigh relationships between all tokens in a sequence simultaneously, regardless of distance. This enables parallelism that RNNs could not achieve.

Tokenization splits text into subword units using BPE or WordPiece. For example, `"unhappines…

This walkthrough focuses on Self-Attention: How Every Token Reads the Room.

Before making a design choice, clarify:

• For each token, the model creates three vectors via separate learned linear projections: - Q (Query): "What information am I looking for?" - K (Key): "What content do I advertise?" - V (Value): "What do I contribute if…
• Softmax then converts raw scores into a probability distribution over all tokens.
• The final output is a weighted sum of Value vectors: $$\text{Attention}(Q, K, V) = \text{softmax}!\left(\frac{QK^T}{\sqrt{d k}}\right) V$$ Concrete coreference example: In "The animal didn't cross the street because it…

Concepts in play: deep-learning-foundations, vector-embeddings-and-similarity, core-algorithm-patterns, database-indexing

Stage 1 — Deep Dive: Under the Hood: Multi-Head Attention, Complexity, and the Math

• Internals Running attention once captures only a single type of token relationship.
• Multi-Head Attention runs $h$ independent attention operations in parallel, each with its own $W^Q i$, $W^K i$, $W^V i$ projection matrices, then concatenates their outputs and projects back to model dimension: $$\text{…
• The projection matrix $W^O$ merges all head outputs back into a single model-width vector before passing downstream.
• 📊 Attention Mechanism: Q × K → Weighted V Reading left-to-right, this diagram traces the exact sequence of operations defined by $\text{softmax}(QK^T/\sqrt{d k})V$: the Query and Key vectors are dot-producted to produc…

flowchart LR
    Q["Query (Q) What am I looking for?"] --> S[Q  K dot product scores]
    K["Key (K) What do I advertise?"] --> S
    S --> SC[" sqrt(d_k) stabilize gradients"]
    SC --> SM[Softmax attention weights]
    SM --> WV["Weights  Value (V)"]
    V["Value (V) What I contribute"] --> WV
    WV --> OUT[Attention Output context-aware token repr]

sequenceDiagram
    participant I as Input Tokens
    participant S1 as Head 1 (syntax)
    participant S2 as Head 2 (coreref)
    participant SH as Head h (semantic)
    participant C as Concat
    participant L as Linear Projection W_O

    I->>S1: split + project Q1,K1,V1
    I->>S2: split + project Q2,K2,V2
    I->>SH: split + project Qh,Kh,Vh
    S1-->>C: head_1 output
    S2-->>C: head_2 output
    SH-->>C: head_h output
    C->>L: concat all heads
    L-->>I: model-dim output (residual add)

Observable outcome: 📊 Attention Mechanism: Q × K → Weighted V Reading left-to-right, this diagram traces the exact sequence of operations defined by $\text{softmax}(QK^T/\sqrt{d k})V$: the Query and Key vectors are dot-producted to produc…

Stage 2 — Encoder-Only, Decoder-Only, Encoder-Decoder: Picking the Right Variant

• The original "Attention Is All You Need" (Vaswani et al., 2017) used a full Encoder-Decoder architecture for machine translation.
• Since then, three specialized families have emerged with distinct strengths.
• Encoder-only (BERT, RoBERTa): Processes input bidirectionally — every token attends to every other token.
• Best for classification, NER, semantic search, and embeddings.

Observable outcome: Best for classification, NER, semantic search, and embeddings.

Stage 3 — From Raw Text to Next-Token Prediction: The Full Pipeline

• Following the diagram top-to-bottom reveals precisely where each architectural decision lives in the forward pass: the tokenizer and embedding layer are a one-time setup cost; the real computation is the Transformer Blo…
• The final projection from hidden dimension to vocabulary size (50,257 tokens for GPT-2; 32,000 for LLaMA-3) produces raw logits, and the softmax converts those into a probability distribution from which the next token i…
• Figure: End-to-end token-to-prediction pipeline through a Transformer.
• The Transformer Block (steps E–I) is stacked N times: 12 for BERT-base, 96 for GPT-3.

flowchart TD
    A[Input Text e.g. 'The cat sat'] --> B[Tokenizer BPE / WordPiece]
    B --> C[Token IDs e.g. [464, 3797, 3332]]
    C --> D[Token Embeddings + Positional Encoding]
    D --> E[Transformer Block x N layers]
    E --> F["Self-Attention softmax(QK / sqrt_dk) times V"]
    F --> G["Add and Norm LayerNorm(x + Attention output)"]
    G --> H["Feed-Forward Network GELU(xW1)W2"]
    H --> I["Add and Norm LayerNorm(h + FFN output)"]
    I --> J[Output Logits vocab_size dimensions]
    J --> K[Softmax - Probability Distribution]
    K --> L[Next Token Prediction]

Observable outcome: The Transformer Block (steps E–I) is stacked N times: 12 for BERT-base, 96 for GPT-3.

Stage 4 — Real-World Applications: Where Transformers Run the World: GPT-4 to AlphaFold2

• Transformers have escaped the NLP lab and now power applications across every major domain: - GPT-4 (OpenAI): Decoder-only model used for chat, code generation, and multi-step reasoning.
• Its 128K context window handles entire codebases in a single prompt.
• • BERT (Google): Encoder-only model powering Google Search's passage ranking and query understanding since 2019.
• Fine-tuned variants run NER, sentiment analysis, and extractive QA at web scale.

Observable outcome: Fine-tuned variants run NER, sentiment analysis, and extractive QA at web scale.

Stage 5 — Implementing Self-Attention and Inspecting It in a Real Model

• Understanding the architecture means being able to both implement the core formula and observe it running in a real model.
• The two patterns below do exactly that — connecting the math directly to running code.
• Part 1 — Implement scaled dot-product attention from scratch: Part 2 — Inspect real BERT attention weights to verify pronoun resolution: The post explains that in "The animal didn't cross the street because it was too t…
• Here is how to check that empirically: The same pipeline API hides the underlying masking difference: BERT runs with causal mask=False (full bidirectional attention); GPT-2 runs with causal mask=True (left-to-right only…

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(
    Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
    causal_mask: bool = False
) -> tuple:
    """
    Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V

    Q, K, V shape: (batch, heads, seq_len, d_k)
    causal_mask=False → encoder (BERT):  each token attends to the full sequence
    causal_mask=True  → decoder (GPT):   each token only attends to past positions
    """
    d_k    = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)  # QKᵀ / √d_k

    if causal_mask:   # block future positions — required for autoregressive generation
        seq_len = scores.size(-1)
        mask    = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
        scores  = scores.masked_fill(mask, float("-inf"))

    weights = F.softmax(scores, dim=-1)   # probability distribution over attended positions
    output  = torch.matmul(weights, V)    # weighted blend of every value vector
    return output, weights

# Toy 3-token sequence: one head, d_k=8
Q = torch.randn(1, 1, 3, 8)
K = torch.randn(1, 1, 3, 8)
V = torch.randn(1, 1, 3, 8)

_, enc_weights = scaled_dot_product_attention(Q, K, V, causal_mask=False)
_, dec_weights = scaled_dot_product_attention(Q, K, V, causal_mask=True)

print("Encoder (BERT-style) — each token attends to all positions:")
print(enc_weights[0, 0].detach().round(decimals=2))
# row i = probability mass token i places on each position j
# e.g. [[0.3, 0.4, 0.3], [0.3, 0.4, 0.3], [0.3, 0.4, 0.3]]

print("\nDecoder (GPT-style) — upper triangle masked to -inf before softmax:")
print(dec_weights[0, 0].detach().round(decimals=2))
# token 0 can only see itself, token 1 sees [0,1], token 2 sees all
# e.g. [[1.0, 0.0, 0.0], [0.4, 0.6, 0.0], [0.2, 0.3, 0.5]]

from transformers import AutoTokenizer, AutoModel, pipeline

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model     = AutoModel.from_pretrained("bert-base-uncased", output_attentions=True)

sentence  = "The animal didn't cross the street because it was too tired"
inputs    = tokenizer(sentence, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

# outputs.attentions: 12 tensors, one per layer, each (1, 12_heads, seq_len, seq_len)
tokens      = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
it_idx      = tokens.index("it")

# Average attention FROM "it" across all 12 heads in the deepest layer
attn_from_it = outputs.attentions[-1][0, :, it_idx, :].mean(dim=0)  # (seq_len,)
top2         = attn_from_it.topk(2)

print("\nTop 2 tokens 'it' attends to (layer 11, all heads averaged):")
for score, idx in zip(top2.values, top2.indices):
    print(f"  '{tokens[idx.item()]}' → weight {score.item():.3f}")
# Expected: 'animal' ranks above 'street' — attention weight resolves the coreference

# Confirm the encoder vs. decoder distinction with the high-level API
classifier = pipeline("sentiment-analysis",
                       model="distilbert-base-uncased-finetuned-sst-2-english")
generator  = pipeline("text-generation", model="gpt2")

print(classifier("The attention mechanism is remarkably elegant."))
# [{'label': 'POSITIVE', 'score': 0.9998}] — BERT encoder, causal_mask=False above

print(generator("The Transformer architecture changed",
                 max_new_tokens=20, do_sample=True, temperature=0.8)[0]["generated_text"])
# GPT-2 decoder, causal_mask=True — each token only sees prior context

Observable outcome: Here is how to check that empirically: The same pipeline API hides the underlying masking difference: BERT runs with causal mask=False (full bidirectional attention); GPT-2 runs with causal mask=True (left-to-right only…

Stage 6 — HuggingFace Transformers and FlashAttention: Running Scaled Dot-Product Attention in Practice

• HuggingFace Transformers is the de-facto open-source library for loading, fine-tuning, and serving pretrained Transformer models — from BERT and GPT-2 to LLaMA and Mistral.
• It handles weight loading, tokenization, and device placement so you can go from a model name to a running forward pass in minutes.
• How it solves the problem in this post: Rather than implementing $\text{softmax}(QK^T/\sqrt{d k})V$ from scratch, transformers gives you pretrained weights for every layer — including multi-head attention projections an…
• The snippet below uses torch.nn.functional.scaled dot product attention , which automatically dispatches to FlashAttention (Dao et al., 2022) when CUDA is available, replacing the naive $O(n^2)$ kernel with an I/O-aware…

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel

# --- 1. Load a pretrained encoder and its tokenizer ---
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, output_attentions=True)
model.eval()

# --- 2. Tokenise a sentence pair and get hidden states ---
text = "The animal didn't cross the street because it was too tired."
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

# last_hidden_state: (batch=1, seq_len, hidden_dim=768)
hidden = outputs.last_hidden_state
print("Hidden state shape:", hidden.shape)

# --- 3. Manual scaled dot-product attention (single head) using Flash-ready API ---
# Project to Q / K / V with small demo dimensions
seq_len, d_model = hidden.shape[1], hidden.shape[2]
d_k = 64  # one attention head

Wq = torch.randn(d_model, d_k)
Wk = torch.randn(d_model, d_k)
Wv = torch.randn(d_model, d_k)

Q = hidden @ Wq          # (1, seq_len, d_k)
K = hidden @ Wk
V = hidden @ Wv

# scaled_dot_product_attention uses FlashAttention kernel on CUDA automatically
attn_output = F.scaled_dot_product_attention(Q, K, V, scale=d_k**-0.5)
print("Attention output shape:", attn_output.shape)  # (1, seq_len, 64)
# → tensor of shape [1, 25, 64]: each token now holds context from all other tokens

Observable outcome: The snippet below uses torch.nn.functional.scaled dot product attention , which automatically dispatches to FlashAttention (Dao et al., 2022) when CUDA is available, replacing the naive $O(n^2)$ kernel with an I/O-aware…

End-to-end outcome

For each token, the model creates three vectors via separate learned linear projections: - Q (Query): "What information am I looking for?" - K (Key): "What content do I advertise?" - V (Value): "What do I contribute if selected?" The attention score measuring how relevant token $j$ is to token $i$ is: $$\text{score}(i, j) = \frac{Q i \cdot K j^T}{\sqrt{d k}}$$ Dividing by $\sqrt{d k}$ (e.g., $\sqrt{64} = 8$ for 64-dim keys) prevents dot products from growing so large that softmax gradients saturate and vanish.

Alternatives considered

Decision Guide: BERT, GPT, or T5: Choosing the Right Transformer for Your Task

Decision Guide: BERT, GPT, or T5: Choosing the Right Transformer for Your Task

Failure and degraded modes

Trade-offs & Failure Modes: The Hidden Costs: Memory, Stability, and Context Length Trade-offs

• Detection: The most underestimated trade-off is depth vs.
• User or operator impact: The most underestimated trade-off is depth vs.