TensorFlow
loss computation
machine learning
deep learning
neural networks

TensorFlow Performing this loss computation

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

When people ask how TensorFlow is "performing this loss computation," they are usually trying to answer one of three questions: what value is being reduced across the batch, whether logits or probabilities are expected, and which extra terms are silently added before optimization. Once those pieces are clear, most loss bugs become straightforward to diagnose.

Loss in TensorFlow Is Usually a Composite

In practice, the training loss is rarely just one formula copied from a paper. TensorFlow often combines:

  • A per-example loss such as cross-entropy or mean squared error
  • A reduction step across the batch
  • Optional regularization losses collected from layers

For example, a binary classification loss built from logits looks like this:

python
1import tensorflow as tf
2
3logits = tf.constant([[2.0], [-1.0], [0.5]], dtype=tf.float32)
4labels = tf.constant([[1.0], [0.0], [1.0]], dtype=tf.float32)
5
6per_example = tf.nn.sigmoid_cross_entropy_with_logits(
7    labels=labels,
8    logits=logits
9)
10
11loss = tf.reduce_mean(per_example)
12print(per_example.numpy())
13print(loss.numpy())

The first tensor contains one loss value per example. tf.reduce_mean turns that vector into the scalar that the optimizer minimizes.

Logits Versus Probabilities

One of the most common mistakes is passing already-sigmoid or already-softmaxed outputs into a loss function that expects raw logits. TensorFlow provides numerically stable "with logits" helpers precisely so you do not need to apply the activation first.

Correct pattern:

python
1model = tf.keras.Sequential([
2    tf.keras.layers.Dense(1)
3])
4
5loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)

Incorrect pattern:

python
1model = tf.keras.Sequential([
2    tf.keras.layers.Dense(1, activation="sigmoid")
3])
4
5loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)

That mismatch often produces strange gradients and unstable training because the loss function interprets probabilities as logits.

Manual Computation Versus tf.keras

Keras layers and losses wrap the same idea with cleaner defaults. Here is a minimal example:

python
1import tensorflow as tf
2
3features = tf.random.normal([8, 4])
4labels = tf.constant([[1.0], [0.0], [1.0], [1.0], [0.0], [0.0], [1.0], [0.0]])
5
6model = tf.keras.Sequential([
7    tf.keras.layers.Dense(
8        16,
9        activation="relu",
10        kernel_regularizer=tf.keras.regularizers.l2(1e-4)
11    ),
12    tf.keras.layers.Dense(1)
13])
14
15loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)
16
17with tf.GradientTape() as tape:
18    logits = model(features, training=True)
19    data_loss = loss_fn(labels, logits)
20    total_loss = data_loss + tf.add_n(model.losses)
21
22grads = tape.gradient(total_loss, model.trainable_variables)

Two important details appear here:

  • 'data_loss comes from the declared loss function'
  • 'model.losses contains regularization terms added by layers'

If you ignore model.losses, you are not optimizing the loss that the model definition implies.

Understanding the Reduction Step

TensorFlow loss APIs often default to averaging across the batch. That is convenient, but it matters when you compare values across different batch sizes or when you implement custom sample weights.

Here is a custom weighted loss:

python
per_example = tf.keras.losses.binary_crossentropy(labels, logits, from_logits=True)
weights = tf.constant([1.0, 1.0, 2.0, 0.5, 1.0, 1.0, 2.0, 0.5])
weighted_loss = tf.reduce_sum(per_example * weights) / tf.reduce_sum(weights)

This makes the reduction explicit instead of relying on a hidden default.

How to Debug a Suspicious Loss

If the number looks wrong, inspect the pipeline in order:

  1. Print the model output before the loss
  2. Confirm whether the loss expects logits or probabilities
  3. Inspect the unreduced per-example values
  4. Check whether regularization terms are included
  5. Verify any masking or sample weighting

That sequence usually pinpoints the issue faster than staring at the final scalar.

Common Pitfalls

  • Passing probabilities into a loss configured with from_logits=True.
  • Comparing unreduced per-example loss values against a reduced batch mean.
  • Forgetting that Keras regularizers add terms through model.losses.
  • Mixing labels with the wrong shape, such as (batch,) versus (batch, 1), and masking the real bug.

Summary

  • TensorFlow loss computation is usually a combination of per-example loss, reduction, and optional regularization.
  • Always verify whether your loss function expects logits or probabilities.
  • Inspect unreduced loss values when debugging.
  • In Keras, include model.losses if your layers use regularization.
  • Most "mysterious" loss values come from a mismatch in one of those steps, not from TensorFlow doing something hidden.

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

All Rights Reserved.