Package: uncertainty-loss
Uncertainty LossLoss functions for uncertainty quantification in deep learning.This package implements loss functions from the following papersEvidential Deep Learning to Quantify Classification Uncertaintyul.evidential_lossInformation Aware Max-Norm Dirichlet Networks for Predictive Uncertainty Estimationul.maxnorm_lossThese loss functions can be used as drop in replacements fortorch.nn.functional.cross_entropy.  See QuickStart and Examples below.QuickstartInstall the package with pippipinstalluncertainty-lossThen use the loss in a training pipeline. For example:importuncertainty_lossasulimporttorchdeffit_step(model,x,targets,reg_factor=0):"""Runs a single training step and retuns the loss for the batch.Note the inputs to the uncertainty loss function need to benon-negative.  Any transformation will work (exp, relu, softplus,etc) but we have found that exp works best (in agreement with theoriginal papers).  For convenience we provide a clamped exp functionto avoid overflow."""logits=model(x)evidence=ul.clamped_exp(logits)# non-negative transformloss=ul.maxnorm_loss(evidence,targets,reg_factor)returnlossExamplesReplacefromtorch.nnimportfunctionalasFloss=F.cross_entropy(x,y)Withimportuncertainy_lossasulloss=ul.evidential_loss(x,y)# orloss=ul.maxnorm_loss(x,y)The loss functions also accept a reduction parameter with the same
properties as thecross_entropyloss.ImportantFor each loss function is a regularization term that is shown to be
beneficial for learning to quantify uncertainty.  In practice,
to ensure that the regularization term does not dominate early
in training, we ramp up the regularization term from 0 to a max factor
e.g. 0->1.  It is up to the user to ensure this happens.  Each loss
function takes an additional parameterreg_factor.  During training
one can incrementreg_factorto accomplish this ramp up.  By
defaultreg_factor==0so there is no regularization unless
explicitly "turned on"Example with Regularization Annealingimportuncertainty_lossasulreg_steps=1000reg_step_size=1/reg_stepsreg_factor=0forepochinrange(epochs):forx,yindataloader:logits=model(x)evidence=ul.clamped_exp(logits)loss=ul.maxnorm_loss(evidence,y,reg_factor=reg_factor)reg_factor=min(reg_factor+reg_step_size,1)MotivationUncertainty quantification has important applications in AI Safety and active learning.  Neural networks trained with a traditional cross entropy loss are often over-confident in unfamiliar situations.  It's easy to see why this can be disastrous: An AI surgeon making a confident but wrong incision in an unfamilar situation, a self-driving car making a confident but wrong turn, an AI investor making a confident but wrong buy/sell decision.There have been several methods proposed for uncertainty quantification.  Many of the popular methods require specific network architectures (e.g. Monte Carlo Dropout requires dropout layers) or require expensive inference (Monte Carlo dropout requires multiple runs through the same model, ensemble methods require multiple models).Recently methods for uncertainty quantification have been proposed that do not require any changes to the network architecture and have no inference overhead.  Instead they propose to learn parameters of a "higher order distribution" and use this distribution to quantify the uncertainty in the prediction.  They have been shown to be effective.Unfortunately, these methods haven't been integrated into any of the main deep learning packages and the heavy math makes the implementation a bit tricky.For these reasons we have created theuncertainty-losspackage.