Workaround for removal of add_loss

TensorFlow

machine learning

deep learning

model training

technical guide

Workaround for removal of add_loss

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

Deep Learning TensorFlow & Keras Machine Learning

Introduction

The add_loss() method in Keras allows layers and models to register auxiliary losses (like regularization penalties, KL divergence in VAEs, or custom constraints) that are automatically included in the total loss during training. While add_loss() still exists in TensorFlow/Keras, its behavior has changed across versions, and some workflows (particularly with tf.function and custom training loops) require alternative approaches. The main workarounds are computing losses directly in the training step, using the losses property, or adding losses through activity_regularizer.

How add_loss Works

python

1import tensorflow as tf
2
3class KLDivergenceLayer(tf.keras.layers.Layer):
4    def call(self, inputs):
5        mu, log_var = inputs
6        kl_loss = -0.5 * tf.reduce_mean(1 + log_var - tf.square(mu) - tf.exp(log_var))
7        self.add_loss(kl_loss)
8        return mu
9
10# The loss is automatically added to model.losses
11model = build_vae_model()
12print(model.losses)  # [<tf.Tensor: KL divergence value>]

During model.fit(), Keras automatically adds sum(model.losses) to the main loss.

Workaround 1: Compute Loss in Custom Training Step

The most reliable approach — compute all losses explicitly:

python

1class VAE(tf.keras.Model):
2    def __init__(self):
3        super().__init__()
4        self.encoder = self.build_encoder()
5        self.decoder = self.build_decoder()
6
7    def train_step(self, data):
8        with tf.GradientTape() as tape:
9            # Forward pass
10            mu, log_var, z = self.encoder(data)
11            reconstruction = self.decoder(z)
12
13            # Compute all losses explicitly
14            recon_loss = tf.reduce_mean(
15                tf.keras.losses.binary_crossentropy(data, reconstruction)
16            )
17            kl_loss = -0.5 * tf.reduce_mean(
18                1 + log_var - tf.square(mu) - tf.exp(log_var)
19            )
20            total_loss = recon_loss + kl_loss
21
22        grads = tape.gradient(total_loss, self.trainable_variables)
23        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
24
25        return {
26            'loss': total_loss,
27            'recon_loss': recon_loss,
28            'kl_loss': kl_loss
29        }

Workaround 2: Use model.losses in Custom Loop

Collect losses registered via add_loss() and add them manually:

python

1model = build_model()  # Model with layers that call add_loss()
2optimizer = tf.keras.optimizers.Adam(1e-3)
3
4@tf.function
5def train_step(inputs, targets):
6    with tf.GradientTape() as tape:
7        predictions = model(inputs, training=True)
8        main_loss = tf.keras.losses.mse(targets, predictions)
9
10        # Add all auxiliary losses registered by layers
11        aux_loss = tf.add_n(model.losses) if model.losses else 0.0
12        total_loss = main_loss + aux_loss
13
14    grads = tape.gradient(total_loss, model.trainable_variables)
15    optimizer.apply_gradients(zip(grads, model.trainable_variables))
16    return total_loss

Workaround 3: activity_regularizer

For regularization losses, use the activity_regularizer parameter instead of add_loss:

python

1# Instead of add_loss for L1 regularization on activations:
2model = tf.keras.Sequential([
3    tf.keras.layers.Dense(128, activation='relu',
4                          activity_regularizer=tf.keras.regularizers.l1(1e-5)),
5    tf.keras.layers.Dense(64, activation='relu',
6                          activity_regularizer=tf.keras.regularizers.l2(1e-4)),
7    tf.keras.layers.Dense(10)
8])
9
10# Keras automatically includes activity_regularizer losses in model.fit()
11model.compile(optimizer='adam', loss='mse')
12model.fit(x_train, y_train, epochs=10)

Workaround 4: Kernel and Bias Regularizers

For weight regularization (L1/L2), use built-in regularizer parameters:

python

1# Weight regularization without add_loss
2model = tf.keras.Sequential([
3    tf.keras.layers.Dense(128,
4                          kernel_regularizer=tf.keras.regularizers.l2(1e-4),
5                          bias_regularizer=tf.keras.regularizers.l1(1e-5)),
6    tf.keras.layers.Dense(10)
7])
8
9# All regularization losses are in model.losses
10print(model.losses)  # List of regularization tensors

Workaround 5: Custom Loss Function with Extra Terms

