4.2  Implementation of Multilayer Perceptrons

Multilayer perceptrons (MLPs) are not much more complex to implement than simple linear models. The key conceptual difference is that we now concatenate multiple layers.

from d2l import torch as d2l
import torch
from torch import nn
from d2l import tensorflow as d2l
import tensorflow as tf
from d2l import jax as d2l
from flax import nnx
from jax import numpy as jnp
from d2l import mxnet as d2l
from mxnet import np, npx
from mxnet.gluon import nn
npx.set_np()

4.2.1 Implementation from Scratch

Let’s begin again by implementing such a network from scratch.

4.2.1.1 Initializing Model Parameters

Recall that Fashion-MNIST contains 10 classes, and that each image consists of a \(28 \times 28 = 784\) grid of grayscale pixel values. As before we will disregard the spatial structure among the pixels for now, so we can think of this as a classification dataset with 784 input features and 10 classes. To begin, we will implement an MLP with one hidden layer and 256 hidden units (Figure 4.2.1). Both the number of layers and their width are adjustable (they are considered hyperparameters). Typically, we choose each layer width to be a power of 2, which is efficient for hardware memory addressing.

Figure 4.2.1: The two-layer MLP of this section: a batched input is flattened to 784 features, mapped by an affine layer to a 256-dimensional hidden representation, passed through a ReLU, then mapped by a second affine layer to 10 logits.

Again, we will represent our parameters with several tensors. Note that for every layer, we must keep track of one weight matrix and one bias vector. As always, we allocate memory for the gradients of the loss with respect to these parameters. We use small Gaussian noise (\(\sigma = 0.01\)) as a simple starting point; principled strategies for choosing this scale are the subject of Section 4.4.

In the code below we use nn.Parameter to automatically register a class attribute as a parameter to be tracked by autograd (Section 1.5).

In the code below we use tf.Variable to define the model parameter.

In the code below, nnx.Param marks each array as a trainable parameter.

In the code below, we first define and initialize the parameters and then enable gradient tracking.

class MLPScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, num_hiddens, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        self.W1 = nn.Parameter(torch.randn(num_inputs, num_hiddens) * sigma)
        self.b1 = nn.Parameter(torch.zeros(num_hiddens))
        self.W2 = nn.Parameter(torch.randn(num_hiddens, num_outputs) * sigma)
        self.b2 = nn.Parameter(torch.zeros(num_outputs))
class MLPScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, num_hiddens, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        self.W1 = tf.Variable(
            tf.random.normal((num_inputs, num_hiddens)) * sigma)
        self.b1 = tf.Variable(tf.zeros(num_hiddens))
        self.W2 = tf.Variable(
            tf.random.normal((num_hiddens, num_outputs)) * sigma)
        self.b2 = tf.Variable(tf.zeros(num_outputs))
class MLPScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, num_hiddens, lr,
                 sigma=0.01, rngs=None):
        super().__init__()
        self.save_hyperparameters(ignore=['rngs'])
        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
        self.W1 = nnx.Param(
            rngs.params.normal((num_inputs, num_hiddens)) * sigma)
        self.b1 = nnx.Param(jnp.zeros(num_hiddens))
        self.W2 = nnx.Param(
            rngs.params.normal((num_hiddens, num_outputs)) * sigma)
        self.b2 = nnx.Param(jnp.zeros(num_outputs))
class MLPScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, num_hiddens, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        self.W1 = np.random.randn(num_inputs, num_hiddens) * sigma
        self.b1 = np.zeros(num_hiddens)
        self.W2 = np.random.randn(num_hiddens, num_outputs) * sigma
        self.b2 = np.zeros(num_outputs)
        for param in self.get_scratch_params():
            param.attach_grad()

4.2.1.2 Model

To make sure we know how everything works, we will implement the ReLU activation ourselves rather than invoking the built-in relu function directly.

def relu(X):
    return torch.maximum(X, torch.zeros_like(X))
def relu(X):
    return tf.math.maximum(X, 0)
def relu(X):
    return jnp.maximum(X, 0)
def relu(X):
    return np.maximum(X, 0)

Since we are disregarding spatial structure, we reshape each two-dimensional image into a flat vector of length num_inputs. Finally, we implement our model with just a few lines of code. Since autograd handles the backward pass, this is all that it takes.

@d2l.add_to_class(MLPScratch)
def forward(self, X):
    X = d2l.reshape(X, (-1, self.num_inputs))
    H = relu(d2l.matmul(X, self.W1) + self.b1)
    return d2l.matmul(H, self.W2) + self.b2

4.2.1.3 Training

Fortunately, the training loop for MLPs is exactly the same as for softmax regression. We define the model, data, and trainer, then finally invoke the fit method on model and data.

model = MLPScratch(num_inputs=784, num_outputs=10, num_hiddens=256, lr=0.1)
data = d2l.FashionMNIST(batch_size=256)
trainer = d2l.Trainer(max_epochs=30)
trainer.fit(model, data)

You should see the validation accuracy settle at typically around \(0.87\), a modest improvement over the softmax regression baseline of Section 3.4 on the same data, bought by the hidden layer and its ReLU.

4.2.2 Concise Implementation

As you might expect, by relying on the high-level APIs, we can implement MLPs even more concisely.

4.2.2.1 Model

Compared with our concise implementation of softmax regression (Section 3.5), the only difference is that we add two fully connected layers where we previously added only one. The first is the hidden layer, the second is the output layer.

class MLP(d2l.Classifier):
    def __init__(self, num_outputs, num_hiddens, lr):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Sequential(nn.Flatten(), nn.LazyLinear(num_hiddens),
                                 nn.ReLU(), nn.LazyLinear(num_outputs))
