Softmax

Softmax is a function that takes a vector of logits (real-valued scores) and converts them into a probability distribution: each output is between 0 and 1, and all outputs sum to 1. In neural networks, softmax is typically used in the final layer for multi-class classification so that the network’s raw outputs can be interpreted as class probabilities and used with a loss like cross-entropy.

Mathematically, for an input vector $x = [x_1, x_2, \dots, x_n]$, the softmax of component $x_i$ is:

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

This is a normalized exponential: the numerator magnifies differences between scores (larger logits get much larger exponentials), and the denominator ensures the outputs sum to 1. In practice, implementations must also consider numerical stability, because $e^{x}$ grows very fast and can overflow for large $x$.

---

1. Background Knowledge

• Logits vs probabilities Neural networks usually output unbounded real numbers (logits). Softmax converts these into a probability vector:

• $p_i \in (0, 1)$ for each class $i$

• $\sum_i p_i = 1$ This makes them suitable for probabilistic interpretation and for computing losses like cross-entropy.

• Shift invariance and stability Softmax is invariant to adding a constant to all inputs:

$$\text{softmax}(x_i) = \text{softmax}(x_i + c)$$

for any scalar $c$. This property is crucial for implementing a numerically stable version by subtracting the maximum logit from all logits before exponentiating.

---

2. Algorithm / Approach

The general algorithm pattern:

1. Optionally stabilize the inputs by subtracting their maximum value.
2. Exponentiate each (possibly shifted) input element.
3. Compute the sum of all exponentials.
4. Normalize each exponential by dividing by the sum.
5. Return the resulting vector of probabilities.

Depending on the problem setting, you must:

• Handle 1D inputs (single vector of logits).
• Possibly handle 2D inputs (batch of vectors), applying the same logic row-wise (or along a chosen axis).

---

3. Step-by-Step Strategy

Assume a 1D array x of length n:

1. Find the maximum value

x_max = max(x)

This is for numerical stability.

2. Shift the inputs by subtracting x_max from each element:

shifted = [xi - x_max for xi in x]

3. Exponentiate each shifted value:

exps = [math.exp(si) for si in shifted]

4. Compute the sum of exponentials:

sum_exps = sum(exps)

5. Normalize each exponential to get probabilities:

softmax = [ei / sum_exps for ei in exps]

6. Return the softmax array.

For a 2D input (e.g., shape (batch_size, num_classes)), apply the above for each row independently.

---

4. Common Pitfalls

• Numerical overflow/underflow
• Directly computing exp(xi) when xi is large (e.g., 1000) can overflow.
• Always do the max-subtraction trick:

$$y_i = x_i - \max_j x_j$$

and then compute $\exp(y_i)$.

• Not normalizing correctly

• Ensure you divide each exponential by the sum over the correct axis (e.g., per row in a batch), not by the sum over the entire array.

• In-place modifications

• If reusing input arrays, be careful about modifying them in place before you’ve finished computing the sum of exponentials.

• Precision issues with very small probabilities

• When logits are very negative relative to the max, exponentials may underflow to zero. This is expected in extreme cases; the max-subtraction helps but cannot avoid all underflow.

---

5. Time & Space Complexity

Let $n$ be the number of elements in the input vector (or per row for batched input):

• Time complexity

• Finding the max: $O(n)$

• Computing exponentials: $O(n)$

• Summing exponentials: $O(n)$

• Normalizing: $O(n)$

• Overall: $O(n)$ per vector (or $O(B \cdot n)$ for batch size $B$).

• Space complexity

• If you store exponentials and output separately: $O(n)$ extra space.

• If you allow in-place and overwrite intermediate arrays carefully, you can keep extra space close to $O(1)$ plus output storage.