Build A Large Language Model 读书笔记 - Virex

On Memory

The right rule is: only keep information that still matters across sessions and cannot be cheaply re-derived from the current workspace.

Memory/Task/Plan/CLAUDE.md

Short rule of thumb: only useful for this task — use task or plan. Useful next session too — use memory. Long-lived instruction text — use CLAUDE.md.

The next-word prediction task

is a form of self-supervised learning, which is a form of self-labeling. This means that we don’t need to collect labels for the training data explicitly but can use the structure of the data itself: we can use the next word in a sentence or document as the label that the model is supposed to predict. Since this next-word prediction task allows us to create labels “on the fly,” it is possible to use massive unlabeled text datasets to train LLMs

The general GPT architecture is relatively simple. Essentially, it’s just the decoder part without the encoder.

Working with text data

2.1 Understanding word embeddings

The advantage of optimizing the embeddings as part of the LLM training instead of using Word2Vec is that the embeddings are optimized to the specific task and data at hand.

2.7 Creating token embeddings

The weight matrix of the embedding layer contains small, random values. These values are optimized during LLM training as part of the LLM optimization itself.

the embedding layer is essentially a lookup operation that retrieves rows from the embedding layer’s weight matrix via a token ID.

2.8 Encoding word positions

The embedding layer converts a token ID into the same vector representation regardless of where it is located in the input sequence. For example, the token ID 5, whether it’s in the first or fourth position in the token ID input vector, will result in the same embedding vector.

Two broad categories of position-aware embeddings:

1. Relative positional embeddings. - The model can generalize better to sequences of varying lengths, even if it hasn’t seen such lengths during training.
2. Absolute positional embeddings. - Associated with specific positions in a sequence. （OpenAI’s GPT models used）

Summary

Coding attention mechanisms

We will implement four different variants of attention mechanisms:

1. Simplified self-attention: A simplified self-attention technique to introduce the broader idea
2. Self-attention: Self-attention with trainable weights that forms the basis of the mechanism used in LLMs
3. Causal self-attention: A type of self-attention used in LLMs that allows a model to consider only previous and current inputs in a sequence, ensuring temporal order during the text generation
4. Multi-head attention: An extension of self-attention and causal attention that enables the model to simultaneously attend to information from different representation subspace

3.1 The problem with modeling long sequences

The job of the encoder is to first read in and process the entire text, and the decoder then produces the translated text

The big limitation of encoder–decoder RNNs is that the RNN can’t directly access earlier hidden states from the encoder during the decoding phase. Consequently, it relies solely on the current hidden state, which encapsulates all relevant information. This can lead to a loss of context, especially in complex sentences where dependencies might span long distances.

3.3.2 Computing attention weights for all input tokens

1. Compute attention scores: Compute the attention scores as dot products between the inputs.
2. Compute attention weights: The attention weights are a normalized version of the attention scores. (attention weights determine the extent to which a context vector depends on the different parts of the input (i.e., to what extent the network focuses on different parts of the input).)
3. Compute context vectors: The context vectors are computed as a weighted sum over the inputs.

3.4.1 Computing the attention weights step by step

Why query, key, and value?

• A query is analogous to a search query in a database. It represents the current item (e.g., a word or token in a sentence) the model focuses on or tries to understand. The query is used to probe the other parts of the input sequence to determine how much attention to pay to them.
• The key is like a database key used for indexing and searching. In the attention mechanism, each item in the input sequence (e.g., each word in a sentence) has an associated key. These keys are used to match the query.
• The value in this context is similar to the value in a key-value pair in a database. It represents the actual content or representation of the input items. Once the model determines which keys (and thus which parts of the input) are most relevant to the query (the current focus item), it retrieves the corresponding values.

A compact self-attention class

import torch.nn as nn
class SelfAttention_v1(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.W_query = nn.Parameter(torch.rand(d_in, d_out))
        self.W_key = nn.Parameter(torch.rand(d_in, d_out))
        self.W_value = nn.Parameter(torch.rand(d_in, d_out))

    def forward(self, x):
        queries = x @ self.W_query
        keys = x @ self.W_key
        values = x @ self.W_value
        attn_scores = queries @ keys.T # omega
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1
        )
        context_vec = attn_weights @ values
        return context_vec

Hiding future words with causal attention

Causal attention, also known as masked attention, restricts a model to only consider previous and current inputs in a sequence when processing any given token when computing attention scores.