Introduction

When loss in TensorFlow.js shoots to Infinity or turns into NaN, the model is almost always numerically unstable. The framework is usually doing exactly what your setup asked it to do; the real issue is a bad combination of input scale, optimizer step size, target encoding, or loss/output mismatch.

The Usual Causes

The most common reason is a learning rate that is too large. Each update overshoots, the weights explode, and the loss quickly leaves the finite numeric range.

The second common cause is unnormalized data. If one feature is tiny and another is enormous, gradient magnitudes become hard to control.

The third is a mismatch between the model head and the loss. For example:

• linear output with a classification loss
• sigmoid output with targets encoded for another objective
• logits passed into a loss that expects probabilities

Any of those can destabilize training quickly.

A Stable Baseline Example

1import * as tf from '@tensorflow/tfjs';
2
3async function trainStable() {
4  const xs = tf.tensor2d([[0], [1], [2], [3], [4]], [5, 1]);
5  const ys = tf.tensor2d([[1], [3], [5], [7], [9]], [5, 1]);
6
7  const model = tf.sequential();
8  model.add(tf.layers.dense({ units: 8, activation: 'relu', inputShape: [1] }));
9  model.add(tf.layers.dense({ units: 1 }));
10
11  model.compile({
12    optimizer: tf.train.adam(0.01),
13    loss: 'meanSquaredError',
14  });
15
16  await model.fit(xs, ys, {
17    epochs: 100,
18    callbacks: {
19      onEpochEnd: async (epoch, logs) => {
20        console.log(epoch, logs.loss);
21      }
22    }
23  });
24}
25
26trainStable();

This setup stays stable because the data is small and well-scaled, the task is regression, and the learning rate is conservative.

An Unstable Version

1import * as tf from '@tensorflow/tfjs';
2
3async function trainUnstable() {
4  const xs = tf.tensor2d([[1000], [2000], [3000], [4000], [5000]], [5, 1]);
5  const ys = tf.tensor2d([[1], [2], [3], [4], [5]], [5, 1]);
6
7  const model = tf.sequential();
8  model.add(tf.layers.dense({ units: 64, activation: 'relu', inputShape: [1] }));
9  model.add(tf.layers.dense({ units: 64, activation: 'relu' }));
10  model.add(tf.layers.dense({ units: 1 }));
11
12  model.compile({
13    optimizer: tf.train.sgd(1.0),
14    loss: 'meanSquaredError',
15  });
16
17  await model.fit(xs, ys, { epochs: 20 });
18}
19
20trainUnstable();

This is much more likely to diverge because the features are large and the learning rate is aggressive.

Normalize First, Then Tune

A very effective first fix is to normalize inputs and, for regression, sometimes targets too.

const xMean = xs.mean();
const xStd = tf.moments(xs).variance.sqrt();
const xsNorm = xs.sub(xMean).div(xStd);

Normalization is not just a preprocessing preference. It directly changes the scale of gradients and often determines whether training is stable.

Verify the Data Is Finite

Do not assume the tensors are valid just because the code created them. Browser-side pipelines often read values from forms, JSON, CSV, or canvas operations, and bad parsing can silently introduce NaN.

1function assertFinite(name, tensor) {
2  const bad = tf.logicalNot(tf.isFinite(tensor)).any().dataSync()[0];
3  if (bad) {
4    throw new Error(`${name} contains NaN or Infinity`);
5  }
6}

Use checks like this on inputs, targets, and sometimes predictions during debugging.

Match the Head to the Loss

A stable model also needs the correct final-layer and loss pairing.

Typical safe combinations are:

• regression: linear output plus mean squared error or mean absolute error
• binary classification: sigmoid output plus binary cross-entropy
• multiclass classification: softmax output plus categorical cross-entropy

If the task and head disagree, training can become numerically unstable even when the dataset is fine.

Common Pitfalls

The biggest mistake is changing the architecture repeatedly before checking the data scale and learning rate.

Another mistake is feeding unnormalized or malformed inputs from browser-side parsing code.

A third issue is using an output layer that does not match the loss and target encoding.

Finally, do not debug with a large model first. A tiny stable baseline is much easier to reason about than a deep network that fails in five different ways at once.

Summary

• 'Infinity loss in TensorFlow.js is usually a numerical-stability problem.'
• Lower the learning rate and normalize data before trying more complex fixes.
• Verify that inputs and targets are finite.
• Make sure the model output and loss function match the task.
• Use a tiny baseline model to isolate the failure mode.
• Treat exploding loss as a signal to simplify and inspect, not to add more layers.