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
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:
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:
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:
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:
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
Wrap everything into a single loss function:
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
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
