Build a Multilayer Perceptron

A multilayer perceptron, usually shortened to MLP, is the simplest useful feed-forward neural network you will build in PyTorch. Its purpose is to map a fixed-size vector of input features to one or more outputs by stacking linear transformations and nonlinear activation functions. In this lesson the outcome is concrete: you should be able to design an MLP for tabular or flattened data, predict the tensor shapes through every layer, choose the right loss for the output, train it with autograd, and debug the common ways it breaks.

Within this PyTorch course, the MLP is the bridge between tensor arithmetic and trainable models. Convolutional and transformer models add stronger structure, but they still rely on the same ideas used here: parameters live in modules, the forward pass builds a computation graph, the loss turns predictions into a scalar objective, and an optimizer updates weights from gradients.

What an MLP Computes

An MLP is a composition of layers. A linear layer computes x @ W.T + b, where x is a batch of input rows, W is a learned weight matrix, and b is a learned bias vector. A nonlinearity such as ReLU, Tanh, or GELU is placed between linear layers so the whole model can represent curved decision boundaries instead of collapsing into a single linear function.

The first dimension of most supervised tensors is the batch dimension. If a batch has shape [32, 10], it contains 32 examples and each example has 10 features. A layer nn.Linear(10, 64) accepts that tensor and returns [32, 64]. The next layer must use 64 as its input feature count. The final layer size is determined by the task: one logit for binary classification with BCEWithLogitsLoss, one value for scalar regression with MSELoss, or one logit per class for multiclass classification with CrossEntropyLoss.

PyTorch Anatomy

In PyTorch, an MLP is commonly built with either nn.Sequential or a custom subclass of nn.Module. nn.Sequential is concise when data flows straight through each layer. A subclass is better when the forward pass branches, returns intermediate values, or needs named components for inspection. In both cases, assigning layers as module attributes registers their parameters so model.parameters() can hand them to an optimizer.

Piece	Role in an MLP
nn.Linear	Stores trainable weights and biases for an affine transform.
activation	Adds nonlinearity so stacked layers become more expressive.
loss_fn	Converts model outputs and targets into a scalar objective.
optimizer	Mutates parameters using gradients accumulated by autograd.
train() / eval()	Switches behavior for layers such as dropout and batch normalization.

The standard training step has a strict order. Compute logits, compute loss, clear old gradients with optimizer.zero_grad(), call loss.backward(), then call optimizer.step(). Clearing gradients matters because PyTorch accumulates gradients by default; this is useful for advanced accumulation workflows but surprising in a first MLP.

Example 1: Shapes and Parameters

The first example builds a small MLP that accepts four features and emits three class logits. The deterministic output is the shape information and parameter count. The parameter count is 4 * 6 + 6 for the first linear layer plus 6 * 3 + 3 for the second, giving 51 trainable scalar values.

import torch
from torch import nn

torch.manual_seed(7)

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

x = torch.randn(5, 4)
logits = model(x)
parameter_count = sum(p.numel() for p in model.parameters())

print("input shape:", tuple(x.shape))
print("logit shape:", tuple(logits.shape))
print("parameter count:", parameter_count)

This model does not apply softmax inside the network. For multiclass training, nn.CrossEntropyLoss expects raw logits and internally applies the numerically stable log-softmax operation. Adding softmax too early can make training less stable and is one of the most common beginner mistakes.

Example 2: Learning XOR

XOR is a compact demonstration of why hidden nonlinear layers matter. A single linear layer cannot separate the positive and negative cases because the decision boundary is not a straight line. A hidden layer with Tanh gives the model enough flexibility to assign high probability to [0, 1] and [1, 0] while assigning low probability to [0, 0] and [1, 1].

import torch
from torch import nn

torch.manual_seed(3)

x = torch.tensor([
    [0.0, 0.0],
    [0.0, 1.0],
    [1.0, 0.0],
    [1.0, 1.0],
])
y = torch.tensor([[0.0], [1.0], [1.0], [0.0]])

model = nn.Sequential(
    nn.Linear(2, 4),
    nn.Tanh(),
    nn.Linear(4, 1),
)
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

