[Paper Review] Multi-Task Variational Information Bottleneck

Weizhu Qian, Bowei Chen+3 more|arXiv (Cornell University)|Jul 1, 2020

Adversarial Robustness in Machine Learning62 references4 citations

TL;DR

This paper proposes a Multi-Task Variational Information Bottleneck (MTVIB) model that integrates variational inference with task-dependent uncertainty weighting to learn robust, disentangled latent representations for multi-task learning. By jointly optimizing for compressed, task-relevant representations and adaptive loss weighting, MTVIB achieves state-of-the-art accuracy and superior robustness under adversarial attacks across three benchmark datasets.

ABSTRACT

Multi-task learning (MTL) is an important subject in machine learning and artificial intelligence. Its applications to computer vision, signal processing, and speech recognition are ubiquitous. Although this subject has attracted considerable attention recently, the performance and robustness of the existing models to different tasks have not been well balanced. This article proposes an MTL model based on the architecture of the variational information bottleneck (VIB), which can provide a more effective latent representation of the input features for the downstream tasks. Extensive observations on three public data sets under adversarial attacks show that the proposed model is competitive to the state-of-the-art algorithms concerning the prediction accuracy. Experimental results suggest that combining the VIB and the task-dependent uncertainties is a very effective way to abstract valid information from the input features for accomplishing multiple tasks.

Motivation & Objective

• To address the imbalance in performance and robustness across multiple tasks in existing multi-task learning (MTL) models.
• To improve latent representation quality by combining the Variational Information Bottleneck (VIB) with task-specific uncertainty estimation.
• To dynamically balance loss contributions across tasks using likelihood-based uncertainty weighting, avoiding manual or fixed hyperparameter tuning.
• To enhance model robustness against adversarial attacks while maintaining high predictive accuracy on multiple downstream tasks.
• To frame multi-task learning as a constrained multi-objective optimization problem using information bottleneck principles.

Proposed method

• The model uses a variational inference framework to learn a posterior distribution over latent codes, enabling regularized and noise-robust representations.
• A shared encoder maps input features to a latent distribution, while task-specific decoders generate predictions for each downstream task.
• Task-dependent uncertainties are learned as parameters of the likelihood distribution, enabling automatic, adaptive weighting of task-specific losses.
• The objective function minimizes a variational approximation of the information bottleneck, balancing mutual information between input and latent code, and between latent code and output.
• The framework is trained end-to-end using stochastic gradient descent, with the total loss being a weighted sum of task-specific cross-entropy losses, where weights are derived from learned uncertainties.
• The model is evaluated under adversarial attacks with varying perturbation levels (ε) to assess robustness and generalization.

Experimental results

Research questions

• RQ1 Can a variational information bottleneck framework improve the quality and robustness of shared latent representations in multi-task learning?
• RQ2 Does learning task-specific uncertainties lead to better dynamic loss weighting and improved performance balance across tasks?
• RQ3 How does the MTVIB model compare to state-of-the-art MTL models in terms of accuracy and robustness under adversarial attacks?
• RQ4 To what extent does the disentangled latent representation learned via VIB enhance generalization across multiple tasks?
• RQ5 Can the joint optimization of representation compression and loss weighting lead to superior performance on multi-task benchmarks?

Key findings

• On the MultiMNIST dataset, MTVIB achieved 93.7% test accuracy under ε=0.05, outperforming STL (87.1%) and STVIB (86.7%), and matching or exceeding other MTL baselines.
• On the MultiFashion dataset, MTVIB achieved 58.2% accuracy at ε=0.05, significantly outperforming STL (59.1%), STVIB (60.1%), and GS (57.1%), showing improved robustness.
• On the MTFL dataset, MTVIB achieved 34.4% accuracy at ε=0.05 for task 1 and 22.6% at ε=0.30 for task 2, demonstrating strong resilience under high-perturbation attacks.
• MTVIB consistently outperformed or matched the best-performing models (e.g., UWL, GS) across all datasets and adversarial levels, with the highest accuracy on MTFL (34.4% at ε=0.05 for task 1).
• The model showed superior robustness, maintaining higher accuracy under increasing adversarial noise (ε) compared to other MTL models, indicating effective regularization through VIB and uncertainty weighting.
• The ablation study confirmed that combining VIB with uncertainty weighting yields better performance than using either component alone, validating the synergy of the two mechanisms.