The Mathematics

The Four Pillars

Understanding LLMs deeply requires four areas of mathematics. Don't worry — we'll build intuition first, then rigor. Each topic connects directly to how models work.

Linear Algebra: The Language of Deep Learning

Neural networks are fundamentally about linear transformations (matrix multiplication) followed by non-linear activations. Understanding linear algebra is understanding how information flows through networks.

Matrix Multiplication as Composition

When we multiply matrices, we're composing transformations:

This says: transform vector x using transformation W. In neural networks:

• x is the input (e.g., token embedding)
• W contains the learned weights
• y is the transformed representation

Attention as Matrix Operations

The attention mechanism is elegantly expressed as matrix operations:

Calculus: How Networks Learn

Training neural networks is optimization — finding weights that minimize loss. This requires computing how changing each weight affects the loss: the gradient.

The Derivative

The derivative tells us the rate of change:

In neural networks, we need partial derivatives — how does the loss change with respect to each individual weight?

The Chain Rule

Neural networks are compositions of functions. The chain rule tells us how to differentiate them:

Probability: Modeling Uncertainty

Language is inherently uncertain. Probability gives us the tools to model this uncertainty and make predictions under uncertainty.

Maximum Likelihood Estimation

Language models are trained with maximum likelihood: find parameters that make the training data most probable.

For language modeling, this becomes predicting the next token:

Conditional Probability

The chain rule of probability lets us decompose joint distributions:

This is exactly what autoregressive language models do! They model P(xᵢ | x<ᵢ) for each position.

Sampling from Models

To generate text, we sample from the predicted distribution:

Different sampling strategies (greedy, temperature, top-k, top-p) modify this distribution to control randomness vs. coherence.

Information Theory: Measuring Information

Information theory, founded by Claude Shannon, gives us precise ways to measure information and uncertainty. It's fundamental to understanding loss functions.

Entropy

Entropy measures uncertainty in a distribution:

Higher entropy = more uncertainty. A uniform distribution has maximum entropy; a deterministic distribution has zero entropy.

Cross-Entropy

Cross-entropy measures how well distribution q approximates distribution p:

In language modeling:

• p is the true distribution (one-hot: the actual next token)
• q is the model's predicted distribution
• Minimizing cross-entropy = making model predictions match reality

Perplexity

Perplexity is the standard metric for language models:

KL Divergence

KL divergence measures how one distribution differs from another:

Used in RLHF to keep the fine-tuned model close to the base model (preventing catastrophic forgetting).

Linear Algebra: Worked Examples

Let's work through the key linear algebra operations that appear in transformers, step by step.

Matrix Multiplication in Attention

Consider computing QK^T for a 3-token sequence with d_k = 2:

Eigenvalues and Attention

The attention matrix A = softmax(QK^T/√d_k) is a row-stochastic matrix (rows sum to 1). Its eigenvalues reveal important properties:

• Largest eigenvalue λ₁ = 1: Corresponds to the stationary distribution — which positions receive the most attention overall
• Other eigenvalues |λᵢ| < 1: Control how quickly attention patterns converge under repeated application
• Spectral gap = 1 - |λ₂|: A larger spectral gap means the attention is more "focused"

Low-Rank Approximation & LoRA

LoRA exploits the fact that weight updates during fine-tuning are typically low-rank. If ΔW ∈ ℝ^(d×d) has rank r ≪ d, we can factorize:

Parameters: d×r + r×d = 2dr instead of d². For d=4096, r=8: 65,536 instead of 16,777,216 — a 256× reduction!

Calculus: Derivations for Neural Networks

Derivative of the Softmax

Softmax is used everywhere in transformers. Let's derive its Jacobian:

The partial derivative ∂σ_i/∂z_j has two cases:

This is a crucial result: the softmax gradient connects the output probabilities directly to the input logits. When we combine this with cross-entropy loss, we get a beautifully simple gradient.

Cross-Entropy + Softmax Gradient Derivation

This is the most important gradient in LLM training. Let's derive it step by step.

The Gradient of ReLU

