Speeding Up MLP Training in PyTorch and TensorFlow

Framework overview and optimization techniques

Overview

Last time, we trained a neural network by hand with pen and paper. Imagine running forward and backward passes 500 times with the same approach. Gross! Now it’s time to hand that work off to a software framework. Naturally, the data we’re using will be more complex, since we want to showcase what these tools can actually do.

For the XOR gate, we had 9 parameters, made up of 6 weights and 3 biases. This time, both the data volume and the network grow many times over. We’ll use 1,000 data points instead of 4, and the network itself will have 337 parameters. Despite the larger scale, the core functionality still boils down to the same three critical pieces: input, weights, and loss function.

That basic two-layer perceptron example ended with five open questions: stochastic gradient descent, mini-batches, overfitting, dropout, and early stopping. In a nutshell, they’re all techniques that make training more accurate and more efficient. They’ll become far more intuitive as we work through the example.

The Framework

Being an expert programmer isn’t a prerequisite, but some foundational Python skills are required. A good gut check is whether you can build a simple to-do list.

NOTE: There are plenty of free resources to help build that foundation first, such as this Khan Academy series.

It might be obvious at this point that PyTorch (from Meta) and TensorFlow (from Google) are the two frameworks used everywhere, from large-scale LLMs to ML models that run on tiny devices. Under the hood, the core workflow is the same, and it consists of the following pieces:

• Tensors – Matrices that know how to run on a GPU. A rank-2 tensor, for instance, is just a 2-dimensional matrix.
• Autograd – Automatic differentiation, so nobody has to compute \(\frac{\partial L}{\partial w}\) by hand.
• Prebuilt layers, losses, and optimizers – The pieces required to actually train a model.

One thing worth mentioning is that, given a similar skill level, creativity is what differentiates an exceptional ML practitioner from an average one.

Since there are a plethora of articles out there, including on the official PyTorch and TensorFlow websites, about using these frameworks, we’ll keep the rest relatively brief.

Understand the Data

The two-moons dataset looks intimidating when plotted out, but it’s just a slightly more complex version of the XOR example. XOR has simple binary inputs (1 or 0). In the case of two-moons, the inputs cover a continuous range. For example, a point at \(x_1 = 0.5\) and \(x_2 = -0.5\) would be classified as 1 in this case.

The two-moons dataset

In the XOR example, it was impossible to draw the prediction boundary with a single line (one neuron). How many neurons and hidden layers will it take to make an accurate prediction here? By looking at the plot, you might estimate the S-curve needs 8-10 straight-line segments to trace it. However, as data gets bigger, it becomes impossible to plot and determine architecture by sight. Since there’s no set formula for calculating the architecture, we could study similar examples, then set up logic for rapid trial-and-error to home in on the optimal parameters.

scikit-learn is a classical machine learning library. The term “classical” refers to non-neural-net techniques, such as linear regression and k-nearest neighbors (kNN). We use it to generate training data:

from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split

