Loss Functions: Softmax, Cross-Entropy & Numerical Stability

The Softmax function acts as a differentiable categorical normalizer:

1. The Unconstrained Logit Space:

A neural network's final linear layer produces a vector of raw logits z = [z₁, z₂, ṡ, z_K] ∈ R^K, where entries can be negative, zero, or arbitrarily large.

1. The Positivity Guarantee (Exponentiation):

Because the base of the natural exponential e > 0, exponentiating any real number zᵢ ∈ R strictly produces a positive quantity:

e^(zᵢ) > 0 (∀ zᵢ ∈ R)

This guarantees that negative logits never produce invalid negative probabilities.

1. The Sum-to-One Normalization (Probability Simplex):

Dividing each exponentiated term by the partition sum ∑_(j=1)^K e^(zⱼ) enforces the probability simplex constraint:

P(y = i | z) = (e^(zᵢ)) / (∑_(j=1)^K e^(zⱼ))

Where ∑_(i=1)^K P(y = i | z) = 1.0 and P(y = i | z) ∈ (0, 1).

Softmax is the final layer in all multi-class classification networks and the core weighting mechanism in self-attention matrices (A = softmax(Q K^T / √(dₖ))).

The mathematical formula of Softmax originates in statistical mechanics (Ludwig Boltzmann, 1868), where it describes the Maxwell-Boltzmann distribution of energy states in a gas. John S. Bridle adapted it to neural networks in 1989 for multi-class classification.

• Exponentiate to make positive; divide to make sum to one: The two-stroke engine of Softmax.
• Logits to the simplex: Softmax projects unconstrained real space into a valid probability manifold.

#### 1. The Catastrophic Overflow Trap (IEEE 754 FP32)

In standard 32-bit floating point arithmetic (FP32), the maximum representable value is:

MAX_(FP32) ≈ 3.4028 × 10^(38)

Because e^(88.7) ≈ 3.4028 × 10^(38), any logit zᵢ > 88.7 causes e^(zᵢ) to overflow to `+inf`!

When calculating Softmax:

(e^(1000)) / (∑ e^(zⱼ)) = inf/inf = NaN (Not a Number)

A single large logit instantly corrupts the entire neural network with NaNs, crashing training permanently.

Let c = max_(1 ≤ j ≤ K) zⱼ. We subtract c from every logit before exponentiating:

p̃ᵢ = (e^(zᵢ - c)) / (∑_(j=1)^K e^(zⱼ - c))

Using the algebraic exponent rule e^(a - b) = e^a · e^(−b):

p̃ᵢ = (e^(zᵢ) · e^(−c)) / (∑_(j=1)^K (e^(zⱼ) · e^(−c))) = (e^(−c) · e^(zᵢ)) / (e^(−c) · ∑_(j=1)^K e^(zⱼ)) = (e^(zᵢ)) / (∑_(j=1)^K e^(zⱼ)) = pᵢ

The scalar e^(−c) factors out identically from both numerator and denominator!

Stability Guarantee: Because zᵢ - c ≤ 0 for all i, the largest exponent is e⁰ = 1.0, and all other terms are e^(≤ 0) ∈ (0, 1]. Overflow is mathematically impossible! \blacksquare

Given extreme logits z = [1000.0, 1001.0, 1002.0]:

• Naive Softmax:

e^(1000) = inf, e^(1001) = inf, e^(1002) = inf ⟹ inf / inf = NaN.

• Numerical Max-Trick:

c = max(z) = 1002.0.

Shifted logits: z̃ = z - c = [1000 - 1002, 1001 - 1002, 1002 - 1002] = [-2.0, -1.0, 0.0].

Exponentials:

• e^(−2.0) ≈ 0.135335
• e^(−1.0) ≈ 0.367879
• e^(0.0) = 1.000000

Sum: d = 0.135335 + 0.367879 + 1.000000 = 1.503214.

Probabilities:

• p₁ = 0.135335 / 1.503214 = 0.0900
• p₂ = 0.367879 / 1.503214 = 0.2447
• p₃ = 1.000000 / 1.503214 = 0.6653

(Check: 0.0900 + 0.2447 + 0.6653 = 1.0000).

Clean, stable probabilities computed without a single NaN.

