Gumbel-Softmax Selective Networks

TL;DR

This paper introduces a novel training method for selective neural networks using the Gumbel-softmax trick, enabling end-to-end differentiable optimization for models that can abstain when uncertain.

Contribution

It proposes a general Gumbel-softmax based approach for training selective networks, addressing the non-differentiability challenge of the abstain decision.

Findings

Effective in selective regression and classification tasks

Enables end-to-end training of abstain options

Shows promising results on public datasets

Abstract

ML models often operate within the context of a larger system that can adapt its response when the ML model is uncertain, such as falling back on safe defaults or a human in the loop. This commonly encountered operational context calls for principled techniques for training ML models with the option to abstain from predicting when uncertain. Selective neural networks are trained with an integrated option to abstain, allowing them to learn to recognize and optimize for the subset of the data distribution for which confident predictions can be made. However, optimizing selective networks is challenging due to the non-differentiability of the binary selection function (the discrete decision of whether to predict or abstain). This paper presents a general method for training selective networks that leverages the Gumbel-softmax reparameterization trick to enable selection within an end-to-end differentiable training framework. Experiments on public datasets demonstrate the potential of Gumbel-softmax selective networks for selective regression and classification.