Conceptio › Archive › arXiv CS

arXiv CSopen access

A Mutual Information Lower Bound for Multimodal Regression Active Learning

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

A Mutual Information Lower Bound for Multimodal Regression Active Learning

arXiv:2605.14917v1 [cs.LG] 14 May 2026

Leonardo Ferreira Guilhoto∗ Graduate Group on Applied Mathematics & Computational Science Philadelphia, PA 19104 [email protected] Akshat Kaushal∗ Department of Computer and Information Science University of Pennsylvania Philadelphia, PA 19104 [email protected] Paris Perdikaris Department of Mechanical Engineering and Applied Mechanics University of Pennsylvania Philadelphia, PA 19104 [email protected]

Abstract Active learning for continuous regression has lacked an acquisition function that targets epistemic uncertainty when the predictive distribution is multimodal: variance misses modal disagreement, and information-theoretic targets like BALD are designed for discrete outputs. We introduce a Two-Index framework that makes this separation explicit: one stochastic index selects among competing model hypotheses (epistemic source), while a second governs within-hypothesis randomness (aleatoric source). An entropy decomposition within the framework identifies the mutual information between the output and the epistemic index as a principled acquisition objective, and we prove this quantity vanishes as the model is trained on growing datasets, confirming that it captures exactly the uncertainty data can resolve. Because this mutual information is intractable for continuous outputs, we derive the Mutual Information Lower Bound (MI-LB) acquisition function, a closed-form approximation for Mixture Density Network ensembles. On benchmarks featuring multimodal systems, MI-LB matches or beats every baseline evaluated and is the only method to do so consistently – geometric and Fisher-based baselines compete only when the input space already encodes the multimodality, and collapse otherwise.

Introduction

Active learning in continuous, multivariate settings requires acquisition functions that can distinguish uncertainty the model can reduce by collecting more data (epistemic) from uncertainty intrinsic to the data-generating process (aleatoric). Existing approaches either operate in restricted settings, such as discrete classification [1] or unimodal Gaussian outputs [2], or rely on scalar summaries like ∗ Equal contribution.

Preprint.

predictive variance that conflate the two sources. The problem is worse in multimodal settings: two model hypotheses may assign the same mean and variance to an output while disagreeing on the number and location of modes. Variance alone cannot detect this kind of distributional disagreement. We address this gap with a framework that keeps the two sources of uncertainty separate by construction. The framework extends the Epistemic Neural Network formalism [3] by introducing two independent stochastic indices: an epistemic index Z ∼ PZ that parameterizes a family of hypotheses, and an aleatoric index ϵ ∼ Pϵ that governs within-hypothesis stochasticity. A learned map gθ (x; z, ϵ) transforms these indices and the input into the output space. This approach provides a common vocabulary for a broad class of uncertainty quantification methods, facilitating the exchange of analyses among model families. The two-index structure yields an entropy decomposition that separates predictive uncertainty into an epistemic term, the mutual information I(Y ; Z | x), and an aleatoric term, the expected conditional entropy EZ [H(Y | x, Z)]. Despite being a natural acquisition target, the mutual information requires computing differential entropies that are intractable for general continuous outputs. When the model family is an ensemble of Mixture Density Networks (MDNs), the predictive and conditional distributions are both Gaussian mixtures, and known upper and lower bounds on Gaussian-mixture entropy [4] can be combined to produce a tractable lower bound on I(Y ; Z | x). We call this the Mutual Information Lower Bound (MI-LB) acquisition function. Across three benchmarks with controllable multimodality, MI-LB matches or beats every baseline evaluated and is the only method to do so consistently. Geometric and Fisher-based baselines compete only when the input geometry already encodes the multimodal structure, and collapse otherwise. To the best of our knowledge, MI-LB is the first acquisition function specifically designed for multimodal continuous settings. The main contributions of this paper are: • The MI-LB acquisition function for active learning in continuous multivariate settings, with a proof that it lower bounds the true epistemic mutual information I(Y ; Z | x) and a closed-form expression that is efficient to evaluate under standard MDN assumptions. • The Two-Index framework for disentangling epistemic and aleatoric uncertainty, applicable to a broad class of model families, with a proof that the epistemic term I(Y ; Z | x) vanishes with sufficient data under standard well-specification and consistency conditions (Appendix D). • Experiments on three multimodal benchmarks showing that MI-LB matches or beats every baseline evaluated (Random, Variance, BAIT, Core-Set), and is the only method to do so consistently – geometric and Fisher-based baselines compete only when the input geometry already encodes the multimodal structure, and collapse when the modal disagreement lives in output space. 1.1

Related Work

Decomposing predictive uncertainty into aleatoric and epistemic components has been a longstanding goal in the machine learning literature [2, 5]. Foundational approaches include Bayesian Neural Networks [6], Deep Ensembles [7], and the Epistemic Neural Network framework [3], which provides a unified interface that does not require explicit Bayesian inference [5, 3]. Despite these advances, recent benchmarking shows that practical disentanglement remains elusive: modern evidential, variational, and deterministic methods exhibit rank correlations often exceeding 0.94 between their aleatoric and epistemic estimates [8]. This problem is particularly pronounced in continuous, multimodal regression tasks, where fitting a single Gaussian to diverse target modes inflates predictive variance, motivating the use of Mixture Density Networks (MDNs) [9, 10]. In the context of active learning, selecting informative data points to minimize labeling costs has traditionally relied on geometric or gradient-based heuristics [11]. Geometric methods such as Core-Set minimize L2 distances in embedding space [12], while Fisher-based methods like BAIT optimize bounds on the maximum likelihood estimator error using network gradients [13]. Both have important failure modes in multimodal settings: Core-Set ensures coverage in the learned embedding space and can fail when multimodality manifests in the output distribution rather than in the geometry of the inputs [12], and BAIT’s last-layer Fisher embedding provides poor candidate rankings when the predictive likelihood is highly multimodal, since the gradient signal at a single Monte-Carlo sample is 2

dominated by within-mode variance. In these regimes, black-box acquisition strategies based on the empirical predictive covariance of ensembles often outperform white-box gradient methods [14, 11]. Information-theoretic approaches such as Bayesian Active Learning by Disagreement (BALD) use mutual information to target epistemic uncertainty directly [1, 15]. BALD is effective in discrete classification, but extending it to continuous regression requires evaluating differential entropy for distributions like Gaussian mixtures, which lacks a closed-form solution [4].

The Two-Index Approach for Disentangling Uncertainties

2.1

The Two-Index Generative Model

We build on the Epistemic Neural Network (ENN) framework [3] and the distribution network formalism of [10] to construct a unified model that explicitly represents both epistemic and aleatoric uncertainty through two independent stochastic indices. In the ENN framework, a model is specified by a parameterized function fθ and a fixed reference distribution PZ . An epistemic index Z ∼ PZ is sampled and passed to the network alongside the input x, producing a prediction fθ (x; Z). Crucially, PZ is not updated during training. Instead, training optimizes θ so that the mapping from Z to predictions becomes increasingly invariant under PZ as more data is collected. Sensitivity of predictions to Z reflects epistemic uncertainty: the model has not yet identified a unique function consistent with the data. We extend this framework by introducing the aleatoric index ϵ ∼ Pϵ , a second fixed reference distribution independent of both Z and θ. For a given epistemic index Z = z and input x, the model defines a conditional distribution over outputs Y through the generative mapping Y = gθ (x; z, ϵ),

Z ∼ PZ ,

ϵ ∼ Pϵ ,

Z ⊥ ϵ.

(1)

Here gθ is a learned transformation that maps the input x and two independent sources of randomness into the output space. The role of ϵ is to represent irreducible stochasticity: for any fixed z, sampling ϵ ∼ Pϵ produces draws from the aleatoric conditional distribution pθ (y | x, z). The role of Z is to represent reducible uncertainty: variation in predictions across different realizations of Z reflects hypotheses about the data-generating process that have not yet been eliminated by the available data. The network gθ learns a transformation from the tractable distributions PZ and Pϵ into a flexible, potentially multimodal distribution over the output Y . Induced distributions. The joint generative model in (1) induces a hierarchy of distributions, made precise using the push-forward operator. For a measurable space (Y, B(Y)) and a measurable map h : E → Y, the push-forward of a measure µ on E under h is the measure h# µ on Y defined by (h# µ)(A) := µ {ϵ ∈ E : h(ϵ) ∈ A} , A ∈ B(Y). (2) For a fixed epistemic realization z ∈ Z and input x ∈ X , the aleatoric conditional distribution is defined as the push-forward of Pϵ under the map gθ (x; z, ·) : E → Y: pθ (· | x, z) := gθ (x; z, ·) # Pϵ . (3) Marginalizing equation (3) over the epistemic index yields the predictive distribution Z pθ (· | x) := pθ (· | x, z) dPZ (z).

(4)

The predictive distribution conflates both sources of uncertainty. Disentangling them requires tracking the two indices separately, as formalized in the decompositions of Section 2.2. Interpretation. Each realization Z = z instantiates a hypothesis pθ (· | x, z) for the true conditional p∗ (· | x); within a fixed hypothesis, ϵ governs irreducible stochasticity. Disagreement across realizations of Z is the origin of epistemic uncertainty. The predictive distribution has the same mathematical form as Bayesian Model Averaging [16], but the framework does not require a strict Bayesian formulation of PZ : as illustrated in Appendix B, it accommodates ensembles, BNNs, conditional flow matching, and VAEs under a common vocabulary. 3

2.2

Quantifying the Two Sources of Uncertainty

The two-index structure enables exact decompositions of predictive uncertainty into its epistemic and aleatoric components. We provide two such decompositions via variance and entropy. Variance-Based Decomposition. In the scalar output (Y ∈ R) case [5], the law of total variance yields an exact additive decomposition. Writing Y = gθ (x; Z, ϵ) with Z ⊥ ϵ, VarZ,ϵ (Y |x) = EZ Varϵ (Y | Z, x) + VarZ Eϵ [Y | Z, x] . (5) | {z } | {z } | {z } total variance

aleatoric variance

epistemic variance

The aleatoric term EZ [Varϵ (Y | Z, x)] averages the output variance within each hypothesis over epistemic realizations; it measures the average irreducible spread. The epistemic term VarZ (Eϵ [Y | Z, x]) measures the variance of the conditional mean across hypotheses; it captures disagreement between different epistemic realizations of the model. This decomposition is exact and interpretable, but it summarizes uncertainty through second moments alone. As noted in the introduction to this section, two hypotheses can share the same conditional mean and variance while assigning probability mass to entirely different modes. The variance-based epistemic term would report an incomplete summary in such a case. For multimodal or non-Gaussian distributions, a richer summary is needed, as further discussed in Appendix C. Entropy-Based Decomposition. We complement (equation 5) with an entropy-based decomposition that captures distributional disagreement beyond second moments. Define the total predictive uncertainty as the differential entropy of pθ (y | x): HZ,ϵ (Y | x) = EZ,ϵ − log pθ (Y | x) . (6) Applying the chain rule of mutual information, this decomposes as HZ,ϵ (Y | x) = EZ Hϵ (Y | Z, x) + I(Y ; Z | x) , (7) | {z } | {z } aleatoric uncertainty

epistemic uncertainty

where I(Y ; Z | x) is the mutual information between the output Y and the epistemic index Z, conditioned on the input x. The aleatoric term EZ [Hϵ (Y | Z, x)] is the expected entropy of the conditional distribution pθ (· | x, z), averaged over the epistemic index. The epistemic term I(Y ; Z | x) quantifies how much knowing Z reduces uncertainty about Y : it is large when different epistemic indices yield meaningfully different conditional distributions, and vanishes when all hypotheses agree. Because mutual information is sensitive to the full shape of each hypothesis distribution, it detects the kind of modal disagreement that variance misses. 2.3

Asymptotic Guarantees

The decompositions (5) and (7) have the correct asymptotic behavior by construction. As the size of the training dataset grows, θ is optimized to produce predictions that are consistent with the data under all probable realizations of Z. In the limit of infinite data and a well-specified model, the learned θ renders gθ (x; z, ϵ) approximately independent of z for PZ -almost all z, so that: n→∞ n→∞ I(Y ; Z | x, Dn ) −−−−→ 0, EZ Hϵ (Y | Z, x, Dn ) −−−−→ H(p∗ (· | x)), (8) where p∗ (· | x) is the true conditional distribution of the data-generating process. Epistemic uncertainty vanishes as the model identifies the correct hypothesis, while aleatoric uncertainty converges to the irreducible entropy of the true process. Conversely, at an out-of-distribution input xOOD , training data provides no constraints on Z, so I(Y ; Z | xOOD ) can remain large, providing a principled basis for out-of-distribution detection. We note that the convergence of the aleatoric term to H(p∗ (· | x)) requires that the aleatoric model class — the family of distributions representable by gθ (x; z, ·) for fixed z — be sufficiently expressive to capture p∗ . This is a non-trivial condition that motivates the use of flexible density estimators, such as Mixture Density Networks or implicit generative models, rather than restricting to unimodal Gaussian outputs. This is stated informally as the following theorem, which is made precise and proved in Appendix D: Theorem 2.1 (Informal: Epistemic Uncertainty Vanishes with Data). If a model can represent the true data-generating process p∗ (· | x) and its parameters converge to the correct values as the dataset grows to infinity, then I(Y ; Z | x) → 0. In other words, a well-trained model’s predictions eventually agree across all epistemic index values, and all remaining uncertainty is aleatoric. 4

