PyTorch · Loss Functions

CrossEntropyLoss vs LogSoftmax + NLLLoss

A stage-by-stage mathematical comparison of two equivalent ways to compute the same loss.

Assume the neural network predicts 3 classes.

True class = Dog (Class 1)

One-hot label

y = \begin{bmatrix} 0 \\ 1 \\ 0 \end{bmatrix}

Raw output (logits) from the network

z = \begin{bmatrix} 0.9 \\ 2.6 \\ 0.1 \end{bmatrix}

Stage-by-Stage Comparison

Stage	CrossEntropyLoss Approach	LogSoftmax + NLLLoss Approach
Output Layer	No activation	LogSoftmax(dim=1)
Network Output	Raw logits $z = [0.9, 2.6, 0.1]$	Log probabilities $\log(p)$
Meaning of Output	Arbitrary scores (not probabilities)	Logarithm of probabilities
Mathematical Form	$z_i = W_i x + b_i$	$\log(p_i) = z_i - \log\left(\sum_j e^{z_j}\right)$
Numerical Output	$[0.9, 2.6, 0.1]$	$[-1.931, -0.234, -2.734]$
Loss Function	`CrossEntropyLoss()`	`NLLLoss()`
What Loss Receives	Raw logits	Log probabilities
Internal Computation	LogSoftmax + NLLLoss	Only NLLLoss
Formula Used by Loss	$L = -\log\left(\dfrac{e^{z_y}}{\sum_j e^{z_j}}\right)$	$L = -\log(p_y)$
Calculation	$-\log(13.46 / 17.025)$	$-(-0.234)$
Loss Value	$0.234$	$0.234$
Output During Inference	Raw logits	Log probabilities
Convert to Probabilities	Softmax	Exponential
Formula	$p_i = \dfrac{e^{z_i}}{\sum_j e^{z_j}}$	$p_i = e^{\log(p_i)}$
Recovered Probabilities	$[0.145, 0.791, 0.064]$	$[0.145, 0.791, 0.064]$
Prediction	$\arg\max(\text{logits}) = \text{Class } 1$	$\arg\max(\text{log probabilities}) = \text{Class } 1$
Probability of Prediction	Softmax → 79.1%	exp() → 79.1%
Recommended?	✅ Standard PyTorch practice	Used mainly for learning or specialized models

What Happens Internally?

Approach 1: CrossEntropyLoss

Input
   │
   ▼
Neural Network
   │
   ▼
Raw Logits
[0.9, 2.6, 0.1]
   │
   ▼
CrossEntropyLoss
   │
   ├── LogSoftmax
   │
   ├── NLLLoss
   │
   ▼
Loss = 0.234

Mathematically,

z = \begin{bmatrix} 0.9 \\ 2.6 \\ 0.1 \end{bmatrix} \;\;\xrightarrow{\text{Softmax}}\;\; \begin{bmatrix} 0.145 \\ 0.791 \\ 0.064 \end{bmatrix} \;\;\xrightarrow{\text{Cross Entropy}}\;\; -\sum_i y_i \log(p_i) \;\;\longrightarrow\;\; -\log(0.791) = 0.234

Approach 2: LogSoftmax + NLLLoss

Input
   │
   ▼
Neural Network
   │
   ▼
Raw Logits
[0.9, 2.6, 0.1]
   │
   ▼
LogSoftmax
   │
   ▼
Log Probabilities
[-1.931, -0.234, -2.734]
   │
   ▼
NLLLoss
   │
   ▼
Loss = 0.234

Mathematically,

\text{Raw logits} = \begin{bmatrix} 0.9 \\ 2.6 \\ 0.1 \end{bmatrix} \;\;\xrightarrow{\text{LogSoftmax}}\;\; \log\left(\frac{e^{z_i}}{\sum_j e^{z_j}}\right) = \begin{bmatrix} -1.931 \\ -0.234 \\ -2.734 \end{bmatrix} \;\;\xrightarrow{\text{NLLLoss}}\;\; L = -\log(p_y) \;\;\longrightarrow\;\; L = -(-0.234) = 0.234

The Big Picture

Quantity	Raw Logits	Softmax Output	LogSoftmax Output
Class 0	0.9	0.145	-1.931
Class 1	2.6	0.791	-0.234
Class 2	0.1	0.064	-2.734
Sum	3.6 (meaningless)	1.0	Not constrained
Interpretation	Scores	Probabilities	Log probabilities

A Simple Way to Remember

               SAME NEURAL NETWORK
                      │
                Raw Logits (z)
             [0.9, 2.6, 0.1]
               /             \
              /               \
     CrossEntropyLoss      LogSoftmax
      (does LogSoftmax)         │
              │                 │
              ▼                 ▼
      Computes Loss      Log Probabilities
              │                 │
              ▼                 ▼
           Loss=0.234       NLLLoss
                                  │
                                  ▼
                              Loss=0.234

