Deep Learning
TensorFlow
Multitask Learning
Artificial Intelligence
Machine Learning

Multitask deep learning with Tensorflow

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

Multitask learning trains one model to solve several related tasks at the same time. In TensorFlow and Keras, the usual pattern is to build a shared feature extractor, then add separate output heads for each task. This can improve data efficiency and generalization, but only if the tasks are related closely enough and their losses are balanced carefully.

Why Multitask Learning Helps

The main idea is that some features are useful across tasks. For example, one image model might predict both object category and bounding-box coordinates, or one text model might predict both sentiment and topic.

A multitask model usually has:

  • shared layers that learn common structure
  • one output head per task
  • one loss function per output
  • optional loss weights to control task influence

The shared trunk reduces duplication. Instead of training two separate models from scratch, the model learns one representation that feeds both tasks.

Build a Shared Trunk with Multiple Heads

The Keras Functional API is the cleanest way to express this architecture.

python
1import tensorflow as tf
2from tensorflow import keras
3from tensorflow.keras import layers
4
5inputs = keras.Input(shape=(20,), name="features")
6x = layers.Dense(64, activation="relu")(inputs)
7x = layers.Dense(32, activation="relu")(x)
8
9class_output = layers.Dense(3, activation="softmax", name="class_output")(x)
10score_output = layers.Dense(1, name="score_output")(x)
11
12model = keras.Model(
13    inputs=inputs,
14    outputs={
15        "class_output": class_output,
16        "score_output": score_output,
17    },
18)
19
20model.compile(
21    optimizer="adam",
22    loss={
23        "class_output": "sparse_categorical_crossentropy",
24        "score_output": "mse",
25    },
26    metrics={
27        "class_output": ["accuracy"],
28        "score_output": ["mae"],
29    },
30)
31
32model.summary()

This model does classification and regression at once. The shared dense layers learn features that both tasks can use.

Train with Label Dictionaries

When a model has named outputs, training data is usually passed as a dictionary keyed by output name.

python
1import numpy as np
2
3x_train = np.random.random((128, 20)).astype("float32")
4y_class = np.random.randint(0, 3, size=(128,))
5y_score = np.random.random((128, 1)).astype("float32")
6
7model.fit(
8    x_train,
9    {
10        "class_output": y_class,
11        "score_output": y_score,
12    },
13    epochs=3,
14    batch_size=16,
15    verbose=0,
16)

This structure becomes especially useful when using tf.data.Dataset, because each batch can yield one input tensor plus a dictionary of targets.

Balance Task Influence with Loss Weights

One of the hardest parts of multitask learning is preventing one task from dominating the shared representation.

You can adjust task influence with loss_weights:

python
1model.compile(
2    optimizer="adam",
3    loss={
4        "class_output": "sparse_categorical_crossentropy",
5        "score_output": "mse",
6    },
7    loss_weights={
8        "class_output": 1.0,
9        "score_output": 0.3,
10    },
11)

This matters because different losses can have very different numeric scales. Without weighting, the regression task might overwhelm the classification task, or the reverse.

A good first step is to monitor each task's loss separately and adjust weights only when one task clearly dominates training.

Multitask learning is not automatically better than separate models. It works best when the tasks genuinely share useful structure.

Good examples:

  • image classification plus attribute prediction
  • sentiment plus intent classification
  • detection plus box regression

Bad combinations are tasks that compete for very different representations. In those cases, the shared trunk can hurt both tasks instead of helping them.

A tf.data Input Example

The same idea works naturally with datasets:

python
1import tensorflow as tf
2import numpy as np
3
4x = np.random.random((64, 20)).astype("float32")
5y1 = np.random.randint(0, 3, size=(64,))
6y2 = np.random.random((64, 1)).astype("float32")
7
8dataset = tf.data.Dataset.from_tensor_slices(
9    (x, {"class_output": y1, "score_output": y2})
10).batch(8)
11
12model.fit(dataset, epochs=2, verbose=0)

This is usually the cleanest path once your training data becomes too large for simple in-memory arrays.

Common Pitfalls

One common mistake is combining unrelated tasks and expecting the shared layers to help both. Multitask learning is useful only when the tasks have genuinely overlapping structure.

Another mistake is ignoring loss scale. If one task has much larger gradients or much larger numeric loss values, it can dominate the shared representation.

Developers also sometimes forget to name outputs clearly, which makes training dictionaries and metric logs harder to read.

Finally, multitask models are harder to debug than single-task models. Always monitor per-task metrics, not just the total loss, or you can miss the fact that one task improved while another quietly regressed.

Summary

  • Multitask learning in TensorFlow usually means one shared trunk plus multiple output heads.
  • The Keras Functional API is the clearest way to build these models.
  • Train with one loss per output and pass targets as a dictionary.
  • Use loss_weights when one task starts dominating the others.
  • Multitask learning helps most when the tasks are related enough to share useful features.

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.