Neural Networks Explained: Why ReLU Beats Sigmoid & How to Build One in PyTorch

1. Neural Networks

Neural networks draw inspiration from the human brain's nervous system: each logistic regression unit corresponds to a neuron, and layers of such units connected together form a neural network . The sigmoid function is called an activation function because it crudely mimics neuron signal transmission.

Logistic regression is an effective classification model, but when features increase, performance does not improve and instead shows underfitting . Experts initially optimized this by relying on expert knowledge to abstract high-level features , keeping the feature count in a suitable range. However, this approach lacks generality and depends too heavily on domain expertise.

The natural next idea: let the model itself abstract high-level features , treating it as a black box. This is exactly what a neural network does.

1.1 Hidden Layer Activation Functions

For hidden-layer logistic regression units, the activation function introduces non-linearity, enabling neural networks to approximate arbitrary functions . While sigmoid suits the output layer for binary classification, it is problematic for hidden layers for two reasons:

Vanishing Gradient Problem : The sigmoid derivative peaks at 0.25. During backpropagation, gradients are multiplied layer by layer via the chain rule , causing exponential decay. Layers near the input receive near-zero gradients and fail to learn effectively.

Non-Zero Centered Outputs : Sigmoid outputs are always positive, so the next layer's inputs are all positive. By the chain rule, weight gradients depend entirely on the sign of the upstream gradient, forcing all weights to update in the same direction (all positive or all negative). This creates a zigzag optimization path and slows convergence.

The solution is ReLU (Rectified Linear Unit) [1], the de facto industrial standard. It offers high computational efficiency and solves both issues. Its only flaw is the Dying ReLU problem: when input is negative, the gradient is exactly zero, potentially killing a neuron permanently. This is easily fixed with Leaky ReLU / PReLU .

1.2 From Binary to Multi-Class Classification

Neural networks are overkill for binary classification alone; tasks like handwritten digit recognition require 10-class classification. Sigmoid is unsuitable here. Instead, Softmax [2] is used. The pre-activation outputs are called logits . Softmax computes probabilities while amplifying differences between logits, accelerating training. The corresponding loss is Cross Entropy Loss , which simplifies the binary cross-entropy formulation for multiple classes.

Formula: For N samples and C classes, loss = -1/N Σ Σ y_ic log(p_ic), where p_ic is the softmax output.

2. PyTorch Implementation

2.1 Module Definition

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(28 * 28, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10)
        )
    def forward(self, x):
        return self.model(x)

This defines a four-hidden-layer network (784→128→128→128→64→10) with ReLU activations. Notably, no output activation is specified because PyTorch's nn.CrossEntropyLoss internally applies Softmax.

2.2 Training Loop

learning_rate = 0.1
model = NeuralNetwork().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# Training process
model.train()
for epoch in range(num_epochs):
    total_loss = 0
    correct = 0
    total = 0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        # Compute training accuracy
        preds = torch.argmax(outputs, dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
    avg_loss = total_loss / len(train_loader)
    train_acc = 100 * correct / total
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {avg_loss:.4f}, Train Accuracy: {train_acc:.2f}%")

PyTorch encapsulates loss computation and gradient descent thoroughly. Parameter initialization (e.g., weights as zero-mean, std-dev random variables) is also handled automatically, but training data normalization and standardization must still be implemented manually .

Complete code is available at https://gitee.com/shichaoyuan/some-torch [3].

References: [1] https://en.wikipedia.org/wiki/Rectified_linear_unit [2] https://en.wikipedia.org/wiki/Softmax_function [3] https://gitee.com/shichaoyuan/some-torch