Introduction

Multitask learning in Keras works well when tasks share part of a model, but the data setup matters. If the two tasks have different numbers of samples, you cannot simply pass mismatched arrays into one model.fit call and expect Keras to align them automatically.

Same Examples Versus Different Datasets

There are two very different cases that often get mixed together:

• one dataset where each sample has multiple targets
• two separate datasets with different sample counts

If every sample has both task labels, a normal multi-output model is enough. If task A has one dataset and task B has another dataset with a different number of rows, you need a custom training strategy because there is no one-to-one pairing between samples.

That is the real issue behind the "different sample size" question.

Multi-Output Keras Works Only When Batches Align

A standard multitask model in Keras might look like this:

1from tensorflow import keras
2
3inputs = keras.Input(shape=(32,))
4x = keras.layers.Dense(64, activation="relu")(inputs)
5
6class_output = keras.layers.Dense(3, activation="softmax", name="class_output")(x)
7score_output = keras.layers.Dense(1, name="score_output")(x)
8
9model = keras.Model(inputs=inputs, outputs=[class_output, score_output])
10model.compile(
11    optimizer="adam",
12    loss={
13        "class_output": "sparse_categorical_crossentropy",
14        "score_output": "mse",
15    },
16)

This design assumes the same input batch produces both outputs. So x_train, y_class, and y_score must all have the same first dimension. If they do not, model.fit will reject the data because it has no rule for matching sample 17 from one task with sample 17 from another task.

Use Separate Task Datasets With a Custom Training Step

When tasks truly have different sample counts, a common solution is to share part of the network and train each task on its own dataset inside a custom train_step.

1import tensorflow as tf
2from tensorflow import keras
3
4shared = keras.Sequential(
5    [
6        keras.layers.Dense(64, activation="relu"),
7        keras.layers.Dense(32, activation="relu"),
8    ]
9)
10
11class_head = keras.layers.Dense(3, activation="softmax")
12score_head = keras.layers.Dense(1)
13
14optimizer = keras.optimizers.Adam()
15class_loss_fn = keras.losses.SparseCategoricalCrossentropy()
16score_loss_fn = keras.losses.MeanSquaredError()
17
18@tf.function
19def train_step(batch_a, batch_b):
20    x_a, y_a = batch_a
21    x_b, y_b = batch_b
22
23    with tf.GradientTape() as tape:
24        features_a = shared(x_a, training=True)
25        features_b = shared(x_b, training=True)
26
27        pred_a = class_head(features_a, training=True)
28        pred_b = score_head(features_b, training=True)
29
30        loss_a = class_loss_fn(y_a, pred_a)
31        loss_b = score_loss_fn(y_b, pred_b)
32        total_loss = loss_a + 0.5 * loss_b
33
34    variables = shared.trainable_variables + class_head.trainable_variables + score_head.trainable_variables
35    gradients = tape.gradient(total_loss, variables)
36    optimizer.apply_gradients(zip(gradients, variables))
37    return loss_a, loss_b

Now each task can use its own dataset, its own label shape, and even its own batch size.

Feed the Two Datasets Intentionally

With separate datasets, you also need to decide how often each task contributes updates. One simple pattern is to repeat the smaller dataset and zip the two streams:

1task_a_ds = tf.data.Dataset.from_tensor_slices((x_task_a, y_task_a)).batch(32).repeat()
2task_b_ds = tf.data.Dataset.from_tensor_slices((x_task_b, y_task_b)).batch(16).repeat()
3
4for batch_a, batch_b in tf.data.Dataset.zip((task_a_ds, task_b_ds)).take(100):
5    loss_a, loss_b = train_step(batch_a, batch_b)
6    print(float(loss_a), float(loss_b))

This keeps both tasks active even when one dataset is much smaller. You can also adjust task weights, batch sizes, or sampling frequency if one task starts dominating the shared representation.

Different Feature Shapes Are Fine

If the two tasks also have different input shapes, give each task its own input branch before the shared or merged layers. Keras handles multiple input tensors well. The hard part is not the feature shape. The hard part is sample alignment when the datasets do not describe the same examples.

So the rule is:

• different feature shapes are fine with multiple Input layers
• different sample counts need separate batching or a custom training loop

Common Pitfalls

• Passing arrays with different first dimensions into one model.fit call and expecting Keras to align them magically.
• Treating two unrelated datasets as if row i in both datasets represents the same training example.
• Forgetting to rebalance the task losses, which can let one head dominate training.
• Repeating the smaller dataset forever without considering whether it leads to overfitting that task.
• Confusing different sample counts with different feature shapes. They are separate design issues.

Summary

• Standard Keras multitask models require aligned samples across all outputs in a batch.
• Different dataset sizes usually mean you need separate task datasets and a custom training loop or train_step.
• Shared layers can still be trained jointly even when the tasks do not share identical examples.
• Task weighting and sampling frequency matter when one dataset is much larger than the other.
• Solve sample-count mismatch and feature-shape mismatch as two separate problems.