Under the Hood: The Mathematics and Architecture of Language Models

What is a language Model? A Large Language Model is a sophisticated mathematical function that predicts what word comes next for any prompt or document. Like a magic 8 ball that can finish your sentence.

Examples: Anthropic’s Claude, Google’s Gemini, and OpenAI’s GPT

LLM = Large Language Model

There are two types of Language models. Auto-regressive models can predict the next future token by knowing the past tokens or the future tokens, but not both. Auto-encoding models can learn the entire sequence by predicting tokens from knowing the past and future tokens. BERT is an example of auto-encoding models.

How do you train a Language Model?

Tokenization → model architecture → training

Tokenization

A converter that takes in bytes and encodes them to a sequence of integers ([987, 65, 32]). The sequence represents the tokens. The input in a human language is turned into tokens that a computer can understand.

Example of a tokenizer: Byte-Pair encoding (BPE)

Purpose: break input into chunks

Significance: Can take a long byte stream (900 bytes → ~200 tokens) and reduce the context length

Tokenization:

• split text into tokens → look each token up in a fixed vocabulary → get an integer ID

Example: “the cat sat” → [“the”, “cat”, “sat”] → [40, 1723, 908]

• Those numbers are token IDs (an index into the vocabulary)

Model Architecture

What is a language model made of? Transformers

A transformer is a neural network architecture. It is the engine that gets stacked N times to build the full language model. One single transformer architecture with multiple stacked layers and attention heads.

Why do we like transformers? attention mechanism. Older architectures (RNNs) process text one token at a time in order. Attention mechanism lets the model weigh the relevance of every other token in a sequence when processing a given token. (parallel motion - tokens processed simultaneously)

Purpose? Attention mechanism tells the relationship between two words and adds context

---

Attention, Context, and Feed-Forward Networks

Attention: three vectors — Query (Q — what am I looking for?), Key (K — what do I contain?), and Value (V — what information do I actually offer?)

Equation to measure attention

Each token’s Query is compared against every other token’s Key → produce a score → using softmax, scores are turned into weights (so they sum to 1) → output: weighted sum of all tokens’ values = weighted by how relevant each one is → token with the highest weight is the next predicted word

• Context window: the max number of tokens the model can ingest — Attention operates within this window
• Feed-forward network (MLP): attention mixes information across tokens (like a group discussion-everyone share info with each other) → each token individually passes through a small neural network that processes it independently of other tokens (each person privately processes what was learned, updating their own notes)

MLP is a small neural network with 3 steps:

1. Expand: take the word’s vector → stretch into a bigger vector
2. Nonlinearity: apply a simple math function that lets the network learn non-straight-line, complex patterns (without this step, stacking layers wouldn’t add any real power)
3. Compress: shrink it back down to the original size

Mixture of Experts (MoE) replaces a single MLP with multiple “expert” MLPs and a small router network that decides which expert(s) each token gets sent to.

Significance: model can have a large parameter count while only activating a few for a single token — improves compute efficiency for the model

Attention — every word asks “who in this sentence matters to me?”

Feed-forward layer — what should I do with what I found?

Significance: repeats across every layer → more context-aware the deeper into the network

Prediction: final layer’s output for the last token → convert into probability distribution = model’s prediction for the next token

Given all this, what happens to a sentence?

1. Tokenization + Embedding: input sentence is split into tokens → token converted to a vector (an embedding)

2. Positional Encoding: attention has no built-in sense of order → RoPE (Rotary Positional Embeddings) rotates Q/K vectors based on position

   • model knows “cat sat on mat” is different from “mat sat on cat,” even though both use the same words

Significance: position embeddings added to token embeddings to provide info about the position of each token within a sequence

3. Self-Attention Mechanism: every word asks “who in this sentence matters to me?”

4. Multi-Head Attention: Attention mechanism done multiple times in parallel (“heads”) → each head captures different relationships or patterns in the data → The outputs of all heads get concatenated and combined → model has a deeper understanding of the training data

5. Residual Connection + Normalization: The attention output gets added back to the original input (a “residual” or “skip” connection — helps preserve info and prevents vanishing gradients in deep networks ) → then normalized (rescales activations to stabilize training — without it values grow or shrink uncontrollably as they pass through multiple layers)

6. Feed-Forward Network (MLP): Each token with context from attention passes through the same small neural network independently. Another residual connection and normalization follow.

7. Stack and Repeat: Steps 3–6 together make up one transformer layer/block. Real models stack many of these (GPT-3 uses 96). Each layer builds a deeper representation of the input, informed by all the layers before it

8. Output Prediction: After the final layer, the representation of the last token position is projected onto the vocabulary and turned into a probability distribution — this is the model’s prediction for the next token. That token gets generated, appended to the sequence, and the whole process runs again for the next one (this is why generation is sequential even though processing within each step is parallel).

---

Training

Training is the process of adjusting weights to reduce the loss function.

Weights are numbers that tell the model which words in a sentence to focus on when processing or translating a specific token.

Important why? What are the parameters to consider when building a model? These make the difference between a run that blows up and a model that runs efficiently.

Parameters:

(1) Loss function — num that measures how inaccurate the model’s predictions are

