Loss functions

Deep dive

How it works

In: model outputs (logits [N, K] or predictions [N]) and targets · Out: a scalar loss and its gradient w.r.t. the outputs · State: none · Cost: O(N·K), memory-bound on the logits tensor

1. Log-sum-exp. logsumexp(z) = m + log Σ exp(z_i − m) with m = max z. The largest term becomes exp(0) = 1, so the sum is in [1, K] and never overflows or underflows to zero.
2. Fusing loss and softmax. Cross-entropy and softmax share the normalizer. Computing them together gives the gradient p − y with no division, and the softmax backward never runs on its own.
3. Soft targets. Label smoothing replaces the one-hot target with (1 − ε)·onehot + ε/K (PyTorch spreads ε over all K classes, including the true one). The gradient is still p − target. Distillation uses the teacher's softmax at temperature T as the target and multiplies the loss by T² so gradient scale does not shrink with T.
4. Reduction. Average over non-ignored tokens, not over the padded tensor. Under gradient accumulation, divide by the total token count across micro-batches, not per micro-batch.

Code

/// Row-wise log-softmax of logits z[n*k] in place. Max subtraction keeps exp() in (0, 1].
fn log_softmax(z: &mut [f32], k: usize) {
    for row in z.chunks_mut(k) {
        let m = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let s: f32 = row.iter().map(|&v| (v - m).exp()).sum();
        let lse = m + s.ln();
        row.iter_mut().for_each(|v| *v -= lse);
    }
}

/// Mean softmax cross-entropy over non-ignored rows; writes dL/dz into `grad`.
pub fn cross_entropy(z: &[f32], y: &[i64], k: usize, smoothing: f32, ignore: i64, grad: &mut [f32]) -> f32 {
    let mut logp = z.to_vec();
    log_softmax(&mut logp, k);
    let count = y.iter().filter(|&&t| t != ignore).count().max(1) as f32; // mean over real tokens
    let mut loss = 0.0f64;
    for (r, &t) in y.iter().enumerate() {
        let (lp, g) = (&logp[r * k..(r + 1) * k], &mut grad[r * k..(r + 1) * k]);
        if t == ignore { g.fill(0.0); continue; }
        for c in 0..k {
            let target = smoothing / k as f32 + if c as i64 == t { 1.0 - smoothing } else { 0.0 };
            loss -= (target * lp[c]) as f64;
            g[c] = (lp[c].exp() - target) / count;              // softmax − target
        }
    }
    (loss / count as f64) as f32
}

/// Binary CE on logits: max(x,0) − x·y + ln(1 + e^−|x|). Finite for any x.
pub fn bce_with_logits(x: &[f32], y: &[f32], grad: &mut [f32]) -> f32 {
    let n = x.len() as f32;
    let mut loss = 0.0;
    for i in 0..x.len() {
        let e = (-x[i].abs()).exp();
        loss += x[i].max(0.0) - x[i] * y[i] + e.ln_1p();
        let sig = if x[i] >= 0.0 { 1.0 / (1.0 + e) } else { e / (1.0 + e) }; // no overflow either side
        grad[i] = (sig - y[i]) / n;
    }
    loss / n
}

pub fn huber(pred: &[f32], y: &[f32], delta: f32, grad: &mut [f32]) -> f32 {
    let n = pred.len() as f32;
    let mut loss = 0.0;
    for i in 0..pred.len() {
        let r = pred[i] - y[i];
        loss += if r.abs() <= delta { 0.5 * r * r } else { delta * (r.abs() - 0.5 * delta) };
        grad[i] = r.clamp(-delta, delta) / n;
    }
    loss / n
}

/// Focal loss for one logit x and label y ∈ {0,1}.
pub fn focal(x: f32, y: f32, gamma: f32, alpha: f32) -> f32 {
    let log_pt = -(x.max(0.0) - x * y + (-x.abs()).exp().ln_1p()); // stable ln p_t
    let a_t = if y == 1.0 { alpha } else { 1.0 - alpha };
    -a_t * (1.0 - log_pt.exp()).powf(gamma) * log_pt
}

fn main() {
    let (n, k) = (4, 7);
    let mut z: Vec<f32> = (0..n * k).map(|i| ((i * 37 % 13) as f32 - 6.0) * 0.8).collect();
    let y = [0i64, 3, -100, 6];
    let mut g = vec![0.0; n * k];
    let loss = cross_entropy(&z, &y, k, 0.1, -100, &mut g);
    let mut scratch = vec![0.0; n * k];
    for i in [0usize, 10, 26] {                                  // central finite differences
        let h = 1e-2;
        z[i] += h; let lp = cross_entropy(&z, &y, k, 0.1, -100, &mut scratch);
        z[i] -= 2.0 * h; let lm = cross_entropy(&z, &y, k, 0.1, -100, &mut scratch);
        z[i] += h;
        assert!(((lp - lm) / (2.0 * h) - g[i]).abs() < 1e-3);
    }
    let zero = vec![0.0f32; 1000];
    let mut g1000 = vec![0.0; 1000];
    println!("CE {loss:.4}; zero logits, 1000 classes: {:.4} (= ln 1000)",
             cross_entropy(&zero, &[0], 1000, 0.0, -100, &mut g1000));
    let mut g2 = vec![0.0; 2];
    println!("CE at logit 1e4 stays finite: {}", cross_entropy(&[1e4, 0.0], &[1], 2, 0.0, -100, &mut g2));
    let mut gb = vec![0.0; 3];
    println!("BCE at x = ±200: {}", bce_with_logits(&[200.0, -200.0, 0.0], &[0.0, 1.0, 1.0], &mut gb));
    let mut gh = vec![0.0; 2];
    huber(&[5.0, 0.3], &[0.0, 0.0], 1.0, &mut gh);
    assert_eq!(gh, vec![0.5, 0.15]);                             // clipped at δ, divided by n
    println!("easy positive: focal {:.2e} vs bce {:.2e}", focal(5.0, 1.0, 2.0, 0.25),
             bce_with_logits(&[5.0], &[1.0], &mut [0.0]));
}