X, y = make_moons(n_samples=1000, noise=0.20, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(
    X, y, test_size=0.2, random_state=42
)

Implementation

We generated 1,000 samples: 800 for the training set and 200 for the validation set. When determining accuracy, it’s important to use data the model has never seen before. This will help us determine whether the model can make generalized predictions in the real world.

Step 1 - Design the Network

Earlier we mentioned that there will be 337 parameters:

• 2 inputs times 16 neurons plus 16 biases for the first hidden layer: \(2 \times 16 + 16 = 48\)
• 16 inputs (the outputs of the first hidden layer) times 16 neurons plus 16 biases: \(16 \times 16 + 16 = 272\)
• 16 inputs times 1 neuron plus 1 bias: \(16 \times 1 + 1 = 17\)
• Total parameters: \(48 + 272 + 17 = 337\)

For example, this matrix operation will create 16 outputs for the first hidden layer: \[ \underbrace{\begin{bmatrix} x_1 & x_2 \end{bmatrix}}_{1 \times 2} \cdot \underbrace{\begin{bmatrix} w_{1,1} & w_{1,2} & \cdots & w_{1,16} \\ w_{2,1} & w_{2,2} & \cdots & w_{2,16} \end{bmatrix}}_{2 \times 16} + \underbrace{\begin{bmatrix} b_1 & b_2 & \cdots & b_{16} \end{bmatrix}}_{1 \times 16} \]

Each column of \(x\) is an input, each column of \(w\) corresponds to a neuron, and each column of \(b\) is a bias.

Step 2 - Activation Function

For the hidden layers, we’ll use a different activation function called ReLU:

\[ \mathrm{ReLU}(z) = \max(0, z) \]

It preserves positive values and outputs zero for negative ones, giving it a derivative of 0 for negative input and 1 for positive input.

We are using ReLU for the sake of expanding our tool belt, but it also helps with an issue called vanishing gradients that shows up when there are many hidden layers.

Imagine we have a 10-layer perceptron. Part of backpropagation contains:

\[ \frac{\partial a_{10}}{\partial z_{10}} \frac{\partial z_{10}}{\partial a_{9}} \dotsm \frac{\partial a_{1}}{\partial z_{1}} \]

If the activations sit in the ballpark of 0.95, the two functions give us:

\[ \text{sigmoid} = a_{10}(1-a_{10}) W_{10} \dotsm a_{1}(1-a_{1}) \quad \text{vs.} \quad \mathrm{ReLU} = 1 \cdot W_{10} \dotsm 1 \]

Sigmoid multiplies one small factor by another small factor: \([0.95(1-0.95)]^{10} = 0.0475^{10} \approx 5.8 \cdot 10^{-14}\). The gradient reaching the earliest weights becomes vanishingly small, so those weights barely change each step according to the gradient descent formula (\(w \leftarrow w - \eta \frac{\partial L}{\partial w}\)). ReLU avoids this source of shrinkage because the derivative is just 1.

We’ll encounter this issue again in more detail when we explore RNNs and LSTMs/GRUs.

Putting It in Sequence

Network flow should be clear at this point:

input \(\times\) weights → ReLU → input \(\times\) weights → ReLU → input \(\times\) weights → \(\sigma\)

We already created data and separated them into training and validation set:

• Training set: X_train and y_train
• Validation set: X_val and y_val

Now let’s put them in sequence in PyTorch and TensorFlow side by side, because they’re quite similar.

1. Library Import

This should be pretty self explanatory. If we want to use a tool for a job, we have to make sure it’s in our toolbox.

PyTorch

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

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers

Now it’s time to define model architecture, which include setup for weights, hidden layers, activation function, and loss function.

2. Define Model Architecture

We define each layer and it’s output by putitng them in a sequence.

PyTorch

model = nn.Sequential(
    nn.Linear(2, 16),         # input layer -> hidden layer 1
    nn.ReLU(),
    nn.Linear(16, 16),        # hidden layer 1 -> hidden layer 2
    nn.ReLU(),
    nn.Linear(16, 1),         # hidden layer 2 -> output neuron
    nn.Sigmoid(),             # squash to a probability, same as part one
)

loss_fn   = nn.BCELoss()      # binary cross-entropy
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

TensorFlow

model = tf.keras.Sequential([
    layers.Input(shape=(2,)),
    layers.Dense(16, activation="relu"),
    layers.Dense(16, activation="relu"),
    layers.Dense(1, activation="sigmoid"),
])

model.compile(
    optimizer=tf.keras.optimizers.SGD(learning_rate=0.5),
    loss="binary_crossentropy",
    metrics=["accuracy"],
)

One thing that might stand out is optimizer. .SGD in the method name stands for Stochastic Gradient Descent, whose benefit we’ll see when we get to the batch_size section, but all this does for now is just set the learning rate.

3. Use the Correct Data Structure

In PyTorch, data must be converted into tensors, the framework’s core data structure. Tensors support automatic differentiation and can run on a GPU. It’s a mandatory component for training.

PyTorch

# NumPy arrays -> float32 tensors, labels shaped as a column vector
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).unsqueeze(1)
X_val   = torch.tensor(X_val,   dtype=torch.float32)
y_val   = torch.tensor(y_val,   dtype=torch.float32).unsqueeze(1)

On the other hand TensorFlow acceps numpy or python arrays.

4. Training the Model

PyTorch

train_loader = DataLoader(
    TensorDataset(X_train, y_train), batch_size=32, shuffle=True
)

for epoch in range(100):
    model.train()
    for xb, yb in train_loader:          # one mini-batch at a time
        optimizer.zero_grad()            # clear old gradients
        y_hat = model(xb)                # forward pass
        loss = loss_fn(y_hat, yb)        # how wrong are we?
        loss.backward()                  # backpropagation
        optimizer.step()                 # w <- w - lr * dL/dw

    if (epoch + 1) % 20 == 0:
        model.eval()
        with torch.no_grad():
            val_pred = model(X_val)
            val_loss = loss_fn(val_pred, y_val)
            val_acc = ((val_pred > 0.5).float() == y_val).float().mean()
        print(f"epoch {epoch+1:3d}  val_loss {val_loss:.4f}  val_acc {val_acc:.3f}")

TensorFlow

history = model.fit(
    X_train, y_train,
    batch_size=32,
    epochs=100,
    validation_data=(X_val, y_val),
    verbose=0,               # silence the per-epoch progress bar
)

val_loss, val_acc = model.evaluate(X_val, y_val, verbose=0)
print(f"val_loss {val_loss:.4f}  val_acc {val_acc:.3f}")

batch_size controls how many samples are used to compute each gradient update. Given 800 samples, one epoch (one full pass through the data) can be organized three ways:

a. Pure SGD (batch_size=1) — 800 updates per epoch, 1 sample each

• Shuffle the 800 samples.
• Take sample #1 → compute its gradient → update all 337 parameters.
• Take sample #2 → compute its gradient → update all 337 parameters.
• … repeat for all 800 samples.

