EvidentialFlux

A Julia/Flux implementation of the Evidential Deep Learning framework. Estimate heteroskedastic aleatoric and epistemic uncertainty in a single forward pass.

Installing

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

Features

EvidentialFlux provides evidential output layers for regression, classification, ordinal, and count data. All layers are subtypes of AbstractEvidentialLayer, which provides generic predict and split_params dispatch.

Layer	Use case	Output	Uncertainty
`NIG(in => out)`	Regression	γ, ν, α, β (4 × out)	Aleatoric + epistemic
`PG(in => out)`	Count regression	α, β (2 × out)	Aleatoric + epistemic
`EG(in => out)`	Positive continuous regression	α, β (2 × out)	Aleatoric + epistemic
`BB(in => out)`	Proportion/success-rate estimation	α, β (2 × out)	Aleatoric + epistemic
`BNB(in => out)`	Overdispersed count regression	r, α, β (3 × out)	Aleatoric + epistemic
`ZIP(in => out)`	Zero-inflated count regression	α_π, β_π, α_λ, β_λ (4 × out)	Aleatoric + epistemic
`VM(in => out)`	Directional/circular regression	μ₀, κ₀, κ (3 × out)	Aleatoric + epistemic
`DIR(in => out)`	Classification	Dirichlet concentrations (out)	Epistemic
`FDIR(in => out)`	Classification / ordinal	α, p, τ (2 × out + 1)	Aleatoric + epistemic
`MVE(in => out)`	Regression	μ, σ (2 × out)	Aleatoric

Ordinal targets (ordered classes) reuse the DIR/FDIR layer with an order-aware loss (ofdirloss) — see Evidential ordinal regression.

Loss functions

Function	Description
`nigloss(y, γ, ν, α, β, λ, ϵ)`	Standard evidential regression loss (Amini et al. 2020)
`nigloss_scaled(y, γ, ν, α, β, λ, p)`	Corrected DER loss (Meinert et al. 2022)
`nigloss_ureg(y, γ, ν, α, β, λ, λ₁)`	Uncertainty regularized loss (Ye et al. 2024)
`dirloss(y, α, t)`	Dirichlet classification loss with KL regularization, returns `(1, B)`
`dirloss_cor(y, α, t)`	Dirichlet loss + correct evidence regularization (Pandey et al. 2025)
`dirmultloss(y, α)`	Dirichlet-Multinomial NLL for count vector targets (reuses `DIR` layer)
`fdirloss(y, α, p, τ)`	Flexible Dirichlet loss (Yoon & Kim 2025)
`ofdirloss(y, α, p, τ; weights)`	Ordinal Flexible Dirichlet loss — expected Ranked Probability Score; optional per-class weights for imbalance
`pgloss(y, α, β, λ)`	Poisson-Gamma count regression loss (NLL + regularizer)
`egloss(y, α, β, λ)`	Exponential-Gamma positive regression loss (NLL + regularizer)
`bbloss(k, n, α, β, λ)`	Binomial-Beta proportion estimation loss (NLL + regularizer)
`bnbloss(y, r, α, β, λ)`	Beta-Negative Binomial count regression loss (NLL + regularizer)
`ziploss(y, α_π, β_π, α_λ, β_λ, λ)`	Zero-Inflated Poisson count regression loss (NLL + regularizer)
`vmloss(θ, μ₀, κ₀, κ, λ)`	Von Mises directional regression loss (NLL + regularizer)
`nllpg(y, α, β)`	Negative Binomial marginal NLL
`nlleg(y, α, β)`	Lomax (Pareto Type II) marginal NLL
`nllbb(k, n, α, β)`	Beta-Binomial marginal NLL
`nllbnb(y, r, α, β)`	Beta-Negative Binomial marginal NLL
`nllzip(y, α_π, β_π, α_λ, β_λ)`	Zero-Inflated Negative Binomial marginal NLL
`nllvm(θ, μ₀, κ₀, κ)`	Von Mises marginal NLL on the circle
`mveloss(y, μ, σ)`	Gaussian negative log-likelihood
`nllstudent(y, γ, ν, α, β)`	Student-T negative log-likelihood

Utilities

