Softmax Function

mediumFREEClassificationMachine Learning 1

Implement the softmax function for multi-class classification.

Given a list of logits (raw scores) $z = [z_1, z_2, ..., z_k]$ , compute the softmax probabilities:

$\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{k} e^{z_j}}$

For numerical stability, subtract the maximum value from all logits before exponentiating: $\text{softmax}(z_i) = \frac{e^{z_i - \max(z)}}{\sum_{j=1}^{k} e^{z_j - \max(z)}}$

Return the probability distribution as a list, rounded to 4 decimal places.

Example

Input

logits = [2.0, 1.0, 0.1]

Output

[0.6590, 0.2424, 0.0986]

Reasoning

logits holds one raw score per class, $z = [2.0, 1.0, 0.1]$ . Scores can be any real numbers, so they are not probabilities yet.
Exponentiating makes every score positive: $e^{2.0} \approx 7.389$ , $e^{1.0} \approx 2.718$ , $e^{0.1} \approx 1.105$ , total $\approx 11.213$ . Dividing each by the total gives $7.389 / 11.213 \approx 0.6590$ , and so on.
The stability trick shifts every logit by the largest one, $2.0$ , giving $[0.0, -1.0, -1.9]$ and exponentials $[1, 0.3679, 0.1496]$ with total $1.5174$ . Each term was scaled by the same factor $e^{-2}$ , so the ratios are unchanged: $1 / 1.5174 \approx 0.6590$ , $0.3679 / 1.5174 \approx 0.2424$ , $0.1496 / 1.5174 \approx 0.0986$ .
The output [0.6590, 0.2424, 0.0986] has one probability per logit, in the same order, rounded to 4 decimal places, and they sum to $1$ .
Notice that a logit gap of $1.0$ multiplies the probability by $e \approx 2.72$ ( $0.6590 / 0.2424 \approx 2.72$ ). Equal logits such as $[100, 100, 100]$ all shift to $0$ , so each gets $1/3$ ; for logits near $1000$ , $e^{1000}$ would overflow without the shift.

Constraints

Input is a list of floats (logits)
Return a list of probabilities that sum to 1 (approximately)
Round each value to 4 decimal places
Use the numerical stability trick (subtract max)

Learn this in

ClassificationML Study Plan Neural NetworksML Study Plan

More from Machine Learning 1

RBF Kernel Matrixhard Polynomial Kernelmedium Hinge Losseasy Linear Decision Boundaryeasy Gradient Boosting Residual Stephard

Concept primer

Background Knowledge

The softmax function is a crucial component in multi-class classification problems, where the goal is to predict one of multiple classes or labels. It takes a list of logits (raw, unnormalized scores) as input and outputs a probability distribution over all classes. The softmax function is often used in the final layer of neural networks to ensure that the output values are valid probabilities, i.e., non-negative and summing up to 1.

In the context of machine learning, the softmax function is used to model the probability of each class given the input features. The exponential function is used to map the logits to non-negative values, and the normalization step ensures that the probabilities sum up to 1. The softmax function is widely used in various applications, including image classification, natural language processing, and recommendation systems.

To ensure numerical stability, it's essential to subtract the maximum value from all logits before exponentiating. This trick helps prevent overflow issues when dealing with large input values. By doing so, we can avoid NaN (Not a Number) or inf (infinity) values in the output, which can occur when the exponential function is applied to large numbers.

Algorithm/Approach

The general approach to solving this problem involves implementing the softmax function using the given formula. The key steps include:

Computing the maximum value of the input logits
Subtracting the maximum value from all logits to ensure numerical stability
Applying the exponential function to the normalized logits
Computing the sum of the exponentiated values
Normalizing the exponentiated values by dividing them by the sum

This approach can be implemented using a variety of programming languages and libraries, including Python and NumPy.

Step-by-Step Strategy

To implement the softmax function, follow these steps:

Compute the maximum value of the input logits: max_value = max(z)
Subtract the maximum value from all logits: z_normalized = [z_i - max_value for z_i in z]
Apply the exponential function to the normalized logits: exp_values = [ez_i for z_i in z_normalized]**
Compute the sum of the exponentiated values: sum_exp_values = sum(exp_values)
Normalize the exponentiated values: probabilities = [exp_value / sum_exp_values for exp_value in exp_values]
Round the probabilities to 4 decimal places: probabilities = [round(p, 4) for p in probabilities]

Common Pitfalls

When implementing the softmax function, watch out for the following common pitfalls:

Forgetting to subtract the maximum value from all logits, which can lead to numerical instability
Using the wrong formula or implementation, which can result in incorrect output
Not handling edge cases, such as input lists with zero or negative values

Time & Space Complexity

The time complexity of the softmax function is O(k), where k is the number of classes or logits. This is because we need to iterate over the input list to compute the maximum value, apply the exponential function, and normalize the values.

The space complexity is also O(k), as we need to store the normalized logits, exponentiated values, and output probabilities. However, the space complexity can be reduced by using in-place computations and avoiding unnecessary memory allocations.

solution.py

Python 3.13

Test results