Bridging the Gap between Variational Inference and Stochastic Gradient MCMC in Function Space
Authors: Mengjing Wu
Junyu Xuan
Jie Lu
Venue: ICLR 2025 Poster (2025)
Keywords: Bayesian deep learning
functional variational inference
functional MCMC
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Abstract: Traditional parameter-space posterior inference for Bayesian neural networks faces several challenges, such as the difficulty in specifying meaningful prior, the potential pathologies in deep models and the intractability for multi-modal posterior. To address these issues, functional variational inference (fVI) and functional Markov Chain Monte Carlo (fMCMC) are two recently emerged Bayesian inference schemes that perform posterior inference directly in function space by incorporating more informative functional priors. Similar to their parameter-space counterparts, fVI and fMCMC have their own strengths and weaknesses. For instance, fVI is computationally efficient but imposes strong distributional assumptions, while fMCMC is asymptotically exact but suffers from slow mixing in high dimensions. To inherit the complementary benefits of both schemes, this work proposes a novel hybrid inference method for functional posterior inference. Specifically, it combines fVI and fMCMC successively by an elaborate linking mechanism to form an alternating approximation process. We also provide theoretical justification for the soundness of such a hybrid inference through the lens of Wasserstein gradient flows in the function space. We evaluate our method on several benchmark tasks and observe improvements in both predictive accuracy and uncertainty quantification compared to parameter/function-space VI and MCMC.

Published as a conference paper at ICLR 2025
BRIDGING THE GAP BETWEEN VARIATIONAL INFER- ENCE AND STOCHASTIC GRADIENT MCMC IN FUNC- TION SPACE
Mengjing Wu, Junyu Xuan, Jie Lu Australian Artificial Intelligence Institute, University of Technology Sydney, Australia mengjing, wu8 st udent. uus. edu.au; {junyu. xuan, jie. @ut s. edu au
ABSTRACT
Traditional parameter-space posterior inference for Bayesian neural networks faces several challenges, such the difficulty in specifying meaningful prior, the potential pathologies in deep models and the intractability for multi-modal posterior. To address these issues, functional variational inference (fVI) and fune- tional Markov Chain Monte Carlo (fMCMC) are two recently emerged Bayesian inference schemes that perform posterior inference directly in function space by incorporating more informative functional priors. Similar to their parameter-space counterparts, IVI and IMCMC have their own strengths and weaknesses. For in- stance, fVI is computationally efficient but imposes strong distributional assump- tions, while fMCMC is asymptotically exact but suffers from slow mixing in high dimensions. To inherit the complementary benefits of both schemes, this work proposes a novel hybrid inference method for the functional posterior inference. Specifically, it combines IVI and fMCMC successively by an elaborate linking mechanism to form an alternating approximation process. We also provide theo- retical justification for the soundness of such a hybrid inference through the lens of Wasserstein gradient flows in the function space. We evaluate our method on sev- eral benchmark tasks and observe improvements in both predictive accuracy and uncertainty quantification compared to parameter/functioonsppec VI and MCMC.
INTRODUCTION
Over the past few decades, Bayesian inference has emerged as promising direction for deep learn- ing (also known as Bayesian deep learning (Wilson & Izmailov] 2020)), offering improved pre- dictive performance, principled estimation of epistemic uncertainty, and enhanced generalization (Neal [995] Gal et al. 2016, Wilson Izmailov] 2020; Pielok et al. 2023). As one of the most representative models in Bayesian deep learning. Bayesian neural networks (BNNs) provide natu- ral support for continual learning, reinforcement learning, and decision-making tasks. BNNs have found wide application in various safety-critich domains, including medical diagnosis such Di- abetic Retinopathy (Filos et al.) 2019 Band et al.] 2021). the dissolution prediction of planetary systems (Cranmer et al. 2021), and the classification of radio galaxies (Mohan & Scaife 2024).
Due to the high-dimensional and multi-modal distributional properties, posterior inference for model parameters remarkably challenging in modern BNNs (Izmailov al. (021). Typically, there two popular classes of inference schemes: deterninistic approximations, such as the variational inference (VI) (Blundell al. 2015 Gal & Ghahramani 2016), and sampling methods like Markov Neal] 2012, Chen et al. 2014). However, the tricky prior issues in their parameter-space inference, including the non-interpretability and potential pathologies in deep models, e.g., the function samples of isotropy Gaussian prior over parameters tend to horizontal as the depth of the network increases (Matthews et al. 2018 Tran al. 2022), can severely impact the performance of BNNs in practical applications (Fortuin et al. 2022 Wild
Chain Monte Carlo (MCMC) (Welling &
jet al. 2022). These challenges have motivated researchers to explore Bayesian inference directly function space. For example, based on the Gaussian Processes prior (Rasmussen & Williams) 2006 Titsias 2009) which can effectively encode prior information about the periodicity, regularity, smoothness through the corresponding kernel functions, there several functional variational

---