Function	Description
`predictive(model, x)`	Inference-time output: `(ŷ, epistemic, aleatoric, params)`
`predictive_mean(Type, params)`	Point prediction in data space
`predict(model, x)`	Raw distributional parameters (for training loops)
`split_params(LayerType, y)`	Generic output decomposition into a NamedTuple
`splitnig(y)`	Split concatenated NIG output into (γ, ν, α, β)
`splitmve(y)`	Split concatenated MVE output into (μ, σ)
`splitpg(y)`	Split concatenated PG output into (α, β)
`spliteg(y)`	Split concatenated EG output into (α, β)
`splitbb(y)`	Split concatenated BB output into (α, β)
`splitbnb(y)`	Split concatenated BNB output into (r, α, β)
`splitzip(y)`	Split concatenated ZIP output into (α_π, β_π, α_λ, β_λ)
`splitvm(y)`	Split concatenated VM output into (μ₀, κ₀, κ)
`splitfdir(y)`	Split concatenated FDIR output into (α, p, τ)
`evidence(ν, α)`	NIG total evidence: 2ν + α
`evidence(α)`	DIR evidence: α - 1

Inference with `predictive`

For inference, use predictive to get the point prediction, uncertainties, and raw parameters in a single call:

r = predictive(model, x)
r.ŷ          # point prediction in data space
r.epistemic  # epistemic uncertainty (nothing if N/A)
r.aleatoric  # aleatoric uncertainty (nothing if N/A)
r.params     # raw distributional parameters for advanced use

The point prediction (ŷ) depends on the layer type:

Layer	`ŷ`	Formula
NIG	Mean of posterior predictive (Student-T)	γ
PG	Expected count	α/β
EG	Expected duration	β/(α-1)
BB	Expected count (n=1 default)	n·α/(α+β)
BNB	Expected count at Beta mean	r·α/β
ZIP	Expected count	β_π/(α_π+β_π)·α_λ/β_λ
VM	Mean direction	μ₀
DIR	Expected counts (n=1 default)	n·α/Σα
FDIR	Expected counts (n=1 default)	n·(α + τp)/(Σα + τ)
MVE	Predicted mean	μ

Training with `predict`

For training, use predict which returns raw distributional parameters needed by the loss functions. These are NamedTuples you can destructure:

γ, ν, α, β = predict(model, x)  # NIG — destructure for loss computation
loss = sum(nigloss(y, γ, ν, α, β))

Uncertainty

All layers support a unified type-dispatched API for uncertainty decomposition. Pass the layer type as the first argument:

eu = epistemic(NIG, ν, α, β)
au = aleatoric(NIG, ν, α, β)

Layer	`epistemic(Type, ...)`	`aleatoric(Type, ...)`
`NIG`	`epistemic(NIG, ν, α, β)` = 1/√ν	`aleatoric(NIG, ν, α, β)` = β(1+ν)/(να)
`DIR`	`epistemic(DIR, α)` = K/Σα	—
`MVE`	—	`aleatoric(MVE, σ)` = σ
`PG`	`epistemic(PG, α, β)` = α/β²	`aleatoric(PG, α, β)` = α/β
`EG`	`epistemic(EG, α, β)` = β²/((α-1)²(α-2))	`aleatoric(EG, α, β)` = β²/((α-1)(α-2))
`BB`	`epistemic(BB, α, β)` = αβ/((α+β)²(α+β+1))	`aleatoric(BB, α, β)` = αβ/((α+β)(α+β+1))
`BNB`	`epistemic(BNB, r, α, β)` = r²α(α+β-1)/((β-1)²(β-2))	`aleatoric(BNB, r, α, β)` = rα(α+β-1)/((β-1)(β-2))
`ZIP`	`epistemic(ZIP, α_π, β_π, α_λ, β_λ)` = Var[(1-π)λ]	`aleatoric(ZIP, α_π, β_π, α_λ, β_λ)` = E[Var[Y\|π,λ]]
`VM`	`epistemic(VM, κ₀)` = 1 - I₁(κ₀)/I₀(κ₀)	`aleatoric(VM, κ)` = 1 - I₁(κ)/I₀(κ)
`FDIR`	`epistemic(FDIR, α, p, τ)`	`aleatoric(FDIR, α, p, τ)` = TU - EU

For NIG, the legacy arity-dispatched functions (uncertainty(ν, α, β), uncertainty(α, β), epistemic(ν), aleatoric(ν, α, β)) remain available for backward compatibility.

