Neurons, layers and the MLP

Problem

Logistic regression and linear regression are limited by their features: they can only draw straight boundaries in the space you give them. For decades the fix was hand-made features: products of inputs, polynomials, edge detectors, word counts. Someone had to guess which features made the problem linear.

A neural network learns those features. Its early layers compute new features from the input, and its last layer is an ordinary linear or logistic model on top of them. Everything is trained together, by gradient descent, so the features are whatever makes the final layer's job easy. The multilayer perceptron (MLP) is the simplest network of this kind, and it is still everywhere: as the feed-forward block inside every Transformer layer, as the classification head on top of vision and language models, and as the standard model for tabular embeddings.

Intuition

One neuron takes a weighted sum of its inputs, adds a bias, and passes the result through a nonlinear activation: h = φ(w·x + b). With φ the sigmoid, that is logistic regression exactly. Modern networks mostly use the ReLU, φ(z) = max(0, z), which passes positive evidence through and blocks negative evidence.

A layer is many neurons reading the same input, each with its own weights. Stack the weight vectors as the rows of a matrix W and the whole layer is one matrix multiply, h = φ(Wx + b), followed by the activation applied element by element. A network is layers fed into layers. Each hidden layer re-describes the input in a new set of coordinates, and the output layer only has to draw a straight line in the last of them.

Left: a neuron is a dot product, a bias and a nonlinearity. Right: a layer is several neurons on the same input. Each arrow is one weight, and each layer's arrows together are its weight matrix.

Mechanics

The forward pass

For a batch X of shape (N, d₀), an MLP with layer widths d₀, d₁, …, d_L computes

H₀ = X,   H_ℓ = φ(H_ℓ−1W_ℓᵀ + b_ℓ) for ℓ = 1 … L−1,   output = H_L−1W_Lᵀ + b_L

where W_ℓ has shape (d_ℓ, d_ℓ−1), PyTorch's (out, in) convention, and b_ℓ is broadcast across the batch's rows (see tensors and shapes). The last layer has no activation: its outputs are logits or regression values, and the loss supplies the sigmoid or softmax, exactly as on the loss-functions page. "Depth" is the number of weight layers, L; "width" is the size of the hidden layers.

Why the nonlinearity is needed

Remove the activations and the network collapses. Two linear layers compose into one:

W₂(W₁x + b₁) + b₂ = (W₂W₁)x + (W₂b₁ + b₂)

which is again just Wx + b. However many layers you stack, a network without activations can only represent linear functions, so depth buys nothing. (It does change the training dynamics, and a product of matrices can have lower rank than any of its factors; but the set of functions is the same.)

The classic counterexample is XOR: output 1 when exactly one of two binary inputs is 1. No line separates {(0, 1), (1, 0)} from {(0, 0), (1, 1)}; the best least-squares linear fit predicts 0.50 for all four. Two ReLU neurons fix it:

h₁ = ReLU(x₁ + x₂),   h₂ = ReLU(x₁ + x₂ − 1),   y = h₁ − 2h₂

The hidden layer maps the four corners to (0, 0), (1, 0), (1, 0) and (2, 1). In those coordinates, the two positives land on the same point and a line separates them from the negatives. The ReLU did the work: h₂ is zero for three corners and switches on only for (1, 1), a bend no linear map can make.

The hidden layer moves the points so that a line works. Filled points are XOR's 1s and hollow points its 0s.

Universal approximation

How much can one hidden layer do? The universal approximation theorem (Cybenko 1989 for sigmoids; Hornik 1991 and Leshno et al. 1993 for any non-polynomial activation) says: any continuous function on a bounded region can be approximated to any accuracy by a network with one hidden layer, if the layer is wide enough.

For ReLU in one dimension the proof is a picture. Each unit ReLU(x − k) is a hinge that bends the output at x = k, so a sum of n hinges is a piecewise-linear curve with n bends. Put the bends at evenly spaced points, choose each output weight to be the change in slope there, and the network joins the dots of the target function:

y(x) = f(k₀) + Σ_i c_i ReLU(x − k_i),   c_i = slope_i − slope_i−1

More units, more bends, smaller error. What the theorem does not say matters as much. It doesn't say how wide "wide enough" is (in d dimensions this construction needs a number of units exponential in d); it doesn't say gradient descent will find the weights; and it says nothing about generalizing beyond the training data. In practice depth is what makes networks efficient: a deep network can reuse features, and some functions need exponentially more units when computed with one layer instead of several.

Each hidden unit contributes one bend. With 3 units the network is a crude zigzag; with 8 the error is 6.21 times smaller. For a smooth target, error falls roughly as 1/n² for n evenly spaced bends.

Counting parameters

A layer from width d_in to d_out has d_out × d_in weights and d_out biases. Activations have none. So an MLP with widths d₀, …, d_L has

Σ_ℓ (d_ℓ−1 + 1) · d_ℓ parameters

dominated by the widest adjacent pair. The cost of a forward pass is about two floating-point operations per weight per example (one multiply, one add), so parameters and compute go together for an MLP. That's not true of every network: a convolution reuses its few weights at every position.

How PyTorch builds one: nn.Module

