Neural Posterior Estimation with Latent Basis Expansions
Authors: Declan McNamara
Yicun Duan
Jeffrey Regier
Venue: ICLR 2026 Poster (2026)
Keywords: forward KL divergence; simulation-based inference; variational inference; exponential family
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: Our approach parameterizes amortized variational distributions as linear combinations of basis functions; both the coefficients and basis functions themselves are fit by targeting the same variational objective.

Abstract: Neural posterior estimation (NPE) is a likelihood-free amortized variational inference method that approximates projections of the posterior distribution. To date, NPE variational families have been either simple and interpretable (such as the Gaussian family) or highly flexible but black-box and potentially difficult to optimize (such as normalizing flows). In this work, we parameterize variational families via basis expansions of the latent variables. The log density of our variational distribution is a linear combination of latent basis functions (LBFs), which may be fixed a priori or adapted to the problem class of interest. 
Our training and inference procedures are computationally efficient even for problems with high-dimensional latent spaces, provided only a low-dimensional projection of the posterior is of interest, owing to NPE's automatic marginalization capabilities.
In numerous inference problems, the proposed variational family exhibits better performance than existing variational families used with NPE, including mixtures of Gaussians (mixture density networks) and normalizing flows, as well as outperforming an existing basis expansion method for variational inference.

