Keras
loss functions
multi-output model
deep learning
machine learning

Using different loss functions for different outputs simultaneously Keras?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Multi-output models are common when one network must solve more than one prediction task at once. A typical example is predicting a class label and a numeric score from the same input features. Keras supports this directly, and each output head can use its own loss and metric as long as model outputs and training targets are mapped consistently.

Build a Model with Named Output Heads

The most important design decision is naming each output layer clearly. Named heads make compile configuration and debugging much easier than relying on positional lists.

python
1import numpy as np
2import tensorflow as tf
3from tensorflow.keras import layers
4
5inputs = layers.Input(shape=(20,), name='features')
6shared = layers.Dense(64, activation='relu')(inputs)
7shared = layers.Dense(32, activation='relu')(shared)
8
9class_head = layers.Dense(4, activation='softmax', name='class_head')(shared)
10price_head = layers.Dense(1, activation='linear', name='price_head')(shared)
11
12model = tf.keras.Model(inputs=inputs, outputs=[class_head, price_head])
13model.summary()

This pattern creates one shared representation, then task-specific heads. Shared layers learn reusable signal, while each head focuses on its own objective.

Compile with Different Loss Functions

Use a loss dictionary keyed by output layer names. This lets classification and regression coexist cleanly.

python
1model.compile(
2    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
3    loss={
4        'class_head': tf.keras.losses.SparseCategoricalCrossentropy(),
5        'price_head': tf.keras.losses.MeanSquaredError(),
6    },
7    metrics={
8        'class_head': [tf.keras.metrics.SparseCategoricalAccuracy(name='acc')],
9        'price_head': [tf.keras.metrics.MeanAbsoluteError(name='mae')],
10    },
11    loss_weights={
12        'class_head': 1.0,
13        'price_head': 0.3,
14    },
15)

Loss weighting matters because different objectives can have different scales. If the regression loss is numerically much larger, it can dominate gradient updates and hurt classification quality. Start simple with equal or near-equal weights, then tune based on validation metrics per head.

Feed Targets that Match Output Names

When calling fit, pass target data using the same output keys used in compile. If names do not match, Keras raises mapping errors or silently trains the wrong objective by position.

python
1# Toy data
2rng = np.random.default_rng(7)
3X = rng.normal(size=(1024, 20)).astype('float32')
4y_class = rng.integers(0, 4, size=(1024,)).astype('int32')
5y_price = rng.normal(loc=100.0, scale=15.0, size=(1024, 1)).astype('float32')
6
7history = model.fit(
8    X,
9    {
10        'class_head': y_class,
11        'price_head': y_price,
12    },
13    epochs=8,
14    batch_size=32,
15    validation_split=0.2,
16    verbose=2,
17)

A practical workflow is to plot each head metric over epochs. A healthy run shows both tasks improving without one head collapsing. If one objective stalls while the other improves quickly, adjust learning rate, head depth, or loss weights.

Add a Custom Loss for One Head

You can mix built-in and custom losses. This is useful when domain constraints require asymmetric penalties or bounded error behavior.

python
1def clipped_mse(y_true, y_pred):
2    err = tf.square(y_true - y_pred)
3    err = tf.clip_by_value(err, 0.0, 25.0)
4    return tf.reduce_mean(err)
5
6model.compile(
7    optimizer='adam',
8    loss={
9        'class_head': 'sparse_categorical_crossentropy',
10        'price_head': clipped_mse,
11    },
12    metrics={
13        'class_head': ['accuracy'],
14        'price_head': ['mae']},
15)

Keep custom losses deterministic and numerically stable. If a custom loss introduces discontinuities or very large gradients, total training can become unstable for all heads.

Inference and Contract Management

At prediction time, output order and names become an interface contract for downstream code. Treat this as API design. Save model version, output names, expected shapes, and post-processing rules together.

python
pred_class, pred_price = model.predict(X[:5])
print(pred_class.shape)  # (5, 4)
print(pred_price.shape)  # (5, 1)

In production systems, many integration bugs come from head renaming or output order changes. A simple compatibility test that loads the model and validates output schema before deployment prevents most of these failures.

Common Pitfalls

  • Using unnamed output layers and then confusing loss assignment by index.
  • Passing training targets in a structure that does not match output names.
  • Ignoring loss magnitude imbalance and monitoring only total combined loss.
  • Applying a custom loss that is not scale-stable, causing gradient spikes.
  • Changing output head names between model versions without updating consumers.

Summary

  • Keras natively supports different losses for each output head in one model.
  • Name output layers and use dictionary mappings for losses, metrics, and targets.
  • Tune loss_weights so one task does not dominate shared representation learning.
  • Use custom losses only when needed and validate numerical stability.
  • Treat output schema as a versioned contract during inference and deployment.

Course illustration
Course illustration

All Rights Reserved.