3.3. Loss Functions

A loss function is the scalar objective the training loop drives toward zero. It is also the source of the first gradient that flows backward through the network. RustyML ships 5 losses in rustyml::neural_network::losses. MeanSquaredError and MeanAbsoluteError handle regression. BinaryCrossEntropy handles binary and multi-label classification. CategoricalCrossEntropy and SparseCategoricalCrossEntropy handle mutually exclusive multi-class problems. Each loss is a small, unit-like struct that implements the Loss trait. This lets every loss compose the same way inside a Sequential model. You can also call a loss directly, for evaluation or for a hand-rolled training loop.

All 5 losses carry no per-example state. They are cheap to construct and to copy (Debug + Clone + Copy + PartialEq). They differ in the shapes they accept, how they normalize, and how they stay numerically stable at the edges of their domains. This page states those facts straight from the implementation, so you know what a loss does before you compile it into a model.

3.3.1. The Loss trait and averaging conventions

Every loss implements two methods:

pub trait Loss {
    fn compute_loss(&self, y_true: &Tensor, y_pred: &Tensor) -> Result<f32, Error>;
    fn compute_grad(&self, y_true: &Tensor, y_pred: &Tensor) -> Result<Tensor, Error>;
}

Tensor is an alias for ndarray::ArrayD<f32>, a dynamically dimensioned f32 array. Note the argument order: y_true comes first, then y_pred. Both methods return a Result. Shape and label validation happen up front, instead of the code panicking mid-computation. See Error Handling for the error types. compute_grad returns the gradient of compute_loss with respect to the predictions. The two are always exactly consistent, so the gradient is not an approximation.

The losses do not all normalize the same way, on purpose. The regression and binary losses average over every element of the tensor (y.len()), and treat each output as an independent target. CategoricalCrossEntropy instead sums over the trailing class axis and averages over the prediction sites, the product of every axis before it. For the usual [batch, classes] target, that divisor is the batch size. For the [batch, height, width, classes] output of a channels-last Conv2D softmax head, the divisor is batch * height * width, 1 prediction per pixel. This matches Keras’ default sum_over_batch_size reduction. SparseCategoricalCrossEntropy accepts rank-2 predictions only, so its divisor is always the batch size. Switching a model from, for example, MeanSquaredError to CategoricalCrossEntropy rescales the gradient magnitude by roughly the number of classes. This changes the effective learning rate. If a learning rate that worked well before suddenly diverges or stalls after you change the loss, check the rescaling first. Adjust the rate on the optimizer to compensate.

The input shapes also differ. Getting them wrong is the most common cause of a rejected fit call. The table below lists every contract, checked against each loss’s validation code.

3.3.2. Mean Squared Error

MSE is the default loss for regression. The forward value is mean((y_pred - y_true)^2) over all elements. The gradient is 2 * (y_pred - y_true) / n, where n is the total element count. The factor of 2 comes from the derivative of the squared term. RustyML folds that factor into the gradient, not into the loss, so compute_grad is the exact gradient of compute_loss.

Because the error term is squared, MSE weights large residuals quadratically. A prediction that is off by 10 contributes 100 times more to the loss than one off by 1. MSE is the right choice when large errors are truly worse and your targets cluster around the true value in a roughly Gaussian way. MSE is the wrong choice when your data has heavy-tailed noise or outliers that you do not want the model to fit closely.

use rustyml::neural_network::Tensor;
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::traits::Loss;
use ndarray::Array;

fn main() {
    let mse = MeanSquaredError::new();

    let y_true: Tensor = Array::from_shape_vec(vec![2, 2], vec![1.0_f32, 2.0, 3.0, 4.0])
        .unwrap()
        .into_dyn();
    let y_pred: Tensor = Array::from_shape_vec(vec![2, 2], vec![1.0_f32, 3.0, 5.0, 4.0])
        .unwrap()
        .into_dyn();

    let loss = mse.compute_loss(&y_true, &y_pred).unwrap();
    let grad = mse.compute_grad(&y_true, &y_pred).unwrap();

    println!("MSE loss: {loss:.4}");        // mean of squared diffs
    println!("grad shape: {:?}", grad.shape());
}

If the two shapes differ, both methods return Err(Error::ShapeMismatch { .. }). This stops an ndarray broadcast panic from escaping. For the metrics you report after training, such as R2 and explained variance, see Regression Metrics. The loss here is the training objective, not the evaluation report.