Published as a conference paper at ICLR 2026
NEURAL POSTERIOR ESTIMATION WITH LATENT BASIS EXPANSIONS
Declan McNamara, Yicun Duan, Jeffrey Regier Department of Statistics University of Michigan Ann Arbor, MI 48109, USA (declan, pduan, regier) regier)@umich.. @umich. edu
ABSTRACT
Neural posterior estimation (NPE) likelihood-fres amortized variational in- ference method that approximates projections of the posterior distribution. To date, NPE variational families have been either simple and interpretable (such as the Gaussian family) highly flexible but black-box and potentially difficult optimize (such as normalizing flows). In this work, we parameterize variational families via basis expansions of the latent variables. The log density of our varia- tional distribution is a linear combination of latent basis functions (LBFs), which may be fixed priori adapted the problem class of interest. Our training and inference procedures are computationally efficient even for problems with nigh-dimensional latent spaces, provided only a low-dimensional projection of the posterior is of interest, owing to NPE's automatic marginalization capabilities. In numerous inference problems, the proposed variational family outperforms exist- ing variational families used with NPE, including mixtures of Gaussians (mixture density networks) and normalizing flows, and also outperforms an existing basis expansion method for variational inference.
INTRODUCTION
Neural Posterior Estimation (NPF) is an increasingly popular approach to Bayesian inference PPa pamakarios Murray 2016 Cranmer et al., 2020 Dax al. [021: Ward al., 2022). In NPE, neural network is trained exclusively on synthetic latent variables drawn from the prior paired with observations- learn the inverse mapping from observations to latent variables. Once trained, this network can produce posterior approximations for real data in single forward pass. In contrast to traditional BELBO-based) variational inference, NPE does not require likelihood evaluations. Furthermore, when the generative model contains both parameters of interest and nuisance variables, NPE can automatically marginalize over the nuisance parameters during training: by simulating complete data and then discarding the nuisance variables create training pairs, the method infers posterior projections for the parameters of interest (Ambrogioni et al. 2019).
Despite these advantages, NPE shares with traditional ELBO-based variational inference a funda- mental trade-off between the flexibility of the variational family and the tractability of optimization. Simple variational families, such Gaussians, enable stable optimization but often lack the ex- pressiveness needed for complex posterior geometries. More flexible alternatives, such as Gaussian mixture models and normalizing flows. offer greater flexibility but can create difficult optimization landscapes with shallow local optima. Recent theoretical work (McNamara et al. (024a) has estab- lished conditions for global convergence in NPE that hold for simple Gaussian variational families, but these results do not extend to the more expressive families commonly deployed in practice.
In this work, we propose a variational family specialized for NPE that leverages NPE's automatic marginalization and likelihood-free nature. Unlike standard variational inference, where nuisance variables must be modeled in the variational distribution, NPE applications typically require posteriors over just few scientifically relevant parameters: high- nigh-dimensional posterior samples do not directly aid interpretation but must instead be post-proceste to estimate www-dimensional interpretable quantities. Because NPE does require likelihood evaluations, this post- poot-processing often be

---

Published as a conference paper at ICLR 2026
incorporated directly into the Bayesian model, resulting in a low-dimensional latent space of interest. In these oowdiiensioonl settings, even numerical integration is computationally feasible, freeing us from the usual requirement of closed- form normalization.
We leverage this freedom by parameterizing the log density of variational distributions through latent basis expansions: neural network processes observations to produce coefficients for linear combinations of basis functions over the latent space. This approach yields distributions in the expo- nential family-among the most flexible classes available (Pacchiardi & Dutta 2022] Khemakhem maintaining favorable optimization properties. The
et al. 2020 Sripeeumbudur et al. resulting method, which we refer as Latent Basis Function NPE (LBF-NPE), optimizes over the class of all exponential families of a fixed dimension K by adaptively fitting the basis functions, denoted 8,(2) E TR,K, Simultaneously, amortization is performed by fitting a separate network (1) E IRKK that maps observations z to coefficients of these basis functions (Section
In Section 4 we introduce and analyze several variants of LBF-NPE. For a variant with fixed basis functions, such as B-splines or wavelets, optimization is convex despite the log- nogmoaaaa,izer, providing stable training that has proven elusive for more complex variational families, and ensuring stable convergence to global optima under the conditions presented by McNamara et al. (2024a). Alter- natively, both coefficients (I) and basis functions (=) can be fitted jointly through alternating optimization that exploits marginal convexity in each component. We employ stereographic projec- tion reparameterization to address identifiability issues in this adaptive setting, constraining outputs to the unit hypersphere and stabilizing training.
We demonstrate superior performance of LBF-NPE across diverse inference problems, from synthetic benchmarks to real scientific applications (Section (). LBF-NPE consistently converges to global optima on multimodal problems where mixture density networks converge shallow local minima. LBF-NPE with just 20 basis functions achieves wrder-of-magnitude improvements in KI. divergence over both MDNs and normalizing flows on complex 2D posteriors. The method successfully captures multimodal posteriors in astronomical object detection and substantially outperforms MDN baselines on cosmological redshift estimation using the LSST DESC DC2 survey dataset.
BACKGROUND
2.1 ELBO-BASED VARIATIONAL INFERENCE
In variational inference (VI), numerical optimization is used to select an approximation q(=) of the posterior distribution of some model p(z, z) on observables æ and latent variables z. The most common variational objective, the evidence lower bound (ELBO), targets minimization of the reverse KI. divergence (Blei et al. 2017: Zhang et al. 2019 Kingma & Welling 2019) by constructing the variational quantity
p(2, x) q(2;,)
ELBO (1) (g(a))) log
< log)(x)
and performing maximization in y for a fixed choice of x. Optimizing the ELBO remains long- standing challenge VI. Firstly, even simple variational families, such as the Gaussian, can exhibit problematic optimization landscapes without careful parameterization. Targeting the ELBO. even in the non-amortizee amootized setting, is generally a nonconvex problem (Liu et al. 2023 Domke 2020 Domke al. 2023). Secondly, and most significantly in our setting. ELBO-based variational inference does not provide a way to marginalize over nuisance latent variables &. For a model p(2. &, x) on latent variables {2, } and observed variables æ; inference 2 conditional on is common target of interest. Yet ELBO-based methods typically cannot compute the quantity p(2, x) = p(2, x)) xg required target the objective function and instead typically require explicitly modeling the nuisance variables.
2.2 NEURAL POSTERIOR ESTIMATION
In amortized variational inference, the shared parameters of a deep neural network define variational approximations for arbitrary observations (Ganguly et al. 2024 McNamara et al.] (202a)). Precisely, a variational approximation for latent variable 2 is given by q(;;(( 77) with parameters n7. Rather than fitting these parameters separately for each æ, the amortized approach sets 1) (x(x) for neural
2

---

Published as a conference paper at ICLR 2026
network fø with parameters 6. This inference network, once fit. yields the variational posterior g(2; fo (x) for arbitrary z by single forward pass (Kingma & Welling 2019 Ambrogioni et al.
2019p. Neural posterior estimation (NPE) (Papamakarios & Murray 2016 targets an expectation of the forward KI. divergence for amortized VI:
ÊNPE ()) = ppx) KL (p(2 x))
g(> x))
Here, the integral over p(r), the marginal of the model p(z, E, r), indicates that the NPE objective averages over all possible draws from the generative process instead of averaging over observations from some finite training set. The objective is equivalent (up to a constant) to
(o) - Ep(..r)) log fo()))
where go(2; = g(=; fo(r)) for any æ. Equation (3 admits unbiased estimation of its gradient with stochastic draws (z, x)) ~ p(z, x)) from the joint model, readily obtained by ancestral sampling, even in the presence of nuisance latent variables Ç. For example, in hierarchical model we can sample (z, x) ~ p(z, 2) by simulating an entire sequence (()) <)p(E2 (1)....PEE EL-1)p(( Ez) and discarding variables that are not the target for inference. The expected forward KIL objective has been independently derived and analyzed in several related works (Bornschein & Bengio 2015 Ambrogioni et al. 2019).
3. LBF-NPE: BASIS EXPANSIONS FOR AMORTIZED LOG DENSITY ESTIMATION
We propose amortized method to fit complex multimodal variational distributions, called Latent Basis Function NPE (LBF-NPE). For latent variables z in latent space I, the method fits a basis function network 8p - RK which evaluates K basis functions (1) ...., [ any point
The inference network fo x + rj E R-K maps observations to coefficients of the basis functions. The number of basis functions K is fixed ahead of time, and larger K may be used to increase expressivity (see Appendix E). The variational parameters to fit are the neural network weights ø and & for the networks la and 8g. Below, we first construct an exponential family defined by the basis functions (Section 3.1 and then give the fitting routine for fo and sy (Section 3.2
3.1 THE VARIATIONAL FAMILY
Fix x E X. Let C Rdd denote the latent space and let {s ()) (2)) be collection of basis functions
defined by parameters v. Selecting more basis functions (larger K) leads to a more expressive variational distribution, but also higher-dimensional optimization problem. For a value E I, let
(z) [
(K) (K)TT= R Our variational family is parameterized by coefficients 7) E Æ
and has the density function
((;))) x h(z) exp (1) sp sz))
We aim to select 7 and such that ((2;))) se p(z x). The log density is
log log v) = log h(()) + 1/ (=) C,
(5)
where is the log of the normalizing constant. This variational family (Equation 4 is exponential family /Wainwright & Jordan 2008 Srivastava et al., 2014). The vector of basis functions (2) is the sufficient statistic, 17 the natural parameter vector, and h22) is any finite base measure on the latent space I (i.e..) J, M=))dz is finite).
We represent the K -i-imensionnl quantity (z) as the output of a deep neural network with input Because the number and form of the basis functions (z) are arbitrary, the expressivity of this variational family is far greater than that of "classical" exponential families, such as the Gaussian family. As K + oc, the set of all exponential family distributions is arbitrarily rich: any distribution can be represented as an infiniteddimessionn 2020 Srrieeuubuuurr et al. 2017).
exponential family distribution (Khemakhem et al.
3

---

Published as a conference paper at ICLR 2026
3.2 THE AMORTIZED VARIATIONAL OBJECTIVE & GRADIENT ESTIMATOR
The formulation of Section 3.1 is non-amorlized: it requires selecting a single for g(=; v)) to approximate the posterior for single In NPE, we consider the amortized problem, where define the posterior for an arbitrary 2. We set = fr(r), and thus require fitting two separate networks, fo and sy. We fit the variational parameters, & and v, by minimizing the NPE objective function, as given in general in Section 2.2 Our specific variational objective (up to constants) is thus
LLBF-NPE(O, 4) -Ep(z.) (fo(z)]ss) (x) log (fesp (fo(z)) s (2))
(6)
where the log- log-normalizer takes the form the log of an integral with respect to the base measure h. The nngnoormllzzrr cannot be estimated without bias by Monte Carlo sampling due to the Jensen gap (Adil Khan et al., 2015). For training, we only require stochastic gradients, which can be computed using importance sampling. Focusing on the log-normalizer for now, we let ko.v (z, 2) = fo fz) (z). We suppress the dependence on x for now, as it is fixed in the integral. Let [á, v) = S exp ke.v (2)dh (3). The gradient can be computed by estimating an expectation with respect to an exponentially tilted transformation of h. Let Ve. denote the gradient with respect either & or v. Then, we have
1 V log (0) (6 Vonul(o v))
(7)
J(o. v) / (e)p
exp
(8)
/ V [orkka)() (
( (x')
dž
(9)
/ (V.kko (2)d
(10)
where we recognize an expectation with respect to ld.v, exponentially tilted density that depends on the current values of é, v.. This integral can be estimated by the use of self self-norma importance sampling (SNIS) with proposal distribution r(ž) Owen 2013):
V (e / Vorc go.u (E)d
E w((t)
where w(2) exp(ka. (()))A() and 2 21...., id The gradient estimator is thus biased for the true r(3) gradient, similar to other gradient estimators targeting the forward KL from the family of "wake-sleep"
algorithms, but consistent as P - oo (Le et al. 2019 Bornschein & Bengio 2015 McNamara et al. 2024b).
Algorithm details our gradient computation procedure.
4 VARIANTS & PROPERTIES OF LBF-NPE
We now motivate the construction of the LBF-NPE variational family by examining aspects of its optimization routine. As a result parameterizing and targeting the log-density (cf. [5] 7). both
construction and optimization routine depend entirely on the inner product fo(x)" sy (z). We elaborate on some key properties and variants LBF-NPE that stem from this observation.
4.1 AFFINE GRADIENTS
3.2 we showed that the gradient of the objective for LBF-NPE depends only on the gradient of (x, = fo(a)" ((())T(())) (z),, via the relation
In Section
V n. (OD xb) Ex(x.)) [V Vo. Eg.. IV No.khis (xr 2)]]
(11)
4

---

Published as a conference paper at ICLR 2026
Algorithm 1: Gradient Computation for LBF-NPE
Inputs: Sampling model p(2, z); networks fo and sys proposal distribution r(z).
sid
Sample batch {(2), 2,)); ~ p(2, x)
/* Gradient for log-normalizer +/
Propose 1 Ep
Compute k) = (1: 23) = fo(x;) (x))) E [B],j [
Compute unnormalized weights wi exp (Kt) h (2,), i [B].] [P]
Compute U... = - i Vo.
i [B]
Σn wy.
/* Gradient non-tilted inner product */
Compute Ve Vi a) = Va.
/* Compute comb mnd gradient */
Return Ve. - Σ
Vi )2.
where d denotes the exponentially tilted density constructed in Section 3.2 of gradient updates for the LBF-NPE procedure is extremely simple; in Tact, holding ve constant, the gradient with respect to o is that of an affine function of the network outputs fa. The same relationship holds when taking gradients for ş holding & constant. Targeting such simple functions for optimization, besides being simple to implement, benefits from convex formulation (see Section 4.2 below). The invariance of the inner product kg. under arbitrary rescalings of and 8, on the other hand, complicates optimization: this motivates a variant of LBF-NPE that reparameterizes outputs to unit norm (see Section
Accordingly, the form
4.2 CONVEXITY
McNamara al. (0224) show that neural posterior estimation (NPE) optimizes convex functional objective function provided that the variational family log-concave in f. the inference network. This ensures the forward KL objective of NPE (ef. Equation 3 is convex f. Recent advances in the study of wide networks via the neural tangent kernel (NTK) (Jacot et al. (018) have shown that fitting network parameters to minimize convex loss functionals (e.g., mean squared error) follows kernel gradient descent dynamics to global optimum in the infinite-width limit (Jacot et al. 2018 McNamara et al. 2024a)
The amortized forward KL objective function that we target (cf. Equation 6b benefits from these same properties. For an arbitrary collection basis functions, the objective function LLBENPE remains convex functional in f. Likewise, for fixed f, LIBEENPE is a convex functional in s. We formalize this the proposition below.
Proposition 1. The functional
L(f,*) = -Ep(e.)) (I(=) T(2)) log (foop s(3) M())d
is marginally convex in the arguments f and respectively.
A proof and additional discussion are provided in Appendix B Proposition 1 shows that in the case where either f or is fixed, the resulting objective function is fully convex in the remaining argument. This observation motivates variant of LBF-NPE where the basis functions are fixed priori (see Section 4.3
4.3 FIXED BASIS FUNCTIONS
Rather than adaptively fitting basis functions sp, the practitioner may simply use fixed basis defined ahead time. This approach motivated by the convexity of the resulting functional in f. as well as the approaches of related work based basis expansion parameterizations, which use fixed orthonormal eigenfunctions (cf. Section 5). In this variant of LBF-NPE, the objective function L( v)) collapses the marginal T.(2) for optimization. As elaborate in Section 4.2 LBF-NPE
5

---

Published as a conference paper at ICLR 2026
has convex formulation in this setting, which empirically results in advantageous optimization trajectories relative to competing methods (we demonstrate this in Section 6.1
Several choices of basis may be of interest to practitioners, Eigen VI, related basis vasis-expansion method for VI (cf. Scction 5). utilizes a (truncated) orthonormal basis of cigenfunctions, such as Bernstein, Legendre, Hermite polynomials (Cai et al. 2024). Selecting large K improves faithfulness the complete basis, but doing so increases the dimension of the optimization problem, exponentially so in multiple dimensions. Further, as generally such basis functions are global (i.e., nonzero on all of the latent space 1). in this design every basis function contributes to the density value ((z) at every point 2; this may make it difficult control the local behavior of the fitted density.
An alternative approach is to model log (2; = log ((; fa(r)) via local basis expansion. We specialize B-splines (Appendix A.I) and wavelets AAppendix[A2. in our experiments. two rich families that we recommend for practitioners. In this framework, each basis function is nonzero only in small neighborhood of the latent support. Locality of the basis functions simplifies optimization by inducing sparser problem than set of global basis functions would. For a single Monte Carlo draw (z" ,x*) ~ p(z, x). the gradient - v, log q(=";") In=fa(=*) is nonzero at only a few indices because many basis functions are zero at any given z".
4.4 REDUCING DEGENERACY TIRROUHH STEREOGRAPHIE PROJECTION
As noted in Section 4.1 both gradients and the log-density itself only depend on the inner product fa(x) (z). Adaptively learning both the inference network fo and basis function network sy thus suffers from an inherent lack of identifiability: different rescalings or rotations of the vectors defined by fo and Sy can lead to the same loss function values, since the loss function (Equation 6 only depends on the inner product. To mitigate this degeneracy, we propose variant of LBF-NPE that uses stereographic projection reparameterization to normalize the output tensor onto the unit hypersphere. This resolves the rescaling degeneracy, but rotational degeneracy persists even after the projection. Specifically, for K --dimensional coefficient or basis function vector, we construct neural network that outputs (K 1>)-dmmessonnal vector We then apply the stereographic
projection reparameterization: y =(I+1u) [+ 1+
2u
which maps u E R to a unit vector y E RK such that || 1. This normalization mitigates identifiability issues, and the reparameterization yields strong results in our experiments (see Ap- pendix for additional discussion). Our loss function takes the following form when we apply this reparameterization:
LLBENPELOO xb) -Bpt..)) (-wfa(t)Taa (z( + log ((ss(wussle) ( 35(3)) M(E))) (12)
where fo()) and (-) are reparameterized network outputs, and is a fixed scaling hyperparameter.
RELATED WORK
Exponential family distributions are common class of distributions for both traditional variational inference and NPE. In the simplest cases, Gaussian, Bernoulli, and other "simple" exponential families are used (Liu al. 2023 Cranmer al.. 2020 Blei al., 2017). Typically, however, these distributions are not parameterized in canonical form (where 1j To (r) is the natural parameter of the family). However, for NPE, McNamara et al. (2024a) recommend using the canonical parameteriza- tion, even for simple families such as the Gaussian, to benefit from convex loss (Section 4.2). General exponential families parameterized by neural networks were first proposed in iPaachiardii & Dutta (2022) to represent the likelihood function for likelihood-free settings. Akin to Approximate Bayesian Computation (ABC) methods, this approach aims to learn low-dimensional summary statistics to represent the likelihood, and subsequently performs inference with potentially expensive MCMC or ABC routines. We compare to this approach in Appendix E.6 To our knowledge, LBF-NPE is the first method to utilize neural exponential families to represent the posterior distribution and to use this family within the amortized inference setting. LBF-NPE is also unique in exploiting the low dimensionality of the posterior projections of interest.

---

Published as a conference paper at ICLR 2026
Parameterizing variational distributions via basis expansions is relatively line of research; recent non-amortized approach, EigenVi VI (Cai et al. 2024). presents an algorithm for optimizing a score-based divergence with variational family parameterized via linear combination of orthogonal eigenfunctions. Key limitations of this approach are i) the lack of amortization and ii) the necessity of utilizing orthogonal, fixed eigenfunctions the basis: truncation of these bases necessarily introduces approximation error. In our case, is unrestricted: the basis functions can be arbitrary, and so fixed number K may be sufficient for some classes of posteriors (cf. Section 6.2
Mixtures of Gaussians and normalizing flows are other common choices of variational families for NPE (Gershman et al. 2012, Papamakarios & Murray, 2016 Papamakarios et al., 2021; Rezende & Mohamed 2015). Although more flexible than simple exponential families, these parameteriza- tions may suffer from convergence to shallow local optima during optimization (cf. Section 6. We compare LBF-NPE to mixtures of Gaussians, normalizing flows, as well as Eigen VI in our experiments.
EXPERIMENTS
In numerical experiments, we fit variety of complex posterior distributions using LBF-NPE. In Sections 6. and 6.4 we infer one-dimensional posterior projections using the variant of our method with fixed basis functions, whereas in Sections 6.2 and 6.3 we infer two-dimensional posterior projections using adaptive basis functions. Additional details about each of these experiments appear Appendix D
NPE with various alternative variational families serves as our primary benchmarks; we can compare to these methods quantitatively by assessing the NPE objective with each choice of variational distribution. We benchmark NPE with variational families based mixture density networks
(MDNs), RealNVP, and neural spline flows (Papamakarios & Murray 2016 Durkan et al. 2019). In addition to the results in this section, additional results appear in Appendix E including results from comparisons to two non-NPP NEE--PPE-aaa based variational inference methods: Eigen VI (Appendix E.5 and a
score-matching neural- neural-likeliben www.-iikieee-bbbaeeeeeeee method for likelihood-free -rree inference (Appendix E.6
6.1
Toy EXAMPLE: SINUSOIDAL LIKELIHOOD
LBF-NPE (ours)
MDN
Negative Log Likelihood
0.6
Nogative Log Lilelihood
0.6
0.5
0.5
0.4
0.4
0.3
0.3
0.2
0.2
0
10k
20k
30k
40k
50k
0
10k
20k
30k
40k
50k
Training Step
Training Step
Figure Negative log likelihood of our method and the MDN, Each model is trained with 20 different random seeds, and the records are smoothed using Gaussian filter with 20.0.
We first exhibit the advantages of LBF-NPE's convex variational objective by demonstrating consistent convergence on a highly multimodal problem with fixed basis functions (cf. Section 4.3p. The model draws an angle
Unif(O, 2n] followed by z | N (sin(22), o²) for fixed a² 1. The exact posteriors p(z x)) have up to four modes, depending on the realization x. We compare LBF-NPE, using a fixed collection 14 B-spline functions degree two on mesh of [0, 2n (see Appendix A. for additional detail on B-splines), and mixture density network (MDN) (Bishop)
1994 Papamakarios & Murray 2016), using a mixture of five Gaussian distributions. Both variational distributions have 14 distributional parameters per observation (there four mixture weight parameters for the MDN due to the simplex constraint). Additional experimental details are given in Appendix D. Figure I shows that for 20 different of the optimization routine, L.B-NNEE
7

---

Published as a conference paper at ICLR 2026
Table 1: Forward KI. divergence (FKI.,, reverse KI. divergence (RKL., and negative log- log-likelihood (NLL) of LBF-NPE (ours), NSF (Neural Spline Flow), RealNVP, and MDN on three 2D test cases. Lower values indicate better posterior approximation.
Forward KL Divergence
LBF-NPE
NSF
RealNVP
MDN
Bands
0.0048 (+: 0.0003)
0.016 (# 0.003)
0.015 (#: 0.005)
0.182 (#: 0.01)
Ring
0.0054 (#: 0.0005)
0.017 (# 0.004)
0.024 (= 0.005)
0.205 (1: 0.02)
Spiral
0.187 (+t 0.004)
0.201 (# 0.01)
0.545 (#: 0.07)
0.948 (# 0.09)
Reverse KL Divergence
LBF-NPE
NSF
RealNVP
MDN
Bands
0.0014 (= 0,0004)
0.0099 (# 0.001)
0.01 (# 0.007)
0,156 (#: 0.02)
Ring
0.0027 (= 0.0003)
0.013 (# 0.003)
0.014 0.003)
0,204 (#: 0.01)
Spiral
0.188 (.L. 0.005)
0.322 (1) 0.04)
0.666 (1 0.09)
1.973 (1 0.14)
Negative I Log-Likelihood
LBF-NPE
NSF
RealNVP
MDN
Bands
-0.060 (:t: 0.07)
0.151
0.23)
0.157 0.22)
1.389 (:: 041)
Ring
0.030 (1: 0.03)
0.621 (1: 0.24)
0.733 (1 0.11)
1.031 (1 0.18)
Spiral
0.838 0.13)
0.727 0.25)
0.859 (t 0.32)
2.788 (# 0.31)
Truth LBF-NPE (ours)
NSF
RealNVP
MDN
Bands
Ring
Spirml
Figure 2: Example posteriors of three problems in two dimensions. NSF refers to Neural Spline Flow.
consistently converges to the same solution, whereas the MDN sometimes converges to suboptimal local optimum. Visualizations of posterior approximations from both methods are provided in Appendix 3
6.2 COMPLEX MULTIVARIATE REPRESENTATIONS IN 2D
We showcase LBF-NPE on three test problems in two dimensions, named "banded", "ring", and "spiral," and visualized in the left column of Figure 2 Each model consists of a two-dimensional latent variable E IR2 and observation E R, and in some cases nuisance latent variables as well. Further details of the generative processes for these three problems are provided in Appendix D (Sections D.2 to D.4]
LBF-NPE able to approximate these complex posteriors nearly perfectly using only 20 adaptive basis functions 8y. Both the amortization function fa and the basis functions 8 parameterized
8

---

Published as a conference paper at ICLR 2026
using deep neural networks. For additional implementation details, we refer the reader to Appendix D We follow Algorithm| and evaluate both the variational posterior q(2; (z), (w))),()))))) (2)) and the exact posterior p(= x))
fine mesh grid that covers the support of the posterior. Qualitative results appear in Figure 2 and quantitative results appear in Table T for held-out test points. LBF-NPE outperforms the MDN and multiple types of normalizing flow in nearly all cases. We provide additional visualizations of the variational approximations found by LBF-NPE and its competitors in Appendix B
6.3 OBJECT DETECTION
We apply LBF-NPE with adaptive basis functions to the problem of object detection in astronomical images.
We use a generative model resembling the scientific model of Liu et al. (0023). In brief, this generative pro- cess first independently samples star locations 41,6 Unif (10, 16] 00, 16]) and star fluxes fi, f2 ~ (n, (?) for two objects. Afterward, a latent noise-free image is rendered by convolving these point sources with a point- spread function (PSF). Finally, given I, the intensity of each pixel (j. k), for j, (I,...; 163 [ is independently drawn as Ijk ~ Poisson [sk), reflecting Poisson shot noise.
Figure 3 shows two examples of the noisy observations x, along with the posterior approximations for the loca- tions of each. The posterior distribution for this problem multimodal with a high degree separation between modes. LBF-NPE parameterizes this shape effectively. Inference on locations of objects would initially appear to method does not directly parameterize location parameters: instead, the learned basis functions must be expressive enough parameter- ize arbitrary pairs of separated modes. In our ablations, we vary the number of basis functions K 9, 20, 36, 64 for LBF-NPE and provide visualizations Appendix E.3 Examining the form of the fitted basis functions across varying K for this problem is particularly illustrative of the expressivity of the adaptive approach fitting the basis functions compared a fixed basis function framework.
Figure 3: Two example posteriors (left) conditional on the observed images (right). In each case, LBF-NPE correctly recovers the locations of the two objects.
be difficult for LBF-NPE.
6.4 CASE STUDY: REDSHIFT ESTIMATION
The redshift of galaxies is key quantity of interest as it characterizes their distances from Earth. Redshift measures the ex- tent to which electromagnetic waves are "stretched" to redder wavelengths as objects move away from Earth. The distribution of
Table 2: Total held-out NLL of the true redshift z.
LBF-NPE
NSF
MDN
NLL
-57,220 (= 152) -55,389 (1: 379)
-50,648 (4: 322)
redshifts across many objects a powerful probe of cosmology (Malz & Hogg 2022). Redshift estimation from photometric data (images) is referred as photo-2 estimation. We extend the methodology the Bayesian Light Source Separator (BLISS) package (Liu et al. 2023), state-of-the-ar package for probabilistic object detection in astronomical images, for this task, by adding a redshift density estimation "head" to the existing BLISS network. To each detected object, we associate redshift probability density function fitted by LBF-NPE with a fixed B-spline basis family. LSST DESC DC2 Simulated Sky Survey dataset (LSST Dark Energy Science Collaboration et al. 2021) serves the generative model, providing simulated (2,z) pairs, where 2 denotes redshift and x are the astronomical images. This highly realistic dataset consists of mock catalogs of astronomical images generated by directly modeling known physical quantities of the universe using empirical priors and physics-informet informed modeling. Appendix D provides further details of the experimental setting.

---

Published as a conference paper at ICLR 2026
We compare LBF-NPE with variational families based neural spline flow (NSF) and mixture density network (MDN), all embedded with the BLISS framework. The MDN uses five Gaussian components, in keeping with ofateoo--the-rr work on photo-2 estimation (Merz et al. 2025); adding more components did not improve the quality of fit. The only difference between the two approaches is the choice of parameterization for the variational family. We compute the negative log- lieliikl (NLL) of a held-out test set of 153, 30,000 astronomical objects. Table Shows that LBF-NPE with the B-spline variational family parameterization outperforms both the MDN and the NSF
DISCUSSION & LIMITATIONS
LBF-NPE models the log density of the variational distribution linear combination of expressive basis functions. which is beneficial for several reasons. First, log-space modeling results in multiplicative influence of different basis functions. Regions of latent space can effectively be "zeroed-out" more easily in this context compared to performing the modeling in density space directly. Second, our model of the log density results an unconstrained optimization problem in fú and sy: the coefficients and basis functions may be either positive or negative, whereas other density estimation methods may require nonnegativity or other constraints on the coefficients or basis functions to obtain valid density that integrates to unity (Cai et al. 2024] Koo & Kim 1996. Kirkby et al. 2023). Third, unlike other methods, LBB-NNEE may be applied to the growing
class of "likelihood frre inference" problems (Cranmer al. 2020 Thomas et al. (202) fitting the variational posteriors only requires samples from the joint model z) The advantageous marginalization properties of the forward KL criterion pair nicely with this setting: the generative
model can have arbitrary nuisance latent variables that are implicitly marginalized out (Ambrogioni al. 2019). We provide an additional example marginalizing over the parameters of Bayesian neural network (BNN) in Appendix E.8
Using basis expansions to parameterize variational distributions is recent and exciting innovation in variational inference. Relative to Eigen VI, LBF-NPE performs better with fewer basis functions (Appendix E.5) This may be due removing orthogonality constraints and adaptively fitting basis functions: by averaging across posteriors for arbitrary our approach implicitly regularizes the basis functions, preventing overfitting to any single instance. Appendix contains additional discussion and visualization of fitted basis functions Sy.
The main limitation of LBF-NPE is the difficulty of sampling from the variational distribution. LBF-NPE directly fits the log density of the variational distribution, but samples from this density are typically needed for inference. The low dimensionality of many NPE targets ensures that sampling straightforward: inverse transform sampling is readily applicable (cf. Appendix ( For higher- dimensional targets, importance sampling may be adequate to estimate functionals with respect to the variational distribution.
Despite this sampling challenge, our approach demonstrates that basis expansion methods offer a compelling middle ground between optimization simplicity and expressivity for NPE, Future work could explore new bilevel optimization approaches (Xiao & Chen (he5) to jointly learn adaptive basis functions alongside their coefficients. Additionally, extending our approach beyond low-dimensional targets to high nigh-dimensional targets with simplifying structure- suchttrre-uchhaae as known or assumed conditional independencies could broaden the applicability of basis expansion methods.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We thank the reviewers for their helpful comments and suggestions. This material is based on work supported by the National Science Foundation under Grant Nos. 2209720 (OAC) and 224 144 (DGE), and the U.S. Department of Energy, Office of Science, Office of High Energy Physics under Award Number DE-SC0023714.
REFERENCES
M. Adil Khan, G. Ali Khan, T. Ali, and A. Kilicman. On the refinement of Jensen's inequality, Applied Mathematics and Compulation, 262:128-1355 2015.
Luca Ambrogioni, Umut Güęlūü Julia Berezutskaya, Eva van den Borne, Yagmur Güęlütürk, Max Hinne, Eric Maris, and Marcel van Gerven. Forward amortized inference for likelihood-free variational marginalization. In International Conference on Artificial Intelligence and Statistics, 2019.
Jason Ansel, Edward Yang. Horace He, et al. PyTorch Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In ACM International Conference on Architectural Support for Programming Languages and Operating Systems, 2024.
Francis Bach. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research, 18(19):1-53 2017.
Yoshua Bengio, Nicolas Roux, Pascal Vincent, Olivier Delalleau, and Patrice Marcotte. Convex neural networks. In Conference on Neural Information Processing Systems, 2005.
Christopher M. Bishop. Mixture density networks. Technical Report NCRG/94/004, Aston University, Birmingham, UK, 1994.
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 1 112(518):859-88)) 2017.
Jörg Bornschein and Yoshua Bengio. Reweighted wake-sleep. In International Conference
Learning Representations, 2015.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Me and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
Diana Cai, Chirag Modi, Charles C. Margossian, Robert M. Gower, David M. Blei, and Lawrence K. Saul. VI: score-based variational inference with orthogonal function expansions. In Confer- ence on Neural Information Processing Systems, 2024.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. International Conference on Machine Learning. 2020.
Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation-bas based inference. Proceedings of the National Academy of Sciences, 1 117(48):30555-3002 2020.
Carl d. Boor. A Practical Guide Splines. Springer Verlag, New York, 1978.
Maximilian Dax, Stephen R Green, Jonathan Gair, Jakob Macke, Alessandra Buonanno, and Bernhard Schölkopf. Real-time gravitational wave science with neural posterior estimation. Physical review letters, 127(24):24110% 2021.
Jiankang Deng, Jia Guo, Niannan Xue, and Stefanos Zafeiriou. ArcFace:
Additive angular margin loss for deep face recognition. In Conference on Computer Vision and Pattern Recognition, 2019.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. In International Conference Learning Representations, 2017.
11

---

Published as a conference paper at ICLR 2026
Justin Domke. Provable smoothness guarantees for black-box variational inference. In International Conference on Machine Learning, 2020.
Justin Domke, Robert M. Gower, and Guillaume Garrigos. Provable convergence guarantees for
black-box variational inference. In Conference on Neural Information Processing Systems, 2023.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Conference Neural Information Processing Systems, 2019.
Paul H. C. Eilers and Brian D. Marx. Flexible smoothing with B-splines and penalties. Statistical Science, 1 11(2):89-221, 1996.
Ankush Ganguly, Sanjana Jain, and Ukrit Watchareeruetai. Amortized variational inference: A systematie review. Journal of Artificial Intelligence Research, 78:1-49, 1 2024.
Samuel Gershman, Matt Hoffman, and David Blei. Nonparametric variational inference. In Interna- fional Conference on Machine Learning, 2012.
Derek Hansen, Ismael Mendoza, Runjing Liu, Ziteng Pang, Zhe Zhao, Camille Avestruz, and Jeffrey Regier. Scalable Bayesian inference for detection and deblending in astronomical images. In ICML Workshop on Machine Learning for Astrophysics, 2022.
Elad Hoffer and Nir Ailon. Deep metric learning using triplet network. In Similarity-Baue -Based Pattern Recognition, 2015.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel:
Convergence and generalization in neural networks. In Conference on Neural Information Processing Systems, 2018.
Glenn Jocher et al. YOLOv5. /tDDDS//tttee//ggttuu.../////// http:///gg.tatuooooouuuaaaaaaaa/
2020.
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvärinen. Variational autoencoders and nonlinear ICA: A unifying framework. In International Conference on Artificial Intelligence and Statistics, 2020.
Patrick Kidger and Cristian Garcia. Equinox: neural networks JAX via callable PyTrees and filtered transformations. Disserentiable Programming Workshop at Neural Information Processing Systems, 2021.
Diederik P. Kingma and Max Welling, An introduction to variational autoencoders. Foundations and Trendy in Machine Learning, 12(4):307332,, 2019.
J. Lars Kirkby, Álvaro Leitao, and Duy Nguyen. Spline local basis methods for nonparametric density estimation. Statistics Surveys, 17:75-118, 2023.
Ja-Yong Koo and Woo-Chul Kim. Wavelet density estimation by approximation of log-densities. Statistics & Probability Letters, 26(3):271-278 1996.
Tuan Anh Le, Adam R. Kosiorek, N. Siddharth, Yee Whye Teh, and Frank Wood. Revisiting reweighted wake-sleep for models with stochastic control flow. In Conference on Uncertainty in Artificial Intelligence, 2019.
Jongmin Lee, Joo Young Choi, Ernest K Ryu, and Albert No. Neural tangent kernel analysis of deep narrow neural networks. In International Conference on Machine Learning. 2022.
Runjing Liu, Jon D. McAuliffe, Jeffrey Regier, and LSST Dark Energy Science Collaboration. Variational inference for deblending crowded starfields. Journal of Machine Learning Research, 24(179):1-36, 2023.
Weiyang Liu, Yandong Wen. Zhiding Yu. Ming Li, Bhiksha Raj, and Le Song. SphereFace: Deep hypersphere embedding for face recognition. In Conference on Computer Vision and Pattern Recognition, 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- on Learning Representations, 2019.
12

---

Published as a conference paper at ICLR 2026
LSST Dark Energy Science Collaboration, Bela Abolfathi, David Alonso, Robert Armstrong, Éric Aubourg, et al. The LSST DESC DC2 simulated sky survey. The Astrophysical Journal Supplement Series, 253(1):31, 2021.
LSST Dark Energy Science Collaboration, Bela Abolfathi, Robert Armstrong, Humna Awan, Yadu N. Babuji, et al. DESC DC2 data release note, 2022.
David J. C. MacKay. A practical Bayesian framework for backpropagation networks. Neural Computation, 4(3):448445 1992a.
J. C. MacKay. Bayesian interpolation. Neural Computation, 4(3):415 1992b.
Alex I. Malz and David W. Hogg. How to obtain the redshift distribution from probabilistic redshift estimates. The Astrophysical Journal, 928(2): 2022.
Declan McNamara, Jackson Loper, and Jeffrey Regier. Globally convergent variational inference. In Conference on Neural Information Processing Systems, 2024a.
Declan McNamara, Jackson Loper, and Jeffrey Regier. Sequential Monte Carlo for inclusive KL minimization in amortized variational inference. In International Conference on Artificial Intelligence and Statistics, 2024b.
Grant Merz, Xin Liu, Samuel Schmidt, Alex I. Malz, Tianqing Zhang, al. DeepDiSC-photoz: Deep dearning-based photometric redshift estimation for Rubin LSST. The Open Journal of Astrophysics, 8, 2025.
Radford M. Neal. Bayesian Learning for Neural Networks. Springer-Verlag, Berlin, Heidelberg. 1996.
Art B. Owen. Monte Carlo Theory, Methods and Examples, chapter 9, pp. 265-294. Stanford University, 2013. Chapter on Importance Sampling. Available at domains/me//
Lorenzo Pacchiardi and Ritabrata Dutta. Score matched neural exponential families for likelihood-free inference. Journal of Machine Learning Research, 23(38);1-71, 2022.
George Papamakarios and lain Murray. Fast e-free inference of simulation models with Bayesian conditional density estimation. In Conference on Neurul Information Processing Systems, 2016.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22((7);: 1-64, 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, et al. Py Torch: An imperative style, high- gigg-performance deep learning library. In Conference Neural Information Processing Systems, 2019.
Aakash Patel, Tianqing Zhang, Camille Avestruz, Jeffrey Regier, and The LSST Dark Energy Science Collaboration. Neural posterior estimation for cataloging astronomical images with spatially varying backgrounds and point spread functions. The Astronomical Journal, 170(3):155, 2025.
Danilo Rezende and Shakir Mohamed, Variational inference with normalizing flows. International Conference Machine Learning, 2015.
Kihyuk Sohn. Improved deep metric learning with multi-class N-pair loss objective. In Conference on Neural Information Processing Systems, 2016.
Yang Song. Jascha Sohl- Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations, 2021.
Bharath Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Aapo Hyvärinen, and Revant Kumar. Density estimation in infinite dimensional exponential families. Journal of Machine Learning Research, 18(57):1-59 2017.
13

---

Published as a conference paper at ICLR 2026
Manoj Kumar Srivastava, Abdul Hamid Khan, and Namita Srivastava. Statistical Inference: Theory of Estimation. PHI Learning, 2014.
J. Michael Steele. Stochastic calculus and financial applications. Springer, New York, 2010.
Owen Thomas, Ritabrata Dutta, Jukka Corander, Samuel Kaski, and Michael U. Gutmann. Likelihood- free inference by ratio estimation. Bayesian Analysis, 17(1):1-31, 2022.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. In European Conference Computer Vision, 2020.
Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 111-2):1-305, 2008.
Hao Wang, Yitong Wang, Zheng Zhou, Xing Ji, Dihong Gong, Jingchao Zhou, Zhifeng Li, and Wei Liu. CosFace: Large margin cosine loss for deep face recognition. IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2018.
Xun Wang, Xintong Han, Weilin Huang, Dengke Dong, and Matthew R. Scott. Multi-similarity loss with general pair weighting for deep metric learning. Conference on Computer Vision and Pattern Recognition, 2019.
Daniel Ward, Patrick Cannon, Mark Beaumont, Matteo Fasiolo, and Sebastian Schmon. Robust neural posterior estimation and statistical model criticism. Conference on Neural Information Processing Systems, 2022.
Stephan Wojtowytsch. On the convergence of gradient descent training for two-layer Rel.U-networks in the mean field regime, 2020.
Quan Xiao and Tianyi Chen. Unlocking global optimality in bilevel optimization: A pilot study. In International Conference on Learning Representations, 2025.
Mang Ye, Xu Zhang, Pong C. Yuen, and Shih-Fu Chang. Unsupervised embedding learning via in- variant and spreading instance feature. In Conference on Computer Vision and Pattern Recognition, 2019.
Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE Transactions on Pattern Analysis and Muchine Intelligence, 41 2019.
Dingyi Zhang, Yingming Li, and Zhongfei Zhang. Deep metric learning with spherical embedding. In Conference on Neural Information Processing Systems, 2020.
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Pa- pademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. In Conference on Newral Information Processing Systems, 2020.
14

---

Published as a conference paper at ICLR 2026
A B-SPLINE & WAVELET BASIS FUNCTIONS
We give examples of two classes of local basis functions defined on the unit interval [0, 1]: B-splines and wavelets. Without loss of generality, these definitions can be extended to construct the family on any interval [a, bj for a, b E I.. These families are potential candidates for a local basis parameterization of the variational posterior.
A.1 B-SPLINEES
For a choice degree and a uniformly spaced set of points {t+)== the B-splines are a set of local basis functions {b} (4) (2)] that are defined recursively follows (d. Boor 1978
Eilers &
Marx I996):
6. (2) = on [t,,, 1-11 o.w. for all 2, if
(13)
b. (d) (z)
2 t b(a-1)
(d-1)
4 4-2 2
1tt1
(ti+d - t;) titd11 ti t+dd1 + tetd - titd d 1.
(14)
The B-spline basis functions 6. (2) are thus individually piecewise polynomial splines of degree d. While each function 6!-) is symmetric about the ith knot t (or midpoint of two knots), it is nonzero for range of only 2d knots, in accordance with the locality aspect described above (see Figure 4
Basis Function Value
1.0
Basis Function Value
1.0
0.5
0.5
0.0
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Figure 4: Example visualizations of B-Spline basis functions of degree 1 (left) and (right). For brevity, show only subset of the basis functions.
A.2 WAVELETS
A collection of wavelet local basis functions on [0, 1] is defined relative to "mother wavelet" function denoted И. For case of exposition, we consider the Haar wavelet (Koo & Kim [996 Steele 2010) given by
1,0< -
((2)
(15)
0, otherwise
Thereafter, the set of local basis functions b; is defined recursively as follows: writing each i uniquely as 2 + k. 2 0,0 2", we have
6₂(2) 20 H2 k).
(16)
The local basis functions are thus defined as shifted and scaled versions of the mother wavelet H. One can check that b;(z) is nonzero only the interval [K 2-j, (k + 1) 2- for i 2 + k. so the basis functions become nonzero on increasingly local regions even for moderate values of (say, 200). Wavelets are often described as being local with respect both space and frequency as b, becomes increasingly "spiky" as well due the coefficient 2//2 (Steele 2010).
15

---

Published as a conference paper at ICLR 2026
B CONVEXITY
B.1 CONVEXITY & CONVERGENCE OF LBF-NPE WITH FIXED BASIS FUNCTIONS
As referenced in Section 4.2 in the setting where the basis functions s6 fixed ahead of time, the objective function of LB F-EE becomes a convex functional of the amortization network f. In this case, our setting can be shown be globally convergent under suitable regularity conditions on the network architecture, in the asymptotic limit as the network width tends arbitrarily large. We restate this result, proven in McNamara et al. (2024a). below.
Proposition. Let X S IR and y C RK Let fo = f(:6): :x:6):X :x - RK be parameterized as a scaled two-layer ReLU network of width p. i.e. [((1;)) - 1 the loss functional
Tb En-1aroo(z wj) for I,...: K. Define
x) g(2;1));
and allow parameters o {aiy w,},i 1- = 1,.... K = L....p to evolve via the gradient flow ODE ((t))).
= -v EoE ((t(t)) Then we have the following results (under regularity conditions (A6)-;
a) (((()) is precisely Eptx (x;; optimized by the gradient flow above. Further, the functional MLBF-NPE(F) : + (x, f(x)) a convex functional in f. with global optimum f*
b) For the parameterization above, with o following the gradient flow ODE, we have that there exists T >> such that
lim MIAFNNE(fTT) S MIEE + p-000
where fr f(;
The proposition above, proven in (McNamara et al. 2024a), states that gradient descent on the convex functional (x, (x;; σ)) converges arbitrarily close to its optimum in the infinite-width limit, relying on universality results of shallow networks and NTK theory (Jacot et al. 2018 Lee et al. 2022). Our parameterization of the variational distribution q as an exponential family in LBF-NPE falls into this setting when the basis functions are fixed, allowing us to directly apply results from MMNNamara et 2024a). We refer the proofs therein rather than restatement here.
Regularity conditions sufficient for the above to hold are provided below. They assume a bell-behaved functional M, a particular initialization of the width-p p network, and uniform boundedness conditions on gradients along the optimization trajectory. Although we emphasize that NTK-style results only approximately explain the success of our method or other neural posterior estimation methods in practice (the infinite-wiee width limit can only be approximated by finite networks, and the continuous gradient flow ODE is approximated by stochastic gradient descent), results like these prove that LBF-NPE benefits from advantageous optimization landscape asymptotically.
(A1) The data space X' is compact.
(A2) Weights are initialized as dij 0.5 0, wy ind
(O Ia).
(A3) The neural tangent kernel initialization, Ka(x, =')= ()))(=)))) at o= is dominated by some integrable random variable G, uniformly over x, x'.
(A4) The gradient V,,E(a, {x, v) is uniformly bounded for x, r).
(A5) The limiting NTK Kx limp- K3 is desitive-defintt (we note that the existence of the limit is guaranteed).
(A6) The functional MIBENFE is bounded below; and its minimizer fo has finite norm with respect to the RKHS norm of the limiting NTK K
(A7) The function f(x, 77) is C-smooth in y for some C'
16

---

Published as a conference paper at ICLR 2026
B.2 CONVEXITY OF LBF-NPE
In this subsection, we prove the convexity of the joint functional L(f, s) in Proposition | Neural network theory and NTK-based analysis of this objective in s are beyond the scope of this work. We prove marginal convexity in s, and appeal to previous NTK-based results (such as the above). which motivate our empirical results: optimization of convex functionals of neural network outputs is ad- vantageous compared to the optimization of nonconvex functionals due to the preferable optimization landscape of the former (Bach 2017 Bengio et al. 2005; Jacot et al. 2018 Wojtowytschi 2020).
We first present Hölder's inequality, which we'll use in the proof.
Lemma 1 (Hölder). If S is measurable subset of R (with respect to Lebesgue measure), and and g are measurable real-valued functions on S, then Hölder's inequality is
| [f())g(z)dd ((1SL))) dæ) (f dx)
for any p. satisfying 1 p =
Hölder's inequality will be used to prove Proposition restated below. Proposition 1. The functional
T(f,8) = -Ep(z.)) ((()TT(>)) log (fozp) (F(x))
s()) (xxx)(xxx((() (f(z) (()))))))))))))))(()))))))))())
is marginally convex in the arguments f and respectively.
Proof. Ignore the outer expectation for now, and consider fixed drawn from p(2, v). The inner product - (x) s(2) is clearly convex s. We now turn to the more complicated expression, log (f cxp (f())=s(( s(2") dz'). Note that the integral over z' this expression and doesn't 't depend on the realization 2. It still depends on s, however.
We prove this function convex follows. Let E (0, 1), and consider functions $1, $2. Then
(f exp (f(x) (asi (1-a)))(1)))))))
a)s2(z())) dh(z')
log
( exp (af(x) 81 exp ((1 (x) ( 62(21)) dhi(z))
log
(f (exp (f(x) T1 sz(2)) (exp (f(z)* s2(x\)) dh(z')
log
(fu(t))) ((())))(((())++)))))))))))
(1(2') ddh=')) dh(z))
log
where we have defined u(=') = (x)( (z')) and v(z') = exp(f(x) integral above, and apply Hölder's inequality with 1/p a and 1/g These sum to one as required. We take f to be [m(=))] and to be [=(=')
82()))) Consider the
Continuing, have
17

