ai_ml / deep learning / 01_neural_network_basics.md

Neural network basics

6 interview angles 5 min read source

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 across backward() 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; forgetting model.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.