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. Bag-of-Words
    2. Static Word Embeddings
    3. RNNs and Attention
  2. Transformer Architecture
    1. Transformer Encoder
    2. Transformer Decoder
    3. Encoder-Only and Decoder-Only Transformers
    4. Context Length
  3. Processing Pipeline
    1. Tokenization
    2. Embeddings
    3. Transformer Processing
  4. Self-Attention
    1. From Hidden States to Q, K, and V
    2. Scaled Dot-Product Attention
    3. Multi-Head Attention and Efficient Variants
  5. Mixture of Experts
    1. Experts and Routers
    2. Total and Active Parameters
    3. Advantages and Limitations
  6. References

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, which are incorporated with positional information, into contextualized representations. Each encoder layer contains:

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, 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

The model uses each token ID to retrieve a learned vector from its embedding table. It then incorporates positional information to distinguish the order of the tokens.

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

4. Self-Attention

Self-attention updates each token representation 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$

At an attention layer, $X \in \mathbb{R}^{N \times d_{\text{model}}}$ contains one hidden-state vector per token. Particularly, $N$ is the sequence length and $d_{\text{model}}$ is the dimension of each token embedding and hidden representation. In standard multi-head attention with $h$ heads, this dimension is divided among the heads, so each head usually has $d_k = d_{\text{model}}/h$ query and key dimensions .For the first layer, $X$ comes from token embeddings, with positional information incorporated according to the architecture. In later layers, it is the output of the preceding layer.

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

\[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.

4.2. Scaled Dot-Product Attention

One attention head computes 2:

\[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 normally run several attention heads in parallel. Each head can learn different relationships; their outputs are concatenated and mixed through an output projection $W^O$:

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

5. Mixture of Experts

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

MoE overview

5.1. Experts and Routers

MoE layer with a router and experts

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.
  4. Is more complex to train and deploy than a dense model.

6. References

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

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

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