Multi-Task Learning

Multi-task learning (MTL) represents a fundamental paradigm shift in how we approach machine learning problems. Rather than training separate models for each task in isolation, MTL trains a single model on multiple tasks simultaneously, leveraging the inductive bias that tasks share underlying structure. At the heart of this paradigm lies the concept of shared representations—the idea that different tasks can benefit from learning and utilizing common feature representations.

This page provides a rigorous exploration of shared representations: why they work, how to formalize them mathematically, what architectural choices enable effective sharing, and how this principle manifests across different domains from computer vision to natural language processing. Understanding shared representations is essential because they form the theoretical and practical foundation upon which all multi-task learning architectures are built.

Consider how humans learn. When a child learns to recognize chairs, they simultaneously develop visual processing capabilities that transfer to recognizing tables, sofas, and countless other objects. The child doesn't develop entirely separate visual systems for each object category—instead, they build a hierarchical representation that captures edges, textures, shapes, and spatial relationships that are useful across many recognition tasks.

This observation motivates the core hypothesis of multi-task learning: related tasks share underlying structure, and by learning this structure jointly, we can achieve better generalization than learning each task independently.

The statistical argument:

When we train a model on a single task with limited data, we risk overfitting to task-specific noise. The model may learn spurious correlations that happen to work on the training set but fail to generalize. By training on multiple related tasks simultaneously, we impose an implicit regularization: the model must learn representations that are useful across all tasks, effectively filtering out task-specific noise and preserving genuinely useful features.

Mathematically, if we denote the true underlying representation as $h^*(x)$ and our learned representation as $\hat{h}(x)$, single-task learning estimates:

$$\hat{h}{\text{single}}(x) = h^*(x) + \epsilon{\text{task}} + \epsilon_{\text{noise}}$$

where $\epsilon_{\text{task}}$ represents task-specific biases and $\epsilon_{\text{noise}}$ represents fitting noise. Multi-task learning, by averaging across tasks, can reduce both error terms:

$$\hat{h}{\text{multi}}(x) \approx h^*(x) + \frac{1}{T}\sum{t=1}^{T}\epsilon_{\text{task}}^{(t)} + \frac{\epsilon_{\text{noise}}}{\sqrt{T}}$$

The task-specific biases may cancel or average out, and the noise term decreases with more tasks.

The information-theoretic perspective:

From an information-theoretic viewpoint, shared representations compress information that is relevant across tasks while discarding task-irrelevant details. Consider a representation $h(x)$ learned for tasks ${T_1, T_2, ..., T_k}$. An ideal shared representation maximizes:

$$I(h(X); Y_1, Y_2, ..., Y_k) - \beta \cdot I(h(X); X)$$

where $I(\cdot; \cdot)$ denotes mutual information. The first term encourages the representation to be predictive of all task outputs, while the second term (with coefficient $\beta$) encourages compressing away information from the input that isn't needed. This is a multi-task extension of the Information Bottleneck principle.

A good shared representation:

1. Retains information relevant to predicting all task labels
2. Discards information that is irrelevant or misleading for task predictions
3. Generalizes by focusing on invariant, causal features rather than spurious correlations

To rigorously understand shared representations, we need a formal mathematical framework. Let's establish the notation and key concepts that underpin multi-task learning theory.

Setup and Notation:

Suppose we have $T$ tasks, each defined by a distribution $\mathcal{D}t$ over input-output pairs $(x, y_t)$. For each task $t \in {1, 2, ..., T}$, we have a training set $S_t = {(x_i^{(t)}, y_i^{(t)})}{i=1}^{n_t}$.

In multi-task learning with shared representations, we decompose our model into:

1. Shared representation function: $h: \mathcal{X} \rightarrow \mathcal{H}$, mapping inputs to a representation space
2. Task-specific heads: $f_t: \mathcal{H} \rightarrow \mathcal{Y}_t$ for each task $t$

The prediction for task $t$ given input $x$ is:

$$\hat{y}_t = f_t(h(x))$$

The shared representation $h$ is parameterized by weights $\theta_{\text{shared}}$, and each task head $f_t$ is parameterized by $\theta_t$.

The Multi-Task Objective:

The standard multi-task learning objective combines task-specific losses:

$$\mathcal{L}{\text{MTL}}(\theta{\text{shared}}, \theta_1, ..., \theta_T) = \sum_{t=1}^{T} \lambda_t \mathcal{L}t(\theta{\text{shared}}, \theta_t)$$

where:

• $\mathcal{L}_t$ is the empirical loss for task $t$: $\mathcal{L}t = \frac{1}{n_t}\sum{i=1}^{n_t} \ell(f_t(h(x_i^{(t)})), y_i^{(t)})$
• $\lambda_t$ are task weights controlling the relative importance of each task
• $\ell(\cdot, \cdot)$ is a suitable loss function (e.g., cross-entropy, MSE)

Gradient Flow Through Shared Representations:

During backpropagation, the shared parameters receive gradients from all tasks:

