Neural network basics
Enough structure to reason about transformers, since that’s where this leads. A feed-forward network is a stack of matrix multiplies with non-linearities between them.
The building block
h = activation(x @ W + b)
That’s a layer. W is (in_features, out_features), b is (out_features,). Stack them:
import torch.nn as nn
model = nn.Sequential(
nn.Linear(784, 256), nn.ReLU(),
nn.Linear(256, 128), nn.ReLU(),
nn.Linear(128, 10), # logits - no softmax here
)
No softmax on the output layer. nn.CrossEntropyLoss applies log-softmax internally for numerical stability. Applying softmax yourself and then passing it to CrossEntropyLoss double-applies it — a real and common bug that makes training mysteriously slow. See ../00_math_foundations/04_information_theory.md.
Why the non-linearity matters
Without an activation function, stacked linear layers collapse:
W2 @ (W1 @ x) = (W2 @ W1) @ x = W_combined @ x
A hundred linear layers equal one linear layer. The activation is what makes depth meaningful — it’s the single most important thing to be able to say about network structure.
Universal approximation, and why it’s not the point
A network with one hidden layer can approximate any continuous function to arbitrary precision, given enough width. True and largely useless: the theorem says nothing about how many neurons, or whether gradient descent will find them.
The practical fact is that depth is exponentially more efficient than width for many function families — a deep narrow network represents things a shallow wide one needs exponentially more units for. That’s why architecture research is about depth and structure, not width.
Parameter counting
Worth being able to do in your head, because it drives memory:
Linear(in, out) -> in * out + out # weights + biases
sum(p.numel() for p in model.parameters() if p.requires_grad)
For the network above: 784*256 + 256 + 256*128 + 128 + 128*10 + 10 ≈ 235,000.
Memory during training is roughly 4x the parameter count in floats: weights, gradients, and two Adam moment tensors — before activations. That’s why fine-tuning a 7B model needs far more than 7B×4 bytes, and why LoRA exists. See ../00_math_foundations/03_calculus_optimization.md.
The training loop
for epoch in range(epochs):
model.train()
for xb, yb in train_loader:
optimizer.zero_grad() # gradients accumulate by default in PyTorch
loss = criterion(model(xb), yb)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
model.eval() # switches dropout and batchnorm
with torch.no_grad(): # no graph, less memory
...
Four things people forget, all of which produce silent misbehaviour rather than errors:
zero_grad()— PyTorch accumulates gradients acrossbackward()calls. Omitting it means every step uses the sum of all previous gradients. (It’s occasionally intentional, for gradient accumulation across micro-batches.)model.eval()— turns off dropout and switches batch norm to running statistics. Forgetting it makes evaluation non-deterministic and worse.model.train()— turn it back on. Forgetting this after a validation pass silently disables dropout for the rest of training.torch.no_grad()— no autograd graph during inference. Without it you build a graph you never use and can run out of memory.
Batch size, epochs, steps
- Batch — samples per gradient update.
- Step / iteration — one gradient update.
- Epoch — one pass over the training set.
steps_per_epoch = n_samples / batch_size. Learning-rate schedules are usually defined in steps, not epochs, which matters when you change the batch size.
Gradient accumulation simulates a large batch on small hardware:
for i, (xb, yb) in enumerate(loader):
loss = criterion(model(xb), yb) / accum_steps
loss.backward() # accumulate
if (i + 1) % accum_steps == 0:
optimizer.step()
optimizer.zero_grad()
Note the division by accum_steps — without it the effective learning rate scales up by that factor.
Output layer by task
| Task | Final layer | Loss |
|---|---|---|
| Binary classification | Linear(h, 1) |
BCEWithLogitsLoss |
| Multi-class | Linear(h, n_classes) |
CrossEntropyLoss |
| Multi-label | Linear(h, n_labels) |
BCEWithLogitsLoss |
| Regression | Linear(h, 1) |
MSELoss / HuberLoss |
The WithLogits variants fuse the sigmoid into the loss for numerical stability. Prefer them over applying sigmoid then BCELoss.
Multi-label uses independent sigmoids, not softmax — softmax forces the outputs to compete. See ../01_ml_foundations/01_ml_problem_types.md.
When to reach for a neural network
| Use one | Use gradient boosting |
|---|---|
| text, images, audio, video | tabular data |
| very large datasets | small to medium data |
| transfer learning available | no pretrained model fits |
| multi-modal fusion | homogeneous features |
| you need learned embeddings | you need explainability |
For tabular data, gradient boosting still wins in 2026. Being clear about that rather than defaulting to deep learning is the stronger position — see ../02_classical_ml/05_gradient_boosting.md.
Interview angle
- “Why do neural networks need activation functions?” — without them, stacked linear layers compose into a single linear transform, so depth adds nothing. The non-linearity is what makes a deep network more expressive than a shallow one.
- “What does the universal approximation theorem tell you?” — that one hidden layer suffices in principle. It says nothing about width required or whether optimisation finds the solution, so it doesn’t justify shallow architectures. Depth is exponentially more parameter-efficient in practice.
- “Why is
optimizer.zero_grad()necessary?” — PyTorch accumulates gradients across backward passes. Without it each step uses the running sum of all previous gradients. The accumulation behaviour is deliberate, so you can simulate large batches. - “What does
model.eval()change?” — disables dropout and switches batch norm to running statistics rather than batch statistics. Forgetting it makes evaluation noisy and worse; forgettingmodel.train()afterwards silently disables regularisation. - “How much memory does training a model need?” — roughly 4x the parameters for weights, gradients and Adam’s two moments, plus activations retained for the backward pass. That’s the motivation for gradient checkpointing, 8-bit optimisers and LoRA.
- “Would you use a neural network for tabular data?” — usually not. Gradient boosting handles heterogeneous features, needs less data and less tuning, and still outperforms on most structured problems. Neural nets win when you need embeddings of high-cardinality entities or multi-modal fusion.