Bayesian Uncertainty for Gradient Aggregation in Multi-Task Learning

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Achituve, Idan, Diamant, Idit, Netzer, Arnon, Chechik, Gal, Fetaya, Ethan
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917663584813056
author Achituve, Idan
Diamant, Idit
Netzer, Arnon
Chechik, Gal
Fetaya, Ethan
author_facet Achituve, Idan
Diamant, Idit
Netzer, Arnon
Chechik, Gal
Fetaya, Ethan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2402_04005
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Bayesian Uncertainty for Gradient Aggregation in Multi-Task Learning
Achituve, Idan
Diamant, Idit
Netzer, Arnon
Chechik, Gal
Fetaya, Ethan
Machine Learning
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.
title Bayesian Uncertainty for Gradient Aggregation in Multi-Task Learning
topic Machine Learning
url https://arxiv.org/abs/2402.04005