ReLU's gradient is either 0 or 1. This is what makes deep networks trainable — the gradient doesn't shrink as it passes through ReLU (unlike sigmoid, where the max gradient is 0.25, causing vanishing gradients in deep networks).

Gradient of Layer Normalization

LayerNorm is more complex than it appears. Given y = γ · (x - μ)/√(σ² + ε) + β:

The gradient w.r.t. x has three terms: the direct gradient, a correction for the mean, and a correction for the variance. Thankfully, autograd handles this automatically.

Probability: Language Modeling as Probability

The Chain Rule of Probability

The most fundamental equation in autoregressive language modeling. For any joint distribution, we can factorize:

This is an identity, not an approximation! Every joint distribution can be decomposed this way. Autoregressive language models parameterize each conditional:

where f_θ is the transformer network that produces logits for every token in the vocabulary.

Maximum Likelihood Estimation — Derivation

We want to find θ that maximizes the probability of the training data. For a dataset D = {document₁, document₂, ...}:

Conditional Independence and the Markov Property

Computing P(xᵢ | x<ᵢ) requires conditioning on all previous tokens. This is expensive for long sequences. If we make a Markov assumption:

This approximation is used in n-gram models (k = n-1). Transformer models make no Markov assumption — they condition on the full context. This is what gives them their power but also their O(T²) computational cost.

Bayes' Theorem in LLMs

Bayesian thinking appears in several places in LLM training:

• Weight decay as a prior: L2 regularization is equivalent to placing a Gaussian prior on weights
• Dropout as Bayesian approximation: Monte Carlo dropout approximates Bayesian inference
• RLHF reward model: P(human_preference | model_A, model_B) uses Bayesian ranking
• In-context learning: The model implicitly performs Bayesian inference over task hypotheses given the prompt

Information Theory: Proofs and Connections

Why Cross-Entropy?

Why is cross-entropy the "right" loss function? It's not arbitrary — it's the unique loss function with certain desirable properties.

Perplexity: Information-Theoretic Interpretation

Perplexity PPL = 2^H(p, q) (for log base 2) has a clean interpretation:

• PPL = 1: Perfect prediction (every token gets probability 1). The model is certain about every next token.
• PPL = |V|: Uniform distribution. The model is maximally uncertain.
• PPL = 10: The model is, on average, as uncertain as if choosing uniformly among 10 tokens.

Mutual Information and Attention

There's a fascinating connection between mutual information and attention:

Attention weights can be viewed as estimates of how much mutual information one token shares with another. Tokens that share high mutual information (i.e., knowing one helps predict the other) should have high attention scores.

Optimization Theory: Why Training Works

Convex vs. Non-Convex Landscape

The loss landscape of a neural network is non-convex — it has many local minima, saddle points, and flat regions. Yet neural networks train successfully. Why?

• Most local minima are good: In high dimensions, local minima tend to have similar loss values. "Bad" local minima (with much higher loss) are rare.
• Saddle points are the real enemy: Most stationary points in high-dimensional spaces are saddle points (not minima). SGD with momentum escapes them effectively.
• Gradient noise helps: The stochasticity of mini-batch gradients acts as implicit regularization, helping escape sharp minima and find flat ones.
• Overparameterization: Modern networks have far more parameters than needed. This creates many connected "valleys" in the loss landscape, almost all with similar loss.

SGD and Generalization

A surprising empirical finding: SGD with small batches generalizes better than SGD with large batches (even when they see the same total data). Why?

The noise in small-batch SGD acts as a form of implicit regularization:

• It pushes the optimization toward flat minima (wide valleys in loss landscape)
• Flat minima generalize better because small perturbations in parameters don't cause large loss increases
• Large batches find sharper minima that don't generalize as well

Convergence Guarantees

For convex functions, SGD with learning rate η_t converges if:

Common choices: η_t = 1/t (converges but slow) or η_t = 1/√t (converges with rate O(1/√t)). For non-convex functions (neural networks), convergence to a global minimum is not guaranteed, but SGD typically converges to a "good enough" local minimum.

Putting It All Together

These mathematical tools combine to create modern LLMs:

Complete Training Loop — Mathematical View