Multi-Task Training AI Prompt

Implement a multi-task learning training setup for a model that simultaneously optimizes {{task_1}} and {{task_2}}.

1. Model architecture:
   - Shared backbone: {{backbone}} that extracts shared representations
   - Task-specific heads: separate output heads for each task
   - Gradient isolation: ensure gradients from one task head do not corrupt features needed by another

2. Loss combination strategies — implement and compare:
   a. Fixed weighting: total_loss = w1 × loss_1 + w2 × loss_2
   b. Uncertainty weighting (Kendall et al. 2018): learn task weights as trainable parameters based on homoscedastic uncertainty
   c. GradNorm (Chen et al. 2018): dynamically adjust weights based on relative gradient magnitudes

3. Task imbalance handling:
   - Normalize each task loss to similar scale before combining
   - Monitor per-task gradient norms — large imbalance indicates weighting issues

4. Training strategy:
   - Option A: alternate between tasks each batch
   - Option B: sample tasks proportionally by dataset size
   - Option C: train all tasks simultaneously in each batch

5. Evaluation:
   - Log per-task metrics separately during validation
   - Use a combined score (e.g. average of normalized per-task metrics) to select the best checkpoint

6. Gradient surgery: implement PCGrad to project conflicting gradients to prevent task interference

Return: multi-task model code, loss combination implementations, and training loop with per-task logging.