Particle Flow Bayes' Rule

TL;DR

This paper introduces a neural ODE-based particle flow method for Bayesian inference that generalizes across different priors and observations, enabling efficient sequential Bayesian updates.

Contribution

It presents a novel neural ODE operator for particle flow Bayes' rule, capable of generalizing across various priors and observations, trained via meta-learning.

Findings

Successfully applied to high-dimensional examples

Demonstrated generalization across different priors and observations

Proved the existence of the neural ODE operator

Abstract

We present a particle flow realization of Bayes' rule, where an ODE-based neural operator is used to transport particles from a prior to its posterior after a new observation. We prove that such an ODE operator exists. Its neural parameterization can be trained in a meta-learning framework, allowing this operator to reason about the effect of an individual observation on the posterior, and thus generalize across different priors, observations and to sequential Bayesian inference. We demonstrated the generalization ability of our particle flow Bayes operator in several canonical and high dimensional examples.

Figures

Tables

Table 1. Table 1: Numbers are averaged over 25 sequences with 25 observations each. For PFBR, gpu is faster when #particles is larger.

Algo	#particles	cpu time (s)	gpu time (s)	cross-entropy
PFBR	256	0.23	0.26	16.56
SMC	256	0.07	0.02	26.78
ASMC-mlp	256	0.17	0.07	19.66
ASMC-gru	256	0.18	0.07	19.38
ASMC-mlp	4096	2.23	0.25	17.63
ASMC-gru	4096	2.26	0.26	17.24
SMC	8192	3.87	0.12	17.60

Equations

$$\stackrel{\text{updated posterior}}{\overbrace{p(\mathbf{x}|{\mathcal{O}}_{m+1})}}\propto \stackrel{\text{current posterior}}{\overbrace{p(\mathbf{x}|{\mathcal{O}}_m)}}\stackrel{\text{likelihood}}{\overbrace{p(o_{m+1}|\mathbf{x})}}.$$

$$\stackrel{\text{updated embedding}}{\overbrace{{\mu }_{m+1}}}=\mathcal{K}(\stackrel{\text{current embedding}}{\overbrace{{\mu }_m}},o_{m+1}).$$

$$\stackrel{\text{updated particle}}{\overbrace{{\mathbf{x}}_{m+1}^n}}$$