3.3.3. Mean Absolute Error

MAE is the loss that tolerates outliers better than MSE. The forward value is mean(|y_pred - y_true|), so residuals contribute linearly, not quadratically. An outlier that is off by 10 adds 10 times as much to the loss as one off by 1, not 100 times. Use MAE when your targets contain outliers that you do not want to dominate training.

The gradient needs more care. The derivative of |x| is sign(x). The MAE gradient is +1/n where the prediction is above the target and -1/n where it is below. At an exact tie, where y_pred == y_true, RustyML returns 0.0 instead of an arbitrary +1/n or -1/n. This is a deliberate subgradient choice. A perfect prediction gives a true zero gradient. The kink at zero is real. The gradient size does not shrink as the prediction approaches the target, unlike MSE, whose gradient scales with the residual. Plain SGD on MAE can therefore oscillate around the optimum. Pairing MAE with an adaptive optimizer that scales steps by gradient history reduces this oscillation in practice.

The sign logic has a final branch. It returns f32::NAN for any residual that is neither greater than 0, less than 0, nor equal to 0. This branch only fires for a NaN residual. Finite inputs never reach it. If a NaN value has already entered your predictions, MAE does not convert it to a clean number. This matches the library’s policy: surface non-finite values instead of hiding them.

use rustyml::neural_network::Tensor;
use rustyml::neural_network::losses::{MeanAbsoluteError, MeanSquaredError};
use rustyml::neural_network::traits::Loss;
use ndarray::Array;

fn main() {
    // 4 clean points, plus 1 large outlier in the last position.
    let y_true: Tensor = Array::from_shape_vec(vec![5], vec![1.0_f32, 2.0, 3.0, 4.0, 5.0])
        .unwrap()
        .into_dyn();
    let y_pred: Tensor = Array::from_shape_vec(vec![5], vec![1.0_f32, 2.0, 3.0, 4.0, 15.0])
        .unwrap()
        .into_dyn();

    let mse = MeanSquaredError::new();
    let mae = MeanAbsoluteError::new();

    // The 10-unit error adds about 100/5 to MSE, but only about 10/5 to MAE.
    // The outlier dominates MSE. It does not dominate MAE.
    println!("MSE: {:.4}", mse.compute_loss(&y_true, &y_pred).unwrap());
    println!("MAE: {:.4}", mae.compute_loss(&y_true, &y_pred).unwrap());
}

3.3.4. Binary Cross-Entropy

BinaryCrossEntropy fits problems where each output is an independent yes/no decision. This covers 1 binary label per sample, or several independent binary labels per sample (multi-label classification). The forward value is mean(-[y * ln(p) + (1 - y) * ln(1 - p)]) over every element. Here, y is a 0.0/1.0 label and p is a predicted probability. Because it averages over every element, a multi-label output of shape [batch, labels] averages over batch * labels. This is the conventional mean binary cross-entropy.

Get 2 things right. First, the label format: y_true must hold 0.0 or 1.0 floats, not class indices or logits. Second, y_pred must be a probability in (0, 1). This loss has no logits mode. You must squash your network’s output into (0, 1) yourself, in practice with a Sigmoid final layer (see Dense Layers and Activations). This is a real difference from Keras. Keras’ BinaryCrossentropy accepts a from_logits flag. RustyML’s does not, so the Sigmoid layer is mandatory, not optional.

A prediction of exactly 0.0 or 1.0 could make ln(0) and division by zero produce NaN or Inf. To prevent this, the loss clips every probability into [1e-7, 1 - 1e-7] before it takes a log, in both the forward and gradient paths. Consider a prediction of exactly 1.0 for a positive label, a fully confident and correct case. It yields a tiny but finite loss (about 1e-7) and a finite gradient, not a zero loss and a NaN gradient.

use rustyml::neural_network::Tensor;
use rustyml::neural_network::losses::BinaryCrossEntropy;
use rustyml::neural_network::traits::Loss;
use ndarray::Array;

