Softmax, Cross-Entropy & Optimization

The previous lesson predicted one continuous scalar: extra decode latency in milliseconds. Half squared error could say "raise the prediction by four units." A review triage bot faces a different challenge: pick exactly one category among bug, docs, and security.

Suppose the incident note reads, "This docs page is stale," while the model emits raw scores bug=3.0, docs=1.0, security=0.0. It picks bug, missing the correct label. What training signal should backpropagation send back through the network?

A binary right-or-wrong check gives the same result for a narrow near-miss and a confident mistake. Assigning docs an 11% probability should cost significantly more than assigning it 49%, even if bug wins the top spot in both runs. We need a mathematical bridge that turns classification scores into normalized probability shares, converts missed confidence into an exact penalty, and sends back a gradient on those scores. Normalization makes probabilities sum to one; calibration requires comparing predicted confidence with observed outcomes.

A decision begins as three raw scores

Our final linear layer outputs one real number per possible class. These numbers are logits: unconstrained scores. A linear classification head computes each score as a dot product between the internal representation and a class weight vector, plus a bias if present. Softmax later converts these real-valued scores into probabilities.

A larger logit means the network prefers that label over competitors. The value 3.0 is a score with no direct percentage interpretation. Read the table as an uncalibrated ranking: bug leads docs by two score points, but we haven't assigned any percentage share to any option.

Logits can be negative, zero, or wildly positive without breaking mathematics. A finite categorical distribution over mutually exclusive, exhaustive choices must have every probability $p_i$ in $[0, 1]$ and sum to $1$.

Predict the three checks in inspect-raw-logits.py before running it. bug takes the highest score, the raw sum is $4.0$, and the list fails the probability test. The scores rank labels; an information-theoretic loss needs their normalized distribution.

labels = ["bug", "docs", "security"]
logits = [3.0, 1.0, 0.0]

best_index = max(range(len(logits)), key=logits.__getitem__)
print("largest logit:", labels[best_index])
print("raw score sum:", sum(logits))
print("valid probability distribution:", all(0 <= z <= 1 for z in logits) and sum(logits) == 1)

largest logit: bug
raw score sum: 4.0
valid probability distribution: False

The raw scores identify a winner, but they don't yet say how much probability the correct label received.

Softmax turns score gaps into probability mass

Softmax turns each unconstrained score into a positive weight, then scales the weights so they sum to one. It does that by exponentiating each score and normalizing the results.

First, exponentiate each logit. Since $e^z > 0$ for every real number $z \in (-\infty, +\infty)$, this turns even a negative score into a positive weight. Exponentiation also magnifies differences: a score lead of $+2.0$ becomes a weight multiplier of $e^{2.0} \approx 7.389$.

Second, it sums those positive weights into a single shared denominator, called the partition function:

Dividing each exponential weight by this sum forces the output shares to sum to $1.0$.

For our three triage scores, walk through the arithmetic: $e^{3.0} \approx 20.086$, $e^{1.0} \approx 2.718$, and $e^{0.0} = 1.000$. The shared sum is $23.804$.

The general equation applies this two-step transformation to any vector $z \in \mathbb{R}^K$:

Here $p_i$ is the probability assigned to class $i$, $z_i$ is its logit, and $K$ is the number of competing classes. Notice the distribution: the model assigns 0.844 to the incorrect bug label, leaving only 0.114 for the correct docs label. That value is the handoff to the loss function.

View larger

Before running stable-softmax.py, confirm the invariants for these small scores: the class ranking stays preserved, every entry is positive, and the values sum to 1.0. Mathematically, softmax is strictly positive for finite logits; floating-point implementations can underflow tiny shares to zero and introduce rounding error in the sum.

import math

def stable_softmax(logits: list[float]) -> list[float]:
    peak = max(logits)
    weights = [math.exp(z - peak) for z in logits]
    total = sum(weights)
    return [weight / total for weight in weights]

labels = ["bug", "docs", "security"]
logits = [3.0, 1.0, 0.0]
probabilities = stable_softmax(logits)

for label, probability in zip(labels, probabilities):
    print(f"{label:8s} {probability:.3f}")
print("sum     ", round(sum(probabilities), 3))

bug      0.844
docs     0.114
security 0.042
sum      1.0

Stable arithmetic prevents floating-point overflow

Our gentle logits [3.0, 1.0, 0.0] fit easily inside standard floating-point variables. What happens when a deep network produces raw scores like [1000.0, 998.0, 997.0]?

If exponentiation runs in float32, its largest finite result is approximately $3.4 \times 10^{38}$, so exponential overflow starts around a logit of $88.72$. For float16, the largest finite value is $65504$, and the threshold is about $11.09$: even $e^{12}$ cannot fit. bfloat16 has roughly float32's range, with fewer precision bits.^{[1]Reference 1Improve Your Model's Performance with bfloat16https://docs.cloud.google.com/tpu/docs/bfloat16} An operation may compute in a wider dtype than its input; the implementation's compute dtype determines the exponent range used for exponentiation.