2.4

Implications for Data Acquisition

The decomposition from (7) and Thm. 2.1 together identify I(Y ; Z | x) as the natural acquisition target for active learning. It isolates exactly the component of predictive uncertainty that additional data can reduce: epistemic disagreement among hypotheses. The variance-based epistemic term from (5) targets the same goal in principle, but summarizes disagreement through second moments alone; it can miss regions where hypotheses differ in modal structure rather than spread. The mutual information I(Y ; Z | x), by contrast, is sensitive to the full distribution of each hypothesis, making it the right objective when the conditional distribution p∗ (y | x) may be multimodal. The remaining challenge is computational. Evaluating I(Y ; Z | x) requires computing differential entropies of the predictive and conditional distributions, which have no closed-form expression for general continuous outputs. In the next section, we show that when the model family is an ensemble of Gaussian mixtures, analytical entropy bounds can be used to construct a tractable lower bound on I(Y ; Z | x) that serves as an effective acquisition function.

The Mutual Information Lower Bound (MI-LB) Acquisition Function

3.1

Entropy Intractability in Continuous Settings

Re-arranging the entropy-based decomposition from (7), the mutual information can be written as (9) I(Y ; Z | x) = H(Y | x) − EZ H(Y | x, Z) . | {z } | {z } (i) marginal entropy

(ii) expected aleatoric entropy

Computing this quantity requires evaluating two differential entropy terms: term (i), the entropy of the predictive distribution pθ (· | x) from equation 4, and term (ii), the entropy of Y | x, Z = z for a fixed z, averaged over PZ . For implicit distribution networks [10] — models that produce samples of Y via the map gθ (x; z, ϵ) but do not provide an analytical form for pθ (· | x, z) — both terms must be estimated from samples using kernel density estimators, which may scale poorly with the dimension of Y. For explicit distribution networks, the density pθ (· | x, z) is available analytically, enabling structured approximation of both terms. We exploit this structure in the case of Mixture Density Networks, where each pθ (· | x, z) is a mixture of Gaussians. 3.2

MDN Structure and the Marginal Mixture

Assume Z takes values in the discrete index set {1, . . . , nens } with PZ (Z = z) = wz , and suppose that for each z, the aleatoric conditional pθ (· | x, z) is a Gaussian mixture with K components: pθ (y | x, z) =

K X

(z) (z) (z) αi (x) N y; µi (x), Ci (x) ,

(10)

i=1

P (z) (z) (z) (z) where αi (x) > 0, i αi (x) = 1, µi (x) ∈ RN , and Ci (x) ∈ RN ×N is a symmetric positive definite covariance matrix. Here N denotes the dimension of the output space Y ⊆ RN . Under this formulation, the predictive marginal distribution (4) is pθ (y | x) =

n ens X z=1

wz pθ (y | x, z) =

n K ens X X

(z) (z) (z) wz αi (x) N y; µi (x), Ci (x) , | {z } z=1 i=1

(11)

=: βz,i (x)

(z)

which is a Gaussian mixture with nens · K components and weights βz,i (x) = wz αi (x). Both terms in (9) therefore require computing the differential entropy of a Gaussian mixture, for which no closed-form expression exists. 3.3

Entropy Bounds for Gaussian Mixtures

PK We make use of two bounds from [4]. Let p(y) = i=1 πi N (y; µi , Ci ) be a Gaussian mixture with K components, weights πi , means µi ∈ RN , and covariances Ci ∈ RN ×N . [4] shows that the 5

differential entropy H(p) satisfies Hlower ≤ H(p) ≤ Hupper , where   K K X X Hlower = − πi log πj N (µi ; µj , Ci + Cj ) , i=1

Hupper =

(12)

j=1

K X

1 πi − log πi + log (2πe)N |Ci | . 2 i=1

(13)

Both bounds are computable in O(K 2 ) operations. The upper bound (13) decomposes as the weighted sum of the entropy of individual modes and is tight when the components have negligible overlap. The lower bound (12) accounts for inter-component similarity through pairwise Gaussian evaluations at the component means. Under the standard MDN assumption of diagonal covariances 2 2 Ci = diag(σi,1 , . . . , σi,N ), the Gaussian evaluations in (12) and the determinants in (13) reduce to sums of univariate quantities, making both bounds particularly efficient to evaluate. 3.4

The MI-LB Acquisition Function

We define the Mutual Information Lower Bound (MI-LB) acquisition function as MI-LB(x) := Hlower (Y | x) −

n ens X

wz Hupper (Y | x, Z = z),

(14)

z=1

where Hlower (Y | x) is the lower bound (12) applied to the nens ·K-component marginal mixture (11) (z) (z) with component weights βz,i (x), means µi (x), and covariances Ci (x); and Hupper (Y | x, Z = z) is the upper bound (13) applied to the K-component conditional mixture (10) for each z. Explicitly,   n n K K ens X ens X X X (z) (l) (z) (l) MI-LB(x) = − βz,i log βl,j N µi ; µj , Ci + Cj  z=1 i=1

−

n ens X z=1

l=1 j=1

K X

(z)

αi

(z)

− log αi

i=1

1 (z) log (2πe)N |Ci | , 2

(15)

where we suppress the dependence on x for readability. The fundamental property of MI-LB is that it constitutes a certified lower bound on the true epistemic uncertainty I(Y ; Z | x), which we state formally as follows, with a proof given in Appendix E. Theorem 3.1 (MI-LB is a lower bound on mutual information). Under the model assumptions of Section 2, with Z supported on {1, . . . , nens } and pθ (· | x, z) a Gaussian mixture of the form (10) for each z, it holds that MI-LB(x) ≤ I(Y ; Z | x)

for all x ∈ X .

(16)

The practical consequence of Theorem 3.1 is that maximizing MI-LB over a candidate pool is a conservative acquisition strategy: any point selected by MI-LB is guaranteed to have true epistemic uncertainty at least as large as the reported score. This mirrors the logic of the Evidence Lower BOund (ELBO) in variational inference [17], where optimizing a lower bound on the log-evidence provides a principled surrogate without risking false certification of a poor solution. Consequently, the MI-LB acquisition rule selects x∗ = arg max MI-LB(x).

(17)

x ∈ Xpool

Connection to BALD. The Bayesian Active Learning by Disagreement (BALD) acquisition [1] is defined identically to (9) in the discrete classification setting, where Y follows a categorical distribution and entropy is Shannon entropy, admitting exact computation. MI-LB can therefore be understood as a continuous, multivariate alternative to BALD for regression, with the [4] entropy bounds replacing exact entropy to restore tractability in the absence of a finite output alphabet. 6

Experiments

We benchmark MI-LB, the acquisition function proposed in Section 3, against four baselines on three multimodal problems: Random, Epistemic Variance, BAIT [13] (last-layer Fisher trace), and Core-Set [12] (k-Center-Greedy on the MDN backbone). SBAL [14] and MaxDist [11] batch variants of Variance and MI-LB are deferred to the appendix. Among the three benchmarks, §4.2 (the coupled double-well) is the discriminating one: input geometry (σ, κ) does not encode the modal structure, which lives in output space via Kramers escape. §4.1 and §4.3 have phase boundaries in the input space and serve to characterise where geometric baselines remain competitive. Table 1: Per-benchmark settings. All used a pool of 50,000 inputs, 100 initial labels, and an ensemble of size 8. Benchmark Multimodal conditional(§4.1) Coupled double-well (§4.2) Ternary phases (§4.3)

(a) Multimodal conditional (§4.1)

dim x

dim y

KMDN

rounds × batch

total labels

NLL∗

10 7 8

16 20 1

5 8 4

20 × 50 20 × 50 30 × 15

1,100 1,100 550

22.98 — 1.33

(b) Coupled double-well (§4.2)

Figure 1: Test NLL vs. training-set size on each benchmark for MI-LB against Random, Epistemic Variance, BAIT, and Core-Set in top-k (5 seeds; bands min–max); dashed = oracle NLL∗ where available. MI-LB matches or beats every baseline on every benchmark; geometric methods (Core-Set, BAIT) are competitive only when input geometry encodes the multimodality (a, c) and collapse when it lives in output space (b).

4.1

Synthetic Multimodal Conditional Problem

We construct a synthetic conditional distribution p∗ (y | x) as a mixture of K=3 Gaussians: p∗ (y | x) =

K X

πk (x) N y; µk (x), Σk (x) ,

(18)

k=1

where the mixing weights πk , means µk , and diagonal covariances Σk are input-dependent functions (full specification in Appendix F). Inputs x ∈ R10 lie on a 4-dimensional manifold embedded via x = tanh(Al + b), l ∼ N (0, I4 ), respecting the manifold hypothesis [18], and outputs y ∈ R16 . Because p∗ is known in closed form, the oracle NLL∗ is computable exactly as a calibration reference. A K=5 MDN ensemble trained on the full pool reaches NLL 24.72 vs. oracle 22.98, while a K=1 MDN attains only 47.76 – the conditional is multimodal and a mixture head is necessary (Fig. 2). The mixing weights πk (x) are gated by the radial coordinate r = ∥x1:L ∥: for r ≲ r0 a single component dominates and p∗ (y | x) is effectively unimodal, while for r ≳ r0 multiple components carry comparable mass and p∗ becomes genuinely multimodal, producing a sharp unimodal–multimodal transition at r ≈ r0 (full form in App. F.2). Why this is a useful benchmark. The coexistence of unimodal and multimodal regimes makes this problem a natural testbed for the Two-Index framework: I(Y ; Z | x) should be small in the interior, where all ensemble members agree on a single-component conditional, and large near the transition boundary ∥x1:L ∥ ≈ r0 , where members must resolve the weights and locations of multiple 7

Figure 2: Predicted vs. true samples in the (y0 , y1 ) plane on held-out inputs. Left: draws from the oracle p∗ (y | x), displaying the multimodal structure of the target conditional. Middle: MDN ensemble samples; the recovered geometry closely matches the oracle, with calibration gap ∆ = 1.74. Right: single-Gaussian MDN samples collapses to an isotropic blob that cannot represent disjoint modes, yielding ∆ = 24.78.

components. A well-calibrated acquisition function should therefore concentrate queries on this boundary. As a lower bound on I(Y ; Z | x), MI-LB is designed to detect this distributional disagreement, while variance-based scores summarize ensemble disagreement through second moments alone and underweight regions where hypotheses differ in modal structure rather than spread. Results. Figure 1a compares MI-LB against the four baselines. From the very first acquisition rounds MI-LB pulls clearly ahead of Random, Variance, and BAIT and tracks the strongest baseline (Core-Set) throughout the budget, ending within seed bands at 31.1 ± 0.3 vs. 30.6 ± 0.6 at n = 1100. MI-LB matches the best geometric baseline on this benchmark with the tightest seed-to-seed std among all methods, tying Core-Set under top-k and beating it under MaxDist batch selection (Appendix F.2, Table 5). Core-Set is competitive here for a benchmark-specific reason: the input manifold x = tanh(Al + b) is 4-dimensional and the radial coordinate ∥x1:L ∥ already encodes the unimodal–multimodal gate, so k-Center-Greedy in feature space implicitly samples the boundary that MI-LB targets explicitly. Sections 4.2–4.3 test this dependence and show that MI-LB retains its advantage when geometry no longer encodes multimodality, whereas Core-Set degrades sharply. Spatially, MI-LB concentrates on the unimodal–multimodal boundary, Variance is weaker on the same region, Core-Set covers the manifold uniformly, Random spreads everywhere (App. F.2, Fig. 6). 4.2

Coupled Double-Well System

We consider a chain of P = 5 particles evolving in coupled one-dimensional double-well potentials under overdamped Langevin dynamics, h i X dqi = a (qi − qi3 ) + κ (qj − qi ) dt + σ dWi , i = 1, . . . , P, (19) | {z } j∈nn(i) double-well force | {z } nearest-neighbour coupling