b. Mini-batch SGD (batch_size=32) — 25 updates per epoch, 32 samples each

• Shuffle the 800 samples, split into 25 batches of 32.
• Batch 1 (samples 1–32) → average their 32 gradients → update all 337 parameters.
• Batch 2 (samples 33–64) → average → update all 337 parameters.
• … repeat for all 25 batches.

c. Full batch (batch_size=800) — 1 update per epoch, all samples

• Take all 800 samples → average their 800 gradients → update all 337 parameters.

Pure SGD makes the most updates per epoch, but each update is noisy, because it relies on one sample, and processing samples one at a time leaves parallel hardware underused. Full batch produces the most accurate gradient per update, but only one update per epoch, so it typically needs many more epochs to converge. It also requires holding activations for the entire dataset in memory during backprop, which can exceed GPU memory and crash training on large datasets. Mini-batch strikes a balance in between. It has many update per epoch, and the batches are large enough to use hardware efficiently.

Instead of computing the exact gradient using the entire dataset, which could be millions of examples, Stochastic Gradient Descent (SDG) estimate it using small batch chosen at random to update parameters.

Verify Effectiveness

We posed a question earlier: how many lines (neurons) would we need to make a reasonable prediction? After running the code and validating accuracy, we can plot the prediction boundary to visualize the result.

Learned decision boundary

That S-shaped curve shows us that two layers of ReLU neurons bend one boundary through the gap between the crescents. The few misclassified points sit where the noise genuinely mixes the classes; the error stems from messiness in the data, not something the model should compensate for.

Regularization and Optimization

Now we’re equipped with the tools and knowledge to design and run a model with less than 100 lines of code. That’s quite empowering. In practice, even given clean data, overfitting might be one of the biggest issues jeopardizing a model’s accuracy.

Intuitively, it might make sense to train as many epochs as possible to improve accuracy, but at some point, you might notice that training loss becomes stagnant and validation loss explodes. It’s often an indication that the model can only discern the training data, and has lost the ability to generalize outside of the training data.

Training vs validation loss, with and without dropout

Since this problem has been around for a long time, talented researchers have already come up with a few methods to overcome this issue.

Dropout

Dropout attacks memorization by sabotage. During training, every pass through a dropout layer switches each neuron off with probability \(p\):

PyTorch

nn.Sequential(
    nn.Linear(2, 256),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(256, 256),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(256, 1),
    nn.Sigmoid(),
)

TensorFlow

tf.keras.Sequential([
    layers.Dense(256, activation="relu"),
    layers.Dropout(0.5),
    layers.Dense(256, activation="relu"),
    layers.Dropout(0.5),
    layers.Dense(1, activation="sigmoid"),
])

The mechanism is that each activation has about a 50% (\(p=0.5\)) chance to get dropped to zero, and the remaining activations get scaled by \(1/(1-p)\) to compensate for the dropped neurons.

This prevents a specific set of neurons from learning the majority of the task. Since activations get dropped randomly, learning can be better distributed across different neurons.

NOTE: Dropout is used during training and is inactive during inference (prediction).

Early Stopping

In a nutshell, the idea is to stop training before overfitting sets in.

PyTorch

best_val, patience, wait = float("inf"), 20, 0

for epoch in range(400):
    train_one_epoch(model, train_loader)
    val_loss = evaluate(model, X_val, y_val)

    if val_loss < best_val:
        best_val, wait = val_loss, 0
        torch.save(model.state_dict(), "best.pt")
    else:
        wait += 1
        if wait >= patience:
            break

model.load_state_dict(torch.load("best.pt"))

TensorFlow

early_stop = tf.keras.callbacks.EarlyStopping(
    monitor="val_loss", patience=20, restore_best_weights=True
)

model.fit(X_train, y_train, epochs=400,
          validation_data=(X_val, y_val), callbacks=[early_stop])

Dynamic Learning Rate

We used optimizer = torch.optim.SGD(model.parameters(), lr=0.5) earlier. The issue with a fixed learning rate is that we’re stuck guessing: too small and convergence is slow, too big and the parameters hop around the minimum instead of settling. In practice, this is solved from two directions. Adaptive optimizers like Adam/AdamW maintain running statistics of each parameter’s gradients and use them to scale each parameter’s effective step size individually. A learning rate scheduler complements this by adjusting the global base rate over the course of training. For example, we might want to reduce the step size as we approach convergence.

PyTorch

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    factor=0.5,      # halve the lr when triggered
    patience=10,     # wait 10 epochs of no val_loss improvement first
)

Additional Resources

• Gradient Descent, How Neural Networks Learn - 3Blue1Brown
• Learn the Basics - Official PyTorch Tutorial
• TensorFlow 2 Quickstart for Beginners
• Dropout: A Simple Way to Prevent Neural Networks from Overfitting - Srivastava et al.
• A Recipe for Training Neural Networks - Andrej Karpathy