Cross-Entropy Loss

Log-softmax, logsumexp trick

Easy Fundamentals

Problem Description

Implement cross-entropy loss from scratch.

$$\text{CE}(x, y) = -\log\frac{e^{x_y}}{\sum_j e^{x_j}}$$

Signature

def cross_entropy_loss(logits: Tensor, targets: Tensor) -> Tensor:
    # logits: (B, C) float, targets: (B,) long indices
    # Returns: scalar loss (mean over batch)

Rules

• Do NOT use F.cross_entropy or nn.CrossEntropyLoss

• Must be numerically stable (use logsumexp trick)

Template

Implement the function below. Use only basic PyTorch operations.

# ✏️ YOUR IMPLEMENTATION HERE

def cross_entropy_loss(logits, targets):
    pass  # log_probs = logits - logsumexp(...)

Test Your Implementation

Use this code to debug before submitting.

# 🧪 Debug
logits = torch.randn(4, 10)
targets = torch.randint(0, 10, (4,))
print('Loss:', cross_entropy_loss(logits, targets))
print('Ref: ', torch.nn.functional.cross_entropy(logits, targets))

Reference Solution

Try solving it yourself first! Click below to reveal the solution.

# ✅ SOLUTION

def cross_entropy_loss(logits, targets):
    log_probs = logits - torch.logsumexp(logits, dim=-1, keepdim=True)
    return -log_probs[torch.arange(targets.shape[0]), targets].mean()

Tips

Run Locally

For interactive practice with auto-grading, run TorchCode locally:
pip install torch-judge then use check("cross_entropy")

Key Concepts

Log-softmax, logsumexp trick

Softmax Dropout

Cross-Entropy Loss

Description Template Test Solution Tips