• cross-entropy loss = model outputs a probability distribution based on the distance between the training’s next predicted text and the actual next token
• multi-token prediction = model predicts several future tokens simultaneously at each position…more to learn from per token processed

(2) Optimizer — algorithm used to adjust weights (Determined from the gradients computed from the loss function) ex. AdamW, SOAP, Muon

(3) Initialization scale — starting values of the model’s weights (scale refers to how large/small those values are)

• scale too large: activations explode throughout the layers
• scale too small: gradients vanish → early training is slow/unstable

(4) Learning rate schedule — how big is each weight update?

• too high: training diverges too broadly
• too low: training lags and wastes compute
• solution: forget the rate, use a schedule that changes learning rate over time (start learning rate low → lower learning rate after early training stage → fine-tune as it converges)

(5) Regularization — techniques that turn a model away from memorizing training data to learning generalized patterns

• weight decay = slowly decreases weights to smaller values with each update → model does not depend heavily on 1 parameter

(6) Batch size — how many tokens go in each batch (group of training examples processed before weight updates start)

• larger batches = accurate estimate of gradients, can parallelize training
• batch size ↑↑↑ increases → adjust learning rate schedule = batch size and learning rate tuned together

(7) MoE specific — load balancing

Problem: a learned router routes → tokens can be sent to favorite experts leaving some tokens wasted with an undertrained model

Solution: load balancing (multiple techniques) pushes router to distribute tokens evenly across all experts

---

RAG

On their own, LLMs can only use the data from their training, which means their knowledge has a cutoff date and can’t include private or fast-changing information (ex. a company’s internal documents or yesterday’s news). RAG (Retrieval-Augmented Generation) is a technique that gives a language model access to outside information it wasn’t trained on.

RAG adds a retrieval step before the model generates its answer: when a user asks a question, a search system looks through an database and pulls out the most relevant sentences. Those sentences are inserted into the model’s prompt as extra context, so the LLM can generate a response grounded in real source material rather than relying purely on memorized patterns. This helps with reducing hallucination and keeping answers current, without the time and cost of retraining the model itself.

End-to-end RAG pipeline:

1. Tokenization — Break input text into tokens, convert each to a token ID

   • Document chunk: “Healthy foods to build muscle protein includes eggs, Greek yogurt, and chicken breast”

Tokenized: [“Healthy”, “foods ”, “to ”, “build ”, “muscle ”, “protein ”, “includes ”, “eggs”, “Greek ”, “yogurt”, “and”, “chicken”, “breast”]

Token IDs: [1189, 4402, 8821, 312, 221, 6754, 940, 456, 1936, 3201, 324, 569, 4019]

2. Embedding — Convert each token ID into a vector that captures meaning

   • token ID gets passed through a learned embedding layer that outputs a vector (hundreds/thousands of dimensions long, made of decimals)

Token IDs → vector

[0.021, -0.084, 0.132, 0.056, …, -0.019] goes on until 768 numbers long

User queries: “What are healthy foods to build muscle?” → query embedded into [0.019, -0.079, 0.129, 0.058, …, -0.024]

Similar vectors means the sentences have similar meanings.

3. Indexing (done ahead of time) — Every document in the knowledge base is also tokenized, embedded, and stored in a vector database

   • Sentences split into chunks → chunk gets a token ID attached → token ID embedded into a multi-dimensional vector → vector stored in vectorDB

{ id: “doc_042”, vector: [0.021, -0.084, …], text: “Healthy foods to…” }

{ id: “doc_043”, vector: [0.187, 0.003, …], text: “Protein is built…” }

{ id: “doc_044”, vector: [-0.045, 0.098, …], text: “Muscle protein contains…” }

An HNSW (Hierarchical Navigable Small World) index organizes the vectors into a graph structure so the system can search efficiently instead of comparing the query against every single document.

4. Retrieval — user prompts a question/query → query gets embedded → cosine similarity to find the most relevant chunks (how the magic 8 ball prediction works)

   • Cosine similarity = measures how similar two vectors are by looking at the angle between them, rather than their raw distance or magnitude

Vector acts like an arrow pointing in a direction. Cosine similarity finds out if they point in a similar direction. Vectors are the query vector (A) and the stored document vectors (B)

. dot product

x multiplication

||A|| magnitude/length of vector A

Returns a similarity score

1 → vectors point in exactly the same direction (perfectly similar)

0 → vectors are perpendicular / unrelated

-1 → vectors point in exactly opposite directions (perfectly dissimilar)

doc_042 (healthy foods to): 0.94 ← highest match ← selected answer

doc_043 (protein is build): 0.31

doc_044 (Muscle protein contains): 0.08

5. Augmentation — Those retrieved chunks are inserted into the prompt as extra context

   • Gives you a clean answer by following the system prompt template → inserting answer → figures out how to order chunks → cites sources if needed

System: Answer the question using only the context below

Context: “Healthy foods to build muscle protein includes eggs, Greek yogurt, and chicken breast”

Query: “What are healthy foods to build muscle?”

6. Generation — The LLM generates its answer using both the original question and the retrieved chunks

   • LLM processes the prompt → self-attention layers let the model weigh the retrieved context against the query when predicting the next token (word) → forms the answer

RAG is a part of how language models predict the next token or word, although there is much more to these models.

Credit to: Stanford CS336 Language Modeling from Scratch