where dWi are independent Wiener increments and we fix a = 1. Inputs encode the initial particle configuration together with (σ, κ); outputs stack particle positions at four snapshot times. Full configuration, integration scheme, and evaluation protocol in Appendix F.3. Why this is a useful benchmark. The conditional p∗ (y | x) undergoes a sharp change in mode structure controlled by σ 2 /a: at low σ each particle stays in its starting well (p∗ unimodal per particle); at high σ, noise triggers Kramers escape [19], making the single-particle marginal bimodal. Coupling κ raises the effective barrier for collective flips, coupling single-particle marginals into a joint P distribution over {−1, +1}p basins. The Kramers exponent a/(2σ 2 ) at barrier height ∆V = a/4 places the crossover at σ ≳ a/2 ≈ 0.71. A single-particle sweep (Fig. 3) confirms this scaling, and at σ ≳ 1 noise overwhelms the barrier so that mass spreads beyond both wells into the |q| > 1 tails. With σ and κ as input coordinates, the phase boundary lives directly in input space, marking where informative acquisitions should concentrate. A Gaussian head must smear mass across both wells in the bimodal regime, inflating test NLL: a single-Gaussian (K = 1) head pays a 5.51-nat penalty against the K = 8 mixture head used on this benchmark (Appendix F.3, Fig. 8). This penalty is a 8

Figure 3: Terminal position q(T ) histograms for an uncoupled particle (P = 1, κ = 0, q(0) = −0.5) at pfour noise levels; dashed lines mark q = ±1. At σ = 0.3 the particle stays trapped near q = −1; at σ = 0.7 ≈ a/2 Kramers escape fills both wells; for σ ≥ 1 noise dominates the barrier, spreading mass into the |q| > 1 tails.

property of the benchmark, not of any acquisition strategy: any method built on a single-Gaussian head fails here by construction. Results. Figure 1b: MI-LB wins decisively at 71 ± 8 at n = 1100 - 1.7× better than Variance (122 ± 11), 4× better than Core-Set (304 ± 55) and BAIT (281 ± 15), and 7× better than Random (518 ± 60). This is the discriminating benchmark: with the bimodal regime’s mode structure living in output space (Kramers escape between wells), two ensemble members can agree on E[Y | x] while disagreeing on whether mass concentrates near q = +1, q = −1, or both – so feature-space coverage (Core-Set) and last-layer Fisher (BAIT) cannot rank candidates informatively, whereas MI-LB’s entropy decomposition over the full output mixture targets exactly the disagreement that matters. 4.3

Synthetic Phase-Competition Benchmark

Active learning is a standard tool in materials discovery, where labels are expensive and composition– process inputs live on low-dimensional manifolds [20, 21, 22]. As in the previous benchmarks, discrete phases compete across the input space, producing boundary-localised multimodality. We build a synthetic CALPHAD-style [23] benchmark that reproduces this structure. Each input is a composition on the 3-component simplex (with xC = 1 − xA − xB ) concatenated with nproc = 6 continuous process parameters. Per-phase Gibbs free energies Gϕ (xA , xB , xC ) are random quadratic forms in composition (independent of p), and the latent phase distribution is softmax(−Gϕ /τG ) at temperature τG = 0.08, producing sharp phase boundaries on the simplex slice. The scalar response is sampled from a per-phase Gaussian whose mean depends on both composition and process parameters; full simulator parameters in Appendix F.4. Why this is a useful benchmark. The composition (xA , xB , xC ) acts as a categorical latent selecting a phase, confining multimodality to a thin boundary region on the simplex; the process subspace contributes only a smooth, phase-conditional shift. This factorisation isolates a clean test of boundary localisation (realised phase masses and boundary fraction in Appendix F.4). Results. The tighter budget than Sections 4.1 and 4.2 reflects the 1-D scalar output, on which NLL saturates faster than the 16- and 20-dimensional outputs of the other benchmarks. Figure 1c: MI-LB reaches 1.99 ± 0.04 at n = 550, tied with Core-Set (1.99±0.04); Variance (2.06), BAIT (2.19), and Random (2.21) trail. MaxDist on top of MI-LB narrowly wins overall (1.95 ± 0.06). BAIT collapses to Random-level, the 1-D scalar output makes its single-MC last-layer Fisher estimate too noisy to rank candidates. Core-Set’s parity again reflects benchmark geometry (phase boundaries lie on the composition simplex) and Sections 4.2 shows it does not transfer when multimodality is in output space. MI-LB retains the tightest seed-to-seed spread in the data-scarce regime (n ∈ [115, 250]): ∼ 2.6× tighter than Random, ∼ 4.5× than Variance (Appendix Table 11).

Conclusion

Summary. We introduce the Two-Index generative framework to formally disentangle epistemic and aleatoric uncertainty across diverse model families. By leveraging this formalism, we derive the MI-LB acquisition function, providing a principled and computationally efficient Mutual Information lower bound for active learning in continuous, multimodal settings. Across three benchmarks MI-LB 9

matches or beats every baseline we evaluate (Random, Variance, BAIT, Core-Set) and is the only method that does so consistently – the geometric baselines are competitive only on benchmarks whose inputs already encode the multimodality. The discriminating test is the coupled double-well, where Kramers escape places the modal disagreement in output space: Core-Set and BAIT both collapse to within 1.7× of Random, while MI-LB wins by 4× – the empirical signature that distributional acquisition is doing something feature-space coverage cannot. Limitations. MI-LB currently relies on the availability of explicit density components, such as those in Mixture Density Networks. Its performance is also tied to the quality of the entropy bounds, which may loosen in some settings. Furthermore, the framework assumes a well-specified model class capable of recovering the true irreducible stochasticity. Our empirical evaluation is also restricted to synthetic benchmarks with controllable multimodality, since no standard benchmarks exist for active learning in continuous multimodal regression; building such benchmarks is left to future work. Future Work. Future research may focus on extending the MI-LB objective to implicit generative models, such as diffusion and flow matching, where densities are not analytically tractable. Another promising direction is to explore the framework’s utility in safety-critical control tasks where distinguishing between lack of data and inherent noise is vital for risk-aware planning.

Acknowledgments and Disclosure of Funding We would like to acknowledge support from the US Department of Energy under the Advanced Scientific Computing Research program (grant DE-SC0024563) and the US National Science Foundation (NSF) Soft AE Research Traineeship (NRT) Program (NSF grant 2152205). We also thank the developers of software that enabled this research, including JAX [24], Flax [25] Matplotlib [26] and NumPy [27].

Author Contributions L.F.G. developed the theory for the two-index framework and MI-LB acquisition. A.K. developed the benchmarks, carried out experiments, and analyzed the data. P.P. provided funding and supervised this study. All authors helped to write and review the manuscript.

Impact Statement Our contributions enable advances in deep learning and uncertainty quantification. This has the potential to impact a wide range of downstream applications. While we do not anticipate specific negative impacts from this work, as with any powerful predictive tool, there is potential for misuse. We encourage the research community to consider the ethical implications and potential dual-use scenarios when applying these technologies in sensitive domains and to avoid its application altogether to weaponry and other military technologies.

Declaration of LLM Usage LLMs were used during the development of this paper, including for editing the writing, developing ideas, and as code assistants. We believe our use to be within the standard uses of this technology, and authors have verified the integrity of all material contained in this work.

References [1] Neil Houlsby, Ferenc Huszár, Zoubin Ghahramani, and Máté Lengyel. Bayesian active learning for classification and preference learning, 2011. [2] Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017. 10

[3] Ian Osband, Zheng Wen, Mohammad Asghari, Morteza Ibrahimi, Xiyuan Lu, and Benjamin Van Roy. Epistemic neural networks. CoRR, abs/2107.08924, 2021. [4] Marco F. Huber, Tim Bailey, Hugh Durrant-Whyte, and Uwe D. Hanebeck. On entropy approximation for gaussian mixture random vectors. In 2008 IEEE International Conference on Multisensor Fusion and Integration for Intelligent Systems, pages 181–188, 2008. [5] Matias Valdenegro-Toro and Daniel Saromo Mori. A deeper look into aleatoric and epistemic uncertainty disentanglement. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW), pages 1508–1516. IEEE, 2022. [6] Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are bayesian neural network posteriors really like? In International conference on machine learning, pages 4629–4640. PMLR, 2021. [7] Andrew Gordon Wilson and Pavel Izmailov. Deep ensembles as approximate bayesian inference. https://cims.nyu.edu/~andrewgw/deepensembles/, 2021. [8] Bálint Mucsányi, Michael Kirchhof, and Seong Joon Oh. Benchmarking uncertainty disentanglement: Specialized uncertainties for specialized tasks. Advances in neural information processing systems, 37:50972–51038, 2024. [9] Christopher M Bishop. Mixture density networks. Neural Computing Research Group Report, 1994. [10] Leonardo Ferreira Guilhoto, Akshat Kaushal, and Paris Perdikaris. Multimodal scientific learning beyond diffusions and flows, 2026. [11] David Holzmüller, Viktor Zaverkin, Johannes Kästner, and Ingo Steinwart. A framework and benchmark for deep batch active learning for regression. Journal of Machine Learning Research, 24(164):1–81, 2023. [12] Ozan Sener and Silvio Savarese. Active learning for convolutional neural networks: A core-set approach. In International Conference on Learning Representations (ICLR), 2018. [13] Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Sham Kakade. Gone fishing: Neural active learning with fisher embeddings. In Advances in Neural Information Processing Systems (NeurIPS), 2021. [14] Andreas Kirsch, Sebastian Farquhar, and Yarin Gal. A simple baseline for batch active learning with stochastic acquisition functions. CoRR, abs/2106.12059, 2021. [15] Andreas Kirsch, Joost Van Amersfoort, and Yarin Gal. Batchbald: Efficient and diverse batch acquisition for deep bayesian active learning. Advances in neural information processing systems, 32, 2019. [16] Jennifer A Hoeting, David Madigan, Adrian E Raftery, and Chris T Volinsky. Bayesian model averaging: a tutorial (with comments by m. clyde, david draper and ei george, and a rejoinder by the authors. Statistical science, 14(4):382–417, 1999. [17] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013. [18] Lawrence Cayton et al. Algorithms for manifold learning. Univ. of California at San Diego Tech. Rep, 12(1-17):1, 2005. [19] H. A. Kramers. Brownian motion in a field of force and the diffusion model of chemical reactions. Physica, 7(4):284–304, 1940. [20] Turab Lookman, Prasanna V. Balachandran, Dezhen Xue, and Ruihao Yuan. Active learning in materials science with emphasis on adaptive sampling using uncertainties for targeted design. npj Computational Materials, 5(1):21, 2019. 11

[21] Dezhen Xue, Prasanna V. Balachandran, John Hogden, James Theiler, Deqing Xue, and Turab Lookman. Accelerated search for materials with targeted properties by adaptive design. Nature Communications, 7(1):11241, 2016. [22] Eric Stach, Brian DeCost, A. Gilad Kusne, Jason Hattrick-Simpers, Keith A. Brown, Kristofer G. Reyes, Joshua Schrier, Simon Billinge, Tonio Buonassisi, Ian Foster, Carla P. Gomes, John M. Gregoire, Apurva Mehta, Joseph Montoya, Elsa Olivetti, Chiwoo Park, Eli Rotenberg, Semion K. Saikin, Sylvia Smullin, Valentin Stanev, and Benji Maruyama. Autonomous experimentation systems for materials development: A community perspective. Matter, 4(9):2702–2726, 2021. [23] Hans Lukas, Suzana G. Fries, and Bo Sundman. Computational Thermodynamics: The CALPHAD Method. Cambridge University Press, USA, 1st edition, 2007. [24] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. [25] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2023. [26] J. D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering, 9(3):90–95, 2007. [27] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825):357–362, September 2020. [28] Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28, 2015. [29] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022. [30] Leonardo Ferreira Guilhoto and Paris Perdikaris. Composite bayesian optimization in function spaces using neon—neural epistemic operator networks. Scientific Reports, 14(1):29199, 2024. [31] George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems, 2(4):303–314, 1989. [32] Yulong Lu and Jianfeng Lu. A universal approximation theorem of deep neural networks for expressing probability distributions. Advances in neural information processing systems, 33:3094–3105, 2020. [33] DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. [34] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415, 2016.

Mathematical Notation

Table 2 summarizes the symbols and notation used in this work. For operands that involve expectations, such as expectation E, variance Var and entropy H, a subindex indicates what is the random variable for which the expectation is being taken over. When the context is clear, this sub-index is often omitted. Table 2: Summary of the symbols and notation used in this paper.

Symbol x∈X

Meaning Input to the model

Y ∈Y

Output random variable

θ∈Θ

Parameters of a neural network

Z ∼ PZ

Epistemic index over model hypotheses

ϵ ∼ Pϵ

Aleatoric index for irreducible stochasticity

gθ (x; z, ϵ)

Learned generative map to the output space

pθ (· | x, z)

Aleatoric conditional distribution for fixed z from eq. (3)

