Thanh Linh Nguyen

PhD candidate at Trinity College Dublin, The University of Dublin

About Research Publications Projects Awards Teaching Blog Misc

How Transformer LLMs Work | Thanh Linh Nguyen

How Transformer LLMs Work

I have wanted to understand Transformers and related advances in a structured, in-depth way for a long time, but I struggled to find the right resources. I therefore decided to study the topic systematically. These notes synthesize the How Transformer LLMs Work and Attention in Transformers: Concepts and Code in PyTorch courses by DeepLearning.AI with material from other courses and readings I did. I hope they also help others build a clear understanding of Transformers, and I will continue updating them over time.

NOTE: At some points, you and (I) will find it difficult, but believe me, you will find the answer by reading further (I hide it from you and future me in later parts @@).

Table of Contents

1. Language Models: Evolution and Fundamentals

1.1. Bag-of-Words

Limitation: Bag-of-Words does not directly represent word order or semantic meaning.

1.2. Static Word Embeddings

Neural methods such as Word2Vec learn dense vector embeddings that capture semantic relationships between words. We can compare these vectors to estimate the similarity between words.

Limitation: Word2Vec creates static embeddings. For example, it produces the same embedding for bank in river bank and financial bank, regardless of the surrounding context.

1.3. RNNs and Attention

A single fixed-size context vector may fail to preserve all relevant information from a long or complex input sequence.

2. Transformer Architecture

A Transformer consists of a stack of (N) structurally identical layers with separate learned parameters. Each layer uses attention to incorporate/attend information from permitted token positions and a position-wise feed-forward network to transform each token representation independently. The original encoder–decoder Transformer has separate encoder and decoder stacks, with decoder layers additionally using cross-attention over encoder outputs.

Transformer encoder-decoder architecture

2.1. Transformer Encoder

The Transformer encoder converts input token embeddings into contextualized representations using positional information and a stack of encoder layers. Its main components are:

2.2. Transformer Decoder

In the original encoder-decoder Transformer, the decoder generates the output sequence using previously available output tokens and the encoder’s contextualized input representations. Each decoder layer contains:

Later, researchers have found out that using encoder-only and decoder-only can work well for different applications.

Source: https://www.linkedin.com/posts/cwolferesearch_cross-attention-is-a-fundamental-idea-that-activity-7310657138467446784-iyvV/

A final linear layer and softmax convert a decoder hidden state into a probability distribution over the vocabulary.

To avoid confusion among concepts of attention, we can remember them as follows: Bidirectional self-attention allows each token to attend to all unmasked tokens (e.g., BERT), masked/causal self-attention blocks each token to attend to future tokens (e.g., GPT), encoder-decoder attention/cross-attention lets decoder representations (i.e., queries) to attend to encoder outputs (keys and values) (e.g., multimodal).

2.3. Encoder-Only and Decoder-Only Transformers

2.4. Context Length

The context length is the maximum number of tokens that a model can process in one forward pass. During generation, the context contains the prompt and previously generated tokens.

3. Processing Pipeline

3.1. Tokenization

A tokenizer divides the input into tokens and maps each token to a discrete token ID. Tokenizers may operate at the word, subword, character, or byte level. Subword tokenization is widely used.

Text:      "Transformers are useful"
Tokens:    ["Transform", "ers", "are", "useful"]
Token IDs: [5812, 1047, 389, 7421]

Each tokenizer has a fixed vocabulary size of $v_{\text{tab}}$, which determines how many distinct tokens it can represent.

A larger vocabulary can represent common words or phrases using fewer tokens. However, it also increases the size of the embedding table and, in a language model, the vocabulary-output layer.

3.2. Embeddings

After tokenization, the model uses each token ID to look up a learned vector, called a token embedding, in an embedding table of shape $v_{\text{tab}} \times d_{\text{model}}$, where $v_{\text{tab}}$ is the vocabulary size. For a sequence of $n$ tokens, the lookup produces a matrix of token embeddings in $\mathbb{R}^{n \times d_{\text{model}}}$. Each row is one token’s embedding, and each column is one embedding feature across the tokens.

The transformer also needs positional information to distinguish token order. In models that use additive positional information, the model adds one position vector to each token embedding. For $n$ tokens, the position vectors form a matrix in $\mathbb{R}^{n \times d_{\text{model}}}$, so the resulting input embeddings have the same shape. Other models, such as those using RoPE, apply positional information within the attention mechanism.

Token embeddings are learned during training. Position vectors can be learned too, or they can be fixed. For example, the original Transformer in Attention Is All You Need used fixed positional encodings: sine and cosine functions produce a vector for each position, which is added to the token embedding and does not change during training.