fn main() {
    let bce = BinaryCrossEntropy::new();

    // Labels are 0.0 or 1.0. Predictions are probabilities in (0, 1).
    let y_true: Tensor = Array::from_shape_vec(vec![4], vec![0.0_f32, 1.0, 1.0, 0.0])
        .unwrap()
        .into_dyn();
    let y_pred: Tensor = Array::from_shape_vec(vec![4], vec![0.1_f32, 0.9, 0.8, 0.2])
        .unwrap()
        .into_dyn();

    println!("BCE loss: {:.4}", bce.compute_loss(&y_true, &y_pred).unwrap());

    // Even at the boundary, clipping keeps the loss finite.
    let hard_pred: Tensor = Array::from_shape_vec(vec![4], vec![0.0_f32, 1.0, 1.0, 0.0])
        .unwrap()
        .into_dyn();
    let loss = bce.compute_loss(&y_true, &hard_pred).unwrap();
    assert!(loss.is_finite());
    println!("BCE at boundary predictions: {loss:.6}");
}

3.3.5. Categorical Cross-Entropy

CategoricalCrossEntropy handles mutually exclusive multi-class classification with one-hot targets. The last axis is always the class axis. Every axis before it indexes an independent prediction site. y_true and y_pred of shape [batch, classes] give 1 prediction per sample. The [batch, height, width, classes] output of a channels-last Conv2D softmax head gives 1 prediction per pixel, the per-pixel segmentation case. The loss sums the cross-entropy over the class axis. It then divides by the total number of prediction sites, the product of the leading axes. That divisor is batch at rank 2, and batch * height * width at rank 4.

The loss also requires at least a 2-D, non-empty input. A 1-D input is rejected with Error::InvalidInput. A 1-D tensor has no leading axis, which would leave a divisor of 1 and a softmax over the only axis present. An empty input is rejected with Error::EmptyInput.

The constructor takes 1 boolean argument, from_logits. It controls 2 distinct modes:

CategoricalCrossEntropy::new(false) // y_pred is already a probability distribution
CategoricalCrossEntropy::new(true)  // y_pred is raw logits, the loss applies softmax internally

With from_logits = false, the default, y_pred should be a probability distribution along its last axis, the output of a Softmax layer. Following Keras, the loss first renormalizes each row with y_pred / sum(y_pred, axis=-1). Only then does it clip into [1e-7, 1 - 1e-7], giving -sum(y_true * ln(q_clipped)) / sites.

A softmax row already sums to 1, so this division does not change the loss value. The division is still part of what compute_grad differentiates. The gradient is (sum(y_true) - y_true / q_clipped) / (row_sum * sites). The whole bracket is divided by the row sum. Written in terms of the unnormalized p, this is (sum(y_true) / row_sum - y_true / p) / sites. That is the familiar -y/p term, plus a term that stays constant across each row. A softmax backward pass cancels any row-constant term, so a softmax head trains the same either way. The extra term only matters when this loss reads a head that is not a softmax.

The renormalization has 1 practical benefit. A head whose rows do not quite sum to 1 gets scored as the distribution it implies, instead of being penalized for its scale.

With from_logits = true, the loss treats y_pred as raw, unnormalized logits. It applies softmax itself, using a numerically stable log-softmax. Prefer this mode for training, for 2 reasons.

The first reason is stability. Computing softmax(z) and then ln(...) naively overflows exp for large logits. It also loses precision when it takes the log of an already-clipped near-zero probability. The internal path instead subtracts each prediction site’s maximum before it exponentiates, so exp never overflows. It then computes log_softmax = z - logsumexp(z) directly, so it never logs a clipped probability. This normalization runs within 1 prediction site, never across sites. The pixels of a rank-4 conv head therefore do not compete with each other for probability mass.

The second reason is efficiency. The gradient with respect to the logits collapses to the fused form (softmax(z) - y_true) / sites. This form is cheaper to compute. It also avoids the -y/p division, which explodes when p is near 0.

This has 1 practical consequence. The gradient in logits mode is with respect to the logits. Your network’s final layer must therefore output logits, not probabilities. Use a Linear final activation, with no Softmax layer. If you keep a Softmax layer and also pass from_logits = true, the model applies softmax twice, which corrupts training.

use rustyml::neural_network::Tensor;
use rustyml::neural_network::losses::CategoricalCrossEntropy;
use rustyml::neural_network::traits::Loss;
use ndarray::Array;