losses = []
for _ in range(250):
    logits = model(x)
    loss = loss_fn(logits, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    losses.append(loss.item())

with torch.no_grad():
    probabilities = torch.sigmoid(model(x))
    predictions = (probabilities >= 0.5).float()

print("first loss:", round(losses[0], 3))
print("last loss:", round(losses[-1], 3))
print("predictions:", predictions.squeeze().tolist())

The expected behavior is that the last loss is lower than the first and the predictions approach [0.0, 1.0, 1.0, 0.0]. Because the seed, data, architecture, and optimizer are fixed, the trend is reproducible on a normal PyTorch installation. This example also shows the binary-classification pairing: one output logit with BCEWithLogitsLoss and floating targets shaped like the logits.

Example 3: A Custom Module for Inference

A custom module is useful when the architecture deserves a name or will be reused. This example includes dropout, then switches the model to evaluation mode before inference. In training mode, dropout randomly zeros activations. In evaluation mode, dropout is disabled, so repeated inference calls are stable for the same input and parameters.

import torch
from torch import nn

class CustomerMLP(nn.Module):
    def __init__(self, in_features: int, hidden: int, classes: int) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, hidden),
            nn.ReLU(),
            nn.Dropout(p=0.25),
            nn.Linear(hidden, classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

model = CustomerMLP(in_features=8, hidden=12, classes=2)
model.eval()

batch = torch.randn(4, 8)
with torch.no_grad():
    logits = model(batch)
    probabilities = torch.softmax(logits, dim=1)

print("training flag:", model.training)
print("probability row sums:", torch.round(probabilities.sum(dim=1), decimals=4).tolist())

The row sums of the softmax probabilities are 1.0, up to floating-point rounding, because softmax(dim=1) normalizes across the class dimension for each example. The call to torch.no_grad() avoids building an autograd graph during inference, reducing memory use and preventing accidental gradient work.

Design Choices and Trade-Offs

Depth, width, activation choice, and regularization control the behavior of an MLP. A wider hidden layer can fit more complex interactions, but it increases memory, compute, and overfitting risk. More hidden layers can represent hierarchical combinations of features, but deep plain MLPs can become harder to optimize without normalization, residual connections, or careful initialization. ReLU is a strong default because it is cheap and usually optimizes well, but inactive ReLU units can stop learning for inputs that keep them below zero. Tanh can be helpful for tiny centered problems like XOR but saturates at large magnitudes.

For tabular data, feature scaling is often as important as architecture. If one input is measured in cents and another is a 0-to-1 ratio, the larger numerical scale can dominate early optimization. Standardizing continuous features and encoding categorical features consistently between training and inference makes the optimization problem better conditioned.

The output layer should match the target and loss. Use raw logits for classification losses, not probabilities. Use float targets for regression and binary logistic losses, and long class indices for CrossEntropyLoss. A mismatch may produce a clear runtime error, or worse, a model that trains while learning the wrong objective.

Failure Modes and Troubleshooting

Shape mismatch. The symptom is an error such as matrix shapes cannot be multiplied. The cause is usually that a linear layer’s in_features does not match the previous tensor’s last dimension. Print or assert x.shape after each transformation, then correct the next nn.Linear input size. Do not guess from the dataset column count after preprocessing; inspect the actual tensor.

Loss does not decrease. Symptoms include a flat loss curve and predictions near a constant value. Causes include an excessive learning rate, unscaled features, missing nonlinearities, wrong target dtype, or calling optimizer.step() before backward(). Diagnose by overfitting a tiny batch, checking gradient norms, and confirming that at least one parameter changes after a step. Correct by scaling inputs, reducing the learning rate, pairing output and loss correctly, and verifying the training-step order.

Training accuracy is high but validation accuracy is poor. This points to overfitting or data leakage. Inspect the split procedure, remove duplicated examples across splits, and compare performance by meaningful slices. Correct with a smaller model, weight decay, dropout, more data, early stopping, or a leakage-safe validation split.

Inference changes between calls. If the same input returns different outputs, the model may still be in training mode with dropout enabled. Call model.eval() before inference and use torch.no_grad(). If randomness is still expected, isolate it and document it; a plain MLP classifier should usually be deterministic for fixed weights and inputs.

Performance and Reliability

MLPs are fast compared with many deep architectures, but dense layers scale with batch_size * in_features * out_features. Large hidden widths can dominate both training time and memory. Batching improves hardware utilization, but very large batches can reduce generalization and exceed memory. Reliability comes from saving the model state, the preprocessing steps, and the class mapping together. A correct state_dict is not enough if inference uses a different feature order or category encoding than training.

Security concerns are mostly about data handling and model loading. Avoid untrusted serialized model objects. Prefer loading a known state_dict into source-controlled model code, and treat training data, labels, and logs as potentially sensitive. For user-facing inference, validate feature counts, numeric ranges, and missing values before constructing tensors.

Hands-On Lab: Train a Two-Class MLP

Prerequisites: a Python environment with PyTorch installed, basic command-line access, and permission to run a local script. The lab creates synthetic two-dimensional data, so it does not require downloading a dataset.

1. Create a new scratch Python file or notebook cell for the lab code.
2. Paste the code below and run it once without modification.
3. Verify that it prints a test accuracy of at least 0.95; the assertion enforces this threshold.
4. Change the hidden width from 16 to 2 and rerun. Note whether optimization becomes less consistent.
5. Set the learning rate to 3.0 and rerun. If the loss becomes unstable, restore 0.03.
6. Cleanup is simply deleting the scratch file or notebook cell; no persistent model artifact is written.

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(11)

class BlobMLP(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(2, 16),
            nn.ReLU(),
            nn.Linear(16, 16),
            nn.ReLU(),
            nn.Linear(16, 2),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)

class0 = torch.randn(120, 2) * 0.45 + torch.tensor([-1.2, -1.0])
class1 = torch.randn(120, 2) * 0.45 + torch.tensor([1.1, 1.0])
features = torch.cat([class0, class1]).float()
targets = torch.cat([
    torch.zeros(120, dtype=torch.long),
    torch.ones(120, dtype=torch.long),
])

permutation = torch.randperm(features.shape[0])
features = features[permutation]
targets = targets[permutation]
train_x, test_x = features[:180], features[180:]
train_y, test_y = targets[:180], targets[180:]
loader = DataLoader(TensorDataset(train_x, train_y), batch_size=32, shuffle=True)

model = BlobMLP()
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.03)

