Introduction

A custom loss function in Keras is useful when standard losses (for example MSE or MAE) do not reflect the true cost of prediction errors in your domain. For multivariate regression in particular, not all output dimensions are independent, and not all dimensions should be penalized equally.

Mahalanobis distance gives you a covariance-aware error metric. Instead of treating each target dimension as orthogonal (as plain Euclidean distance does), it scales and rotates the error by the inverse covariance structure. This is often better when outputs are correlated, have different variance scales, or have known uncertainty geometry.

Mahalanobis Distance

Mathematical Foundation

Mahalanobis distance between vectors x and y with covariance matrix S:

In training losses, we usually minimize the squared form (drop the square root) for smoother gradients and lower computational overhead:

To keep the matrix invertible in practice, add regularization:

Then use S_lambda^{-1} in the loss.

When This Beats MSE

Mahalanobis-style losses are especially helpful when:

• Anomaly Detection: Outliers get larger distance under learned covariance.
• Correlated Outputs: Errors along high-correlation directions are weighted correctly.
• Heteroscedastic Targets: High-variance target dimensions do not dominate loss unfairly.
• Physics/Geometry Constraints: Distance in transformed statistical space matters more than raw coordinate error.

Keras Custom Loss Function

The standard setup is:

1. Estimate covariance from target data (or residuals).
2. Regularize and invert it.
3. Build a custom loss that computes the quadratic form per sample.
4. Train with normal Keras compile/fit flow.

Step 1: Import Libraries

1import tensorflow as tf
2from tensorflow import keras
3from tensorflow.keras import layers
4import numpy as np

Step 2: Estimate Inverse Covariance

1def estimate_inv_cov(y_train, reg_lambda=1e-3):
2    """
3    y_train: numpy array with shape [n_samples, d]
4    returns inv_cov with shape [d, d]
5    """
6    y_train = np.asarray(y_train, dtype=np.float32)
7    cov = np.cov(y_train, rowvar=False)  # [d, d]
8    cov = cov + reg_lambda * np.eye(cov.shape[0], dtype=np.float32)
9    inv_cov = np.linalg.inv(cov).astype(np.float32)
10    return inv_cov

If your targets are time-dependent or regime-dependent, you can estimate covariance per segment and train separate models, or use a dynamic weighting strategy.

Step 3: Build a Mahalanobis Loss Factory

1def mahalanobis_loss_factory(inv_cov):
2    # inv_cov shape: [d, d]
3    inv_cov = tf.convert_to_tensor(inv_cov, dtype=tf.float32)
4
5    def loss(y_true, y_pred):
6        # diff: [batch, d]
7        diff = y_pred - y_true
8        # quadratic form per sample: diff^T * inv_cov * diff
9        # einsum result shape: [batch]
10        quad = tf.einsum("bi,ij,bj->b", diff, inv_cov, diff)
11        return tf.reduce_mean(quad)
12
13    return loss

Step 4: Create and Compile a DNN Model

1# Example: 16 input features -> 3 regression outputs
2inputs = keras.Input(shape=(16,))
3x = layers.Dense(64, activation="relu")(inputs)
4x = layers.BatchNormalization()(x)
5x = layers.Dense(32, activation="relu")(x)
6outputs = layers.Dense(3)(x)
7model = keras.Model(inputs, outputs)
8
9# inv_cov = estimate_inv_cov(y_train, reg_lambda=1e-3)
10# For illustration:
11inv_cov = np.eye(3, dtype=np.float32)
12
13model.compile(
14    optimizer=keras.optimizers.Adam(1e-3),
15    loss=mahalanobis_loss_factory(inv_cov),
16    metrics=[
17        keras.metrics.MeanAbsoluteError(name="mae"),
18        keras.metrics.RootMeanSquaredError(name="rmse"),
19    ],
20)

Step 5: Train

1callbacks = [
2    keras.callbacks.EarlyStopping(
3        monitor="val_loss",
4        patience=8,
5        restore_best_weights=True,
6    )
7]
8
9history = model.fit(
10    X_train,
11    y_train,
12    validation_data=(X_val, y_val),
13    epochs=80,
14    batch_size=32,
15    callbacks=callbacks,
16)

CNN Variant for Regression

For image-to-vector regression, only the backbone changes. The loss can stay the same.

1img_in = keras.Input(shape=(128, 128, 3))
2z = layers.Conv2D(32, 3, activation="relu")(img_in)
3z = layers.MaxPooling2D()(z)
4z = layers.Conv2D(64, 3, activation="relu")(z)
5z = layers.MaxPooling2D()(z)
6z = layers.Conv2D(128, 3, activation="relu")(z)
7z = layers.GlobalAveragePooling2D()(z)
8z = layers.Dense(64, activation="relu")(z)
9img_out = layers.Dense(3)(z)  # same output dimension d=3
10
11cnn_model = keras.Model(img_in, img_out)
12cnn_model.compile(
13    optimizer=keras.optimizers.Adam(1e-3),
14    loss=mahalanobis_loss_factory(inv_cov),
15    metrics=[keras.metrics.MeanAbsoluteError(name="mae")],
16)

Numerical Stability and Performance

For high output dimensions, direct inverse can be noisy. Prefer stable decomposition methods where possible:

• Compute S_lambda with regularization.
• Use Cholesky decomposition.
• Solve linear systems instead of explicitly inverting matrices.

If training becomes slow, profile the custom loss and consider:

• Lower output dimension via learned projection.
• Mixed precision where safe.
• Precomputing static covariance once per training run.

Common Failure Modes

1. Singular Covariance Matrix
   Symptoms: NaNs, exploding loss, inversion errors.
   Fix: add stronger diagonal regularization (lambda * I).
2. Incorrect Target Shape
   Symptoms: einsum shape errors.
   Fix: enforce [batch, d] target and prediction shapes.
3. Scale Instability
   Symptoms: very large gradients, noisy convergence.
   Fix: normalize targets before covariance estimation.
4. Overfitting with Complex Backbones
   Symptoms: training loss down, validation loss up.
   Fix: dropout, L2 regularization, early stopping.
5. Mismatch Between Train and Inference Pipelines
   Symptoms: good offline metrics, bad production behavior.
   Fix: identical preprocessing, target transforms, and postprocessing.

Model Saving and Loading

Because this is a custom loss, load with custom_objects:

1custom_loss = mahalanobis_loss_factory(inv_cov)
2model.save("mahalanobis_model.keras")
3loaded = keras.models.load_model(
4    "mahalanobis_model.keras",
5    custom_objects={"loss": custom_loss}
6)

For production workflows, package covariance and loss creation together so loading is deterministic.

Practical Checklist

• Covariance Regularization: If covariance is near-singular, add lambda * I before inversion.
• Scale Sensitivity: Standardize targets/features before estimating covariance.
• Target Dimension Consistency: Keep output layer size equal to covariance dimension d.
• Metric Clarity: Track MAE/RMSE alongside custom loss for interpretability.
• Validation Discipline: Validate on held-out data from production-like distribution.