Transformer Deep Dive

In 2017, Google published the paper "Attention Is All You Need", introducing the Transformer architecture. No one expected that this paper would completely change the AI field.

Today's GPT, Claude, Gemini, Llama... almost all mainstream large language models are essentially variants of the Transformer.

Understand the Transformer, and you understand 90% of modern LLMs.

In this article, we will dive into every component of the Transformer: self-attention, multi-head attention, positional encoding, feed-forward networks, layer normalization... covering not only "what" but also "why".

This is one of the most technically in-depth modules. We will use code to demonstrate the core computations, ensuring that you not only understand the formulas but can also implement them yourself.

Transformer Prequel: Limitations of RNN

Before the Transformer, sequence tasks (such as translation and text generation) were mainly handled by RNNs (Recurrent Neural Networks) and their variants LSTM and GRU.

How RNN Works

The core idea of RNN is "process one by one": the input sequence enters the network word by word, and the output at each step contains information from all previous words.

RNN seems reasonable, but it has three fatal problems:

Problem 1: Vanishing Gradients, Difficulty Remembering Long-Range Dependencies

RNN is a "chain" structure; the gradient at each step must backpropagate to the first step. After repeated multiplications, the gradient decays exponentially, becoming almost 0.

For example, in the sentence "I went to Paris in 2010, ... that was my favorite city" — "city" needs to refer to "Paris", but there are too many words in between, so RNNs struggle to learn such long-range dependencies.

LSTM and GRU alleviate this problem, but do not completely solve it.

Problem 2: Cannot Compute in Parallel, Slow Training

An RNN must wait for step t-1 to finish before it can compute step t. This means:

• Even with 100 GPUs, you can only compute one word at a time
• The longer the sequence, the longer the training time
• Difficult to scale to very large datasets

Problem 3: Earlier Information Is Easily "Overwritten"

The hidden state of an RNN is updated step by step, and later information keeps overwriting earlier information. Important information at the beginning of a sentence may be largely "diluted" by the end.

The emergence of Transformer solved all three problems at once.

Self-Attention Mechanism

Self-attention is the core of Transformer. Its idea is simple:Each word must "communicate" with all words in the sentence, see who is important to itself, and then do a weighted sum by importance.

First look at the overall architecture diagram:

Intuitive Understanding of Q, K, V

Self-attention uses three vectors to describe each word:

• Q (Query)——"Who am I looking for?"—— which word this word wants to attend to
• K (Key)——"Who am I?"—— this word's identity
• V (Value)——"What do I have?"—— the actual content of this word

The computation process is:

1. Each word uses its Q to "match" with the K of all words, obtaining attention scores
2. Use Softmax to normalize the scores so they sum to 1
3. Use the normalized scores to weight and sum the V of all words

The diagram makes it clearer:

Attention Score Calculation Formula

The complete formula is:

Attention(Q, K, V) = softmax( Q·Kᵀ / √dₖ ) · V

Let's implement it step by step in pure Python:

Run this code and you will see the complete computation process: Q/K/V generation, attention score calculation, scaling, Softmax, and finally the weighted output.

Why Divide by √dₖ?

This is an important detail. Assuming each element of Q and K is a random variable with mean 0 and variance 1, the variance of their dot product is dₖ.

If dₖ is large (e.g., 512), the dot product values become very large, causing Softmax to enter the "saturation region"—the gradient is nearly 0, and training cannot proceed.

After dividing by √dₖ, the variance returns to 1, and the Softmax gradient becomes healthy again.

Scaled dot-product attention is one of the key techniques for stable Transformer training.

Multi-Head Attention

Single-head self-attention is powerful, but it has one problem: each word can only "attend" once. What if we want a word to attend to multiple different aspects at the same time?

For example, in the sentence "In the zoo, the tiger is chasing the rabbit":

• "Tiger" might need one head to attend to "zoo" (location)
• Another head to attend to "rabbit" (object of action)
• The third head focuses on "chasing" (action).

Multi-head attention splits Q/K/V into multiple "heads", each learning a different attention pattern, and finally concatenates them.

Steps to Compute Multi-Head Attention

1. Split Q/K/V into h parts along the last dimension (h is the number of heads).
2. Each head independently performs self-attention
3. Concatenate the outputs of all heads
4. Pass through a linear projection layer to get the final output

The formula is:

MultiHead(Q, K, V) = Concat(head₁, head₂, ..., headₕ) · Wₒ
  其中 headᵢ = Attention(Q·Wqᵢ, K·Wkᵢ, V·Wvᵢ)