Notes:

DIR and MVE only expose one uncertainty type (epistemic and aleatoric, respectively)
EG requires α > 2 for the moments to exist; values are clamped internally
BNB requires β > 2 for the moments to exist; values are clamped internally
FDIR uncertainties are per-sample (1, B), derived from the FD mixture-of-Dirichlets decomposition (Yoon & Kim 2025)
ZIP uncertainties are derived via the law of total variance over independent Beta and Gamma priors
VM uncertainties are circular variances in [0, 1], where 0 = certain and 1 = uniform on the circle
PG and BNB uncertainties are per-output (O, B), derived via the law of total variance

Quick start

Evidential regression (NIG)

using Flux, EvidentialFlux, Statistics

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

model = Chain(Dense(1 => 100, relu), Dense(100 => 100, relu), NIG(100 => 1))
opt_state = Flux.setup(AdamW(1e-3), model)

for epoch in 1:3000
    loss, grads = Flux.withgradient(model) do m
        γ, ν, α, β = splitnig(m(x'))
        mean(nigloss_scaled(y, γ, ν, α, β, 0.01))
    end
    Flux.update!(opt_state, model, grads[1])
end

# Inference: get predictions and uncertainty in one call
r = predictive(model, x')
r.ŷ          # predicted value (γ)
r.epistemic  # high when extrapolating beyond training data
r.aleatoric  # high when data is inherently noisy

Evidential classification (DIR)

model = Chain(Dense(2 => 64, relu), Dense(64 => 64, relu), DIR(64 => 3))
opt_state = Flux.setup(AdamW(1e-2), model)

for epoch in 1:500
    loss, grads = Flux.withgradient(model) do m
        sum(dirloss(y_onehot, m(x), epoch))
    end
    Flux.update!(opt_state, model, grads[1])
end

r = predictive(model, x_test)
r.ŷ          # class probabilities
r.epistemic  # high for out-of-distribution inputs

Evidential ordinal regression (FDIR)

When the K classes have a natural order (Very Low < Low < … < Very High), reuse the FDIR layer with ofdirloss. The loss scores the cumulative distribution (expected Ranked Probability Score under the Flexible Dirichlet), so predicting a level far in rank from the truth costs more than a near one. Because FDIR is a mixture of Dirichlets it also represents bimodal ordinal conditionals — mass piled at both extremes with a dip in the middle — which structurally unimodal ordinal models (Beta-Binomial, CORAL) cannot.

model = Chain(Dense(nfeat => 64, relu), Dense(64 => 64, relu), FDIR(64 => 5))
opt_state = Flux.setup(AdamW(1e-2), model)

# Optional: per-class weights (length K) to counter class imbalance.
# Use inverse class frequency; each sample is scaled by its true-class weight.
weights = Float32.(1 ./ (classfreq .+ 1f-3))

for epoch in 1:1000
    loss, grads = Flux.withgradient(model) do m
        α, p, τ = splitfdir(m(x))
        sum(ofdirloss(y_onehot, α, p, τ; weights)) / size(x, 2)
    end
    Flux.update!(opt_state, model, grads[1])
end

r = predictive(model, x_test)
r.ŷ          # per-level probability vector (K × B)
r.aleatoric  # data ambiguity (high where the conditional is bimodal)
r.epistemic  # model uncertainty

# Point ordinal level = expectation over levels (report the full r.ŷ for
# bimodal conditionals — the expected level can land on a rare middle class).
levels = Float32.(0:(size(r.ŷ, 1) - 1))
ŷ_level = sum(levels .* r.ŷ, dims = 1)

Mean-variance estimation (MVE)

model = Chain(Dense(1 => 100, relu), Dense(100 => 100, relu), MVE(100 => 1))
opt_state = Flux.setup(AdamW(1e-3), model)

for epoch in 1:3000
    loss, grads = Flux.withgradient(model) do m
        μ, σ = splitmve(m(x'))
        mean(mveloss(y, μ, σ))
    end
    Flux.update!(opt_state, model, grads[1])
end

See the docs guide for examples of all layer types and advice on choosing the right one.

Classification

Deep evidential modeling works for classification as well. The plot below shows epistemic uncertainty when separating three Gaussians in 2D. See classification.jl.

Regression

For regression, the NormalInverseGamma distribution models a type II likelihood that explicitly captures aleatoric and epistemic uncertainty. See regression.jl.

Examples

The examples/ folder contains complete working examples:

regression.jl -- NIG with nigloss
regression2.jl -- NIG with nigloss_scaled (corrected DER)
regression3.jl -- NIG with LayerNorm and evidence tracking
regression4.jl -- MVE with parameter freezing/thawing
classification.jl -- DIR for multi-class classification
ordinal.jl -- FDIR with ofdirloss for ordinal targets, incl. bimodal conditionals

GPU support

All layers and losses work on both CPU and GPU via standard Flux conventions. Move a model and data to the GPU with gpu:

using CUDA

model = Chain(Dense(1 => 100, relu), NIG(100 => 1)) |> gpu
x_gpu = cu(x)
γ, ν, α, β = predict(model, x_gpu)  # returns CuArrays

Gradient computation, predict, split_params, and all loss functions are GPU-compatible. The test suite includes GPU-specific tests that run automatically when CUDA.functional() is true.

Adding a new distributional output

All output layers subtype AbstractEvidentialLayer. To add a new distribution (e.g. LogNormal), implement four things:

1. Layer struct and forward pass (src/dense.jl):

struct MyLayer{F, M <: AbstractMatrix, B} <: AbstractEvidentialLayer
    W::M; b::B; σ::F
end

function (a::MyLayer)(x::AbstractVecOrMat)
    o = a.W * x .+ a.b
    chunk1, chunk2 = _split_equal(o, 2)       # reuse generic splitter
    return vcat(chunk1, a.σ.(chunk2))
end

2. Parameter decomposition (src/utils.jl):

split_params(::Type{<:MyLayer}, y) = let (a, b) = _split_equal(y, 2)
    (a = a, b = b)
end

predict then works automatically -- no additional dispatch needed.

3. Loss function(s) (src/losses.jl):

Define loss functions that operate on the split parameters.

4. Uncertainty/evidence methods (src/utils.jl, optional):

Add uncertainty and/or evidence methods via multiple dispatch if your distribution supports them.

References

Amini, A., Schwarting, W., Soleimany, A. & Rus, D. Deep Evidential Regression. NeurIPS (2020).
Meinert, N., Gawlikowski, J. & Lavin, A. The Unreasonable Effectiveness of Deep Evidential Regression. arXiv (2022).
Ye, K., Chen, T., Wei, H. & Zhan, L. Uncertainty Regularized Evidential Regression. AAAI 38 (2024).
Sensoy, M., Kaplan, L. & Kandemir, M. Evidential Deep Learning to Quantify Classification Uncertainty. NeurIPS (2018).
Pandey, D. S., Choi, H. & Yu, Q. Generalized Regularized Evidential Deep Learning Models. arXiv (2025).
Yoon, T. & Kim, H. Uncertainty Estimation by Flexible Evidential Deep Learning. arXiv (2025).

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.

Name		Name	Last commit message	Last commit date
.github/workflows		.github/workflows
docs		docs
examples		examples
images		images
src		src
test		test
.gitignore		.gitignore
.versionrc		.versionrc
AGENTS.md		AGENTS.md
CHANGELOG.md		CHANGELOG.md
LICENSE		LICENSE
Makefile		Makefile
Project.toml		Project.toml
README.md		README.md
bump.sh		bump.sh

Navigation Menu

Repository files navigation

EvidentialFlux

Installing

Features

Loss functions

Utilities

Inference with `predictive`

Training with `predict`

Uncertainty

Quick start

Evidential regression (NIG)

Evidential classification (DIR)

Evidential ordinal regression (FDIR)

Mean-variance estimation (MVE)

Classification

Regression

Examples

GPU support

Adding a new distributional output

References

Summary

About

Releases

Packages

Used by

Contributors

Languages

Latest commit

History

Folders and files

Repository files navigation

EvidentialFlux

Installing

Features

Loss functions

Utilities

Inference with predictive

Training with predict

Uncertainty

Quick start

Evidential regression (NIG)

Evidential classification (DIR)

Evidential ordinal regression (FDIR)

Mean-variance estimation (MVE)

Classification

Regression

Examples

GPU support

Adding a new distributional output

References

Summary

About

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Used by

Contributors

Languages

Inference with `predictive`

Training with `predict`