fn main() {
    // Probability mode: y_pred must already be a softmax distribution.
    let cce_probs = CategoricalCrossEntropy::new(false);
    let y_true: Tensor = Array::from_shape_vec(
        vec![3, 3],
        vec![1.0_f32, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0],
    )
    .unwrap()
    .into_dyn();
    let probs: Tensor = Array::from_shape_vec(
        vec![3, 3],
        vec![0.8_f32, 0.1, 0.1, 0.2, 0.7, 0.1, 0.1, 0.2, 0.7],
    )
    .unwrap()
    .into_dyn();
    println!("CCE (probs): {:.4}", cce_probs.compute_loss(&y_true, &probs).unwrap());

    // Logits mode: y_pred is raw logits, and the loss applies softmax internally.
    let cce_logits = CategoricalCrossEntropy::new(true);
    let logits: Tensor = Array::from_shape_vec(vec![1, 3], vec![1.0_f32, 2.0, 0.5])
        .unwrap()
        .into_dyn();
    let label: Tensor = Array::from_shape_vec(vec![1, 3], vec![0.0_f32, 1.0, 0.0])
        .unwrap()
        .into_dyn();
    let loss = cce_logits.compute_loss(&label, &logits).unwrap();
    let grad = cce_logits.compute_grad(&label, &logits).unwrap();
    // grad is the fused (softmax(logits) - one_hot) / prediction sites.
    println!("CCE (logits): {loss:.4}, grad shape {:?}", grad.shape());
}

3.3.6. Sparse Categorical Cross-Entropy

SparseCategoricalCrossEntropy computes the same loss as CategoricalCrossEntropy, but it takes integer class labels instead of one-hot vectors. y_true is a [batch, 1] tensor of class indices, stored as f32 and rounded to the nearest integer. y_pred is a [batch, num_classes] tensor of probabilities or logits. The tests confirm the equivalence. For the same predictions, the sparse loss on integer labels and the dense CCE on the matching one-hot encoding agree to floating-point round-off. Their gradients agree too.

The shape rules are strict. Without them, a malformed label would cause an opaque index-out-of-bounds panic. y_pred must be exactly 2-D. A 3-D prediction is rejected with Error::InvalidInput. y_true must be exactly [batch, 1]. A 1-D label vector or a [batch, 2] shape is rejected. This is a real difference from Keras, whose sparse categorical crossentropy takes a 1-D [batch] label. Each label must be finite, non-negative, and strictly less than num_classes. Negative or out-of-range labels come back as Error::InvalidInput. A batch-size mismatch between labels and predictions comes back as Error::DimensionMismatch. Validation happens once, up front. It extracts a Vec<usize> of class indices before any arithmetic runs.

The memory savings grow with the number of classes. A dense one-hot target is [batch, num_classes] of f32, or batch * num_classes * 4 bytes, almost all of it zeros. The sparse label is [batch, 1], or batch * 4 bytes. This is a num_classes-fold reduction on the target tensor, and you never build the one-hot matrix at all. For a vocabulary of tens of thousands of classes, this is the difference between a trivial label tensor and one larger than the predictions themselves. Use the sparse loss when your labels are naturally integers and num_classes is large. Use the dense CCE when you already have one-hot, or soft, targets.

The from_logits flag works the same way as it does for dense CCE. When false, the loss expects per-row probabilities. It renormalizes each row before the clip, exactly as CategoricalCrossEntropy does, so it scores an unnormalized head as the distribution it implies. The row-normalizer is also part of what compute_grad differentiates. Because of this, every class picks up a shared 1 / (batch * row_sum) term in the gradient, not only the labeled class. A softmax head cancels this shared term in its backward pass, the same way it does for dense CCE. When true, the loss expects logits and applies the same stable log-softmax. The fused gradient is softmax with 1 subtracted at the true class, divided by the batch.

The probability path also sums the per-sample losses serially, instead of with a parallel reduction. This keeps the reported loss from drifting with thread scheduling. See Reproducibility and Random Seeds.

use rustyml::neural_network::Tensor;
use rustyml::neural_network::losses::SparseCategoricalCrossEntropy;
use rustyml::neural_network::traits::Loss;
use ndarray::Array;

fn main() {
    let scce = SparseCategoricalCrossEntropy::new(false);

    // Labels are class indices in a [batch, 1] column, not one-hot rows.
    let y_true: Tensor = Array::from_shape_vec(vec![3, 1], vec![0.0_f32, 1.0, 2.0])
        .unwrap()
        .into_dyn();
    let y_pred: Tensor = Array::from_shape_vec(
        vec![3, 3],
        vec![0.8_f32, 0.1, 0.1, 0.2, 0.7, 0.1, 0.1, 0.2, 0.7],
    )
    .unwrap()
    .into_dyn();

    let loss = scce.compute_loss(&y_true, &y_pred).unwrap();
    let grad = scce.compute_grad(&y_true, &y_pred).unwrap();
    println!("SCCE loss: {loss:.4}");        // equals dense CCE on the one-hot equivalent
    println!("grad shape: {:?}", grad.shape());

    // RustyML rejects out-of-range and negative labels, instead of wrapping them silently.
    let bad: Tensor = Array::from_shape_vec(vec![3, 1], vec![0.0_f32, 1.0, 9.0])
        .unwrap()
        .into_dyn();
    assert!(scce.compute_loss(&bad, &y_pred).is_err());
}

