mixle.inference.uncertainty module

Epistemic / aleatoric uncertainty decomposition for any predictive.

Current LLMs emit a point estimate, so their “confidence” cannot be split into what is irreducibly noisy (aleatoric) versus what more data would resolve (epistemic). A model that carries a posterior can. This module makes that split a first-class, model-agnostic operation – generalizing mixle.stats.graphs.knowledge_graph.KnowledgeGraphEnsemble.epistemic_tail_uncertainty() (which did it only for knowledge-graph tails) to any predictive.

Two exact decompositions, matched to the two answer types:

  • discrete outcomes -> the entropy (BALD mutual-information) split:

    total     = H( mean_m p_m )          predictive entropy (all uncertainty)
aleatoric = mean_m H( p_m )          expected member entropy (genuine ambiguity)
epistemic = total - aleatoric >= 0   mutual information (disagreement among members)

    The epistemic term is the Bayesian-active-learning-by-disagreement (BALD) score: it is zero when every posterior draw agrees, and large where they disagree – i.e. where more data helps.

  • continuous outcomes -> the law-of-total-variance split:

    aleatoric = mean_m Var_m             expected member variance (irreducible noise)
epistemic = Var_m( mean_m )          variance of member means (model uncertainty)
total     = aleatoric + epistemic    total predictive variance

“Members” are draws from q(theta | data) (parameter uncertainty, via ParameterPosterior) or an explicit ensemble of fitted models (as in a deep ensemble / bagged fit). Both splits are exact given the members; the only approximation is the finite number of members used to represent the posterior.

class UncertaintyDecomposition(total, aleatoric, epistemic, kind)[source]

Bases: object

A predictive uncertainty split into aleatoric + epistemic (summing to total).

kind is "entropy" (values in nats) or "variance" (values in the outcome’s squared units). Each field is a scalar for a single query point, or an array over query points.

Parameters:
total: ndarray
aleatoric: ndarray
epistemic: ndarray
kind: str
property fraction_epistemic: ndarray

Share of the total uncertainty that is epistemic (reducible by more data), in [0, 1].

item()[source]

Collapse size-1 arrays to Python floats (convenience for single-point decompositions).

Return type:

UncertaintyDecomposition

decompose_entropy(member_probs)[source]

BALD entropy split of a discrete predictive.

Parameters:

member_probs (Any) – array (M, ..., K) – M posterior draws / ensemble members, each a categorical predictive over K outcomes (optionally batched over query points in the middle axes). Rows need not be normalized; each is renormalized over the last axis.

Returns:

An UncertaintyDecomposition with kind="entropy" (nats). epistemic is the mutual information H(mean) - mean H and is clamped to >= 0 (it is non-negative in exact arithmetic; the clamp only removes tiny floating-point negatives).

Return type:

UncertaintyDecomposition

decompose_variance(member_means, member_vars=None)[source]

Law-of-total-variance split of a continuous predictive.

Parameters:

  • member_means (Any) – array (M, ...) – each member’s predictive mean E[y | theta_m].

  • member_vars (Any) – array (M, ...) – each member’s predictive variance Var[y | theta_m]. If None, aleatoric noise is taken as zero (members are point predictors) and the decomposition reports only the epistemic spread of the means.

Returns:

aleatoric = mean_m Var_m, epistemic = Var_m(mean_m), total their sum.

Return type:

An UncertaintyDecomposition with kind="variance"

predictive_distribution(members, support)[source]

Evaluate an iterable of fitted distributions over a discrete support -> (M, K) probs.

Each member’s log_density is evaluated at every point of support and softmax-normalized over the support, giving one categorical row per member. Feed the result to decompose_entropy().

Parameters:
Return type:

ndarray

posterior_ensemble(param_post, build, n=200, rng=None)[source]

Materialize n models from a parameter posterior – an ensemble representing q(theta|data).

build maps one parameter draw (whatever ParameterPosterior.sample() returns) to a fitted distribution, mirroring from_parameter_posterior(). The returned list is the “members” the decomposition integrates over – so epistemic uncertainty here is genuine parameter uncertainty, not just ensemble disagreement.

Parameters:
Return type:

list[Any]

decompose_uncertainty(*, probs=None, means=None, variances=None)[source]

Front door: decompose a predictive into aleatoric + epistemic uncertainty.

Pass exactly one representation of the per-member predictive:

  • probs=(M, ..., K) – categorical predictives -> BALD entropy split (decompose_entropy());

  • means=(M, ...) (with optional variances=(M, ...)) – continuous predictives -> law-of-total-variance split (decompose_variance()).

To decompose a fitted model’s parameter uncertainty, build the members first with posterior_ensemble() (then predictive_distribution() for the discrete case).

Parameters:
Return type:

UncertaintyDecomposition

class Clustering(representatives, probs, labels)[source]

Bases: object

Samples grouped into equivalence classes: representatives, class probs, per-sample labels.

Parameters:
representatives: list[Any]
probs: ndarray
labels: ndarray
cluster_samples(samples, equivalent=None)[source]

Group samples into equivalence classes under equivalent (default exact ==).

For discrete draws whose surface form varies but meaning does not – e.g. LLM generations (“Paris”, “It’s Paris.”, “The capital is Paris”) – pass a semantic equivalent (embedding similarity, an entailment check, normalized match). Greedy single-linkage against each cluster’s first member; returns the class distribution and per-sample assignments.

Parameters:
Return type:

Clustering

marginalize_meaning(items, equivalent=None, *, log_probs=None, weights=None)[source]

The distribution over meanings = the string distribution marginalized over each meaning class.

A generative model puts probability on strings; a meaning is an equivalence class of strings (equivalent decides sameness). The probability of a meaning c is the pushforward under the quotient – you sum the string probabilities over the class: P(c) = sum_{s in c} P(s). This returns that marginal (as a Clustering, probs = P(c)).

How the per-string probability enters:

  • log_probs – the model’s sequence log-probabilities log P(s) for each item; classes are combined by logsumexp (exact marginalization, numerically stable). Use this when the items are distinct strings whose probabilities you know – it corrects the counting form’s hidden “every string in a class is equiprobable” assumption.

  • weights – explicit non-negative masses per item (summed within class).

  • neither – uniform mass, i.e. counting: P(c) = count_c / N. For i.i.d. samples from the model this is the unbiased Monte-Carlo estimate of the same marginal.

Parameters:
Return type:

Clustering

semantic_entropy(samples, equivalent=None, *, log_probs=None, weights=None)[source]

Entropy (nats) over the meaning classes of samples – the model’s predictive uncertainty.

Sample a stochastic generator (an LLM at temperature) n times, marginalize the string distribution over meaning classes (marginalize_meaning()), and take the entropy of that marginal. High semantic entropy means the model disagrees with itself about what the answer is (a hallucination signal), as opposed to merely phrasing one answer many ways (which collapses to low entropy). Pass log_probs (the sequence log-likelihoods) to marginalize with the actual string probabilities rather than by sample counting. Feed the clusters’ per-member distributions to decompose_entropy() for an epistemic/aleatoric split.

Parameters:
Return type:

float