pθ (· | x)

Predictive distribution, marginal over PZ from eq. (4)

∗

p (· | x)

True conditional distribution

h# µ

Push-forward of measure µ under map h from eq. (2)

Expectation operator

Var(·)

Variance operator

H(·)

Differential entropy

I(Y ; Z | x)

Epistemic mutual information

KL(p ∥ q)

KL divergence from q to p

TV(p, q)

Total variation distance between p and q

N (µ, Σ)

Normal distribution with mean µ and covariance Σ

N (y; µ, Σ)

PDF of N (µ, Σ) evaluated at y

nens

Number of ensemble members

Number of components in a Gaussian mixture

(z)

αi (x)

Mixture weight of component i for member z

(z) µi (x)

Mean of component i for member z

(z) Ci (x)

Covariance of component i for member z

Weight of ensemble member z under PZ

MI-LB(x)

Mutual Information Lower Bound acquisition function

Training dataset of size n

Identity matrix of size d

Concrete Examples of the Two Index Framework

To ground the two-index framework in practice, we show how three families of models already in common use fit naturally into the formalism of equation (1). In each case we identify the epistemic 13

index Z, the aleatoric index ϵ, the learned map gθ , and the resulting conditional and predictive distributions. Ensemble of Conditional Variational Autoencoders (C-VAEs). Consider an ensemble of nens independently trained C-VAE [28], indexed by k ∈ {1, . . . , nens }. Each model k consists of a decoder dθk : X × L → Y, where L is the latent space. At inference time, a latent code is sampled from the prior ϵ ∼ Pϵ := N (0, I) and passed through the decoder to produce a prediction. In the two-index framework, we set: Z ∼ PZ := Uniform{1, . . . , nens }, ϵ ∼ Pϵ := N (0, Id ), gθ (x; z, ϵ) := dθz (x, ϵ), where θ = (θ1 , . . . , θnens ) collects all decoder parameters. The epistemic index Z = k selects a member of the ensemble, capturing uncertainty about which model best represents the data-generating process. The aleatoric index ϵ is the latent noise injected into the decoder of the selected member, capturing the intrinsic stochasticity of the output given a fixed hypothesis. The aleatoric conditional from eq. (3) becomes pθ (· | x, k) = (dθk (x, ·))# N (0, Id ), which is the generative distribution of the k-th C-VAE. The predictive distribution equation 4 is the equally-weighted mixture of these per-member distributions: pθ (· | x) =

nens 1 X

nens

pθ (· | x, k).

k=1

Epistemic uncertainty, measured by I(Y ; Z | x), is large when the ensemble members disagree on their generative distributions, and vanishes when all decoders produce the same output distribution regardless of which member is selected. Ensemble of Conditional Flow Matching Models. Flow matching models [29] learn a deterministic vector field that transports an initial noise sample ϵ ∼ Pϵ to a target distribution over Y, conditioned on input x. At inference time, the output is obtained by integrating the learned vector field from time t = 0 to t = 1, starting from ϵ. An ensemble of nens such models captures epistemic uncertainty through disagreement across members. The two-index instantiation mirrors the C-VAE case, with a key structural difference: the aleatoric index ϵ is not a latent code injected at an intermediate layer but rather the initial noise supplied to the flow at inference time. Letting Φθk (·; x) : Y → Y denote the learned flow map of member k (i.e., the solution operator of the ODE from t = 0 to t = 1), we set: Z ∼ PZ := Uniform{1, . . . , nens }, ϵ ∼ Pϵ := N (0, Id ), gθ (x; z, ϵ) := Φθz (ϵ; x). The aleatoric conditional is pθ (· | x, k) = (Φθk (· ; x))# N (0, Id ), which, for an exactly learned flow, equals the target conditional distribution of member k. Bayesian Neural Network for Classification. Bayesian Neural Networks (BNNs) are a special case of ENNs [3], with the epistemic index Z playing the role of a sample from the weight posterior p(θ | D). We instantiate this within the two-index framework for a classification setting with ncls classes. Let fw : X → ∆ncls −1 be a neural network with weights w, mapping inputs to the probability simplex. In a BNN, weights are treated as random variables with posterior p(w | D) approximated by a distribution Qθ (w) with learnable parameters θ (e.g., a mean-field Gaussian). The epistemic index Z is a draw from this approximate posterior. The aleatoric index ϵ serves as the source of randomness 14

for sampling a class label from the predicted categorical distribution. Concretely: Z ∼ PZ := Qθ (w), ϵ ∼ Pϵ := Uniform(0, 1), gθ (x; z, ϵ) := min c ∈ {1, . . . , ncls } : Fz (x) c ≥ ϵ , Pc where (Fz (x))c := j=1 [fz (x)]j is the c-th entry of the cumulative sum of the predicted class probabilities under weights z. This is the inverse CDF (quantile) transform: gθ (x; z, ϵ) samples a class label from the categorical distribution Cat(fz (x)) using ϵ as the uniform source of randomness. The aleatoric conditional is then pθ (· | x, z) = Cat(fz (x)), and the predictive distribution is equal to the Bayesian Model Average (BMA): pθ (· | x) = EZ∼Qθ Cat(fZ (x)) . The epistemic term I(Y ; Z | x) measures disagreement among weight posterior samples about the predicted class probabilities, recovering the standard notion of epistemic uncertainty in BNN classification [3]. The aleatoric term EZ [H(fZ (x))] is the expected entropy of each categorical prediction, averaged over the weight posterior. It measures the label ambiguity that remains under a fixed weight configuration. Epinet for Regression. The epinet [3] augments a base network with a small auxiliary network that takes an epistemic index z as an additional input, so that varying z at a fixed input traces out a family of predictions from a single trained model. Paired with a Gaussian aleatoric head, an epinet captures both sources of uncertainty within one network: the epinet generates the epistemic spread, and the predicted variance captures the aleatoric stochasticity. Let µbase : X → RN be the base mean network, ηθ : X × Zepi → RN be the epinet, and θ 2 σθ : X → R N >0 be the predicted aleatoric variance. The mean prediction at epistemic index z is µθ (x, z) := µbase (x) + ηθ (x, z). We set: θ Z ∼ PZ := N (0, Idepi ), ϵ ∼ Pϵ := N (0, IN ), gθ (x; z, ϵ) := µθ (x, z) + σθ (x) ⊙ ϵ, where ⊙ denotes element-wise multiplication. The aleatoric conditional is pθ (· | x, z) = N µθ (x, z), diag σθ2 (x) , a Gaussian whose mean depends on z but whose covariance does not. The predictive distribution is the continuous mixture pθ (· | x) = EZ∼N (0,Idepi ) N µθ (x, Z), diag σθ2 (x) , which is a Gaussian-mean mixture: the mean varies with Z while the covariance is fixed at diag(σθ2 (x)). The epistemic term I(Y ; Z | x) measures how much the mean prediction µθ (x, Z) varies with the epistemic index. The aleatoric term has the closed form EZ [Hϵ (Y | x, Z)] = QN 2 1 N i=1 σθ,i (x) , since the conditional covariance does not depend on Z. The two 2 log (2πe) sources of uncertainty are therefore parameterized by disjoint sub-networks: the epinet ηθ governs the epistemic spread, and the variance head σθ governs the aleatoric noise. The same construction can be extended to the operator learning setting, where inputs and outputs are functions rather than vectors [30].

Limitations of Variance for Multimodal Distributions

The Epistemic Variance acquisition function (Section 2.2), also known as Uncertainty Sampling, is the standard choice for active learning with real-valued predictions. There are settings, however, where it assigns identical scores to distributions with very different uncertainty profiles. We give a concrete example on the unit sphere that illustrates this failure mode. 15

For a distribution p over Rm with mean ȳ = Ep [Y ], we define the multivariate variance as the trace of the covariance matrix: Var(p) := tr Covp (Y ) = Ep ∥Y − ȳ∥22 .

(20)

This is the expected squared ℓ2 distance from a sample to the mean, and it is the quantity that variance-based acquisition functions rank candidates by. Consider the unit sphere S m−1 := {y ∈ Rm : ∥y∥2 = 1} and two inputs x, x′ ∈ Rn . For x, suppose pθ (y | x) is the uniform distribution on S m−1 . For x′ , let Cδ (e1 ) and Cδ (−e1 ) denote the spherical caps of geodesic radius δ > 0 centered at e1 and −e1 respectively, and define the polar distribution pθ (y | x′ ) as the distribution that selects one of the two caps with equal probability and then draws uniformly within it. Both distributions are continuous with respect to the surface measure on S m−1 , both have mean ȳ = 0 by symmetry, and every sample lies on S m−1 . Since ∥Y ∥22 = 1 almost surely and ȳ = 0 for both distributions, equation (20) gives Var(p) = 1 in both cases, exactly and for all m (see Figure 4 for an illustration). Any acquisition function that ranks inputs by variance assigns the same score to x and x′ .

Figure 4: Two distributions on the unit circle with identical variance but different entropy. Left: the uniform distribution spreads samples evenly across the circle. Right: the polar distribution concentrates samples near two antipodal points. Both distributions have mean zero and multivariate variance equal to 1, yet the uniform distribution has strictly higher entropy.

The two distributions are not equally uncertain, however. The polar distribution confines the output to two small patches of the sphere, while the uniform distribution spreads mass over the entire surface. The entropy-based decomposition of Section 2.2 captures this gap. Writing Am = Area(S m−1 ) for the total surface area and Aδ = Area(Cδ ) for the area of each cap, the differential entropies with respect to the surface measure are h pθ (· | x) = log Am ,

h pθ (· | x′ ) = log 2 + log Aδ .

(21)

The entropy gap is log(Am /(2Aδ )), which grows without bound as m → ∞ for any fixed cap radius δ. This follows from concentration of measure on the sphere: the fraction Aδ /Am of the sphere covered by a cap of fixed geodesic radius δ < π/2 vanishes exponentially in m, because the mass of the integrand sinm−2 (ϕ) in the surface area element concentrates near the equator ϕ = π/2. Entropy-based acquisition functions detect this growing gap, while variance-based scores cannot. 16

Proof of Asymptotic Guarantees

In what follows, we formalize the insights from Section 2.3 and theorem 2.1 in order to provide guarantees on the asymptotic behavior under infinite data. First, we prove that epistemic uncertainty measured as the mutual information I(Y ; Z|x) converges to zero. Then, we show that the predicted aleatoric uncertainty EZ Hϵ (Y | Z, x, Dn ) converges to the true quantity H(p∗ (· | x)). As a corollary of these two results, we also observe that the total predictive entropy also converges to the true entropy of the data-generating process. D.1

Mutual Information Collapse

We begin by stating and proving a lemma that will be helpful in proving the main theorem. Lemma D.1 (KL divergence–total variation bound under bounded likelihood ratio). Let p and q be probability measures on a measurable space (Y, B(Y)) with p ≪ q. Suppose there exists a constant C ≥ 1 such that dp/dq ≤ C holds q-almost everywhere. Then KL(p ∥ q) ≤ 2MC · TV(p, q), where C log C − C + 1 MC := max 1, C −1

(22)

for C > 1, and M1 := 1. Proof. Let h := dp/dq, so that 0 ≤ h ≤ C holds q-a.e. Since written as Z Z KL(p ∥ q) =

h log h dq =

h dq = 1, the KL divergence can be

ϕ(h) dq,

(23)

where ϕ(t) := t log t − t + 1. The function ϕ is non-negative and convex on [0, ∞) with ϕ(1) = 0 and ϕ′ (1) = 0. We claim that ϕ(t) ≤ MC |t − 1| for all t ∈ [0, C]. On the interval [0, 1]: since t log t ≤ 0 for t ∈ [0, 1], we have ϕ(t) = t log t + (1 − t) ≤ 1 − t = |t − 1|. On the interval [1, C]: convexity of ϕ t−1 together with ϕ(1) = 0 gives ϕ(t) ≤ ϕ(C) · C−1 , where ϕ(C) = C log C − C + 1. Combining both cases yields ϕ(t) ≤ MC |t − 1| on [0, C]. Integrating against q and applying the bound gives Z Z KL(p ∥ q) = ϕ(h) dq ≤ MC |h − 1| dq = 2MC · TV(p, q), R R where the last equality uses TV(p, q) = 21 |dp − dq| = 21 |h − 1| dq.

(24)