3.3.7. Pairing activations with losses

The final-layer activation and the loss form a matched pair. Mismatch them, and you either feed the loss the wrong domain (probabilities where it wants logits, or the reverse), or you apply softmax twice. The table below lists the correct combinations, and why each one works.

The distinction between the two classification rows matters. Sigmoid plus BinaryCrossEntropy fits independent binary outputs, where a sample can belong to several labels at once. Softmax plus CategoricalCrossEntropy fits mutually exclusive classes, where the probabilities across classes sum to 1. Do not use softmax where the labels are independent. Do not use sigmoid where the labels compete.

The model below builds the preferred multi-class setup: a Linear output feeding CategoricalCrossEntropy in logits mode. It has no Softmax layer, because the loss applies the softmax itself.

use rustyml::neural_network::sequential::Sequential;
use rustyml::neural_network::layers::{Activation, Dense};
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::losses::CategoricalCrossEntropy;
use ndarray::Array;

fn main() {
    // 4 samples, 3 features -> 3 classes.
    let x = Array::from_shape_vec(
        vec![4, 3],
        vec![
            0.1_f32, 0.2, 0.7, 0.9, 0.1, 0.0, 0.2, 0.8, 0.1, 0.7, 0.2, 0.1,
        ],
    )
    .unwrap()
    .into_dyn();
    // One-hot targets: classes 2, 0, 1, 0.
    let y = Array::from_shape_vec(
        vec![4, 3],
        vec![
            0.0_f32, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0,
        ],
    )
    .unwrap()
    .into_dyn();

    let mut model = Sequential::new();
    model
        .add(Dense::new(3, 8, Activation::ReLU).unwrap())
        .add(Dense::new(8, 3, Activation::Linear).unwrap()); // logits, no Softmax

    model.compile(
        Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        CategoricalCrossEntropy::new(true), // from_logits = true
    );

    model.fit(&x, &y, 5).unwrap();
    let preds = model.predict(&x).unwrap();
    println!("prediction shape: {:?}", preds.shape());
}

3.3.8. Numerical stability and edge behavior

3 mechanisms keep these losses finite at the edges of their domains. Each guards a different failure.

Probability clipping guards the cross-entropy losses whenever they take probabilities as input. This covers BinaryCrossEntropy, and CategoricalCrossEntropy / SparseCategoricalCrossEntropy in from_logits = false mode. Every probability is clamped into [1e-7, 1 - 1e-7] before any ln. A prediction of exactly 0.0 or 1.0 therefore produces a large but finite loss and a finite gradient, instead of NaN or Inf.

The clip guards against log(0). It does not gate the gradient. Unlike Keras’ autodiff, compute_grad evaluates at the clipped probability. It does not zero out clipped positions. Loss and gradient agree everywhere inside the interval.

For the categorical losses, the row renormalization happens before the clip. A row that sums to zero still yields a non-finite result, exactly as it does in Keras. The tests check this directly: fully confident predictions and worst-case boundary predictions both stay finite.

The stable log-softmax guards the logits paths. Instead of forming softmax(z) and logging it, the loss subtracts each prediction site’s maximum before it exponentiates. It then computes log_softmax = z - logsumexp(z) directly. This makes from_logits = true more than a convenience. It is a more reliable way to train a classifier. It never overflows exp. It never logs a clipped probability. It returns the well-conditioned fused gradient softmax(z) - y. Prefer this mode when you have a choice.

Nothing sanitizes NaN. MAE’s sign branch returns f32::NAN for a NaN residual, and a NaN prediction propagates through the other losses too. A non-finite value stays non-finite, by design, so divergence surfaces loudly instead of hiding. To tame large but finite gradients before they diverge, clip on the optimizer with clip-by-global-norm, instead of clamping inside the loss.

After training, to score a model, use the dedicated evaluators in Classification Metrics and Regression Metrics. A loss is tuned to be a smooth training signal, not a human-readable report.