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.
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.
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.
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:
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.
Use Separate Heads Only When Tasks Are Related
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:
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_weightswhen one task starts dominating the others. - Multitask learning helps most when the tasks are related enough to share useful features.
Related reading
- Multivariate LSTM with missing values
- My LSTM learns, loss decreases, but Numerical Gradients don''t match Analytical Gradients
- NaN from sparse_softmax_cross_entropy_with_logits in Tensorflow
- nan values in loss in keras model
- Multithreading in tensorflow/keras
- NARX implementation using keras
- Multivariate polynomial best fit curve in python?
- Multivariate Regression Neural Network \`Loss\` Function
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.