$$\left\{\begin{tabular}[]{l}$\frac{d{\bm{x}}}{dt}=f({\mathcal{X}}_{0},o_{1},{\bm{x}}(t),t),~{}\forall t\in(0,T]$\cr${\bm{x}}(0)={\bm{x}}_{0}^{n}$\end{tabular}\right.\xRightarrow{\text{gives}}{\bm{x}}_{1}^{n}={\bm{x}}(T).$$

$${\mathbf{x}}_{m+1}^n$$

$$:={\mathbf{x}}_m^n+{\int }_0^{T}f({\mathcal{X}}_m,o_{m+1},\mathbf{x}(t),t)dt.$$

$$\partial q(\mathbf{x},t)/\partial t=-{\nabla }_x\cdot (qf),$$

$$d\log (q(\mathbf{x},t))/dt=-{\nabla }_x\cdot f.$$

$$dq/dt=\partial q/\partial t+{\nabla }_{x}q\cdot d\mathbf{x}/dt.$$

$$\log q_{m+1}(x_{m+1}^n)=\log q_m(x_m^n)-{\int }_0^T{\nabla }_x\cdot fdt.$$

$$d\mathbf{x}(t)=$$

$$\partial q/\partial t=$$

$$+{\Delta }_{x}q(\mathbf{x},t),$$

$$q(\mathbf{x},\infty )=e^{-\Psi }/Z=p(\mathbf{x}|{\mathcal{O}}_m)p(o_{m+1}|\mathbf{x})/Z.$$

$$f={\nabla }_x\log p(\mathbf{x}|{\mathcal{O}}_m)p(o_{m+1}|\mathbf{x})-{\nabla }_x\log q(\mathbf{x},t)$$

$$\underset{w}{\min }$$

$$\frac{d\mathbf{x}}{dt}={\nabla }_x\log p(\mathbf{x}|{\mathcal{O}}_m)p(o_{m+1}|\mathbf{x})-w,$$

$${\nabla }_x\log p(\mathbf{x}|{\mathcal{O}}_m)p(o_{m+1}|\mathbf{x})-w^{\ast }(p(\mathbf{x}|{\mathcal{O}}_m),t),$$

$$f=\mathbf{h}\left(\frac{1}{N}{\sum }_{n=1}^N\varphi ({\mathbf{x}}_m^n),o_{m+1},\mathbf{x}(t),t\right),$$

$$\mathcal{T}:=\left(\pi (\mathbf{x}),p(\cdot |\mathbf{x}),{\mathcal{O}}_M:=\{o_1,\dots ,o_M\}\right)\in {\mathcal{D}}_{train}$$

$$\tau :=(p(\mathbf{x}|{\mathcal{O}}_m),p(\cdot |\mathbf{x}),o_{m+1})\in (\pi (\mathbf{x}),p(\cdot |\mathbf{x}),{\mathcal{O}}_M).$$

$$\mathcal{L}(\mathcal{T})=\sum _{m=1}^M\sum _{n=1}^N\left(\log q_m({\mathbf{x}}_m^n)-\log p({\mathbf{x}}_m^n,{\mathcal{O}}_m)\right).$$

$$\hat{\pi }(\mathbf{x};{\mathcal{X}}_m)=\frac{1}{N}{\sum }_{n=1}^N\frac{1}{\sqrt{2\pi }{\sigma }^d}{e}^{-\frac{\Vert \mathbf{x}-{\mathbf{x}}_m^n{\Vert }^2}{2{\sigma }^2}},$$

$$(\pi (\mathbf{x}),p(\cdot |\mathbf{x}),o_{1:m})\cup (\hat{\pi }(\mathbf{x};{\mathcal{X}}_m),p(\cdot |\mathbf{x}).o_{m+1:M}),$$

$$\mathcal{L}({\mathcal{D}}_{train})=\frac{1}{|{\mathcal{D}}_{train}|}\sum _{\mathcal{T}\in {\mathcal{D}}_{train}}\mathcal{L}(\mathcal{T}),$$

$$\log p({\mathbf{x}}_m^n,{\mathcal{O}}_m)=\log \pi ({\mathbf{x}}_m^n)+{\sum }_{t=1}^m\log p(o_t|{\mathbf{x}}_m^n).$$

$$f=\mathbf{h}\left(\frac{1}{N}{\sum }_{n=1}^N\varphi ({\mathbf{x}}_m^n),\frac{1}{L}{\sum }_{l=1}^{L}g(o_{m+1}^l),\mathbf{x}(t),t\right).$$

$$\log p(\mathbf{x},{\mathcal{O}}_m)\approx \log \hat{\pi }({\mathbf{x}}_m^n;{\mathcal{X}}_{m^{\ast }})+\sum _{t=m^{\ast }+1}^m\log p(o_t|{\mathbf{x}}_m^n).$$

$${\text{MMD}}^2={\mathbb{E}}_{x,x^{\prime }\sim p;y,y^{\prime }\sim q}[k(x,x^{\prime })-2k(x,y)+k(y,y^{\prime })];$$

$${\mathbf{x}}_m=A{\mathbf{x}}_{m-1}+{ϵ}_m,o_m=B{\mathbf{x}}_m+{\delta }_m,$$

$$\text{(i)}{\tilde{\mathbf{x}}}_m^n=A{\mathbf{x}}_{m-1}^n+{ϵ}_m,\text{(ii)}{\mathbf{x}}_m^n=\mathcal{F}({\tilde{\mathcal{X}}}_m,{\tilde{\mathbf{x}}}_m^n,o_{m+1}),$$

$${\sum }_{n=1}^N\log q_m({\mathbf{x}}_m^n)-\log p(o_m|{\mathbf{x}}_m^n)-\log \hat{\pi }({\mathbf{x}}_m^n;{\tilde{\mathcal{X}}}_m),$$

Full text

ODE, Bayesian Inference, Particle Flow, Sequential Monte Carlo, ICML

1 Introduction

In many data analysis tasks, it is important to estimate unknown quantities ${\bm{x}}\in\mathbb{R}^{d}$ from observations ${\mathcal{O}}_{m}:=\{o_{1},\cdots,o_{m}\}$. Given prior knowledge $\pi({\bm{x}})$ and likelihood functions $p(o_{t}|{\bm{x}})$, the essence of Bayesian inference is to compute the posterior $p({\bm{x}}|{\mathcal{O}}_{m})\propto\pi({\bm{x}})\prod_{t=1}^{m}p(o_{t}|{\bm{x}})$ by Bayes’ rule. For many nontrivial models, the prior might not be conjugate to the likelihood, making the posterior not in a closed form. Therefore, computing the posterior often results in intractable integration and poses significant challenges. Typically, one resorts to approximate inference methods such as sampling (e.g., MCMC) (Andrieu et al., 2003) or variational inference (Wainwright & Jordan, 2003).

In many real problems, observations arrive sequentially online, and Bayesian inference needs be performed recursively,

[TABLE]

That is the estimation of $p({\bm{x}}|{\mathcal{O}}_{m+1})$ should be computed based on the estimation of $p({\bm{x}}|{\mathcal{O}}_{m})$ obtained from the last iteration and the presence of the new observation $o_{m+1}$. It therefore requires algorithms which allow for efficient online inference. In this case, both standard MCMC and variational inference become inefficient, since the former requires a complete scan of the dataset in each iteration, and the latter requires solving an optimization for every new observation. Thus, sequential Monte Carlo (SMC) (Doucet et al., 2001; Balakrishnan & Madigan, 2006) or stochastic approximations, such as stochastic gradient Langevin dynamics (Welling & Teh, 2011) and stochastic variational inference (Hoffman et al., 2012), are developed to improve the efficiency. However, SMC suffers from the path degeneracy problems in high dimensions (Daum & Huang, 2003; Snyder et al., 2008), and rejuvenation steps are designed but may violate the online sequential update requirement (Canini et al., 2009; Chopin et al., 2013). Stochastic approximation methods are prescribed algorithms that cannot exploit the structure of the problem for further improvement.

To address these challenges, the seminal work of Kernel Bayes Rule (KBR) views the Bayes update as an operator in reproducing kernel Hilbert spaces (RKHS) which can be learned and directly produce the posterior from prior after each observation (Fukumizu et al., 2012). In the KBR framework, the posterior is represented as an embedding $\mu_{m}:=\mathbb{E}_{p({\bm{x}}|{\mathcal{O}}_{m})}[\phi({\bm{x}})]$ using a feature map $\phi(\cdot)$ associated with a kernel function; then the kernel Bayes operator ${\mathcal{K}}(\cdot,o)$ will take this embedding as input and produce the embedding of the updated posterior,

[TABLE]

Another novel aspect of KBR method is that it contains a training phase and a testing phase, where the structure of the problem at hand (e.g., the likelihood) is taken into account in the training phase, and in the testing phase, the learned operator ${\mathcal{K}}$ will directly operate on the current posterior $\mu_{m}$ to produce the output. However, despite the nice concepts of KBR operator, it only works well for a limited range of problems due to its strong theoretical assumptions.

In this work, we aim to lift the limitation of KBR operator, and will design a novel continuous particle flow operator ${\mathcal{F}}$ to realize the Bayes update, for which we call it particle flow Bayes’ rule (PFBR). In the PFBR framework (Fig. 1), a prior distribution $\pi({\bm{x}})$, or, the current posterior $\pi_{m}({\bm{x}}):=p({\bm{x}}|{\mathcal{O}}_{m})$ is approximated by a set of $N$ equally weighted particles ${\mathcal{X}}_{m}=\{{\bm{x}}_{m}^{1},\ldots,{\bm{x}}_{m}^{N}\}$; and then, given an observation $o_{m+1}$, the flow operator ${\mathcal{F}}({\mathcal{X}}_{m},{\bm{x}}_{m}^{n},o_{m+1})$ will transport each particle ${\bm{x}}_{m}^{n}$ to a new particle ${\bm{x}}_{m+1}^{n}$ to approximate the new posterior $p({\bm{x}}|{\mathcal{O}}_{m+1})$. That is,

[TABLE]

where ${\mathcal{X}}_{m+1}=\{{\bm{x}}_{m+1}^{1},\ldots,{\bm{x}}_{m+1}^{N}\}$ will be used as samples from the new posterior $p({\bm{x}}|{\mathcal{O}}_{m+1})$. Furthermore, this PFBR operator ${\mathcal{F}}$ can be applied recursively to ${\mathcal{X}}_{m+1}$ and a new observation $o_{m+2}$ to produce ${\mathcal{X}}_{m+2}$, and so on.

In a high-level, we model the PFBR operator ${\mathcal{F}}$ as a continuous deterministic flow, which propagates the locations of particles and the values of their probability density simultaneously through a dynamical system described by ordinary differential equations (ODEs). A natural question is whether a fixed ODE-based Bayes operator applicable to different prior distributions and likelihood functions exists. In our paper, we resolve this important theoretical question by making a novel connection between PFBR operator and the Fokker-Planck equation of Langevin dynamics. The proof of existence also provides a basis for our parameterization of PFBR operator using DeepSet (Zaheer et al., 2017).

Similar to KBR, PFBR have a training phase and a testing phase. However, the training procedure is very different as it adopts a meta-learning framework (Andrychowicz et al., 2016), where multiple related Bayesian inference tasks with different priors and observations are created. PFBR operator ${\mathcal{F}}$ will learn from these tasks how to update the posteriors given new observations. During test phase, the learned PFBR will be applied directly to new observations without either re-optimization or storing previous observations. We conduct various experiments to show that the learned PFBR operator can generalize to new Bayesian inference tasks.

Related work. There is a recent surge of interests in ODE-based Bayesian inference (Chen et al., 2018; Zhang et al., 2018; Grathwohl et al., 2018; Lei et al., 2017). These works focus on fitting a single target distribution. Consequently, the learned flow can not generalize directly to a new dataset, a new prior or to sequential setting without re-optimization.

2 Bayesian Inference as Particle Flow

We present details in this section from four aspects: (1) How to map sequential Bayes inference to particle flow? (2) What is the property of such particle flow? (3) Does a shared flow-based Bayesian operator ${\mathcal{F}}$ exist? (4) How to parameterize the flow operator ${\mathcal{F}}$?

2.1 PFBR: Particle Flow Bayes’ Rule

The problem mapping from sequential Bayesian inference to particle flow goes as follows. Initially, $N$ particles ${\mathcal{X}}_{0}=\{{\bm{x}}_{0}^{1},\ldots,{\bm{x}}_{0}^{N}\}$ are sampled i.i.d. from a prior $\pi({\bm{x}})$. Given an observation $o_{1}$, the operator ${\mathcal{F}}$ will transport the particles to ${\mathcal{X}}_{1}=\{{\bm{x}}_{1}^{1},\ldots,{\bm{x}}_{1}^{N}\}$ to estimate the posterior $p({\bm{x}}|{\mathcal{O}}_{1})\propto\pi({\bm{x}})p(o_{1}|{\bm{x}})$. We define this transformation as the solution of an ODE. That is, $\forall n$,

[TABLE]

The flow velocity $f$ takes observation $o_{1}$ as input and determines both direction and speed of the change of ${\bm{x}}(t)$. In this ODE model, each particle ${\bm{x}}_{0}^{n}$ sampled from the prior gives an initial value ${\bm{x}}(0)$, and then the flow velocity $f$ will evolve the particle continuously and deterministically. At terminate time $T$, we will take solution ${\bm{x}}(T)$ as the transformed particle ${\bm{x}}_{1}^{n}$ for estimating the posterior.

Applying this ODE-based transformation sequentially as new observations $o_{2},o_{3},\ldots$ arrive, we can define a recursive particle flow Bayes operator, called PFBR, as

[TABLE]

The set of obtained particles ${\mathcal{X}}_{m+1}$ can be used to perform Bayesian inference such as estimating the mean and quantifying the uncertainty of any test function by averaging over these particles.

At this moment, we assume $f$ has a form of $f({\mathcal{X}},o,{\bm{x}}(t),t)$, and will be shared across different sequential stages. In section 2.3, a rigorous discussion on the existence of a shared flow velocity of this form will be made. Next, we will discuss further properties of this continuous particle flow which will help us study the existence of such operator for Bayesian inference, and design the parameterization and the learning for the flow velocity.

2.2 Property of Continuous Deterministic Flow

The continuous transformation of ${\bm{x}}(t)$ described by ODE $d{\bm{x}}/dt=f$ defines a deterministic flow for each particle. Let $q({\bm{x}},t)$ be the probability density of the continuous random variable ${\bm{x}}(t)$. The change of this density is also determined by $f$. More specifically, $q$ follows the continuity equation (Batchelor, 2000):

[TABLE]

where $\nabla_{x}\cdot$ is the divergence operator. Continuity equation is the mathematical expression for the law of local conservation of mass - mass can neither be created nor destroyed, nor can it ”teleport” from one place to another.

Given continuity equation, one can describe the change of log-density by another differential equation (Theorem 2.1).

Theorem 2.1**.**

If $d{\bm{x}}/dt=f$, then the change in log-density follows the differential equation (Chen et al., 2018)

[TABLE]

Since for any physical quantity $q({\bm{x}},t)$, the distinguish between material derivative $d/dt$ and partial derivative $\partial/\partial t$ is important, we clarify the definition before the proof of this theorem.

Definition 2.1**.**

Material derivative of $q({\bm{x}},t)$ is defined as

[TABLE]

Note that $dq/dt$ defines the rate of change of $q$ in a given particle as it moves along its trajectory ${\bm{x}}={\bm{x}}(t)$ in the flow, while $\partial q/\partial t$ means the rate of change of $q$ at a particular point ${\bm{x}}$ that is fixed in the space.

Proof of Theorem 2.1.

By continuity equation, $\frac{\partial q}{\partial t}=-\nabla_{x}q\cdot f-q\nabla_{x}\cdot f\Rightarrow\frac{dq}{dt}=-q\nabla_{x}\cdot f$. By chain rule, we have $\frac{d\log q}{dt}=\frac{1}{q}\frac{dq}{dt}=\frac{1}{q}(-q\nabla_{x}\cdot f)=-\nabla_{x}\cdot f$. ∎

Theorem 2.1 gives the same result as the Instantaneous Change of Variables Theorem stated by Chen et al. (2018). However, our statement is more accurate using the notation of material and partial derivatives. Our proof is simpler and intuitively clearer using continuity equation. This also helps us to see the connection to other physics problems such as fluid dynamics and electromagnetism.

With theorem 2.1, we can compute the log-density of the particles by integrating across $(0,T]$ for each $n$:

[TABLE]

2.3 Existence of Flow-based Bayes’ Rule

Does a unified flow velocity $f$ exist for different Bayesian inference tasks involving different priors and observations? If it does, what is the form of this function? These questions are non-trivial even for simple Gaussian case.

For instance, let the prior $\pi(x)={\mathcal{N}}(0,\sigma_{x})$ and the likelihood $p(o|x)={\mathcal{N}}(x,\sigma)$ both be one dimensional Gaussian distributions. Given an observation $o=0$, the posterior distribution is ${\mathcal{N}}(0,(\sigma\cdot\sigma_{x})/(\sigma+\sigma_{x}))$. It means the ODE $d{\bm{x}}/dt=f$ needs to push a zero mean Gaussian distribution with covariance $\sigma_{x}$ to another zero mean Gaussian distribution with covariance $(\sigma\cdot\sigma_{x})/(\sigma+\sigma_{x})$ for any $\sigma_{x}$. It is not clear whether such a unified flow velocity function $f$ exists and what is the form for it.

To resolve the existence issue, we will first establish a connection between the deterministic flow in Section 2.2 and the stochastic flow: Langevin dynamics. Then we will leverage the connection between closed-loop control and open-loop control to show the existence of a unified $f$.

2.3.1 Connection to Stochastic Flow

Langevin dynamics is a stochastic process

[TABLE]

where $d{\bm{w}}(t)$ is a standard Brownian motion. Given a fixed initial location ${\bm{x}}(0)$, multiple runs of the Langevin dynamics to time $t$ will result in multiple random locations of ${\bm{x}}(t)$ due to the randomness of ${\bm{w}}(t)$. This stochastic flow is very different in nature comparing to the deterministic flow in Section 2.2, where a fixed location ${\bm{x}}(0)$ will always end up with the same location ${\bm{x}}(t)$.

Nonetheless, while Langevin dynamics is a stochastic flow of a continuous random variable ${\bm{x}}(t)$, the probability density $q({\bm{x}},t)$ of ${\bm{x}}(t)$ follows a deterministic evolution according to the associated Fokker-Planck equation (Jordan et al., 1998)

[TABLE]

where $\Delta_{x}=\nabla_{x}\cdot\nabla_{x}$ is the Laplace operator. Furthermore, if the so-called potential function $\Psi({\bm{x}}):=-\log p({\bm{x}}|{\mathcal{O}}_{m})p(o_{m+1}|{\bm{x}})$ is smooth and $e^{-\Psi}\in L^{1}(\mathbb{R}^{d})$, the Fokker-Planck equation has a unique stationary solution in the form of a Gibbs distribution (Jordan et al., 1998),

[TABLE]

Clearly, this stationary solution is the posterior distribution $p({\bm{x}}|{\mathcal{O}}_{m+1})$.

Now we will rewrite the Fokker-Planck equation in Eq. (2.3.1) into the form of the deterministic flow in Eq. (5) from Section 2.2, and hence identify the corresponding flow velocity.

Theorem 2.2**.**

Assume the deterministic transformation of a continuous random variable ${\bm{x}}(t)$ is $d{\bm{x}}/dt=f$, where

[TABLE]

and $q({\bm{x}},t)$ is the probability density of ${\bm{x}}(t)$. If the potential function $\Psi$ is smooth and $e^{-\Psi}\in L^{1}(\mathbb{R}^{d})$, then $q({\bm{x}},t)$ converges to $p({\bm{x}}|{\mathcal{O}}_{m+1})$ as $t\rightarrow\infty$.

Proof.

By continuity equation in Eq. (5), the probability density $q({\bm{x}},t)$ of ${\bm{x}}(t)$ satisfies {\partial q}/{\partial t}=-\nabla_{x}\cdot\big{(}q\big{(}-\nabla_{x}\Psi({\bm{x}})-\nabla_{x}\log q({\bm{x}},t)\big{)}\big{)}. It is easy to see this equation is the same as the Fokker-Planck equation in Eq. (2.3.1), by decomposing the Laplace as $\Delta_{x}q=\nabla_{x}\cdot(q\nabla_{x}\log q)$. Under the conditions for the potential function $\Psi$, since Fokker-Planck equation has a unique stationary distribution $q({\bm{x}},\infty)$ equal to the posterior distribution $p({\bm{x}}|O_{m+1})$, the deterministic flow in Eq. (12) will also converge to $p({\bm{x}}|O_{m+1})$.

∎

The implication of Theorem 2.2 is that we can construct a deterministic flow of particles to obtain the posterior and hence establish the existence. However, the flow velocity in Eq. (12) depends on the intermediate density $q({\bm{x}},t)$ which changes over time. This seemingly suggests that $f$ can not be expressed as a fixed function of $p({\bm{x}}|{\mathcal{O}}_{m})$ and $p(o_{m+1}|{\bm{x}})$. In the next section, we show that this dependence on $q({\bm{x}},t)$ can be removed using theory of optimal control for deterministic systems.

2.3.2 Closed-loop to Open-loop Conversion

Now the question is: whether the term $\nabla_{x}\log q({\bm{x}},t)$ in Eq. (12) can be made independent of $q({\bm{x}},t)$, or whether there is a equivalent form which can achieve the same flow. To investigate this question, we consider the the following deterministic optimal control problem

[TABLE]

where $d$ can be any metric defined on the set of densities over $\mathbb{R}^{d}$. In Eq. (13), we are optimizing over $w$, which is usually called the control. By Theorem 2.2, $w=\nabla_{x}\log q({\bm{x}},t)$ is apparently an optimal solution. Furthermore, the corresponding flow velocity derived by Fokker-Planck equation can be regarded as the continuous steepest descent of KL$(q({\bm{x}},t)||p({\bm{x}}|{\mathcal{O}}_{m+1}))$ under Wasserstein distance (Jordan et al., 1998). We are seeking an alternative expression to the above optimal solution which only depends on $p({\bm{x}},{\mathcal{O}}_{m})$ and $p(o_{m+1}|{\bm{x}})$. First, we introduce the terminology below from optimal control literature.

Definition 2.2**.**

In optimal control literature, a control in a feed-back form $w=w(q({\bm{x}},t),t)$ is called closed-loop. In contrast, another type of control $w=w(q({\bm{x}},0),t)$ is called open-loop. An open-loop control is determined when the initial state $q({\bm{x}},0)$ is observed, whereas, a closed-loop control can adapt to the encountered states $q({\bm{x}},t)$.

Theorem 2.3**.**

For the optimal control problem in Eq. (13) and Eq. (14), there exists an open-loop control $w^{*}=w^{*}(q({\bm{x}},0),t)$ such that the induced state $q^{*}({\bm{x}},t)$ satisfies $q^{*}({\bm{x}},\infty)=p({\bm{x}}|{\mathcal{O}}_{m+1})$. Moreover, $w^{*}$ has a fixed expression with respect to $p({\bm{x}}|{\mathcal{O}}_{m})$ and $p(o_{m+1}|{\bm{x}})$ across different $m$.

Proof.

By Theorem 2.2, $\tilde{w}^{*}(q({\bm{x}},t),t):=\nabla_{x}\log q({\bm{x}},t)$ can induce the optimal state $\tilde{q}^{*}({\bm{x}},\infty)=p({\bm{x}}|{\mathcal{O}}_{m+1})$ and achieve a zero loss, $d=0$. Hence, $\tilde{w}^{*}$ is an optimal closed-loop control for this problem.

Although in general closed-loop control has a stronger characterization to the solution, in a deterministic system like Eq. (14), the optimal closed-loop control and the optimal open-loop control will give the same control law and achieve the same optimal loss (Dreyfus, 1964). Hence, there exists an optimal open-loop control $w^{*}=w^{*}(q({\bm{x}},0),t)$ such that the induced state also gives a zero loss and thus $q^{*}({\bm{x}},\infty)=p({\bm{x}}|{\mathcal{O}}_{m+1})$. More details are provided in Appendix A to express $w^{*}$ as a fixed function of $q({\bm{x}},0)$, $p({\bm{x}}|{\mathcal{O}}_{m})$, $p(o_{m+1}|{\bm{x}})$ and $t$. ∎

Conclusion of a unified $f$. In sequential Bayesian inference, we will set $q({\bm{x}},0)$ as $p({\bm{x}}|{\mathcal{O}}_{m})$. Therefore, Theorem 2.3 shows that there exists a fixed and deterministic flow velocity $f$ of the form

[TABLE]

which can transform $p({\bm{x}}|{\mathcal{O}}_{m})$ to $p({\bm{x}}|{\mathcal{O}}_{m+1})$ and in turns define a unified particle flow Bayes operator ${\mathcal{F}}$.

2.4 Parametrization

We design a practical parameterization of $f$ based on the expression of the unified flow $f$ in Eq. (15).

(i) $p({\bm{x}}|{\mathcal{O}}_{m})\Rightarrow{\mathcal{X}}_{m}$: Since we do not have full access to the density $p({\bm{x}}|{\mathcal{O}}_{m})$ but have samples ${\mathcal{X}}_{m}=\{{\bm{x}}_{m}^{1},\ldots,{\bm{x}}_{m}^{N}\}$ from it, we can use these samples as surrogates for $p({\bm{x}}|{\mathcal{O}}_{m})$. A related example is feature space embedding of distributions (Smola et al., 2007), $\mu_{{\mathcal{X}}}(p):=\textstyle{\int_{{\mathcal{X}}}}\phi({\bm{x}})p({\bm{x}})\,d{\bm{x}}\approx\textstyle{\frac{1}{N}\sum_{n=1}^{N}}\phi({\bm{x}}^{n}),~{}{\bm{x}}^{n}\sim p$. Ideally, if $\mu_{{\mathcal{X}}}$ is an injective mapping from the space of probability measures over ${\mathcal{X}}$ to the feature space, the resulting embedding can be treated as a sufficient statistic of the density and any information we need from $p({\bm{x}}|{\mathcal{O}}_{m})$ can be preserved. Hence, we represent $p({\bm{x}}|{\mathcal{O}}_{m})$ by $\frac{1}{N}\sum_{n=1}^{N}\phi({\bm{x}}_{m}^{n})$, where $\phi(\cdot)$ is a nonlinear feature mapping to be learned. Since we use a neural version of $\phi(\cdot)$, this representation can also be regarded as a DeepSet (Zaheer et al., 2017).

(ii) $p(o_{m+1}|{\bm{x}})\Rightarrow(o_{m+1},{\bm{x}}(t))$: In both Langevin dynamics and Eq. (15), the only term containing the likelihood is $\nabla_{x}\log p(o_{m+1}|{\bm{x}})$. Consequently, we can use this term as an input to $f$. In the case when the likelihood function is fixed, we can also simply use the observation $o_{m+1}$, which results in similar performance in our experiments.

Overall we parameterize the flow velocity as

[TABLE]

where ${\bm{h}}$ and $\phi$ are neural networks (See specific architecture we use in Appendix C.1). Let $\theta\in\Theta$ be their parameters which are independent of $t$. From now on, we write $f=f_{\theta}({\mathcal{X}}_{m},o_{m+1},{\bm{x}}(t),t)$. In the next section, we will propose a meta learning framework for learning these parameters.

3 Learning Algorithm

Since we aim to learn a generalizable Bayesian operator, we need to create multiple inference tasks as the training set and design the corresponding learning algorithm.

Multi-task Framework. The training set ${\mathcal{D}}_{train}$ contains multiple inference tasks. Each task is a tuple

[TABLE]

which consists of a prior distribution, a likelihood function and a sequence of $M$ observations. A task with $M$ sequential observations can also be interpreted as a sequence of $M$ sub-tasks with 1 observation:

[TABLE]

Therefore, each task is a sequential Bayesian inference and each sub-task corresponds to one step Bayesian update.

Cumulative Loss Function. For each sub-task we define a loss $\text{KL}(q_{m}({\bm{x}})||p({\bm{x}}|{\mathcal{O}}_{m+1}))$, where $q_{m}({\bm{x}})$ is the distribution transported by ${\mathcal{F}}$ at $m$-th stage and $p({\bm{x}}|{\mathcal{O}}_{m+1})$ is the target posterior (see Fig. 1 for illustration). Meanwhile, the loss for the corresponding sequential task will be $\textstyle{\sum_{m=1}^{M}}\text{KL}(q_{m}({\bm{x}})||p({\bm{x}}|{\mathcal{O}}_{m}))$, which sums up the losses of all intermediate stages. Since its optimality is independent of normalizing constants, it is equivalent to minimize the negative evidence lower bound (ELBO)

[TABLE]

The above expression is an empirical estimation using particles ${\bm{x}}_{m}^{n}$. In each iteration during training phase, we will randomly sample a task from ${\mathcal{D}}_{train}$ and update the PFBR operator ${\mathcal{F}}$ by the loss gradient.

3.1 Training Tasks Creation

Similar to some meta learning problems, the distribution of training tasks will essentially affect learner’s performance on testing tasks (Dai et al., 2017). Depending on the nature of the Bayesian problem, we propose two approaches to construct the multitask training set: one is data driven, and the other is based on generative models. The general principle is that the collection of training priors have to be diverse enough or representative of those may be seen in testing time.

Data-Driven Approach. We can use the posterior distribution obtained from last Bayesian inference step as the new prior distribution. If the posterior distribution has an analytical (but unnormalized) expression, we will directly use this expression.

More precisely, since each Bayesian inference step will generate a set of particles, ${\mathcal{X}}_{m}=\{{\bm{x}}_{m}^{1},\ldots,{\bm{x}}_{m}^{N}\}$, corresponding to the posterior distribution $p({\bm{x}}|{\mathcal{O}}_{m})$, we can apply a kernel density estimator on top of these samples to obtain an empirical density function

[TABLE]

where $\sigma$ is the kernel bandwidth. Then this density function with a set of samples from it will be used as the new prior for the next training task. This approach has two attractive properties. First, it does not require any prior knowledge of the model and thus is generally applicable to most problems. Second, it provides us a way of breaking a long sequence $(\pi({\bm{x}}),p(\cdot|{\bm{x}}),{\mathcal{O}}_{M})$ with large $M$ into multiple tasks with shorter sequences

[TABLE]

This will help make the training procedure more efficient. This approach will be particularly suited to sequential Bayesian inference and is used in later experiments.

Generative Model Approach. Another possibility is to sample priors from a flexible generative model, such as a Dirichlet process Gaussian mixture model (Antoniak, 1974). We will leave the experimental evaluation of this approach for future investigation.

3.2 Overall Learning Algorithm

Learning the PFBR operator ${\mathcal{F}}$ is learning the parameters $\theta$ of $f_{\theta}$ to minimize the following loss for multiple tasks

[TABLE]

where ${\mathcal{L}}({\mathcal{T}})$ is the loss for a single task defined in Eq. (17). Both the particle location ${\bm{x}}_{m}^{n}$ and its density value $q_{m}({\bm{x}}_{m}^{n})$ in ${\mathcal{L}}({\mathcal{T}})$ are results of forward propagation determined by $f_{\theta}$ according to ODEs in Eq. (4) and Eq. (8). The training procedure is similar to other deep learning optimizations, except that an ODE technique called adjoint method (Chen et al., 2018) is used to compute the gradient $d{\mathcal{L}}/d\theta$ with very low memory cost. The overall algorithm under a meta-learning framework is summarized in Algorithm 1.

3.3 Efficient Learning Tricks

We introduce two techniques to improve training efficiency for large scale problems. Its application to an experiment which contains millions of data points is demonstrated in Section 4. These two techniques can be summarized as mini-batch embedding and sequence-segmentation.

The loss function ${\mathcal{L}}({\mathcal{T}})$ in Eq. (17) contains a summation from $m=1$ to $M$ and also a hidden inner-loop summation

[TABLE]

Thus the evaluation cost of ${\mathcal{L}}({\mathcal{T}})$ is quadratic with respect to the length $M$ of the observation sequences. Therefore, we need to reduce the length for large scale problems.

Mini-batch embedding. Previously we defined $o_{m}$ as a single observation. However, we can also view it as a batch of $L$ observations, i.e., $o_{m}=\{o_{m}^{1},\ldots,o_{m}^{L}\}$. Each Bayesian update corresponding to this mini-batch will become $p({\bm{x}}|{\mathcal{O}}_{m+1})\propto p({\bm{x}}|{\mathcal{O}}_{m})\prod_{l=1}^{L}p(o_{m}^{l}|{\bm{x}})$. If we rewrite $p(o_{m}|{\bm{x}})=\prod_{l=1}^{L}p(o_{m}^{l}|{\bm{x}})$, this is essentially the same as our previous expression. Therefore, we can replace $o_{m}$ by $\{o_{m}^{l}\}_{l=1}^{L}$ and input these samples simultaneously as a set to the flow $f_{\theta}$. To reduce the input dimension, we resort to a set-embedding $\frac{1}{L}\sum_{l=1}^{L}g(o_{m}^{l})$, where $g$ is a neuralized nonlinear function to be learned. Depending on the structure of the model, one define the embedding as a Deepset (Zaheer et al., 2017), or, if the posterior is not invariant with respect to the order of the observations (e.g., hidden Markov models), one need to use a set-embedding that is not order invariant. To conclude, for a mini-batch Bayesian update, the parameterization of flow velocity can be modified as

[TABLE]

Sequence-segmentation. We will use the approach in Eq. (19) to break a long sequence into short ones. More precisely, suppose we have particles $\{{\bm{x}}_{m^{*}}^{n}\}_{n=1}^{N}$ at $m^{*}$-th stage. We can cut the sequence at position $m^{*}$ and generate a new task using the second half. The prior for the new sequence will be an empirical density estimation $\hat{\pi}({\bm{x}};{\mathcal{X}}_{m^{*}})$ as defined in Eq. (18). Then, for all stages $m>m^{*}$, the terms in Eq. (20) becomes

[TABLE]

We can apply this technique for multiple times to split a long sequence into multiple segments.

4 Experiments

We conduct experiments on multivariate Gaussian model, hidden Markov model and Bayesian logistic regression to demonstrate the generalization ability of PFBR and also its accuracy for posterior estimation.

Evaluation metric. For multivariate Gaussian model and Gaussian linear dynamical system, we could calculate the true posterior. Therefore, we can evaluate:

• (i)

Cross-entropy $\mathbb{E}_{x\sim p}-\log q(x)$;

• (ii)

Maximum mean discrepancy (Gretton et al., 2012)

[TABLE]

• (iii)

Integral evaluation discrepancy $\|\mathbb{E}_{p}h(x)-\mathbb{E}_{q}h(x)\|_{2}$;

where we use Monte Carlo method to compute $\mathbb{E}_{x\sim p}[\cdot]$ for the first two metrics. For the integral evaluation, we choose some test functions $h$ where the exact value of $\mathbb{E}_{p}h(x)$ has a closed-form expression. For the experiments on real-world dataset, we estimate the commonly used prediction accuracy due to the intractability of the posterior.

Multivariate Gaussian Model. The prior ${\bm{x}}\sim{\mathcal{N}}(\mu_{x},\Sigma_{x})$, the observation conditioned on prior $o|{\bm{x}}\sim{\mathcal{N}}({\bm{x}},\Sigma_{o})$ and the posterior all follow Gaussian distributions. In our experiment, we use $\mu_{x}=\vec{0},\Sigma_{x}=I$ and $\Sigma_{o}=3I$. We test the learned PFBR on different sequences of 100 observations ${\mathcal{O}}_{100}$, while the training set only contains sequences of 10 observations, which are ten times shorter than sequences in the testing set. However, since we construct a set of different prior distributions to train the operator, the diversity of the prior distributions allows the learned ${\mathcal{F}}$ to generalize to compute posterior distributions in a longer sequence.

We compare the performances with KBR (Fukumizu et al., 2012) and one-pass SMC (Balakrishnan & Madigan, 2006). Both PFBR and KBR are learned from the training set and then used as algorithms on the testing set, which consists of 25 sequences of 100 observations $\{{\mathcal{O}}_{100}^{j}\sim{\mathcal{N}}({\bm{x}}^{j},\Sigma_{o})\}_{j=1}^{25}$ conditioned on 25 different ${\bm{x}}^{j}$ sampled from ${\mathcal{N}}(\mathbf{0},I)$. We compare estimations of $p({\bm{x}}|{\mathcal{O}}_{m}^{j})$ across stages from $m=1$ to 100 and the results are plotted in Fig. 2. Since KBR’s performance reveals to be less comparable in this case, we leave its results to Appendix D. We can see from Fig. 2 that as the dimension of the model increases, our PFBR has more advantages over one-pass SMC.

Gaussian Mixture Model. Following the same setting as Welling & Teh (2011), we conduct an experiment on an interesting mixture model, where the observations o\sim\frac{1}{2}\big{(}{\mathcal{N}}(x_{1},1.0)+{\mathcal{N}}(x_{1}+x_{2},1.0)\big{)} and the prior $x_{1},x_{2}\sim{\mathcal{N}}(0,1)$. The same as Dai et al. (2016), we set $(x_{1},x_{2})=(1,-2)$ so that the resulting posterior will have two modes: $(1,-2)$ and $(-1,2)$. During training, multiple sequences are generated, each of which consists of 30 observations.

Compared to Welling & Teh (2011) and Dai et al. (2016), our experiment is more challenging. First, they are only fitting one posterior given one sequence of observations via optimization, while PFBR will operate on sequences that are NOT observed during training without re-optimization. Second, they are only estimating the final posterior, while we aim at fitting every intermediate posteriors. Even for one fixed sequence, it is not easy to fit the posterior and capture both two modes, which can be seem from the results reported by Dai et al. (2016). However, our learned PFBR can operate on testing sequences and the resulting posteriors look closed enough to true posterior (See Fig. 3).

Hidden Markov Model — LDS. Our PFBR method can be easily adapted to hidden Markov models, where we will estimate the marginal posterior distribution $p({\bm{x}}_{m}|{\mathcal{O}}_{m})$. For instance, consider a Gaussian linear dynamical system (LDS):

[TABLE]

where $\mathbf{\epsilon}_{m}\sim{\mathcal{N}}(0,\Sigma_{1})$ and $\mathbf{\delta}_{m}\sim{\mathcal{N}}(0,\Sigma_{2})$. The particles can be updated by recursively applying two steps:

[TABLE]

where $\tilde{{\mathcal{X}}}_{m}:=\{\tilde{{\bm{x}}}_{m}^{n}\}_{n=1}^{N}$. The second step is a Bayesian update from $p({\bm{x}}_{m}|{\mathcal{O}}_{m-1})$ to $p({\bm{x}}_{m}|{\mathcal{O}}_{m})$ given the likelihood $p(o_{m}|{\bm{x}}_{m})$. The only tricky part is we do not have a tractable loss function in this case because of the integration $p({\bm{x}}_{m}|{\mathcal{O}}_{m-1})=\int p({\bm{x}}_{m}|{\bm{x}}_{m-1})p({\bm{x}}_{m-1}|{\mathcal{O}}_{m-1})\,d{\bm{x}}_{m-1}$. Hence, at each stage $m$, we use the particles $\tilde{{\mathcal{X}}}_{m}$ to construct an empirical density estimation $\hat{\pi}({\bm{x}};\tilde{{\mathcal{X}}}_{m})$ as defined in Eq. (18) and then define the loss at each stage $m$ as

[TABLE]

where we replace the intractable density $p({\bm{x}}_{m}|{\mathcal{O}}_{m-1})$ by $\hat{\pi}$. Given this loss, the PFBR operator ${\mathcal{F}}$ can be learned.

In the experiment, we sample a pair of 2 dimensional random matrices $A$ and $B$ for the LDS model. We learn both PFBR and KBR from a training set containing multiple different sequences of observations. For KBR, we use an adapted version called Kernel Monte Carlo Filter (Kanagawa et al., 2013), which is designed for state-space models. We use Kalman filter (Welch & Bishop, 2006) to compute the true marginal posterior for evaluation purpose.

Fig. 4 compares our method with KBR and SMC (importance sampling with resampling) on a testing set containing 25 new sequences of ordered observations. We see that our learned PFBR can generalize to test sequences and achieve better and stabler performances.

Comparison to Variational SMC. Autoencoding SMC (AESMC), Filtering Variational Objectives, and Variational SMC are three recent variational inference approaches that approximate the posterior based on SMC and a neural proposal (Le et al., 2018; Maddison et al., 2017; Naesseth et al., 2018). Since they share similar ideas, we implemented AESMC as a representative. We tried both MLP and GRU as mentioned in these papers. A comparison is made for 10-dimensional LDS (Table 1), which shows PFBR is better even with much fewer particles.

Inference Time Comparison. Table 1 also shows the inference time for updating posterior given one new observation. Though PFBR takes more time for the same #particles (e.g., 256), to get closer to PFBR’s performance, others need to use much more particles (e.g., 4096).

Bayesian Logistic Regression (BLR). We consider logistic regression for digits classification on the MNIST8M 8 vs. 6 dataset which contains about 1.6M training samples and 1932 testing samples. We reduce the dimension of the images to 50 by PCA, following the same setting as Dai et al. (2016). Two experiments are conducted on this dataset. For both experiments, we compare our method with SMC, SVI (stochastic variational inference (Hoffman et al., 2012)), PMD (particle mirror descent (Dai et al., 2016)), SGD Langevin (stochastic gradient Langevin dynamics (Welling & Teh, 2011)) and SGD NPV (stochastic version of nonparametric variational inference (Gershman et al., 2012)). This is a large dataset, so we use the techniques discussed in Section 3.3 to facilitate training efficiency.

BLR-Meta Learning. In the first experiment, we create a multi-task environment by rotating the first and second components of the features reduced by PCA through an angle $\psi$ uniformly sampled from $[-15\degree,15\degree]$. Note that the first two components account for more variability in the data. With a different rotation angle $\psi$, the classification boundary will change and thus a different classification task will be created. Also, a different sequence of image samples will result in different posterior distributions and thus corresponds to a different inference task.

We learn the operator ${\mathcal{F}}$ from a set of training tasks, where each task corresponds to a different rotation angle $\psi$. After that, we use ${\mathcal{F}}$ as Bayes’ Rule for testing tasks and compare its performances with other stochastic methods or sampling methods. Test is done in an online fashion: all algorithms start with a set of particles sampled from the prior (hence the prediction accuracy at 0-th step is around 0.5). Each algorithm will make a prediction to the encountered batch of 32 images, and then observe their true labels. After that, each algorithm will update the particles and make a prediction to the next batch of images. Ideally we should compare the estimation of posteriors. However, since it is intractable, we evaluate the average prediction accuracy at each stage. Results are shown in Fig. 5.

Note that we have conducted a sanity check to confirm the learned operator ${\mathcal{F}}$ does not ignore the first 2 rotated dimensions and use the rest 48 components to make predictions. More precisely, if we zero out the first two components of the data and learn ${\mathcal{F}}$ on them. The accuracy of the particles dropps to around 65%. This further verifies that the learned PFBR indeed can generalize across different tasks.

BLR-Variational Inference. For the second experiment on MNIST, we use PFBR as a variational inference method. That is to say, instead of accurately learning a generalizable Bayesian operator, the estimation of the posterior $p({\bm{x}}|{\mathcal{O}}_{train})$ is of more importance. Thus here we do not use the loss function in Eq. (17) summing over all intermediate states, but emphasize more on the final error $\text{KL}(q({\bm{x}})||p({\bm{x}}|{\mathcal{O}}_{train}))$. After the training is finished, we only use the transported particles to perform classification on the test set but do not further use the operator ${\mathcal{F}}$. The result is shown in Fig. 6, where $x$-axis shows number of visited samples during training. Since we use a batch size of 128 and consider 10 stages, the first gradient step of our method starts after around $10^{3}$ samples are visited.

5 Conclusion and Future Work

In this paper, we have explored the possibility of learning an ODE-based Bayesian operator that can perform online Bayesian inference in testing phase and verified its generalization ability through both synthetic and real-world experiments. Further investigation on the parameterization of flow velocity $f$ (e.g., use a stable neural architecture (Haber & Ruthotto, 2018) with HyperNetwork (Ha et al., 2017)) and generating diverse prior distributions through a Dirichlet process can be made to explore a potentially better solution to this challenging problem.

Acknowledgements

This project was supported in part by NSF IIS-1218749, NIH BIGDATA 1R01GM108341, NSF CAREER IIS-1350983, NSF IIS-1639792 EAGER, NSF IIS-1841351 EAGER, NSF CNS-1704701, ONR N00014-15-1-2340, Intel ISTC, NVIDIA, Google and Amazon AWS.

Appendix A Existence of Unified Flow Operator

See 2.3

Proof.

By Theorem 2.2, $\tilde{w}^{*}(q({\bm{x}},t),t):=\nabla_{x}\log q({\bm{x}},t)$ can induce the optimal state $\tilde{q}^{*}({\bm{x}},\infty)=p({\bm{x}}|{\mathcal{O}}_{m+1})$ and achieve a zero loss, $d=0$. Hence, $\tilde{w}^{*}$ is an optimal closed-loop control for this problem.

Although in general closed-loop control has a stronger characterization to the solution, in a deterministic system like Eq. (14), the optimal closed-loop control and the optimal open-loop control will give the same control law and thus the same optimality to the loss function(Dreyfus, 1964). Hence, there exists an optimal open-loop control $w^{*}=w^{*}(q({\bm{x}},0),t)$ such that the induced optimal state also gives a zero loss and thus $q^{*}({\bm{x}},\infty)=p({\bm{x}}|{\mathcal{O}}_{m+1})$.

More specifically, when the system is deterministic, a state $q({\bm{x}},t)$ is just a deterministic result of the initial state $q({\bm{x}},0)$ and the dynamics. The optimal flow determined by $\tilde{w}^{*}(q({\bm{x}},t),t)$ is

[TABLE]

The continuity equation gives

[TABLE]

Hence, for any ${\bm{x}}$,

[TABLE]

The dynamcis $g$ is a fixed function of $p({\bm{x}}|{\mathcal{O}}_{m})$, $p(o_{m+1}|{\bm{x}})$ and $q({\bm{x}},t)$, so the solution of this initial value problem(IVP) $q({\bm{x}},t)$ is a fixed function of $p({\bm{x}}|{\mathcal{O}}_{m})$, $p(o_{m+1}|{\bm{x}})$, $q({\bm{x}},0)$ and $t$, which can be written as

[TABLE]

Finally, we can write the optimal open-loop control as

[TABLE]

Hence, $w^{*}(q({\bm{x}},0),t)$ has a fixed form across different $m$.

∎

Appendix B Adjoint Method

To explain it more clearly, let us denote the evolution of the $n$-th particles at the $m$-th stage by ${\bm{x}}^{n}_{m}(t)$ for $t\in[0,T]$. Note that ${\bm{x}}_{m}^{n}(T)={\bm{x}}_{m+1}^{n}(0)$. (Then the notation ${\bm{x}}_{m}^{n}$ in the main text will become ${\bm{x}}_{m}^{n}(T)$.)

Recall the loss for each task:

[TABLE]

The loss of one particle ${\bm{x}}^{n}$ is

[TABLE]

where

[TABLE]

and $y_{m}^{n}(t):=-\log q^{n}_{m}(x_{m}^{n}(t),t)$.

First, an adjoint process is defined as

[TABLE]

Denote $f_{m}({\bm{x}}(t),\theta)=f_{\theta}({\mathcal{X}}_{m},o_{m+1},{\bm{x}}(t),t)$. During the $m$-th stage, the adjoint process follows the following differential equation

[TABLE]

Note that

[TABLE]

Claim: The gradient of the loss is the solution of a backward ODE. That is to say, $\frac{\partial L^{n}}{\partial\theta}={\bm{z}}_{1}(0)$, if ${\bm{z}}_{M}(T)=\mathbf{0}$ and

[TABLE]

and ${\bm{z}}_{m}(T)={\bm{z}}_{m+1}(0)$, for $m=0,\cdots,M-1$.

Proof.

First, we can compute $\frac{d}{dt}\frac{\partial L^{n}}{\partial\theta}$:

[TABLE]

Next, we have

[TABLE]

Hence, $\frac{\partial L^{n}}{\partial\theta}={\bm{z}}_{1}(0)$ if ${\bm{z}}_{M}(T)=\mathbf{0}$. ∎

An algorithm for computing $\frac{\partial L}{\partial\theta}$ is summarized in Algorithm 2. A nice python package of realizing this algorithm is provided by Chen et al. (2018).

Appendix C Experiment Details

C.1 Parameterization

Overall we parameterize the flow velocity as

[TABLE]

where both $\bm{\phi}$ and ${\bm{h}}$ are neural networks. For instance, let $\text{ctx}=[\textstyle{\frac{1}{N}\sum_{n=1}^{N}}\bm{\phi}({\bm{x}}_{m}^{n})^{\top},o_{m+1}^{\top}]$ be the context of this conditional flow, where $\bm{\phi}$ is a dense feed-forward neural network, a specific neural architecture we use in the experiment is

[TABLE]

where $*$ is element-wise multiplication. The number of layers $k$ can be tuned, but in general ${\bm{h}}$ is a shallow network.

C.2 Evaluation Metric

MMD2

The maximum mean discrepancy (MMD) of the true posterior $p$ and the estimated posterior $q$ is defined as

[TABLE]

When ${\mathcal{F}}$ is a unit ball in a characteristic RKHS,  Gretton et al. (2012) showed that the squared MMD is

[TABLE]

where $x,x^{\prime}\sim p$ and $y,y^{\prime}\sim q$.

Cross-entropy

Evaluating the KL divergence is equivalent to evaluating the cross-entropy.

[TABLE]

where $q(x)$ is approximated by kernel density estimation on the set of particles obtained from different sampling methods.

Integral Evaluation

When the true posterior is a Gaussian distribution ${\mathcal{N}}(\mu,\Sigma)$, the expectation of the following test functions have closed-form expressions.

• •

$\mathbb{E}[{\bm{x}}]=\mu$

• •

$\mathbb{E}[{\bm{x}}^{\top}A{\bm{x}}]=tr(A\Sigma)+\mu^{\top}A\mu$

• •

$\mathbb{E}[(A{\bm{x}}+{\bm{a}})^{\top}(B{\bm{x}}+{\bm{b}})]=tr(A\Sigma B^{\top})+(A\mu+{\bm{a}})^{\top}(B\mu+{\bm{b}})$

Appendix D More Experimental Results

D.1 Multivariate Guassian Model

D.2 LDS Model