In NumPy and PyTorch, this is the universal production implementation:

```python

def stable_softmax(z):

exp_z = np.exp(z - np.max(z, axis=-1, keepdims=True))

return exp_z / np.sum(exp_z, axis=-1, keepdims=True)

```

The coordinate shift invariance of Softmax is a classic mathematical property of projective geometry and Gibbs distributions. Its adoption as the standard "Max-Trick" in scientific computing was codified by early LAPACK and SciPy developers in the 1990s.

• Subtract the max before you exp: The golden rule of numerical deep learning.
• e⁰ = 1.0 is the ceiling: Shift so the maximum exponent is always 1; overflow becomes impossible.

Categorical Cross-Entropy (CCE) is the gold standard loss function for multi-class classification:

1. The Information-Theoretic Definition:

Given true distribution y and predicted distribution ŷ:

H(y, ŷ) = - ∑_(k=1)^K yₖ ln(ŷₖ)

1. One-Hot Simplification:

In single-label classification, the target is a one-hot vector where y_c = 1 for the correct class c, and yₖ = 0 for all k ≠ c. The sum collapses to a single term:

Lᵢ = - ln(ŷ_(i, c))

The loss depends exclusively on the probability assigned to the true class.

1. The Penalty Curve Dynamics:

• Perfect prediction (ŷ_c = 1.0): L = -ln(1.0) = 0.0.
• Confident error (ŷ_c → 0): L = -ln(0) → +∞.

The logarithmic curve provides an infinite penalty wall against confident mistakes, forcing gradients to aggressively correct erroneous weights.

Categorical cross-entropy powers image classifiers (ResNet on ImageNet-1k), speech recognition models, and every foundation LLM. In PyTorch, `nn.CrossEntropyLoss` combines Softmax and Cross-Entropy into a single numerically stable module.

The cross-entropy concept was developed by Claude Shannon in his 1948 paper establishing Information Theory. In 1951, Solomon Kullback and Richard Leibler formalized the KL divergence, proving that minimizing cross-entropy is mathematically identical to minimizing relative entropy between models and empirical truth.

• Loss is negative log of the correct class: Forget all other classes; find the true label and take -ln(p).
• Zero loss requires 100\% confidence: You only achieve zero loss when your predicted probability is exactly 1.0.

#### 1. The ln(0) Catastrophic Underflow Trap

If a neural network assigns a predicted probability of exactly ŷ_(i, c) = 0.0 to the correct class (due to floating-point rounding), evaluating the loss triggers an unrecoverable mathematical error:

Lᵢ = -ln(0.0) = -(-∞) = +∞

Subsequent gradient calculations produce NaNs:

(∂ L) / (∂ z) = NaN

Training crashes instantly.

To safeguard numerical stability, predicted probabilities ŷ are clipped to a closed interval [ε, 1 - ε] (typically ε = 10^(-7)):

ŷ_(clipped) = clip(ŷ, 10^(-7), 1 - 10^(-7))

• Even if the model outputs ŷ = 0.0, the clipped value is 10^(-7), yielding a safe finite loss:

L_(max) = -ln(10^(-7)) = -(-16.11809565) ≈ 16.118

• The maximum possible loss is strictly bounded, preventing gradient explosions.

For a mini-batch of N samples, the total batch loss is the arithmetic mean:

L_(batch) = 1/N ∑_(i=1)^N Lᵢ = - 1/N ∑_(i=1)^N ln(ŷ_(i, cᵢ))

Given a batch of N = 3 samples across K = 3 classes:

• Target indices: y = [0, 1, 1] (Classes: Cat, Dog, Dog).
• Predicted probabilities:

ŷ =

0.70 & 0.20 & 0.10 \

0.10 & 0.50 & 0.40 \

0.02 & 0.90 & 0.08

Step 1: Extract Correct Class Probabilities:

• Sample 1 (c₁ = 0): ŷ_(1, 0) = 0.70
• Sample 2 (c₂ = 1): ŷ_(2, 1) = 0.50
• Sample 3 (c₃ = 1): ŷ_(3, 1) = 0.90

Step 2: Compute Sample Losses (with 10^(-7) clipping):

