doktormike / evidentialflux.jl

Evidential Deep Learning Layers for Flux

EvidentialFlux

This is a Julia implementation in Flux of the Evidential Deep Learning framework. It strives to estimate heteroskedastic aleatoric uncertainty as well as epistemic uncertainty along with every prediction made. All of it calculated in one glorious forward pass. Boom!

Installing

If you want bleeding edge you can install it directly from my repo like this:

using Pkg; Pkg.add(url="https://github.com/DoktorMike/EvidentialFlux.jl")

Otherwise just do

using Pkg; Pkg.add("EvidentialFlux.jl")

For the impatient

Below is an example of how to train Deep Evidential Regression model, extract the predictions as well as the epistemic and aleatoric uncertainty. For a more elaborate example have a look in the examples folder.

using Flux
using EvidentialFlux

x = Float32.(-4:0.1:4)
y = x .^3 .+ randn(Float32, length(x)) .* 3

lr = 0.0005
m = Chain(Dense(1 => 100, relu), Dense(100 => 100, relu), Dense(100 => 100, relu), NIG(100 => 1))
opt = AdamW(lr, (0.89, 0.995), 0.001)
pars = Flux.params(m)
for epoch in 1:500
    grads = Flux.gradient(pars) do
        ŷ = m(x') 
        γ, ν, α, β = ŷ[1, :], ŷ[2, :], ŷ[3, :], ŷ[4, :]
        trnloss = Statistics.mean(nigloss2(y, γ, ν, α, β, 0.01, 2))
        trnloss
    end
    Flux.Optimise.update!(opt, pars, grads)
end

γ, ν, α, β = predict(m, x)
eu = epistemic(ν)
au = aleatoric(ν, α, β)

Classification

Deep evidential modeling works for classification as well as for regression. In the plot below you can see the epistemic uncertainty as a consequence of position in the plot. The task is to separate three Gaussians in 2D. The code for this example can be found in classification.jl.

Regression

In the case of a regression problem, we utilize the NormalInverseGamma distribution to model a type II likelihood function that then explicitly models the aleatoric and epistemic uncertainty. The code for the example producing the plot below can be found in regression.jl.

Summary

Uncertainty is crucial for the deployment and utilization of robust machine learning in production. No model is perfect and each one of them has strengths and weaknesses, but as a minimum requirement, we should all at least demand that our models report uncertainty in every prediction.

evidentialflux.jl's Issues

Implement corrections to DER

This is from a recent paper on Arxiv. It provides corrections for the loss function as well as the equations for epistemic and aleatoric uncertainty.

Add additional distributions to Deep evidential regression

Currently we assume that all regression cases will want to use a y~Normal(mu, sigma) model. This is definitely not always true and there may be many other distributions which could be useful for a deep learning setting. For instance.

• Poisson
• Binomial
• NegativeBinomial
• StudentT (This might be a bit flaky since we do end up with a model evidence following a StudentT either way)
• Cauchy

TagBot trigger issue

This issue is used to trigger TagBot; feel free to unsubscribe.

If you haven't already, you should update your TagBot.yml to include issue comment triggers.
Please see this post on Discourse for instructions and more details.

If you'd like for me to do this for you, comment TagBot fix on this issue.
I'll open a PR within a few hours, please be patient!

Implement MVE network

We should support a simple MVE network, i.e., a Normal distribution where we estimate mu and sigma separately.

Create a comparison between VI and DER

I would like to see a head to head comparison of the uncertainties produced by Variational Inference and Deep Evidential Regression for a moderately complex neural network.

This could be anything related to regression which is non trivial.

Implement the evidence function for the deep evidential regression case

Implement deep evidential classification