3.3. Transformer Processing

Transformer layers convert the input embeddings into contextualized token representations.

Source: Attention in Transformers: Concepts and Code in PyTorch - deeplearning.ai

3.4. Layer Norm

4. Self-Attention

Self-attention updates each token representation (one row of the input matrix $X$) by combining information from other permitted positions in the same sequence. Which positions are permitted depends on the attention mask: bidirectional models (e.g., BERT) can usually use both earlier and later tokens, while causal models (e.g., GPT) cannot use/attend future tokens.

Workflow from tokenization to self-attention

The diagram shows additive positional information and unrestricted attention for simplicity. Some architectures instead apply position information to $Q$ and $K$; for example, RoPE rotates them before their dot products are calculated. You can find an example for single-head (masked) self-attention class below:

class MaskedSelfAttention(nn.Module):                     
    def __init__(self, d_model=2, # dimension for each input token embedding 
                 row_dim=0, # row and column indexes
                 col_dim=1): 
        
        super().__init__()
        
        self.W_q = nn.Linear(in_features=d_model, out_features=d_model, bias=False)
        self.W_k = nn.Linear(in_features=d_model, out_features=d_model, bias=False)
        self.W_v = nn.Linear(in_features=d_model, out_features=d_model, bias=False)
        
        self.row_dim = row_dim
        self.col_dim = col_dim

    # This method is to calculate the (masked) self-attention values for each token    
    def forward(self, token_encodings, mask=None):

        q = self.W_q(token_encodings)
        k = self.W_k(token_encodings)
        v = self.W_v(token_encodings)

        sims = torch.matmul(q, k.transpose(dim0=self.row_dim, dim1=self.col_dim))

        scaled_sims = sims / torch.tensor(k.size(self.col_dim)**0.5)

        if mask is not None:
            ## Here we are masking out things we don't want to pay attention to
            ##
            ## We replace values we wanted masked out
            ## with a very small negative number so that the SoftMax() function
            ## will give all masked elements an output value (or "probability") of 0.
            scaled_sims = scaled_sims.masked_fill(mask=mask, value=-1e9) # I've also seen -1e20 and -9e15 used in masking

        attention_percents = F.softmax(scaled_sims, dim=self.col_dim)

        attention_scores = torch.matmul(attention_percents, v)

        return attention_scores

4.1. From Hidden States to $Q$, $K$, and $V$

In self-attention, $X \in \mathbb{R}^{N \times d_{\text{model}}}$ contains one token representation per sequence position, where $N$ is the sequence length and $d_{\text{model}}$ is the representation size. In the first Transformer block, $X$ is derived from token embeddings plus any positional vectors added at the input. In later blocks, it is derived from the preceding block’s representations. In a pre-norm block, $X$ is the LayerNorm output of those representations. Thus, $Q$, $K$, and $V$ are projections of the representations entering the current self-attention layer. Methods such as RoPE apply positional information to $Q$ and $K$ after the projections. In standard multi-head attention with $h$ heads, each head usually has $d_k=d_{\text{model}}/h$ query and key dimensions. See 3 for a visualization.

For each token in one attention head, three learned linear projections produce6:

\[Q = XW^Q, \qquad K = XW^K, \qquad V = XW^V\]

where $W^Q,W^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$. Therefore, $Q,K \in \mathbb{R}^{N \times d_k}$ and $V \in \mathbb{R}^{N \times d_v}$.

  The question it answers Its role
Query $Q$ “What am I looking for?” The token doing the looking
Key $K$ “What am I?” How a token advertises itself
Value $V$ “What do I contribute if chosen?” The content actually retrieved

Separate projections let the model learn different representations for matching and for transferring information. In particular, $QK^T$ can express directional relationships, whereas $XX^T$ is symmetric.

Training note: When training from scratch, the token embeddings and projection matrices normally begin with random values and are updated by backpropagation. Learned positional embeddings are updated too; fixed positional encodings are not. By contrast, $X$, $Q$, $K$, and $V$ are temporary activations. They are recomputed on every forward pass and change as the model parameters change.

Additive positions in the attention score: in the first block of the original (post-norm) Transformer, row $i$ of $X$ is $x_i=e_i+p_i$. The query at position $i$ is then $q_i=x_iW^Q$ (row $i$ of $Q$), the key at position $j$ is $k_j=x_jW^K$ (row $j$ of $K$), and the raw score $(QK^T)_{ij}=q_ik_j^T$ splits into content-content, content-position, position-content, and position-position terms:

\[q_ik_j^T=e_iW^Q{W^K}^Te_j^T+e_iW^Q{W^K}^Tp_j^T+p_iW^Q{W^K}^Te_j^T+p_iW^Q{W^K}^Tp_j^T.\]

