The forward pass

A layer is three operations

A dense layer takes a vector, multiplies it by a weight matrix, adds a bias vector, and passes the result through an activation. It is dense because every input connects to every output. The activation is a fixed nonlinear function, and the parameters are the numbers training adjusts. If the input has three features and the output has two, the matrix is two by three. The bias then has two entries, and the whole layer has eight parameters. A batch of inputs turns the vector into a matrix. The same weight matrix is reused for every row.

The linear part changes basis: it projects the input onto a new set of directions. The bias shifts the origin. That shift matters because many datasets are not centred on it. Without a bias, every boundary has to pass through zero. The activation then bends the result. All three parts are necessary, and only one of them is nonlinear.

The parameter count is the other thing to read from the shape. A layer with four inputs and eight outputs has thirty-two weights and eight biases. The activations are not parameters, but they are the memory. When a model is too large, the number that matters is usually the activation count, not the weight count.

Why the nonlinearity is the point

Without an activation, two layers collapse into one. The composition of two linear maps is another linear map. A hundred stacked matrices therefore have exactly the expressive power of a single matrix. The parameters are wasted, the depth is cosmetic, and the model still cannot represent an XOR. XOR is true when exactly one of two inputs is true.

ReLU is the default because it is cheap. The name is short for rectified linear unit, and it returns the input when it is positive and zero otherwise. Its gradient, the slope of the output with respect to the input, is either zero or one. Sigmoid and tanh saturate: for a large positive or negative input, the derivative is almost zero. The gradient that reaches earlier layers is then tiny, which is the vanishing gradient problem in its simplest form. ReLU avoids it on the positive side, at the cost of units that can get stuck at zero.

GELU and SiLU are smooth alternatives. They behave like ReLU for large inputs and like a soft gate near zero. They cost more to evaluate, and they often buy a small accuracy gain in large models. For a first network the choice barely matters. What matters is that there is a nonlinearity at all, and that it comes after the linear part rather than before it.

A stack in PyTorch

import torch
from torch import nn

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 3),
)

x = torch.randn(16, 4)
logits = model(x)
print(logits.shape)                                  # (16, 3)
print(sum(p.numel() for p in model.parameters()))    # 67

Logits, probabilities and decisions

The last layer usually returns logits: raw scores, one per class, with no normalisation. They are not probabilities, and they do not need to be. The loss function applies the softmax internally. Softmax turns scores into a distribution that sums to one, and the fused form is numerically stable. The prediction is the argmax, which picks the largest score and does not care about normalisation at all.

Only at inference, when a human reads the output, do you turn logits into probabilities. Doing it inside the model is a common mistake. The loss then receives probabilities and applies a second softmax, which flattens the distribution and slows learning. Keep the model linear at the end, and let the loss own the transformation.

Temperature is the exception that proves the rule. Dividing logits by a temperature before the softmax sharpens or flattens the distribution. That is useful for calibration and for sampling. It is a post-hoc operation, though, not a layer the loss expects. If you add it, keep it out of the training graph unless you have a reason.

The forward pass is also a cache

Training keeps every intermediate activation. The backward pass needs it to compute the local gradient. That cache is why activation memory grows with depth and batch size. It is also why the largest batch a GPU can hold is often set by activations rather than by weights. Inference does not need the cache, and torch.no_grad() switches it off.

Initialisation is part of the forward pass too. Initialisation means the starting values of the weights. If every weight is drawn from a distribution whose variance ignores the layer width, the signal shrinks or explodes as it travels. He initialisation scales by the fan-in for ReLU, where fan-in is the number of inputs to a unit. Xavier scales for tanh. A network with a bad init can look broken before the first update.

A quick check is to print the mean and standard deviation of each layer’s output on a random batch. If they collapse toward zero by the last layer, or explode past a few units, the initialisation or the normalization is wrong. Fix that before training. It saves a day of tuning a model that could never have learned.