The power of multi-head attention lies in:Different heads automatically learn different types of attention patterns.

For example, in translation tasks:

• Some heads focus on syntactic relationships (subject-verb-object)
• Some heads focus on coreference (which noun a pronoun refers to)
• Some heads focus on local neighboring words

These patterns are not manually designed; they are automatically learned by the model during training.

Positional Encoding

Self-attention has a problem: it is a "bag of words" model—it does not care about word order. "I love you" and "You love me" may look the same to self-attention.

But order is crucial to language. We need to "inject" positional information into the model.

Sinusoidal Positional Encoding

The original Transformer uses sinusoidal positional encoding:

PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

Where:

• pos is the position (0, 1, 2, ...)
• i is the dimension index
• d_model is the model dimension

This design is ingenious:

• Each position has a unique encoding
• Relative positions can be represented by linear combinations (sin(a+b) = sin a cos b + cos a sin b)
• It can extrapolate to longer sequences than those seen during training

RoPE: Rotary Positional Encoding

Although sinusoidal positional encoding is classic, better solutions emerged later. The current mainstream large models (such as LLaMA) useRoPE（Rotary Positional Embedding）。

The idea of RoPE is:Rotate position information into Query and Key.

Specifically, rotate each pair of dimensions of Q and K:

[ q₁, q₂, q₃, q₄, ... ]
  ↓ 旋转
[ q₁·cos mθ - q₂·sin mθ,
  q₁·sin mθ + q₂·cos mθ,
  q₃·cos mθ' - q₄·sin mθ',
  q₃·sin mθ' + q₄·cos mθ',
  ... ]

Advantages of RoPE:

• Naturally contains relative position information
• Good extrapolation (works for lengths beyond training)
• Usually performs better than sinusoidal positional encoding

ALiBi: Decaying Attention by Distance

Another simple and effective approach isALiBi（Attention with Linear Biases）。

It does not add positional encoding, but directly adds a bias to the attention scores: the farther apart, the lower the score.

The advantage of ALiBi is its simplicity, and it is especially friendly to length extrapolation.

Feed-Forward Network (FFN) Layer

In each Transformer layer, after attention, there is also a feed-forward network.

Structure of FFN

The FFN of the original Transformer is:

FFN(x) = max(0, x·W₁ + b₁) · W₂ + b₂

That is: linear layer → ReLU activation → linear layer.

Usually the intermediate dimension is larger than the input dimension (e.g., d_model=512, intermediate dimension=2048).

ReLU vs SwiGLU

The original used ReLU, but now the mainstream isSwiGLU：

SwiGLU(x) = (x·W₁ ⊗ sigmoid(x·V₁)) · W₂

where ⊗ is element-wise multiplication.

SwiGLU performs better than ReLU on many tasks; LLaMA, PaLM, etc., all use it.

Why is FFN important?

• Attention is "aggregating information", FFN is "processing information"
• FFN has a large number of parameters, and is the main place where the model "memorizes" knowledge
• Without FFN, the expressive capability of Transformer would be greatly reduced

Layer Normalization (Layer Norm)

There is another important component in Transformer: Layer Normalization.

What Layer Normalization Does

Layer normalization normalizes the features of each sample to mean 0 and variance 1:

LN(x) = γ ⊗ ((x - μ) / √σ²) + β

where:

• μ is the mean of x, σ² is the variance of x
• γ (scale) and β (shift) are learnable parameters
• ⊗ is element-wise multiplication

Pre-Norm vs Post-Norm

Layer normalization has two placement methods:

• Post-Norm: sublayer → residual → normalization (original Transformer)
• Pre-Norm: normalization → sublayer → residual (now mainstream)

Pre-Norm training is more stable, and almost all large models use it now.

RMS Norm

RMS NormIt is a simplified version of layer normalization, only calculating root mean square, without subtracting the mean:

RMS(x) = sqrt(mean(x²))
RMSNorm(x) = γ ⊗ (x / RMS(x))

RMS Norm is faster to compute, with similar results; models like LLaMA use it.

Encoder-Decoder vs Decoder-Only Architecture

The original Transformer is an Encoder-Decoder architecture, but most large language models today are Decoder-Only.

Encoder-Decoder: Used for Tasks Like Translation

Encoder-Decoder architecture:

• Encoder: reads the input sequence (e.g., source language sentence) and generates a context representation
• Decoder: generates the target sequence (e.g., translation result) based on the Encoder's output