Complexity and constants

• Forward plus backward: about 3 passes over the N×K logits (max, exp-sum, gradient). Memory-bound: at a few TB/s, a 16.8 GB fp32 logit tensor costs order of 5-10 ms per pass on one GPU.
• The output projection (d × K matmul) costs far more FLOPs than the loss. The loss costs memory.
• Chunked cross-entropy computes logits and gradient for 1-8k tokens at a time and never materializes the full N×K tensor. That cuts loss memory by about N/chunk.

Efficiency

• Cost: forward plus backward is about 3 passes over the N×K logits, memory-bound. A 32k-token × 128k-vocab fp32 logit tensor is 16.8 GB (above), so one pass at about 3 TB/s HBM is order of 5-6 ms. The output matmul costs far more FLOPs than the loss; the loss costs memory.
• Levers:
  • Chunk the loss over 1-8k tokens — logit memory drops by about N/chunk (order of) — large vocabularies.
  • Fused linear+cross-entropy kernels — Cut Cross-Entropy cut the Gemma 2 (2B) loss memory from 24 GB to 1 MB (Wijmans et al. 2024) — vocab ≥ 100k.
  • Single-pass max and exp-sum — one fewer read of each row (order of) — every softmax loss.
  • Average over non-ignored tokens and skip padded rows — saves the padding fraction (order of) — variable-length batches.
  • Sampled softmax / NCE — O(samples) instead of O(K) — recommendation, retrieval with millions of classes.
• Naive — softmax, then log, then NLL as separate kernels, N×K probabilities materialized · 3+ full reads and writes of the tensor
• Standard — fused log-softmax + NLL (F.cross_entropy) · about 3 passes, N×K logits resident
• State of the art — fused chunked linear+loss that never materializes the logits · loss memory O(chunk×K) or less

Implementation notes

• Run the loss in fp32 even under bf16 training. bf16 keeps 8 significant bits (2-3 decimal digits), too coarse for a 128k-term exp-sum. See mixed precision.
• Fused kernels (Liger kernel, Apple's Cut Cross-Entropy) compute the loss and the logits gradient inside the output matmul, so the N×K tensor never reaches HBM.
• PaLM's z-loss adds 1e-4 · logsumexp(z)² to keep the normalizer near 0 and stop logit drift in bf16.
• For BCE, keep the inputs as logits (BCEWithLogitsLoss), never sigmoid then BCELoss: log(sigmoid(x)) for x = −100 gives −inf in fp32.
• Class imbalance: weight per class by about inverse frequency, or use focal loss. RetinaNet found γ = 2, α = 0.25 best.
• Mask padding with ignore_index, not a zero target vector. A zero target still pulls logits through the normalizer.

Pitfalls

• Applying softmax in the model and then CrossEntropyLoss, which applies log-softmax again. Training still runs, but gradients are squashed.
• Averaging per micro-batch under gradient accumulation with variable-length sequences: short sequences get over-weighted.
• MSE on targets of wildly different scales: one output dominates. Normalize targets or weight per output.
• Label smoothing hurts calibration-sensitive uses and distillation. Teachers trained with it transfer worse (Müller et al., 2019).
• Huber δ in raw target units: δ = 1 on targets in the thousands is just L1.

Variants and alternatives

In production

torch.nn.functional.cross_entropy (fused log-softmax plus NLL), binary_cross_entropy_with_logits, huber_loss; focal loss in object detectors (RetinaNet); distillation losses (Hinton et al.). JAX/Optax softmax_cross_entropy_with_integer_labels. Liger kernel LigerFusedLinearCrossEntropyLoss. XGBoost and LightGBM objectives take gradient and Hessian of these losses directly.

Implementations

Sources

• Lin, Goyal, Girshick, He, Dollár, Focal Loss for Dense Object Detection, ICCV 2017. paper
• Szegedy, Vanhoucke, Ioffe, Shlens, Wojna, Rethinking the Inception Architecture for Computer Vision, CVPR 2016 (label smoothing). paper
• Hinton, Vinyals, Dean, Distilling the Knowledge in a Neural Network, 2015. paper
• Müller, Kornblith, Hinton, When Does Label Smoothing Help?, NeurIPS 2019. paper
• Chowdhery et al., PaLM: Scaling Language Modeling with Pathways, 2022 (z-loss). paper
• Wijmans et al., Cut Your Losses in Large-Vocabulary Language Models, 2024 (Cut Cross-Entropy). paper · code
• Hsu et al., Liger Kernel: Efficient Triton Kernels for LLM Training, 2024. paper · code