---

Published as a conference paper at ICLR 2026
(f [ulet)) [o(=')] dh(22)
=log
(
[this line repeated for clarity)
(I/ (Iu()))1o) a an(z?)
dh(22) )
[/
1/0-a)
1-a
s log
[Hölder]
( ((())ah(z))) +1 1a--a)oo log (ʃ (fote)ahle)
log
(ʃ exp(((()) s((2))d((22)) (1 a) (exp(f(a)) (esp(f(a)ssae"")ddcc) sa(z"))d(z'))
aooo
This all that is required to show that the mapping - log (Jexp exp (f s=2)) dz') is convex. As the sum of two convex functions is convex, we've shown convexity the integrand for any fixed draw 2, x 2),,)). To conclude, observe that by linearity of integration, this holds for the integral as well. The argument for f identical by symmetry of the inner product.
C SAMPLING
Similar VI (Cai et al. 2024), LBF-NPE does easily admit sampling from the fitted variational density. This is one limitation the nonparametrie nature the density model in both of these approaches variational inference.
In the low-dimensional case, sampling can be performed by inverse transform sampling. In this approach, one uses the cumulative distribution function Q of g, defined as
Q(=")= ( f ((2)dz
where ((2) is the fitted variational density (say, conditional on some datum x of interest). The function is invertible. Sampling be performed by drawing U ~ Unif(0, 1), and thereafter computing
Z ==
The result of this procedure is draw Z ~ ((). Computing and inverting cumulative distri- bution functions low dimensions is fairly straightforward. As outlined in the main text, this low-dimensional setting is one we commonly find to be of use to practitioners, especially for the types of problems we consider in our experiments.
To sample from high high-dimensional posteriors, several different approaches are available. One approach, as outlined (Cai et al. 2024), is sequential sampling, whereby one samples
((=1), ((2 21), q(23 21, 22), .9(2 21, 21...., Zd-1)
in order. Each individual density above can be sampled using the inverse transform sampling approach outlined above; conditioning within the exponential family parameterization is easily accomplished by freezing the already sampled indices and changing the variables of integration in the log integral. More generally, one could also utilize rejection sampling other Monte Carlo sampling algorithms. As the unnormalized variational density has an extremely simple form, i.e., q(2) x exp(1) b(2)). Markov chain Monte Carlo algorithms could also be efficient way to sample from the distribution defined by the fitted density. We present some selected results of Langevin sampling and inverse transform sampling Appendix C.I to illustrate the utility of these approaches.
C. RESULTS OF LANGEVIN DYNAMICS AND INVERSE TRANSFORM SAMPLING
We present the sampling results obtained via Langevin dynamics (Sogg al. 2021) and inverse transform sampling for the three 2D case studies: bands, z zing, and spiral. addition to the inverse transform sampling described in the preceding section, explore the of Langevin
18

---

Published as a conference paper at ICLR 2026
dynamics, widely adopted method for generating samples in score-based generative models. This approach iteratively updates a set particles according to:
dz = vv logp(zz x:dtt + 2edW,,
where 0.001, dW, (V 1), and (1.2..... 1000). Since model provides differen- tiable approximation of log p(z; z) through fa(x) (2), the gradient V. log x)) can be directly estimated. This enables to apply Langevin dynamics for posterior sampling.
For each of the bands, r ring, and spiral case studies, we generate 10,000 samples and visualize both the samples and their marginal densities in Figures [ to I0 illustrated in these figures, both Langevin dynamics and inverse transform sampling yield samples that closely match the estimated posterior distributions.
Truth
Estimated
LD; dim
LD; dim 2
1.3e-2
1.5e=2
1.0e-2
5.0e-3
5.0e-3
Samples (LD)
Samples (Inv)
Inv; dim
Inv; dim 2
1.5e22
1.5e-2
1.0e-2
1.0cx2
5.0w-9
5.0c-3
Figure 5: Sampling results for bands. "LD" and "Inv" denote Langevin dynamics and inverse trans- form sampling, respectively.
Figure 6: Marginal density of samples for bands. The green line indicates the estimated posterior density.
Truth
Estinated
LD; dim
LD; dim 2
1.0c-2
1.0c=2
0.0
Samples (LD)
Samples (Inv)
Inv; dim
Inv; dim 2
2.0c-2
2.0c-2
1 1.0e-22
1.0e22
Figure 7: Sampling results for ring. "LD" and "Inv" denote Langevin dynamics and inverse trans- form sampling, respectively.
Figure 8: Marginal density of samples for ring. The green line indicates the estimated posterior density.
19

---

Published as a conference paper at ICLR 2026
Truth
Estimated
LD; dim
LD; dim 2
8.0e-2
8.0e-2
4 4.0e-2
0.9
Samples (LD)
Samples (Inv)
Inv; dim
Inv; dim 2
8. 8.0e-2
8.0e-2
4. 4.0c-2
4,0r-2
Figure 9: Sampling results for spiral. "LD" Figure 10: Marginal density of samples for and "Inv" denote Langevin dynamics and inverse spiral. The green line indicates the estimated transform sampling, respectively.
posterior density.
D EXPERIMENTAL DETAILS
Code to reproduce results is provided at
com hwtm///icuDuunnua... ch/
LEF-NPE We use PyTorch (Paszke et al.] 2019 Ansel et al. 2024) and JAX (Bradbury al.
2018). We also use the equ nox library for deep learning in JAX (Kidger & Garcia 2021) l). For Section 6.4, the DC2 Simulated Sky Survey data is publicly available from the LSST Dark Energy Science Collaboration (LSST DESC) (LSST Dark Energy Science Collaboration et al. 2021 2022).
All experiments are conducted on an Ubuntu 22.04 server equipped with an NVIDIA RTX 2080 Ti
GPU.
D.1 TOY EXAMPLE: SINUSOIDAL LIKELIHIOO
In Section 6.1 we design simple Bayesian model to evaluate the convergence behavior of LBF-NPE and MDN. This statistical model is
Unif(O, 2m],
2 ((sin((2), o²),
where ²² 1. This model induces a multimodal posterior:
P( æ²) exp (-(sin(2z) - =))² =2)²((20)) I(0 S 2m),
which exhibits two modes when or 1,, and four modes otherwise.
For LBF-NPE, we construct multilayer perceptron (MLP) to predict the coefficient vector 1) fz(x). and the sufficient statisties 6(2) IR K are computed using B-spline basis functions with K 14. The MIP architecture consists of an input layer, four hidden layers, and an output layer, mapping x E IR. to E K Each hidden layer includes a full connection layer of 128 units, layer normalization, and ReLU activation. The output layer linear. The B-spline basis comprises 14 degree-2 basis functions, with knots at [0, O, linspace(O, 2n, num ), 2n,, 2T.. Although the B-spline basis 6(z) can be evaluated recursively
described in Section A. we precompute it grid to avoid redundant computation during training. Specifically, we pick 1000 uniformly spaced points in the interval [0. 2n and approximate the integral exp(11, b(z)) dz using the trapezoidal sum. For the term fa( b(2). we use the basis vector corresponding to the grid point closest to the true latent variable Zimee-
For MDN, a mixture of Gaussian distributions. The output vector has 10 parameters for means and variances, and 4 additional parameters for the mixture weights. The MDN objective
the same MIP architecture, with the output adapted to represent the parameters of
LMDN(T)
Ep((r,)) log q(2; t, (z)),
(17)
20

---

Published as a conference paper at ICLR 2026
where y the neural network parameters, (x) denotes the predicted distribution parameters, and ((;-)) is the corresponding density.
The training procedures for LBF-NPE and MDN are identical, except for the loss function. At each step, we sample 1024 latent-observed pairs (z, x)) from the generative model and update model parameters using the AdaBelief optimizer (Zhuang et al. (2020) with learning rate of 0.001. Training proceeds for 50,000 steps and completes within one hour for both methods. Peak GPU memory usage approximately 8300MB. We hold out 1000 (z,z:) pairs and track their negative log- llliikelihood (NLL) over training, as shown in Figure [] We apply Gaussian smoothing to the NLL curves with standard deviation o 20. This results in a smoothing kernel of size 161 4 20 x + with weights given by G; exp( -12 12/ )2o²)) for i E {-80..... 80}. With the normalization constants
for G; omitted, the smoothed NLL at step j computed as
NLLsmoothh
Σ
NLL C,.
D.2 2D CASE STUDIES: BANDS
The statistical model for the bands test case, introduced in Section 6.2 is
21, 22 ~ Unif[-1, 1],
(2),
x
(( - 22 02))
where 0² = 10-2, The resulting posterior forms two elongated bands the 2D latent space P(= v) ix exp (- 22 - x))))/)/////)) (20²)) I(-1 S 21; F2 S 1), with its maxima occurring along the lines where |2] 22| x.
As the latent variable is now wwo-dimensional, LBF-NPE encounters increased complexity due to the larger number of basis functions required. In our LBF-NPE framework, both the coefficient network Sa and the sufficient statistic network 89. are implemented as multilayer perceptrons (MLPs) with four hidden layers, each containing 128 units. All layers use layer normalization to stabilize optimization and ReL.U activations. The network So maps input x E TR to a coefficient vector in R while 86 maps z E TR22 sufficient statistics in K We set other 2D case studies, though even K = 2 suffices to capture the posterior structure in this example (see Appendix [E.3). The loss function for LBF-NPE follows Algorithm T where the integral term /exp(nn sy (2) dz is approximated using trapezoidal sum over a 100 x 100 uniform grid spanning 1|1 During training, we alternate between updating So and sy: we train Ss for 1000 steps while holding sy fixed, then train sy for 1000 steps with fo fixed, and repeat this process until the total training budget is exhausted. The choice of 1000 steps per phase is empirical; we observe diminishing returns in the loss reduction beyond 1000 steps, indicating that each sub-network has reached near-optimal solution given the other is fixed. In addition, we stereographie projection to reparameterize the output of fo and 89.
20 for consistency with
For the MDN baseline, we use MLP with the same architecture as fø. except that it outputs 50-dimensional vector representing the parameters of a mixture of 10 Gaussian components. Each component is parameterized by five values: two for the mean, two for the (diagonal) variance (assuming zero covariance), and one for the mixture weight. The loss function is identical to that described in Appendix D.I
For the normalizing flow baseline, we adapt the classic coupling flow from Diihh et al. 2017) to model the conditional posterior p(2 2). Each coupling layer includes translation and scaling subnetworks conditioned on E R. These sub-networks are implemented as MLPs, each taking as input the concatenation of the masked latent variable z and the conditioning variable x. Each MLP consists of a single hidden layer with 128 units. We use 10 coupling layers to ensure sufficient expressiveness. The resulting conditional density is given by:
q(= x) qN (x; x)) |det J/,
where (-) denotes the standard Gaussian density, h, æ)) is the transformed latent variable via the flow, and det. J is the product of the Jacobian determinants from each flow transformation.
21

