Why Multi-Task Learning?

The Bilingual Child

Children raised bilingual learn each language slightly more slowly in the early years — and then catch up, often surpassing monolingual peers on a battery of cognitive measures by adulthood. The brain is not a finite vocabulary cup; learning a second language forces it to extract patterns at a higher level — phonology, grammar, abstraction — that benefit BOTH languages.

That same phenomenon, applied to neural networks, is multi-task learning (MTL). One backbone, several related tasks, all sharing parameters. The auxiliary task acts as a regulariser; it forces the shared layers to learn representations that generalise across both objectives. For RUL prediction the main task is regression (predict the cycle count to failure); the auxiliary task is classification (predict the discrete health state — normal, degrading, critical). The auxiliary task sharpens the same shared features the regressor reads from.

Multi-Task Learning, Formally

Suppose we have $K$ related tasks. For each task $k$ there is a dataset $\mathcal{D}_k = \{(\mathbf{x}_i, y_i^{(k)})\}$ and a per-task loss $\mathcal{L}_k$. The model is a function

$f_{\theta_s, \theta_1, \ldots, \theta_K}(\mathbf{x}) \;=\; \bigl(g_1(s(\mathbf{x}; \theta_s); \theta_1),\, \ldots,\, g_K(s(\mathbf{x}; \theta_s); \theta_K)\bigr).$

Here $s(\cdot; \theta_s)$ is the shared encoder (the “backbone”) with parameters $\theta_s$, and $g_k(\cdot; \theta_k)$ is the head for task $k$. The combined training objective is

$\mathcal{L}_{\text{MTL}}(\theta_s, \theta_1, \ldots, \theta_K) \;=\; \sum_{k=1}^{K} \lambda_k \, \mathcal{L}_k.$

The weights $\lambda_k$ determine how much each task influences the shared parameters. Choosing them — statically, dynamically, or adaptively from gradients — is the central problem of Sections 4.3, 14, 17, and 21.

Three Reasons MTL Helps

On C-MAPSS FD002, switching from a single-task RUL regressor to a naive 0.5/0.5 MTL improves RMSE from 8.11 to 7.37 — a 9% reduction with zero architectural change beyond the auxiliary head. The gradient-aware variants (GABA / GRACE) push that further; the whole rest of this book is about extracting the maximum benefit from this single architectural idea.

Interactive: Single-Task vs MTL on FD002

Toggle between single-task and the two MTL variants. The diagram highlights which heads are active; the bar chart shows the actual FD002 numbers from the paper's Table I.

Python: A Tiny Multi-Task Network

Twenty lines of NumPy show the architectural pattern: one shared forward pass, two task-specific projections.

1import numpy as np
2
3# ----- Hand-rolled multi-task forward pass -----
4# Shared trunk: 17 -> 32 features. Two heads: regression (1) + classification (3).
5np.random.seed(0)
6
7# Shared parameters (would be learned in practice)
8W_shared = np.random.randn(17, 32).astype(np.float32) * 0.1
9b_shared = np.zeros(32, dtype=np.float32)
10
11# Task-specific parameters
12W_rul    = np.random.randn(32, 1).astype(np.float32) * 0.1
13b_rul    = np.zeros(1, dtype=np.float32)
14W_health = np.random.randn(32, 3).astype(np.float32) * 0.1
15b_health = np.zeros(3, dtype=np.float32)
16
17
18def forward(x):
19    """x: (B, 17) - one sensor reading per engine.
20    Returns (rul_pred, health_logits) - both tasks at once."""
21    h = np.maximum(0, x @ W_shared + b_shared)        # ReLU activation
22    rul    = (h @ W_rul + b_rul).squeeze(-1)          # (B,)
23    health = h @ W_health + b_health                   # (B, 3)
24    return rul, health
25
26
27# ----- One forward pass on a batch of 4 engines -----
28x = np.random.randn(4, 17).astype(np.float32)
29rul, health = forward(x)
30
31print("input  x.shape           :", x.shape)
32print("shared output  h.shape   :", (4, 32))
33print("RUL head    rul.shape    :", rul.shape)
34print("Health head health.shape :", health.shape)
35print("rul                       :", rul.round(3).tolist())
36
37# input  x.shape           : (4, 17)
38# shared output  h.shape   : (4, 32)
39# RUL head    rul.shape    : (4,)
40# Health head health.shape : (4, 3)

PyTorch: nn.Module With Two Heads

The PyTorch idiom is one nn.Module with three sub-modules: shared, head_rul, head_health. The forward returns a 2-tuple. This same skeleton scales up to the full CNN-BiLSTM-Attention backbone in Chapter 11.

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class DualTaskMLP(nn.Module):
6    """The skeleton of every model in this book: shared trunk + two heads."""
7
8    def __init__(self, in_dim: int = 17, hidden: int = 32, n_classes: int = 3):
9        super().__init__()
10        # Shared trunk
11        self.shared = nn.Sequential(
12            nn.Linear(in_dim, hidden),
13            nn.ReLU(),
14        )
15        # Task-specific heads
16        self.head_rul    = nn.Linear(hidden, 1)
17        self.head_health = nn.Linear(hidden, n_classes)
18
19    def forward(self, x: torch.Tensor):
20        h = self.shared(x)                          # (B, hidden) - SHARED
21        rul    = self.head_rul(h).squeeze(-1)       # (B,)        - TASK 1
22        health = self.head_health(h)                # (B, 3)      - TASK 2
23        return rul, health
24
25
26# ----- Use it -----
27torch.manual_seed(0)
28model = DualTaskMLP()
29x = torch.randn(4, 17)
30rul, health = model(x)
31
32print("rul.shape    :", tuple(rul.shape))     # (4,)
33print("health.shape :", tuple(health.shape))  # (4, 3)
34print("# params     :", sum(p.numel() for p in model.parameters()))
35# # params     : 711  (= 544 shared + 33 RUL head + 99 health head + 35 biases)

MTL Beyond RUL

Every row of the table is solved with the same structure we just coded: shared encoder plus task-specific heads, one combined loss. The architectural commit you make in Chapter 11 transfers to all of them.

When MTL Hurts Instead of Helps

The chapter's theme. Multi-task learning promises better generalisation through parameter sharing. The promise is real but conditional - the auxiliary task must be related, the loss weighting must be sensible, and the gradients must be balanced. Sections 4.2 - 4.4 unpack each of those conditions.

Takeaway

• MTL = shared backbone + task-specific heads. One forward pass produces multiple outputs.
• The shared parameters get gradients from every task. Task-specific parameters only see their own task.
• The mechanisms are regularisation, augmentation, and transfer. Empirically, MTL improves FD002 RMSE from 8.11 (single-task) to 7.37 (naive MTL).
• The combined loss is $\mathcal{L} = \sum_k \lambda_k \mathcal{L}_k$. Choosing $\lambda_k$ well is the rest of the book.
• The PyTorch skeleton is one Module with two heads. Same pattern, larger backbone in Chapter 11.