for epoch in range(40):
    model.train()
    for xb, yb in loader:
        logits = model(xb)
        loss = loss_fn(logits, yb)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

model.eval()
with torch.no_grad():
    test_logits = model(test_x)
    predictions = test_logits.argmax(dim=1)
    accuracy = (predictions == test_y).float().mean().item()

print("test accuracy:", round(accuracy, 3))
assert accuracy >= 0.95

Verification should include more than the final number. Confirm that train_x has two columns, test_logits has two columns, and predictions contains class indices rather than probabilities. Those checks prove the data, final layer, and loss function agree.

Assessment Exercises

1. An MLP receives tensors of shape [64, 20]. You add nn.Linear(20, 32), nn.ReLU(), and nn.Linear(16, 3). What fails, and how would you repair it?
2. Why should a multiclass MLP trained with CrossEntropyLoss return logits instead of applying softmax in its forward method?
3. You can perfectly fit 12 training examples but fail on validation data. Name two model changes and two data-splitting checks that would make your diagnosis more credible.
4. Modify the lab to perform binary classification with one output logit. Which loss, target dtype, and prediction threshold would you use?
5. Explain why forgetting optimizer.zero_grad() changes training even though the code still runs.

Summary

A PyTorch MLP is a stack of registered modules that transforms batches of feature vectors into logits or values. The important mechanics are tensor shape flow, affine parameters, nonlinear activations, loss selection, gradient accumulation, optimizer updates, and mode changes for training versus inference. Build the smallest architecture that matches the task, scale and encode features consistently, pair outputs with the correct loss, and debug from observable tensors rather than intuition.