Ch 3.5 CrossEntropyLoss: Softmax Loss vs. MNR LossApply

Same PyTorch loss function, two very different goals, inputs, and targets.

Core concepts

  • One loss function, two jobs. Both Softmax loss (Ch 3) and MNR loss (Ch 4) are ultimately computed with PyTorch’s nn.CrossEntropyLoss — but they hand it very different logits and targets.
  • Softmax loss classifies a pair. The logits are 3 class scores (entailment / neutral / contradiction) for one sentence pair; the target is a fixed label.
  • MNR loss classifies a row of a similarity matrix. The logits are the cosine-similarity scores between one anchor and every candidate in the batch; the target is “the positive is at this index.”
  • Why this matters. Recognizing CrossEntropyLoss underneath both objectives demystifies “contrastive loss” — it’s classification, just reframed so the batch itself supplies the classes.

Softmax Loss (Ch 3)

Goal
Classify a sentence pair as entailment / neutral / contradiction, shaping the embedding space as a side effect.
Input to CrossEntropyLoss
logits: shape (batch, 3) — raw scores from the classifier head over concat(u, v, |u−v|).
Target
labels: shape (batch,) — the ground-truth NLI class index (0, 1, or 2) from the dataset.
import torch.nn as nn

# u, v: sentence embeddings from the shared BERT encoder, shape (batch, dim)
features = torch.cat([u, v, torch.abs(u - v)], dim=1)
logits = classifier_head(features)          # (batch, 3)
labels = batch["nli_label"]                 # (batch,) in {0, 1, 2}

loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, labels)

MNR Loss (Ch 4)

Goal
Rank each anchor’s true positive above every other in-batch candidate, directly optimizing retrieval quality.
Input to CrossEntropyLoss
logits: shape (batch, batch) — cosine similarity of every anchor against every positive in the batch.
Target
labels: shape (batch,) — [0, 1, 2, …, batch-1], since each anchor’s positive sits on the diagonal.
import torch
import torch.nn as nn
import torch.nn.functional as F

# a, p: anchor and positive embeddings, shape (batch, dim)
a_norm = F.normalize(a, dim=1)
p_norm = F.normalize(p, dim=1)

logits = a_norm @ p_norm.T * scale          # (batch, batch) similarity matrix
labels = torch.arange(logits.size(0))       # diagonal = correct positive

loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, labels)

The pattern to remember: CrossEntropyLoss always wants (logits, target_class_indices). Softmax loss fixes the number of classes at 3 and reads the label from the dataset; MNR loss lets the batch size define the number of classes and derives the label from position (the diagonal) — no human labelling required.

What you must master

  • Explain that both objectives ultimately reduce to a classification cross-entropy call Level 1
  • Identify the shape and meaning of the logits and targets each loss expects Level 2
  • Read PyTorch training code and recognize which loss (Softmax vs. MNR) is in use Level 2

Architect’s lens

When you review a client’s fine-tuning script, the tell is in what gets passed into CrossEntropyLoss: a small fixed-width classifier head means Softmax loss (older, weaker for retrieval); a batch × batch similarity matrix with arange labels means MNR loss (the modern default). Spotting this in code review tells you instantly whether the embedder was trained for classification or for ranking — and therefore whether it will perform well in your vector search stack.