# Bayesian Uncertainty for Gradient Aggregation in Multi-Task Learning

arXiv: https://arxiv.org/abs/2402.04005 · HTML: https://mlsanity.com/paper/2402.04005/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 57 of its 84 references that are in the map.
- First posted on arXiv 6 Feb 2024.

## Builds on (most cited first, 30 of 57)

- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases](https://mlsanity.com/paper/1705.02315/) (2017)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [An Overview of Multi-Task Learning in Deep Neural Networks](https://mlsanity.com/paper/1706.05098/) (2017)
- [Merging Models with Fisher-Weighted Averaging](https://mlsanity.com/paper/2111.09832/) (2021)
- [End-to-End Multi-Task Learning with Attention](https://mlsanity.com/paper/1803.10704/) (2018)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [Cross-stitch Networks for Multi-task Learning](https://mlsanity.com/paper/1604.03539/) (2016)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [Age Progression/Regression by Conditional Adversarial Autoencoder](https://mlsanity.com/paper/1702.08423/) (2017)
- [Which Tasks Should Be Learned Together in Multi-task Learning?](https://mlsanity.com/paper/1905.07553/) (2019)
- [A Comprehensive Survey on Pretrained Foundation Models: A History from BERT to ChatGPT](https://mlsanity.com/paper/2302.09419/) (2023)
- [Laplace Redux -- Effortless Bayesian Deep Learning](https://mlsanity.com/paper/2106.14806/) (2021)
- [Multi-Task Deep Neural Networks for Natural Language Understanding](https://mlsanity.com/paper/1901.11504/) (2019)
- [Scalable Bayesian Optimization Using Deep Neural Networks](https://mlsanity.com/paper/1502.05700/) (2015)
- [Transformers in Medical Imaging: A Survey](https://mlsanity.com/paper/2201.09873/) (2022)
- [Order Matters: Sequence to sequence for sets](https://mlsanity.com/paper/1511.06391/) (2015)
- [Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks](https://mlsanity.com/paper/2002.10118/) (2020)
- [Multi-Task Learning as a Bargaining Game](https://mlsanity.com/paper/2202.01017/) (2022)
- [Deep Learning-Based Human Pose Estimation: A Survey](https://mlsanity.com/paper/2012.13392/) (2020)
- [Expectation Propagation for approximate Bayesian inference](https://mlsanity.com/paper/1301.2294/) (2013)

## Built on by (most cited first, 11 of 11)

- [MTL-UE: Learning to Learn Nothing for Multi-Task Learning](https://mlsanity.com/paper/2505.05279/) (2025)
- [LDC-MTL: Balancing Multi-Task Learning through Scalable Loss Discrepancy Control](https://mlsanity.com/paper/2502.08585/) (2025)
- [Go Beyond Your Means: Unlearning with Per-Sample Gradient Orthogonalization](https://mlsanity.com/paper/2503.02312/) (2025)
- [Aligned Multi Objective Optimization](https://mlsanity.com/paper/2502.14096/) (2025)
- [DeepChest: Dynamic Gradient-Free Task Weighting for Effective Multi-Task Learning in Chest X-ray Classification](https://mlsanity.com/paper/2505.23595/) (2025)
- [AutoScale: Linear Scalarization Guided by Multi-Task Optimization Metrics](https://mlsanity.com/paper/2508.13979/) (2025)
- [Self Supervised Correlation-based Permutations for Multi-View Clustering](https://mlsanity.com/paper/2402.16383/) (2024)
- [TapWeight: Reweighting Pretraining Objectives for Task-Adaptive Pretraining](https://mlsanity.com/paper/2410.10006/) (2024)
- [GCDance: Genre-Controlled Music-Driven 3D Full Body Dance Generation](https://mlsanity.com/paper/2502.18309/) (2025)
- [Simple Optimizers for Convex Aligned Multi-Objective Optimization](https://mlsanity.com/paper/2509.05811/) (2025)
- [Uniform Loss vs. Specialized Optimization: A Comparative Analysis in Multi-Task Learning](https://mlsanity.com/paper/2505.10347/) (2025)

## Abstract (arXiv)

As machine learning becomes more prominent there is a growing demand to perform several inference tasks in parallel. Running a dedicated model for each task is computationally expensive and therefore there is a great interest in multi-task learning (MTL). MTL aims at learning a single model that solves several tasks efficiently. Optimizing MTL models is often achieved by computing a single gradient per task and aggregating them for obtaining a combined update direction. However, these approaches do not consider an important aspect, the sensitivity in the gradient dimensions. Here, we introduce a novel gradient aggregation approach using Bayesian inference. We place a probability distribution over the task-specific parameters, which in turn induce a distribution over the gradients of the tasks. This additional valuable information allows us to quantify the uncertainty in each of the gradients dimensions, which can then be factored in when aggregating them. We empirically demonstrate the benefits of our approach in a variety of datasets, achieving state-of-the-art performance.