---

II(-1 s Published as a conference paper at ICLR 2026
We train LBF-NPE, MDN, and the normalizing flow using the AdamW optimizer (Loshchilov & Hutter, 2019) with a learning rate of 10 for 50,000 steps. The batch size set to 1024, and training completes in approximately 2 hours. Maximum GPU memory usage is around 8400MB. For evaluation, we use a held-out set of 1000 (z, x)) pairs to compute the average forward and reverse KI. divergences. For each test observation x, LBF-NPE, MDN, and the normalizing flow estimate the density g z) over a 100 x 100 uniform grid on 112 These estimated posteriors are normalized such that their integral over the grid equals The true posterior p(z a) computed analytically over the same grid, enabling pointwise comparison. We then calculate the forward and reverse KL divergences between the estimated and true posteriors and report the average over all 1000 test cases in [ For the illustrative posterior plots shown in Figure 2 we fix 0.7 and visualize the estimated density g(= xr) from each method over the same 100 100 grid.
D.3 2D CASE STUDIES: RING
The statistical model for the ring case study in Section 6.2 defined as:
21 22 ~ Unif[- - 1],
(21, 22),
x ([[[==1²,
where 0² = -10- The resulting posterior, P(z x) x exp
21. $ 1), forms ring-shaped distribution in the latent space, with radius approximately vT.
The network architectures and training configurations used in this case are identical to those described in Appendix D.2 An example posterior q 0.7) is visualized in Figure 2
D.4 2D CASE STUDIES: SPIRAL
The spiral model defined as follows:
b ~ 00
Unif[0.0, (2x)
0 s, (d) bi
= r (Os
22 (0)
Z (#1, 22)
x ~ Nb,, x
where o² = 10-4, and 86 (0) = (0V1 + 8 + sinh (0)). The posterior is P(= x) o
exp -=)==))((200²) 1(0 < 0 < 2 0.10 <r < 0.50).
Most training settings follow those in Appendix D.2 except that increase the number of coupling layers in normalizing flow to 16. We observe minimal performance gain beyond this depth. For visualization Figure 2
we display the estimated posterior q(z x 0.35) over the 100 100 grid.
D.S ОВJECT DETECTION
We define the image generative model as follows:
4, l Unif([O, 16) [0, 16)).
f1, f2 NN,, 0²).
Image( (4, 12), {f1, f2), PSF),
Fj.k Poisson Iz,));
22

---

Published as a conference paper at ICLR 2026
where 2000, o² 4002, and Image(-) and PSF(-) defined below. Note that in our implementation, flux values are constrained to be positive.
Algorithm Image
Algorithm 3: PSF
Inputs: list of source locations {41, 12) list of source fluxes {f1 fz
Inputs: relative position matrix rl.. Compute d: (rl.....,00 0]2 rll!....
point-spread function PSF.
Initialize pixel location matrix pl (/1,22] 1, {2} do
for 44, ʃa in zip((1
Compute d, cxpp(d)/suum(xp((d) Return d,.
Compute relative location rly = pl -4 Compute PSF density de (rF()4)
Compute I; fi x d;
end
Compute / I1 12
Return I.
The pixel location matrix pl is a mesh grid of shape (H, W, 2) defined over [0.5, 1....... H 0.5] x [0.5, 1.5,.... W 0.5). where H and W the height and width of the image, respectively. Each source location L is a 2D) vector, and the term rl.,.., 0|2 + rl.[.., "al...,1" 1]2 is a matrix of shape (H, W). We use Jंrr 1. Each source flux f. scalar. Before passing the image to the network, normalize it using min-max scaling: I = (æ min(x)) (mmxx))/(max(() min(x)).
Since the input is a 16 16 image, we employ convolutional layer in our network reduce computational cost. The first layer of the model, fa. is a 2D convolutional layer with a kernel size of increasing the channel dimension from 3. This is followed by a 2D max pooling layer and ReLU activation. The output is then flattened and passed through four MLP layers, each with 128 hidden units, layer normalization, and ReI.U activations. Another model, 8p, is an MLP with four hidden layers, each with 128 units, layer normalization, and ReLU activations. The outputs of so and Sy are reparameterized via stereographic projection.
As each image contains two astronomical sources, compute the loss separately for each source. For the source located at 4, the first term in the loss is -wf (4). Only $ (4)) needs to be evaluated per source; shared terms such as (x) and the integral term can be reused across both. We approximate the integral using Monte Carlo integration with 22,500 random samples. The final loss is the sum of the losses for both sources. For alternating optimization, we train one of Sa or 8. for 300 steps at time (shorter than the 1000-step updates used in the 2D case studies; see Appendix D.2) since convergence is typically achieved more rapidly in this setting. Optimization is performed using the Adam W optimizer (Loshchilov & Hutter 2019) with a learning rate of 0.001. The total number of training steps is 45,000, with overall training time under two hours.
For posterior visualization, Appendix D.2 but evaluate the posterior over a 200 200 grid on the domain [0, 16] The estimated posterior for certain image is shown in Figure To generate this posterior, leverage model trained with K = 64. For results with other values of K (e.g., 9, 20, 36), we provide further discussion in Appendix E.3
adopt the same procedure in the previous 2D studies (see
D.6 REDSIIFT ESTIMATION
Our redshift experiment extends the methodology of the Bayesian Light Source Separator (BLISS) (Liu al. 2023 Hansen et al.) 2022 Patel et al. 2025). For given generative model of astronomical images and latent quantities (locations: fluxes: type of object; redshift; etc.), BLISS utilizes neural posterior estimation (Papamakarios & Murray] 2016) to perform amortized variational inference. The network architecture for BLISS operates files of images, returning distributional parameters for each object detected per tile. The architecture is thus convolutional, except for several additional image normalizations and other design choices suitable for astronomical image processing.
For samples of the generative model, we use images from two tracts of the LSST DESC DC2 Simulated Sky Survey (LSST Dark Energy Science Collaboration al. 2021 2022). numbers 3828
23

---

Published as a conference paper at ICLR 2026
and 3829. LBF-NPE does not sample the generative model on-the-fly fly in this setting, but only has access to a finite number of draws from the training sets.
We use the BLISS preprocessing routines to produce training, validation, and test image sets, along with ground-truth catalogs. Images, each of size 80 x 80, are processed in batches 64 by the BLISS inference network, which further splits these into x 4-pixel tiles. The network is fit the training set to minimize the forward KL divergence using a learning rate of 0.001. All nuisance latent variables are marginalized over, and we only score redshift variational posteriors; BLISS also allows easy addition of posteriors on other latent quantities to the computed NLL loss, should the user desire to perform inference on them.
We adapt the neural network architecture from BLISS for redshift estimation. The complete architec- tural details and parameter configurations are provided in Table As shown in the table, the input and output shapes of each layer are expressed as tuples, e.g., (bands, h, w) or (64, h, w) where bands denotes the number of bands in the input astronomical images. In the DC2 dataset, there are six bands: g. i, y. The variables h and represent the image's height and width, respectively, and are both set to 80 in our experiments. The model is composed primarily of three types of layers: Conv2DB1 lock, C3Block, and Upsample, A Conv2DBlock is composite module consisting of 2D convolution, group normalization, and SiLU activation function. The C3Block is adapted from the YOLOv5 architecture (Jocher al. 2020). It comprises three convo- lutional layers with kernel size and includes skip connections implemented via multiple bottleneck blocks (parameterized by n). The Upsampl 1 layer performs spatial upscaling of the input tensor by specified factor. The architecture follows a U-shaped design with four downsampling steps followed by two upsampling steps. To denote skip connections and input dependencies between layers, we use the "Input From" " column. For example, the entry "[-1, 9]" indicates that the current layer takes as input the concatenation of the outputs from the previous layer and layer 9. The final layer is a convolutional module with kernel size 1, producing an output of shape (n (n_coefff, h/ 4, w/ 4), where n_ is the number of predicted coefficients per tile.
The forward KI divergence framework prescribes that predictions are only scored for true objects. Accordingly, for each ground-truth redshift in the training set, we score the predicted NLL computed from the variational distribution for the 4x4 pixel tile containing that object. BLISS makes this transdimensional inference problem (a result priori) tractable by sharing parameters across objects within the same 4 x 4- 4-pixel tile, at the cost of bias introduced by this approximation. For both the MDN and B-spline parameterizations, we fit to the training data for 30 epochs and use the model weights with the lowest held-out NI.I the validation set to compute the test-set NLL.. Training the inference network Sa takes approximately 12 hours a single NVIDIA GeForce RTX 2080 Ti GPU. We note that, due the approximations involved in using a finite training set rather than true "simulated" draws, we can easily overfit to the training and validation sets, The procedure outlined above aims mitigate these issues to the extent possible.
the number of objects per tile being unknown
24

---

Published as a conference paper at ICLR 2026
Layer #
Input From
Input Shape
Layer Type
Config
Output Shape
in_chrbandar out_chhm6d
(banda, b, W)
Conv2DBlock
(64, h, MI
vernel_sfzed
in_ch-64; out_chh-4;;
2
4
(64, h, w1
Com-2DBlock
(64, h, MI
kecnel_sizer5
Sequence of Como2BBoock
in_sh-64) out__ch-64;
3
-
(64, hy #1
(64, h, wl
kecnel_siuer5; sequence_learg
in_sh=64; out_oh-64;
4
4
(64, h, x1
Cono2DBlock
(64, h/2, w/2)
kernel_sizem3, stridem2
in_oh=64 out_ch-64;
5
-1
(64, h/2, w/2)
C3Block
(64, h/2, w/2)
n=3
in_sh=61; out_chr128;
6
(64, h/2, w/2)
Conv2DBlock
(128, h/4,
kernel_size-3p stride-2
in_sh=128; oat_chh128;
(128, h/4, w/41
C3Block
(128, h/4,
in_ch=128; cut_cam2567
(128, h/4, w/4)
Cony2BBlock
(256, h/8, w/81
pernet_stze-3p stride=?
in_sh-256; eut_ch-256r
9
(256, h/8, w/81
C3Block
(256, h/8, w/81
n=3
an_sh-256) out_ch-512;
10
(256, h/8, w/81
Como2BBook
(512, h/16, w/16)
Recnel_sluer3, stcidem2
in_ch=512) out_ch-2567
(512, h/16, w/16)
C3Block
(256, h/16, w/16)
n=3
scale=2;
1
(256, h/16, w/16)
Upsample
(256, h/8, w/81
mederfaesrest"
an_ah-512; out__ch-25;;
[1.9]
(512,
C3Block
(256, h/8,
n-3
scalum2;
(256, h/8, w/81
Upsample
(256, h/4, w/41
"nearest"
in_oh=384; out_eh-256;
[-166
(384, h/4, w/4)
C3Block
(256, h/4, w/41
n=3
in_sh-256;
eut_cärr__ceffi
.1
(256, h/4, w/4)
Comv2D
(r_coeff, h/4, w/41
parnal_size-1
Table 3: Neural network architecture for redshift estimation.
25

---

Published as a conference paper at ICLR 2026
D.7 ANGULAR DISTANCE OPTIMIZATION
As discussed in Section 4.4 our method can be interpreted as performing angular-distance op- timization, but with a loss and gradient derived from a probabilistic space. This interpretation becomes evident we decouple the magnitude and directional components of the output tensors fa(z), (2) E ERKK through normalization techniques such as L2 normalization or stereographic pro- jection reparameterization. Angular distance optimization is a common objective in modern machine learning pipelines, improving performance and ensuring consistent alignment between training and testing metrics. Several widely-used loss functions, including the triplet loss (Hoffer & Ailon 2015).
N-pair loss (Sohn) 2016), and multi-simillrity loss (Wang et al., 2019). incorporate angular distance in their formulation. Cosine-based softmax loss functions are widely used for face recognition (Liu et al. 2017: Wang et al. 2018 Deng et al., 2019), and many contrastive learning algorithms (Chen et al. 2020 Tian et al. 22020 Ye et al.) 2019)) employ angular objectives to maximize the cosine similarity between embeddings from positive pairs.
Our variational objective in Equation
(6) is related to cosine-based softmax loss, whose general form
L = S.. + log exp(wSS...) + Σ cxp(wS..))
(18)
as described in Section 4.4 and suggesting more general form for our variational objective, namely
LLBE-NPE(O, v)- = Ex(x,x) fo + log (wxolaaoaa) bo d )=
(19)
where again fo(-) and (-) are normalized outputs of neural networks (i.e., with unit norm) and w is a scaling factor.
The key differences from the cosine-based softmax formulation are: (1) the summation is replaced by an integral over the continuous latent space, and (2) the angular distance is computed between coeffi- cient vectors and basis functions, rather than between learned embeddings. This connection offers two main ad advantages. First, our theoretical guarantees on convexity and convergence may extend to angular distance optimization problems, suggesting broader applicability. Second, our method can leverage off-the thel improvements developed for angular optimization, such as SEC (Zhang et al.) 2020), which regularizes gradient updates to stabilize and accelerate training. Given that even simple stereographic normalization already yields strong empirical results, we leave the integration of these enhancements to future work.
In our experiments, hypersphere. Precisely, u E R is transformed to R via
stereographic reparameterization to project the output tensor onto the unit
|u122) y 1+ | 1 12
2u
(20)
ensuring that ||y/|| = 1. This projection serves as a smooth and bijective transformation from Euclidean space IR -1 onto the K- 1-sphere 5 {a E E²x ||==|| = 13. Although this transformation changes the form of the variational objective in the neural network outputs, and thus violates some assumptions of our NTK framing from the perspective of convexity, strong empirical results suggest the benefits of reparameterization, and also the importance of future work in understanding the success of neural posterior estimation (NPE) techniques under arbitrary parameterizations. We hypothesize that the advantageous properties of this normalization stem from the smooth gradient trajectories and mapping to a compact space, discussed more detail below.
We illustrate the stereographic normalization process in 2D case, as shown in Figure 11 In this setting, a scalar input u E IR3 is projected onto a vector y E IRa lying on the I-sphere (i.e., the unit circle). For any given there exists unique line connecting the point (u, 0) and the north pole N = (0, 1). This line intersects the 1-sphere at a single point, which serves as the projection of u. By drawing vector from the origin to this intersection point, obtain unit vector y on the 1-sphere. Notably, the location of the intersection reflects the magnitude of u: if Ill > 1, the intersection lies on the upper half of the circle; if ||u|| < falls on the lower half.
26

---

Published as a conference paper at ICLR 2026
1''
Figure 1I: Visualization of stereographic projection 2D. scalar u is mapped to a point on the unit circle via intersection of the line connecting (u, 0) and the north pole N = (0, 1).
This reparameterization offers several advantages. First, the stereographic projection differentiable everywhere and provides well-behaved gradients throughout the domain. Second, the projection naturally enforces unit-norm constraints without requiring additional normalization layers or manual clipping, thus making training more stable and efficient.
E
ADDITIONAL. EXPERIMENTAL RESULTS
E. EFFECT OF NORMALIZATION
We compare the neural network's ' output basis functions with and without stereographic projection normalization, demonstrating that normalization helps the network learn clearer boundaries between regions of the parameter space. Figure 12 and Figure 13 show the values of the 20-dimensional basis functions (i.e.. (2)]s. where E {1, ......... 20)) evaluated over the plane E [-3, 3] x [-3,3] for the spiral case study. is evident that the model with normalization exhibits more distinguishable and structured partition boundaries in 2-space, while the model without normalization suffers from blurry transitions and over-exposure artifacts, as seen in Figure 13) This highlights key drawback of the non-normalized approach: it struggles to effectively disentangle the parameter space. Normalization also enhances interpretability by promoting better separation among basis functions. The estimated posterior density expressed as a weighted linear combination of these basis function densities. For given target spiral, the neural network increases the weights (i.e., [=)]) for dimensions whose corresponding basis functions have high overlap with the target density, and decreases weights for dimensions with low overlap.
27

---

Published as a conference paper at ICLR 2026
10
11
12
13
14
18
Basis Function Value
Figure 12:
Density plot of 20-dim basis function (w/ stereographic projection normalization) over plane z. Each subplot represents the density plot of a certain dimension.
3
10
16
17
19
20
Basis Function Value
Figure 13: Density plot of 20-dim basis function (w/o stercographic projection normalization) over plane Each subplot represents the density plot of certain dimension.
28

---

Published as a conference paper at ICLR 2026
E.2
POSTERIOR COMPARISON
Truth
LBF-NPE (ours)
NSF
Rc
MDN
Obs: 0.05
Obs: 0.26
Obs: 0.47
Obs: 0.68
Obs: 0.89
Obs: 1.11
Obs: 1.32
Obs: 1.53
Obs: 1.74
Obs: 1.95
Figure 14: Visual comparison of posterior densities estimated by different methods (LBF-NPE,
Normalizing Flow, MDN) against ground truth for ten representative observations.
29

---

Published as a conference paper at ICLR 2026
Truth
LBF-NPE (ours)
NSF
RealNVP
MDN
Obs: 0.05
Obs: 0.15
0
Obs: 0.25
Obs: 0.35
Obs: 0.45
Obs: 0.55
Obs: 0.65
Obs: 0.75
Obs: 0.85
Obs: 0.95
Figure 15: Visual comparison of posterior densities estimated by different methods (LBF-NPE,
Normalizing Flow, MDN) against ground truth for ten representative observations.
30

---

Published as a conference paper at ICLR 2026
Truth
LBF-NPE (ours)
NSF
RealNVP
MDN
N
Obs: 0.10
Obs: 0.14
C
Obs: 0.19
Obs: 0.23
Obs: 0.28
Obs: 0.32
Obs: 0.37
Obs: 0.41
Obs: 0.16
Obs: 0.50
Figure 16: Visual comparison of posterior densities estimated by different methods (LBF-NPE,
Comalliiigg Flow, MDN) against ground truth for ten representative observations.

---

Published as a conference paper at ICLR 2026
E.3 DIMENSIONS OF BASIS FUNCTIONS & FLEXIBILITY
The flexibility of our method is positively correlated with the dimensionality of the basis functions. We demonstrate this by analyzing both the forward and reverse KL divergences, well basis-function density plots, for models using basis functions of varying dimensions in the object detection case study. To quantify this, we compute the forward and reverse KL divergence between the estimated posterior distribution and target mixture of Gaussians defined by the true underlying locations 4 l2:
0.5 0.5.N(14, +
(21)
The results, shown in Table indicate that both forward and reverse KL divergences decrease as the dimensionality of the basis functions increases. For example, the 64-dimensional basis functions achieve the lowest forward KL divergence (1 (1.524), while the 9-dimensional basis functions result in the highest (3.397). A similar trend holds for reverse KI. divergence. However, the marginal gain from increasing dimensionality diminishes as the number of basis functions grows. Increasing from 9 20 dimensions yields significant improvement forward/reverse KL divergence (1.187/0.511),, but the improvement from 36 to 64 dimensions is relatively minor (0.246/0.28)) 2877). This suggests that, for a complex task like object detection, basis function dimensionality under 100 is sufficient to achieve near-optimal performance.
9-dim
20-dim
36-iim
64-dim
Forward KL Divergenee
3.397
2.210
1.770
1.524
A Forward KL Divergence
-1.187
-0.440
-0.246
Reverse KL Divergence
2,380
1.869
1.360
1.073
A Reverse KI. Divergence
-0.511
-0,509
-0.287
Table 4: Object detection: forward/reverse KL divergence for models of different basis function dimensions.
The basis function density plots provide further intuition for this trend. As shown in Figures 17 to 20 the 64- and 36-dimensional basis functions partition the z-space into fine mne-grained regions, capturing detailed structure. In contrast, 20- and 9-dimensional basis functions fail to do so, resulting in coarser approximations and reduced representational capacity.
Basis Function Value
Figure 17: Object detection: density plot of 9-dim basis function over plane Each subplot represents the density plot of certain dimension.
32

---

Published as a conference paper at (CLR 2026
10
11
12
13
144
1.5
19
20
Basis Function Value
Figure 18: Object detection: density plot of 20-dim basis function over plane
3
11
12
14
15
16
177
18
19
20
Basis Function Value
Figure 19: Object detection: density plot of 36-dim basis function over plane For brevity, we only show the first 20 dimensions.
33

---

Published as a conference paper at ICLR 2026
10
11
13
14
15
1.5
-1.0
-0.5
0.0
0.5
1.0
1.5
Basis Function
Value
Figure 20: Object detection: density plot of 64-dim basis function over plane For brevity, we only show the first 20 dimensions.
9-dim
20-dim
36-dim
64-dim
Ohs Image
Obs:
Obs: 2
Obst 3
Obs:
Figure 21: Object detection: 9/20/36/64-dim basis functions and the corresponding estimated posterior density.
34

---

Published as a conference paper at ICLR 2026
Interestingly, in the ring study, observe that even 2-dimensional basis function is sufficient to recover the ring-shaped posterior. As shown Figure 23 the estimated posterior using a 2- dimensional basis function is visually nearly indistinguishable from the true posterior, with only minor artifacts appearing when the observation x approaches extreme values (e.g., x = 1.95). This observation is quantitatively supported the KI. divergence results in Table where both forward and reverse KL values are low (0.032/0.031) for the 2-dimensional case. These results demonstrate that for simpler posterior structures, our method achieve accurate inference with remarkably Low-dimensiee basis functions.
2-dim
4-dim
9-dim
20-dim
Forward Divergence
0.032
0.0057
0.0056
0.0054
Forward KL Divergence
-.0063
-0.0001
-0.0002
Reverse KI. Divergence
0.031
0.0032
0.0028
0.0027
Reverse KL Divergence
-0.0278
-0.0004
-0.0001
Table 5: Ring: forward/reverse KL divergence for models of different basis function dimensions.
2
4
2
2
Basis Function Value
Basis
Function
Value
Figure 22: Ring: density plot of 2-dim basis func- tion over plane z. Each subplot represents the density plot of certain dimension.
Figure 23: Ring: density plot of 4-dim basis func- tion over plane z.
4
2
Basis Function Value
Figure 24: Ring: density plot of 9-dim basis function over plane z.
35

---

Published as a conference paper at ICLR 2026
2
5
7
11
12
13
14
1.6
17
18
19
20
-2
2
Basis Function Value
Figure 25; Ring: density plot of 20-dim basis function over plane
2-dim
4-dim
9-dim
20-dim
Truth
Obs: 0.05
Obs: 0.68
Obs: 1.32
Obs: 1.95
Figure 26: Ring: 2/4/9/20-dim basis functions and the corresponding estimated posterior density.
36

---

Published as a conference paper at ICLR 2026
E.4 TOWARDS HIGH DIMENSION
To evaluate the capability of our model in predicting high-dimensional posteriors, we construct 50-dimensional annulus model defined as:
21, 22,
250 ~ Unif 1],
(21,...; 250).
2 N(11=|?,,
where x² -10
In Figures 27 and 28 we randomly select two pairs of dimensions (3, 14) and (43, 47). and visualize
the estimated posterior over these subspaces, i.e., g(23, 214 2)) and ((243 247 x). We discretize each pair onto 100 x 100 grid and perform Monte Carlo integration over the remaining 48 dimensions to obtain an estimate of the posterior in the chosen 2D subspace. The results show that the estimated posteriors closely match the true posteriors, with minor discrepancies likely due to Monte Carlo integration variance. The final two columns cach plot show the marginal densities, demonstrating that our model captures the true marginal posterior distributions. This good performance also be verified by quantitative metrics. Our model attains 0.018 forward KI. divergence and 0.022 reverse KL divergence, average, across 50 dimensions.
Truth
Estimated
Density (dim 3) Density (dim 14)
Obs: 13.9
Obs: 16.2
Obs: 18.1
Obs: 22. 7
Figure 27: True and estimated posterior density over dimensions 3 and 14. The y-axis the last two columns represents the marginal density.
37

---

Published as a conference paper at ICLR 2026
Truth
Estimated
Density (dim 43) Density (dim 47)
Obe: 13.9
Obe: 16.2
Obs: 18.1
Obs: 22.
Figure 28: True and estimated posterior density over dimensions 43 and 47. The y-axis in the last two columns represents the marginal density.
38

---

Published as a conference paper at ICLR 2026
E.5 VI ON 2D CASE STUDIES
We evaluate Eigen VI (Cai et al. 2024) on three two-dimensional targets with thin or curved posterior density patterns. We use tensor product expansion with K 16 basis functions per axis (K2 256 coefficients) and 50, 000 importance samples for fitting. The reconstructions capture only coarse structure: for the diagonal bands, EigenVI recovers orientation, but the two ridges are blurry; for the ring. it fills the central hole and collapses mass inward; for the spiral, it loses the manifold and yields blurry lobes. These failures arise from the spectral bias of orthogonal expansions, which under-represent the figh-freyuency content required by thin strongly curved posterior. While increasing K can help, computational and statistical costs scale as Kd (here 2), making adequate resolution impractical. In sum, with practical K, Eigen is adequate for smooth densities but inadequate for multimodal or topologically nontrivial wwo-dimensional targets.
Truth
LBF-NPE (ours)
Eigen VI
Figure 29: Eigen results for three 2D test cases. Along each axis, we fit 16 basis functions (i.e.,
K 16 in their original paper); For their importance sampling, we draw 50,000 samples.
39

---

Published as a conference paper at ICLR 2026
E.6 SCORE MATCHED NEURAL EXPONENTIAL FAMILIES ON 2D CASE STUDIES
We reproduce the method proposed in Score Matched Neural Exponential Families for Likelihood- Free Inference (Pacchiardi & Dutta 2022) and evaluate its sampling quality three two-dimensional case studies. Their approach estimates the unnormalized probability p(a =), rather than our focus on p(2 .)) and employs Exchange MCMC to draw posterior samples. As illustrated in the figure below, the Exchange MCMC samples are suboptimal, often overdispersed and misaligned with the true density. In the bands case study, the samples fail to align with the ridges and instead spread into low-density regions. In the ring case study, many samples are scattered inside the ring rather than concentrating on its boundary, where the density peaks. In the spiral case study, the bias is most evident, with samples deviating substantially from the high-density spiral structure. We also apply inverse transform sampling, which was not considered in the original paper. This approach produces samples that more closely follow the true density across all three case studies, though in the spiral example, a residual bias remains.
Truth
LBF-NPE
SMNEF
SMNEF ITS
Figure 30: Sampling results of Score Matched Neural Exponential Families (SMNEF) for Likelihood- Free Inference on three 2D case studies: SMNEF uses Exchange MCMC as its default sampling setting. SMNEF+ITS is our variant of SMNEF, customized for low-dimensional settings. For our method, we use inverse transform sampling.
40

---

Published as a conference paper at ICLR 2026
E.7 COMPUTATIONAL COST
the computational efficiency of several baseline models with the proposed LBF- NPE across the 2D case studies, object detection, and redshift estimation. While LBF-NPE shows slightly higher per-step runtime and memory usage than some baselines, its principal advantage is its substantially faster convergence, requiring markedly fewer training steps, such as 8k in the 2D case and 48k in the redshift task. This accelerated convergence results in competitive or superior total training time across tasks, indicating that LBF-NPE, achieves an effective optimization trajectory. notable detail is that, in the redshift experiments, the computational costs of all three methods appear very similar. This is largely due to the dominant overhead from the convolutional U-shape network used for image processing, which outweighs differences in the loss computation. Another point of clarification is that the GPU memory usage reported here is lower than the values in Section D.2 (e.g.. the peak usage is approximately 8400 MB). This discrepancy arises because, for the computational cost evaluation, we disable GPU memory preallocation in JAX, which otherwise reserves roughly 75% of the available GPU memory.
Time per step (s/step)
GPU Memory (MB)
Converge (steps)
Case study
Method
Time until converge (s)
Sinusoidal likelihood
LBF-NPE
0.038
312
4k
152
(batch size: 1024)
MDN
0.021
248
10k
210
LBF-NPE
0.127
970
8k
1016
studies (batch size: 1024)
NSF
0.084
2436
RealNVP
0.082
25k
2050
MDN
0.044
Object detection
LBF-NPE
0.143
2230
15k
2145
(batch size: 1024)
LBF-NPE
0.28
7319
48k
13440
Redshift
NSF 0.28
7012
80k
22400
(batch size: 32)
MDN
0.26
6988
54k
14040
Table 6: Computational cost: For 2D case studies, we only report the computational cost for the spiral case study, because the other two case studies have similar computational costs. The "Converge at (steps)" refers the maximum number of training steps required to achieve the performance reported in the paper. NSF is the abbreviation of Neural Spline Flow.
41

---

Published as a conference paper at ICLR 2026
E.8 BAYESIAN NEURAL NETWORKS AS GENERATIVE MODELS
A Bayesian neural network (BNN) is a neural network in which each weight (and bias) is treated probability distribution rather than fixed value (MacKay 1992a by Neal 1996). When making predictions, it marginalizes over these distributions to produce not just a prediction, but also an estimate of uncertainty. We consider BNNs as generative models illustrate that LBF-NPE can perform posterior predictive inference while implicitly marginalizing over a high-dimensional parameter space. particular, we consider two-layer fully connected BNN:
(N) Lineare, Lineare (Lineaa (z))), 01,02 N/O, I),
(22)
where 0, and 02 are the weight matrices of the first and second linear layers, respectively, and ReLU(-) denotes the rectified linear unit activation. The network takes a one-dimensional input and produces one-dimensional output through hidden layer of width 16. Throughout this experiment we restrict attention to x E [0, 10] and y E [-
Let f (81, (2) and let the dataset be
D { 1)))=
V: = (N:)(r:)) + +G, n n=
for a fixed draw of from (N I), and noise e, ~ /O0 1). The posterior predictive distribution for pair (w', v/) (')rr is
pl [) = plu a (4 p(0 d0,
(23)
which involves integration over the high- nigh-dimensional weight vector 0. Our goal in this section is to show that LBF-NPE can approximate the conditional density in 23 without ever explicitly sampling or optimizing over 0.
network 8o and coefficient network fa- The coefficient network fo takes as input the query variate x' together with the conditioning set D, while the basis-function network 8 takes y' as input. Together they define
In our implementation, LBF-NPE parameterizes the conditional density via a
LBF-NPE(y x', D) = p(y 2D,
Both and fø are implemented as four-layer multilayer perceptrons with 128 hidden units per layer; the resulting basis-function and coefficient vectors have dimension 20. We optimize the parameters (tb, Φ) with the AdamW optimizer (Loshchilov & Hutter 2019), using a learning rate of 10 -3 for 10,000 gradient steps.
and are trained using Algorithm I
The qualitative behavior of the learned posterior predictive distributions is shown in Figure 31 Across the four panels, the underlying BNN functions (orange curves) exhibit markedly different slopes, curvatures, and ranges, yet LBF-NPE recovers their overall shape from only Moreover, the posterior highest-denstty interval (HDI) is wide in regions with sparse observations and narrow in regions with many data points. For example, when observation is available near x E [0, 2] (top-left panel), the HDI is wide, whereas in regions densely populated with observations (e.g., E [2.5, 5.0] in the bottom-left panel) the HDI becomes narrow. This pattern indicates that LBF-NPE successfully captures the epistemic uncertainty induced by marginalization over the BNN weights.
5 observations.
42

---

Published as a conference paper at ICLR 2026
8
1
4
Prediction
90% HDI BNN(x)
-4
-4
Observations
-8
0.0
2.5
5.0
7.5
10.0
0.0
2.5
5.0
7.5
10.0
X
X
8
8
4
0
0
-4
-4
8
0.0
2.5
5.0
7.5
10.0
0.0
2.5
5.0
7.5
10.0
X
X
Figure 31: Posterior predictive distributions obtained from LBF-NPE on four synthetic regression problems generated by Bayesian neural network (BNN). Each panel corresponds to different draw of BNN weights. The orange curve is the true mapping x + BNN+(x), and the red dots denote the n 5 observed data points used for inference. The solid purple curve shows the pointwise mode of the posterior predictive density produced by LBF-NPE, and the shaded region indicates the associated 90% nighest-densig interval (HDI).
43


Reviews:
[{"review_id": "Wy5qYy17hL", "reviewer": "Reviewer_ZrDB", "summary": "This paper proposes a variational family LBF-NPE that is specialized for neural posterior estimation, which is a simulation-based (i.e., likelihood-free) inference method of the Bayesian posterior. Specifically, the variational family is an exponential family where the natural parameter $\\eta$ is given by a function $f_{\\phi}(x)$, providing amortization across observations $x$, and the sufficient statistic is a vector of $K$ basis functions $s_{\\psi}(z)$. This means the optimization depends only on the inner product between $f_{\\phi}(x)$ and $s_{\\psi}(z)$, leading to several properties, including affine gradients, convexity in $f$ and $s$, and the ability to fix the basis functions. Redundancy between $\\phi$ and $\\psi$ can be mitigated by a stereographic projection approach. The paper provides four experiments on four datasets, two are low dimensional toys and two are more realistic datasets related to astronomy, showing improvements in NLL and/or KL. The paper includes detailed explanations of the experiments in the appendix.", "questions": "- Regarding sections 4.1 and 4.2, can you comment on the importance of convexity in $f$ and $s$ given that they are potentially nonlinear functions of $\\phi$ and $\\psi$, respectively? For example, if you fix the basis functions, then the objective in Eq 6 may still be nonconvex in $\\phi$ (e.g., if $f_\\phi$ is a neural network), is that correct? If so, can you expand on how being convex in $f_\\phi$  but not $\\phi$ is still helpful? \n- In the Discussion & Limitations section, you write that regions of the latent space can be “zeroed-out”. I understand this to mean that some of the coefficients $\\eta = f_\\phi$ are zero after training. Can you comment on why this is useful and which experiments show this behavior? It’s interesting this can happen without anything encouraging sparsity.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Overall this is a well-written and thorough paper that provides theoretical and empirical backing for the proposed method. Given recent interest in basis-expansion methods in VI (e.g., eigenVI), I think it provides a nice contribution. It’s interesting that LBF-NPE performs better than eigenVI with fewer basis functions. I appreciate the attempt at a middle-ground between, say, mixtures of Gaussians and normalizing flows. This gives the user a few levers (e.g., fixing the basis) to ease the complexity of the optimization problem. The experiments showed improvements over recent baselines.", "weaknesses": "- The paper aims to strike a balance between simplicity and expressivity in the variational family. However, since the amortization $f_\\phi$ is a neural network in parameters $\\phi$, one concern is that it doesn’t make the optimization problem that much easier (see my question below regarding the importance of convexity in the function $f$). \n- As the authors point out, sampling from the variational distribution in LBF-NPE is challenging.\n- The results of the object detection experiment are buried in the appendix. As I understand it, the takeaway is that more basis functions reduces the KL to the chosen target but at a diminishing rate.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DB8GRdxOHn", "reviewer": "Reviewer_L4UB", "summary": "This paper introduces LBF-NPE, a neural posterior estimation method that parameterizes variational distributions using latent basis expansions. The method models log-densities as linear combinations of basis functions, forming an exponential family that balances expressivity with optimization stability. The authors propose several variants, including fixed basis functions and adaptive basis learning with stereographic projection for identifiability.", "questions": "* How to determine the number of samples and the number of basis $K$, and how sensitive is the performance to them?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Novel variational family that bridges the gap between simple distributions and complex black-box models.\n* Strong theoretical foundation with convexity guarantees in certain settings.\n* Comprehensive experiments across synthetic and real-world scientific problems.\n* Outperforms MDNs and other baselines on multiple benchmarks.\n* Effective handling of multimodal and complex posterior geometries, as demonstrated by the experiments.", "weaknesses": "* There seems to be no discussion of computational cost compared to baselines.\n* If experiments are repeated several times with different random seeds, error bars should be reported in the plot and tables.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k8ojH4xn0D", "reviewer": "Reviewer_hPZ3", "summary": "The paper proposes a new variational family for Neural Posterior Estimation (NPE) that parameterizes the log-density of the variational distribution through basis expansions of latent variables. Instead of the traditional mixture of Gaussian or flow-based families, the authors parametrize the log density as a linear combination of latent basis functions, with coefficients modeled by an inference network and basis function network either fixed (e.g., B-splines, wavelets) or learned adaptively. The authors have claimed convexity and stable optimization in the fixed-basis case, and flexibility and expressivity comparable to normalizing flows in the adaptively learning case. The proposed method shows strong empirical results across both synthetic and real-world applications, consistently outperforming baselines such as MDN, RealNVP, and NSF on KL divergence and NLL.", "questions": "1. Could the authors elaborate on whether LBF-NPE can be extended to high-dimensional latent spaces (e.g., posteriors of Bayesian neural networks) without prohibitive integration cost? Given time constraints, you can illustrate this only using a small toy regression experiment.\n2. How does the computational efficiency of the proposed LBF-NPE compare to other baselines? Can you provide inference time comparisons?\n3. The current gradient estimator is biased but consistent. Have you quantified this bias, or explored control variates or reparameterization tricks to reduce it?\n4. While Proposition 1 ensures convexity when either $s_\\psi$ or $f_\\phi$ is fixed, what guarantees or empirical evidence exist when both networks are trained jointly? Have the authors observed non-convex behavior or mode collapse in practice?\n5. Can you provide quantitative results comparing fixed (B-spline/wavelet) versus adaptive bases in terms of convergence speed, stability, and overfitting? How do you decide whether to use B-splines, wavelets, or neural bases across different tasks? Is there an automatic basis selection mechanism?\n6. The choice of the scaling factor $w$ is ad-hoc. Is it tuned per task, or could it be learned automatically? \n7. Section 7 and Appendix C mention inverse-transform and Langevin sampling. How efficient are these in high-dimensional cases, and have you compared them to Hamiltonian Monte Carlo or other samplers for posterior recovery?\n8. Since each latent basis has an interpretable contribution to the posterior, can you comment on whether the learned bases correspond to meaningful latent modes (especially in the astrophysical tasks)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The latent basis function formulation for the variational family is original, offering a middle ground between the simple Gaussian family and black-box flows.\n2. Convexity and convergence properties are explicitly proven (App. B) and related to NTK theory, a rare level of rigor for NPE papers. The fixed-basis variant provides stable inference under certain conditions where global convexity is crucial.\n3. Benchmarks cover synthetic, multimodal, and real astrophysical datasets, including redshift estimation (DC2 dataset). The proposed LBF-NPE achieves lower forward/reverse KL and NLL than baselines.", "weaknesses": "1. Most experiments are in ≤ 2D latent spaces. The high-dimensional example (App. E.4) is also synthetic (50-dimensional); stronger evidence on real higher-dimensional tasks would strengthen the claim of scalability.\n2. While the impact of basis dimensionality and normalization is explored (App. E.1–E.3), ablations on the effect of stereographic projection and alternating optimization are somewhat fragmented.\n3. Given that both the basis function $s_\\psi$ and inference network $f_\\phi$ are parameterized as deep neural networks, computational efficiency remains a concern for this algorithm.\n4. Although Algorithm 1 provides a gradient estimator for the objective function, it lacks pseudocode for the entire algorithm. Including this would clarify the overall process.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)