Published as a conference paper at ICLR 2025
inference (fVI) approaches (Sun al. 2019: Ma & Hernández-Lobato 2021 Rudner et al. 2022)
have demonstrated improved prediction and uncertainty quantification compared to the parameter- space VI typically relies non-informative isotropic Gaussian priors. In the realm of MCMC, Wu et al. 2024a) proposed a functional stochastic gradient MCMC (fMCMC) scheme, which lifts the parameter-space dynamics of BNNs onto function space, enabling the incorporation of more informative functional priors.
Despite their advantages, like their parameter-space counterparts, IVI and fMCMC have comple- mentary strengths and limitations. I'I is fast and relatively easy to implement due to its explicit variational objective. but it is prone to bias and often relies on strong distributional assumptions about the structure and form of the variational posterior over parameters. such as using an oversim- plified fully-factorized Gaussian for large models. Another drawback is that it tends underestimate posterior variance (Zhang et al.) 2018), leading to poor uncertainty estimation. On the other hand, IMCMC non-aarametric and can asymptotically generate exact samples from the true posterior measure, but it is computationally expensive and suffers from slow mixing in high-dimensional mod- els. While stochastic gradient variants MCMC are more scalable, they can introduce significant bias in posterior expectations (Izmailov et al. 2021).
Naturally, it is prospective to combine fVI and fMCMC to leverage their complementary strengths and further improve posterior inference in function space. There has been some work focused on how to connect VI and MCMC (Ahn et al. 2012 Domke 2017: Habib & Barber (019). However, all these methods performed in the parameter space, extending them function spaces is not straightforward. For example, the structured MCMC proposed by Alexos et al. (2022) is based on an artificial partitioning structure of variational posterior over parameters, which is hardly applicable to infinieedimmensional posterior over functions. The tighter variational bound proposed by Salimans al.) (2015) in an expanded space is constructed by auxiliary random variables embedded from the intermediate MCMC transitions. However, extending this expanded variational posterior over parameters directly to an infiite-diimensional variational posterior over functions is not reasonable.
In this work, we propose a new hybrid inference method for posterior over functions that combines ideas from both fVI and fMCMC called Functional Variational Inference MCMC (FVIMC). We bridge the theoretical gap between IVI and IMCMC and design sophisticated linking mechanism that can alternate between variational optimization and sampling in a seamless manner. Specifically, explore the relationship between these two schemes through the lens of Wasserstein gradient flows in the function space, revealing that they share the same probability evolution marginals in posterior inference. Thus, innovatively alternate between fVT and fMCMC without changing the evolutionary path of the posterior, linking such two inference schemes together successively based on effective probability measure transformations form hybrid approximation process. This newly developed FVIMC incorporates meaningful functional priors and combines the strengths of both schemes: (1) the ability of I'I to quickly converge local high-probability target regions and (2) the non- paraaammetric nature of fMCMC, which allows it to explore multi multi-modal modal surfaces and prevent the posterior from collapsing to a local solution. Our main contributions are as follows:
We develop novel hybrid inference method that can efficiently link fVI and fMCMC successively to form an alternating approximation process, which can delicately inherit the benefits of both inference schemes and allow arbitrarily flexible and complex posterior over functions.
We prove that the functional Langevin stochastic differential equation (SDE) behind the fMCMC and the probability flow ordinary differential equation (ODE) derived from the corresponding Wasserstein gradient flows of the fVI share the same probability evolution marginals, which can support the justification of our method.
We perform a wide range of experiments to show the improved predictive performance and reliable uncertainty estimations of our method consistently compared to the competing VI and MCMC methods.
2 PRELIMINARIES
Bayesian neural networks (BNNs) Given a dataset D = (zt, Yi N training inputs I; X C TRP and the corresponding targets yı E y
{Xp, Yp). where the
C IRC Bayesian neural
2

---

Published as a conference paper at ICLR 2025
network is stochastic neural network characterized by random network parameters (weights and biases) denoted by a multivariate random variable w E IR* defined on a probability space (S2, A, P). Let f(: w) X' S Y/ be the random function (product measurable) defined by BNN on an infinite-dimensional function (Polish) space IHT with Borel o-algebra ((TH), which is A measurable for every x X. The prior distribution for the network parameters is denoted as po(w.. The likelihood is defined by a function as 3 x HI + (0), (),, f[Z;[[(]] + f( w)). where Y' S Rº is Borel measurable. The likelihood evaluated on the training data is pl D w)). The posterior distribution over parameters is then inferred with the unnormalised form as pl w (D) x of (ww). Given test data x²,, the predictive distribution for y" is obtained by inte- grating over the posterior as p(y" x=",, D) = p(v ( w))p(w(D)dw.
Function-space VI Suppose P(HI) represents the space of Borel probability measures on B(H.). To incorporate more meaningful prior information into posterior inference, IVI assigns a functional prior measure Po P(H)), such the classical Gaussian Processes (GP) prior denoted by Po GP(m, K) and perform variational inference in function spaces. The posterior measure PAID E P(H)), induced by the functional prior and likelihood. is defined by the Radon-Nikodym derivative as (cf) p( ((YplffC fxlffXp: (ddf) (Matthews al. 2016 Lambley 2023). The functional variational objective is to minimize the Kullback-Leibler (KL) divergence minq, KL[Q,
where Q E ((H)) distribution over parameters denoted by g(w: f). with variational parameters 9, and q(w; (dww) Qd(df). The pushforward of q( 0). denoted by T=g(ww 0) Q₁. where T# (-) refers to the
the approximate posterior measure induced by the approximate posterior
pushforward measure (Wild et al. 2022; Rudner et al. 2022). The functional evidence lower bound
(ELBO) to be maximized in this framework is
La, Eq, [log pl YD f(XD: - [Po]
Eo, [log (YD f(Xx; w])] sup
KT
(I)
XeX
where X3 UnEN (X Xn X₂ S RPPp) the set of all finite marginal measurement points in the input domain, fx is the function value evaluated at X (Sun et al. 2019).
Function-space MCMC Wu et al. (2024a) explored the dynamics of BNNs in function space and proposed novel fMCMC scheme. This functional MCMC scheme based on the potential energy functional designed in terms of the posterior over functions Prip. which effectively incorpo- rates functional priors while ensuring that the target posterior over functions remains the stationary distribution. The specific functional Langevin dynamics for [(; w) on Hi is given by:
dfíl: w) p(fel: w) ((fe(::)))d
- fa ( log (xx: + Io(f)) + Hwfe dt +
(2)
where lo( denotes the Onsager- (OM) functional for Po heuristically interpreted as the negative logarithm of Po (Lambley) ).23)) B is Wiener process (Brownian motion), W is the Fréchet derivative of f(; w) w.r.t. network parameters w, and Hw denotes the second-order Fréchet derivative. Additionally, leveraging the Itô Lemma (see in Appendix (A), they derived the corresponding diffusion process for the network parameters. The specific discretization update rule for sampling the network parameters under this functional stochastic gradient Langevin dynamics (fSGLD) is
W++1 = Wt - - V.
V log
- V pM log Po(f.XM) + 22
(3)
where é, is the decreasing step-size, 7h is standard Gaussian noise, is the mini-batch size, XMM def _x,...,,xx XM) are the finite measurement points, and pndxxxx XMM are the corresponding function values evaluated X.M
Note that without loss of generality, we assume that Q, E P(HH) is dominated by Po (therefore also dominated by p)(p) to avoid technical difficulties in the definition.
3

---

Published as a conference paper at ICLR 2025
3 OUR METHOD
In this section, we propose to bridge the theoretical gap between these two functional inference schemes and develop novel hybrid inference method for posterior measure over functions. We start from an initial random function mapping defined by a BNN denoted by san(; Wso). given simple initial approximate posterior over parameters q-o g-o(w (Wa:: 0s0), e.g., a factorized Gaussian has 0.0 {u Es which induces the posterior measure over functions denoted by Qr... Note that this factorized Gaussian will be modified later rather than being fixed, as in I'I, it does not impose a strict restriction. Our functional hybrid inference method is formed by the following alternating approximation stages:
V-Stage: The first stage starts by executing the IVI through maximizing the functional ELBO in Equation ① to obtain the updated approximate posterior measure as Q and the corresponding variational posterior distribution over parameters is 9a1 (W&1 0.1) with 0 {Me1, E-1)
M-Stage: In this stage, starting from the randomly sampled initial point w,1 N(Hs1 Σ.1) from the last variational stage, we perform stochastic gradient fMCMC using fSGLD Equation (3). Let (x)(1.x) denotes the collected N samples of network parame- ters, the corresponding function samples is as fe2 The distribution over parameters for the collected network parameters samples is denoted by 4s2 4=2(w (Ws2), which is distributional assumption free and implicit, and the corresponding probability measure over functions is as ff
Our idea is to alternate between the two stages, but one challenge arises: 9.2( (2)) from the M-Stage is con-aaametrric and lacks the explicit distribution form required by the V-Stage. To address this, we reconsider the relationship between 1 and / where the latter comes from the last V-Stage with explicit parametric 9s1 (W.1: 0.)). We can simply consider Qj the input the M-Stage and Q as the output. Thus, Q,.2 is a refined posterior compared to Q Given this relationship, we propose treating fa transformation of Qi. where the transformation follows nonlinear, pa- rameterized form. By uncovering this transformation, we can derive an explicit parametric qa2 (2) with the aid of qai (Waii 0st). We define the following stage for this transformation learning:
T-Stage: Let Fx (52(s) + fai(; w.1) be an invertible bijective function parametrized by a1 (with inverse F, where fa2 wee) and fai wa1) denote the ob- tained random function mappings from the M-Stage and V-Stage, respectively. According to the change of variable formula, we have
= Qr. (FX (F (fx det OF
OF det
1r (1%, (FxMM)
OF
(fXM)
3 pPMM
(4)
and
FF-1 afxM
log Q1. = log Q, (F (FX + log det
(5)
where X.M [xI...., XM] T are finite randomly sampled measurement points and
P.M are the function evaluations XM- Based on the local linearization (Rudner et al. 2022), the implicit functional Qr is approximated as Gaussian process given by (Paillfaa ( Junn E.ITTT)). where Jun denotes the Jacobian (far).w..))
Dw.1. w.i maat
and the finite marginal distribution of Q reduced to an analytical Gaussian distribu- tion. Given the function samples fx on X.M from M-Stage, we can obtain the optimal
bijective function FX. by maximizing Equation (5) and transform fe2(: W.2). With F in hand, we can obtain the explicit form of Q as Fx; o Q where Fx; o Q denotes the pushforward relationship as shown in Equation
4

---

Published as a conference paper at ICLR 2025
With the help of this T-Stage, successfully link the above V-Stage and M-Stage and form chain as follows:
X. (
V-Stage 244 M-Stage
Lig V-Stage
FxT Px
T-Stage
V-Stage.
The pseudocode for the above procedure (named FVIMC) shown in Algorithm Din Appendix C Let us have an in-depth analysis of this procedure. Because each T-Stage in this hybrid inference process yields a new explicit form for the approximate posterior obtained from the M-Stage based on the approximate posterior from the previous V-Stage, and this will be used for the following V-Stage, the form of the final posterior from the above procedure would no longer be a (initialized) factorized Gaussian, and we have the following result.
Proposition 3.1. The final approximate posterior measure is in a composited form as Exix 13/2 ...Fx;: Fx; Qi with total of K alternating approximate V-Stages and M-Stages, which is arbitrarily complex and preferably flexible.
The derivation is given in Appendix B.I] This proposition demonstrates that, although we ini- tially set a simple approximate posterior, the final result can still be flexible enough to match a complex posterior. This validates our ability to overcome the strong distributional assumption of the variational posterior in fVI. The key to achieving this lies in the M-Stage, which introduces non-parametric updates to the posterior. By alternating between the V-Stage and M-stage with the support of the T-Stage, the process continuously refines the posterior approximation. Specifically, the M-Stage introduces samples that capture the multi-modal nature of the true posterior, while the V-Stage applies optimization to tighten the approximation. Through this alternating and succes- sive process, the approximate posterior evolves to become more expressive and accurate, leveraging the strengths of both schemes-fvi efficiency in local regions and fMCMC's ability explore complex distributions.
Another important question arises: can we safely link the M-Stage and V-Stage? Specifically, how can we ensure consistent convergence paths if these two stages evolve in different directions and ultimately converge to distinct optimums? To address this concern, we present the following result:
Proposition 3.2. The functional Langevin SDE defined in Equation ( and the probability flow ODE derived from the Wasserstein gradient of the functional variational inference in Equation ([) share the same probability measure evolution marginals if they evolve from the same initial point.
The proof is given in Appendix B.2
This proposition guarantees the soundness of our proposed hybrid FVIMC since the approximate posterior measures of FSGLD and I'I share the same evolution marginals. Therefore, it is reasonable to link these two schemes successively to form an alternating approximation process.
The form of the invertible transformations The parametric form of transformation Fx(-) in the T-Stage plays a crucial role in determining the approximation accuracy to the samples from the M- Stage and, consequently, the flexibility of the posterior modelling capability. In general, increasing complexity improves approximation. However, more complex transformations also increase compu- tational difficulty, particularly when calculating the determinant of the Jacobian during optimization. Therefore, designing (-) requires a balance between modelling flexibility and computational fea- sibility. Here, use three classes of flexible and tractable invertible bijections: linear, non-linear Tanh and Sigmoid, respectively:
y 2 [(;; w) 1 bba,ww)
y a2 Tanh (ars (x; + ba (1;4w)) by ba(a;
(8)
y = a2 Sigmoid (al((::)) + b(I;w),
where f(; w) denotes function mappings, is scalar parameter and b(x; denotes a bias func- tion. The computation of the Jacobian for all three transformations is straightforward and does not involve the Jacobian of function b(x; w). so it be arbitrarily complex and we make them deep neural networks with parameters w to achieve high expression ability. Then the parameters of the bijections are denoted as À (a,w). The specific form for each transformation need not be fixed.
5

---

Published as a conference paper at ICLR 2025
For flexibility, implement automatic selection mechanism that chooses the form with the highest likelihood value during the optimization. This allows the model to adaptively select the most suitable transformation for different function mappings. More details about the Jacobians and inverse bijections are shown in Appendix
Avoid the potential problematic KL divergence for I'VI The KL divergence in fVI is found to be potentially problematic (Burt et al., (020)) (see Appendix A). Hence, following the lat- est idea in the area (Knoblauch al. (022), we propose to use the Wasserstein distance to replace the KL divergence in the V-Stage of our FVIMC with the alternative functional varia- tional objective as Lw == -Rq log pYD f(XD: w)]] + (Q Po) + (Q Po), where (i. the I-Wasserstein distance between Q5 and Po that admits the dual form definition as sup | 1x ((x) - ( and M2(Qy, Po) = |ear(Qy)) var ( (Po) is an enhanced second-order moment matching term to preserve the uncertainty information encoded in the func- tional prior, which is found to be weak in the naive 1-Wasserstein distance (Wu et al.. 2024b). Compared to the KL divergence-based IVI, which is severely limited by the requirement for ex- plicitly accessible densities, the Wasserstein distance-based functional variational objective offers significantly more flexibility. This is because it relies entirely on sampling- -based procedure, elim- inating the need for the closed-form of Q, At the same time, Proposition 3.2 still holds under this replacement in certain cases (see Appendix G.5 for more discussions and empirical evidence).
4 RELATED WORK
Functional posterior inference in BNNs Due to unresolved deficiencies of parameter-space in- ference for BNNs, recent work has shifted toward performing Bayesian inference directly in fune- tion space. For functional variational inference, [Sun et al.] 2019) introduced a functional ELBO that leverages more informative GP priors and explicitly minimizes the KI. divergence between the variational posterior and the true posterior in function space. They proved that the KI. divergence between two dimensoaa stochastic processes corresponds to the supremum of the KI. di- vergence over all finite marginal distributions. However, employing KL divergence in function spaces requires careful consideration as it presents several definitional and estimation challenges (Burt al., 2020). For example, the variational posterior measure must be dominated by the func- tional prior to ensure the existence of the Radon Nikodyy derivative, which defines the posterior measure (Matthews et al., 2016: Wild et al. 2022). If this condition is unmet, this functional KL divergence may become ill ill-define, e.g., the KL. divergence between two BNNs formed by different network architectures
be infinite (Ma & Hornández-Looato 202T). T). Additionally, |Rudner et al. 2022) pointed out the analytical intractability of the supremum over marginal KL divergences and proposed a more practical, well-defined functional variational objective. Their method is based on approximating the posterior and prior over functions as Gaussian processes via local linearization. Considering the limitations of functional KL divergence, |Tran et al. (2022) proposed matching a BNN prior with an interpretable GP prior using the I-Wasserstein distance (Kantorovich) (960). while Wild et al. 2022)) introduced a generalized functional variational objective for Gaussian mea- sures based on the 2- Wassersttii distance. For functional MCMC, Wu et al. (222aa) investigated the diffusion process of BNNs in function space and designed a new functional stochastic gradient MCMC scheme, which can incorporate functional priors and guarantee that the stationary distribu- tion of the resulting functional dynamics aligns with the target posterior distribution over functions.
Connections between VI and MCMC in parameter space There has also been some work aimed at hybridizing VI and MCMC in parameter space. Most of these efforts focus on sampling from an approximate posterior, rather than the true posterior, typically on tempered or structured evidence lower bound. Alternatively, they optimize variational parameters using auxiliary variables derived from embedded MCMC schemes to achieve a tighter variational bound. For example, Ahn et al. (2012) proposed a Stochastic Gradient Fisher Scoring (SGFS) algorithm that uses the inverse Fisher information matrix to sample from a Gaussian approximation of the posterior based on stochas- tic gradient Langevin dynamics. [Saliaans et al. (2015) integrated Markov transition chains into variational lower bound as auxiliary variables to improve posterior approximation. [omkee| (2017)
derived two bounds for the variational divergence and introduced hybrid algorithm based on it that directly interpolates between variational inference and Langevin dynamics. Hoffman & Ma (2020)
find that Langevin dynamics can implicitly track the nonparametrie normalizing flow of variational

---

Published as a conference paper at ICLR 2025
10
co
ab
Ch CO
(a) FVIMC
(b) KLBBB
(c) SGLD
(d) (SGLD
10
(c) IVI
(f) FBNN
(g) GWI
(h) IFBNN
Figure 1: Learning polynomial curves. The green line is the ground true function and the blue lines correspond to mean approximate posterior predictions. Black dots denote 20 training points; shadow areas represent the predictive standard deviations. The top row is our FVIMC, parameter- space VI/SGLD and functional FSGLD, and the bottom row shows the results of four function-space VI baselines. For more experimental details, see Appendix F The complete inference process of FVIMC shown in Appendix FF]
inference, reinterpreting black-box variational inference parametric approximation Langevin dynamics. Moreover, Liu et al. (2019) developed framework on the Wasserstein space, casting general MCMC dynamics on as the fiber-gradient Hamiltonian flow. Alexos et al. (2022) pro- posed structured MCMC scheme that samples from an approximate posterior under distributional factorization constraints. However, all of these methods are restricted parameter space, making challenging to extend them directly to function space.
EXPERIMENTS
We evaluate FVIMC on a variety of benchmark tasks, including a synthetic extrapolation, multi- variate regressions on UCI datasets, contextual bandits and image classification tasks. We compare the performance of FVIMC with several competing parameterfiunction-spac VI and MCMC base- lines, including the benchmark parameter-space VI method BBB (Blundell al., 2015) denoted by KLBBB, and four competing function-space VI methods: FBNN (KL divergence-based IVI) (Sun et al. 2019). GWI(Wild al. 2022)) IFBNN(WU al. 2023 and our proposed Wasser- stein distance-basse IVI denoted as IVI. Also, for MCMC methods, we compare with the classic parameter-space stochastic gradient Langevin dynamics denoted by SGLD (Welling & Teh 201T) and the functional fSGLD (Wu et al., 2024a). We obtain the improved predictive performance and uncertainty quantification, indicating that FVIMC can efficiently incorporate meaningful knowledge from the functional prior and inherit the strengths of both fVI and fSGLD schemes.
5.1 EXTRAPOLATION ON SYNTHETIC DATA
We begin by verifying the data fitting and uncertainty quantification capabilities of our method through a I-D oscillation curve extrapolation experiment using synthetic dataset. Consider polynomial target function: y sin(3 + 0.3 cos(9 + 0.5 sin((z)) + e with noise - -0.25) and another 10
(0, 0.52), we randomly sampled 10 points from Uniform(-0.7 points from Uniform(0.25, 0.75) as input data. The model used is fully-connected neural network with two hidden layers. We apply the same GP prior with the RBF kernel, pre-trained the input dataset, for all functional methods. For all parameter-spase methods that cannot specify a meaning- network parameters. For FVIMC,
ful prior, place the default isotropic Gaussian prior we employ a hybrid inference process consisting of 21 alternating V-Stages and M-Stages, along with 10 intermediate T-Stages. Specifically, this process includes 11 V-Stages and 10 M-Stages,
7

---

Published as a conference paper at ICLR 2025
Table I: The table shows the results of average RMSE for multivariate regression UCI datasets. We split each dataset randomly into 90% training data and 10% test data, and this process is repeated 10 times ensure validity, Bold indicates statistically significant best results (p- < 0.01 with t-test).
RMSE
YACHT
BOSTON
CONCRETE
ENERGY
WINE
KIN8NM
PROTEIN
FVIMC
0.37 1 0.13
0.36 1 0.07
0.40 1 0.04
0.23 1: 0.02
0.68 1 0.05
0.57 1 0.02
0.95 0.03
FVI
0.45 0.10
0.40 0.08
0.04
0.27 - 0.02
0.69 + 0.05
0.61 0.02
1,01 0.02
GWI
2.20 2.00
1.74 0.05
+.30
1.46
1.68 + 0.06
1.19 0.02
1.33 0.01
FBNN
1.52 0.08
1.68
0.12
1.27 0.05
1.35
1.53
1.50
FFNN
1.110
1.44 0.09
1.07 0.07
1.19
1.21 1 0.05
1.12 1 0.02
1.16 1 0.01
KLBBB
2.13 0.09
0.07
1.78
1.78=0.07
1.86 + 0.07
1.79 1.79 = 0.03
1.80 0.01
SGLD
1.09 0.100
1.23 0.006
0.07
1.04 0.06
1,08 1 0.09
1.20 0.02
1.12 0.01
ESGLD
0.41 0.10
0.36 0.09
0.45 0.004
0.24 - 0.03
0.71 1 0.05
0.74 0.02
0.84 0.01
with each V-Stage trained for 1000 epochs and each M-Stage for 500 iterations. We use two dif- ferent bijections for the transformations of probability measures in the T-Stage after each M-Stage:
the linear transformation in Equation 7 and the non-linear Tanh bijection in Equation (8). The bias function for both bijections is parametrized by neural network with two hidden layers. We implement the likelihood-based auto-selection mechanism choose between these two invertible transformations. For fair comparisons, we run 16000 epochs for all other parameter/function-spae VI methods. For all sampling methods, we use 2000 burn-in iterations and 14000 iterations for 140 samples (to reduce correlations between samples, we draw separated samples every 100 epochs). The results are shown in Figure The parameter-space VI method KL.BBB in Figure (bb) fails to fit the target function. The two MCMC methods could recover the key polynomial trend of the curve in the observation range, while the uncertainty quantization in the unseen areas is not justified. Sim- ilarly, the three function-space VI methods in Figure I(f) Figure [(g) and Figure ( also severely underestimates the predictive uncertainty, which is a typical downside of VI ZZang et al. 2018 Alexos et al.. 2022). In contrast, FVIMC and the Wasserstein distance- based fVI show better predictive performances and improved uncertainty estimations. In particular, the hybrid FVIMC in exhibits exhibbbbits stronger ability than the pure variational fVI in Figure [(ee] e.g., the smoother fitting curves, and can capture critical features of the non-observed middle range [-0.25, 0.25), which demonstrates the powerful advantages of combining I'I and fMCMC for a hybrid inference. See Appendix H for more detailed ablation studies about the impact of kernel functions in functional GP priors and the form of the bijections for probability transformations.
Figure T
5.2 UCI REGRESSION
In this experiment, we evaluate the predictive performance of FVIMC for multivariate regression tasks on 7 real-world UCI datasets: Yacht, Boston, Concrete, Energy. Wine, Kin8nm and Protein. We use 2- hidden-layer fully connected neural network. For FVIMC, we employ total of five alternating V-Stages and M-Stages to form the hybrid approximation process. Each V-Stage and M- Stage is trained for 400 epochs, resulting in total of 2000 iterations. Therefore, for all other VI methods for fair comparison. For the two sampling methods, we run 500 iterations for the burn-in stage and collect 15 samples in the following 1500 iterations. We report the average root mean square еттor (RMSE) results in Table I Our FVIMC achieves better accuracy than all other VI and MCMC baselines 6/7 datasets, which illustrates the impressive predictive power of our hybrid method. For NLL results shown in Table 2 FVIMC outperforms other methods on 5/7 datasets, demonstrating its competitive uncertainty quantification ability. Moreover, the analyses of the mixing time, computational complexity and transformation errors in the T-Stage are presented in Appendix G where our FVIMC exhibits high convergence speed and training stability.
2000 epochs
5.3 CONTEXTUAL BANDITS
In this section, we examine the ability to guide exploration-exxploittton in contextual bandits, where uncertainty modelling is crucial in sequential decision-makinn for such downstream tasks. In this problem, the agent interacts with unknown environment repeatedly and chooses optimal action to maximize the reward given the context in each round of interaction. Thompson sampling is a clas- sic technique applicable to this scenario Thompson (1933). It samples model configuration from
8

---

Published as a conference paper at ICLR 2025
Table 2: The table shows the results of average NLL for multivariate regression UCI datasets. We split each dataset randomly into 90% training data and 10% test data, and this process is repeated 10 times ensure validity. Bold indicates statistically significant best results (p 0.01 with t-test).
YACHI
BOSTON
CONCRETE
ENERGY
WINE
KIN8NM
PROTBIN
FVIMC
33.03
1.13
.3.44 1 0.48
5500
0.34
5.70 1 0.39
-3.77
0.40
0031
0.65
-.07 =: 0.36
FVI
-3.07 0.66
-2.66
0.27
-3.03
0.22
4.20
-2.25
0.55
-2.30
0.30
GW1
0.76
-1.04
0.68
-0.68 0.49
-2.03
0.70
-2.60 0.224
-1.58
0.23
0.86
-1.50
-..14
-2.65
IFBNN
-1.25 1.21
0.32 0.30
.0.39 0.33
-1.78
0.336
0.26
0.15
-1.01 1.114
22.13 0.34
KLBBB
2.51=:0.16
2.07
0.12
2.61
2.17
2.13
2.61
2.02
SGLD
-0.37 0.08
-0.00
0.11
-0.54
0.005
-0.58
0.09
-0.90
0.12
-0.60
-0.60
0.01
FSGLD
-2.46 0.22
-2.20
0.20
-1.63
1.63 0.14
-2.25 4 0.25
-2.33
0.15
-..24
-2.10
0.08
3000
22.9
0M
11:23
1880 569
50.D
FVPC
Ieeatinn
lteration
lieeston
(a) 0.4
(b) p= 0.5
(c) p=
Figure 2: Comparisons of cumulative regrets of FVIMC, IVI, FBNN, GWI. KLBBB, SGLD for contextual bandit task on the Mushroom dataset. Lower represents better performance.
the current posterior first and adaptively chooses an optimal action under this sampled configuration for the current context, then updates the posterior based on the corresponding observed reward.
We compare FVIMC with several baselines on the UCI Mushroom dataset with 8124 instances. In each instance, the mushroom is labelled as edible or poisonous and has 22 features as the input context. The agent observe mushroom features in each interaction and choose eat or reject a mushroom.
follow the basic settings by (Blundell et al. 2015) and consider three different reward patterns: if the agent eats edible mushroom, it receives reward of 5; otherwise, reward of if the agent rejects the edible mushroom. For the poisonous situation, if the agent chooses to cat, it receives reward of -35 with probabilities of 0.4, 0.5, and 0.6, respectively, to form three different patterns; otherwise, a reward of () for rejection. Suppose an oracle always eats an edible mushroom and rejects the poisonous ones. The cumulative regrets with respect the reward achieved by the oracle can measure the exploration-exploinee alternating V-Stages and M-Stages with total of 7000 iterations. And run 7000 epochs for all other baselines with batch size 64. The cumulative regrets for three reward patterns are shown in Figure 2 where our FVIMC consistently achieves the lowest regrets, which indicates its reliable uncertainty quantification in decision-making.
ability of an agent. For FVIMC, we run five
5.4 IMAGINE CLASSIFCCAION AND OOD DETECTION
We demonstrate the scalability of FVIMC on high- digh-dimensional image classification tasks. We test the in-distribution predictive performance and out-of-distribution (OOD) detection ability on MNIST (LeCun et al.) 2010) and FashionMNIST (Xiao et al.) (2017). For our FVIMC, we use five alternating V-Stages and M-Stages with total of 320 iterations. Therefore, we run 320 epochs for all other VI methods. For sampling methods, we run for 300 burn-in iterations and then collect 20 samples. We report the test classification errors (%) for predictive performance and the area under the curve (AUC) of OOD detection pairs FashionMNISTTMNNISS MNIIT/FashionMNNSS based on predictive entropies Table 3 Our FVIMC outperforms all parameter/function-nae VI and MCMC baselines for classification accuracy and demonstrates competitive OOD detection ability.

---

Published as a conference paper at ICLR 2025
Table 3: Image classification and OOD detection performance.
MNIST
FMNIST
MODEL
TEST ERROR
AUC
TEST
ERROR
AUC
FVIMC
4.95 + 0.00
0.832 = 0.00
14.57 = 0.00
0.800 = 0.00
FVI
5.49 0.00
0.815
14.99 + 0.00
0.831 0.05
GWT
6.23
0.826+0.00
15.24 0.00
0.459 0.07
FBNN
5.06
0.787
+.000
0.809 1009+000 0.03
IFBNN
5.00
0.703
1.004
0.833
KLBBB
5.30
0.827 0.00
1..26
0.782
FSGLD
1.000
0.756
15.99
0.810 0.00
CONCLUSION
We propose hybrid inference method for BNNs in function space, combining the strengths of both functional variational inference (fVI) and functional MCMC (fMCMC) schemes. To achieve this, we develop a hybrid approximation process that alternately and successively links these two schemes, allowing the process to inherit the advantages of both while constructing arbitrarily flexible and complex approximate posteriors. We prove that the functional Langevin SDE and the probability flow ODE derived from the Wasserstein gradient flows of the functional variational inference share the same probability evolution marginals. This ensures the theoretical soundness of the hybrid ap- proach. Empirically, we show the improved predictive performance and uncertainty estimation of our method on a range tasks.
ACKNOWLEDGMENTS
This work is supported by the Australian Research Council under the Discovery Early Career Re- searcher Award DE200100245.
REFERENCES
Sungjin Ahn, Anoop Korattikara, and Max Welling. Bayesian posterior sampling via stochastic gradient fisher scoring. In International Conference on Machine Learning, 2012.
Antonios Alexos, Alex Boyd, and Stephan Mandt. Structured stochastic gradient meme. In Inter- national Conference on Machine Learning, pp. 414-434. PMLR, 2022.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008
Neil Band, Tim GJ Rudner, Qixuan Feng, Angelos Filos, Zachary Nado, Michael W Dusenberry. Ghassen Jerfel, Dustin Tran, and Yarin Gal. Benchmarking bayesian deep learning diabetic retinopathy detection tasks. In 35th Conference on Neural Information Processing Systems, 2021.
Roman V Belavkin. Relation between the wewpooovviccwwssessteeee metric and the kullback-leible divergence. In Information Geometry and Its Applications: On the Occasion of Shun-ichi Amari's 80th Birthday, IGAIA IV Liblice, Czech Republic, June 2016, pp. 363-373. Springer, 2018.
Charles Blundell, Julien Corneiise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International Conference on Machine Learning, pp. 1613-1622. PMLR, 2015.
Zdzislaw Brzežniak, Jan MAM van Neerven, Mark C Veraar, and Lutz Weis. Itô's formula in umd banach spaces and regularity of solutions of the zakai equation. Journal of Differential Equations, 245(1):30-58, 2008.
David R Burt, Sebastian W Ober, Adrià GarrigaAAlnss and Mark van der Wilk. Understand- ing variational inference in function-space. In Third Symposium on Advances Approximate Bayesian Inference, pp. 1-17, 2020.
10

---

Published as a conference paper at ICLR 2025
Changyou Chen, Ruiyi Zhang, Wenlin Wang, Bai Li, and Liqun Chen. A unified particle- optimization framework for scalable bayesian sampling. In 34th Conference on Uncertainty in Artificial Intelligence, 2018.
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International Conference Machine Learning, pp. 1683-1691. PMLR, 2014.
Miles Cranmer, Daniel Tamayo, Hanno Rein, Peter Battaglia, Samuel Hadden, Philip J Armitage, Shirley Ho, and David N Spergel. A bayesian neural network predicts the dissolution of compact planetary systems. Proceedings of the National Academy of Sciences, 118(40)::20260531188 2021.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. In International Conference Learning Representations, 2017.
Justin Domke. A divergence bound for hybrids of meme and variational inference and an application to langevin dynamics and sgvi. In International Conference on Machine Learning, pp. 1029- 1038. PMLR, 2017.
Angelos Filos, Sebastian Farquhar, Aidan Gomez, Tim GJ Rudner, Zachary Kenton, Lewis Smith, Milad Alizadeh, Arnoud De Kroon, and Yarin Gal. A systematic comparison of bayesian deep learning robustness in diabetic retinopathy tasks. 4th workshop on Bayesian Deep Learning (NeurIPS 2019), 2019.
Vincent Fortuin, Adrià Garriga-Alonso, Sebastian W. Ober, Florian Wenzel, Gunnar Ratsch, Richard E Turner, Mark van der Wilk, and Laurence Aitchison. Bayesian neural network pri- ors revisited. In International Conference on Learning Representations, 2022.
Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In International Conference on Machine Learning, pp. 1050-1059.
PMLR, 2016.
Yarin Gal et al. Uncertainty in deep learning. PhD thesis, University of Cambridge, 2016.
Arka P Ghosh. Backward and forward equations for diffusion processes. Wiley Encyclopedia of Operations Research and Management Science, 2010.
Henry Gouk, Eibe Frank, Bernhard Pfahringer, and Michael Cree. Regularisation of neural net-
works by enforcing lipschitz continuity. Machine Learning, 11:3933-41 -1993-41,,,1 2021.
Robert M Gray. Entropy and information theory. Springer Science & Business Media, 2011.
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Im- proved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
Raza Habib and David Barber. Auxiliary variational meme. 7th International Conference on Learning Representations, pp. 1-13. PMLR, 2019.
Matthew Hoffman and Yian Ma. Black-box variational inference as a parametric approximation to langevin dynamics. In International Conference Machine Learning, pp. 4324-4341. PMLR, 2020.
Kiyosi Itô. On formula concerning stochastic differentials. Nagoya Mathematical Journal, 3: 55-65, 1951.
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are bayesian neural network posteriors really like? In International Conference on Machine Learning, pp. 4629-4640. PMLR, 2021.
Leonid Kantorovich. Mathematical methods of organizing and planning production. Management science, 6(4):366 422, 1960.
11

---

Published as a conference paper at ICLR 2025
Jeremias Knoblauch, Jack Jewson, and Theodoros Damoulas. An optimization-centeee centric view bayes' rule: Reviewing and generalizing variational inference. Journal of Machine Learning Research, 23(13)):1-109, 2022.
Ivan Kobyzev, Simon JD Prince, and Mareus Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 43 (11):3964-3979 2020.
Hefin Lambley. Strong maximum a posteriori estimation in banach spaces with gaussian priors. Inverse Problems, 39(12):125010,, 2023.
Yann LeCun, Corinna Cortes, Chris Burges, et Mnist handwritten digit database, 2010.
Chang Liu, Jingwei Zhuo, and Jun Zhu. Understanding meme dynamics as flows on the wasserstein space. In International Conference on Machine Learning, pp. 4093-4103, PMLR, 2019.
Chao Ma and José Miguel Hernández-Lobate Functional variational inference based stochastic process generators. Advances in Neural Information Processing Systems, 34:21795-21807, 2021.
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient meme, Advances in Neural Information Processing Systems, 28, 2015.
Alexander G de G Matthews, James Hensman, Richard Turner, and Zoubin Ghahramani. On sparse variational methods and the kullback-leibler divergence between stochastic processes. In Artificial Intelligence and Statistics, pp. 231-239. PMLR, 2016.
Alexander G de G Matthews, Jiri Hron, Mark Rowland, Richard E Turner, and Zoubin Ghahra- mani. Gaussian process behaviour in wide deep neural networks. International Conference on Learning Representations, 2018.
Devina Mohan and Anna M M Scaife. Evaluating bayesian deep learning for radio galaxy classifi- cation. In The 40th Conference on Uncertainty in Artificial Intelligence, 2024.
Radford M Neal. Bayesian learning for neural networks. PhD thesis, University of Toronto, 1995.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
Bernt Øksendal. Stochastic differential equations. Springer, 2003.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lak- shminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57): 1-64, 2021.
Tobias Pielok, Bernd Bischl, and David Rügamer. Approximate bayesian inference with stein func- tional variational gradient descent. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Carl Edward Rasmussen and Christopher KI Williams. Gaussian processes for machine learning, volume 2. MIT press Cambridge, MA, 2006.
Hannes Risken. fokker-planck equation. Springer, 1996.
Tim GJ Rudner, Zonghao Chen, Yee Whye Teh, and Yarin Gal. Tractable function-space variational inference in bayesian neural networks. Advances in Neural Information Processing Systems, 35: 22686-22698, 2022.
Tim Salimans, Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In International Conference on Machine Learning, pp. 1218-1222 PMLR, 2015.
Jianghong Shi, Tianqi Chen, Ruoshi Yuan, Bo Yuan, and Ping Ao. Relation of interpretation of stochastic differential equations to ito process. Journal of Statistical physics, 148:579-599 2012.
12

---

Published as a conference paper at ICLR 2025
Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. Functional variational bayesian
neural networks. In International Conference on Learning Representations, 2019.
Esteban G Tabak and Cristina V Turner. A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics, 66(2):145-164, 2013.
William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-29) 294, 1933.
Michalis Titsias. Variational learning of inducing variables in sparse gaussian processes. In Artificial Intelligence and Statistics, pp. 567-574. PMLR, 2009.
Ba-Hien Tran, Simone Rossi, Dimitrios Milios, and Maurizio Filippone. All you need is a good functional prior for bayesian deep learning. The Journal of Machine Learning Research, 23(1): 3210-3265, 2022.
Ziyu Wang, Tongzheng Ren, Jun Zhu, and Bo Zhang. Function space particle optimization for bayesian neural networks. International Conference on Learning Representations, 2019.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning, pp. 681-688, 2011.
Veit David Wild, Robert Hu, and Dino Sejdinovic. Generalized variational inference in function spaces: Gaussian measures meet bayesian deep learning. Advances in Neural Information Pro- cessing Systems, 35:3716-3732 2022.
Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. Advances in Neural Information Processing Systems, 33:4697-4708, 2020.
Mengjing Wu, Junyu Xuan, and Jie Lu. Indirect functional bayesian neural networks In Fifth
Symposium on Advances in Approximate Bayesian Inference, 2023.
Mengjing Wu, Junyu Xuan, and Jie Lu. Functional stochastic gradient meme for bayesian neural networks. arXiv preprini arXiv:2409. 240.:2632.6662 2024a.
Mengjing Wu, Junyu Xuan, and Jie Lu. Functional wasserstein bridge inference for bayesian deep learning. In 40th Conference on Uncertainty in Artificial Intelligence, 2024b.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint arXiv:1708.07744: 2017.
Mingxuan Yi and Song Liu. Bridging the gap between variational inference and wasserstein gradient flows. arXiv preprint arXiv:2310. 10.20090, 2023.
L Yin and P Ao. Existence and construction of dynamical potential in nonequilibrium processes without detailed balance. Journal of Physics A: Mathematical and General, 39(27):8593, 2006.
Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE Transactions on Pattern Analysis and Machine Intelligence, 41(8):2008-202: 2018.
Ruiyi Zhang, Zheng Wen, Changyou Chen, and Lawrence Carin. Scalable thompson sampling via optimal transport. In International Conference on Artificial Intelligence and Statistics, 2019.
13

---

Published as a conference paper at ICLR 2025
A FURTHER BACKGROUND
Itō Lemma Itö Lemma IIō 1951 a fundamental result in stochastic calculus find the differ- ential of function of stochastic process, which serves as the stochastic calculus counterpart of the chain rule. For an Itō diffusion, let f(x)) be an arbitrary twice differentiable scalar function real variables Xx, then the differential of f())) can be derived from the Taylor series expansion of the function as
(m(X,)
0f (x,,) 8- dr 2
di + (Xx)) Of
as(x;)
dB,
(10)
On.
dr
which implies that f(X,) is itself an Itô diffusion (Brzežniak et al. 2008
Change of variable formula and normalizing flows normalizing flow a powerful transfor- mation that generate highly flexible and complex probability density from simple initial one through a sequence of bijective invertible mappings (Tabak & Turner 2013 Kobyzev et al. 2020 Papamakarios et al. 2021). Let Ee be random variable with probability density P.(5), and
F Rd - Rd is invertible mapping (with inverse F-!). Then, the density of the transformed variable = F() can be obtained using the change of variable formula (Dinh et al. 2017) as
0F-1 On
Or ac
(n) = pe(() det
(E( det
(11)
By successively transforming initial variable Eo with density Pto through sequence of K invert- ible mappings F6 as εx = Fx o...o of & by successively applying equation II
o F22 o Fi (Eo). we can derived an arbitrarily flexible density Рex
log det 0FF DE
Peg (Ex) = loo Pog (En)
(12)
1=1
Limitations of KL-based functional variational inference Note that the Radon-Nikodym derivative between the variational posterior measure and the prior measure is dQ,
ap, and KL[ Pa)
is defined as log( d&2 aPP )dQ, for the infinite-dimensionaal stochastic processes Gray 2011). In the
case where dQ2, is not absolutely continuous with respect to dPo will lead to k 1 oo (Matthews et al. 2016; Wild et al. 2022), c.g., the KL divergence between stochastic processes defined by two BNNs with different network structures can be infinite (Ma & Hernández-Lobato 2021 Moreover, Sun et al. (2019) proved that KL[Q,]| Po] = supxex, KL[C,((F>))PP)
where X' = UneN [X E Xr X, ASC" R"Xp) is the set of all finite marginal measurement points in the input domain. However, there is no analytical solution for such supremum, and even for finite X, this KL term is intractable for many distributions due to the unavailability of the density (Rudner
al., 2022). Therefore, the KL-based functional variational inference is delicate task that has several limitations in practical applications.
B THEORETICAL. PROOF
B.1 PROOF FOR PROPOSITION 3.1
Based the transformed posterior obtained in the first T-Stage of the main paper for the approx- imate posterior resulting from the second M-Stage, we derive the approximate posteriors for the following alternating stages in detail follows:
V-Stage Based on the transformed Fx; o Q we now can continue to perform the I'I procedure, and the resulting approximate posterior measure denoted as Fx; Q with the corresponding updated posterior parameters qsal (W,3; 0sa)-
14

---

Published as a conference paper at ICLR 2025
M-Stage Similar to the previous M-Stage, continue to run the fSGLD from the randomly sampled initial point was (Muss, Esa) from the V-Stage, and the corresponding initial function Fxz (fea(; was)). Then, the probability measure of the collected function samples Fxy (XX is denoted as Fx; fes and Fx; (:sa(: (:s)) is the random function mapping this stage.
T-Stage In order to perform the IVI procedure in following V-Stage, it is necessary to fit the second invertible bijection Fx1 to transform (fea(: Waa)) based on Fx; (<; Waa)) obatnied from the last V-Stage as Fx; (: formation for probability measure Fx; oQ4,, can be represented by F); o Fx; Q fus with parametric (w 0aa), which is able to run the I'VI smoothly and the resulting updated approximate posterior measure is denoted by Fx; o Ex; o Qs with the corresponding 9a5 (Was: 0.a).
= Fx3 (Fx; ( Then the corresponding trans-
We continue in this fashion by successively alternating between V-Stage and M-Stage, utilizing T-Stage as the intermediary, to create a chain of total K V-Stages and M-Stages. This hybrid ap- proximate process results in final approximate posterior measure that is in composite form: Fxix 11/2 o... Fx; o Fx; o f.x- which is arbitrarily complex and preferably flexible. This composi- tional structure allows for more refined posterior approximation, enhancing both its expressiveness and ability to capture complex distributions.
B.2 PROOF FOR PROPOSITION 3.2
Dynamics-based MCMC sampling techniques are derived from the more general Itô diffusion [Øksendal 2003), which is often used to characterize the evolution of particles. Given a stochastic process 00 20) x Q + I" defined on probability space (S2, P), an Itō diffusion in n- dimensional Euclidean space driven by the standard Wiener process satisties the following specific type of stochastic differential equation (SDE):
dx, = dt + (X,)
(13)
where X₂ is the state of the stochastic process at time t. pa(-) R" + R" a vector field demoting the deterministic drift term, (()) RR -> R=x* is matrix field denotes the diffusion coefficient for X,. Both p(-)) and (-)) are assumed to satisfy the usual Lipschitz continuity condition (Ghosh, 2010). And B E Rn is n-dimensional Wiener process (Brownian motion), dB, is the increment of a Wiener process distributed as dB, ~ (O0 dt I.). Let p(x, t) be the probability density of X,, the evolution of p(z, t) is formed by the following Fokker-Plance (FP) equation (Risken 1996):
3p(x, 1) 0t
82 Or,0n, b.
Σ dz; (n((X))((x,
()]
Dis(s,t)p(c,
(14)
where (i,(x, f) = (x.)))
Maet al. 2015) proved in Theorem that if (xx,) and of (X.)
are restricted to the following form
m(X)= DD(X)
T(X)= = Σ 2
+ ox (X) Qi₂(x)),
(15)
oXX) = -X)((),
where U.) is the potential energy functional, D(.) is a positive semidefinite matrix, and Q(-) is skew-symmetric curl matrix, then the Fokker-Planck equation can be transformed into a more compact form (Yin & Ao 2006; Shi et al. 2012):
0p(æ, t) V" (D(() + Q(X)) + Vy(x,))). Øt
(16)
Wu et al. (2024a) proved that their functional Langevin dynamics in Equation can be cast into the above general framework as
df((:w) -- Hl
Hl. V2 V
(17)
(2 (fff)),B
=
15

---

Published as a conference paper at ICLR 2025
where U(f) = - log p( YDlf w)) + Io(f) is an OM functional for Pyp, D(f)
(V wwJ) (V wf), Q()) 0, r(j) Hw Therefore, the Fokker-Planck equation of the prob-
ability measure p(f, () for fi(:; w) then can be derived as
0p(S.) V". (D() +OII0I vp(f.t)))
(18)
Øt
from which the stationary measure xf) of p(l, can be verified as ()p(--U())) = Pyps that is, the target posterior measure over functions.
The proof of Proposition 3.2 is based on the Wasserstein gradient flows, which an extension version of gradient flows in Euclidean space the space of probability measures. Formally, suppose P(2)) is the space of probability measures on S2 with finite second moments and endowed with Riemannian geometry characterized by the 2- 2-Wasserstein distance (Chen et al. 2018). For functional JF that maps probability measures to real values: P(22) + I. Wasserstein gradient flows describe the evolution of probability measures over time by decreasing the functional F (Yi & Liu, 2023), e.g., F(g) can be the [Tg]ll for probability measure q. Let (P(S2. W2) be metric space of P(2) and be equipped with 2- 2-Wasserstein distance, a curve {gc} E P(22) in the Wasserstein space (P(S2, W2) is the gradient flow of functional FF if it satisfies the following PDE:
3q: div Vw, V. (q
8J (F(a)) ))
(19)
5qc
where F(q) is the Wasserstein gradient of the functional F(q). ( is called the first variation of F at (Ambrosio et al. 2008) and V denotes the Euclidean gradient operation. Let {vr} be family of vector fields space { induced by the Wasserstein gradient and it formulates a probability flow ordinary differential equation (ODE) as
5g
5F((1)) 8g:
dæ, U2 (z,) -
dt,,
(20)
where v(x) VZ 6.F(4)) This ODE characterize the evolution of particle 2, 9i in Q when the 6q
corresponding marginal de evolves to decrease F(4) as in Equation (9
The Wasserstein gradient flows of parameter-space variational inference for posterior inference is to decrease the following functional JF of the approximate posterior q(w)
F(((w)) & U(W)g(W)dw +
q(w) log ((w)dw
(21)
E1
E2
<LLgg (l)p/)/))
where ((w) = log pl Y I[ - Po (w), Feg is the energy functional of a pure Brownian
motion (hhagg et al. 2019h. The first variational of Ej and Eg is
Ty(w] U(w) and SE
4B
ōg[w]
log (( ()) + Then, gradient flow F(g(w)) is as
@q((w)
Vw (c(w)) V (w) + V log
Øt
(22)
Vw (e(w) (w) v (w)(
which reads the Fokker-Planck equation of Langevin dynamics for posterior sampling of network parameters as in Equation (6)) And the corresponding probability ODE is
dw; = - U
V log qe( ((W,)dtf
(23)
which share the same marginals {ge (w) Teso with the Langevin SDE dw; = V if they evolve from the same %o (w) (Yi & Liu 2023).
We now extend this conclusion to the function spaces, that is, to verify whether the functional vari- ation inference and fSGLD share the same marginals. Let P(H) be the space of Borel probability
16

---

Published as a conference paper at ICLR 2025
measures function space H. with finite second moments and endowed with Riemannian geom- etry. The functional Flus for functional variational inference in Equation ①
Ffoi A flognYY
dQ, (n)Q)(()
log dPa
(24)
E1
B2
The Wasserstein gradient flows of the corresponding variational distribution g(w) this Frut is derived as
(q.(w) V V -logp((Y [(Xx; + log aPo (n)]
dQj
(25)
and the corresponding probability ODE for W,
dw; V logp(YD> ft(XDi))) (o)p(Y)l((XX:vvv))) log dQ (fc) (fi) dt
dPo
(26)
= vw V (ogooppYY w)) + log drx (ft) log dQ, dQI((;))) dt, drw.F
dP'
where rH.= denotes the corresponding Hausdorf//Riemaniian measure, and the second equation holds by applying the chain rule for Radon-Nikodyn derivatives (Matthews et al.) (2016). shares the same evolution with the functional Langevin SDE for network parameters (Wu et 2024a)
dw,
V W V2dB.
w IV (- logplo Yo (Xp: w)) + lof((t))) + V2dB,
(27)
Vw (00g w)) + log dPo
(fi)) dt + v2dB,
drw.F
where the OM functional If()) for prior measure Po can be represented by - log dru.F dpo
dfp (f) with
TH.F, and the potential energy functional U(} log (oxp(y///CC (XX w)) log arw (f). More-
over, the Wasserstein gradient flows in Equation 25
corresponds to the following function-space
probability measure evolution (Wang al. 2019
aQj / (/ 01
logg(Yy YYD [[(xi w)) log dQj [n) aPo
(28)
V (V (V
where the second equation recovers the Fokker-Planck equation in Equation 18 of the probability measure p(f,t) t) for fa w) of the functional Langevin SDE in Equation The corresponding probability ODE for is derived
d,f (V (V U(() log dQ;
drw (f)) de,
which, therefore, share the same probability measure evolution marginals with functional Langevin SDE in Equation P Note that this proposition remains valid even under data subsampling. The justification is twofold: 1) The stochastic gradient approximation of the likelihood term does not alter the target functional of the Wasserstein gradient flow in fVI. This ensures that the overall optimization landscape remains unchanged despite subsampling: 2) It is well-established that the stationary distribution of ISGLD converges the target posterior when using decreasing step size. This holds because the stochastic gradient noise decays at faster rate than the injected Gaussian noise, preserving the asymptotic correctness of the inference process. Thus, the validity of this proposition remains independent of the stochastic gradient approximation, ensuring its robustness even when subsampling is applied. The schematically shown in the following Figure
C PSEUDOCODE FOR FVIMC
Algorithm Coorithmmpreeeeee presents the pseudocode for FVIMC.
17

---

Published as a conference paper at ICLR 2025
Algorithm Functional Variational Inference MCMC (FVIMC)
Require: Dataset D (XD, YD), minibatch
{.5.0ūlī= D. functional prior Po. total number of hybrid approximation stages K, U() - log DoggpYYDs(XCC (ols(xp:))) ((())
1: Initialise w00 (uso, Σ.,), reparameterize w =1+ with (0, I), Pso (0,, I),
00 ~ (O, + +, 200 == {1s0: 3
2: while 0 not converged do
3: # V-Stage 1:
4: draw measurement set X.M randomly from input domain
5: draw functional prior functions fxx ~ Po at XM
6: draw neaural network (NN) functions f.M Q XM
7: a0 < Optimizer((Oao. Lw)
8:
0. < 0%
9:
Obtain
qc1(W.1;001)
corresponding
Q
and
approximated
GP(failfal( (3 Mol), Jw.), 2.1
# M-Stage 2:
11:
St x=
randomly sample Wo
NN Ex1)
13:
for / 0 to N do
draw measurement set XMi
15:
Wi+1 < V (fi)) 22 using Equation 3
16:
S SU
17:
end for
18:
Obtain the implicit 4s2 (v2(w,2) corresponding Q
T-Stage 1:
20:
Collect corresponding function samples P.M (1.N)
Qfa2 on XM based the network pa-
rameters samples
21:
Å + Optimizer(AA, L) in Equation (5) based the Qs from the last V-Stage
22:
Obtain the transformed Fx; Q. for Q5a
23:
# Alternating loops
24:
for in r3nge(3, K, 2) do
25:
# V-Stage i:
26:
draw measurement set XMM randomly from input domain
27:
draw functional prior functions pxM Po at X,M
28:
draw NN functions fxM ~ Fxi.1112 13/2 o... Fx o Q at XM
29:
0 + Optimizer((.x-12) Lw)
30:
O.r < 02
31:
Obain q.((Wsi Osi), corresponding Fxt--1/2
Fx Qi and approximated
GP(Ss: [ K do
while Mss)), ) Jnn. ExiTI. (aa) for Qf.:
32:
33:
# M-Stage i 1:
34:
S;
35:
randomly sample wo N(psi, Exi) do
36:
for f
37:
draw measurement set XMi
38:
W + W, V S S SU (w
( 1// (1; (x:))) + v 2e using Equation 3
39:
40:
end for Obtain the implicit qoi (Wei+1), corresponding Fxp-11/2
41:
Fx; Qfain
42:
# T-Stage 1)/2:
43:
Collect corresponding function samples Fy (... (Fxz (f XM
(fiti Ny
))) ~ F Fxie-1y/2
Fx; Qiatt X,M based the network parameters samples S
44:
A(1+1)/2 t Optimizer(2
L) in Equation (5) based on Fx; 1/2 0... Fx; o Qja
obtained from the last V-Stage
45:
Obtain the transformed Exiu-1/2 o Fx 11/2 o
o F2; o Qy for F Exi- o...o F
Q of the M-Stage +
end while
47:
end for
end while
18

---

Published as a conference paper at ICLR 2025
NI
1
TargetPho
-2
Oa
Figure 3: Schematic representation of the same probability measure evolution marginals of ISGLD and the probability ODE derived from the Wasserstein gradient of fVI for posterior inference.
D MORE DETAILS FOR INVERTIBLE TRANSFORMATIONS
The specific forms for Linear transformation follows:
Forward y f(; w) b(x:w),
Inverse: f(x; w) y b(x;w),
(30)
Jacobian
The specific forms for Tanh transformation follows:
Forward: y = 02 Tanh
a1 f (2;; w) + by (x)
+ 62(21;22)
1
arctanh (4 - b2(2) (V-be(zzu2) )
(1r())
Inverse: f(x;
(31)
a1
a2
a1
Tanh(f(a; * as + bI; 21
Jacobian: aj
The specific forms for Sigmoid transformation are follows:
Forward y = ===ag a2 Sigmoid (1 f(x; f(a;w) + b(::),
b(aw)
1 -
Inverse: f(;; w)
log
42
(32)
01
y- by (rrase))
[ax; Jacobian di 02 Sigmoid (ai [f w)
Sigmoid (a1 * (x;; w))]
where f(; w) denotes function mappings, a a scalar parameter and b(z:w) denotes a bias func- tion.
E EXPERIMENTAL SETTING
Extrapolation Synthetic Data In this experiment, we employ two-hidden-layer fully con- nected neural networks with 100 hidden units per layer across all methods. For all functional ap- proaches, the functional GP prior, utilizing an RBF kernel, is pre-trained on 20 training points for 1000 epochs. The bias function used in the bijections during the T-Stage is implemented as 2 x 100 fully connected neural network.
19

---

Published as a conference paper at ICLR 2025
UCI Regression In this experiment, utilize 2 x 10 fully connected neural networks for all methods. Using the RBF kemel,, the functional GP prior pre-trained for 100 epochs. The bias function of the bijections in the T-Stage is formed 2 100 fully connected neural network.
Contextual Bandits In this experiment, all methods use fully connected neural networks (input- 100-100-outpur) The GP prior (RBF kernel) is pre-trained on 1000 randomly sampled points from training data. All methods are trained using the last 4096 input-output tuples in the training buffer with batch size of 64 and training frequency of 64 for each iteration.
Imagine Classification and OOD Detection In this experiment, we use two-hidden-layer fully connected neural networks with 50 hidden units per layer for all methods. The functional prior is birichlet-based GP designed for classification tasks and pre-trained on randomly sampled 1000 points for 500 epochs.
F COMPLETE INFERENCE PROCESS OF FVIMC
The full inference process of FVIMC for the I-D extrapolation experiment is shown in the Figure and 3 the whole process cannot one figure, we split it into two), where the leftmost column is the posteriors of all V-Stages, the middle column shows the posteriors of all M Stages, and the rightmost column is the approximation posterior measures in all T-Stages for the transformation of posterior obtained in all M-Stages. The results are organized in 1-20 alternating V-Stage and M-Stage from top to bottom (the final result from the last V-Stage is shown Figure The final approximate posterior measures in the composited form of Fxjo Fxg Fx; Qpm-
G FURTHER RESULTS
G.1 ANALYSIS OF MIXING TIME
The trajectories of fMCMC in the M-Stage of FVIMC. the parameter-space SGLD and the func- tional ISGLD for the synthetic extrapolation and UCI regression (Yacht) are presented Figure 6 The model used for the extrapolation experiment is two-hidden-layer fully connected neural net- work (100 hidden units in each layer) with (4401-dimesiional network parameters. For the UCI regression, the network is 2 10 fully connected with 191-dimensional parameters. As shown in Figure 6 the top and bottom rows show the results of the toy and UCI experiments, respectively. With the aid of the explicit optimization of fVI in the V-stage, which quickly converges local high-probabiliyy target regions, the fMCMC in the M-stage of our hybrid FVIMC exhibits a rapid mixing rate for high-dimensional parameter space. In the extrapolation example, the fMCMC in the M-stage converges rapidly to the stationary measure within 800 iterations, while fSGLD and SGLD show stabilizing trends until after 2000 iterations. Similarly, for the UCI regression, our FVIMC reaches the stationary in only 400 iterations, while SGLD and fSGLD require more than 500 iterations for mixing.
G.2 COMPUTATIONAL COMPLEXITY
The running time for all 2000 training iterations of all inference methods in the UCI regression on the small Boston dataset and the large Protein dataset are shown in Table 4 There are 455 training points with 13-dimensional input features for the Boston dataset and 41157 training points with 9-dimensional features in the Protein dataset. For all functional inference methods, the GP prior was pre-trained for 100 epochs, which takes only an extra 1s. In general, parameter-space KLBBB and SGLD are faster than the corresponding functional VI and MCMC methods. In the small Boston dataset, our FVIMC shows similar computational efficiency to the other functional [VI, GWI, IFBNN, and is nearly 5 times faster than FBNN. With the fast convergence of the V- Stage to locally high probabilities, our FVIMC demonstrates its computational efficiency advantage on the larger Protein dataset: the rumning time of GWI and FBNN is nearly 1.7-2.5x higher than FVIMC. The mixing rate of ISGLD, which performs well on small datasets, has deteriorated at this time and is significantly slower than hybrid method. Moreover, the convergence processes of
20

---

Published as a conference paper at ICLR 2025
2.4
pocteror (Feriadd (ertee teenel
Tag
35
2.5
35
-35
-35
-55
-16
-13
-15
-15
-13
-Sa
-335
303
0.30
-225
100
-4
-150 -025
C25
(a) V-Stagel
(b) M- M-Stage2
(c) T-Stagel
24
PVINC poste
Tog posteror
Flgs appros maatonn
18-
-35
-5.
-50
-2.30
(d) V-Stage3
(e) M MtStage4
(I) T-Stage2
PVINC poslo
poaderior
asprxx mration
.15-
Pamng
.15
-15
-5. -355 -050
-025
910
030
Ite
-05
<056
-335
330
0.35
110
-2.30
-035
-150
-325
0.5
035
gg) V-Stage5
(h) M-Stage6
(i) T-Stage3
-01
-13
1.0
1.9
-25-
-15
-15
-015
235
-335
312
0.33
0.0
676
100
-A
2036
-335
330
0.35
0.50
635
10
-2.30
-005 -358
-235
0.00
6.5
0.50
635
130
V-Stage7
(k) M-Stage8
()) T-Stage4
2.5
PVIN
pouterio
1.9
16-
13-
15.
35
16-
3.8
33
13
18-
-15
-15
-25
-33
-33
(m) V-Stage9
(n) M-Stage10
(o) T-Stage5
Figure 4: The alternating approximation stages 1-10 in full hybrid inference process of FVIMC.
training loss for all parameter/function--pee methods and the V-Stage of FVIMC are shown in Figure FVIMC achieves high convergence speed and training stability.
21

---

Published as a conference paper at ICLR 2025
2.
PAIC
Togi
1.4
35-
2.5
35-
41-
3.8
-35
-35
-35
-10
-00
-13
-11
-19
-13
-55
-375
-023
303
030
-20
-225
10
-4
-025
C25
(a) V-Stage11
(b) M-Stage 12
(c) T-Stage6
24
PVINC possD
Tog posteror
Flon espros mration
1--
-5.0
-50
-2.30
(d) V-Stage13
(e) M- M-Stagel4
(I) T-Stage7
PVINC poste
1gg poaderoo
asprxx mration
.15-
.15
-L5
-5. -055
-825
910
030
-0
-335
330
0.3
0.50
110
-2.30
-03
-150
-3
0.5
035
V-Stage15
(h) M-Stage16
ii T-Stage8
-0
-13
1.0
1.
-25-
-15
-13
-5
-035
235
-335
312
0.23
0.0
676
100
20
200
-335
330
0.35
0.50
635
10
-2.30
-005 -358
0.00
625
0.50
635
130
V-Stage17
(k) M-Stage18
()) T-Stage9
2.
PVIN
1.9
16-
13
W
15.
3.5
35
16-
3
33
11
18-
-15
-33
-33
(m) V-Stage19
(n) M-Stage20
(0) T-Stage10
Figure 5: The alternating approximation stages 1 11-20 in full hybrid inference process of FVIMC.
22

---

Published as a conference paper at ICLR 2025
0.7
-3.003
10 powweeer
-3.015
0.013
0.005
-0.25
-0.00
.08
-3.23
...05
-0.055
-4.000
-3335
-0.05
1005
2500
755e
12500
Beratin
Meration
Iteration
(a) tyy
(b) ISGLD-toy
(c) SGLD-toy
603
kew
40
24
FVIMC-uei
FSGLD-uci
SGLD-uci
Figure 6: Log-posterion probability versus the number of iterations.
Table 4: Running time comparison on Boston and Protein dataset.
Run time(s)
Boston
Protein
FVIMC
39.000 + 0.000
182.500 1 0.707
fVI
19.000 0.000
104.000 + 2.828
GWI
16.000 1.000
472.667 0.577
FBNN
200.667 = 6.531
318.333 5.774
IFBNN
19.000 + 0.000
78.500 +: 2.121
KLBBB
8.333 + 0.577
8.000 + 0.000
SGLD
5.500 + 0.707
6.500 2.121
ISGLD
9.500 = 0.707
239.500 2.121
(a) FVIMC
(b) IVI
(c) GWI
(d) FBNN
(e) IFBNN
(0 FVIMC
(g) IVI
(h) GWI
(i
FBNN
(i) IFBNN
Figure 7: Convergence of training loss for UCI regressions. The top row the results for the Boston dataset, and the bottom is the results for the Protein dataset.
G.3 CALIBRATION CURVES
The calibration of all parameter/fuuction-piiie inference approaches for the synthetic trapolation example are shown Figure g which assess how well the predicted probabilities of a model match the true frequencies. The horizontal axis represents the predicted cumulative distribu-
23

---

Published as a conference paper at ICLR 2025
Predicted CDF vs Empirical GDF
1.0
Empirical
0.6
KLBBB
0.4
PVIMC
FSNN
M
ISGLD
SOLD
0.
0.6 Predicted CDF
6.4
0.8
1.0
Figure 8: Calibration curves for the synthetic extrapolation. The gray dashed line indicates the perfect calibration.
tion function (CDF), while the vertical axis corresponds to the empirical CDF The 45-degree dashed diagonal line indicates perfect calibration, where FVIMC demonstrates competitive calibration capability.
G.4 TRANSFORMATION ERRORS T-STAGE
To quantify the transformation errors introduced in the T-Stage, we plot the predictive RMSE sults from the posterior in the M-Stage and the transformed posterior obtained in the T-Stage in the UCI regressions (on Yacht and Boston datasets). As shown in the leftmost column Figure (a) and Figure 9(d) the black line represents the RMSE obtained from the M-Stage posterior, the red line denotes the result of the corresponding transformed posterior in the T-Stage, and the blue line is the difference between these two results. With the progress of FVIMC, we can see that the gap between the predicted RMSE by the posterior the M-Stage and the results corresponding to the approximate posterior obtained through the T-Stage transformation is significantly reduced. There- fore, it can be empirically assumed that this transformation error is controllable, and it will gradually converge to 0 as the whole hybrid inference process proceeds. That is, the parameterized posterior obtained by the T-Satge transformation is a reasonable approximation to the implicit posterior of M-Stage. Additionally, the convergence processes of the training losses in the two T-stages in the middle and rightmost columns, respectively, and it can be seen that they converge quickly and smoothly.
G.5 KL DIVERGENCE IN FVIMC
Empirically, we use the KL divergence-based fVI in the V-Stage to compare the KL-based hy- brid FVIMC (denoted as FVIMC(k)) and our alternative Wasserstein distance-based FVIMC used in the main paper on the UCI regressions. The results are shown in the Table [5] where FBNN represents KI.-based pure functional VI. We can see that under our hybrid inference scheme, FVIMC(KI) outperforms the pure functional variational FBNN, demonstrating the advantage of hy- brid inference. However, our FVIMC still achieves superior performance compared to FVIMC(K)), highlighting the benefits of the more flexible Wasserstein distance in capturing complex posterior distributions. Moreover, we present the convergence processes the KL divergence FVIMC(KI) and the Wasserstein distance in FVIMC in Figure 10, where the Wasserstein distance exhibits significantly more rapid and stable convergence.
Table 5: Comparisons between KL--based FVIMC and Wasserstein distance-based FVIMC
FBNN
FVIMC(K)
FVIMC
Yacht
RMSE
1.52
0.92
0.37
NLL
40.77
-0.83
-5.03
Boston
RMSH
1.68
0.99
0.36
NLL
-1.19
-1.62
-3.74
24

---

Published as a conference paper at ICLR 2025
180
180
2.3
160
185
300
109
(a) Yacht
(b) Transform 1
(c) Transform 2
2714
1
800 150
120
130
323
450
20
153
(d) Boston
Transform
(l) Transform 2
Figure 9: Transformation errors in the T-Stage UCI regressions. The top row is the result for the Yacht dataset, and the bottom row is the result for the Boston dataset.
4/J
(a) W1-Yacht
(b) KL-Yacht
(c) WI-Baston
(d) KL-Boston
Figure 10: Convergence processes of the KI. divergence and the Wasserstein distance in FVIMC, where WI denotes the Wasserstein distance, and KL denotes the KL divergence.
For the theoretical guarantee of Proposition 3.2 we observe that in some special cases, KL-based FVIMC and Wasserstein- based FVIMC are equivalent. This suggests that replacing the KI diver- gence with the Wasserstein distance can be interpreted as a relaxation of the original formulation. In the following, we provide a detailed discussion of this relationship.
Let X and Y be two measurable sets, and define P(X) and P(Y) are the sets of all proba- bility measures on and Y, respectively. Similarly, let P(X Y) denote the set of joint probability measures
X x Y. Given two probability measures (distributions) g E P(X) and p P(Y). ) the general dual formulation of the I-Wasserstein distance between and p is given by Wi == ( (g))))ppp))) f(x) distance metric(e.g. Euclidean norm distance), f X + IR and g Y assume that X = (Belavkin) 2018 On the other hand, the KL divergence between g and p can be denoted as KL(glp) EE [log g) E., (log p). Interestingly, this can be seen as a special case of the I-Wasserstein distance when choosing
g(v) I c(x, y)}, where X Y + IR is proper are real functions, and
log q and g 4 logg. Under this condition, the p Wasserstein gradient flows of the objective functional Fkz(g) = KL[g]pp] and (w W (g.p) p) "F div(ge (V, log q: - V. log,)), which reads the Fokker- Planck equation of the corresponding Langevin SDE dæ, V_ log p(r;))d + v2dB, As result, Proposition 3.2 remains valid for FVIMC when using the Wasserstein distance instead of the KL divergence in this special case. Moreover, met-ie.. fxx) g(y) < x - y/ in the general form of the Wasserstein distance-fuure consider- ations and discussions are provided in the following.
share the same continuity equation as 0g.
ensure that the I-Lipschitz continuity condition is
25

---

Published as a conference paper at ICLR 2025
Consider first the special case where q and p are of the same type (e.g., both Gaussian distributions with identical means and variances). Under certain local assumptions, it can be shown that the 1- Lipschitz condition f(x) - gvy) I |z 1 y/ holds. Since the ratio q/p is constant in this case, the rate of change is controlled. Specifically, at this time, f = = log f(x) g(v) = log p(z) - log p(y)- The composition of the logarithmic function with Gaussian distribution is locally 1-L.ischitz. To satisfy the 1-Lipschitz condition, it should have a bounded first derivative as VV log p(x)) =
- 1 by the mean value theorem. Thus, the geg-saussian is 1-Lipschitz in the local range I | S 0². Additionally, we can verify the derivative condition as 1f((()) g'(v)l = | " + * = الت Thus, for |z yl o², the function f(z) - g(y) satisfies 1 I-Lipschitz continuity. This indicates that if the distance between x and is sufficiently small (bounded by the variance o²), the desired Lipschitz property can be satisfied. For more general cases where q and absolutely continuous (differentiable almost everywhere), the Lipschitz condition is influenced by how smoothly q/p varies. If q/p does not exhibit abrupt fluctuations, the behaviour of remains controlled. Additionally, If and are smooth and similar in shape, the variations in log q(x) and (q(y)/(()) log p(y) will be moderate, further supporting the 1-Lipschitz condition. To enforce the Lipschitz condition in practical applications. a gradient norm regularization term can be added to the dual form of the Wasserstein distance (Gulrajani et al. 2017 Gouk et al. 2021) as (f) (170))Epp65 (g) +
(x) g' (1() 1)² where fx(x) ('<x)/((2), g ((0) d'(x)-log p(y)
+ 2(x).p' (w) (1-1og p(2))
p(1)
and À a penalty coefficient.
H ABLATION STUDY
H.1 THE EFFECTS OF THE INVERTIBLE TRANSFORMATIONS
We used two different bijection forms as the base for the auto-selection mechanism in the T-Stage for the probability measure transformation in the 1-D extrapolation experiment: the Linear form and the non-linear Tanh model. To explore the effects of the invertible bijections for the transfor- mation of probability measure after each M-Stage and final on the posterior inference, we consider four other bijection schemes: pure Linear transformation, pure Tanh transformation, pure Sigmoid transformation and the likelihodd-base auto-selection mechanism based on these three different bi- jections. The resulting corresponding FVIMC posteriors are shown in Figure [T] For the pure Linear bijection, the prediction of the FVIMC posterior in the middle unseen range [- 0.25, 0.25] has been destroyed as in Figure T [a) The pure Tanh form in Figure (b exhibits some in-between uncer- tainty quantification pathologies in the non-observation (-0.25, 0.25]. And for the pure Sigmoid scheme in Figure [Ic)), the uncertainty in the three non-observed regions is underestimated, which can not fully cover the ground truth function. In contrast, the posterior from the auto-sclection mech- anism in Figure [Td] significantly outperforms these three mechanisms regarding both predictive accuracy and uncertainty quantification, demonstrating that this log- ligeliihd auto-selection, rather than pre-fixed form, can approximate the posterior more accurately. The bi- jections selected by this automatic mechanism are, in order, the following: [Linear, Tanh, Sigmoid, Linear, Linear, Linear, Linear, Linear, Sigmoid, Linear].
coomiization-based
H.2 THE EFFECTS OF THE KERNEL FUNCTION IN GP PRIOR
Functional Gaussian processes (GPs) priors are able to encode prior knowledge about function prop- erties (e.g., periodicity and smoothness) through corresponding kernel functions. We used the RBF kernel in the 1-D extrapolation experiment, which is suited for modelling polynomial functions. We now consider two alternative kernels: the Matern kernel with the smoothness parameter taking the value of 1/2, which implies that the function will be less smooth; and the Linear kernel, which is not suitable for modelling polynomial oscillatory curves. The results are shown in Figure 12 The top row shows three GP priors with different RBF, Matern and Linear kernels. And the bottom row is the corresponding FVIMC posteriors using these three different GP priors, respectively. Compared to the RBF GP prior in Figure 12(a) the Martern GP prior in Figure 12(b) is slightly less smooth and the variance for the two observation regions| [-0.75, -0.25] and [0.25, 0.75] shows some unrea- sonable expansion. Consequently, the uncertainty estimation of the corresponding FVIMC posterior in Figure [2(c] of these two observation ranges is a little higher, and it can not recover the main trend of the middle unseen region [-0.25, 0.25). The results from the mismatched Linear kernel
26

---

Published as a conference paper at ICLR 2025
2.
PANG
THIN
35-
2.5
3.5
0.5
-35
-35
-55
-1
-15
-19
-13
-15
-075
-023
309
030
-225
110
-4
-156
-023
C25
(a) Pure Linear
(b) Pure Tanh
(c) Pure Sigmoid
2.8
PUNC poster
(d) Auto-sclection
Figure 11: The effects of the bijection forms in the T-Stage.
-26
-13
-19
-15
-15
pm
47s
Die
025
09
120
-28
91
0.
029
191
6.79
030
51
850
F22
195
(a) GP prior-RBi RBF
(b) GP prior- Matern
(c) Linear
2.6
PVING
FINC
2.5
1.8
REAA
45-
88-
-25-
-25
-05
-2.
-1.
-13
-15
-1.
-15
823 213
030
110
-13
C33
(d) FVIMC posteior-RBB
(e) FVIMC poterior-Matern
(1) FVIMC posterior-Linear
Figure 12: The effects of the kernel function in the functional GP prior. The text after the short line in each subheading represents the corresponding kernel functions.
are shown in Figure 12(c) and Figure 12(t) Since the Linear GP prior shows completely lin- ear trend and falls the target function, the predictive accuracy and uncertainty estimation of the corresponding FVIMC posterior deteriorates significantly, which indicates that our method can effectively incorporate functional prior information into posterior inference process.
I NOTATION TABLE
Table is the notation table to demonstrate the notation used this paper.
27

---

Published as a conference paper at ICLR 2025
Table 6: Notation table
Notation
Mcanings
D (Xx, Y
Training dataset
NC R
(p-dimensional) input space
y C Re
(e-dimensional) output space
(0, A, P)
Probability space Rº
dimens-aimeenional function space (Polish space) Borel o-algebra on
B(H)
P(H)
The space of Borel probability measures on ((H)).
X
Finite marginal points
X.
Finite measurement points
w E R*
Random network parameters Variational parameters Parameters of invertible bijections
(M, Σ}
f(; w)
Random function mapping defined by a BNN parameterized by w Prior distribution over network parameters Posterior distribution over network parameters Likelihood function evaluated the training data
Po(w)
p(w D)
p( Y p( w))
((w; 0)
Variational posterior network parameters parametrized by 0 Functional measure
Po
P
Posterior measure over functions
Q5
Approximate posterior measure over functions
Invertible bijection parametrized by with inverse FT
28


Reviews:
[{"review_id": "NDI9nRAco9", "reviewer": "Reviewer_bU3Q", "summary": "This paper proposes a method for functional posterior inference in BNNs, combining the best of functional variational inference and functional MCMC. It alternates the two, with an additional step of fitting a parametric transformation to the MCMC samples to form the analytic approximate posterior for the next round. Experiments show that the method achieve better performance than other BNN techniques.", "questions": "1. Which specific fVI method do you use? I didn't find any details of this in the paper, while the algorithm in appendix is quite inpenetrable with some notation undefined etc. This seems quite an important question, as I don't know how the algorithm works with the complex posterior forms composed of the transformations with Q_fs1.\n\n2. I think Q_fsK in Prop 3.1 should be Q_fs1?\n\n3. The forms of the invertible transformations are very simple. Can they really capture sufficiently complex posteriors? It would be nice to demonstrate this in a simple 2D toy example?\n\n4. I don't think there's a precise statement of Prop 3.2. Can you furnish this? I think it's saying that if you look at gradient flow of the ELBO, you also get the evolution of the marginal distributions under Langevin SDE. I don't think it really says that for a particular (implementable) fVI method, that they are the same, right? I feel that Prop 3.2 is just saying that the fVI and fMCMC are \"going in the same direction\". \n\n5. The transformation step will probably introduce errors, and if we want to careful it is important to quantify these errors, and whether they are controlled or will compound as the algorithm runs.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The method seems technically sound. \n\nThe comparisons to other BNN methods are reasonably broad ranging and convincing.", "weaknesses": "The method and general area are technically very involved. While I have not checked everything, things do make sense at high level though there are some notational inconsistencies.\n\nThe write-up is very involved as a result. I think the method can probably be applied in the more simple non-functional setting, which can make the description more easy to understand and get across.\n\nExperiments show that the method achieve better performance than other BNN techniques. However the paper did not compare against other baselines that are \"not as Bayesian\", though these are often simpler, computationally more efficient, and more performant empirically. These include deep ensembles, SNGPs, DUE, epinets, temperature scaling, or even conformable prediction methods.\n\nIt would also be nice to see experimental results on uncertainty quantification, e.g. reporting expected calibration error for classification, given as UQ is one of the selling points of Bayesian methods.\n\nMy worry with the paper is the technical complexity and computationally expense, that while shown to be better than simpler BNN baselines, has not been shown to be better than non BNN methods on a representative benchmark (e.g. uncertainty quantification). This calls into question the practical impact and usefulness of the method.  Granted this criticism can perhaps be directed at many papers in this area, but I think as researchers it is good for us to be able to answer the question of why/under what circumstance somebody should use the methods/tools that we have developed.", "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": "VvMeDr9YW2", "reviewer": "Reviewer_K6pQ", "summary": "This paper proposes a Bayesian inference method for determining the posterior distribution of neural networks that assume a prior distribution over functions rather than parameters. The specific method introduced is a hybrid between functional variational inference and functional Markov Chain Monte Carlo. This is done by alternating between maximizing the functional ELBO in conjunction with a parametric approximate posterior and taking MCMC samples. The gap between these two paradigms are smoothed over by approximating the empirical distribution of posterior samples with a learned, parametric transformation of the variational distribution.", "questions": "I do not believe this general approach has been taken for parametric Bayesian neural networks, is this the case? If not, do you see any potential for also adapting this scheme (V->M->T->...) for parameter-based inference?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The proposed method does a good job of ameliorating the downsides to both fVI and fMCMC while seemingly retaining the benefits of both. I found the learning of a parametric representation of the empirical distribution (in the T-Stage) with a normalizing flow to be particularly clever at bridging the gap between VI and MCMC. The proposed approach makes a solid contribution towards Bayesian methods for neural networks, as functional priors in general are an attractive approach in this space but previously have been plagued with practical downsides that this technique improves upon.", "weaknesses": "I mainly found parts of the experimental results to be a bit lacking. In the abstract, the paper claims that the proposed method improves upon the restrictive posterior of fVI and the slow mixing times of fMCMC; however, as far as I can tell only the former has been showcased. A more expressive posterior can shown with visualizations (Fig 1) or could be argued with improved predictive performance (rest of Section 5), whereas the mixing times, especially in high dimensions, requires a stand alone investigation. Having another synthetic experiment where the posterior distribution is known (e.g., Gaussian Process) and comparing how many iterations is required to converge to arbitrary precision would demonstrate this well. Additionally, showcasing some loss-curves and other metrics throughout training for the real-world experiments could also showcase this.", "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": "BlKLoC5HLX", "reviewer": "Reviewer_7yD4", "summary": "In this paper, the authors studied Bayesian Neural Networks (BNNs) and proposed a new inference method by combining two existing ones: the functional variational inference (fVI) and functional Markov chain Monte Carlo (fMCMC). Their hybrid approach alters between two stages: a V-stage that performs fVI to maximize the ELBO, and a M-stage that performs stochastic gradient fMCMC. They provide theoretical justification that the functional Langevin SDE and the probability flow ODE share the same probability evolution marginals. In the experimental evaluations, the authors show that their approach outperform against baselines and improved predictive accuracy and uncertainty characterization in multiple tasks.", "questions": "In the section 5.1, the authors said they make fair comparison between their iterative method and the non-iterative methods by specifying certain number of epochs and iterations. I wonder how those number are determined in terms of the same computational budget between all methods?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The idea from this paper is well-motivated and to my best knowledge, this combination of functional variational inference and functional MCMC is a novel approach. In general, this paper is well-structured, and they provided detailed derivation to show that the functional Langevin stochastic differential equation in the fMCMC and the probability flow ordinary differential equation share the same probability marginals. The math derivation is sound in my opinion.", "weaknesses": "I am a bit confused about how to validate the iterations between the V-stage and the M-stage. As the paper mentioned in the page 4, there is a problem when altering between these two stages, because the M-stage outputs a non-parametric distribution while the V-stage requires a parametric one as its input. I wonder how the transformation from $Q_{fs1}$ to $Q_{fs2}$ is defined specifically.", "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": "R4FLANcAq7", "reviewer": "Reviewer_qXT6", "summary": "This paper proposes a novel hybrid posterior inference method that combines the advantages of both functional variational inference and functional MCMC. To frame these two update schemes as an alternating approximation process, the authors aim to learn a bijective transformation of the variational distribution updates by maximizing the log-likelihood functions of the samples collected from functional MCMC steps.\nMoreover, since using KL divergence in functional variational inference (fVI) may be potentially problematic, the authors propose replacing it with the Wasserstein distance in their method, fVIMC. \nIn experiments, the authors validate their method on several benchmark tasks, including the Bayesian neural networks on UCI datasets, exploration-exploitation performance in contextual bandits and classification tasks on MNIST and Fashion MNIST datasets. The results indicate that fVIMC outperforms the primary baseline, like fVI and fMCMC.", "questions": "- Typo: 1. \nIn line 132, the term $\\mathcal{X}^n = \\bigcup \\\\{\\mathbf{X} \\in \\mathcal{X}^n \\mid \\mathcal{X}^n \\in \\mathbb{R}^n \\\\}$ is confusing.\n\n- In line 166, \"this factorized Gaussian will be modified later rather than being fixed, as in fVI, so it does not impose a strict restriction.\" I believe that the choice of variational distribution family $q_{s0}(\\mathbf{w}_{s0},\\boldsymbol{\\theta}_{s0})$ does, in fact, constitute a restriction for fVI. Can you clarify this point more clearly?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper proposes an interesting framework that combines variational inference and MCMC in functional space.\n- The details of toy experiments are well-explained, and the real-data tasks are comprehensive. The discussion of related work is also very clear.", "weaknesses": "The fVIMC method proposed in this paper is natural and interesting, but I have some specific concerns about fVIMC:\n- The mechanism of using the 1-Wasserstein distance to replace KL divergence in functional priors seems appropriate. However, this may cause the variational distribution update in fVI step no longer corresponding to the gradient flow of the functional $KL[q(w)|p(w|D)]$. In that case, would the result in Proposition 3.2 still hold?\n- In the results for BNNs, the RMSE scale seems inconsistent with previous BNN results reported in [1] and [2].\n- In the experimental results, an ablation study comparing KL divergence and 1-Wasserstein distance in fVIMC would be helpful.\n\n\n[1] Liu, Qiang, and Dilin Wang. \"Stein variational gradient descent: A general purpose bayesian inference algorithm.\" Advances in neural information processing systems 29 (2016).\n\n[2] Sun, Shengyang, et al. \"Functional variational Bayesian neural networks.\" ICLR 2019.", "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": "gPwvgD53GQ", "reviewer": "Reviewer_aL6E", "summary": "The paper presents a hybrid functional inference method called Functional Variational Inference MCMC (FVIMC) for posterior approximation. FVIMC alternatives between the V-phase, which optimizes a functional variational objective, and M-phase, which samples a functional Langevin dynamic. The method relies on a learnable transformation to transport the approximate posterior from the M-phase to a density function usable in the V-phase. The paper makes the following claims:\n  1. FVIMC is a novel hybrid method that links functional variational inference (FVI) with functional Markov chain Monte Carlo (FMCMC).\n  2. FVI and FMCMC share the exact evolution of marginals, which justifies the proposed alternating method.\n  3. It provides strong empirical evidence that the proposed method improves predictive performance and uncertainty estimation compared to baselines.", "questions": "1. Proposition 3.2 connects FVI and FMCMC when the likelihood is evaluated on $\\mathcal{D}$; however, the FVIMC uses data-subsampling. Could you clarify whether proposition 3.2 still holds with data subsampling? If this is the case, are there any constraints on the sampling scheme? If not, please justify that FVIMC still targets the same V- and M-Stage evolution of marginals with subsampling.  \n  2. The paper states that FVIMC can incorporate meaningful prior information into the posterior. The experiments are done through empirical Bayes, i.e., using a GP learned from data as a prior; however, it is not done with BBB. Why can't you cannot use empirical Bayes to find the prior for BBB? If you can, consider including this to make the comparison on equal footing.\n  4. Could you please clarify which pathologies \"potential pathologies in deep learning\" (line 147) are referring to? The statement as presented is vague and, therefore, hard to judge.\n  5. On line 54, the paper calls the isotropic Gaussian prior non-informative. Should this be understood as non-interpretable in the sense introduced by Tran et al. (2022)? If so, consider changing it to non-interpretable instead. Otherwise, please show that an isotropic Gaussian is non-informative for BNNs, like a Jeffreys prior would.\n  6. What is the size of the networks used for experimentation? It is only clear that they are two-layered.\n  7. At what scale of BNNs can we use FVIMC feasibly? Can you make an inference on a ResNet-sized network (11-25k params)? \n  8. What model is used in the image classification and OOD detection experiments? This information should be made available in the article or appendix.\n9. From the remarks in \"Avoid the potential problematic KL divergence for fVI,\" the objective variational objective FVIMC optimizes is not from eq. 2 as the KL divergence is replaced with a Wasserstein-1 distance, and a variance constraint is introduced. Does Proposition 3.2 still apply to FVIMC with these alterations? If so, please add a sentence that proposition 3.2 still holds to the Avoid the potential... paragraph. Otherwise, what justifies the validity of FVIMC with the changed objective?\n10. Why does the collapse occur in FVIMC Fig. 1a? Are there ways to address or mitigate it?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Proposition 3.2 is insightful, novel and sound.\n- The evaluation is done against many baseline models, and the UCI results sustain the claim of improved predictive performance.\n- The paper suggests a likelihood heuristic for selecting the transformation. The idea is simple and cheap and aligns with an MLE principle.", "weaknesses": "## Reason for score\n- From the remarks in \"Avoid the potential problematic KL divergence for fVI,\" the objective FVIMC optimizes is not from eq. 2 as the KL divergence is replaced with a Wasserstein-1 distance, and a variance constraint is introduced. It is not clear that proposition 3.2 still applies to FVIMC with these alterations. See question 9.\n-  Proposition 3.2 connects FVI and FMCMC when the likelihood is evaluated on $\\mathcal{D}$; however, the FVIMC uses data-subsampling. It is unclear whether proposition 3.2 holds with data subsampling. See question 1.\n- The proposed method requires running an MCMC chain, optimizing the variational distribution and optimizing a transformation at each iteration. However, the computational cost of their proposed method is not discussed. It would improve the paper to include complexity analysis or empirical evidence comparing runtime and memory overhead to the baselines.\n- The method uses several hyperparameters that interact in a complex manner. These include the kernel for the GP prior, the number of MCMC samples $N$, the convergence criteria for the V-phase, and the set of transformations and their architecture. However, the paper does not address how to choose these. It would greatly benefit from guidelines, supported by experiments or analysis, for specifying the hyperparameters.\n- The paper does not clearly state that the T-stage preserves the shared evolution of posteriors between the V-stage and M-stage. Including a proposition that this holds would strengthen the claim that proposition 3.2 justifies the validity of FVIMC and would improve the paper's score.\n- The extrapolation experiment has an unexplained variance collapse in Fig. 1 at $x=0$ for FVIMC. This is also present in the fig. 7 for the FVIMC posteriors. Notably, this collapse is not present in baseline models. It invalidates the claim that \"The hybrid fig1(a) exhibits stronger ability than the pure variational fVI in Figure 1(e), e.g., the smoother fitting curves, and can capture critical features of the non-observed middle range [−0.25, 0.25],\" because the collapse point is not on the actual curve. Importantly, there is not even posterior mass covering the curve.", "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)