There are two types of attention in the Decoder:

1. Masked self-attention: can only look at previous words (cannot see the future)
2. Cross-attention: to attend to the Encoder's output

Suitable tasks: machine translation, speech recognition, summarization (requires "rewriting" input into output).

Decoder-Only: GPT Series, Used for Generation

The Decoder-Only architecture has only a Decoder, no Encoder. Its characteristics:

• Only masked self-attention (causal attention)
• Autoregressive generation: outputs one word at a time
• Suitable for the pretraining + fine-tuning paradigm

Current GPT, LLaMA, Claude, Gemini... are almost all Decoder-Only.

Encoder-Only: BERT Series, Used for Understanding

The Encoder-Only architecture has only an Encoder. Its characteristics:

• Bidirectional attention: each word can see all words
• Suitable for understanding tasks: classification, named entity recognition, extractive question answering
• Pretrained with "Masked Language Model" (MLM)

Why did Decoder-Only become mainstream? Because it is simple, and under large-scale pretraining, its generation ability is extremely strong. "One size fits all" actually works better than "specialized design".

Flash Attention: Making Attention Computation Fast and Memory-Efficient

The time and memory complexity of standard self-attention are both O(n²); when the sequence gets long (e.g., n=8192), it is slow and GPU memory is insufficient.

Flash AttentionIt is a key optimization; its emergence in 2022 made long-sequence training feasible.

Memory Bottleneck of Standard Attention

Standard self-attention needs to save intermediate results (attention weight matrices) for backpropagation. For n=8192, this is 8192×8192 ≈ 67 million elements, which alone takes up tens of GB of GPU memory.

Optimization Principles of Flash Attention

Core idea of Flash Attention:

1. Block-wise computation: Split Q/K/V into small blocks and compute block by block
2. Online Softmax: Don't save the full attention matrix; use mathematical tricks to recompute it during backpropagation
3. Leverage SRAM: Put the blocks into the GPU's high-speed cache (SRAM) to reduce HBM access

The result:

• Memory complexity drops from O(n²) to O(n)
• 2-4x faster
• Identical precision (mathematically equivalent)

Almost all large model training now uses Flash Attention or its variants (Flash Attention v2, Memory-Efficient Attention, etc.).

KV Cache: Principles of Inference Acceleration

Once training optimization is done, inference (generation) also needs optimization.KV CacheIt is the key to inference acceleration.

Problem: Repetitive Computation in Autoregressive Generation

During autoregressive generation, every time a new token is generated, the K and V of all previous tokens must be recomputed — this is very wasteful.

For example, when generating the 100th token, the K/V of the first 99 tokens has already been computed and does not need to be recomputed.

How KV Cache Works

The idea of KV Cache is simple:

• Cache the K and V of all previous positions
• When generating a new token, only compute the new token's K/V and append it to the cache
• Use the complete K/V for attention

This way, computation per step drops from O(n²) to O(n), greatly improving generation speed.

KV Cache has a small cost: it requires extra GPU memory to store K and V. For a 7B model with a 4k context, KV Cache takes about 1-2GB of GPU memory — this is worth it.

There are also more advanced optimizations:

• Multi-Query Attention: Multiple heads share K/V
• Grouped Query Attention: Grouped sharing of K/V

These further reduce memory usage and improve inference speed.

Summary: Complete Transformer Pipeline

Let's put all the components together and summarize the complete Transformer flow.

One Decoder-Only Transformer Layer

x → LayerNorm → Masked Multi-Head Attention → 残差连接 →
    → LayerNorm → FFN → 残差连接 → 输出

Complete Decoder-Only Model

1. Input token IDs → word embedding
2. Add positional encoding
3. Pass through N Transformer layers
4. Final LayerNorm
5. Linear layer mapping to vocabulary size
6. Softmax to get the probability of the next token

Key Hyperparameters

Code Implementation Guide

If you want to implement a simplified Transformer from scratch, the steps are:

1. Implement layer normalization
2. Implement self-attention (with masking)
3. Implement multi-head attention
4. Implement FFN
5. Implement Transformer layer
6. Implement positional encoding
7. Assemble the complete model
8. Implement the training loop

In practice, we use frameworks like PyTorch rather than pure Python. The pure Python implementation here is for demonstrating the principles and helping you understand what each step does.

The best way to understand Transformer is to implement it yourself. Even a simplified version will give you a fresh perspective on many details.