The machinery: how a neural network computes and learns · Transformers

A neural network is a function with adjustable numbers in it, and learning means adjusting those numbers to make a loss smaller. Everything in this guide, attention included, is built from three operations: matrix multiplication, simple elementwise nonlinear functions, and normalization. Training uses one algorithm: gradient descent, with gradients computed by backpropagation.

Tensors and shapes

A tensor is an n-dimensional array of numbers. Shapes are the grammar of deep learning: a batch of token ids is (B, T), after embedding it is (B, T, d), and attention scores are (B, h, T, T). When reading code or papers, annotate every line with its shape.

The linear layer: a batch of dot products

A linear (dense) layer computes y = xW + b. With x of shape (d_in) and W of shape (d_in, d_out), each output y_j is a dot product of x with column j of W, plus a bias.

y_j = \sum_{i=1}^{d_{\text{in}}} x_i W_{ij} + b_j

A dot product measures alignment: large when two vectors point the same way, zero when orthogonal, negative when opposed. So each output unit asks “how much does the input look like my pattern?” Hold on to this: attention asks the same question between tokens.

Cost: multiplying an (m × k) matrix by a (k × n) matrix takes about 2·m·k·n floating-point operations (one multiply and one add per term). This rule drives every compute estimate later.

Why nonlinearity

Stacking linear layers gives another linear layer: W₂(W₁x) = (W₂W₁)x. Depth buys nothing without something nonlinear in between. An activation function applied elementwise breaks the collapse.

Activation	Formula	Where you see it
ReLU	max(0, x)	classic CNNs, early transformers
GELU	x · Φ(x), Φ = Gaussian CDF	BERT, GPT-2, GPT-3
SiLU (Swish)	x · σ(x)	inside SwiGLU in most modern LLMs

With nonlinearities, a wide enough network can approximate any reasonable function (the universal approximation theorem). That says good weights exist, not that training will find them or how many you need. The rest of deep learning is about making them findable.

Loss: turning “good” into a number

For a language model the task at each position is classification: which of V vocabulary tokens comes next? The network outputs V raw scores called logits, z. Softmax turns logits into probabilities:

p_i = \frac{e^{z_i}}{\sum_{j=1}^{V} e^{z_j}}

Cross-entropy loss is the negative log-probability assigned to the correct answer y:

\mathcal{L} = -\log p_y

Three readings of the same number. It is the surprise of the true token. Minimizing its average is maximum-likelihood estimation. And exp(average loss) is perplexity, the effective number of tokens the model is choosing between: a loss of 2.0 nats means a perplexity of about 7.4.

A useful fact: the gradient of cross-entropy with respect to the logits is simply p − onehot(y), the predicted distribution minus the truth. Learning raises the score of the right token and lowers the others in proportion to the probability they took.

Gradient descent

The gradient ∇L is the vector of partial derivatives of the loss with respect to every parameter; it points uphill. Step downhill:

\theta_{t+1} = \theta_t - \eta \, \nabla_\theta \mathcal{L}(\theta_t)

η is the learning rate. The gradient over the full dataset is too expensive, so we estimate it on a mini-batch (stochastic gradient descent). The estimate is noisy, which tends to help generalization but limits how big a step is safe.

Backpropagation: the chain rule, organized

A network is a chain of simple functions. The chain rule says the derivative of a composition is the product of local derivatives:

\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial y} \cdot \frac{\partial y}{\partial x}

Backprop runs the network forward, storing intermediate results (activations). It then walks the computation graph backward, multiplying local derivatives and accumulating each parameter’s gradient. This is reverse-mode automatic differentiation; PyTorch’s autograd records the graph as you compute and runs it backward when you call loss.backward().

Two consequences matter at scale:

Compute: the backward pass costs about twice the forward pass. Forward is about 2 FLOPs per parameter per token and backward about 4, so training costs about 6 FLOPs per parameter per token (Part 8).
Memory: forward activations must be kept until the backward pass uses them. Activation memory, not weights, often limits batch size and sequence length.

The training loop

Sample a batch of examples.
Forward pass: compute predictions and the loss.
Backward pass: compute gradients.
Optimizer step: update parameters (Part 2 covers better optimizers than plain SGD).
Repeat for many steps, watching training and validation loss.

Questions you should now be able to answer

Why can’t a stack of linear layers learn XOR, and what single change fixes it?
What is the shape of W in a layer mapping 768-d vectors to 3072-d vectors, and how many FLOPs does it take per token?
Why is the gradient of softmax cross-entropy p − y, and what does that say about how learning moves probabilities?
Why does training need far more memory than inference for the same model?
If validation loss is 1.5 nats per token, what is the perplexity?