TensorFlow
Optimizers
Machine Learning
Deep Learning
`Loss` Functions

Tensorflow Optimizers - multiple loss values passed to minimize?

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

Neural networks often optimize more than one objective at a time. A model might predict a class label and a numeric value together, or it might combine a task loss with regularization, consistency, or adversarial terms.

In TensorFlow, the important rule is simple: an optimizer step updates variables from gradients of a scalar objective. You can compute that scalar from multiple loss values, or you can run separate optimizer steps for different variable groups, but you do not hand a raw list of unrelated losses to one minimize call and expect TensorFlow to guess your intent.

What minimize Actually Wants

In older TensorFlow code you may see optimizer.minimize(loss). In modern TensorFlow and Keras custom loops, the equivalent pattern is usually:

  1. Compute one or more loss components.
  2. Combine them into a scalar total loss.
  3. Record operations with tf.GradientTape.
  4. Call tape.gradient on the scalar total.
  5. Apply gradients with optimizer.apply_gradients.

That is the right mental model. TensorFlow can differentiate a weighted sum just fine, but it needs a single tensor to backpropagate from.

Combine Multiple Losses Into One Scalar

The most common case is multi-task learning. Suppose one head predicts a class and another predicts a continuous value:

python
1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(8,))
4x = tf.keras.layers.Dense(16, activation="relu")(inputs)
5class_output = tf.keras.layers.Dense(3, activation="softmax", name="class_head")(x)
6reg_output = tf.keras.layers.Dense(1, name="reg_head")(x)
7
8model = tf.keras.Model(inputs=inputs, outputs=[class_output, reg_output])
9optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
10
11ce_loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
12mse_loss_fn = tf.keras.losses.MeanSquaredError()
13
14@tf.function
15def train_step(x_batch, y_class, y_reg):
16    with tf.GradientTape() as tape:
17        pred_class, pred_reg = model(x_batch, training=True)
18        class_loss = ce_loss_fn(y_class, pred_class)
19        reg_loss = mse_loss_fn(y_reg, pred_reg)
20        total_loss = class_loss + 0.2 * reg_loss
21
22    grads = tape.gradient(total_loss, model.trainable_variables)
23    optimizer.apply_gradients(zip(grads, model.trainable_variables))
24    return class_loss, reg_loss, total_loss

The weighting term matters. If one loss is numerically much larger than another, it can dominate training unless you rescale it.

When Separate Optimizers Make More Sense

Sometimes different losses target different parameter sets. A common example is a GAN, where the generator and discriminator have distinct objectives. In that situation, use separate gradient computations and often separate optimizers:

python
1import tensorflow as tf
2
3generator = tf.keras.Sequential([tf.keras.layers.Dense(4)])
4discriminator = tf.keras.Sequential([tf.keras.layers.Dense(1)])
5
6gen_optimizer = tf.keras.optimizers.Adam(1e-4)
7disc_optimizer = tf.keras.optimizers.Adam(1e-4)
8loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)
9
10@tf.function
11def train_step(noise, real_samples):
12    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
13        fake_samples = generator(noise, training=True)
14        real_logits = discriminator(real_samples, training=True)
15        fake_logits = discriminator(fake_samples, training=True)
16
17        gen_loss = loss_fn(tf.ones_like(fake_logits), fake_logits)
18        disc_real_loss = loss_fn(tf.ones_like(real_logits), real_logits)
19        disc_fake_loss = loss_fn(tf.zeros_like(fake_logits), fake_logits)
20        disc_loss = disc_real_loss + disc_fake_loss
21
22    gen_grads = gen_tape.gradient(gen_loss, generator.trainable_variables)
23    disc_grads = disc_tape.gradient(disc_loss, discriminator.trainable_variables)
24
25    gen_optimizer.apply_gradients(zip(gen_grads, generator.trainable_variables))
26    disc_optimizer.apply_gradients(zip(disc_grads, discriminator.trainable_variables))

Here there is no reason to force everything into one optimizer call, because the parameter groups and objectives are intentionally different.

Using model.compile With Multiple Outputs

If you are using the high-level Keras API, you can declare multiple losses directly and let Keras build the scalar objective for you:

python
1model.compile(
2    optimizer="adam",
3    loss={
4        "class_head": tf.keras.losses.SparseCategoricalCrossentropy(),
5        "reg_head": tf.keras.losses.MeanSquaredError(),
6    },
7    loss_weights={
8        "class_head": 1.0,
9        "reg_head": 0.2,
10    },
11)

This is convenient for standard supervised training. Drop to a custom loop when you need conditional logic, manual scheduling, or separate variable updates.

Common Pitfalls

The biggest mistake is passing multiple raw losses without reducing them to one scalar. An optimizer step needs one objective per variable update path.

Another common problem is forgetting that Keras losses may return per-example vectors if reduction is changed. If that happens, explicitly reduce with tf.reduce_mean or another intentional aggregation before computing gradients.

Loss balancing is also easy to get wrong. If one term has a scale near 100 and another stays near 0.01, the smaller one may have almost no influence. Inspect magnitudes during training instead of picking weights blindly.

Finally, a non-persistent GradientTape can only be consumed once. If you need multiple gradient calls from the same recorded operations, either restructure the code or create the tape with persistent=True and clean it up afterward.

Summary

  • TensorFlow optimizers update variables from gradients of a scalar objective.
  • Multiple loss values are usually combined into one weighted total loss.
  • Separate optimizers are appropriate when different losses act on different variable groups.
  • 'model.compile already supports multi-output models with named losses and loss weights.'
  • Watch for reduction issues, bad loss scaling, and one-time GradientTape usage.

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.