# Variational Multi-Task Learning with Gumbel-Softmax Priors

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

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 59 references that are in the map.
- First posted on arXiv 9 Nov 2021.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Moment Matching for Multi-Source Domain Adaptation](https://mlsanity.com/paper/1812.01754/) (2018)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Deep Hashing Network for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1706.07522/) (2017)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [A Survey on Multi-Task Learning](https://mlsanity.com/paper/1707.08114/) (2017)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Generating Sentences from a Continuous Space](https://mlsanity.com/paper/1511.06349/) (2015)
- [End-to-End Multi-Task Learning with Attention](https://mlsanity.com/paper/1803.10704/) (2018)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [Cross-stitch Networks for Multi-task Learning](https://mlsanity.com/paper/1604.03539/) (2016)
- [Attentive Neural Processes](https://mlsanity.com/paper/1901.05761/) (2019)
- [Variational Dropout Sparsifies Deep Neural Networks](https://mlsanity.com/paper/1701.05369/) (2017)
- [Probabilistic Model-Agnostic Meta-Learning](https://mlsanity.com/paper/1806.02817/) (2018)
- [VAE with a VampPrior](https://mlsanity.com/paper/1705.07120/) (2017)
- [Variational Mixture-of-Experts Autoencoders for Multi-Modal Deep Generative Models](https://mlsanity.com/paper/1911.03393/) (2019)
- [AdaShare: Learning What To Share For Efficient Deep Multi-Task Learning](https://mlsanity.com/paper/1911.12423/) (2019)
- [Inference Suboptimality in Variational Autoencoders](https://mlsanity.com/paper/1801.03558/) (2018)
- [Learning to Branch for Multi-Task Learning](https://mlsanity.com/paper/2006.01895/) (2020)
- [Meta-Learning Probabilistic Inference For Prediction](https://mlsanity.com/paper/1805.09921/) (2018)
- [Many Task Learning with Task Routing](https://mlsanity.com/paper/1903.12117/) (2019)
- [MTL-NAS: Task-Agnostic Neural Architecture Search towards General-Purpose Multi-Task Learning](https://mlsanity.com/paper/2003.14058/) (2020)
- [Stochastic Filter Groups for Multi-Task CNNs: Learning Specialist and Generalist Convolution Kernels](https://mlsanity.com/paper/1908.09597/) (2019)
- [Variational Autoencoder with Implicit Optimal Priors](https://mlsanity.com/paper/1809.05284/) (2018)
- [Efficient Output Kernel Learning for Multiple Tasks](https://mlsanity.com/paper/1511.05706/) (2015)

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

- [FAMO: Fast Adaptive Multitask Optimization](https://mlsanity.com/paper/2306.03792/) (2023)
- [Recon: Reducing Conflicting Gradients from the Root for Multi-Task Learning](https://mlsanity.com/paper/2302.11289/) (2023)
- [MmAP : Multi-modal Alignment Prompt for Cross-domain Multi-task Learning](https://mlsanity.com/paper/2312.08636/) (2023)
- [Pareto Manifold Learning: Tackling multiple tasks via ensembles of single-task models](https://mlsanity.com/paper/2210.09759/) (2022)
- [GRIDS: Grouped Multiple-Degradation Restoration with Image Degradation Similarity](https://mlsanity.com/paper/2407.12273/) (2024)
- [GO4Align: Group Optimization for Multi-Task Alignment](https://mlsanity.com/paper/2404.06486/) (2024)
- [A Simple Yet Effective Strategy to Robustify the Meta Learning Paradigm](https://mlsanity.com/paper/2310.00708/) (2023)
- [Episodic Multi-Task Learning with Heterogeneous Neural Processes](https://mlsanity.com/paper/2310.18713/) (2023)
- [Learning Generalizable Skills from Offline Multi-Task Data for Multi-Agent Cooperation](https://mlsanity.com/paper/2503.21200/) (2025)
- [Gradient Similarity Surgery in Multi-Task Deep Learning](https://mlsanity.com/paper/2506.06130/) (2025)
- [FunOTTA: On-the-Fly Adaptation on Cross-Domain Fundus Image via Stable Test-time Training](https://mlsanity.com/paper/2407.04396/) (2024)
- [Learning a Consensus Sub-Network with Polarization Regularization and One Pass Training](https://mlsanity.com/paper/2302.10798/) (2023)
- [An Information-theoretic Multi-task Representation Learning Framework for Natural Language Understanding](https://mlsanity.com/paper/2503.04667/) (2025)
- [Balancing Efficiency and Quality: MoEISR for Arbitrary-Scale Image Super-Resolution](https://mlsanity.com/paper/2311.12077/) (2023)
- [Generalized Radius and Integrated Codebook Transforms for Differentiable Vector Quantization](https://mlsanity.com/paper/2602.01140/) (2026)

## Abstract (arXiv)

Multi-task learning aims to explore task relatedness to improve individual tasks, which is of particular significance in the challenging scenario that only limited data is available for each task. To tackle this challenge, we propose variational multi-task learning (VMTL), a general probabilistic inference framework for learning multiple related tasks. We cast multi-task learning as a variational Bayesian inference problem, in which task relatedness is explored in a unified manner by specifying priors. To incorporate shared knowledge into each task, we design the prior of a task to be a learnable mixture of the variational posteriors of other related tasks, which is learned by the Gumbel-Softmax technique. In contrast to previous methods, our VMTL can exploit task relatedness for both representations and classifiers in a principled way by jointly inferring their posteriors. This enables individual tasks to fully leverage inductive biases provided by related tasks, therefore improving the overall performance of all tasks. Experimental results demonstrate that the proposed VMTL is able to effectively tackle a variety of challenging multi-task learning settings with limited training data for both classification and regression. Our method consistently surpasses previous methods, including strong Bayesian approaches, and achieves state-of-the-art performance on five benchmark datasets.