Inference
Stable softmax: normalize each row without overflowing exponentials
Subtract each row maximum before exponentiating logits, keep the reduction dimensions, and test both probability sums and invalid inputs.
Evidence and testing
Review record: Tensor Almanac automated evidence gate (no independent human review). Example versions, recorded outputs and limitations appear below.
Shift the logits before taking exponentials
For a batch shaped (examples, classes), calculate softmax separately across the classes of each example. Subtract the maximum of each row, exponentiate those shifted values, and divide by that row's sum. The softmax definition divides each exponential by the sum of exponentials along the chosen dimension. PyTorch documents that dimension explicitly; the NumPy example below implements the arithmetic so its intermediate choices are visible. [1]
The shift follows directly from that definition: exp(x_i - m) / sum_j exp(x_j - m) equals exp(x_i) / sum_j exp(x_j), because the common factor exp(-m) cancels in exact arithmetic. Choosing m as the row maximum makes every shifted value nonpositive and at least one zero. Its largest exponential is then one. This derivation explains the numerical choice; it is not a claim that a particular framework uses this exact implementation.
NumPy exp operates element by element. In the executed fixture, applying it directly to logits near 1000 raises a floating-point overflow exception. The shifted calculation instead returns finite probabilities. The second row is near -1000, but has the same within-row differences, so it should produce the same distribution. [2]
Keep one denominator for each row
The input has two rows and three classes. axis=1 reduces across classes, while keepdims=True retains a size-one class dimension. The denominator therefore has shape (2, 1) and broadcasts across its own row. Omitting axis would sum the entire batch; that would answer a different question. These axis and retained-dimension behaviors are documented by NumPy. [3]
The assertions check row sums, output bounds, shape preservation and agreement with the small reference row containing 0, 1 and 2. They also check that NaN, infinity, an empty row and a one-dimensional input are rejected. Printing rounded values makes the result readable; assertions use the unrounded arrays. Run the file with python examples/stable_softmax_rows.py.
Executed with Python 3.12.14 and NumPy 2.3.5 on a bounded float64 fixture; actual stdout and source-byte checksum are preserved in private evidence.
"""Bounded float64 softmax fixture: overflow, row axis and input checks."""
import numpy as np
def softmax_rows(values):
x = np.asarray(values, dtype=np.float64)
if x.ndim != 2 or x.shape[0] == 0 or x.shape[1] == 0:
raise ValueError('expected a nonempty two-dimensional array')
if not np.isfinite(x).all():
raise ValueError('expected finite logits')
with np.errstate(over='raise', invalid='raise', divide='raise', under='ignore'):
shifted = x - np.max(x, axis=1, keepdims=True)
weights = np.exp(shifted)
return weights / np.sum(weights, axis=1, keepdims=True)
logits = np.array([[1000., 1001., 1002.], [-1000., -999., -998.]])
try:
with np.errstate(over='raise'):
np.exp(logits)
except FloatingPointError:
print('direct exponential: overflow detected')
else:
raise AssertionError('expected overflow in the direct fixture')
probabilities = softmax_rows(logits)
reference = softmax_rows([[0., 1., 2.]])
assert np.allclose(probabilities, reference, rtol=1e-12, atol=1e-12)
assert np.allclose(probabilities.sum(axis=1), 1.0)
assert (probabilities >= 0).all() and (probabilities <= 1).all()
assert probabilities.shape == logits.shape
print('probabilities:', probabilities.round(6).tolist())
print('row sums:', probabilities.sum(axis=1).round(12).tolist())
print('common-shift check:', bool(np.allclose(probabilities, reference)))
rejected = 0
for invalid in ([[np.nan, 0.]], [[np.inf, 0.]], [[]], [1., 2.]):
try:
softmax_rows(invalid)
except ValueError:
rejected += 1
assert rejected == 4
print('invalid fixtures rejected:', rejected)
Recorded output
direct exponential: overflow detected probabilities: [[0.090031, 0.244728, 0.665241], [0.090031, 0.244728, 0.665241]] row sums: [1.0, 1.0] common-shift check: True invalid fixtures rejected: 4
References: [3]
Define the input boundary instead of hiding bad values
This helper accepts nonempty two-dimensional inputs that can be converted to float64 and contain only finite values. It deliberately rejects infinite logits rather than inventing an application policy for them. Extreme finite values can still overflow during subtraction; the explicit error policy raises in that case. Small exponential terms may underflow to zero, which this example permits. It does not promise accurate tiny tail probabilities across all floating-point inputs.
The tested result is a small arithmetic example, not a speed comparison, model-quality result or replacement for a framework operator. If a loss expects log probabilities, use the appropriate log-space operation: the PyTorch Softmax documentation specifically points NLLLoss users to LogSoftmax. A row summing to one establishes normalization, not that its probabilities are calibrated to real outcomes. [1]
References: [1]