$$ abla_{\theta_{\text{shared}}} \mathcal{L}{\text{MTL}} = \sum{t=1}^{T} \lambda_t abla_{\theta_{\text{shared}}} \mathcal{L}_t$$

This means the shared representation is shaped by the learning signals from all tasks. If tasks have conflicting gradients (pointing in different directions in parameter space), the shared representation must find a compromise—this is both a strength (regularization) and a challenge (gradient interference).

Representation Learning Dynamics:

The optimization dynamics of shared representations are governed by how task gradients interact. Let $g_t = abla_{\theta_{\text{shared}}} \mathcal{L}_t$ be the gradient for task $t$. The effective update to shared parameters is:

$$\Delta \theta_{\text{shared}} \propto -\sum_{t} \lambda_t g_t$$

Three scenarios can occur:

1. Aligned gradients ($g_i \cdot g_j > 0$): Tasks agree on the update direction. Learning each task reinforces the others.

2. Orthogonal gradients ($g_i \cdot g_j \approx 0$): Tasks require different features but don't conflict. The representation can learn features for both.

3. Conflicting gradients ($g_i \cdot g_j < 0$): Tasks disagree on the update direction. This creates negative transfer where improving one task hurts another.

Understanding these dynamics is crucial for designing effective MTL systems and diagnosing training issues.

The primary motivation for using shared representations is improved generalization. But why exactly does sharing lead to better generalization? This section provides a rigorous analysis of the statistical and computational benefits.

Sample Complexity Reduction:

Consider learning a representation $h$ of complexity $C_h$ (e.g., measured by VC dimension or Rademacher complexity). In single-task learning with $n$ samples:

$$\text{Generalization Error} \leq \hat{\mathcal{L}} + \mathcal{O}\left(\sqrt{\frac{C_h + C_f}{n}}\right)$$

In multi-task learning with $T$ tasks and $n$ samples per task:

$$\text{Generalization Error} \leq \hat{\mathcal{L}} + \mathcal{O}\left(\sqrt{\frac{C_h}{Tn}} + \sqrt{\frac{C_f}{n}}\right)$$

The shared representation complexity $C_h$ is now amortized across $T$ tasks, reducing its contribution to generalization error by a factor of $\sqrt{T}$. This is the fundamental statistical benefit: we learn the expensive representation using data from all tasks, while only the cheap task-specific heads need per-task data.

Inductive Bias as Implicit Regularization:

Beyond sample complexity, shared representations impose a beneficial inductive bias. By forcing the representation to be useful across multiple tasks, we bias the learning toward features that capture invariant, causal structure rather than spurious correlations.

Formally, suppose each task has a true underlying function $f_t^* = f_t \circ h^$ where $h^$ is the ideal shared representation. Single-task learning might find any $\tilde{h}$ such that $f_t \circ \tilde{h} \approx f_t^*$ on the training data—there are many such $\tilde{h}$, and most don't generalize. Multi-task learning restricts the search to representations that work for all tasks, greatly reducing the hypothesis space.

Theoretical Guarantees (Ben-David et al.):

Seminal work by Ben-David et al. provides formal guarantees for multi-task learning. Define the task relatedness as:

$$\mathcal{R}(T_1, T_2) = \sup_{h \in \mathcal{H}} |\mathcal{L}{T_1}(h) - \mathcal{L}{T_2}(h)|$$

This measures the maximum difference in loss any representation can achieve across the two tasks. When $\mathcal{R}$ is small, tasks are related and share optimal representations.

The generalization bound becomes:

$$\mathcal{L}{\text{test}}(h, f_t) \leq \hat{\mathcal{L}}{\text{train}} + \mathcal{R} + \text{complexity term}$$

When tasks are highly related ($\mathcal{R} \approx 0$), multi-task learning provides strong generalization guarantees.

Eavesdropping and Representation Bootstrapping:

An underappreciated benefit is that tasks can "eavesdrop" on each other's learning signals. A task $T_1$ that is difficult to learn directly might become easier when co-trained with a related task $T_2$, because the representation learned for $T_2$ provides useful features.

For example:

• Learning to parse sentences (syntax) helps with sentiment analysis (semantics)
• Learning to detect edges (low-level vision) helps with object recognition (high-level vision)
• Learning to predict next words (language modeling) helps with translation, QA, summarization, etc.

This creates a form of curriculum learning where easier tasks bootstrap harder tasks through shared representations.

Understanding what shared representations actually look like—their internal structure, geometry, and properties—provides insight into why they work and how to design better MTL systems.

Hierarchical Feature Organization:

In deep networks, shared representations typically exhibit hierarchical structure:

• Lower layers: Learn generic, broadly applicable features (edges, textures, basic phonemes, common word patterns)
• Middle layers: Learn more abstract, compositional features (object parts, phrase structures)
• Upper layers: Learn task-specific features (unless explicitly shared)

This hierarchy reflects the compositional structure of the world: complex concepts are built from simpler primitives that recur across domains.

