Investigating Uncertainty Weighting for Multi-Task Learning: Insights and Analytical Alternative

Abstract Multi-task learning (MTL) enables a single neural network to solve multiple tasks simultaneously, offering efficiency and improved generalization potential through shared representations. A central challenge in MTL is balancing task-specific losses during training to avoid performance degradation. While uncertainty-based loss weighting (UW) is a popular and competitive approach, we argue that it suffers from several limitations, including overfitting, rigid homoscedastic assumptions, and a lack of theoretical grounding for various loss functions. Therefore, we propose Soft Optimal Uncertainty Weighting (UW-SO), a novel loss weighting method that builds on UW by deriving analytically optimal weights and applying softmax normalization with adaptable temperature parameter, thereby alleviating several of the shortcomings of UW. Through extensive experiments across diverse datasets and architectures, we show that UW-SO achieves superior and robust performance compared to a variety of existing loss weighting methods. Additionally, we provide insights into the effects of temperature selection and propose measures to reduce computational demand.