Output Layer Activation Functions: Softmax and Linear in MLPs

Introduction

Welcome back to our course on The MLP Architecture: Activations & Initialization! You're making excellent progress, having now completed two lessons in which we built a flexible MLP architecture and implemented the powerful ReLU activation function. In this third lesson, we'll focus on a critical aspect of neural networks: output layer activation functions . While we've been using activation functions in the hidden layers to introduce nonlinearity and enhance the network's learning capabilities, the activation function in the output layer serves a different purpose. The output layer activation function determines the type of prediction your network can make, and choosing the appropriate one is essential for your model's success. We'll explore two key output activation functions: Softmax : For multi-class classification problems, converting raw outputs into probabilities Linear : For regression problems, allowing the model to predict unbounded continuous values By the end of this lesson, you'll understand when and why to use these activation functions, implement them efficiently, and apply them in different neural network architectures for classification and regression tasks.

Understanding Output Layer Activation Functions

The activation function in the output layer plays a fundamentally different role compared to those in hidden layers. While hidden layer activations primarily introduce nonlinearity to help the network learn complex patterns, output layer activations transform the network's raw outputs into the desired format for your specific task . The choice of output activation depends on the type of problem you're solving: Classification problems : We need outputs that represent probabilities or confidence scores. Binary classification: Sigmoid activation (which we've already implemented) squashes values to the range [0,1]. This means the output can be interpreted as the probability of the input belonging to the positive class, making it easy to set a threshold (like 0.5) for decision-making. Multi-class classification: Softmax activation converts raw scores into a probability distribution across all classes. Each output neuron represents a class, and the softmax ensures the outputs sum to 1, so you can directly interpret them as the model's confidence in each class. Regression problems : We need to predict continuous unbounded values. Linear activation (or no activation) preserves the raw output of the network. This allows the network to predict any real-valued number, which is essential for tasks where the target variable is continuous and unbounded, such as predicting prices or measurements. Understanding this distinction is crucial because using the wrong output activation can lead to poor model performance, even if the rest of your network architecture is sound. For example, using a sigmoid activation for regression would limit your predictions to the range [0,1], which would be problematic if you're trying to predict values like house prices or temperatures. Let's implement these output activation functions and see how they transform our MLP's capabilities.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

The Softmax Activation Function

The Softmax activation function is the natural choice for multi-class classification problems. It converts a vector of real numbers (often called "logits") into a probability distribution over multiple classes.

Mathematically, the softmax function is defined as:

\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}

where $x_i$ is the input value for class $i$ , and $n$ is the total number of classes.

Key properties of softmax:

All output values are between 0 and 1.
The sum of all outputs equals 1, making it a valid probability distribution.
The function amplifies the highest input values and suppresses the lower ones.

When implementing softmax, we need to be careful about numerical stability. The exponential function can lead to extremely large numbers, potentially causing overflow. A common technique is to subtract the maximum value from all inputs before applying the exponential function, which doesn't change the final result but prevents numerical issues.

Here's how you can implement a numerically stable softmax function in JavaScript:

function softmax(x) {
    // x: mathjs matrix or array of shape [nSamples, nFeatures]
    // Numerically stable: subtract row-wise max before exponentiating
    const arr = math.isMatrix(x) ? x.toArray() : x;
    const maxValues = arr.map(row => Math.max(...row));
    const shifted = arr.map((row, i) => row.map(v => v - maxValues[i]));
    const exps = shifted.map(row => row.map(Math.exp));
    const sumExps = exps.map(row => math.sum(row));
    const softmaxArr = exps.map((row, i) => row.map(v => v / sumExps[i]));
    return math.matrix(softmaxArr);
}

Explanation:

If the input is a mathjs matrix, it is converted to a regular array for easier manipulation.
For each row (sample), the maximum value is found and subtracted from every element in that row for numerical stability.
The exponentials of the shifted values are computed.
Each exponentiated value is divided by the sum of exponentials in its row, ensuring the outputs sum to 1.
The result is converted back to a mathjs matrix, where each row is a probability distribution over the classes.