• L₁ = -ln(0.70) ≈ -(-0.3567) = 0.3567
• L₂ = -ln(0.50) ≈ -(-0.6931) = 0.6931
• L₃ = -ln(0.90) ≈ -(-0.1054) = 0.1054

Step 3: Compute Batch Mean:

L_(batch) = (0.3567 + 0.6931 + 0.1054) / (3) = (1.1552) / (3) ≈ 0.3851

In NumPy, this is implemented as:

```python

y_pred_clipped = np.clip(y_pred, 1e-7, 1.0 - 1e-7)

correct_confidences = y_pred_clipped[range(len(y_pred)), y_true]

negative_log_likelihoods = -np.log(correct_confidences)

loss = np.mean(negative_log_likelihoods)

```

The epsilon clipping technique was popularized in early machine learning libraries like scikit-learn and Keras (circa 2015) to prevent numerical divergence in binary and categorical cross-entropy implementations.

• Clip between 1e-7 and 1 - 1e-7: Never let the logarithm touch zero.
• Average across the batch: The batch loss is the mean of the individual sample penalties.

To prevent numerical divergence when p → 0 or p → 1, probabilities are clipped to a safe epsilon range [ε, 1 - ε]:

p_(safe) = clip(p, 10^(-7), 1 - 10^(-7))

L = -ln(p_(safe))

This bounds maximum scalar loss to -ln(10^(-7)) ≈ 16.118, preventing catastrophic floating-point NaNs.

Standard defensive engineering inside PyTorch, Keras, and scikit-learn loss routines.

Developed through decades of numerical analysis in floating-point IEEE-754 arithmetic standards (Kahan, 1985).

Clip with epsilon to banish infinity: never let probability touch exact zero.

For a Bernoulli random variable with parameter p = P(X = 1):

H_b(p) = -p log₂ p - (1 - p) log₂(1 - p)

Analyzing its calculus and geometry:

1. Boundary Limits:

lim_(p → 0^+) p log₂ p = 0, lim_(p → 1^-) (1 - p) log₂(1 - p) = 0 ⟹ H_b(0) = H_b(1) = 0

1. First Derivative (Slope):

(d H_b) / (dp) = -log₂ p - (1) / (ln 2) + log₂(1 - p) + (1) / (ln 2) = log₂( (1 - p) / (p) )

Setting (d H_b) / (dp) = 0:

(1 - p) / (p) = 1 ⟹ 1 - p = p ⟹ p = 0.5

1. Second Derivative (Strict Concavity):

(d² H_b) / (dp²) = -(1) / (ln 2) ( 1/p + (1) / (1 - p) ) = -(1) / (p(1 - p) ln 2) < 0, ∀ p ∈ (0, 1)

Since the second derivative is strictly negative everywhere, H_b(p) is strictly concave, and p = 0.5 is the unique global maximum:

H_b(0.5) = -0.5 log₂(0.5) - 0.5 log₂(0.5) = 0.5(1) + 0.5(1) = 1.0 bit

Binary cross-entropy loss L_(BCE) = -y log p - (1-y) log(1-p) in logistic regression and binary classification networks directly reflects this curve.

Claude Shannon plotted this exact symmetric arch as Figure 7 in his 1948 Bell Labs paper, introducing the universal graphic icon of information theory.

• Zero at the edges; 1 bit at the 50/50 peak.
• A fair coin has maximum uncertainty.

Let p be the true target distribution (reality) and q be the predicted distribution (the model's codebook).

Cross-entropy is the expected surprisal under p using code lengths -log₂ qᵢ:

H(p, q) = E_(X ~ p)[-log₂ q(X)] = -∑_(i=1)^K pᵢ log₂ qᵢ

We decompose cross-entropy into true entropy plus excess penalty:

H(p, q) = -∑_(i=1)^K pᵢ log₂ qᵢ = -∑_(i=1)^K pᵢ log₂ pᵢ + ∑_(i=1)^K pᵢ log₂( pᵢ/qᵢ ) = H(p) + D_(KL)(p ∥ q)

Where:

• H(p) is the irreducible intrinsic entropy of reality (noise floor).
• D_(KL)(p ∥ q) ≥ 0 is the Kullback-Leibler divergence (the penalty for using model q).

By Gibbs' inequality, H(p, q) ≥ H(p), with equality if and only if q = p everywhere.

Cross-entropy is the universal loss function for classification and language generation. Minimizing H(p, q) during SGD pushes model q to match data distribution p perfectly, driving D_(KL)(p ∥ q) → 0.

Solomon Kullback and Richard Leibler formalized the relative cross-entropy divergence in 1951, providing the statistical bridge between Shannon's information theory and Fisher's hypothesis testing.

• Cross-entropy = True Entropy + KL Divergence.
• The cost of reality encoded through an imperfect model.

For a 2D tensor A ∈ R^(M × N) with row index i and column index j:

1. Row Collapse (`axis=0`):

Collapses along vertical rows:

s_(col) = ∑_(i=1)^M A_(i, j) ∈ R^N

• Shape: (M, N) →(i, j) } (sum across class columns)