Everything with parameters in PyTorch is an nn.Module. A module does two jobs. It holds state, and it defines forward. The state bookkeeping is automatic: Module.__setattr__ intercepts attribute assignment, so writing self.fc = nn.Linear(4, 3) inside __init__ registers a submodule, and assigning an nn.Parameter registers a parameter. That registry is what parameters(), state_dict(), .to(device), train()/eval() and zero_grad() walk, recursively through submodules. A plain tensor attribute is not registered and silently won't be trained, moved or saved; use register_buffer for non-trained state.

Calling a module, model(x), runs Module.__call__, which checks for hooks and otherwise goes straight to your forward. Call the module, not forward directly, or hooks are skipped.

nn.Linear(in_features, out_features) stores weight with shape (out, in) and bias with shape (out,), computes x @ weight.T + bias for any number of leading batch dimensions, and initializes both uniformly in ±1/√in_features. nn.Sequential(*modules) registers its arguments as submodules named "0", "1", … and its forward is a four-line loop that feeds each module's output to the next. It covers any straight-line stack; anything with branches or skip connections needs a Module subclass with its own forward, and nn.ModuleList to hold a variable number of layers.

Worked example

Implementation

import torch
from torch import nn

# XOR with the hand-set weights from this page
xor = nn.Sequential(nn.Linear(2, 2), nn.ReLU(), nn.Linear(2, 1))
with torch.no_grad():
    xor[0].weight.copy_(torch.tensor([[1., 1.], [1., 1.]]))
    xor[0].bias.copy_(torch.tensor([0., -1.]))
    xor[2].weight.copy_(torch.tensor([[1., -2.]]))
    xor[2].bias.zero_()
X = torch.tensor([[0., 0.], [0., 1.], [1., 0.], [1., 1.]])
xor(X).squeeze(1)                              # tensor([0., 1., 1., 0.])

# the MNIST-sized MLP, written as a Module subclass
class MLP(nn.Module):
    def __init__(self, sizes=(784, 256, 128, 10)):
        super().__init__()
        self.layers = nn.ModuleList(
            nn.Linear(a, b) for a, b in zip(sizes, sizes[1:])
        )

    def forward(self, x):                      # x: [batch, 784]
        for layer in self.layers[:-1]:
            x = torch.relu(layer(x))
        return self.layers[-1](x)              # logits: [batch, 10]

model = MLP()
sum(p.numel() for p in model.parameters())     # 235146
[name for name, _ in model.named_parameters()][:2]
# ['layers.0.weight', 'layers.0.bias']
model(torch.randn(32, 784)).shape              # torch.Size([32, 10])

The same forward pass in plain Python:

def relu(values):
    return [max(0, v) for v in values]

def linear(weight, bias, x):                   # weight: out rows of in columns
    return [sum(w * xi for w, xi in zip(row, x)) + b for row, b in zip(weight, bias)]

def mlp(layers, x):
    for i, (weight, bias) in enumerate(layers):
        x = linear(weight, bias, x)
        if i < len(layers) - 1:                  # no activation on the output layer
            x = relu(x)
    return x

layers = [([[1, 1], [1, 1]], [0, -1]), ([[1, -2]], [0])]
[mlp(layers, x)[0] for x in ([0, 0], [0, 1], [1, 0], [1, 1])]   # [0, 1, 1, 0]

Tradeoffs

• Width versus depth. One wide layer can approximate anything in principle, but deep, narrower networks usually reach the same accuracy with far fewer parameters, because later layers reuse features computed by earlier ones. Depth makes optimization harder: gradients have to pass through every layer, which is why deep networks need careful initialization, normalization and residual connections.
• No built-in structure. An MLP treats its input as an unordered vector. It knows nothing about neighbouring pixels or word order, so on images it needs far more parameters than a convolutional network, which shares weights across positions. Its strength is in places where the structure has already been handled: on top of learned embeddings, or as the per-token block inside a Transformer.
• Parameters cost memory and data. A dense layer between two wide vectors is expensive: the first layer of a 4096 → 4096 MLP alone has 16.8 million weights. More parameters need more data or more regularization to avoid overfitting.
• Non-convex training. Unlike the linear and logistic models, an MLP's loss has many minima and saddle points, and permuting hidden units gives identical networks. In practice gradient descent from a good random initialization finds good solutions, but there is no closed form and no guarantee.

Check yourself

Connections

References

• Rosenblatt. The Perceptron: A Probabilistic Model for Information Storage and Organization in the Brain. Psychological Review, 1958.
• Minsky and Papert. Perceptrons. MIT Press, 1969 (the XOR limitation of single-layer networks).
• Rumelhart, Hinton and Williams. Learning Representations by Back-Propagating Errors. Nature, 1986.
• Cybenko. Approximation by Superpositions of a Sigmoidal Function. Mathematics of Control, Signals and Systems, 1989.
• Hornik. Approximation Capabilities of Multilayer Feedforward Networks. Neural Networks, 1991.
• Leshno, Lin, Pinkus and Schocken. Multilayer Feedforward Networks with a Nonpolynomial Activation Function Can Approximate Any Function. Neural Networks, 1993.
• Telgarsky. Benefits of Depth in Neural Networks. COLT, 2016 (arXiv 1602.04485).
• Goodfellow, Bengio and Courville. Deep Learning, chapter 6. MIT Press, 2016.
• PyTorch v2.15.0: torch/nn/modules/module.py, torch/nn/modules/linear.py and torch/nn/modules/container.py.