Learnable Bernoulli Dropout for Bayesian Deep Learning

TL;DR

This paper introduces learnable Bernoulli dropout (LBD), a model-agnostic scheme that optimizes dropout rates jointly with model parameters, enhancing robustness and uncertainty quantification in deep learning, especially when combined with VAEs.

Contribution

It proposes a novel learnable dropout method optimized via ARM, enabling flexible semi-implicit posterior representations in VAEs and improving performance across various tasks.

Findings

LBD outperforms traditional dropout schemes in accuracy and uncertainty estimation.

SIVAE achieves state-of-the-art results in collaborative filtering tasks.

The method provides robust predictions and better uncertainty quantification.

Abstract

In this work, we propose learnable Bernoulli dropout (LBD), a new model-agnostic dropout scheme that considers the dropout rates as parameters jointly optimized with other model parameters. By probabilistic modeling of Bernoulli dropout, our method enables more robust prediction and uncertainty quantification in deep models. Especially, when combined with variational auto-encoders (VAEs), LBD enables flexible semi-implicit posterior representations, leading to new semi-implicit VAE~(SIVAE) models. We solve the optimization for training with respect to the dropout parameters using Augment-REINFORCE-Merge (ARM), an unbiased and low-variance gradient estimator. Our experiments on a range of tasks show the superior performance of our approach compared with other commonly used dropout schemes. Overall, LBD leads to improved accuracy and uncertainty estimates in image classification and semantic segmentation. Moreover, using SIVAE, we can achieve state-of-the-art performance on collaborative filtering for implicit feedback on several public datasets.