The second and third terms are cross terms: each pairs one token’s content with the other token’s absolute position ($p_j$ or $p_i$), not with the offset $j-i$. Without projection, the sinusoidal encodings from Section 2.1 satisfy $p_ip_j^T=\sum_r\cos\big((j-i)\omega_r\big)$, a function of the offset $j-i$ alone, but the learned $W^Q{W^K}^T$ generally breaks this. RoPE avoids the problem by rotating $q_i$ and $k_j$ so that their dot product depends on position only through $j-i$ (see Section 4.3.5).

Example: suppose the model should learn the rule “attend to the previous token”, i.e., the offset $j-i=-1$ for query-key pairs such as $(3,2)$, $(7,6)$, and $(100,99)$.

  • Additive: each pair uses different position vectors ($p_3,p_2$ vs. $p_7,p_6$ vs. $p_{100},p_{99}$), so the scores differ even for identical words. The model must learn from training data at many positions that these pairs mean the same thing.
  • RoPE: every pair uses the same rotation $R_{-1}$, so identical words give identical scores. A rule learned at one position holds at every position.

4.2. Scaled Dot-Product Attention

One attention head computes 7:

\[S = \frac{QK^T}{\sqrt{d_k}} + M, \qquad A = \operatorname{softmax}(S), \qquad Z = AV\]

