Multi-Task & Transfer Learning
Multi-task learning (MTL) jointly trains a model on multiple related tasks, sharing representations to improve generalization. Transfer learning leverages knowledge from source tasks to improve performance on target tasks. This module covers sharing strategies, gradient conflicts, and modern adaptation methods.
1. Problem Formulation
Multi-Task Objective
Minimize a weighted combination of task losses:
where are shared parameters and are task-specific parameters.
Transfer Learning Objective
Adapt a pretrained model to a new task with data :
2. Sharing Strategies
2.1 Hard Parameter Sharing
All tasks share all hidden layers:
Benefits: Reduces overfitting, parameter efficient Limitations: Negative transfer when tasks conflict
2.2 Soft Parameter Sharing
Each task has its own network with regularization:
Benefits: Task-specific adaptation, reduced negative transfer Limitations: More parameters, less sharing
2.3 Cross-Stitch Networks
Learn linear combinations of task features:
where are learned mixing coefficients.
2.4 Multi-Task Attention Networks
Attention-based sharing:
where is the task query and are shared keys/values.
2.5 Adapters
Lightweight modules inserted into pretrained networks:
where , , .
3. Negative Transfer
3.1 Definition
Negative transfer occurs when learning task hurts performance on task :
3.2 Detection
Task similarity metric: Compute correlation between task gradients:
Negative transfer occurs when .
3.3 Mitigation
Selective sharing: Only share layers where gradients are aligned.
Task grouping: Cluster tasks by gradient similarity:
4. Task Weighting Methods
4.1 GradNorm
GradNorm (Chen et al., 2018) balances gradient magnitudes across tasks:
Gradient magnitude target:
where is the average gradient norm.
Update rule:
where is the learning rate for task weights.
4.2 PCGrad: Projecting Conflicting Gradients
PCGrad (Yu et al., 2020) resolves gradient conflicts:
Conflict detection: Gradients conflict if .
Projection: For conflicting pair :
This projects onto the normal plane of .
PCGrad with random selection: Randomly pair conflicting gradients and project iteratively.
4.3 Nash-MTL
Model multi-task learning as a Nash game:
Solve via alternating optimization.
4.4 CAGrad
Confidence-Aware Gradient (Li et al., 2021): Weight gradients by task confidence:
5. Transfer Learning
5.1 Fine-Tuning Strategies
Full fine-tuning: Update all parameters:
Feature extraction: Freeze pretrained layers, train only new layers:
Progressive unfreezing: Gradually unfreeze layers during training.
5.2 Domain Adaptation
Adapt from source domain to target domain :
Covariate shift: but
Importance weighting:
5.3 Prompt-Based Transfer
Pretrained language models with task-specific prompts:
The model predicts the mask token as the task output.
6. Implementation
import torch
import torch.nn as nn
import torch.nn.functional as F
class GradNorm(nn.Module):
def __init__(self, num_tasks, alpha=1.5):
super().__init__()
self.num_tasks = num_tasks
self.alpha = alpha
self.task_weights = nn.Parameter(torch.ones(num_tasks))
self.initial_losses = None
def forward(self, losses):
if self.initial_losses is None:
self.initial_losses = [l.item() for l in losses]
weighted_loss = sum(
self.task_weights[t] * losses[t]
for t in range(self.num_tasks)
)
# Compute gradient norms
grad_norms = []
for t in range(self.num_tasks):
grad = torch.autograd.grad(
losses[t], self.task_weights, retain_graph=True
)[0]
grad_norms.append(grad.abs().item())
avg_norm = sum(grad_norms) / self.num_tasks
# Compute target norms
loss_ratio = [losses[t] / sum(losses) for t in range(self.num_tasks)]
init_ratio = [self.initial_losses[t] / sum(self.initial_losses)
for t in range(self.num_tasks)]
target_norms = [
avg_norm * (loss_ratio[t] / init_ratio[t]) ** self.alpha
for t in range(self.num_tasks)
]
# GradNorm loss
grad_norm_loss = sum(
abs(grad_norms[t] - target_norms[t])
for t in range(self.num_tasks)
)
return weighted_loss + grad_norm_loss
class PCGrad:
def __init__(self, optimizer):
self.optimizer = optimizer
def pc_backward(self, losses):
grads = []
for loss in losses:
self.optimizer.zero_grad()
loss.backward(retain_graph=True)
grads.append([p.grad.clone() for p in self.optimizer.param_groups[0]['params']])
# For each pair of tasks
num_tasks = len(losses)
for i in range(num_tasks):
for j in range(i + 1, num_tasks):
# Check for gradient conflict
dot_product = sum(
(grads[i][k] * grads[j][k]).sum()
for k in range(len(grads[i]))
)
norm_j = sum(grads[j][k].norm().pow(2) for k in range(len(grads[j])))
if dot_product < 0:
# Project grad_i onto normal plane of grad_j
scale = dot_product / (norm_j + 1e-8)
for k in range(len(grads[i])):
grads[i][k] -= scale * grads[j][k]
# Apply projected gradients
for k, param in enumerate(self.optimizer.param_groups[0]['params']):
param.grad = sum(grads[t][k] for t in range(num_tasks)) / num_tasks
self.optimizer.step()
class Adapter(nn.Module):
def __init__(self, input_dim, adapter_dim=64):
super().__init__()
self.down_proj = nn.Linear(input_dim, adapter_dim)
self.up_proj = nn.Linear(adapter_dim, input_dim)
self.activation = nn.GELU()
self.scale = nn.Parameter(torch.ones(1))
def forward(self, x):
residual = x
h = self.activation(self.down_proj(x))
h = self.up_proj(h)
return residual + self.scale * h
class MultiTaskModel(nn.Module):
def __init__(self, pretrained_model, num_tasks, adapter_dim=64):
super().__init__()
self.encoder = pretrained_model
self.adapters = nn.ModuleList([
Adapter(pretrained_model.hidden_size, adapter_dim)
for _ in range(num_tasks)
])
self.task_heads = nn.ModuleList([
nn.Linear(pretrained_model.hidden_size, num_classes)
for num_classes in task_num_classes
])
def forward(self, x, task_id):
features = self.encoder(x)
adapted = self.adapters[task_id](features)
return self.task_heads[task_id](adapted)
7. SVG: Multi-Task Architecture
8. SVG: Negative Transfer Illustration
9. Comparison of Methods
| Method | Gradient Handling | Task Weights | Computational Cost | Scalability |
|---|---|---|---|---|
| Equal Weighting | None | Fixed | Low | Good |
| GradNorm | Gradient magnitude | Learned | Medium | Good |
| PCGrad | Gradient projection | None | High | Moderate |
| CAGrad | Confidence-weighted | Learned | Medium | Good |
| Nash-MTL | Game-theoretic | Learned | High | Moderate |
| Uncertainty | Task variance | Learned | Medium | Good |
10. Open Problems
- Dynamic task addition: Adding new tasks without retraining
- Task imbalance: Handling varying dataset sizes across tasks
- Multi-task pre-training: Choosing pre-training tasks for downstream transfer
- Curriculum MTL: Ordering tasks during training
- Theoretical understanding: When does MTL help vs hurt?
References
- Caruana, R. (1997). Multitask Learning. Machine Learning.
- Chen, Z., et al. (2018). GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks. ICML.
- Yu, T., et al. (2020). Gradient Surgery for Multi-Task Learning. NeurIPS.
- Liu, B., et al. (2021). Gradient Vaccine: Investigating and Improving Multi-task Optimization. ICML.
- Pfeiffer, J., et al. (2021). AdapterFusion: Non-Destructive Task Composition. EACL.