Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes
Bayesian Uncertainty for Gradient Aggregation in Multi-Task Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes BayesAgg-MTL, which places a Bayesian posterior over last task-specific layer weights to induce Gaussian distributions over shared-representation gradients and aggregates them per dimension and example by precision weighting.
Problem
Gradient aggregation in multi-task learning ignores per-dimension sensitivity of task gradients, treating each task's gradient as a single deterministic vector.
Approach
Treat the last task-specific linear layer as Bayesian, propagate the posterior to a Gaussian approximation of each task's gradient w.r.t. the shared hidden layer (closed-form moments for regression, Laplace/GGN-based posterior with Monte Carlo moments for classification), and combine task gradients with per-dimension, per-example precision weights derived from a diagonal-Gaussian likelihood produc
Key results
- On QM9, BayesAgg-MTL reaches Delta_m 53.2 ± 7.1 versus 62.0 ± 1.4 for Nash-MTL (lower is better).
- On CIFAR-MTL, it reaches 59.97% accuracy, above MGDA's 59.74% and RLW's 59.30%.
- On UTKFace, it reaches Delta_m -2.23 ± 0.76 and is best on age and ethnicity, outperforming all baselines.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use BayesAgg-MTL instead of Nash-MTL | MTL on QM9 · QM9 test | Δm%: 62.0 ± 1.4 -> 53.2 ± 7.1 | Helped |
| Use BayesAgg-MTL instead of MGDA | MTL on CIFAR-MTL · CIFAR-MTL test | accuracy: 59.74 ± .07 -> 59.97 ± 0.4 | Helped |
| Use BayesAgg-MTL instead of DWA | MTL on ChestX-ray14 · ChestX-ray14 test | Δm%: -14.79 ± .07 -> -14.96 ± 0.1 | Helped |
| Use BayesAgg-MTL instead of MGDA | MTL on UTKFace · UTKFace test | Δm%: -1.39 ± 0.50 -> -2.23 ± 0.76 | Helped |
| Use BayesAgg-MTL instead of a 10-network deep ensemble | MTL on QM9 · QM9 | Δm%: ~144.5 ± 0.3 -> ~~53.2 ± 7.1 | Helped |
| Use BayesAgg-MTL instead of a 10-network deep ensemble | MTL on UTKFace · UTKFace | Δm%: -0.13 ± 0.39 -> -2.23 ± 0.76 | Helped |
Surprising findings
- Standard deep ensembles with a shared backbone perform roughly the same as linear scalarization, indicating that uncertainty must be used in gradient aggregation rather than only in training.
- On ChestX-ray14, almost all gradient-balancing baselines underperform naive linear scalarization, with only BayesAgg-MTL and DWA beating it.
Limitations
- Restricted to linear task-specific heads; posterior and gradient moment estimation is harder for nonlinear heads.
- Uses a diagonal Gaussian approximation and a precision scaling hyper-parameter s to avoid instability, indicating model misspecification.
- Classification requires Monte Carlo sampling (J=1024), which increases training cost relative to simpler methods.
Experimental setup
Datasets: QM9 CIFAR-100 (coarse labels, CIFAR-MTL) ChestX-ray14 UTKFace
Benchmarks: QM9 CIFAR-MTL ChestX-ray14 UTKFace
Metrics: Delta_m% accuracy AUC-ROC MAE expected calibration error Brier score
Baselines: Single Task Learning Linear Scalarization Scale-Invariant Random Loss Weighting Dynamic Weight Average Uncertainty weighting MGDA PCGrad CAGrad IMTL-G Nash-MTL IGBv2
Builds on: Nash-MTL Laplace approximation Generalized Gauss-Newton MGDA PCGrad
Closest papers by meaning
- Towards Consistent Multi-Task Learning: Unlocking the Potential of Task-Specific Parameters (CVPR 2025 · similarity 0.89)
- Delve into the Applicability of Advanced Optimizers for Multi-Task Learning (NeurIPS 2026 · similarity 0.88)
- Fast Rate Bounds for Multi-Task and Meta-Learning with Different Sample Sizes (NeurIPS 2025 · similarity 0.88)
- Revisiting Fairness in Multitask Learning: A Performance-Driven Approach for Variance Reduction (CVPR 2025 · similarity 0.87)
- Fair Resource Allocation in Multi-Task Learning (ICML 2024 · similarity 0.87)
- Scalable Bayesian Learning with posteriors (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.