Representation Geometry:

The geometry of shared representations reveals how tasks relate in the learned feature space. Key geometric properties include:

1. Linear Decodability: A well-structured shared representation allows each task to be solved by a simple (often linear) transformation. This suggests the representation organizes information in a way that separates task-relevant dimensions.

$$y_t \approx W_t \cdot h(x) + b_t$$

where $W_t$ and $b_t$ are the task-specific head parameters.

2. Subspace Organization: Tasks often utilize overlapping but distinct subspaces of the representation. Analysis shows that:

• Shared variance captures common features
• Task-unique variance captures task-specific information
• The ratio of shared to unique variance correlates with transfer effectiveness

3. Manifold Structure: Data from different tasks may lie on different manifolds within the representation space, connected by shared regions corresponding to common features.

Disentanglement and Factorization:

Ideal shared representations exhibit disentanglement: different factors of variation are encoded in separate dimensions. This makes it easier for task-specific heads to extract relevant information.

For multi-task learning, a useful property is task-aware factorization:

$$h(x) = [h_{\text{shared}}(x), h_{\text{task-specific}}(x)]$$

where $h_{\text{shared}}$ captures factors common across all tasks, and $h_{\text{task-specific}}$ captures factors relevant only to specific tasks. While full disentanglement is generally impossible without inductive biases or supervision, MTL naturally encourages partial disentanglement by forcing representations to serve multiple purposes.

The nature of shared representations varies dramatically across domains. Understanding these domain-specific patterns is essential for designing effective MTL systems.

Computer Vision:

In vision, shared representations follow the hierarchical feature learning paradigm established by deep convolutional networks:

• Layer 1: Gabor-like edge detectors, color blobs
• Layers 2-3: Textures, patterns, corners, conjunctions of edges
• Layers 4-6: Object parts, deformable templates
• Higher layers: Object-level representations, scene context

Typical MTL tasks (object detection, segmentation, depth estimation, surface normal prediction) share early convolutional layers and diverge at later stages. Empirical studies show that sharing up to layer 3-4 of networks like ResNet provides benefit, while sharing beyond that can hurt performance for dissimilar tasks.

Natural Language Processing:

In NLP, shared representations have evolved dramatically with the transformer era:

Pre-transformer: Shared word embeddings (Word2Vec, GloVe) provided basic semantic representations. Recurrent layers learned to compose these into sentence representations, shared across tasks like sentiment analysis, NER, and parsing.

Transformer era: Models like BERT represent the ultimate shared representation—a single pre-trained model provides contextual embeddings that transfer to virtually all NLP tasks. The representation is:

• Contextual (word meaning depends on surrounding context)
• Multi-layer (different layers capture different linguistic phenomena)
• Attention-based (explicitly models relationships between tokens)

Key insight: In NLP, the shared representation often is the model. Fine-tuning adds minimal task-specific parameters on top of massive shared representations.

Reinforcement Learning:

In RL, shared representations are particularly valuable because:

1. Sample efficiency is critical (environment interaction is expensive)
2. Related tasks share environment dynamics and reward structures
3. Skills learned in one task often transfer to others

Common sharing patterns:

• State representation: Learn a shared encoding of observations that captures relevant state features across tasks
• Policy networks: Share feature extraction layers, with task-specific policy heads
• Value functions: Shared representations for value estimation across task families

Multi-modal Learning:

Modern systems increasingly learn shared representations across modalities (text, images, audio). Key approaches:

• Joint embedding spaces: Project different modalities into a common representation space (e.g., CLIP aligns images and text)
• Cross-modal attention: Allow modalities to attend to each other through shared transformer architectures
• Fusion networks: Combine modality-specific representations into unified multimodal representations

The challenge is bridging very different data types while preserving modality-specific information.

Based on both theory and extensive empirical evidence, we can distill key principles for designing effective shared representations in multi-task learning systems.

Practical Implementation Guidelines:

1. Start with established backbones: For vision, use pre-trained ResNet/ViT; for NLP, use pre-trained transformers. These provide strong initial shared representations.

2. Use normalization carefully: Batch normalization can be problematic in MTL because statistics differ across tasks. Consider layer normalization or task-specific batch norm statistics.

3. Initialize properly: Task-specific heads should be initialized such that initial predictions are reasonable for all tasks. Bad initialization can cause one task to dominate early training.

4. Validate on all tasks: Multi-task performance should be evaluated holistically. Improvement on average with severe degradation on one task may be unacceptable.

5. Analyze representations: Periodically inspect the learned representations using the techniques described earlier. This provides insight into whether sharing is working as intended.

Shared representations are the conceptual and mathematical foundation of multi-task learning. By learning features that are useful across multiple tasks, we achieve better generalization, improved sample efficiency, and more robust models.

Next Steps:

With a solid understanding of shared representations, we're ready to explore specific mechanisms for implementing them. The next page covers Hard vs Soft Parameter Sharing—two fundamental paradigms for how parameters are shared across tasks in multi-task learning architectures.