Introduction

NaN loss in a custom Keras loss function is usually caused by numerical instability, invalid input values, or gradient explosion. It is rarely random. The fastest fix is to isolate loss math, add finite-value checks, and stabilize operations that involve logs, exponentials, and division.

Common Sources of NaN in Custom Losses

Most NaN failures come from a small set of patterns:

• 'log(0) or log(negative)'
• division by zero or near-zero values
• overflow in exp operations
• incompatible dtypes and implicit casts
• unstable gradients from too-large learning rates

Because custom losses run every step, small numeric flaws become catastrophic quickly.

Reproduce the Problem in a Minimal Example

Here is a deliberately unstable loss:

1import tensorflow as tf
2
3@tf.function
4def unstable_loss(y_true, y_pred):
5    # This can produce log(0)
6    return tf.reduce_mean(tf.math.log(y_pred))
7
8x = tf.random.normal((64, 4))
9y = tf.cast(tf.random.uniform((64, 1), maxval=2, dtype=tf.int32), tf.float32)
10
11model = tf.keras.Sequential([
12    tf.keras.layers.Input(shape=(4,)),
13    tf.keras.layers.Dense(1, activation="sigmoid")
14])
15
16model.compile(optimizer="adam", loss=unstable_loss)
17model.fit(x, y, epochs=1, verbose=0)

This demonstrates how unsafe math can trigger NaN immediately.

Stabilize Loss Math with Clipping and Epsilon

Numerical guards are the first step.

1import tensorflow as tf
2
3@tf.function
4def stable_binary_loss(y_true, y_pred):
5    eps = tf.keras.backend.epsilon()
6    y_true = tf.cast(y_true, tf.float32)
7    y_pred = tf.cast(y_pred, tf.float32)
8    y_pred = tf.clip_by_value(y_pred, eps, 1.0 - eps)
9
10    loss = -(y_true * tf.math.log(y_pred) +
11             (1.0 - y_true) * tf.math.log(1.0 - y_pred))
12    out = tf.reduce_mean(loss)
13
14    tf.debugging.assert_all_finite(out, "Loss has NaN or Inf")
15    return out

Clipping ensures dangerous domains are avoided.

Validate the Loss in Isolation

Test custom loss on known tensors before full training.

1import tensorflow as tf
2
3y_true = tf.constant([[1.0], [0.0]], dtype=tf.float32)
4y_pred = tf.constant([[0.9], [0.2]], dtype=tf.float32)
5
6print(stable_binary_loss(y_true, y_pred).numpy())

If this produces finite outputs, your core formula is probably safe.

Stabilize Training Dynamics

Even a stable formula can become NaN during training if optimizer settings are too aggressive.

1import tensorflow as tf
2
3opt = tf.keras.optimizers.Adam(learning_rate=1e-4, clipnorm=1.0)
4
5model = tf.keras.Sequential([
6    tf.keras.layers.Input(shape=(4,)),
7    tf.keras.layers.Dense(16, activation="relu"),
8    tf.keras.layers.Dense(1, activation="sigmoid")
9])
10
11model.compile(optimizer=opt, loss=stable_binary_loss)

Gradient clipping and lower learning rates often eliminate late-epoch NaN explosions.

Check Data Pipeline for Invalid Values

NaNs in input features or labels will propagate through the model.

1import numpy as np
2
3x_np = np.random.randn(128, 4).astype("float32")
4y_np = np.random.randint(0, 2, size=(128, 1)).astype("float32")
5
6print("x has NaN:", np.isnan(x_np).any())
7print("x has Inf:", np.isinf(x_np).any())
8print("y has NaN:", np.isnan(y_np).any())

Validate this immediately after dataset loading and after each major preprocessing transform.

Inspect Gradients During Debugging

If NaN appears after some training steps, gradient magnitude tracking helps identify instability.

1import tensorflow as tf
2
3x = tf.random.normal((32, 4))
4y = tf.random.uniform((32, 1), maxval=2, dtype=tf.int32)
5y = tf.cast(y, tf.float32)
6
7with tf.GradientTape() as tape:
8    preds = model(x, training=True)
9    loss = stable_binary_loss(y, preds)
10
11grads = tape.gradient(loss, model.trainable_variables)
12max_grad = max(float(tf.reduce_max(tf.abs(g)).numpy()) for g in grads if g is not None)
13print("max grad:", max_grad)

Very large gradient values are a strong signal to tune optimizer or normalize features.

Mixed Precision and Custom Losses

With mixed precision enabled, some operations may overflow or underflow in lower precision. Cast operands to float32 inside custom losses if needed.

1@tf.function
2def mp_safe_mse(y_true, y_pred):
3    y_true = tf.cast(y_true, tf.float32)
4    y_pred = tf.cast(y_pred, tf.float32)
5    out = tf.reduce_mean(tf.square(y_true - y_pred))
6    tf.debugging.assert_all_finite(out, "MSE invalid")
7    return out

This avoids precision surprises while keeping global mixed-precision benefits.

Common Pitfalls

A common pitfall is assuming NaN means model architecture is wrong when the actual issue is unsafe loss math.

Another issue is debugging only training loops without validating raw data for NaN and Inf values.

Teams often change many settings at once, which makes root cause unclear. Apply one stabilization change at a time and retest.

Ignoring dtype consistency is also frequent, especially when mixing integer labels and floating-point predictions.

Finally, custom losses are rarely unit-tested, so regressions appear late in long training jobs.

Summary

• NaN custom loss usually comes from unsafe math, invalid data, or unstable gradients.
• Add clipping, epsilon, and finite-value checks inside the loss.
• Validate loss on fixed tensors before running full training.
• Control optimizer dynamics with lower learning rate and gradient clipping.
• Check data and dtypes early to prevent hidden pipeline issues.