TensorFlow
deep learning
loss function
neural networks
NaN error

Adding multiple layers to TensorFlow causes loss function to become Nan

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

If a TensorFlow model starts producing NaN loss after you add more layers, the extra depth is usually exposing a numerical stability problem that was already latent. More layers can amplify bad scaling, poor initialization, exploding activations, or an incorrect loss configuration until training breaks visibly.

The Most Common Reasons Loss Becomes NaN

In practice, NaN loss often comes from one of these:

  • learning rate is too high
  • input data already contains NaN or very large values
  • final activation does not match the loss configuration
  • gradients explode in a deeper network
  • custom math takes log(0), divides by zero, or overflows

Adding layers does not magically create NaN. It just makes unstable training easier to trigger.

Check the Output Layer and Loss Pairing

One common bug is mixing logits and probabilities incorrectly. For example, if the last layer already applies softmax, then the loss should usually use from_logits=False. If the last layer returns raw scores, use from_logits=True.

A safe pattern is:

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(64, activation="relu"),
5    tf.keras.layers.Dense(64, activation="relu"),
6    tf.keras.layers.Dense(3),
7])
8
9loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
10
11model.compile(
12    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
13    loss=loss_fn,
14    metrics=["accuracy"],
15)

If you instead add a softmax layer on top, then from_logits should normally be False.

Lower the Learning Rate and Clip Gradients

Deeper networks are more sensitive to large updates. A quick stabilization step is to reduce the learning rate and clip gradients:

python
1optimizer = tf.keras.optimizers.Adam(
2    learning_rate=1e-4,
3    clipnorm=1.0,
4)
5
6model.compile(optimizer=optimizer, loss=loss_fn)

Gradient clipping does not fix every issue, but it often prevents one bad step from blowing the model into NaN.

Validate the Data Before Training

Always inspect the inputs and labels:

python
1import numpy as np
2
3print(np.isnan(x_train).any(), np.isinf(x_train).any())
4print(np.isnan(y_train).any(), np.isinf(y_train).any())
5print(x_train.min(), x_train.max())

If features are extremely large or unnormalized, deeper layers can amplify them into overflow. Standardizing inputs or scaling them into a reasonable range often helps immediately.

Use TensorFlow's Numerics Checks

TensorFlow provides debugging tools for this exact situation. Enabling numerics checks can stop execution near the first bad tensor instead of letting the whole training loop silently drift into NaN:

python
import tensorflow as tf

tf.debugging.enable_check_numerics()

You can also inspect specific tensors:

python
tensor = tf.constant([1.0, float("nan")])
tf.debugging.check_numerics(tensor, "bad tensor")

These tools help you find whether the first NaN appears in the inputs, activations, gradients, or custom loss.

Depth Changes Initialization Pressure

More layers mean the initialization and activation choices matter more. relu often works well with default Keras initializers, but stacking many dense layers or using saturating activations can still produce unstable gradients.

If the model deepened significantly, consider:

  • normalizing inputs
  • adding batch normalization where appropriate
  • reducing depth until the smallest failing change is isolated
  • training on a tiny batch first to see when the first bad value appears

Debugging becomes much easier when you shrink the problem instead of changing ten things at once.

Common Pitfalls

The biggest mistake is assuming the number of layers is the direct bug. Usually the real issue is unstable numerics, and the extra layers only make it visible.

Another common issue is misconfiguring the final layer and loss, especially around logits versus probabilities.

A third problem is ignoring bad input data. If the training set already contains NaN, no optimizer setting will rescue the run.

Summary

  • 'NaN loss in a deeper TensorFlow model usually indicates numerical instability, not "too many layers" by itself.'
  • Check the output layer and loss configuration first, especially logits versus softmax.
  • Lower the learning rate and consider gradient clipping.
  • Validate the input data for NaN, Inf, and bad scaling.
  • Use tf.debugging.enable_check_numerics() or tf.debugging.check_numerics() to catch the first failing tensor.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design