Once an exponential becomes +inf, the denominator can become +inf. Dividing inf / inf yields NaN (Not a Number). If you backpropagate through a nonfinite loss and step the optimizer, the invalid values can spread into parameters. Python's math.exp, used below, computes with ordinary Python floats (normally binary64) and raises OverflowError for math.exp(1000.0) instead of returning infinity.

Softmax possesses a key algebraic property: shift invariance.^{[2]Reference 2Deep Learning.https://www.deeplearningbook.org/} Adding or subtracting any arbitrary scalar constant $c$ from every logit leaves the resulting probabilities completely unchanged:

The common factor $e^{-c}$ factors out of the sum and cancels between numerator and denominator. This identity unlocks the classic numerical stabilization trick: choose $c = m = \max_j(z_j)$, the maximum logit across the vector.

For a nonempty vector of finite logits, subtracting $m$ makes the largest shifted value zero ($z_{\text{max}} - m = 0$). Every other shifted logit is nonpositive ($z_j - m \le 0$); tied maxima are also zero. Exponentiating these values gives:

1. The largest exponential is $e^0 = 1.0$. Overflow is mathematically impossible.
2. Every other exponential lies in $(0, 1]$.
3. The denominator sum is at least $1.0$, preventing division by zero.

Very negative shifted logits (such as $-1000$) can underflow to 0.0, approximating a tiny probability by zero. Keep the loss in log space to avoid taking the logarithm of that rounded zero. Max-subtraction stabilizes finite values; validate inputs for NaN or +inf separately. A masked candidate may have logit -inf if at least one candidate has a finite score; an all-masked vector has no valid distribution.

import math

def stable_softmax(logits: list[float]) -> list[float]:
    peak = max(logits)
    weights = [math.exp(z - peak) for z in logits]
    total = sum(weights)
    return [weight / total for weight in weights]

large_logits = [1000.0, 998.0, 997.0]
try:
    _ = [math.exp(z) for z in large_logits]
    naive_ok = True
except OverflowError:
    naive_ok = False

stable = stable_softmax(large_logits)
print("naive overflowed:", not naive_ok)
print("stable:", [round(probability, 3) for probability in stable])
print("stable sum:", round(sum(stable), 3))

naive overflowed: True
stable: [0.844, 0.114, 0.042]
stable sum: 1.0

For a NumPy array in float32 or float64, naive exponentiation of these scores produces inf, and dividing the infinite entries by their infinite sum produces NaN (normally with warnings). Max-subtraction avoids that overflow and preserves the mathematical distribution, subject to floating-point rounding.

Cross-entropy prices missed predictions through surprise

Softmax gave us a probability distribution where docs received $p_{\text{docs}} \approx 0.114$. How do we score that prediction?

Cross-entropy measures the quality of a predicted distribution against ground truth. It's grounded in Maximum Likelihood Estimation (MLE) and information theory.

In information theory, observing an event with probability $p$ yields an information content, or surprise, of:

We use natural logarithms throughout, measuring surprise in nats (using base-two logarithms would measure surprise in bits). If an event is certain ($p = 1.0$), observing it brings zero surprise: $-\ln(1.0) = 0$. If an event is extremely rare ($p \to 0$), observing it brings near-infinite surprise: $-\ln(0.001) \approx 6.91$ nats.

When training a supervised model, the ground truth for an input is a target distribution $y$. For single-label classification, $y$ is a one-hot vector where the correct class index $c$ gets weight $1.0$ and all competing classes get $0.0$. For our triage note:

Cross-entropy calculates the expected surprise across classes under the supplied target distribution:

Because $y_k = 0$ for all non-target classes and $y_c = 1$ for the target, every term except the correct class multiplies by zero. Cross-entropy collapses to the Negative Log-Likelihood (NLL) of the true target class:

For our incident note where $c = \text{docs}$:

If the network had assigned 0.844 to docs, the loss would be only $-\log(0.844) \approx 0.170$ nats. If it had placed 0.999 on docs, the loss would plummet to $0.001$ nats.

The log-sum-exp formulation prevents underflow crashes

Computing p = softmax(z) and then calling math.log(p[c]) introduces a subtle failure mode. With logits = [1000.0, 0.0, -1.0] and target docs, the target's shifted logit is $-1000$. Its exponential underflows to 0.0. Python's math.log(0.0) raises ValueError; NumPy and PyTorch logarithms return -inf, giving an infinite negative log-likelihood. The finite logit difference gives a finite mathematical loss; floating-point underflow caused the failure.

The loss can be expressed directly in terms of raw logits, bypassing intermediate probabilities entirely.^{[2]Reference 2Deep Learning.https://www.deeplearningbook.org/} Expanding $p_c$:

The first term is log-sum-exp. We stabilize it using the same maximum shift $m = \max_j(z_j)$:

Substituting this back into the loss expression:

Grouping $(m - z_c)$ first prevents catastrophic cancellation. When $z_c = 0.0$ and $m = 1000.0$, the first term evaluates cleanly to $1000.0$, while the log-sum evaluates to $\log(1.0 + \dots) \approx 0.0$. The loss returns $1000.0$ without ever taking the logarithm of zero.

import math

def cross_entropy_from_logits(logits: list[float], target_index: int) -> float:
    peak = max(logits)
    log_total = math.log(sum(math.exp(z - peak) for z in logits))
    return (peak - logits[target_index]) + log_total

ordinary = [3.0, 1.0, 0.0]
shifted_high = [z + 1000.0 for z in ordinary]

print("ordinary docs loss:", round(cross_entropy_from_logits(ordinary, 1), 3))
print("large-offset loss: ", round(cross_entropy_from_logits(shifted_high, 1), 3))
print("underflowed target weight:", math.exp(-1000.0))
print("finite target loss:", cross_entropy_from_logits([1000.0, 0.0, -1.0], 1))
print("equal huge scores:", round(cross_entropy_from_logits([1e20, 1e20, 1e20], 1), 3))
assert math.isclose(cross_entropy_from_logits([1e20] * 3, 1), math.log(3))

ordinary docs loss: 2.17
large-offset loss:  2.17
underflowed target weight: 0.0
finite target loss: 1000.0
equal huge scores: 1.099

Entropy is reference uncertainty and KL divergence is prediction error

Before moving on to gradients, keep three related quantities separate. They use the same logarithms, but answer different questions about uncertainty and prediction error.

Entropy $H(P)$ measures the inherent uncertainty already present inside a reference probability distribution $P$:

Cross-entropy $H(P, Q)$ measures the total cost of representing distribution $P$ using model distribution $Q$. Kullback-Leibler divergence $D_{\text{KL}}(P \parallel Q)$ measures the extra penalty caused specifically by the mismatch between $Q$ and $P$:

For our one-hot docs target, $P = [0, 1, 0]$ has zero entropy ($H(P) = 0$): this annotation assigns all target mass to one label. Every fraction of a nat in the $2.170$ loss is mismatch with that target ($D_{\text{KL}} = 2.170$). The annotation records its selected label; ambiguity and label noise require separate evidence.

When targets are soft (such as in distillation or human annotations with disagreement, like $P = [0.5, 0.5, 0]$), the reference distribution has inherent entropy $H(P) = \ln(2) \approx 0.693$ nats. Even an ideal model ($Q = P$) will report a loss of $0.693$ nats, with $D_{\text{KL}} = 0$. Don't mistake inherent target uncertainty for model failure.

from math import exp, inf, log

def entropy(probabilities: list[float]) -> float:
    return sum(-p * log(p) for p in probabilities if p > 0)

def cross_entropy(target: list[float], predicted: list[float]) -> float:
    # Inputs must be normalized, nonnegative distributions over the same classes.
    pairs = list(zip(target, predicted, strict=True))
    if any(p > 0 and q == 0 for p, q in pairs):
        return inf
    return -sum(p * log(q) for p, q in pairs if p > 0)

weights = [exp(value) for value in [3.0, 1.0, 0.0]]
predicted = [weight / sum(weights) for weight in weights]
one_hot_target = [0.0, 1.0, 0.0]
uncertain_target = [0.5, 0.5, 0.0]

one_hot_loss = cross_entropy(one_hot_target, predicted)
matched_loss = cross_entropy(uncertain_target, uncertain_target)

print(f"one-hot target entropy: {entropy(one_hot_target):.3f}")
print(f"one-hot cross-entropy:  {one_hot_loss:.3f}")
print(f"one-hot KL mismatch:    {one_hot_loss - entropy(one_hot_target):.3f}")
print(f"soft target entropy:    {entropy(uncertain_target):.3f}")
print(f"matching cross-entropy: {matched_loss:.3f}")
print(f"matching KL mismatch:   {matched_loss - entropy(uncertain_target):.3f}")

one-hot target entropy: 0.000
one-hot cross-entropy:  2.170
one-hot KL mismatch:    2.170
soft target entropy:    0.693
matching cross-entropy: 0.693
matching KL mismatch:   0.000

These functions assume valid probability distributions; zip(..., strict=True) prevents a mismatched class list from silently losing terms. We use the convention $0\log 0 = 0$. If $P(k)>0$ but $Q(k)=0$, cross-entropy and KL divergence are infinite. For valid distributions, KL is nonnegative and the minimum cross-entropy is the target entropy. Soft targets may express disagreement, distillation, or deliberate smoothing, so that entropy isn't always irreducible noise in the real task.

The combined gradient reveals a clean error signal

Loss tells us how poorly the model performed. To update weights, backpropagation needs the gradient: the partial derivative of the loss with respect to each logit, $\frac{\partial L}{\partial z_k}$. The open question is whether that derivative stays readable after softmax, or whether the probability Jacobian buries the error.

Softmax and cross-entropy simplify together. The logarithm and the softmax derivative cancel, leaving a logit gradient of predicted probability minus target.^{[2]Reference 2Deep Learning.https://www.deeplearningbook.org/}

Recall the loss in terms of raw logits:

Differentiate with respect to any logit $z_k$:

By the chain rule, the derivative of $\log(u)$ is $\frac{1}{u} \frac{\partial u}{\partial z_k}$. With $u = \sum_j e^{z_j}$:

The derivative of the log-sum-exp term is softmax itself.

Meanwhile, the second term $\frac{\partial z_c}{\partial z_k}$ equals $1$ if $k = c$ (the target class) and $0$ otherwise. That is exactly the target weight $y_k$:

Vectorized across all classes:

Predicted probability minus true target. Look at the numbers for our triage example:

The zero-sum conservation property

Sum all elements of the gradient vector:

The gradients sum to zero. The upward pull on the correct class ($-0.886$) is exactly balanced by the downward push on competing classes ($+0.844 + 0.042 = +0.886$). Softmax distributes a fixed unit of probability mass; learning to favor one class inevitably pulls mass away from others.

View larger

In a complete neural network, backpropagation channels this $p - y$ residual back to shared weights via the chain rule. Testing on logits directly makes the local mechanic transparent.

import math

def probabilities_and_loss(logits: list[float], target_index: int) -> tuple[list[float], float]:
    peak = max(logits)
    shifted = [z - peak for z in logits]
    log_total = math.log(sum(math.exp(value) for value in shifted))
    log_probs = [value - log_total for value in shifted]
    probabilities = [math.exp(value) for value in log_probs]
    return probabilities, -log_probs[target_index]

labels = ["bug", "docs", "security"]
logits = [3.0, 1.0, 0.0]
target_index = labels.index("docs")
target = [1.0 if index == target_index else 0.0 for index in range(3)]

before, before_loss = probabilities_and_loss(logits, target_index)
gradient = [probability - label for probability, label in zip(before, target)]
after_logits = [z - 0.5 * slope for z, slope in zip(logits, gradient)]
after, after_loss = probabilities_and_loss(after_logits, target_index)

print("gradient:", [round(value, 3) for value in gradient])
print("docs probability:", round(before[1], 3), "->", round(after[1], 3))
print("loss:", round(before_loss, 3), "->", round(after_loss, 3))

gradient: [0.844, -0.886, 0.042]
docs probability: 0.114 -> 0.23
loss: 2.17 -> 1.469

Confirm the mathematical slope against numerical finite differences: nudge each logit by $\epsilon = 10^{-5}$ and check the symmetric difference quotient.

import math

def loss(logits: list[float], target_index: int) -> float:
    peak = max(logits)
    shifted = [z - peak for z in logits]
    return math.log(sum(math.exp(value) for value in shifted)) - shifted[target_index]

logits = [3.0, 1.0, 0.0]
target_index = 1
peak = max(logits)
shifted = [z - peak for z in logits]
total = sum(math.exp(value) for value in shifted)
probabilities = [math.exp(value) / total for value in shifted]
analytic = [
    probability - (1.0 if index == target_index else 0.0)
    for index, probability in enumerate(probabilities)
]

epsilon = 1e-5
numeric = []
for index in range(3):
    plus = logits.copy()
    minus = logits.copy()
    plus[index] += epsilon
    minus[index] -= epsilon
    slope = (loss(plus, target_index) - loss(minus, target_index)) / (2 * epsilon)
    numeric.append(slope)

print("analytic:", [round(value, 6) for value in analytic])
print("numeric: ", [round(value, 6) for value in numeric])
print("match:", all(math.isclose(left, right, abs_tol=1e-6) for left, right in zip(analytic, numeric)))

analytic: [0.843795, -0.885805, 0.04201]
numeric:  [0.843795, -0.885805, 0.04201]
match: True

Why squared probability error can learn slowly from confident mistakes

Why not use squared error on predicted probabilities: $L_{\text{sq}} = \frac{1}{2} \sum_k (p_k - y_k)^2$? This is a scaled Brier loss, a valid proper scoring rule for categorical forecasts.^{[3]Reference 3Strictly Proper Scoring Rules, Prediction, and Estimationhttps://sites.stat.washington.edu/raftery/Research/PDF/Gneiting2007jasa.pdf} It can train classifiers. Here we compare its optimization behavior with cross-entropy; the sum above isn't an average over classes.

Examine how gradients flow through the softmax Jacobian matrix:

When differentiating cross-entropy $L_{\text{CE}} = -\sum_i y_i \log(p_i)$ via the multivariate chain rule:

The factor $-\frac{y_i}{p_i}$ has $p_i$ in the denominator, which cancels the $p_i$ in the softmax Jacobian numerator.

Under squared probability error, $\frac{\partial L_{\text{sq}}}{\partial p_i} = (p_i - y_i)$. There is no $p_i$ denominator to cancel the softmax Jacobian:

When the model is confidently wrong (approximately $p_{\text{bug}} = 0.9993$ and $p_{\text{docs}} = 0.0003$), the softmax Jacobian entries are small: the target diagonal is about $0.0003$, and the wrong favorite's is about $0.0007$. The squared-error logit gradient can become tiny, slowing learning in a saturated region. In this fixture, the values remain nonzero.

With cross-entropy, the target logit gradient is approximately $0.0003 - 1.0 = -0.9997$. This local signal remains strong even when the target probability is tiny; gradients farther back still depend on the network's Jacobians.

import math

def softmax(logits: list[float]) -> list[float]:
    peak = max(logits)
    weights = [math.exp(z - peak) for z in logits]
    total = sum(weights)
    return [weight / total for weight in weights]

def squared_probability_loss(logits: list[float], target: list[float]) -> float:
    error = [probability - label for probability, label in zip(softmax(logits), target)]
    return 0.5 * sum(value * value for value in error)

logits = [8.0, 0.0, 0.0]
target = [0.0, 1.0, 0.0]
probabilities = softmax(logits)
ce_gradient = [probability - label for probability, label in zip(probabilities, target)]

epsilon = 1e-5
mse_gradient = []
for index in range(3):
    plus = logits.copy()
    minus = logits.copy()
    plus[index] += epsilon
    minus[index] -= epsilon
    slope = (
        squared_probability_loss(plus, target)
        - squared_probability_loss(minus, target)
    ) / (2 * epsilon)
    mse_gradient.append(slope)

print("probabilities:", [round(value, 4) for value in probabilities])
print("cross-entropy gradient:", [round(value, 4) for value in ce_gradient])
print("squared-probability gradient:", [round(value, 4) for value in mse_gradient])

probabilities: [0.9993, 0.0003, 0.0003]
cross-entropy gradient: [0.9993, -0.9997, 0.0003]
squared-probability gradient: [0.001, -0.0007, -0.0003]

In this fixture, the target logit gradients differ by about three orders of magnitude: roughly $-1.0$ versus $-0.0007$. That stronger gradient explains why cross-entropy can be easier to optimize when the model is confidently wrong. Squared probability error remains a valid classification loss with different gradient behavior.

PyTorch expects unnormalized logits directly

In PyTorch, nn.CrossEntropyLoss expects unnormalized logits, not probabilities.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html}

For integer class targets without smoothing, this loss is equivalent to stable log_softmax followed by negative log-likelihood (NLLLoss). It works on CPU and GPU. Kernel count and intermediate allocation depend on backend and execution mode.^{[5]Reference 5ATen Cross-Entropy and Negative Log-Likelihood Implementationhttps://github.com/pytorch/pytorch/blob/v2.14.0/aten/src/ATen/native/LossNLL.cpp}

If you insert nn.Softmax() before nn.CrossEntropyLoss, PyTorch treats the resulting probabilities $p$ as logits. The mistaken loss is $-\log(\operatorname{softmax}(p)_c)=\log(\sum_j e^{p_j})-p_c$. Valid shapes won't catch this mistake: it changes the objective and the gradients.

import torch
from torch import nn

logits = torch.tensor([[3.0, 1.0, 0.0]], requires_grad=True)
target = torch.tensor([1])  # docs
loss_fn = nn.CrossEntropyLoss()

loss = loss_fn(logits, target)
loss.backward()
probabilities = torch.softmax(logits.detach(), dim=1)

print("probabilities:", [round(value, 3) for value in probabilities[0].tolist()])
print("loss:", round(loss.item(), 3))
print("gradient:", [round(value, 3) for value in logits.grad[0].tolist()])

probabilities: [0.844, 0.114, 0.042]
loss: 2.17
gradient: [0.844, -0.886, 0.042]

Now reproduce the common bug: softmax first, then CrossEntropyLoss. The shapes match, so Python raises no warning. Compare the target gradients.

import torch
from torch import nn

target = torch.tensor([1])  # docs
loss_fn = nn.CrossEntropyLoss()

correct_input = torch.tensor([[3.0, 1.0, 0.0]], requires_grad=True)
correct_loss = loss_fn(correct_input, target)
correct_loss.backward()

wrong_input = torch.tensor([[3.0, 1.0, 0.0]], requires_grad=True)
wrong_loss = loss_fn(torch.softmax(wrong_input, dim=1), target)
wrong_loss.backward()

print("raw logits loss:", round(correct_loss.item(), 3))
print("probabilities passed as logits:", round(wrong_loss.item(), 3))
print("correct docs gradient:", round(correct_input.grad[0, 1].item(), 3))
print("distorted docs gradient:", round(wrong_input.grad[0, 1].item(), 3))

raw logits loss: 2.17
probabilities passed as logits: 1.387
correct docs gradient: -0.886
distorted docs gradient: -0.127

The corrective gradient for docs shrank from $-0.886$ to $-0.127$. Feed raw logits directly into nn.CrossEntropyLoss. You can compute softmax separately for training metrics, distillation targets, or inference probabilities; the mistake is passing those probabilities as the loss's input. For integer targets, use torch.long class indices. Soft probability targets use floating-point tensors matching the input shape; PyTorch doesn't validate that they are normalized distributions.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html}

Temperature scaling shapes inference entropy

Our classification examples train at temperature $T = 1.0$. Other objectives, including distillation, can use a different temperature during training. Scaling fixed logits by a positive temperature $T > 0$ changes how concentrated the distribution is:

Dividing logits by $T$ changes how sharply probabilities concentrate:

• When $T < 1.0$, logit differences expand. As $T \to 0^+$, probability concentrates on the highest-scoring class if the maximum is unique. Tied maxima instead share the limiting mass equally.
• When $T > 1.0$, logit differences shrink. Mass spreads more evenly across competitors. As $T \to \infty$, the distribution approaches uniform uncertainty ($1/K$).

Positive temperature preserves the mathematical ranking. Rescaling fixed logits doesn't update the network's weights, although a calibration procedure can fit temperature on held-out data. Softmax itself is undefined at $T=0$; a serving API may use zero as a convention for greedy decoding. Greedy selection and drawing a sample from a sharpened distribution are different operations.

In distillation, Hinton et al. soften teacher and student distributions with higher temperatures to expose relationships among non-target classes ("dark knowledge"). Their experiments use different temperatures for different setups, so evaluate the value for the specific distillation task.^{[6]Reference 6Distilling the Knowledge in a Neural Network.https://arxiv.org/abs/1503.02531} Cross-entropy through $\operatorname{softmax}(z/T)$ has raw-logit gradient $(p(T)-y)/T$. Distillation commonly multiplies the soft-target objective by $T^2$ to compensate for its high-temperature gradient scaling. During serving, temperature changes sampling diversity; the decoding policy still decides how a token is selected.

import math

def softmax_at_temperature(logits: list[float], temperature: float) -> list[float]:
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("temperature must be finite and positive")
    if not logits or not all(math.isfinite(z) for z in logits):
        raise ValueError("logits must be nonempty and finite")
    peak = max(logits)
    weights = [math.exp((z - peak) / temperature) for z in logits]
    total = sum(weights)
    return [weight / total for weight in weights]

logits = [3.0, 1.0, 0.0]
for temperature in (0.5, 1.0, 2.0):
    probabilities = softmax_at_temperature(logits, temperature)
    print(f"T={temperature:.1f}", [round(probability, 3) for probability in probabilities])

equal_scores = softmax_at_temperature([1e308, 1e308, 1e308], 1e-308)
print("equal huge scores:", [round(p, 3) for p in equal_scores])
assert all(math.isclose(p, 1 / 3) for p in equal_scores)

T=0.5 [0.98, 0.018, 0.002]
T=1.0 [0.844, 0.114, 0.042]
T=2.0 [0.629, 0.231, 0.14]
equal huge scores: [0.333, 0.333, 0.333]

Subtract the peak before dividing by temperature. Dividing 1e308 by 1e-308 first produces infinity, turning a perfectly valid three-way tie into inf - inf and NaN. The shifted calculation keeps each maximum at zero. Temperature only reshapes the logits you already have. It doesn't stop training from pushing the winning score toward infinity. What objective keeps the best gap finite?

Label smoothing curbs overconfident logits

Nothing in the temperature rescaling keeps that gap finite. Consider optimizing the logits for one example without any regularization. What distribution minimizes the loss?

With a hard target $y_{\text{docs}} = 1.0$, minimizing cross-entropy loss $L = -\log(p_{\text{docs}})$ requires driving $p_{\text{docs}} \to 1.0$. In softmax, $p_{\text{docs}} \to 1.0$ only occurs when:

One-hot cross-entropy approaches zero as the target logit gap grows without bound. On separable training data, an unregularized model may keep increasing score gaps after its classification accuracy reaches 100%. Conflicting labels, weight decay, early stopping, and shared parameter constraints can keep weight norms bounded.

Label smoothing changes the objective by mixing the one-hot target with a distribution over classes.^{[7]Reference 7Rethinking the Inception Architecture for Computer Visionhttps://arxiv.org/abs/1512.00567} With a uniform mixture, parameter $\alpha \in (0, 1)$, and $K$ classes:

For our 3-class triage problem with target docs ($c=1$), $\alpha = 0.1$, and $K = 3$:

• Target class: $\tilde{y}_{\text{docs}} = (1 - 0.1) + \frac{0.1}{3} \approx 0.9333$
• Competitor classes: $\tilde{y}_{\text{bug}} = \tilde{y}_{\text{security}} = \frac{0.1}{3} \approx 0.0333$

What logit gap does the model need now to achieve zero gradient? Instead of an infinite score difference, the target probability reaches equilibrium when:

For this example, unrestricted logits attain their optimum at the smoothed distribution, with a finite gap of about $3.33$. If $p_{\text{docs}}>0.9333\ldots$, its logit gradient becomes positive, so a direct gradient-descent step lowers that score. The minimum loss equals the positive entropy of the smoothed target. The objective controls relative scores, while adding a common constant to every logit leaves it unchanged; parameter norms need separate controls.

Label smoothing can improve generalization and calibration, but its benefits depend on the task and smoothing amount. Müller et al. also found that smoothing a teacher's labels could hurt subsequent distillation.^{[8]Reference 8When Does Label Smoothing Help?https://arxiv.org/abs/1906.02629} Evaluate $\alpha=0.1$ on the target task; its effect depends on the label noise and smoothing amount.

PyTorch provides built-in label smoothing via nn.CrossEntropyLoss(label_smoothing=0.1).

import torch
from torch import nn

logits = torch.tensor([[3.0, 1.0, 0.0]], requires_grad=True)
target = torch.tensor([1])  # docs
alpha = 0.1
num_classes = 3

loss_fn = nn.CrossEntropyLoss(label_smoothing=alpha)
loss = loss_fn(logits, target)
loss.backward()

probs = torch.softmax(logits.detach(), dim=1)[0]
smoothed_target = torch.tensor([
    alpha / num_classes,
    (1.0 - alpha) + (alpha / num_classes),
    alpha / num_classes,
])
manual_loss = -torch.sum(smoothed_target * torch.log_softmax(logits.detach(), dim=1)[0])
grad_manual = probs - smoothed_target

print("probabilities:   ", [round(x, 3) for x in probs.tolist()])
print("smoothed target: ", [round(x, 3) for x in smoothed_target.tolist()])
print("pytorch loss:    ", round(loss.item(), 3))
print("manual loss:     ", round(manual_loss.item(), 3))
print("logit gradient:  ", [round(x, 3) for x in logits.grad[0].tolist()])
print("manual gradient: ", [round(x, 3) for x in grad_manual.tolist()])

probabilities:    [0.844, 0.114, 0.042]
smoothed target:  [0.033, 0.933, 0.033]
pytorch loss:     2.137
manual loss:      2.137
logit gradient:   [0.81, -0.819, 0.009]
manual gradient:  [0.81, -0.819, 0.009]

The loss and gradient agree to the displayed precision. The manual calculation uses log_softmax to stay in log space. The same $p-y$ gradient identity works for normalized soft targets; for a malformed target with sum $a$, the gradient would be $a p-y$, which is another reason to validate probability targets.

Next-token prediction extends the loss across sequences

Our triage classifier evaluated one decision. Autoregressive language models evaluate the exact same categorical decision at every sequence position.

Given context $x_1, x_2, \dots, x_t$, a language model emits a hidden state $h_t \in \mathbb{R}^d$. Its output projection produces vocabulary logits $z_t \in \mathbb{R}^{|V|}$. The projection's weights may be tied to the input embedding matrix, but tying isn't required. The observed next token $x_{t+1}$ provides the target.

For $S+1$ observed tokens, there are $S$ next-token prediction positions when we use each token after the first as a target. With all those positions supervised, their mean negative log-likelihood is:

Reuse our numbers with vocabulary ["bug", "docs", "today"] and observed sequence bug docs today. Teacher forcing supplies inputs [bug, docs] and aligned targets [docs, today]:

• Position 1 sees bug and predicts docs ($c=1$) from logits [3.0, 1.0, 0.0]. Target probability is $0.114$, producing $L_1 = 2.170$ nats.
• Position 2 sees bug docs and predicts today ($c=2$) from logits [0.5, 0.0, 2.0]. Target probability is $0.736$, producing $L_2 = 0.306$ nats.
• Sequence mean loss: $L_{\text{seq}} = \frac{2.170 + 0.306}{2} = 1.238$ nats.

View larger

Notice the proportion: position 1 accounts for $2.170 / 2.476 \approx 87.6\%$ of the sequence loss.

Tensor contracts and vocabulary memory scaling

In PyTorch, a language model batch emits logits shaped [B, S, V] (batch size, sequence length, vocabulary size) and targets shaped [B, S]. PyTorch's multi-dimensional cross-entropy expects classes on dimension 1: [B, V, S]. Alternatively, reshape both into 2D tensors: [B * S, V] with targets [B * S].

Vocabulary size can make [B, S, V] expensive. For a worked memory budget, choose $B=16$, $S=4096$, and $V=128256$: logits alone occupy $16\times4096\times128256\times4=33{,}621{,}540{,}864$ bytes in float32, about 33.62 GB (31.31 GiB). Activations, gradients, optimizer state, and other tensors add to training memory.

Chunking the output projection and loss can reduce the largest intermediate. Specialized implementations such as Cut Cross-Entropy combine the projection with tiled log-sum-exp, avoiding a full token-by-vocabulary logit matrix in global GPU memory.^{[9]Reference 9Cut Your Losses in Large-Vocabulary Language Modelshttps://arxiv.org/abs/2411.09009} Its paper also studies gradient filtering, so approximation settings and supported objectives matter. An ordinary call to CrossEntropyLoss receives logits you've already allocated; projection-level memory savings require a specialized implementation.

import torch
from torch import nn

logits = torch.tensor([[[3.0, 1.0, 0.0], [0.5, 0.0, 2.0]]])
targets = torch.tensor([[1, 2]])  # docs, today
loss_fn = nn.CrossEntropyLoss()

class_axis_loss = loss_fn(logits.transpose(1, 2), targets)
flattened_loss = loss_fn(
    logits.reshape(-1, logits.size(-1)),
    targets.reshape(-1),
)

print("logits shape:", tuple(logits.shape))
print("targets shape:", tuple(targets.shape))
print("class-axis loss:", round(class_axis_loss.item(), 3))
print("flattened loss:", round(flattened_loss.item(), 3))

logits shape: (1, 2, 3)
targets shape: (1, 2)
class-axis loss: 1.238
flattened loss: 1.238

The logits in this example are already aligned to next-token targets. Some model APIs accept unshifted labels and perform the shift internally. Check that contract to avoid training on the current token or shifting twice. Causal attention must also mask future input tokens; a loss mask only controls which positions contribute to the objective.

Padding and prompt positions often aren't supervised. For integer targets, ignore_index=-100 removes those positions from both the loss and its mean denominator.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html} Here we add an ignored third position with deliberately extreme scores. Predict whether it changes the loss or receives a gradient.

import torch
from torch.nn import functional as F

logits = torch.tensor([
    [[3.0, 1.0, 0.0], [0.5, 0.0, 2.0], [100.0, -100.0, 0.0]]
], requires_grad=True)
targets = torch.tensor([[1, 2, -100]], dtype=torch.long)
valid_count = (targets != -100).sum()
if valid_count.item() == 0:
    raise ValueError("skip an optimizer step with no supervised targets")

position_losses = F.cross_entropy(
    logits.transpose(1, 2), targets, ignore_index=-100, reduction="none"
)
loss = position_losses.sum() / valid_count
loss.backward()

print("supervised positions:", valid_count.item())
print("position losses:", [round(x, 3) for x in position_losses[0].tolist()])
print("masked mean:", round(loss.item(), 3))
print("ignored gradient:", logits.grad[0, 2].tolist())
assert torch.count_nonzero(logits.grad[0, 2]).item() == 0

supervised positions: 2
position losses: [2.17, 0.306, 0.0]
masked mean: 1.238
ignored gradient: [0.0, 0.0, 0.0]

With no class weights, the default mean pools all non-ignored targets across the batch. Each supervised target contributes equally: a sequence with one loss of 2 and another with three losses of 0 produce a pooled mean of $2/4=0.5$. Giving each sequence equal weight would yield 1. Choose the weighting your objective needs. An all-ignored batch has an empty denominator; PyTorch's mean loss is NaN, so guard it before updating.

ignore_index masks target positions, not a vocabulary class. Suppose a padding token is class 2. At a position whose target is docs (class 1), the padding logit still competes in the softmax denominator, even with ignore_index=2. Compare that loss with explicitly excluding the padding class:

import torch
from torch.nn import functional as F

logits = torch.tensor([[0.0, 1.0, 10.0]], dtype=torch.float64, requires_grad=True)
target = torch.tensor([1])
loss = F.cross_entropy(logits, target, ignore_index=2)
loss.backward()

allowed_logits = logits.detach().clone()
allowed_logits[:, 2] = -torch.inf
restricted_loss = F.cross_entropy(allowed_logits, target)

print(f"ignore_index=2 loss: {loss.item():.3f}")
print(f"padding-class gradient: {logits.grad[0, 2].item():.6f}")
print(f"padding excluded from candidates: {restricted_loss.item():.3f}")
assert logits.grad[0, 2] > 0.99 and restricted_loss < 0.32

ignore_index=2 loss: 9.000
padding-class gradient: 0.999831
padding excluded from candidates: 0.313

The padding class receives a positive gradient because this supervised prediction put almost all its mass there. Excluding it changes the candidate distribution and the objective. Use a vocabulary mask only when the task really forbids that candidate, keep the target available, and leave at least one finite candidate. This fixture uses a hard target without label smoothing; assigning positive target mass to an excluded class instead makes the mathematical loss infinite.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html}

Class weighting changes reduction too. With integer targets and no smoothing, weight[c] multiplies each target loss, and PyTorch's mean divides by the sum of those weights over non-ignored positions. Floating-point probability targets use a different denominator: the number of positions. Even a one-hot float target can therefore differ from its integer form under class weights. Check the loss contract before combining weighted chunk means.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html}

The same reduction can be implemented without PyTorch using standard library loops:

import math

def per_position_cross_entropy(logits: list[list[float]], targets: list[int]) -> list[float]:
    losses = []
    for row, target_index in zip(logits, targets, strict=True):
        peak = max(row)
        log_total = math.log(sum(math.exp(z - peak) for z in row))
        losses.append((peak - row[target_index]) + log_total)
    return losses

logits = [[3.0, 1.0, 0.0], [0.5, 0.0, 2.0]]
targets = [1, 2]
losses = per_position_cross_entropy(logits, targets)

print("position losses:", [round(value, 3) for value in losses])
print("mean loss:", round(sum(losses) / len(losses), 3))

position losses: [2.17, 0.306]
mean loss: 1.238

Diagnostic failure modes in production pipelines

A bad loss usually fails in a recognizable place. Match the symptom to the arithmetic before you change the model:

Start with the failing example's input values, dtypes, shapes, target alignment, and supervised count. Then compare its stable loss and logit gradients with a small fixture. A sampling change can't repair a mislabeled or misaligned training objective.

Try a second target

Experiment with the running example: switch the correct label to security ($c = 2$) while keeping logits [3.0, 1.0, 0.0].

Before running the code, calculate the expected loss and gradient by hand:

• Target probability: $p_{\text{security}} \approx 0.042$.
• Loss: $-\log(0.0420) \approx 3.170$ nats.
• Gradient $p - y$:
  • bug: $0.844 - 0 = +0.844$ (pushed down)
  • docs: $0.114 - 0 = +0.114$ (pushed down)
  • security: $0.042 - 1.0 = -0.958$ (pulled up strongly) Notice the zero-sum balance: $+0.844 + 0.114 - 0.958 = 0.000$. Subtracting $\eta(p - y)$ raises the security logit while suppressing both competitors.