Gumbel-Softmax Selective Networks
This addresses the problem of integrating uncertainty-aware abstention into ML models for systems that require reliable predictions, such as those with fallback mechanisms, though it appears incremental as it builds on existing selective network concepts.
The paper tackles the challenge of training selective neural networks that can abstain from predicting when uncertain, by introducing a method using the Gumbel-softmax reparameterization trick to make the binary selection function differentiable, enabling end-to-end training.
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.