Equipped with this lemma, we now state the first part of the theorem precisely. Theorem D.2 (Mutual Information Collapse). Let (Y, B(Y)) be a measurable space, let (Z, PZ ) and (E, Pϵ ) be probability spaces, let Θ ⊆ Rd , and let gθ : X × Z × E → Y be measurable for each θ ∈ Θ. Fix x ∈ X , define Rthe aleatoric conditional pθ (· | x, z) := (gθ (x; z, ·))# Pϵ and the predictive marginal pθ (· | x) := pθ (· | x, z) dPZ (z). Suppose the following conditions hold. A1 Well-specification. There exists θ∗ ∈ Θ such that pθ∗ (· | x, z) = p∗ (· | x) for PZ -almost every z. A2 Consistency. θ̂n → θ∗ in probability as n → ∞. A3 Uniform continuity in total variation. The map θ 7→ pθ (· | x, z) is continuous at θ∗ in total variation, uniformly in z: for every ε > 0 there exists δ > 0 such that ∥θ − θ∗ ∥ < δ implies supz∈Z TV(pθ (· | x, z), pθ∗ (· | x, z)) < ε. 17

A4 Bounded likelihood ratio. There exists C ≥ 1 such that for all θ in a neighborhood of θ∗ and PZ -almost every z, the aleatoric conditional is absolutely continuous with respect to the predictive marginal and dpθ (· | x, z) ≤C dpθ (· | x)

pθ (· | x)-a.e.

Then i p h →0 Iθ̂n (Y ; Z | x) := EZ KL pθ̂n (· | x, Z) pθ̂n (· | x) −

as n → ∞.

Before proving this theorem, we describe what assumptions A1–A4 intuitively mean. Assumption A1 relates to universality of the chosen model class: the neural network architecture should be expressive enough to contain a parameter θ∗ that captures the desired map, and is a standard assumption in the machine learning literature, having been proved for many choices of architectures via universal approximation theorems [31, 32]. Assumption A2 relates to well-posedness of training: under infinite data, it should be possible to recover the true optimum of the problem. Assumption A3 is a regularity condition requiring that the aleatoric conditionals respond smoothly to parameter perturbations, uniformly across the epistemic index. Assumption A4 requires a uniform upper bound on the likelihood ratio dpθ (· | x, z)/dpθ (· | x); for finite ensembles with prior P weights wz > 0, this is automatically satisfied with C = 1/ minz wz , since the marginal pθ (· | x) = z wz pθ (· | x, z) ≥ wz pθ (· | x, z) dominates each conditional. For continuous epistemic indices, A4 is a mild regularity condition excluding pathological concentration of the conditional relative to the marginal. Proof. Fix x ∈ X and write Iθ := Iθ (Y ; Z | x), pθ (z) := pθ (· | x, z), and pθ := pθ (· | x) for brevity. The strategy is to show that θ 7→ Iθ is continuous at θ∗ with Iθ∗ = 0, after which the conclusion follows from A2 by the continuous mapping theorem. By A1, pθ∗ (z) = p∗ (· | x) for PZ -a.e. z, and marginalizing gives Z Z pθ∗ = pθ∗ (z) dPZ (z) = p∗ (· | x) dPZ (z) = p∗ (· | x),

(25)

so that pθ∗ (z) = pθ∗ for PZ -a.e. z. Consequently KL(pθ∗ (z) ∥ pθ∗ ) = 0 for PZ -a.e. z, and integrating against PZ yields Iθ∗ = 0. For continuity at θ∗ , fix ε > 0. Since pθ (· | x) is a mixture over PZ with pθ (· | x, z) ≪ pθ (· | x), Lemma D.1 together with A4 gives, for every θ in the neighborhood specified by A4 and PZ -a.e. z, KL(pθ (z) ∥ pθ ) ≤ 2MC · TV(pθ (z), pθ ),

(26)

where MC is the constant from equation 22. It therefore suffices to bound EZ [TV(pθ (Z), pθ )]. The triangle inequality for total variation, applied with pθ∗ (z) and pθ∗ as intermediate points, gives TV(pθ (z), pθ ) ≤ TV(pθ (z), pθ∗ (z)) + TV(pθ∗ (z), pθ∗ ) + TV(pθ∗ , pθ ),

(27)

in which the middle term vanishes for PZ -a.e. z by equation 25. By A3 there exists δ > 0 such that ∥θ − θ∗ ∥ < δ implies ε , (28) sup TV(pθ (z), pθ∗ (z)) < 4M C z∈Z controlling the first term in equation 27 uniformly R in z. For the third term, joint convexity of total variation together with the representation pθ = pθ (z) dPZ (z) yields Z ∗ TV(pθ , pθ ) ≤ TV(pθ∗ (z), pθ (z)) dPZ (z) <

ε , 4MC 18

(29)

where the strict bound again follows from equation 28 at the same δ. Integrating equation 27 against PZ and substituting equation 28 and equation 29 gives ε ε EZ [TV(pθ (Z), pθ )] < +0+ 4MC 4MC ε = , (30) 2MC which combined with equation 26 yields Iθ ≤ 2MC · EZ [TV(pθ (Z), pθ )] < 2MC ·

ε = ε. 2MC

(31)

Hence θ 7→ Iθ is continuous at θ∗ . p

Since θ̂n − → θ∗ by A2 and Iθ is continuous at θ∗ with Iθ∗ = 0, the continuous mapping theorem p gives Iθ̂n − → Iθ∗ = 0. D.2

Aleatoric Entropy Convergence

Theorem D.2 establishes that the epistemic mutual information collapses to zero. We now show that the aleatoric uncertainty estimator converges to the entropy of the true conditional. The argument follows the same three-step template as Theorem D.2: point evaluation at θ∗ , continuity at θ∗ , and the continuous mapping theorem. The new ingredient is a regularity condition that does not follow from A3 alone. Assumption A3 controls θ 7→ pθ (· | x, z) in total variation, but differential entropy is not continuous in TV: distributions can be arbitrarily TV-close with entropies that differ by any prescribed amount. We therefore add the following assumption. A5 Uniform entropy continuity. The map θ 7→ H(pθ (· | x, z)) is continuous at θ∗ , uniformly in z: for every ε > 0 there exists δ > 0 such that ∥θ − θ∗ ∥ < δ implies supz∈Z |H(pθ (· | x, z)) − H(pθ∗ (· | x, z))| < ε. For Gaussian-mixture conditionals of the type used in Section 4.1, A5 follows from continuous parametrisation of the mixture weights, means, and covariances together with a uniform lower bound on component covariance eigenvalues in a neighborhood of θ∗ . For finite ensembles, the supremum over z reduces to a maximum over finitely many continuous functions and is automatic. Theorem D.3 (Aleatoric Entropy Consistency). Adopt the setup of Theorem D.2 and suppose A1, A2, and A5 hold. Then h i p Aθ̂n (x) := EZ Hϵ Y | Z, x; θ̂n − → H(p∗ (· | x)) as n → ∞. Proof. Fix x ∈ X and write Aθ := EZ [H(pθ (· | x, Z))] and pθ (z) := pθ (· | x, z) for brevity. The strategy mirrors Theorem D.2: show that θ 7→ Aθ is continuous at θ∗ with Aθ∗ = H(p∗ (· | x)), then apply the continuous mapping theorem under A2. By A1, pθ∗ (z) = p∗ (· | x) for PZ -a.e. z, so H(pθ∗ (z)) = H(p∗ (· | x)) for PZ -a.e. z. Integrating against PZ gives Aθ∗ = H(p∗ (· | x)). (32) For continuity at θ∗ , fix ε > 0. By A5 there exists δ > 0 such that ∥θ − θ∗ ∥ < δ implies sup |H(pθ (z)) − H(pθ∗ (z))| < ε.

(33)

z∈Z

< ε,

(34) 19

so θ 7→ Aθ is continuous at θ∗ . p

→ θ∗ by A2 and θ 7→ Aθ is continuous at θ∗ , the continuous mapping theorem combined Since θ̂n − p with equation 32 gives Aθ̂n − → H(p∗ (· | x)). Theorems D.2 and D.3 together imply that the total predictive entropy converges to the entropy of the true conditional, which gives a complete consistency picture for the entropy decomposition equation 7. Corollary D.4 (Total Predictive Entropy Consistency). Under A1–A5, p

HZ,ϵ (Y | x; θ̂n ) − → H(p∗ (· | x))

as n → ∞.

Proof. By the entropy decomposition equation 7, HZ,ϵ (Y | x; θ̂n ) = Aθ̂n (x) + Iθ̂n (Y ; Z | x). p

Theorem D.3 gives Aθ̂n (x) − → H(p∗ (· | x)) and Theorem D.2 gives Iθ̂n (Y ; Z | x) − → 0. The continuous mapping theorem applied to the sum yields the result.

Proof of Theorem 3.1

Proof. We wish to show that MI-LB(x) ≤ I(Y ; Z | x) for all x ∈ X . Recalling the decomposition from (9), I(Y ; Z | x) = H(Y | x) − EZ H(Y | x, Z) . (35) We bound the two terms on the right-hand side in opposite directions using the results of [4]. Step 1: Lower bounding H(Y | x). The marginal predictive distribution pθ (· | x) is the Gaussian mixture with nens · K components given in equation 11. By the lower bound of [4], H(Y | x) ≥ Hlower (Y | x),

(36) (z)

where Hlower (Y | x) is equation 12 applied to the marginal mixture with weights βz,i , means µi , (z) and covariances Ci . Step 2: Upper bounding EZ [H(Y | x, Z)]. For each fixed z ∈ {1, . . . , nens }, the aleatoric conditional pθ (· | x, z) is a Gaussian mixture with K components as in equation 10. By the upper bound of [4], H(Y | x, Z = z) ≤ Hupper (Y | x, Z = z)

for each z.

(37)

Multiplying by wz ≥ 0 and summing over z, n ens ens X nX EZ H(Y | x, Z) = wz H(Y | x, Z = z) ≤ wz Hupper (Y | x, Z = z). z=1

(38)

z=1

Step 3: Combining the bounds. Substituting equation 36 and equation 38 into equation 35, and using the fact that f (a, b) = a − b is non-decreasing in a and non-increasing in b, I(Y ; Z | x) = H(Y | x) − EZ H(Y | x, Z) n ens X ≥ Hlower (Y | x) − wz Hupper (Y | x, Z = z) = MI-LB(x). z=1

Experimental Details

F.1

Software

Our code is implemented in JAX [24] using the Flax [25] and Optax [33] libraries to define and train our neural networks. We used the JaxMix library [10] as a starting point for working with MDN ensembles. Our complete codebase is available at https://github.com/ PredictiveIntelligenceLab/JaxMix-AL. F.2

Synthetic Multimodal Conditional Problem

This appendix specifies the full data-generating process, model, trainer, and active-learning protocol used in Section 4.1. All hyperparameters listed here are the values used to produce Figures 1a and 6 and are exposed as defaults in examples/multimodal_conditional/experiment_config.py. Input manifold. Latent codes l ∈ RL are drawn l ∼ N (0, IL ) and embedded into the input space via iid x = tanh A l + bm , A ∈ RD×L , Aij ∼ N 0, 1/L , bm ∼ N (0, ID ). Both A and bm are drawn once using manifold_seed = 1 and held fixed across the benchmark; the tanh bounds x to [−1, 1]D . Random Fourier feature map. All mixture parameters are input-dependent through a fixed random Fourier feature map iid iid h(x) = cos Ω x + φ ∈ RP , Ωij ∼ N (0, 1/D), φi ∼ Uniform(0, 2π). The parameters (Ω, φ) are drawn once with dist_seed = 42 and held fixed. Per-component means and variances. For each of the K mixture components, uncentered means and log-variances are µ̃k (x) = Bk h(x) + ck ,

log Σk (x) = Ck h(x) + bvar k ,

with Bk , Ck ∈ RM ×P having i.i.d. N (0, 1/P ) entries, ck ∼ N (0, c2scale IM ), and bvar = 0. The k per-component offsets ck control how far apart the modes sit in output space. Component means are then centered to enforce E[y | x] = 0: µk (x) = µ̃k (x) −

K X

πj (x) µ̃j (x),

j=1

and the covariance is diagonal, Σk (x) = diag(exp log Σk (x)). The variance-coupling coefficient α in the code is set to 0, so log-variances do not depend on πk . Structured mixing weights. To create a clean unimodal/multimodal phase boundary, we use the structured mixing mode. Letting r = ∥x1:L ∥, define a radial gate and K−1 angular scores, g(x) = 21 1 + tanh β (r − r0 ) , s(x) = softmax γ V x1:L ∈ ∆K−2 , V ∈ R(K−1)×L , Vk,: /∥Vk,: ∥ = Vk,: , where the rows of V are unit vectors drawn from N (0, IL ) and normalised, and then assemble the logits ℓ0 (x) = scale · (1 − g(x)),

ℓk (x) = scale · g(x) · sk (x) (k = 1, . . . , K−1),

and set π(x) = softmax(ℓ(x)). Inside the radius (r < r0 ) the mass concentrates on component 0, making p∗ (y | x) effectively unimodal; outside the radius the angular scores activate components 1, . . . , K−1 in sectors. Hyperparameters. Table 3 lists the distribution hyperparameters and Table 4 the model and active-learning hyperparameters. All values match the defaults in experiment_config.py. 21

