TensorFlow
machine learning
multiple loss functions
training operations
neural networks

Tensorflow Multiple loss functions vs Multiple training ops

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

Multiple loss functions and multiple training operations are related ideas, but they are not the same thing. A model can combine several loss terms into one objective and still use a single optimizer step, while multiple training operations usually mean different parameter groups are being updated separately.

That distinction matters in multi-task learning, GAN training, actor-critic methods, and any model with parts that should not all move in exactly the same way on every step.

Multiple Loss Functions

Using multiple losses means the model is evaluated against more than one objective. The usual pattern is to compute each loss, weight them, and sum them into one scalar that drives backpropagation.

In Keras, a multi-output model can do this directly:

python
1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(16,))
4shared = tf.keras.layers.Dense(32, activation="relu")(inputs)
5price = tf.keras.layers.Dense(1, name="price")(shared)
6category = tf.keras.layers.Dense(3, activation="softmax", name="category")(shared)
7
8model = tf.keras.Model(inputs=inputs, outputs=[price, category])
9
10model.compile(
11    optimizer="adam",
12    loss={
13        "price": "mse",
14        "category": "sparse_categorical_crossentropy",
15    },
16    loss_weights={
17        "price": 0.5,
18        "category": 1.0,
19    },
20)

This setup still uses one optimizer update per batch. The optimizer follows the gradient of the combined objective.

What A Single Training Operation Means

A single training operation is the actual parameter update step. In TensorFlow terms, that usually means one call to optimizer.apply_gradients(...) over a chosen variable set.

You can have several loss components and still keep one training op:

python
1with tf.GradientTape() as tape:
2    predictions = model(x_batch, training=True)
3    loss_a = tf.reduce_mean(tf.square(predictions - y_batch))
4    loss_b = 0.01 * tf.add_n(model.losses)
5    total_loss = loss_a + loss_b
6
7grads = tape.gradient(total_loss, model.trainable_variables)
8optimizer.apply_gradients(zip(grads, model.trainable_variables))

That is still one training operation because there is one gradient computation against one total loss and one optimizer application.

When You Need Multiple Training Ops

Multiple training operations are useful when different parts of the system should be updated differently. Common cases include:

  • GANs, where generator and discriminator train on different steps
  • models with frozen or partially trainable sub-networks
  • systems where different parameter groups need different optimizers or learning rates

Here is a simplified two-optimizer example:

python
1with tf.GradientTape(persistent=True) as tape:
2    g_output = generator(noise, training=True)
3    d_real = discriminator(real_batch, training=True)
4    d_fake = discriminator(g_output, training=True)
5
6    d_loss = discriminator_loss(d_real, d_fake)
7    g_loss = generator_loss(d_fake)
8
9d_grads = tape.gradient(d_loss, discriminator.trainable_variables)
10g_grads = tape.gradient(g_loss, generator.trainable_variables)
11
12d_optimizer.apply_gradients(zip(d_grads, discriminator.trainable_variables))
13g_optimizer.apply_gradients(zip(g_grads, generator.trainable_variables))

That is multiple training ops because there are separate updates for different variable sets and objectives.

How To Choose

If your model has several related objectives but one coherent parameter update, use multiple losses combined into one total loss. That is the standard multi-task learning pattern.

If different subnetworks should be optimized on different schedules or with different optimizers, use multiple training ops. That gives you finer control, but it also makes training logic more complex and easier to destabilize.

Beware Of Conflicting Gradients

Even with one optimizer step, multiple losses can compete. One task may dominate the gradient and prevent another from improving. That is why loss weighting matters.

You may need to tune:

  • fixed loss weights
  • dynamic reweighting strategies
  • separate learning rates for different components

The presence of multiple losses does not automatically mean you need multiple training ops. Often the right fix is better scaling, not more optimizers.

Common Pitfalls

  • Assuming every extra loss term requires its own optimizer step.
  • Combining unrelated variable groups into one update when they should be trained separately.
  • Ignoring loss scale differences and letting one task dominate the others.
  • Using multiple training ops without a clear scheduling rule.
  • Confusing model regularization losses with separate task objectives.

Summary

  • Multiple loss functions describe how many objectives the model is optimizing.
  • A training operation is the actual parameter update step.
  • Several losses can be combined into one total loss and trained with one optimizer step.
  • Multiple training ops are useful when different variable groups need separate updates.
  • Choose the simpler single-op approach unless the model architecture truly needs separate optimization flows.

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.