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). The DECOMPOSITION IDENTITIES are exact given the members; what the numbers mean depends on where the members came from (audit U-1): only for exact posterior draws (the conjugate route) is the finite member count M the sole approximation. A VARIATIONAL q is mode-seeking and systematically understates the posterior spread (the epistemic term inherits that understatement at any M); MCMC members carry burn-in and autocorrelation, so the effective sample size, not M, governs their error. The finite-M estimators are themselves biased: the entropy split’s plug-in mutual information is capped at log M outright and biased low (audit U-2 – at M = 2 a true 2.08-nat disagreement reads at most 0.693), and the variance split’s epistemic term uses the unbiased ddof=1 estimator precisely so its expectation is the population variance at every M (audit U-3 – the ddof=0 plug-in understated it by exactly (M-1)/M, 50% at M = 2).

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. Array fields are owned read-only copies, so the total = aleatoric + epistemic identity the record asserts cannot be broken after the fact (MXR-080-1899).

Parameters:
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 PLUG-IN mutual information H(mean) - mean H, clamped to >= 0 (non-negative in exact arithmetic; the clamp removes floating-point negatives, and total is recomputed as aleatoric + epistemic afterward so the asserted identity holds exactly – audit U-8). ESTIMATOR LIMITS (audit U-2): the plug-in MI over M members is structurally capped at log M whatever the true disagreement – at the permitted minimum M = 2 a true 2.08-nat mutual information reads at most 0.693 – and is biased low below the cap; read it as a lower bound at small M, and prefer M large enough that log M comfortably exceeds the disagreement you need to resolve.

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:

An UncertaintyDecomposition with kind="variance" (or kind="variance-epistemic-only" when member_vars is None – total then contains NO aleatoric noise and is not a predictive variance; audit U-9).

Return type:

UncertaintyDecomposition

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().

PRECONDITION (audit U-7): the softmax treats every support point as carrying EQUAL measure. That is exact for a discrete family evaluated on its full support, and a grid approximation for a continuous density ONLY when the grid is uniform – a non-uniform grid silently reweights the density by the missing cell widths, and any downstream entropy inherits the distortion. Restrict a discrete member’s support at your own risk: the row then represents the conditional given that subset.

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.

The three fields are one interlocking statement – probs[c] is the mass of the class whose member representatives[c] stands for, and labels[i] says which class sample i landed in – so the record owns them: representatives is a tuple and the arrays are write-locked copies (MXR-080-1899). Consumers in this repository only read them (index, iterate, zip, argmax, tolist), which is why the sequence can be a tuple without breaking a caller.

Parameters:
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). True single linkage is computed as the connected components of all pairwise equivalence edges.

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 DISTINCT item; classes are combined by logsumexp and renormalized over the items you passed, so the result is the meaning marginal CONDITIONAL ON the observed strings, not the model’s full marginal (audit U-5) – unobserved tail mass is excluded, and repeated strings would each add their own P(s) (estimating sum P(s)^2, a different quantity), so duplicates are refused in this branch. 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]

PLUG-IN entropy (nats) over the meaning classes of samples.

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).

APPROXIMATION, NOT A MEASUREMENT OF THE GENERATOR (STAT-RR17-14): with n counted samples this is the plug-in estimator, biased LOW by about (K - 1) / (2n) nats (Miller-Madow, K = observed meaning classes) – at the eight-draw default an abstention loop actually uses, a uniform ten-class generator’s true 2.303 nats measured 1.646 on average (bias -0.656): the bias direction makes a HALLUCINATING model look confident, the exact failure the signal exists to catch. Use semantic_entropy_receipt() for the estimate WITH its sample size, observed class count, bias estimate, and the bias-corrected value; abstention thresholds should use the corrected value (it errs toward abstaining). 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