class MLP(d2l.Classifier):
    def __init__(self, num_outputs, num_hiddens, lr):
        super().__init__()
        self.save_hyperparameters()
        self.net = tf.keras.models.Sequential([
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(num_hiddens, activation='relu'),
            tf.keras.layers.Dense(num_outputs)])
class MLP(d2l.Classifier):
    def __init__(self, num_outputs, num_hiddens, lr, num_inputs=784,
                 rngs=None):
        super().__init__()
        self.save_hyperparameters(ignore=['rngs'])
        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
        self.hidden = nnx.Linear(num_inputs, num_hiddens, rngs=rngs)
        self.output = nnx.Linear(num_hiddens, num_outputs, rngs=rngs)

    def forward(self, X):
        X = X.reshape((X.shape[0], -1))  # Flatten
        return self.output(nnx.relu(self.hidden(X)))
class MLP(d2l.Classifier):
    def __init__(self, num_outputs, num_hiddens, lr):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Sequential()
        self.net.add(nn.Dense(num_hiddens, activation='relu'),
                     nn.Dense(num_outputs))
        self.net.initialize()

Previously, we defined forward methods for models to transform input using the model parameters. These operations are essentially a pipeline: you take an input and apply a transformation (e.g., matrix multiplication with weights followed by bias addition), then repetitively use the output of the current transformation as input to the next transformation. However, you may have noticed that no forward method is defined here. In fact, MLP inherits the forward method from the Module class (Section 2.2.2) to simply invoke self.net(X) (X is input), which is now defined as a sequence of transformations via the Sequential class. The Sequential class abstracts the forward process enabling us to focus on the transformations. We will further discuss how the Sequential class works in Section 5.1.2.

4.2.2.2 Training

The training loop is exactly the same as when we implemented softmax regression. This modularity enables us to separate matters concerning the model architecture from orthogonal considerations. Note that while the two versions compute the same architecture, they do not start from the same parameters: the scratch model draws its weights from \(\mathcal{N}(0, 0.01^2)\), whereas the concise version uses the library’s default initializer. Defaults differ across libraries and layer types, so the two versions need not start at comparable scales. Their training trajectories and final accuracies can therefore differ.

model = MLP(num_outputs=10, num_hiddens=256, lr=0.1)
trainer.fit(model, data)

4.2.3 Summary

We have now built and trained a working multilayer perceptron, a network with a hidden layer and a nonlinearity, in both a from-scratch and a concise form. The from-scratch version makes the new ingredients concrete: two weight matrices, two bias vectors, a hand-rolled ReLU, and a two-step forward computation. The concise version shows that nn.Sequential collapses all of that into a four-element stack. The training loop, the loss function, and the data loader are unchanged from softmax regression, a benefit of the modular design.

The from-scratch version also exposes why we reach for the high-level API: naming and tracking parameters by hand quickly becomes awkward. Imagine inserting another layer between layers 42 and 43; we would be stuck renaming or improvising a “layer 42b”. nn.Sequential removes both problems at once.

Three questions remain open, and each is the subject of one of the next sections:

  • How do gradients flow through this stack, and what can go wrong as it gets deeper? (Section 4.3, Section 4.4)
  • Why does such a flexible model generalize to unseen data at all? (Section 4.5)
  • How can we regularize it to generalize better? (Section 4.6)

Answering them turns this small working model into a reliable building block.

4.2.4 Exercises

  1. Change the number of hidden units num_hiddens and plot how its number affects the accuracy of the model. What is the best value of this hyperparameter?
  2. Try adding a hidden layer to see how it affects the results. In particular, add one to the from-scratch model while keeping its \(\sigma = 0.01\) Gaussian initialization: you may find that the deeper network trains worse. Why that happens, and what to do about it, is the subject of Section 4.4.
  3. Why is it a bad idea to insert a hidden layer with a single neuron? What could go wrong?
  4. How does changing the learning rate alter your results? With all other parameters fixed, which learning rate gives you the best results? How does this relate to the number of epochs?
  5. Let’s optimize over all hyperparameters jointly, i.e., learning rate, number of epochs, number of hidden layers, and number of hidden units per layer.
    1. What is the best result you can get by optimizing over all of them?
    2. Why it is much more challenging to deal with multiple hyperparameters?
    3. Describe an efficient strategy for optimizing over multiple parameters jointly.
  6. Compare the speed of the framework and the from-scratch implementation for a challenging problem. How does it change with the complexity of the network?
  7. Measure the speed of tensor–matrix multiplications for well-aligned and misaligned matrices. For instance, test for matrices with dimension 1024, 1025, 1026, 1028, and 1032.
    1. How does this change between GPUs and CPUs?
    2. Determine the memory bus width of your CPU and GPU.
  8. Try out different activation functions. Which one works best on Fashion-MNIST? Compare at least ReLU, tanh, sigmoid, and GELU (torch.nn.functional.gelu in PyTorch, jax.nn.gelu in JAX). (Hint: for sigmoid and tanh you may need to retune the learning rate.) GELU is used in BERT and GPT-2-style Transformers, while many recent language models use gated SiLU/SwiGLU blocks. Does this small image task provide enough evidence to choose among them for a Transformer?
  9. Compare the effect of three initialization scales on training: (a) small Gaussian noise with \(\sigma = 0.001\); (b) the value used in this section, \(\sigma = 0.01\); (c) large Gaussian noise with \(\sigma = 0.1\). Plot the training and validation curves for each. Why does \(\sigma\) matter? (Hint: consider what happens to the activations on the very first forward pass.) The principled answer is developed in Section 4.4.