Table 3: Distribution hyperparameters for the multimodal conditional problem.

Symbol

Value

Meaning

D M L K P cscale α τ β r0 γ scale dist_seed manifold_seed

10 16 4 3 128 10.0 0.0 1.0 8.0 1.3 2.0 3.0 42 1

Input dimension Output dimension Latent manifold dimension Number of true mixture components Random Fourier features Std. of per-component offsets ck Variance-coupling coefficient (disabled) Softmax temperature (unused in structured mode) Transition sharpness Transition radius Angular sharpness Logit magnitude PRNG seed for distribution parameters PRNG seed for manifold map (A, bm )

Oracle NLL. Because p∗ (y | x) is known in closed form, the oracle negative loglikelihood NLL∗ = −E(x,y)∼p∗ [log p∗ (y | x)] can be evaluated exactly on the test set. Under the settings above it equals NLL∗ = 22.98 (computed by compute_true_nll in examples/multimodal_conditional/utils.py). Model architecture. Each ensemble member is a Mixture Density Network with shared backbone: a 2-layer MLP with 128 hidden units per layer and GELU activations [34]. The MLP output is linearly projected to KMDN = 5 mixture components, each with a mean in RM and a diagonal covariance; mixture weights come from a softmax head. Ensemble size is nens = 8, with members initialised from different PRNG seeds and trained independently on the same labeled set (the standard deep-ensemble construction used to instantiate PZ = Uniform{1, . . . , nens }). Training schedule. All members are trained with AdamW (peak learning rate 5 × 10−4 , weight decay 10−2 ) preceded by adaptive gradient clipping at 0.1. The learning rate follows a linear warmup over min(500, niter /5) steps to the peak value, after which it decays exponentially at rate 0.9 per 2,000 steps. The number of gradient steps per active learning round is set adaptively to niter (nlab ) = min 10,000, 10 · nlab , where nlab is the current labeled-set size; this prevents overfitting on small labeled sets in early rounds while allowing full optimisation once nlab ≥ 1000. The training mini-batch size is 128 (full-batch when nlab < 128). Optimiser, architecture, and active-learning hyperparameters are fixed across all acquisition functions; we do not perform per-method tuning. Active-learning protocol. For each of 5 seeds we instantiate a fresh pool of 50,000 candidate inputs, a held-out test set of 2,000 inputs, and an initial labeled set of 100 inputs. Acquisition scores are evaluated on the remaining pool in chunks of 256. Each active-learning run performs 20 rounds, acquiring 50 queries per round, so the final labeled set contains 100 + 20 · 50 = 1100 examples, i.e. 2.2% of the pool. Test NLL is evaluated on the held-out set at the end of every round. Acquisition functions. Five scoring functions are evaluated. Random draws scores i.i.d. Uniform(0, 1). Epistemic Variance uses the trace of the covariance of conditional means across ensemble members. MI-LB is equation 15. BAIT [13] uses the paper’s last-layer mean-head Fisher of ensemble member 0 with one MC sample y ∼ pθ (· | x) per pool point. Core-Set [12] runs k-Center-Greedy on the shared MDN backbone activations of ensemble member 0 (the “final FC layer” recipe of §4.4). The latter two produce no scalar score and bypass the score-to-batch step. Implementation of BAIT and Core-Set in the MDN setting. Both baselines are adapted from their original classification specifications. BAIT requires per-input Fisher embeddings G(x) such that F (x) = G(x)⊤ G(x) approximates the per-sample Fisher information; following the last-layer recipe 22

Table 4: Model and active-learning hyperparameters for the multimodal conditional problem.

Symbol / option

Value

Meaning

nens KMDN hidden features depth activation training batch size peak learning rate weight decay gradient clip niter (cap) iter-per-sample candidate pool size test set size initial labelled AL rounds query batch size acquisition batch size seeds

8 5 128 2 GELU 128 5 × 10−4 10−2 0.1 10,000 10 50,000 2,000 100 20 50 256 {0, 1, 2, 3, 4}

Ensemble size MDN mixture components per member MLP width Number of MLP hidden layers Nonlinearity Mini-batch size for AdamW AdamW peak LR (after warmup) AdamW weight decay Adaptive gradient clipping threshold Max gradient steps per round Slope of adaptive schedule Unlabeled pool |Xpool | Held-out evaluation set Initial labeled budget Number of acquisition rounds Queries acquired per round Chunk size for scoring the pool Data / training seeds for reported curves

of [13] we restrict the Fisher to the mean-head weights of ensemble member 0. Drawing one MonteCarlo sample y ∼ pθ (· | x) and using the closed-form MDN log-likelihood gradient ∇Wµ log p(y | x; θ) = γk (y, x) (y − µk (x))/σk2 (x) ⊗ z(x), where γk are posterior mixture responsibilities and z(x) is the shared backbone activation, the embedding is the flattened outer product, G(x) ∈ R1×hKd . Selection minimises tr((Ftrain + λI)−1 Fcand ) via the forward+backward greedy of [13] with ridge λ = 10−3 and the labelled-set Fisher as the burden matrix. This single-ensemble, single-MC choice is a compute trade-off rather than a capacity claim: extending the Fisher embedding to all nens = 8 members (to capture cross-member disagreement and the mixture-weight signal the mean-head Fisher omits) or to multiple MC samples per pool point would multiply an already dominant selection overhead — BAIT alone consumes ∼ 12 of the ∼ 25 GPU-h total compute budget (Table 13), pushing a richer embedding into a 100+ GPU-h regime no other acquisition here requires. We therefore report the canonical last-layer recipe of [13] and flag the resulting Fisher-estimate variance as a known limitation in the empirical discussion below. Core-Set uses the activations of the shared MDN backbone of ensemble member 0 (the layer immediately before the mixture head) as h-dimensional features, then runs k-Center-Greedy [12] on Euclidean distances initialised at the labelled set. Both methods consume the full pool without subsampling. Selection strategies. Three batch-selection rules convert per-point scores into a query batch of size 50: • Top-k (greedy). Select the 50 points with the highest acquisition scores. Used for all main-text results. • SBAL [14]. We use the softmax variant of SBAL (as opposed to the softrank and power variants in Kirsch et al.), which handles negative scores without modification: sample 50 points without replacement from softmax(score/T ) via the Gumbel-top-k trick. The temperature T interpolates between exploitation (T → 0, recovers top-k) and uniform exploration (T → ∞, recovers Random). All SBAL runs reported here use T = 1.0. • MaxDist [11]. Acquisition-weighted farthest-point sampling in feature space (LCMD-TP variant). Given input features x standardised per dimension and normalised scores s̃ ∈ [0, 1], greedily pick the point maximising dmin (i) · (1 + w s̃(i)), where dmin (i) is the squared distance from point i to its nearest already-selected or already-training point. All MaxDist runs use score-weight w = 1. At w = 0 this reduces to pure farthest-point sampling, and as w → ∞ it approaches top-k on s̃. 23

SBAL and MaxDist only apply to the deterministic acquisitions (Epistemic Variance and MI-LB); combining either with Random is undefined (Random has no meaningful score ranking or signal to weight against). Any acquisition of the form sbal_X or maxdist_X in the experiment log corresponds to applying the named selection strategy on top of base score X. BAIT and Core-Set bypass the score-to-batch step entirely (Fisher-trace forward+backward greedy and k-Center-Greedy respectively) and have no SBAL/MaxDist variants. Results for SBAL and MaxDist variants. Figure 5 and Table 5 report final test NLL at n = 1100 across 5 seeds. At T = 1.0, MI-LB (SBAL) sits roughly halfway between MI-LB (top-k) and Random (33.99 vs 31.12 and 39.73), because softmax sampling at T = 1.0 places non-trivial mass outside the high-score region near the radial gate; lowering T tightens the distribution back onto MI-LB (top-k) and recovers it exactly as T → 0. MI-LB (MaxDist) at w = 1 matches MI-LB (top-k) (30.49 vs 31.12), adding batch diversity without sacrificing score quality; larger w moves it toward top-k on the score, smaller w toward pure farthest-point sampling. BAIT and Core-Set. Core-Set ties MI-LB within seed-to-seed noise (30.57±0.56 vs 31.12±0.34); MI-LB retains the tighter spread. The tie reflects benchmark geometry: pool inputs lie on a 4-D tanh manifold, so k-Center-Greedy spreads queries across that manifold and ends up sampling the gate region that MI-LB targets through entropy disagreement. BAIT effectively ties Variance (32.26±1.55 vs 32.56 ± 1.14), but with the largest seed-to-seed spread of any acquisition we evaluate — the variance of the single-MC Fisher estimate at small budgets. Both baselines bypass the Two-Index decomposition; the next two appendices test whether the apparent tie generalises beyond a benchmark in which input geometry already encodes the multimodality. Table 5: Final test NLL at n = 1100, mean ± std across 5 seeds (min–max in brackets). SBAL uses temperature T = 1.0; MaxDist uses score weight w = 1. Oracle NLL⋆ = 22.98. Acquisition

Mean ± std

Min – Max

Random Epistemic Variance (top-k) Epistemic Variance (SBAL) MI-LB (top-k) MI-LB (SBAL) MI-LB (MaxDist) BAIT Core-Set

39.73 ± 0.90 32.56 ± 1.14 31.01 ± 0.57 31.12 ± 0.34 33.99 ± 0.74 30.49 ± 0.46 32.26 ± 1.55 30.57 ± 0.56

38.89 – 40.97 31.41 – 34.14 30.43 – 31.61 30.58 – 31.37 33.21 – 35.01 29.84 – 30.91 30.40 – 34.34 29.97 – 31.31

Figure 5: Learning curves on the multimodal benchmark for SBAL (T = 1.0) and MaxDist (w = 1) variants of Variance and MI-LB, plus the BAIT and Core-Set baselines, alongside the three top-k curves from Fig. 1a (5 seeds; bands min–max).

F.3

Coupled Double-Well System

This appendix specifies the full data-generating process, model, trainer, and active-learning protocol used in Section 4.2. Simulator. P = 5 particles evolve according to the overdamped Langevin SDE equation 19 under open boundary conditions: particle 1 couples only to particle 2, particle P only to particle P −1, and 24

Figure 6: Spatial distribution of all labeled inputs (initial + acquired) at n = 1100 projected onto (x0 , x1 ) for the five base acquisitions: Random, Epistemic Variance, and MI-LB (top row), and the BAIT and Core-Set baselines (bottom row). By the end of the budget the four informative acquisitions look broadly similar — all concentrate mass along the multimodal edge of the pool (x0 → −1) and the upper/lower boundary — while Random remains the clear outlier, spreading queries uniformly across the support. All five panels inherit the same off-center pool support from the fixed bias b in x = tanh(Al + b) (drawn once at benchmark instantiation; see Input manifold above).

interior particles to both neighbours. Barrier heights are uniform, a = 1, giving per-particle potential V (q) = q 4 /4 − q 2 /2 with minima at q = ±1 and barrier height a/4 = 0.25. We integrate with the Euler–Maruyama scheme X (t) √ (t+dt) (t) (t) (t) (t) (t) (t) qi = qi + a (qi − (qi )3 ) + κ (qj − qi ) dt + σ dt ηi , ηi ∼ N (0, 1), j∈nn(i)

with dt = 0.005 and terminal time T = 5.0 (i.e. nsteps = 1000). Inputs and outputs. Each trajectory is summarised by the 7-dimensional input x (q1 (0), . . . , q5 (0), σ, κ) ∈ RP +2 , drawn component-wise from qi (0) ∼ Uniform(−1.5, 1.5),

σ ∼ Uniform(0.3, 2.0),

κ ∼ Uniform(0, 3.0).

The output is the concatenation of nsnap = 4 snapshots of the particle positions evenly spaced in [T /nsnap , T ] (i.e. at t = 1.25, 2.5, 3.75, 5.0), flattened into y ∈ Rnsnap ·P = R20 . Multi-snapshot outputs give the model direct access to trajectory-level structure and substantially improve training stability compared to single-endpoint outputs, at no additional simulation cost. Hyperparameters. Table 6 lists the simulator and data hyperparameters and Table 7 the model and active-learning hyperparameters. Model architecture. Each ensemble member is an MDN with a 3-hidden-layer MLP backbone, 128 hidden units per layer, GELU activations, and a final linear projection to KMDN = 8 mixture components with diagonal covariances over the 20-dimensional output. The ensemble size is nens = 8, with members initialised from different PRNG seeds and trained independently on the same labeled set. 25

Table 6: Simulator and data hyperparameters for the coupled double-well benchmark.

Symbol / option

Value

Meaning

P a T dt nsnap boundary qi (0) range σ range κ range

