Does the output layer in a deep neural network need an activation function?

Question

I have enrolled in a course that uses only one hidden layer, and that is the only layer that has activation functions. The model can be visualized as follows:

and here is a PyTorch implementation:

class MnistModel(nn.Module):
    """Feedfoward neural network with 1 hidden layer"""
    def __init__(self, in_size, hidden_size, out_size):
        super().__init__()
        # hidden layer
        self.linear1 = nn.Linear(in_size, hidden_size)
        # output layer
        self.linear2 = nn.Linear(hidden_size, out_size)
        
    def forward(self, xb):
        # Flatten the image tensors
        xb = xb.view(xb.size(0), -1)
        # Get intermediate outputs using hidden layer
        out = self.linear1(xb)
        # Apply activation function
        out = F.relu(out)
        # Get predictions using output layer
        out = self.linear2(out)
        return out

shouldn't the output layer also have activation functions?

Answer 1

A neural network layer with no activation function is the same as "linear activation" i.e. $f(x) = x$

This is often used for the output layer in regression problems, where a constrained output like that of sigmoid, hyperbolic tangent or ReLU may not be appropriate.

For an output layer, this is fine, and does not conflict with any theory behind neural networks. It will often be combined with a mean squared error (MSE) loss function that makes the gradient calculation simple at the output layer.

For hidden layers, skipping the activation function can be a problem, since a purely linear layer in the middle of a multi-layer network is redundant - it could be replaced or even removed with no impact. That is because two directly connected linear layers are functionally equivalent to a single linear layer with different parameters, and every hidden layer consists of a linear component plus an activation function. So even one missing activation function on a hidden layer directly connects two linear sub-components, making one of them redundant.

In the case of a classifier, some libraries (notably including PyTorch which you are using) require you to have two variants of your neural network - a trainable version without a final sigmoid or softmax layer, and the "full" version which adds sigmoid or softmax on top (usually just a function that calls the training network and add this last activation). This is done to allow you to use more stable gradient calculations that are based on the logits, the values before applying an activation function.

Answer 2

I'd like to add more information to Neal's answer.

A network implemented in the example does not include activation for the output, because it will be applied during the training:

# ...
def training_step(self, batch):
    images, labels = batch 
    out = self(images)
    # apply activation and calculate loss
    loss = F.cross_entropy(out, labels)
    return loss
# ...

The problem considered in the example is multiclass classification, which is usually solved with softmax activation and cross-entropy loss. F.cross_entropy combines log_softmax and nll_loss in a single function, which is numerically more stable as softmax and NLL loss. See this discussion for more details, and here is an article about different implementations.