Bayesian Learning via Neural Schr¨odinger-F¨ollmer Flows
Abstract
In this work we explore a new framework for approximate Bayesian inference in large datasets based on stochastic control.
We advocate stochastic control as a ﬁnite time alternative to popular steady-state methods such as stochastic gradient Langevin dynamics (SGLD).
Furthermore, we discuss and adapt the existing theoretical guarantees of this framework and establish connections to already existing VI routines in SDE-based models
1. Introduction
Steering a stochastic ﬂow from one distribution to another across the space of probability measures is a well-studied problem initially proposed in Schr¨odinger (1932).
There has been recent interest in the machine learning community in these methods for generative modelling, sampling, dataset imputation and optimal transport (Wang et al., 2021; De Bortoli et al., 2021;
Huang et al., 2021; Bernton et al., 2019; Vargas et al., 2021; Chizat et al., 2020;
Cuturi, 2013; Maoutsa and Opper, 2021; Reich, 2019).
We consider a particular instance of the Schr¨odinger bridge problem (SBP), known as the Schr¨odinger-F¨ollmer process (SFP).
In machine learning, this process has been proposed for sampling and generative modelling (Huang et al., 2021; Tzen and Raginsky, 2019b) and in molecular dynamics for rare event simulation and importance sampling (Hartmann and
Sch¨utte, 2012; Hartmann et al., 2017); here we apply it to Bayesian inference.
We show that a control-based formulation of the SFP has deep-rooted connections to variational inference and is particularly well suited to Bayesian inference in high dimensions.
This capability arises from the SFP’s characterisation as an optimisation problem and its parametrisation through neural networks (Tzen and Raginsky, 2019b).
Finally, due to the variational characterisation that these methods possess, we believe that many low-variance estimators (Richter et al., 2020; N¨usken and Richter, 2021; Roeder et al., 2017) are applicable to the
SFP formulation we consider.
We reformulate the Bayesian inference problem by constructing a stochastic process Θ t which at a ﬁxed time t = 1 will generate samples from a pre-speciﬁed posterior p ( θ | X ) with dataset X = { x i } Ni =1 (i.e. Law Θ 1 = p
( θ | X )), where the model is given by:
[Equation]
and the prior p ( θ ) and the likelihood p ( x i | θ ) are user-speciﬁed.
Our target is π 1 ( θ )
=
p ( X | θ ) p ( θ ) Z , where Z = (cid:82) p ( x i | θ ) p ( θ ) d θ .
This formulation is reminiscent of the setup proposed by Welling and Teh (2011) and covers many Bayesian machine-learning models, but our formulation has an important diﬀerence.
SLGD relies on a diﬀusion that reaches the posterior as time approaches inﬁnity.
In contrast, our dynamics are controlled and the posterior is reached in ﬁnite time.
We note that recently submitted concurrent work (Anonymous, 2022) proposes an algorithm akin to ours based on Dai Pra (1991); Tzen and Raginsky (2019b), however their focus and experiments are on estimating the normalising constant of unormalised densities, while our focus is on Bayesian ML tasks such as Bayesian regression, classiﬁcation and
LVMs, thus our work leads to diﬀerent insights and algorithmic motivations.
1.1. Schr¨odinger-F¨ollmer Processes
Let Q γ 0 be the distribution for the solutions to the stochastic diﬀerential equation (SDE):
[Equation]
Deﬁnition 1 (Schr¨odinger-Bridge Process)
The Schr¨odinger bridge path measure is given by
[Equation]
where D ( π 0 , π 1 ) = { Q : ( Θ 0 )
# Q = π 0 , ( Θ 1 )
# Q = π 1 } is the set of path measures with ﬁxed initial and ﬁnal time-marginals ( π 0 at t = 0 and π 1 at t = 1 ).
Here, Q γ 0 acts as a “prior” and D KL ( ·|· ) represents the Kullback Leibler (KL) divergence.
It is known (L´eonard, 2013), that Q ∗ is induced by an SDE with modiﬁed drift,
[Equation]
the solution of which is called the Schr¨odinger-Bridge Process (SBP).
Deﬁnition 2 (Schr¨odinger-F¨ollmer Process)
The SFP is an SBP where π 0 = δ 0 and Q γ 0 =
W γ is the Wiener measure associated to the process described by d Θ t = √ γd B t , Θ 0 ∼ δ 0 .
The SFP diﬀers to the general SBP in that, rather than constraining the initial value to δ 0 , the SBP considers any initial distribution π 0 .
The SBP also considers more general
Itˆo SDEs associated with Q γ 0 as the dynamical prior, compared to the SFP which considers only a Wiener process as a prior.
The advantage of considering this more limited version of the SBP is that it admits a closed-form characterisation of the solution to the Schr¨odinger system (L´eonard, 2013;
Wang et al., 2021; Pavon et al., 2018), which allows for an unconstrained formulation of the problem.
For accessible introductions to the SBP we suggest (Pavon et al., 2018; Vargas et al., 2021).
Now we will consider instances of the SBP and the SFP where π 1 = p ( θ | X ).
1.1.1. Analytic Solutions and the Heat Semigroup
Prior work (Pavon, 1989; Dai Pra, 1991; Tzen and Raginsky, 2019b; Huang et al., 2021) has explored the properties of SFPs via a closed form formulation of the F¨ollmer drift expressed in terms of expectations of Gaussian random variables known as the heat semigroup.
The seminal works (Pavon, 1989; Dai Pra, 1991; Tzen and Raginsky, 2019b) highlight how this formulation of the F¨ollmer drift characterises an exact sampling scheme for a target distribution and how it could potentially be used in practice.
The recent work by Huang et al.
(2021) builds on Tzen and Raginsky (2019b) and explores estimating the optimal drift in practice via the heat semigroup formulation using a Monte Carlo approximation.
Our work aims to take the next step and scale the estimation of the F¨ollmer drift to high dimensional cases (Graves, 2011; Hoﬀman et al., 2013).
In order to do this we must move away from the heat semigroup and instead consider the dual formulation of the F¨ollmer drift as a stochastic control problem (Tzen and Raginsky, 2019b).
Deﬁnition
The Euclidean heat semigroup Q γt , t ≥ 0 acts on bounded measurable functions f : R d → R as
[Equation]
In the setting where π 0 = δ 0 we can express the optimal SBP drift as follows:
[Equation]
where Θ t is the SBP prior Q γ 0 .
In the SFP case where Q γ 0 = W γ , the optimal drift can be written in terms of the heat semigroup, u ∗ t ( x ) = ∇ x ln Q γt (cid:104)
[Equation]
γ
I
[Equation]
( x ) (cid:105) .
Note that an SDE with the above drift, d Θ u ∗ t
[Equation]
that is at t = 1 these processes are distributed according to our target distribution of interest π 1 .
Huang et al.
(2021) carried out preliminary work on empirically exploring the success of using the heat semigroup formulation of SFPs in combination with the EulerMayurama (EM) discretisation to sample from target distributions in a method they call
Schr¨odinger-F¨ollmer samplers (SFS).
We build on their work by considering a formulation of the Schr¨odinger-F¨ollmer process that is suitable for the high dimensional settings arising in Bayesian ML.
Our work will focus on a dual formulation of the optimal drift that is closer to variational inference and admits the scalable and ﬂexible parametrisations used in ML.
2. Stochastic Control Formulation
In this section, we introduce a particular formulation of the Schr¨odinger-F¨ollmer process in the context of the Bayesian inference problem in Equation 1.
In its most general setting of sampling from a target distribution, this formulation was known to Dai Pra (1991).
Tzen and Raginsky (2019b) study the theoretical properties of this approach in the context of generative models (Kingma et al., 2021; Goodfellow et al., 2014), ﬁnally Opper (2019) applies this formulation to time series modelling.
In contrast our focus is on the estimation of a Bayesian posterior for a broader class of models than Tzen and Raginsky explore.
Corollary 4 The minimiser (where U is the space of admissible controls)
[Equation]
satisﬁes Law Θ u ∗ 1 = p
( X | θ ) p ( θ ) Z , where d Θ u t = u t ( Θ u
[Equation]
The objective in Equation 6 can be estimated using an SDE discretisation, such as the
EM method.
Since the drift u ∗ t is Markov, it can be parametrised by a ﬂexible function estimator such as a neural network, as done in Tzen and Raginsky (2019a,b).
In this work we will refer to the above formulation of the SFP as the Neural Schr¨odinger-F¨ollmer sampler (NSFS) when we parametrise the drift with a neural network and we implement unbiased mini-batched estimator for this objective detailed in Appendix C.This formulation of SFPs has been previously studied in the context of generative modelling/marginal likelihood estimation (Tzen and Raginsky, 2019b), while we focus on Bayesian inference.
2.1. Theoretical Guarantees for Neural SFS
While the focus in Tzen and Raginsky (2019b) is in providing guarantees in generative models of the form x ∼ q φ ( x | Z 1 ) , d Z t =
u φ ( Z t , t ) dt + √ γd B t , Z 0 =0 , their results extend to our setting as they explore approximating the F¨ollmer drift for a generic target π 1 .
Theorem 4 in Tzen and Raginsky (restated as Theorem 9 in Appendix A.2) motivates using neural networks to parametrise the drift in Equation 6 as it provides a guarantee regarding the expressivity of a network parametrised drift via providing an upper bound on the target distribution error in terms of the size of the network.
We will now proceed to highlight how this error is aﬀected by the EM discretisation:
Corollary
Given the network v from Theorem 9 it follows that the Euler-Mayurama discretisation
ˆ Θ v t of Θ v t has a KL-divergence to the target distribution π 1 of:
[Equation]
This result provides us a bound of the error in terms of the depth ∆ t − 1 of the stochastic ﬂow and the size of the network that we parametrise the drift with.
2.2. Structured SVI in Models with Local and Global Variables
We consider the general setting where our model has global and local variables { θ i } , Φ satisfying θ i ⊥⊥ θ j | Φ (Hoﬀman et al., 2013).
This case is particularly challenging as the local variables scale with the size of the dataset and so will the state space.
This is a fundamental setting as many hierachical latent variable models in machine learning admit such dependancy structure, such as Topic models (Pritchard et al., 2000; Blei et al., 2003);
Bayesian factor analysis (Amari et al., 1996; Bishop, 1999; Klami et al., 2013; Daxberger and
Hern´andez-Lobato, 2019); Variational GP Regression (Hensman et al., 2013); and others.
Remark
The heat semigroup does not preserve conditional independence structure in the drift, i.e. the optimal drift does not decouple and thus depends on the full state-space.
Remark 6 tells us that the drift is not structured in a way that admits scalable sampling approaches such as stochastic variational inference (SVI)
(Hoﬀman et al., 2013).
Additionally this also highlights that the method by Huang et al.
(2021) does not scale to models like this as the dimension of the state space will be linear in the size of the dataset.
In a similar fashion to Hoﬀman and Blei (2015), who focussed on structured SVI, we suggest parametrising the drift via [ u t ] θ i =
u θ i ( t, θ i , Φ , x i ); this way the dimension of the drift depends only on their respective local variables and the global variable Φ.
While the F¨ollmer drift does not admit this particular decoupling we can show that this drift is ﬂexible enough to represent general distributions, thus it has the capacity to reach the target distribution.
When parametrised in this form we can carry out sampling in the same fashion as SVI whilst maintaining unbiased gradient estimates.
Remark 7 An SDE parametrised with a decoupled drift [ u t ] θ i =
u θ i ( t, θ i , Φ , x i ) can reach transition densities which do not factor.
3. Connections to Variational Inference in Latent Diﬀusion Models
In this section, we highlight the connection between the objective in Equation 6 to variational inference in models where the latent object is given by an SDE, as studied in Tzen and Raginsky (2019a).
Taking inspiration from the recursive nature of Bayesian updates (Khan and Rue, 2021)
we develop the following observation.
Lemma 8 The SBP inf Q ∈D ( δ 0 , p ( θ | X ))
D KL (cid:0) Q (cid:12)(cid:12)(cid:12)(cid:12) Q γ 0 (cid:1) with reference process Q γ 0 :
[Equation]
[Equation]
corresponds to maximising the ELBO of the model:
[Equation]
In short we can view a variant of the objective in Equation 6 as an instance of variational
Bayesian inference with an SDE prior.
Note that this provides a succinct connection between variational inference and maximum entropy in path space (L´eonard, 2012).
4. Experimental Results
We ran experiments on Bayesian NN regression, classiﬁcation, logistic regression and ICA (Amari et al., 1996), reporting accuracies, log joints (Welling and Teh, 2011; Izmailov et al., 2021) and expected calibration error (ECE) (Guo et al., 2017).
For details on exact experimental setups please see Appendix F. Across experiments we compare to SGLD as [Table]
[Table]
[Table] [Table]
it has been shown to be a competitive baseline in Bayesian deep learning (Izmailov et al., 2021)
[Table]
In the Bayesian NN tasks the likelihood is parametrised via p ( y i | x i , θ )
= p ( y i | f θ ( x i )) where f θ is a neural network and the prior on θ is Gaussian.
Step Function : We ﬁt a 2-hidden-layer neural network with a total of 14876 parameters on a toy step function dataset (see Figure 1).
We can see in Figure how both the SGD and
SGLD ﬁts interpolate the noise, whilst N-SFS has straight lines and thus achieves a better test error, whilst having well calibrated error bars.
We believe it is a great milestone to see how an overparametrised neural network is able to achieve such well calibrated predictions.
Augmented MNIST :We train the standard LeNet5 (LeCun et al., 1998) architecture (with 44426 parameters) on the MNIST dataset (LeCun and Cortes, 2010).
At test time we evaluate the methods on the MNIST test set augmented by random rotations of up to 30 ◦ (Ferianc et al., 2021).
Table 2 shows how N-SFS has the highest accuracy whilst obtaining the lowest calibration error among the considered methods, highlighting that our approach has the most well calibrated and accurate predictions when considering a slightly perturbed test set.
We highlight how LeNet5 falls into an interesting regime as the number of parameters is considerably less than the size of the training set, and thus we can argue it is not in the overparametrised regime.
This regime (Belkin et al., 2019) has been shown to be challenging in achieving good generalisation errors, thus we believe the predictive and calibrated accuracy achieved by N-SFS is a strong milestone.
a9a/MEG Datasets: Following Welling and Teh (2011) we explore a logistic regression model on the a9a dataset and a Bayesian variant of ICA on the MEG-Dataset (Vigario, 1997).
We can observe (Tables 4, 1) that N-SFFS achieves results comparable to SGLD.
5. Discussion and Future Directions
We are competitive to SGLD and obtain better calibrated predictions in high-dimensional (Bayesian DL) settings.
We would like to highlight that these results were achieved without any tuning and simple NN architectures.
We believe that our results illustrate how stochastic control objectives constitute a promising and exciting avenue for approximate Bayesian inference.
We are currently exploring the decoupled drift parametrisations (Section 2.2) for models such as LDA as well as bench-marking the performance of diﬀerent estimators such as STL (Xu et al., 2021) and VarGrad (N¨usken and Richter, 2021).
Additionally we notice that the architecture used in the drift network can inﬂuence results thus we will explore modern architectures for diﬀusions such as the score networks in De Bortoli et al. Acknowledgements.
Francisco Vargas is Funded by
Huawei Technologies Co.
This research has been partially funded by Deutsche Forschungsgemeinschaft (DFG) through the grant CRC 1114 ‘Scaling Cascades in Complex Systems’ (project A02, project number 235221301).
Andrius Ovsianas is funded by EPSRC iCASE Award EP/T517677/1. Mark Girolami is supported by a Royal Academy of Engineering Research Chair and EPSRC grants EP/T000414/1, EP/R018413/2, EP/P020720/2, EP/R034710/1, EP/R004889/1. Appendix A. Main Results A.1. Posterior Drift
Corollary
The minimiser
[Equation]
satisﬁes Law Θ u ∗ 1 = p
( X | θ ) p ( θ ) Z , where
[Equation]
and
[Equation]
Proof This follows directly after substituting the Radon-Nikodym derivative between the
Gaussian distribution and the posterior into Theorem 1 in Tzen and Raginsky (2019b) or
Theorem 3.1 in Dai Pra (1991).
A.2. EM-Discretisation Result
First we would like to introduce the following auxiliary theorem from Tzen and Raginsky (2019b):
Theorem 9 (Tzen and Raginsky, 2019b)
Given the standard regularity assumptions presented for f = dπ 1 d N ( 0 ,γ I ) in Tzen and Raginsky (2019b), let L = max { Lip( f ) , Lip( ∇ f ) } and assume that there exists a constant c ∈ (0 , 1] such that f ≥ c .
Then for any 0 < (cid:15) <
16 L 2 cthere exists a neural net v : R d × [0 , 1] → R d with size polynomial in 1 /(cid:15), d, L, c, 1 /c, γ , such that the activation function of each neuron follows the regularity assumptions in Tzen and
Raginsky (2019b)
(e.g. ReLU , Sigmoid , Softplus ) and
[Equation]
where π v 1 =
Law( Θ v 1 ) is the terminal distribution of the the diﬀusion process
[Equation]
We can now proceed to prove the direct corollary of the above theorem when using the EM scheme for simulation.
Corollary
Given the network v from Theorem 9 it follows that the Euler-Mayurama discretisation ˆ X v t of X v t has a KL-divergence to the target distribution π 1 of:
[Equation]
Proof Consider the path-wise KL between the exact Schr¨odinger-F¨ollmer process and its
EM-discretised neural approximation:
[Equation]
Deﬁning d ( x , y ) :
= (cid:113) 12 γ (cid:82) 10 E || x
( Θ u ∗ t ,t ) − ˆ y
( Θ u ∗ t ,t )
||
2 dt , it is clear that d ( x , y ) satisﬁes the triangle inequality as it is the L 2 ( Q u ∗ ) metric between drifts, thus applying the triangle inequality at the drift level we have that (for simplicitly let γ = 1):
[Equation]
From Tzen and Raginsky (2019b) we can bound the ﬁrst term resulting in:
[Equation]
Now remembering that the EM drift is given by ˆ v √ 1 − t ( Θ t )
= v ( ˆ Θ t , (cid:112) 1 − ∆ t
cid:100) t/ ∆ t (cid:101) ), we can use that v is L’-Lipschitz in both arguments thus:
[Equation]
which using the strong convergence of the EM approximation (Gy¨ongy and Krylov, 1996) implies:
[Equation]
thus:
[Equation]
squaring both sides and applying the data processing inequality completes the proof.
Appendix B. Connections to VI
We ﬁrst start by making the connection in a simpler case – when the prior of our Bayesian model is given by a Gaussian distribution with variance γ , that is p ( θ ) =
N ( θ | 0 , γ I d ).
Observation 1 When p ( θ ) =
N ( θ | 0 , γ I d ) , it follows that the N-SFP objective in Equation 6 corresponds to the negative ELBO of the model:
[Equation]
Proof Substituting p ( θ ) into Equation 6 yields
[Equation]
Then, from (Bou´e and Dupuis, 1998; Tzen and Raginsky, 2019a,b) we know that the term
E (cid:104)(cid:82) 10 || u t || 2 dt − ln p ( X | Θ 1 ) (cid:105) is the negative ELBO of the model speciﬁed in Equation 17.
While the above observation highlights a speciﬁc connection between N-SFP and VI, it is limited to Bayesian models that are speciﬁed with Gaussian priors.
To extend the result, we take inspiration from the recursive nature of Bayesian updates in the following result.
Lemma 8 The SBP
[Equation]
with reference process Q γ 0 :
[Equation]
corresponds to maximising the ELBO of the model:
[Equation]
Proof For brevity let u 0 t ( θ )
=
∇ ln Q γt (cid:104)
[Equation]
First notice that the time-one marginals of Q γ 0 are given by the Bayesian prior:
[Equation]
Now from L´eonard (2012); Pavon et al.
we know that the Schr¨odinger system is given by:
[Equation]
[Equation]
where Equation 22 can be given a rigorous meaning in weak form (that is, by integrating against suitable test functions).
Notice φ 0 = δ 0 and thus it follows that:
[Equation]
[Equation]
then by Pavon (1989); Dai Pra (1991); Pavon et al.
the optimal drift is given by:
[Equation]
where the expectation is taken with respect to the reference process Q γ 0 .
Now if we let v ( θ, t ) = − ln E
[ p ( X | Θ 1 ) | Θ t = θ ] be our value function then via the linearisation of the Hamilton-Bellman-Jacobi Equation through Fleming’s logarithmic transform (Kappen, 2005; Thijssen and Kappen, 2015; Tzen and Raginsky, 2019b)
it follows that said value function satisﬁes:
[Equation]
and thus u ∗ t ( θ ) =
γ ∇ ln E [ p ( X | Θ 1 ) | Θ t
= θ ] is a minimiser to:
[Equation]
Note Lemma 8 induces a new method in which we ﬁrst estimate a prior reference process’s as in Equation
and then we optimise the ELBO for the model in 9, this raises the question on what eﬀect the dynamical prior can have within SBP based frameworks.
Data: data set X
= { x i } Ni =1 , initialized drift neural network u φ , parameter dimension d , number of iterations M , batch size B , number of Euler-Maruyama discretization steps k , diﬀusion coeﬃcient γ .
∆ t ← 1 k t j ← j ∆ t for all j = 0 , . . .
, k for i = 1 , . .
, M do
Initialize Θ s 0 ← 0 ∈ R d for all s = 1 , . .
, S { Θ sφj } kj =1 ← Euler-Maruyama( u φ , Θ s 0 , ∆ t ) for all s = 1 , . . .
, S
Sample x r 1 , . .
, x r B ∼ X
Compute g ← ∇ φ (cid:18)
S (cid:80) Ss =1 (cid:80) kj =
(cid:18)
||
u φ ( Θ sφj , t j ) ||
− ln (cid:18) p
( Θ sφk )
N ( Θ sφk | 0 ,γ I d ) (cid:19) + NB (cid:80)
Bj =1 ln p ( x r j | Θ sφk ) (cid:19)(cid:19)
φ ← Gradient Step( φ, g ) end return u φ
Appendix C. Stochastic Variational Inference
For a Bayesian model having the structure speciﬁed by equation 1 the objective in equationcan be written as follows:
[Equation]
where the last term can be written as:
[Equation]
That is, it is possible to estimate the objective (and its gradients) by subsampling the data with random batches of size B and using the scaling NB .
A version of the algorithm with
Euler-Maruyama discretization of the SDE is given in Algorithm 1.
Appendix D. Decoupled Drift Results
First let us consider the setting where the local variables are fully independent, that is, θ i ⊥⊥ θ j .
Remark 10 The heat semigroup preserves fully factored (mean-ﬁeld) distributions thus the
F¨ollmer drift is decoupled.
In this setting we can parametrise the dimensions of the drift which correspond to local variables in a decoupled manner, [ u t ] θ i =
u θ i ( t, θ i , x i ).
This amortised parametrisation (Kingma and Welling, 2013) allows us to carry out gradient estimates using a mini-batch (Hoﬀman et al., 2013) rather than hold the whole state space in memory.
Remark
The heat semigroup does not preserve conditional independence structure in the drift.
That is, the optimal drift does not decouple and as a result depends on the full state space.
Proof Consider the following distribution:
[Equation]
We want to estimate:
[Equation]
where X, Y, Z ∼ N (0 , √ 1 − t )
[Equation]
we can easily see that the above no longer has conditional independence structure and thus when taking its logarithmic derivative the drift does not decouple.
Remark 7 An SDE parametrised with a decoupled drift [ u t ] θ i = u ( t, θ i , Φ , x i ) can reach transition densities which do not factor.
Proof Consider the linear time-homogeneous SDE:
[Equation]
where:
[Equation]
then this SDE admits a closed form solution:
[Equation]
which is a Gauss-Markov process with 0 mean and covariance matrix:
[Equation]
We can carry out the matrix exponential through the eigendecomposition of A , for simplicity let us consider the 3-dimensional case:
[Equation]
From this we see that:
[Equation]
Integrating wrt to s yields:
[Equation]
The covariance matrix is dense at all times and thus the density Law( Θ t ) =
N ( µ ( t ) , Σ ( t )) does not factor (is a fully joint distribution).
This example motivates that even with the decoupled drift we can reach coupled distributions.
Appendix E. Sticking the Landing and Low Variance Estimators
As with VI (Richter et al., 2020; Roeder et al., 2017), the gradient of the objective in this study admits several low variance estimators (N¨usken and Richter, 2021; Xu et al., 2021).
In this section we formally recap what it means for an estimator to “stick the landing” and we prove that the estimator proposed in Xu et al.
satisﬁes said property.
The full objective being minimised in our approach is:
[Equation]
noticing that in previous formulations we have omitted the Itˆo integral as it has zero expectation (but the integral appears naturally through Girsanov’s theorem).
We call the estimator calculated by taking gradients of the above objective the relative-entropy estimator.
The estimator proposed in Xu et al.
(Sticking the landing estimator) is given by:
[Equation]
where ⊥ means that the gradient is stopped/detached as in Xu et al.
(2021); Roeder et al.
We study perturbations of F around u ∗ by considering u ∗ + ε φ , with φ arbitrary, and ε small.
More precisely, we set out to compute:
[Equation]
through which we deﬁne the deﬁnition of sticking the landing:
Deﬁnition
We say that an estimator sticks the landing when
[Equation]
almost surely, for all smooth and bounded perturbations φ .
Notice that by construction, u ∗ is a global minimiser of J , and hence all directional derivatives vanish,
[Equation]
Deﬁnition 11 additionally demands that this quantity is zero almost surely, and not just on average.
Consequently, “sticking the landing”-estimators will have zero-variance at u ∗ .
Remark 12 The relative-entropy stochastic control estimator does not stick the landing.
Proof
See N¨usken and Richter (2021), Theorem 5.3.1, clause 3, Equation 133 clearly indicates
[Equation]
We can now go ahead and prove that the estimator proposed by Xu et al.
(2021) does indeed stick the landing.
Theorem 13 The STL estimator proposed in (Xu et al., 2021) satisﬁes:
[Equation]
[Equation]
almost surely, for all smooth and bounded perturbations φ .
Proof Let us decompose F in the following way:
[Equation]
where (denoting the terminal cost with g ):
[Equation]
From N¨usken and Richter (2021), Theorem 5.3.1, Equation 133 it follows that:
[Equation]
almost surely, where A t is deﬁned as
[Equation]
and satisﬁes:
[Equation]
Similarly via the chain rule it follows that:
[Equation]
(57) almost surely, combining these results we can see that dd ε F ( u ∗ + ε φ ) (cid:12)(cid:12)(cid:12) ε
=
= 0 almost surely as required.
Appendix F. Experimental Details and Further Results F.1. Method Hyperparameters
In Table F.1 we show the experimental conﬁguration of the trialled algorithms across all datasets. [Table] λ ( i ) = a ( i + b )
[Table] λ ( i ) = a ( i + b ) γ λ ( i ) = a ( i + b ) γ λ ( i ) =
a ( i + b
[Table]
F.2. Step Function Dataset
Here we describe in detail how the step function dataset was generated:
[Equation]
Where:
• σ y = 0 .•
N train = 100, N test =• x train ∈ ( − 3 .
5 , 3 . 5) • x test ∈ ( − 10 , 10)
F.3. F¨ollmer Drift Architecture
Across all experiments we used the same architecture to parametrise the F¨ollmer drift: 1 class SimpleForwardNetBN(torch.nn.
Module):  def __init__(self , input_dim=1, width =20): super(SimpleForwardNetBN , self).__init__()  self.input_dim = input_dim  self.nn = torch.nn.
Sequential( torch.nn.
Linear(input_dim + 1, width), torch.nn.
BatchNorm1d(width , affine=False), torch.nn.
Softplus(), torch.nn.
Linear(width , width), torch.nn.
BatchNorm1d(width , affine=False), torch.nn.
Softplus(), torch.nn.
Linear(width , width), torch.nn.
BatchNorm1d(width , affine=False), torch.nn.
Softplus(), torch.nn.
Linear(width , width), torch.nn.
BatchNorm1d(width , affine=False), torch.nn.
Softplus(), torch.nn.
Linear(width , input_dim) )  self.nn[-1] weight data fill
(0.0) self.nn
[Table]
Note the weights and biases of the ﬁnal layer are initialised to 0 in order to start the process at a Brownian motion matching the SBP prior.
F.4. BNN Architectures
For the step function dataset we used the following architecture: 1 class DNN_StepFunction(torch.nn.
Module):  def __init__(self , input_dim=1, output_dim =1): super(DNN , self).__init__()  self.output_dim = output_dim self.input_dim = input_dim  self.nn = torch.nn.
Sequential( torch.nn.
Linear(input_dim , 100), torch.nn.
ReLU(), torch.nn.
Linear(100, 100), torch.nn.
ReLU(), torch.nn.
Linear(100, output_dim) )
Listing 2: Architecture for step function dataset.
For LeNet5 the exact architecture used was: 1 class LeNet5(torch.nn.
Module):  def __init__(self , n_classes): super(LeNet5 , self).__init__()  self.feature_extractor = torch.nn.
Sequential( torch.nn.Conv2d
( in_channels=1, out_channels=6, kernel_size=5, stride=  ), torch.nn.
Tanh(), torch.nn.
AvgPool2d(kernel_size =2), torch.nn.
Conv2d
( in_channels=6, out_channels=16, kernel_size=5, stride=  ), torch.nn.
Tanh(), torch.nn.
AvgPool2d(kernel_size =2), )  self.classifier = torch.nn.
Sequential( torch.nn.
Linear(in Listing 3: Architecture for MNIST. [Table]
F.5. Likelihood and Prior Hyperparameters
In Table F.5 we describe the hyperparameters of each Bayesian model as well as their priors and likelihood.