5 1 (uniform) 5.0 0.005 4 open [−1.5, 1.5] [0.3, 2.0] [0, 3.0]

Number of coupled particles Barrier-height coefficient Terminal integration time Euler–Maruyama step (1000 steps) Snapshots recorded per trajectory Chain ends couple only to one neighbour Initial-configuration prior Noise-intensity prior Coupling-strength prior

Training schedule. All members are trained with AdamW using the same optimiser configuration as in Appendix F.2: peak learning rate 5 × 10−4 , weight decay 10−2 , linear warmup followed by exponential decay at rate 0.9 per 2,000 steps, and adaptive gradient clipping at 0.1. The number of gradient steps per round is set adaptively to niter (nlab ) = min(10,000, 10 · nlab ), with mini-batch size 128. Active-learning protocol. For each of 5 seeds we instantiate a fresh pool of 50,000 candidate inputs, a held-out test set of 2,000 inputs, and an initial labeled set of 100 inputs. Acquisition scores are evaluated on the remaining pool in chunks of 256. Each run performs 20 rounds of 50 queries each, so the final labeled set contains 100 + 20 · 50 = 1,100 trajectories (2.2% of the pool). Test NLL is evaluated at the end of every round. Seeds {0, 1, 2, 3, 4} are used for all reported curves. Table 7: Model and active-learning hyperparameters for the coupled double-well benchmark.

Symbol / option

Value

Meaning

8 8 128 3 GELU 128 5 × 10−4 10−2 0.1 10,000 10 50,000 2,000 100 20 50 256 {0, 1, 2, 3, 4}

Acquisition functions and selection strategies. Five base acquisitions, defined in Appendix F.2 unchanged: Random, Variance, MI-LB, BAIT [13], Core-Set [12]. Selection strategies: top-k for Random / Variance / MI-LB; SBAL (T = 1.0) for Variance and MI-LB; MaxDist (w = 1) for MI-LB. BAIT and Core-Set bypass the score-to-batch step entirely. Results for SBAL and MaxDist variants. Figure 7 and Table 8 report the full learning curves and the final test NLL at n = 1100 across 5 seeds. Two patterns are worth highlighting. First, MaxDist pairs well with MI-LB: MI-LB (MaxDist) at w = 1 finishes at 131 ± 15, within a factor of two of MI-LB (top-k) and comparable to Epistemic Variance (top-k). Acquisition-weighted farthest-point 26

sampling adds batch diversity without discarding the informative ranking. Second, SBAL severely degrades both base scores on this benchmark. MI-LB (SBAL) ends at 553 ± 92 and Epistemic Variance (SBAL) at 457 ± 46, both within striking distance of the Random baseline (518 ± 60) and roughly 4–8× worse than their respective top-k counterparts (122 and 71). This is the opposite of the multimodal-conditional result in Appendix F.2, where SBAL was competitive or helpful, and it reflects a quantitative rather than a qualitative difference: the total budget (1,100 points) is small relative to the pool, so softmax sampling at T = 1.0 spreads enough mass onto low-score points that the effective fraction of informative queries collapses toward the Random baseline; annealing T toward zero would tighten the distribution onto the top-k result and recover it in the limit T → 0. This is consistent with the original observations of Kirsch et al. [14]: stochastic batch acquisition helps when the top of the score ranking is corrupted by approximation noise and diverse sampling hedges against that noise; it hurts when the top ranking is already approximately correct and exploration trades informative points for uninformative ones. BAIT and Core-Set. Both geometric baselines fail decisively on this benchmark: Core-Set finishes at 304.3 ± 54.8 and BAIT at 281.3 ± 15.1, ∼ 4× worse than MI-LB (70.8 ± 8.0) and within 1.7× of Random (518.1 ± 60.1). Mode structure here lives in output space (Kramers escape between q = ±1): k-Center-Greedy spreads queries uniformly over (σ, κ) regardless of where the bimodal regime sits, and the single-MC last-layer Fisher embedding is dominated by within-mode noise. Their joint collapse is the strongest evidence in the paper that MI-LB’s advantage is not subsumed by feature-space coverage or last-layer information geometry. Table 8: Final test NLL at n = 1100, mean ± std across 5 seeds (min–max in brackets) for the coupled double-well benchmark. All SBAL runs use T = 1.0; MaxDist uses score weight w = 1. Acquisition

Mean ± std

Min – Max

Random Epistemic Variance (top-k) Epistemic Variance (SBAL) MI-LB (top-k) MI-LB (SBAL) MI-LB (MaxDist) BAIT Core-Set

518.1 ± 60.1 122.1 ± 11.3 456.8 ± 45.7 70.8 ± 8.0 553.0 ± 91.9 130.7 ± 14.9 281.3 ± 15.1 304.3 ± 54.8

462.6 – 616.2 109.1 – 138.5 392.9 – 515.3 59.9 – 81.6 423.5 – 630.2 113.5 – 147.0 268.8 – 305.3 249.3 – 377.5

Figure 7: Learning curves for all eight (acquisition, selection-strategy) combinations on the coupled double-well benchmark including the BAIT and Core-Set baselines (5 seeds; bands min–max). MaxDist-MI-LB stays within ∼ 2× of top-k MI-LB; both SBAL variants and the two geometric baselines (BAIT, Core-Set) collapse toward the Random regime.

Why a mixture head: K = 1 vs K = 8. To justify the choice of a mixture head (KMDN = 8) over a single-Gaussian head (K = 1), we train both with the identical architecture used in the 27

AL experiments (depth-3 MLP, 128 hidden units, nens = 8) on a larger offline dataset of 200,100 trajectories drawn from the same input distribution as the active-learning experiments, and evaluate both on a held-out test set of 5,000 trajectories. The ensemble-averaged test NLL is 15.90 for K = 8 versus 21.41 for K = 1, a gap of 5.51 nats on the 20-dimensional output (roughly 0.28 nats per output dimension). Absolute NLL values here are below those in the AL experiments because training uses 200,100 examples rather than the 1,100 queried by the AL loop; the quantity of interest is the K = 8-vs-K = 1 gap, which is driven by the head structure, not the training-set size. The K = 1 model must place a single Gaussian per input, so wherever the conditional is bimodal it is forced to smear mass between the wells and pay an unavoidable log-likelihood penalty; the K = 8 mixture can allocate one component per well. Figure 8 shows the qualitative picture on the test distribution: ground-truth samples of (q0 (T ), q1 (T )) concentrate near (±1, ±1) in the bimodal regime and the K = 8 MDN recovers the same multi-well envelope, while the K = 1 model collapses to a single broad ellipse near the origin. This is the failure mode anticipated in Section 4.2 and the quantitative reason the active-learning experiments on this benchmark use KMDN = 8 throughout.

Figure 8: Coupled double-well benchmark: final-snapshot positions (q0 (T ), q1 (T )) for 4,000 held-out test inputs drawn from the same joint distribution as the AL test set. Left: ground-truth simulator samples; mass p concentrates near (±1, ±1), the signature of the per-particle bimodal regime at σ ≳ a/2 ≈ 0.71, with a diagonal enhancement that reflects neighbour alignment under non-zero κ. Middle: samples from the K = 8 MDN ensemble (net 0); the mixture head preserves the same multi-well envelope (test NLL 15.90). Right: samples from a K = 1 single-Gaussian MDN trained with the same architecture and budget; per-input unimodality forces mass into a single broad ellipse centred near the origin (test NLL 21.41; +5.51 nats over K = 8). Dotted lines mark the well locations q = ±1.

F.4

Material Science Application: Alloy Phase Competition

This appendix details the synthetic phase-competition simulator and the active-learning protocol summarised in Section 4.3, and reports the SBAL and MaxDist batch-diversity variants deferred from the main text. Phase model. The simulator implements a softmin-over-quadratic-free-energies model inspired by the CALPHAD framework [23]. For each of Nϕ = 4 nominal phases ϕ we draw a positive-definite 3×3 Hessian as Hϕ = Rϕ Rϕ⊤ +0.3·I3 with Rϕ a Gaussian matrix (Rϕ,ij ∼ N (0, 0.52 )), and a linear bias bϕ ∈ R3 with bϕ,i ∼ N (0, 22 ). Writing x3 = (xA , xB , xC ) for the composition coordinates extended onto the 3-simplex, the Gibbs free energy of phase ϕ is the quadratic form ⊤ Gϕ (x3 ) = 21 x⊤ 3 Hϕ x3 + bϕ x3 ,

and the phase posterior at temperature τG is π(ϕ | x3 ) ∝ exp(−Gϕ (x3 )/τG ). This posterior does not depend on the process parameters p ∈ Rnproc ; phase assignment is a function of composition alone. The conditional response mean for phase ϕ is ⊤ ⊤ 1 µϕ (x3 , p) = c⊤ ϕ x3 + dϕ + 2 sin(ωϕ x3 ) + Wϕ p,

with cϕ ∈ R3 , dϕ ∈ R, ωϕ ∈ R3 , and Wϕ ∈ Rnproc all drawn from independent Gaussians (entrywise scales cµ-scale = 6 for cϕ , 2 for dϕ and ωϕ , 1 for Wϕ ). The per-phase log-variance is the affine 28

2 2 function log σϕ2 (x3 ) = e⊤ ϕ x3 + fϕ with eϕ,i ∼ N (0, 0.5 ) and fϕ ∼ N (−1, 0.3 ), so per-phase noise scales vary smoothly with composition around a typical σϕ ≈ 0.6.

The output is sampled by first drawing ϕ ∼ π(· | x3 ), then drawing y | ϕ, x3 , p ∼ N (µϕ (x3 , p), σϕ2 (x3 )); the ground-truth conditional p⋆ (y | x) is therefore a Gaussian mixture in y with composition-dependent mixing weights, composition- and-process-dependent component means, and composition-dependent component variances. The number of components with appreciable weight at any given x is determined by the realised phase structure of the chosen system seed, summarised below. Inputs and outputs. The input is the 8-dimensional vector x = (xA , xB , p1 , . . . , p6 ) with (xA , xB , xC ) drawn from a symmetric Dirichlet(α = (1, 1, 1)) on the 3-simplex and pi ∼ Uniform(−1, 1) component-wise. The output is the scalar response y ∈ R. Realised phase structure at the system seed. At the system seed used throughout (system_seed = 12, Nϕ = 4, τG = 0.08), all four nominal phases carry appreciable mass on the simplex: a grid-evaluation of π(ϕ | x3 ) on 20,301 uniform simplex points gives marginal phase masses 48.18%, 25.60%, 18.84%, 7.38% (sorted, summing to 100%). The boundary region defined by maxϕ π(ϕ | x3 ) < 0.7 occupies 11.40% of the simplex slice; tightening to maxϕ π < 0.5 contracts it to 1.14%, and relaxing to maxϕ π < 0.8 expands it to 17.60%. The realised problem is thus a genuine 4-phase competition with curved boundaries between phase domains in the 8-dimensional input space. Inside any single phase domain the property response collapses to a single Gaussian whose mean varies smoothly with the 6-dimensional process subspace, while in the boundary region the conditional is a genuine multi-component mixture. Model architecture. Each ensemble member is an MDN with a 2-hidden-layer MLP backbone, 64 hidden units per layer, GELU activations [34], and a final linear projection to KMDN = 4 mixture components with diagonal covariances over the scalar property output. The ensemble size is nens = 8, with members initialised from different PRNG seeds and trained independently on the same labeled set. Training schedule. All members are trained with AdamW (peak learning rate 2 × 10−4 , weight decay 5 × 10−2 ) preceded by adaptive gradient clipping at 0.1, with the same warmup-then-exponentialdecay schedule as in Appendix F.2. The number of gradient steps per round is set adaptively to niter (nlab ) = min(40,000, 200 · nlab ) with a minimum of 2,000 steps per round, and mini-batch size 64. The smaller peak LR and stronger weight decay (relative to the multimodal and double-well benchmarks) reflect the larger per-round step budget; as in the other benchmarks, optimiser, architecture, and active-learning hyperparameters are fixed across all acquisition functions and are not tuned per method. Hyperparameters. Tables 9 and 10 list the simulator/data and model/AL hyperparameters used for all 30 runs reported below. Active-learning protocol. For each of 5 seeds we instantiate a fresh pool of 50,000 candidate inputs, a held-out test set of 2,000 inputs, and an initial labeled set of 100 inputs. Acquisition scores are evaluated on the remaining pool in chunks of 256. Each run performs 30 rounds of 15 queries each, so the final labeled set contains 100 + 30 · 15 = 550 inputs (1.1% of the pool). Test NLL is evaluated at the end of every round. Acquisition functions and selection strategies. We evaluate the same five base acquisitions as in Appendix F.2 (Random, Epistemic Variance, MI-LB, BAIT [13], Core-Set [12]), paired with three selection strategies for the score-based ones: top-k for Random / Variance / MI-LB, SBAL for Variance and MI-LB, and MaxDist for MI-LB. SBAL temperature is T = 1.0; MaxDist score weight is w = 1. BAIT and Core-Set bypass the score-to-batch step entirely (Fisher-trace forward+backward greedy and k-Center-Greedy on the MDN backbone respectively). Definitions follow Appendix F.2 unchanged. Seed-to-seed spread of the top-k base acquisitions. Section 4.3 reports that MI-LB is markedly more consistent across seeds than Random or Epistemic Variance in the data-scarce regime, with the 29

