Introduction

In Keras, a custom loss function should usually return one loss value per sample, not one manually reduced scalar for the whole batch. Keras can then apply its own reduction logic, sample weighting, and masking correctly on top of those per-sample losses.

The Usual Expectation in Keras

When you write a custom loss function with the simple function form:

def custom_loss(y_true, y_pred):
    ...

the normal expectation is that it returns a tensor whose first dimension matches the batch. Keras then reduces that result according to the configured loss reduction behavior.

A simple example:

1import tensorflow as tf
2
3def custom_mse(y_true, y_pred):
4    return tf.reduce_mean(tf.square(y_true - y_pred), axis=-1)

If y_true and y_pred have shape [batch_size, features], the result here has shape [batch_size], which is exactly what Keras typically wants.

Why Per-Sample Losses Are Better

Returning per-sample losses gives Keras room to do several important things correctly:

• average or sum the batch in a consistent way
• apply sample weights
• respect masking in sequence models
• combine multiple output losses predictably

If you reduce to one scalar too early inside the custom function, you take that flexibility away.

That means a loss like this is often too aggressive:

def custom_loss(y_true, y_pred):
    return tf.reduce_mean(tf.square(y_true - y_pred))

This returns one scalar for the entire batch immediately. It can work in simple cases, but it is usually less aligned with how Keras expects losses to behave internally.

A Good Mental Model

Think of a Keras loss in two steps:

1. compute each sample's loss
2. let Keras reduce across the batch

That separation is why so many built-in losses reduce along the feature axis but not across the batch axis.

For example, if each sample is a vector prediction, it is fine for your loss to reduce that vector to one scalar per sample. What you usually should avoid is collapsing the entire batch to one number too early.

When a Batch-Level Scalar Is Reasonable

There are advanced cases where a batch-level scalar is intentional, such as:

• contrastive losses that depend on relationships across samples
• regularizers based on full-batch statistics
• custom training loops where you control reduction explicitly

In those cases, returning one scalar may be exactly what you want. The point is not that scalar batch losses are forbidden. The point is that for ordinary supervised losses, per-sample values are the better default because they fit Keras's training machinery more naturally.

A Full Example

1import tensorflow as tf
2
3def custom_binary_penalty_loss(y_true, y_pred):
4    base_loss = tf.keras.backend.binary_crossentropy(y_true, y_pred)
5    penalty = 0.1 * tf.abs(y_true - y_pred)
6    return tf.reduce_mean(base_loss + penalty, axis=-1)
7
8model = tf.keras.Sequential(
9    [
10        tf.keras.layers.Input(shape=(4,)),
11        tf.keras.layers.Dense(1, activation="sigmoid"),
12    ]
13)
14
15model.compile(optimizer="adam", loss=custom_binary_penalty_loss)

Here the reduction is over feature dimensions, leaving a per-sample loss tensor for Keras to handle across the batch.

Common Pitfalls

• Reducing across the batch inside the loss function when Keras should do that reduction later.
• Returning a tensor with the wrong shape because the wrong axis was reduced.
• Forgetting that sample weighting and masking work better when the loss remains sample-wise.
• Assuming every custom loss must return exactly one scalar no matter the model shape.
• Using a batch-level scalar loss accidentally when the problem does not actually require cross-sample behavior.

Summary

• In standard Keras usage, a custom loss should usually return one loss value per sample.
• Reduce over feature dimensions if needed, but usually do not reduce over the batch dimension yourself.
• Keras then handles batch reduction, weighting, and masking consistently.
• A single scalar batch loss is valid only when your loss genuinely depends on batch-wide structure.
• If you are unsure, return per-sample losses and let Keras do the final aggregation.