Attention matrix mask (Source: https://krypticmouse.hashnode.dev/attention-is-all-you-need)

For self-attention over $N$ tokens, $S,A \in \mathbb{R}^{N \times N}$ and $Z \in \mathbb{R}^{N \times d_v}$.

A dot product depends on both vector direction and magnitude. Unlike cosine similarity, it is not normalized to the range $[-1,1]$.

4.3. Multi-Head Attention and Efficient Variants

Transformers run several attention heads in parallel. Each head produces an output $Z$ as described in Section 4.2 and can learn different relationships.

4.3.1. Projection

For each token, the head outputs are concatenated along the feature dimension, then mixed by the output projection $W^O$:

\[\operatorname{MultiHead}(X) = \operatorname{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O.\]

4.3.2. Residual Connections

We now have the output of the multi-head self-attention sublayer. We add it element-wise to the sublayer’s input, i.e., the token representations entering the layer (for the first layer, these are the input embeddings). This is called a residual (skip) connection. It is also why $W^O$ projects the concatenated heads back to $d_{\text{model}}$: both terms of the sum must have the same shape.

Why do we need residual connections 8?

  • Gradient flow (backward pass): by the chain rule, the gradient reaching an early layer is a product of the Jacobians of all later layers (weight matrices, activation derivatives, normalization, softmax). If many of these factors shrink the gradient, their product can become very small (vanishing gradients). With a residual connection $y = x + F(x)$, the Jacobian (e..g, which is a set of derivatives using for vectors) becomes $\frac{\partial y}{\partial x} = I + \frac{\partial F}{\partial x}$, so $\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial y} + \frac{\partial \mathcal{L}}{\partial y}\frac{\partial F}{\partial x}$. The identity term passes the gradient back unchanged, even when $\frac{\partial F}{\partial x}$ is small. Stacking blocks gives $x_L = x_l + \sum_{i=l}^{L-1} F_i(x_i)$, so every earlier layer, including the embedding table, receives a direct gradient signal from the output 9.

  • Preserving token information (forward pass): within each head, $AV$ makes every token’s output a weighted average of the value vectors of the tokens it attends to. This lets tokens share context, but it also dilutes each token’s own information and pulls token representations toward each other. Concatenating the heads and applying $W^O$ cannot undo this, because they only mix features within each token’s concatenated head outputs, not across tokens. So when many attention layers are applied one after another without residual connections, token representations become increasingly similar. The residual connection keeps each token’s own representation in the sum, so each layer adds context instead of overwriting it 10.

Note: the clean identity path above holds for pre-norm blocks. In post-norm blocks, LayerNorm sits on the residual path after every addition, so the gradient must also pass through each LayerNorm; deep post-norm Transformers are harder to train and usually rely on learning-rate warm-up 11.

4.3.3. Feed-Forward Network (FFN)

Let $h$ be the output of the attention sublayer after its residual connection. The position-wise FFN (Section 2.1) is a small multi-layer perceptron (MLP): a linear layer expands each token’s vector (typically to $4\,d_{\text{model}}$), a nonlinear activation $\phi$ is applied, and a second linear layer projects it back to $d_{\text{model}}$:

\[\operatorname{FFN}(x) = \phi(xW_1 + b_1)\,W_2 + b_2.\]

The original Transformer uses ReLU for $\phi$; BERT and GPT-2 use GELU; many recent LLMs, such as Llama, use a gated variant (SwiGLU). The FFN is wrapped with LayerNorm and a residual connection in the same way as attention (Section 4.3.2), which gives the output of the Transformer block, i.e., the input to the next block:

Normalization evolution: Nowadays, Pre-norm and RMSNorm are widely used.

4.3.4. Output

In a decoder-only LLM, only the output of the last Transformer block, $H \in \mathbb{R}^{n \times d_{\text{model}}}$ (one row $h_i$ per position), is converted into next-token predictions. A pre-norm model first applies a final LayerNorm (Section 3.4). The LM head, a single linear layer (not an MLP), then maps each row to $v_{\text{tab}}$ scores called logits, one per vocabulary token:

\[\text{logits} = H\,W_{\text{LM}} \in \mathbb{R}^{n \times v_{\text{tab}}}, \qquad W_{\text{LM}} \in \mathbb{R}^{d_{\text{model}} \times v_{\text{tab}}}.\]

Many models, such as GPT-2, reuse the token embedding table $E \in \mathbb{R}^{v_{\text{tab}} \times d_{\text{model}}}$ (Section 3.2) as this matrix, i.e., $W_{\text{LM}} = E^T$ (weight tying). The logit of token $k$ at position $i$ is then simply the dot product $h_i \cdot e_k$ between the hidden state and token $k$’s embedding, so a token scores high when the hidden state is aligned with its embedding. Other models, such as Llama 3 8B, learn a separate $W_{\text{LM}}$ of the same shape.

Softmax is applied to each position’s logits separately, giving $n$ probability distributions over the vocabulary. The distribution at position $i$ predicts token $i+1$: training uses all $n$ of them, while generation uses only the last one (Section 3.3).

Example (weight tying, $v_{\text{tab}} = 4$, $d_{\text{model}} = 2$): for the input “the cat”, suppose the final hidden state at the last position is $h_2 = [1, 2]$. Each logit is the dot product of $h_2$ with one row of $E$:

Token $k$ the cat sat mat
Embedding $e_k$ (row of $E$) $[1, 0]$ $[0, 1]$ $[1, 1]$ $[-1, 0]$
Logit $h_2 \cdot e_k$ $1$ $2$ $\mathbf{3}$ $-1$
Probability (softmax) $0.09$ $0.24$ $\mathbf{0.66}$ $0.01$

The model therefore predicts sat as the next token. Computing all four dot products at once is exactly the matrix product $h_2 E^T$.

4.3.5. Variants

5. Mixture of Experts

Mixture of Experts (MoE) 12 increases a model’s parameter capacity without activating all parameters for every token. In selected Transformer layers, MoE replaces the traditional FFN with multiple FFNs, called experts.

MoE overview

5.1. Experts and Routers

MoE layer with a router and experts

Routing Collapse and Remedy in MoE

5.2. Total and Active Parameters

5.3. Advantages and Limitations

Advantages:

  1. Increases model capacity without increasing computation in direct proportion to the total parameter count.
  2. Allows different experts to learn different patterns.
  3. Can improve performance over dense models under a similar computational budget.

Limitations:

  1. Requires substantial memory to store all experts.
  2. Introduces routing and communication overhead.
  3. May overload some experts while leaving others under-trained (see Routing collapse above for a remedy).
  4. Is more complex to train and deploy than a dense model.

5.4. Recent trends till Sept. 2026

MoE shared experts

6. References

[^10]: [Stanford CME295 Transformers & LLMs Autumn 2026 course](https://cme295.stanford.edu/syllabus/)

  1. Vaswani et al., 2017. Attention Is All You Need ↩

  2. Su et al., 2024. RoFormer: Enhanced Transformer with Rotary Position Embedding ↩

  3. LLM Visualization - For inference but highly recommend to check ↩ ↩2

  4. Devlin et al., 2019. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding ↩

  5. Liu et al., 2019. RoBERTa: A Robustly Optimized BERT Pretraining Approach ↩

  6. LLM Architecture Refresh: Inside a Transformer Block — Attention, Heads, and the FFN ↩ ↩2 ↩3

  7. The Math Behind Multi-Head Attention in Transformers ↩

  8. StackExchange: Why are residual connections needed in transformer architectures? ↩

  9. He et al., 2016. Identity Mappings in Deep Residual Networks ↩

  10. Dong et al., 2021. Attention Is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth ↩

  11. Xiong et al., 2020. On Layer Normalization in the Transformer Architecture ↩

  12. A Visual Guide to Mixture of Experts (MoE) ↩

  13. Fedus et al., 2022. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity ↩