Note: In practice, when we calculate softmax manually or in code, due to floating-point precision limits, the final sum might not be exactly 1, but something like 0.9999999 or 1.0000001. This is because computers use floating-point arithmetic, which means numbers are stored with limited precision and operations like Math.exp() and division can introduce tiny rounding errors. These small discrepancies are normal and generally not a cause for concern.

Enhancing Our DenseLayer Class

Now that we've defined our new activation functions, let's enhance our DenseLayer class to support them. We'll build on the class we updated in the previous lesson to support ReLU.

Here's how we can modify the constructor to handle our new activation functions in JavaScript:

JavaScript

class DenseLayer {
    constructor(nInputs, nNeurons, activationFnName = 'sigmoid') {
        this.weights = math.multiply(math.random([nInputs, nNeurons]), 0.1);
        this.biases = math.zeros(1, nNeurons);
        this.nInputs = nInputs;
        this.nNeurons = nNeurons;
        this.output = null;
        this.activationFnName = activationFnName;

        if (activationFnName === 'sigmoid') {
            this.activationFn = sigmoid;
        } else if (activationFnName === 'relu') {
            this.activationFn = relu;
        } else if (activationFnName === 'softmax') {
            this.activationFn = softmax;
        } else if (activationFnName === 'linear') {
            this.activationFn = linear;
        } else {
            throw new Error(`Unsupported activation function: ${activationFnName}`);
        }
    }

    forward(inputs) {
        const weightedSum = math.multiply(inputs, this.weights);
        const outputBeforeActivation = math.add(weightedSum, this.biases);
        this.output = this.activationFn(outputBeforeActivation);
        return this.output;
    }
}

class DenseLayer {
    constructor(nInputs, nNeurons, activationFnName = 'sigmoid') {
        this.weights = math.multiply(math.random([nInputs, nNeurons]), 0.1);
        this.biases = math.zeros(1, nNeurons);
        this.nInputs = nInputs;
        this.nNeurons = nNeurons;
        this.output = null;
        this.activationFnName = activationFnName;

        if (activationFnName === 'sigmoid') {
            this.activationFn = sigmoid;
        } else if (activationFnName === 'relu') {
            this.activationFn = relu;
        } else if (activationFnName === 'softmax') {
            this.activationFn = softmax;
        } else if (activationFnName === 'linear') {
            this.activationFn = linear;
        } else {
            throw new Error(`Unsupported activation function: ${activationFnName}`);
        }
    }

    forward(inputs) {
        const weightedSum = math.multiply(inputs, this.weights);
        const outputBeforeActivation = math.add(weightedSum, this.biases);
        this.output = this.activationFn(outputBeforeActivation);
        return this.output;
    }
}

We've added softmax and linear as additional options for the activation function.
The forward method remains unchanged, as it already applies whatever activation function was chosen during initialization.
This design allows us to easily extend our neural network framework with new activation functions.

function softmax(x) { // x: mathjs matrix or array of shape [nSamples, nFeatures] // Numerically stable: subtract row-wise max before exponentiating const arr = math.isMatrix(x) ? x.toArray() : x; const maxValues = arr.map(row => Math.max(...row)); const shifted = arr.map((row, i) => row.map(v => v - maxValues[i])); const exps = shifted.map(row => row.map(Math.exp)); const sumExps = exps.map(row => math.sum(row)); const softmaxArr = exps.map((row, i) => row.map(v => v / sumExps[i])); return math.matrix(softmaxArr); }

Introduction

Understanding Output Layer Activation Functions

Join the 1M+ learners on CodeSignal

The Softmax Activation Function

The Linear Activation Function

Enhancing Our DenseLayer Class

Building Multi-Class Classification Networks: Architecture

Building Multi-Class Classification Networks: Output Interpretation

Building Regression Networks: Architecture

Building Regression Networks: Output Interpretation

Conclusion and Next Steps