Core Concepts

Think of CrossEntropyLoss as a "2-in-1 package":

\text{CrossEntropyLoss} = \text{LogSoftmax} + \text{NLLLoss}

So you have two equivalent implementation choices:

Option 1 (recommended): Network → Logits → CrossEntropyLoss
Option 2: Network → LogSoftmax → Log Probabilities → NLLLoss

Both optimize the same objective, produce the same loss ($0.234$), and lead to the same trained model. The only difference is where the LogSoftmax computation takes place.

Gradient Equivalence

The loss values matching is only half the story. The deeper reason these two approaches are treated as interchangeable is that they produce the exact same gradient with respect to the logits during backpropagation:

\frac{\partial L}{\partial z_i} = p_i - y_i

where $p$ is the softmax probability vector and $y$ is the one-hot true label. For our example:

\frac{\partial L}{\partial z} = [0.145,\ 0.791,\ 0.064] - [0,\ 1,\ 0] = [0.145,\ -0.209,\ 0.064]

This clean form — prediction − target — falls out of the chain rule whether you compute it via CrossEntropyLoss directly on logits, or via LogSoftmax followed by NLLLoss. Both paths differentiate down to the same gradient, which is what actually drives identical weight updates and identical trained models.

Why CrossEntropyLoss Is Preferred

If the math is identical, why does PyTorch's documentation and most codebases favor CrossEntropyLoss over manually chaining LogSoftmax + NLLLoss? The answer is numerical stability, not mathematics.

Computing softmax directly involves $e^{z_i}$ terms. When logits are large (e.g. $z_i = 1000$), $e^{1000}$ overflows to infinity in floating point, and the resulting probabilities become NaN. CrossEntropyLoss avoids this by internally fusing the log and the sum using the log-sum-exp trick:

\log\sum_j e^{z_j} = m + \log\sum_j e^{z_j - m}, \quad \text{where } m = \max_j z_j

Subtracting the max logit $m$ before exponentiating keeps every term in the sum bounded between 0 and 1, eliminating overflow entirely — while producing a mathematically identical result. This fusion is only fully available when LogSoftmax and NLLLoss are computed together as one operation, which is exactly what CrossEntropyLoss does internally.

In short: LogSoftmax + NLLLoss is correct and useful for understanding the mechanics, but CrossEntropyLoss is the numerically safe, production-recommended choice — same math, more robust computation.

PyTorch Example: Verifying the Equivalence

The snippet below builds the exact example from this article — logits $z = [0.9, 2.6, 0.1]$, true class 1 — and confirms both approaches give the same loss and the same gradient.

import torch
import torch.nn as nn

# Same raw logits used throughout this article
logits = torch.tensor([[0.9, 2.6, 0.1]], requires_grad=True)
true_class = torch.tensor([1])  # Dog = Class 1

# ---------- Approach 1: CrossEntropyLoss ----------
logits_a = logits.clone().detach().requires_grad_(True)
ce_loss_fn = nn.CrossEntropyLoss()
loss_a = ce_loss_fn(logits_a, true_class)
loss_a.backward()

# ---------- Approach 2: LogSoftmax + NLLLoss ----------
logits_b = logits.clone().detach().requires_grad_(True)
log_softmax = nn.LogSoftmax(dim=1)
nll_loss_fn = nn.NLLLoss()
log_probs = log_softmax(logits_b)
loss_b = nll_loss_fn(log_probs, true_class)
loss_b.backward()

# ---------- Compare ----------
print(f"CrossEntropyLoss:        {loss_a.item():.4f}")
print(f"LogSoftmax + NLLLoss:    {loss_b.item():.4f}")
print(f"Losses match:            {torch.allclose(loss_a, loss_b)}")

print(f"\nGradient (CrossEntropyLoss):     {logits_a.grad}")
print(f"Gradient (LogSoftmax+NLLLoss):   {logits_b.grad}")
print(f"Gradients match:                 {torch.allclose(logits_a.grad, logits_b.grad)}")

# ---------- Expected output ----------
# CrossEntropyLoss:        0.2350
# LogSoftmax + NLLLoss:    0.2350
# Losses match:            True
#
# Gradient (CrossEntropyLoss):     tensor([[ 0.1451, -0.2088,  0.0637]])
# Gradient (LogSoftmax+NLLLoss):   tensor([[ 0.1451, -0.2088,  0.0637]])
# Gradients match:                 True

Both branches converge on the same loss (≈0.235) and the same gradient (p − y), confirming the equivalence proven above — not just in theory, but in what PyTorch actually computes.