Wrap everything into a single loss function:

python

1class CustomLoss(tf.keras.losses.Loss):
2    def __init__(self, model, lambda_kl=1.0):
3        super().__init__()
4        self.model = model
5        self.lambda_kl = lambda_kl
6
7    def call(self, y_true, y_pred):
8        # Main loss
9        recon_loss = tf.reduce_mean(tf.square(y_true - y_pred))
10
11        # Add regularization losses from the model
12        reg_loss = tf.add_n(self.model.losses) if self.model.losses else 0.0
13
14        return recon_loss + self.lambda_kl * reg_loss
15
16model = build_model()
17model.compile(optimizer='adam', loss=CustomLoss(model))

VAE Complete Example Without add_loss

python

1class Encoder(tf.keras.layers.Layer):
2    def __init__(self, latent_dim):
3        super().__init__()
4        self.dense1 = tf.keras.layers.Dense(256, activation='relu')
5        self.mu_layer = tf.keras.layers.Dense(latent_dim)
6        self.log_var_layer = tf.keras.layers.Dense(latent_dim)
7
8    def call(self, x):
9        h = self.dense1(x)
10        mu = self.mu_layer(h)
11        log_var = self.log_var_layer(h)
12        # Reparameterization trick
13        eps = tf.random.normal(tf.shape(mu))
14        z = mu + tf.exp(0.5 * log_var) * eps
15        return mu, log_var, z
16
17class VAE(tf.keras.Model):
18    def __init__(self, latent_dim):
19        super().__init__()
20        self.encoder = Encoder(latent_dim)
21        self.decoder = tf.keras.Sequential([
22            tf.keras.layers.Dense(256, activation='relu'),
23            tf.keras.layers.Dense(784, activation='sigmoid')
24        ])
25        self.loss_tracker = tf.keras.metrics.Mean(name='loss')
26
27    def train_step(self, data):
28        with tf.GradientTape() as tape:
29            mu, log_var, z = self.encoder(data)
30            reconstruction = self.decoder(z)
31            recon_loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(data, reconstruction))
32            kl_loss = -0.5 * tf.reduce_mean(1 + log_var - tf.square(mu) - tf.exp(log_var))
33            total_loss = recon_loss + kl_loss
34
35        grads = tape.gradient(total_loss, self.trainable_variables)
36        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
37        self.loss_tracker.update_state(total_loss)
38        return {'loss': self.loss_tracker.result()}
39
40    @property
41    def metrics(self):
42        return [self.loss_tracker]
43
44vae = VAE(latent_dim=32)
45vae.compile(optimizer='adam')
46vae.fit(x_train, epochs=10, batch_size=128)

Common Pitfalls

model.losses is empty outside forward pass: add_loss() registers losses during the forward pass. If you access model.losses before calling model(inputs), the list is empty. Always read model.losses after the forward pass.
Double-counting losses with model.fit(): If you use add_loss() in a layer AND manually add the same loss in a custom train_step, the loss is counted twice. Choose one approach, not both.
tf.add_n(model.losses) fails on empty list: If no layers call add_loss(), model.losses is empty and tf.add_n([]) raises an error. Guard with if model.losses else 0.0.
add_loss in @tf.function tracing: Losses added during @tf.function tracing are captured once. If the loss depends on dynamic values, ensure the function is traced correctly or use eager mode for debugging.
Keras 3 behavior changes: In Keras 3 (TF 2.16+), add_loss() behavior may differ from Keras 2. Check the Keras 3 migration guide if upgrading, and prefer explicit loss computation in train_step for maximum compatibility.

Summary

add_loss() still works but alternative patterns are more explicit and portable across Keras versions
Compute all losses directly in a custom train_step for full control and clarity
Use model.losses to collect auxiliary losses registered by layers after the forward pass
Use kernel_regularizer and activity_regularizer for standard regularization instead of manual add_loss
Guard tf.add_n(model.losses) with an empty-list check to avoid runtime errors

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

Workaround for removal of add_loss

ML System Design practice on Codemia

Introduction

How add_loss Works

Workaround 1: Compute Loss in Custom Training Step

Workaround 2: Use model.losses in Custom Loop

Workaround 3: activity_regularizer

Workaround 4: Kernel and Bias Regularizers

Workaround 5: Custom Loss Function with Extra Terms

VAE Complete Example Without add_loss

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

Track what you have practised

ML System Design practice on Codemia