Understanding axis reductions prevents catastrophic orientation errors when accumulating distributed gradients across GPUs (`all_reduce`).

The axis reduction conventions (`axis=0, 1, ...`) were established in APL (A Programming Language) by Kenneth Iverson in 1962 and adapted into NumPy and PyTorch.

• axis=0 squashes the rows (downward collapse).
• axis=1 squashes the columns (horizontal collapse).

A 3D perspective geometric diagram on clean ivory linen (#f7f6ee) with an 8px border.

Top section: Title in technical grotesque: "Logits to Simplex: Projecting Real Space onto Probabilities". Subtitle: "Softmax normalizes an unconstrained 3D vector onto a 2D probability manifold."

Middle section: A clean 3D coordinate frame (axes z₁, z₂, z₃).

Show the standard 2-simplex triangle Δ² lying on the plane z₁ + z₂ + z₃ = 1 in the positive octant.

Draw an unconstrained 3D logit point z = [3.0, 1.0, 0.0].

Draw a projection ray from z down onto point p = [0.84, 0.11, 0.05] on the simplex triangle surface.

Show temperature effect: As T -> ∞, point p slides toward the exact center of the triangle (maximum entropy uniform point [1/3, 1/3, 1/3]). As T -> 0, point p snaps to the closest vertex (1, 0, 0).

Bottom section: Formula card: "pᵢ = (e^(zᵢ / T)) / (∑ e^(zⱼ / T)) ⟹ Softmax is an exponential map followed by normalization onto the simplex manifold."

Crisp 3D coordinate lines, airy layout, elegant mathematical styling.

This principle is foundational across deep learning optimization, high-dimensional representation, and modern LLM architectures.

"Logits to Simplex: Projecting Real Space onto Probabilities: Invariance and structure determine stability under spatial and numerical transformations."

A geometric probability projection diagram in square 1:1 format on warm cream parchment (#fcf0db) with a sepia frame.

Top section: Title in navy serif: "Logits to Simplex: Projecting Real Space onto Probabilities". Subtitle: "Softmax non-linearly maps unbounded scores R^K onto the bounded probability simplex Δ^(K-1)."

Middle section: A clean 3D isometric coordinate plane:

Show the unbounded real score space where raw logits z = [z₁, z₂, z₃]^T can range from -∞ to +∞.

Show the triangular probability simplex Δ² slicing through the axes.

Draw three gradient vectors showing how changing logit z₁ pushes the output point toward the corner vertex [1, 0, 0] (Pure Class 1).

Annotate the translation invariance property:

"Softmax(z + c 1) = Softmax(z)".

Subtracting max(z) prevents floating-point exponential overflow without changing output probabilities!

Bottom section: Formula: "pᵢ = (e^(zᵢ - max(z))) / (∑ⱼ e^(zⱼ - max(z))). Stable Softmax guarantees numeric safety across extreme logit scales."

Pristine isometric simplex diagram, clear projection arrows, elegant square composition.

Every transformation preserves structural invariants while mapping coordinate spaces. When operating in high dimensions, numerical precision and computational efficiency require adhering strictly to canonical mathematical laws.

Formal derivation and peer-reviewed analysis:

Goodfellow, Bengio, and Courville (2016), Deep Learning, MIT Press ↗

"Logits to Simplex: Projecting Real Space onto Class Probabilities: Invariance, symmetry, and structure govern mathematical truth across continuous and discrete domains."