Table 9: Simulator and data hyperparameters for the ternary phase-competition benchmark.

Symbol / option

Value

Meaning

Nϕ τG nproc cµ-scale Hϕ raw scale Hϕ regulariser bϕ scale dϕ scale ωϕ scale Wϕ scale eϕ scale fϕ mean / std (xA , xB , xC ) prior pi prior system_seed candidate pool size test set size initial labelled

4 0.08 6 6 0.5 0.3 · I3 2 2 2 1 0.5 −1 / 0.3 Dirichlet(1, 1, 1) Uniform(−1, 1) 12 50,000 2,000 100

Nominal phases (all four active at the chosen seed) Free-energy softmin temperature Process parameters per input Scale of per-phase composition coupling Rϕ,ij ∼ N (0, 0.52 ) PD floor on Hessian Linear-bias scale on composition Per-phase mean offset scale Sinusoidal-modulation frequency scale Per-phase process-coupling scale Composition-dependence of log σϕ2 Per-phase log-variance offset Composition prior on 3-simplex Process-parameter prior PRNG seed for Hϕ , bϕ , cϕ , . . . Unlabeled pool |Xpool | Held-out evaluation set Initial labeled budget

Table 10: Model and active-learning hyperparameters for the ternary phase-competition benchmark.

Symbol / option

Value

Meaning

nens KMDN hidden features depth activation training batch size peak learning rate weight decay gradient clip min iter per round niter (cap) iter-per-sample AL rounds query batch size acquisition batch size seeds

8 4 64 2 GELU 64 2 × 10−4 5 × 10−2 0.1 2,000 40,000 200 30 15 256 {0, 1, 2, 3, 4}

Ensemble size MDN mixture components per member MLP width Number of MLP hidden layers Nonlinearity [34] Mini-batch size for AdamW AdamW peak LR (after warmup) AdamW weight decay Adaptive gradient clipping threshold Floor on adaptive iteration count Max gradient steps per round Slope of adaptive schedule Number of acquisition rounds Queries acquired per round Chunk size for scoring the pool Data / training seeds

gap closing by n ≳ 400. Table 11 gives the underlying numbers: the seed-to-seed min–max range of test NLL averaged over each budget band. In the early band n ∈ [115, 250], MI-LB’s range is ∼ 2.6× tighter than Random’s and ∼ 4.5× tighter than Epistemic Variance’s; the inflated Variance spread comes from seeds in which the second-moment signal selects unhelpful queries before the ensemble has accumulated enough data to disagree informatively over modal structure. By the late band (n ≳ 400) all three ranges sit within a factor of ∼ 2.4 of each other. Table 11: Seed-to-seed min–max range of test NLL (in nats), averaged over each budget band, for the top-k base acquisitions on the synthetic phase-competition benchmark (system_seed = 12, 5 seeds). Acquisition (top-k)

Early (n ∈ [115, 250])

Late (n ≳ 400)

Random Epistemic Variance MI-LB

1.82 3.13 0.69

0.24 0.31 0.13

Results for SBAL and MaxDist variants. Table 12 reports the final test NLL at n = 550 across 5 seeds for all combinations, and Figure 9 shows the corresponding learning curves. Three observations. First, MaxDist is the mean-best acquisition on this benchmark: MI-LB (MaxDist) at w = 1 finishes at 1.945 ± 0.061, narrowly ahead of MI-LB (top-k) at 1.985 ± 0.043. The seed-to-seed spread is large enough that the two are not separated cleanly—per-seed, MaxDist wins on 2 seeds, ties on 1, and is within 0.01 nat of top-k on the other 2—so the safest summary is that MaxDist matches MI-LB (top-k) and offers a small mean improvement, the same qualitative conclusion as on the coupled double-well benchmark of Appendix F.3. Second, SBAL hurts MI-LB: MI-LB (SBAL) at T = 1.0 degrades from 1.985 ± 0.043 to 2.109 ± 0.054, with MI-LB (top-k) winning on all 5 seeds. This is consistent with the small-budget, narrowhigh-information-manifold argument of Appendix F.3: at T = 1.0 enough mass leaks onto low-score points that the effective fraction of informative queries drops, and the realisable improvement from batch diversity is bounded above by the MaxDist result that preserves the informative ranking exactly. Third, SBAL modestly helps Epistemic Variance: 2.029 ± 0.050 for SBAL versus 2.062 ± 0.087 for top-k, with the SBAL variant winning on 4 of 5 seeds and producing a tighter seed-to-seed spread. Because the Variance-based score is itself already a noisier ranking signal than MI-LB on this benchmark, its top-k ordering is less informative, and stochastic relaxation does not cost what it does for MI-LB, at this temperature it actively helps by hedging against unhelpful early queries. BAIT and Core-Set. Core-Set ties MI-LB within seed-to-seed noise (1.989 ± 0.043 vs 1.985 ± 0.043); the tie reflects benchmark geometry, with phase boundaries lying directly on the 2-D composition simplex so that k-Center-Greedy in feature space hits the same boundary MI-LB targets through entropy disagreement. BAIT is the only acquisition whose mean NLL falls within seed-to-seed noise of Random (2.190 ± 0.090 vs 2.209 ± 0.085): the 0.019 improvement is an order of magnitude smaller than the ∼ 0.09 standard deviation of either run, whereas every other acquisition in Table 12 beats Random by at least 0.10 in mean NLL — well outside the noise band. The single-MC last-layer Fisher embedding produces a noisy candidate ranking on the 1-D scalar output, and the 30 × 15 small-budget protocol does not give the forward+backward greedy enough rounds to recover — the same failure mode anticipated in Appendix F.2 (where BAIT had the largest seed-to-seed spread of any reported method) but more pronounced here because of the lower-dimensional output. Remark on per-acquisition simplex visualisation. A per-acquisition visualisation of selected points projected onto the (xA , xB ) simplex slice would complement the boundary-seeking narrative of Section 4.3. The production runs reported in Table 12 did not log query indices to the experiment tracker, so we do not include such a figure here; the example notebook accompanying this benchmark provides a single-seed reproduction that materialises the simplex projection from local state. Why a mixture head: K = 1 vs K = 4. To justify the choice of a mixture head (KMDN = 4) over a single-Gaussian head (K = 1), we train both with the identical architecture used in the AL experiments (depth-2 MLP, 64 hidden units, nens = 8) on a larger offline dataset of 100,000 samples drawn from the same input distribution as the active-learning experiments, and evaluate both on a held-out test set of 10,000 samples. The ensemble-averaged test NLL is 1.34 for K = 4 31

Table 12: Final test NLL at n = 550, mean ± std across 5 seeds {0, 1, 2, 3, 4} (min–max in brackets) for the synthetic phase-competition benchmark at system_seed = 12. SBAL temperature is T = 1.0 for both SBAL runs; MaxDist uses score weight w = 1. Acquisition Random Epistemic Variance (top-k) Epistemic Variance (SBAL) MI-LB (top-k) MI-LB (SBAL) MI-LB (MaxDist) BAIT Core-Set

Mean ± std

Min – Max

2.209 ± 0.085 2.062 ± 0.087 2.029 ± 0.050 1.985 ± 0.043 2.109 ± 0.054 1.945 ± 0.061 2.190 ± 0.090 1.989 ± 0.043

2.136 – 2.346 1.918 – 2.122 1.977 – 2.107 1.911 – 2.017 2.057 – 2.174 1.870 – 2.023 2.087 – 2.265 1.951 – 2.058

Figure 9: Synthetic phase-competition benchmark: test NLL vs. training-set size for all eight (acquisition, selection-strategy) combinations including the BAIT and Core-Set baselines (5 seeds; shaded bands min–max). SBAL T = 1.0, MaxDist w = 1. The dotted line marks the oracle floor NLL∗ ≈ 1.33. The y-axis is clipped to [1.2, 7.0] to suppress early-iteration training-instability spikes for the SBAL variants and BAIT.

versus 1.88 for K = 1, against an analytically computed oracle floor of NLL∗ = 1.33. The K = 4 ensemble therefore sits within 0.01 nat of the oracle—consistent with the fact that the simulator is itself a 4-component Gaussian mixture per input, so the K = 4 MDN model class contains the true distribution. The K = 1 model pays a 0.55-nat penalty at the same budget: forced to place a single Gaussian per input, it cannot resolve phase competition in the boundary region and must smear mass across phase means that differ by several nats (per-phase median means +2.9, +6.9, +3.7, −4.0 on this simulator instance). The penalty is smaller than on the coupled double-well benchmark (5.51 nats, Appendix F.3) because the per-phase response is a unimodal Gaussian with bounded variance, so the smearing cost is bounded; it is still clearly nonzero, and the active-learning experiments on this benchmark therefore use KMDN = 4 throughout. Figure 10 shows the qualitative picture on a 120-point simplex grid with process parameters pinned to zero. The K = 4 ensemble’s predicted conditional mean Ê[Y | x] tracks the ground-truth E∗ [Y | x] to within roughly ±1 nat, with residuals smoothly distributed across the simplex. The K = 1 residual is visibly structured and larger (up to ±4) at the intersections of the phase-boundary network, where the single-Gaussian head cannot simultaneously fit the correct mean and an appropriate variance—exactly the failure mode anticipated in Section 4.3. F.5

Compute Resources

All experiments were run on a workstation with 7× NVIDIA RTX A6000 GPUs (48 GB VRAM each), 128 CPU cores, and approximately 504 GiB of system memory. Each active-learning run uses a single A6000 GPU; peak GPU memory fits comfortably within the 48 GB budget for all configurations (small MLP backbones, ensemble size 8), so GPU memory is not a binding constraint. 32

Figure 10: Synthetic phase-competition benchmark (system_seed = 12): predicted conditional mean on a 120-point simplex grid (process parameters pinned to zero) for the K = 4 MDN ensemble (top row) and a K = 1 single-Gaussian MDN (bottom row), both trained offline on 100,000 samples with the architecture used throughout the AL experiments. Left column: ground-truth E∗ [Y | x]. Middle column: ensemble-predicted Ê[Y | x]. Right column: residual (predicted − true). The K = 4 residuals sit within ∼ ±1 and are smoothly distributed; the K = 1 residuals reach ±4 and concentrate at the intersections of the phase-boundary network, where the single-Gaussian head cannot represent the competing phase components. The ensemble-averaged test NLL on a held-out set of 10,000 samples is 1.34 (K = 4) vs 1.88 (K = 1), against an oracle floor of 1.33.

Per benchmark we evaluate eight (acquisition, selection-strategy) combinations across 5 seeds, giving 40 runs per benchmark and 120 runs in total across the three benchmarks. Per-run wall-clock times, extracted from the W&B logs and reported as the median over 5 seeds, are summarised in Table 13. Total compute across the three benchmarks is approximately 25 GPU-hours, of which ∼ 12 h (roughly half) is the 15 BAIT runs, BAIT is the only acquisition with non-trivial selection overhead, since its forward+backward greedy on the per-input Fisher embeddings of the full 50,000-point pool scales with both pool size and query batch. All other acquisitions, including Core-Set, complete in roughly the same wall time as Random per benchmark. Table 13: Per-run wall-clock time on a single NVIDIA RTX A6000, median over 5 seeds, for the three paper benchmarks. Each entry is the end-to-end runtime of one 20-round (multimodal, double-well) or 30-round (ternary) active-learning run. Total compute across all 120 runs is ∼ 25 GPU-hours, of which ∼ 12 h is BAIT (∼ 8.6 h on the double-well benchmark alone). Acquisition (selection)

Multimodal

Double-well

Ternary

Random (top-k) Epistemic Variance (top-k) MI-LB (top-k) BAIT (top-k) Core-Set (top-k) Epistemic Variance (SBAL) MI-LB (SBAL) MI-LB (MaxDist)

2.7 min 3.2 min 3.1 min 22.4 min 3.5 min 3.1 min 3.1 min 3.7 min

3.1 min 3.4 min 3.4 min 1.35 h 4.4 min 3.4 min 3.4 min 3.9 min

12.8 min 13.6 min 13.1 min 15.0 min 14.2 min 13.2 min 13.2 min 13.6 min

Record · ID 187338 · SHA-256 aad71cebb4881c04
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.