Title: BAYESIAN META SAMPLING FOR FAST UNCERTAINTY ADAPTION

Abstract: Meta learning has been making impressive progress for fast model adaptation. However, limited work has been done on learning fast uncertainty adaption for Bayesian modeling. In this paper, we propose to achieve the goal by placing meta learning on the space of probability measures, inducing the concept of meta sampling for fast uncertainty adaption. Specifically, we propose a Bayesian meta sampling framework consisting of two main components: a meta sampler and a sample adapter. The meta sampler is constructed by adopting a neural-inverse-autoregressive-flow (NIAF) structure, a variant of the recently proposed neural autoregressive flows, to efficiently generate meta samples to be adapted. The sample adapter then pushes meta samples to task-specific samples, based on a newly proposed and general Bayesian sampling technique, called optimal-transport Bayesian sampling. The combination of the two components allows a simple learning procedure for the meta sampler to be developed, which can be efficiently optimized via standard back-propagation. Extensive experimental results demonstrate the efficiency and effectiveness of the proposed framework, obtaining better sample quality and faster uncertainty adaption compared to related methods.

1. INTRODUCTION:

Meta learning (Schmidhuber, 1987; Andrychowicz et al., 2016) is an important topic in modern machine learning. The goal is to learn some abstract concepts from different but related tasks, which can then be adapted and generalized to new tasks and environments that have never been encountered during training. There has been lots of research on this topic. A recent review classifies the methods as metric-based, model-based and optimization-based methods (Weng, 2018) . Among these methods, learning-to-learn seeks to learn a meta optimizer that can be applied to different models, with some task-specific information such as current gradients as input (Andrychowicz et al., 2016) . Model agnostic meta learning (MAML) aims to learn a meta parameter/model from a set of training tasks such that it can quickly adapt to models for new tasks (Finn et al., 2017) . Many follow-up works have been proposed recently, including but not limited to the meta network (Munkhdalai & Yu, 2017) , the meta learner (Ravi & Larochelle, 2017) , the Reptile model (Nichol et al., 2018) , and the lately extensions to an online setting (Finn et al., 2019) , to model hierarchical relation (Yao et al., 2019) and sequential strategies (Ortega et al., 2019) , and to its stable version Antoniou et al. (2019) and to some theoretical analysis (Khodak et al., 2019) .

It is worth noting that all the aforementioned models are designed from an optimization perspective. Bayesian modeling, in parallel with optimization, has also been gaining increasing attention and found various applications in deep learning. Recent research has extended the above meta-learning methods to a Bayesian setting. For example, Bayesian MAML (BMAML) replaces the stochasticgradient-descent (SGD) step with Stein variational gradient descent (SVGD) for posterior sampling (Yoon et al., 2018) . Probabilistic MAML (PMAML) extends standard MAML by incorporating a parameter distribution of the adapted model trained via a variational lower bound (Finn et al., 2018) . Amortized Bayesian Meta Learning extends the idea of MAML to amortized variational inference (Ravi & Beatson, 2019; Choi et al., 2019) . VERSA (Gordon et al., 2019) uses an amortization network to approximate the posterior predictive distributions. Meta particle flow (Chen et al., 2019) realizes Bayes's rule based on ODE neural operator that can be trained in a meta-learning framework.

Though methodologically elegant with many interesting applications, the above methods lack the ability to uncertainty propagation/adaption, in the sense that uncertainty is either not considered (e.g., in MAML) or only considered in the specific task level (e.g., BMAML). Uncertainty modeling is critical when considering especially in few labeled data setting. Lacking of such modeling, the prediction uncertainty would be inaccurate. Model prediction uncertainty is considered and calibrated in Ananya Kumar (2019).

To tackle this problem, we propose to perform meta learning on the space of probability measures, i.e., instead of adapting parameters to a new task, one adapts a meta distribution to new tasks. When implementing distribution adaption in algorithms where distributions are approximated by samples, our distribution-adaptation framework becomes sample-to-sample adaption. In other words, the meta parameter in standard MAML becomes meta samples in our method, where uncertainty can be well encoded. For this reason, we call our framework Bayesian meta sampling.

Specifically, we propose a mathematically elegant framework for Bayesian meta sampling based on the theory of Wasserstein gradient flows (WGF) (Ambrosio et al., 2005) . Our goal is to learn a meta sampler whose samples can be fast adapted to new tasks. Our framework contains two main components: a meta sampler and a sample adapter. For the meta sampler, we adopt the state-ofthe-art flow-based method to learn to transport noise samples to meta samples. Our meta sampler is parameterized by a neural inverse-autoregressive flow (NIAF), an extension of the recently developed neural autoregressive flows (NAFs) (Huang et al., 2018) . The NIAF consists of a meta-sample generator and an autoregressive conditioner model, which outputs the parameters of the meta-sample generator. The NIAF takes some task-specific information (such as gradients of target distributions) and random noise as input and outputs meta samples from its generator. These meta samples are then quickly adapted to task-specific samples of target distributions by feeding them to the sample adapter.

To ensure efficient and accurate adaptations to new task distributions, a novel optimal-transport Bayesian sampling (OT-sampling) scheme, based on Wasserstein gradient flows, is proposed as the adaptation mechanism of the sample adapter. The OT-sampling is general and can ensure samples to be adapted in a way that makes the sample density evolve to a target distribution optimally, thus endowing the property of fast uncertainty adaption. Finally, when one aims to perform specific tasks such as Bayesian classification with a task network, these samples are used to encode uncertainty into modeling. To this end, we further develop an efficient learning algorithm to optimize the task network based on variational inference. Extensive experiments are conducted to test the advantages of the proposed meta-sampling framework, ranging from synthetic-distribution to posterior-distribution adaption and to k-shot learning in Bayesian neural networks and reinforcement learning. Our results demonstrate a better performance of the proposed model compared to related methods.

2. NEURAL INVERSE-AUTOREGRESSIVE BAYESIAN META SAMPLING:

Our model combines ideas from Bayesian sampling, Wasserstein gradient flows and inverseautoregressive flows. A detailed review of these techniques is provided in Section A of the Appendix. The task is denoted with τ . The two components and specific inputs will be described in details.

2.1. THE BASIC SETUP AND OVERALL IDEA:

In meta sampling, one is given a set of related distributions, e.g., posterior distributions of the weights of a set of Bayesian neural networks (BNNs), each of which is used for classification on a different but related dataset. With our notation, each of the network and the related dataset is called a task, which is denoted as τ . Meta sampling aims to learn a meta sampler based on a set of training tasks so that samples from the meta sampler can be fast adapted to samples for an unseen new task.

For each task, the dataset D τ is divided into two sets D tr τ {X tr τ , y tr τ } and D val τ {X val τ , y val τ }. Our model builds on the MAML framework Finn et al. (2017) , consisting of meta initialization and task adaptation. The difference is that instead of adapting parameters as in MAML, we adapt the distribution of parameters in terms of samples. Following Grant et al. (2018) , the probabilistic view of our model can be formulated as a hierarchical Bayes model in equation 1, where W 0 denotes the meta initialization of network parameters, W τ denotes the network parameters for task τ , and adapting from P (W 0 |D tr T ) to P (W τ |D tr τ ) for a specific task τ is denoted as P (W τ |D tr τ , W 0 ):

P (y val T |X val T , D tr T ) = E W0∼P (W0 |D tr T ) τ ∈T P (y val τ | W τ , X val τ )P (W τ |D tr τ , W 0 )d W τ . (1)

Our overall idea to solve the problem is to mimic a hierarchical-sampling procedure but in a much more efficient way. Specifically, we propose to decompose meta sampling into two components: a meta sampler and a sample adapter. The meta sampler is responsible for generating meta samples that characterize common statistics of different tasks, i.e., approximating P (W 0 |D tr T ) on all the training tasks T with meta samples; The sample adapter is designed for fast adaptation of meta samples to task-specific target distributions, i.e., adapt the distribution from P (W 0 |D tr T ) to P (W τ |D tr τ ) in terms of samples for a task τ . The meta sampler is parameterized as a conditional generator, and aggregates all local losses of different tasks to form a final loss for optimization based on optimaltransport theory. Our method allows gradients to be directly backpropagated for meta-sampler updates. The overall idea is illustrated in Figure 1 .

Comparisons with related works We distinguish our model with two mostly related works: the meta NNSGHMC (Gong et al., 2019) and the probabilistic MAML (PMAML) (Finn et al., 2018) . The main differences lie in two aspects: meta representation and model architecture. In terms of meta-model representation, our model adopts data/parameter samples, instead of determinstic parameters, as meta representation, and thus can be considered as a sample-to-sample adaption. Meta NNSGHMC uses samples on different tasks, there is no concept of meta samples. Finally, PMAML fully relies on variational inference whose representation power could be restricted. In terms of model architecture, our model adopts the state-of-the-art autoregressive architectures, which can generate high-quality meta samples. Furthermore, our model adopts a simpler way to define the objective function, which allows gradients to directly flow back for meta-sampler optimization. It is worth noting that our framework reduces to MAML when only one meta sample is considered for sample adaption. Since our methods aim for Bayesian sampling, it is more similar to meta NNSGHMC (Gong et al., 2019) ; whereas MAML aims for point estimation of a model. Finally, we note that the recently proposed neural process (Garnelo et al., 2018) might also be used for meta-learning "few-shot function regression" as stated in the original paper. However, to our knowledge, no specific work has been done for this purpose.

2.2. THE BAYESIAN META SAMPLER:

G ! " # " G ! $ # % Γ Γ Γ One meta sample 0 = [# " , # $ , # % ]

Task network 5(⋅; 9) This section aims to design a meta sampler for efficiently generating meta samples. One idea is to use a nonparametric model to generate meta samples such as standard stochastic gradient MCMC (SG-MCMC) (Welling & Teh, 2011; Ma et al., 2015) and the Stein variational gradient descent (SVGD) (Liu & Wang, 2016a) , where no parameters are consider in the model. However, methods under this setting are typically slow, and the generated samples are usually highly correlated. Most importantly, it would be hard to design nonparametric samplers that can share information between different tasks. As a result, we propose to learn the meta sampler with a parametric model, which is also denoted as a generator.

NIAF:



2.2.1. LEARNING PARAMETERIZED META SAMPLERS WITH NIAFS:

There are two popular options to parameterize the meta sampler: with an explicit generator or with an implicit generator. An explicit generator parameterizes the output (i.e., meta samples) as samples from an explicit distribution with a known density form such as Gaussian, thus limiting the representation power. In the following, we propose to adopt an implicit generator for the meta sampler based on neural inverse-autoregressive flows (NIAF), an inverse extension of the recently proposed NAF (Huang et al., 2018) used for density estimation. As will be seen, NIAF can incorporate task-specific information into an implicit generator and generates samples in an autoregressive manner efficiently.

Finally, meta samples are used in a task network to encode uncertainty for specific tasks such as Bayesian classification. The architecture of the meta sampler is illustrated in Figure 2 . Note the idea of using NIAF to generate network parameter is similar to the hypernetwork (Ha et al., 2016; Krueger et al., 2017) . As an extention, Pradier et al. (2018) perform inference on a lower dimensional latent space. Using hypernetworks to model posterior distributions have also been studied in (Pawlowski et al., 2017; Sheikh et al., 2017) .

Neural inverse-autoregressive flows Directly adopting the NAF for sample generation is inappropriate as it was originally designed for density evaluation. To this end, we propose the NIAF for effective meta-sample generation. Specifically, let z k denote the k-th element of a sample z to be generated; z denotes a sample from last round; Γ denotes the task-specific information. In our case, we set Γ (z, ∇ z log p(z)) with p(•) denoting the target distribution (with possible hidden parameters). NIAF generates the sample z

= (z 1 , • • • , z k , • • •

) via an autoregressive manner, as:

EQUATION

where { k } are noise samples; G(•, •; φ φ φ) is an invertible function (generator) parameterized by φ φ φ and implemented as a DNN; and T is an autoregressive conditioner model parameterized by ψ ψ ψ to generate the parameters of the generator G at each step k, which is itself implemented as a deep sigmoidal flow or a deep dense sigmoidal flow as in (Huang et al., 2018) . According to Huang et al. (2018) , using strictly positive weights and strictly monotonic activation functions for G is sufficient for the entire network to be strictly monotonic, thus invertible.

Remark 1 The original NAFs are not designed for drawing samples, as one needs the inverse function G -1 , which is not analytically solvable when G is implemented as a neural network.

Although it is stated in (Huang et al., 2018 ) that G -1 can be approximated numerically, one needs repeated approximations, making it computationally prohibited. Our proposed NIAF is designed specifically for sample generation by directly transforming the noise with a flow-based network G.

The task network In addition to the NIAF, a task network might be necessary for processing specific learning tasks. In particular, if one is only interested in generating samples from some target distribution, a task network is not necessary as the meta samples will be used to adapt to task-specific samples. However, if one wants to do classification with uncertainty, the task network should be defined as a classification network such as an MLP or CNN. In this case, denoting the weights of the task network as W, we consider the task network as a Bayesian neural network, and propose two ways of parameterization to encode uncertainty of meta samples into the task network:

• Sample parameterization: A sample of the weights of the task network is directly represented by a meta sample from our meta sampler, i.e., W = (z 1 , z 2 , • • • , z p ) with p the parameter dimensionality. • Multiplicative parameterization: Adopting the idea of multiplicative normalizing flows (Louizos & Welling, 2017a) , we define an inference network for the weights as the multiplication of z and a Gaussian variational distribution for W, i.e., the variational distribution is defined as the following semi-implicit distribution to approximate the true posterior distribution for W:

EQUATION

Here and in the following, we consider the task network parameterized as a one-layer MLP for notation simplicity, although our method applies to other network structures; and we have used NIAF({ k }, Γ) to denote the output of meta samples from the NIAF.

Comparing the two parameterizations, sample parameterization directly generates weights of the task network from the meta sampler, thus is more flexible in uncertainty modeling. However, when the task network grows larger to deal with more complex data, this way of parameterization quickly becomes unstable or even intractable due to the high dimentionality of meta samples. Multiplicative parameterization overcomes this issue by associating each element of a meta sample with one node of the task network, reducing the meta-sample dimensionality from

O(N in × N out ) to O(N in )

with N in and N out being the input and output sizes of the task network. As a result, we adopt the multiplicative parameterization when dealing with high dimensional problems in our experiments. Efficient inference for these two cases will be described in Section 2.4. Note a recent work on NAF inference proposes to first sample from a mean-field approximating distribution, which are then transformed by an NAF to a more expressive distribution (Webb et al., 2019) . However, the approach is hard to scale to very high dimensional problems, e.g., posterior distributions of network parameters.

2.3. THE SAMPLE ADAPTER:

The output of the meta sampler contains shared (meta) information of all the tasks. Task-specific samples are expected to be adapted fast from these meta samples. This procedure is called sample adaption. Since there are potentially a large number of tasks, learning task-wise parametric models for sample adaption is impractical. Instead of using standard nonparametric samplers such as SG-MCMC or SVGD, we propose a general Bayesian sampling framework based on optimal-transport theory (Villani, 2008) for new task-sample adaption, where back-propagation can be directly applied.

A general Bayesian sampling framework based on optimal transport Let a task-specific target distribution be p τ (z), indexed by τ . A standard way is to adapt the samples based on a Markov chain whose stationary distribution equals p τ , e.g., via SG-MCMC. However, Markov-chain-based methods might not be efficient enough in practice due to the potentially highly-correlated samples (Chen et al., 2018b) . Furthermore, it is not obvious how to apply backpropagation (BP) in most of sampling algorithms. To deal with these problems, we follow Chen et al. (2018b) and view Bayesian sampling from the Wasserstein-gradient-flow perspective (discussed in Section A.2), i.e., instead of evolving samples, we explicitly evolve the underlying sample density functions. We will see that such a solution allows us to train the proposed meta sampler efficiently via standard BP.

Considering our meta-learning setting. Since we aim to adapt meta samples to new tasks, it is reasonable to define the adaptation via task-wise WGFs, i.e., for each task, there is a WGF with a specific functional energy and the corresponding first variation, denoted respectively as E τ and F τ δEτ δρ with the task index τ . Here ρ denotes the underlying density of the samples. Consequently, ρ will evolve with a variant of the PDE by replacing E with E τ in equation 9 for each task. To solve the corresponding PDE, we prove Theorem 1 based on a discrete approximation of ρ with the evolved meta samples, which is termed optimal-transport Bayesian sampling (OT-Bayesian sampling).

Theorem 1 (Optimal-Transport Bayesian Sampling) Let ρ t at time t be approximated by particles, i.e., ρ t (z)

≈ 1 M M i=1 δ z (i) t (z), where {z (i)

t } at time t = 0 are a set of meta samples from the meta generator * ; the delta-function δ a (b) = 1 if a = b and 0 otherwise. Then sample adaption for task τ can be described with the following differential equations:

EQUATION

)

Based on Theorem 1, particle updates for sampling can be obtained by solving equation 4 numerically, e.g., via the Euler scheme, as z

EQUATION

)

where h is the stepsize. Here with a little abuse of notation but for conciseness, we use z 

k can be considered as ∂Eτ ∂ z (i) k = ∇ z (i) k F τ (z (i) k )

. This is a useful result to derive a learning algorithm for the meta sampler described in Section 2.4. * We use the bold letter z (i) t (or z (i) k ) to denote the i-th meta sample evolved with equation 4 at time t (or equation 5 at iteration k). This should be distinguished from the normal unbold letter z k defined in Section 2.2.1, which denotes the k-th element of z.

Energy functional design Choosing an appropriate energy function E τ is important for efficient sample adaptation. To achieve this, the following conditions should be satisfied: i) E τ (ρ) should be convex w.r.t. ρ; ii) The first variation F τ could be calculated conveniently. A general and convenient functional family is the f -divergence, which is defined, with our notation and a convex function f : R → R such that f (1) = 0, as:

E f τ D f (ρ p τ ) = p τ (z)f ρ(z) pτ (z) d z.

The f -divergence is a general family of divergence metric. With different functions f , it corresponds to different divergences including the popular KL divergence, inverse-KL divergences, and the Jensen-Shannon divergence. For more details, please refer to (Nowozin et al., 2016) . A nice property of f -divergence is that its first variation endows a convenient form as stated in Proposition 2. Proposition 2 Let r(z) ρ(z) pτ (z) . The first variation of the f -divergence endows the following form:

F f τ δE f τ δρ (ρ) = f (r) .

In our experiments, we focus on the KL-divergence, which corresponds to f (r) = r log r. In this case,

∇ z F f τ (z) = ∇ z log p τ (z) -∇ z log ρ(z).

Since the density ρ(z) required in evaluating r is not readily available due to its implicit distribution, we follow Chen et al. (2018b) and use the meta samples {z (i) k } at the k-th step for approximation, resulting in

∇zF f τ (z) = ∇z log pτ (z) - M i=1 ∇ z (i) k κ(z, z (i) k )/ j κ(z (i) k , z (j) k ) - M i=1 ∇ z (i) k κ(z, z (i) k )/ j=1 κ(z, z (j) k ) (6)

where κ(•, •) is a kernel function. The number of adaptation steps k should be set based on problems.

For tasks that vary significantly, a larger k should be chosen to ensure the quality of adapted samples.

To further improve the accuracy, inspired by Chen et al. (2018b) , we combine equation 6 with the first variation of SVGD, resulting in the following form at iteration k:

EQUATION

where λ ≥ 0 is a hyperparameter to balance the two terms.

2.4. MODEL TRAINING AND SAMPLE GENERATION:

We first describe how to train the proposed model under the two kinds of parameterization of the task network defined in Section 2.2.1.

Training in the sample-parameterization setting In this case, one only needs to optimize the conditional model T (•; ψ ψ ψ) as all parameters of other networks are directly generated. Specifically, because the energy functionals for each task τ are designed so that the minima correspond to the target distributions, the objective thus can be defined over the whole task distribution p(τ ) as:

min L E τ ∼p(τ ) [E τ (ρ)] .

For notation simplicity, we will not distinguish among the adapted samples (i.e., z

k ) for different tasks. Since the only parameter is ψ ψ ψ in the autoregressive conditioner model T (see Figure 2 , and note the parameters φ φ φ and W for the meta generator and task network do not need to be learned as they are the outputs of T and G, respectively), its gradient can be directly calculated using chain rule:

EQUATION

where "

( * )

= " follows by the results from Section 2.3 and 2019), with the inference network defined in (3), and writing the implicit distribution of z as qψ ψ ψ (z) and the Gaussian prior distribution of W as p(W), we arrive at the following ELBO:

∇ z (i) k F τ (z (i) k ) is

L(θ, φ φ φ, ψ ψ ψ) E q φ φ φ (z,W) [log p(y| x, W, z) + log p(W) + log r θ (z | W) -log q φ φ φ (W | z) -log qψ ψ ψ (z)]

Different from Louizos & Welling (2017a) , we only update (θ, φ φ φ) by optimizing the above ELBO; while leave the update of ψ ψ ψ with the gradient calculated in equation 8, which reflects gradients of samples from the NIAF. Note also that in the meta learning setting, the task network needs to be adapted for new tasks. This can be done by standard MAML with the above ELBO as the new-task objective. The meta training and testing algorithms are illustrated in the Appendix B. A similar ideas of multiplicative parametrization was proposed recently in (Kristiadi & Fischer, 2019) , which used a compound density network to quantify the predictive uncertainty.

New-task sample generation After training, samples for a new task can be directly generated by feeding the task information Γ and some noise to the meta sampler depicted in Figure 1 . Typically, one needs to run a few numbers of sample-adaption steps to generate good samples for new tasks. Notably, the number of sample-adaption steps required to obtain good accuracy will be shown much less than simply starting a sampler from scratch in the experiments.

3. EXPERIMENTS:

We conduct a series of experiments to evaluate the efficiency and effectiveness of our model, and compare it with related Bayesian sampling algorithms such as SGLD, SVGD and SGHMC. The main compared algorithms for meta learning include the PMAML (Finn et al., 2018) , Amortized Bayesian Meta-Learning (ABML) (Ravi & Beatson, 2019) , and NNSGHMC (Gong et al., 2019) , a recently proposed meta SG-MCMC algorithm. Inspired by Finn et al. (2017) , we denote our algorithm distribution agnostic meta sampling (DAMS). Implementation details are given in Appendix E. Our code is made available at: https://github.com/zheshiyige/meta-sampling.git. We first demonstrate our proposed NIAF-based sampler is able to generate more effective samples compared to the popular Bayesian algorithms such as SVGD, SGLD and SGHMC, in a non-meta-sampling setting. To this end, we apply standard Bayesian Logistic Regression (BLR) on several real datasets from the UCI repository: Australian (15 features, 690 samples), German (25 features, 1000 samples), Heart (14 features, 270 samples). We perform posterior sampling for BLR using our proposed sampler, as well as SVGD, SGLD, SGHMC. For a more detailed investigation of different components in our model, we also test the generator with different architectures, including generators with MLP (DAMS with MLP), IAF (DAMS with IAF), and NIAF (DAMS with NIAF). We follow Liu & Wang (2016a) and apply Gaussian priors for the parameters p 0 (w|α) = N (w; 0, α -1 I) with p 0 (α) = Gamma(α, 1, 0.01). A random selection of 80% data are used for training and the remaining for testing. The testing accuracies are shown in Table 1 . It is observed that DAMS with NIAF achieves the best performance in terms of accuracy.

3.1. EVALUATIONS OF OT-BAYESIAN SAMPLING:

The results also indicate the effectiveness and expressiveness of the proposed NIAF architecture in the OT-Bayesian sampling framework.

3.2. EVALUATIONS OF META-SAMPLE ADAPTABILITY:

In this set of experiments, we aim to demonstrate the excellent meta-sample adaptability of our meta-sampling framework in different tasks. An additional synthetic experiment on meta posterior adaption is presented in Section D.3 of the Appendix.

Gaussian mixture model We first conduct experiments to meta-sample several challenging Gaussian mixture distributions. We consider mixtures of 4, 6 and 20 Gaussians. Detailed distributional forms are given in the Appendix. To setup a meta sampling scenario, we use 2, 3 and 14 Gaussian components with different means and covariance, respectively, for meta-training of the meta sampler. After training, meta samples are adapted to samples from a target Gaussian mixture by following the new-task-sample-generation procedure described in Section 2.4. We plot the convergence of 1000 meta samples to a target distribution versus a number of particle (sample) updates (iterations), measured with the maximum mean discrepancy (MMD) evaluated by samples. For a fair comparison, we use the same number of samples (particles) to evaluation the MMD. The results are shown in Figure 3 . It is clear that our proposed meta sampler DAMS converges much faster and better than other sampling algorithms, especially on the most complicated mixture of 20-Gaussians. The reason for the fast convergence (adaption) is partially due to the learned meta sampler, which provides good initialization for sample adaption. This is further verified by inspecting how samples in the adaption process evolve, which are plotted in Figure 10 , 11 and 12 in the Appendix.

Meta Sampling for Bayesian Neural Networks Finally, we test the proposed DAMS for meta sampling of BNNs on MNIST and CIFAR-10 (Krizhevsky, 2009) . We follow the experimental setting in (Gong et al., 2019) , and split the MNIST and CIFAR10 dataset into two parts for meta training and testing (sample adaption), respectively. As we focus on fast adaptation, we show the accuracy within 200 iterations. To deal with the high-dimensionality issue, we adopt the method of multiplicative parameterization proposed in Section 2.2.1.

We randomly pick 5 classes for training, and the remaining classes for testing. A BNN is trained only on the training data for meta sample (weights of the BNN) generation, with each sample corresponding to a meta BNN. In testing, the meta BNNs are adapted based on the testing data. For the MNIST dataset, we parameterize a BNN as a CNN with two convolutional layers followed by a fully connected layer with 100 hidden units. The kernel and filter sizes of the two conv layers are 3 and 16, respectively. A similar architecture is applied for the CIFAR10 dataset, but with 16 and 50 filters whose kernel sizes are 7 and 5 for the two convolutional layers, respectively. The hidden units of the fully connected layer is 300.

a) Adaptation efficiency: For this purpose, we compare our model with NNSGHMC (Gong et al., 2019) , as well as with a non-meta-learning method to train from scratch. To demonstrate the effectiveness of our NIAF structure for adaptive posterior sampling, we also compare it with the simple conditional version of MNF (Louizos & Welling, 2017a) . Figure 4 plots the learning curves of testing accuracy versus the number of iterations. It is clearly seen that our DAMS adapts the fastest to new tasks, and is able to achieve the highest classification accuracy on all cases due to the effectiveness of uncertainty adaption. To further demonstrate the superiority of DAMS over NNSGHMC, we list the test accuracy at different adaptation steps in Table 2 . The results clearly show faster adaption and higher accuracy of the proposed DAMS compared to NNSGHMC. It is also interesting to see that MNF with adaptation, performs better than NNSGHMC; while our method performs better than both, demonstrating the effectiveness of the proposed NIAF architecture.

b) Sample efficiency::

To demonstrate sample efficiency of our framework, we compare it with both NNSGHMC and the standard Bayesian learning of DNNs with SGHMC. To this end, we use the first five classes of image on CIFAR10 as training data to learn the meta sampler and we randomly select 5%, 20%, 30% of the training data in the remaining image classes of the test task to adapt the meta sampler. Figure 5 shows the corresponding test accuracies for different settings. It is observed that ours, the adaptation-based methods, obtain higher accuracies than the non-adaptive method of SGHMC. Furthermore, our method achieves the best sample efficiency among others.

c) Uncertainty evaluation: Finally, we ablate study the uncertainty estimation of our, the standard SGHMC and NNSGHMC models in terms of negative log-likelihood, also show the corresponding test accuracy. The results on the CIFAR10 dataset are shown in Figure 6 and Table 3 . We follow Louizos & Welling (2017a) to evaluate uncertainty via entropy of out-of-sample predictive distributions (Figure 6 (right)). We observe that uncertainty estimates with DAMS are better than others, since the probability of low entropy prediction is much lower than others. Details are given in Section D.4 of the Appendix. We compare our method with MAML and its variants with uncertainty modeling, including the Amortized Bayesian Meta-Learning (ABML) (Ravi & Beatson, 2019) and Probabilistic MAML (PMAML) (Finn et al., 2018) . To get a better understanding of each component of our framework, we also conduct an ablation study with three variants of our model: MAML-SGLD, MAML-SGHMC and DAMS-SGLD. MAML-SGLD and MAML-SGHMC correspond to the variants where SGLD and SGHMC are used to sample the parameters of the classifier, respectively; and DAMS-SGLD replaces the WGF component of DAMS-NIAF with SGLD. Follow the setting in (Finn et al., 2017; 2018) , the network architecture includes a stacked 4-layer convolutional feature extractor, followed by a meta classifier with one single fully-connected layer using the multiplicative parameterization. Testing results are presented in Table 4 . With our method, we observe significant improvement of the classification accuracy at an early stage compared with MAML. The learning curves are plotted in Figure 7 , further demonstrating the superiority of our method, which can provide an elegant initialization for the classification network. Finally, from the ablation study, the results suggest both the NIAF and the WGF components contribute to the performance gain obtained by our method. Meta sampling for reinforcement learning We next adapt our method for meta reinforcement learning. We test and compare the models on the same Mu-JoCo continuous control tasks (Todorov et al., 2012) as used in (Finn et al., 2017) , including the goal velocity task and goal direction task for cheetah robots. For a fair comparison, we leverage the TRPO-RL (Schulman et al., 2015) framework for meta updating following MAML method. Specifically, we implement the policy network with two hidden layers with ReLu activation followed by a linear layer to produce the mean value of the Gaussian policy.

The first hidden layer is a fully connected layer, and we adopt the multiplicative parameterization for the second hidden layer. As shown in Figure 8 , our method obtains higher rewards compared with MAML on both tasks, indicating the importance of effective uncertainty adaptation in RL.

4. CONCLUSION:

We present a Bayesian meta-sampling framework, called DAMS, consisting of a meta sampler and a sample adapter for effective uncertainty adaption. Our model is based on the recently proposed neural autoregressive flows and related theory from optimal transport, enabling a simple yet effective training procedure. To make the proposed model scalable, an efficient uncertainty parameterization is proposed for the task network, which is trained by variational inference. DAMS is general and can be applied to different scenarios with an ability for fast uncertainty adaptation. Experiments on a series of tasks demonstrate the advantages of the proposed framework over other methods including the recently proposed meta SG-MCMC, in terms of both sample efficiency and fast uncertainty adaption.

A BACKGROUND:

This section provides a review of background on Bayesian sampling, Wasserstein gradient flows and autoregressive flows.

A.1 BAYESIAN SAMPLING Bayesian sampling has been a long-standing tool in Bayesian modeling, with a wide range of applications such as uncertainty modeling (Li et al., 2015) , data generation (Feng et al., 2017; Chen et al., 2018a) and reinforcement learning (Zhang et al., 2018) . Traditional algorithms include but are not limited to Metropolis-Hastings algorithm, importance sampling and Gibbs sampling (Gelman et al., 2004) . Modern machine learning and deep learning have been pushing forward the development of large-scale Bayesian sampling. Popular algorithms in this line of research include the family of stochastic gradient MCMC (SG-MCMC) (Welling & Teh, 2011; Ma et al., 2015) and the Stein variational gradient descent (SVGD) (Liu & Wang, 2016a) . Recently, a particle-optimization sampling framework that unifies SG-MCMC and SVGD has also been proposed (Chen et al., 2018b) , followed by some recent developments (Liu et al., 2019b; a) . Generally speaking, all these methods target at sampling from some particular distributions such as the posterior distribution of the weights of a Bayesian neural network (BNN).

On the other hand, meta learning is a recently developed concept that tries to learn some abstract information from a set of different but related tasks. A natural question by considering these two is: can we design meta sampling algorithms that learns to generate meta samples, which can be adapted to samples of a new task-specific distribution quickly? This paper bridges this gap by proposing a mathematically sound framework for Bayesian meta sampling.

A.2 OPTIMAL TRANSPORT AND WASSERSTEIN GRADIENT FLOWS:

In optimal transport, a density function, ρ t , evolves along time t to a target distribution optimally, i.e., along the shortest path on the space of probability measures P(Ω), with Ω being a subset of R d . The optimality is measured in the sense that ρ t moves along the geodesic of a Riemannian manifold induced by a functional energy, E : P(Ω) → R, under the 2-Wasserstein distance metric. Formally, the trajectory of ρ t is described by the following partial differential equation (PDE):

EQUATION

where

∇ z • f i ∂ ∂ zi f is the divergence operator; and F δE δρt (ρ t

) is called the first variation of E at ρ t (functional derivative on a manifold in P(Ω)). To ensure ρ t to converge to a target distribution p such as the posterior distribution of the model parameters, one must design an appropriate E such that p = arg ρ min E(ρ). A common choice is the popular KL-divergence, KL(ρ, p). We will consider a more general setting in our framework presented later. Note the WGF framework equation 9 allows to view Bayesian sampling from a density-optimization perspective. For example, recent works (Chen et al., 2018b; Liu et al., 2019b; a) consider approximating ρ t with samples and evolve the samples according to equation 9. Parts of our model will follow this sampling setting.

A.3 AUTOREGRESSIVE FLOWS AND VARIANTS:

Our model relies on the concept of autoregressive flows for meta-sampler design. We review some key concepts here. More detailed comparisons are provided in the Appendix.

A normalizing flow defines an invertible transformation from one random variable to another z. A flexible way to implement this is to define it via implicit distributions, meaning sample generation is implemented as: i ∼ q 0 ( i ), z i = G( i ; φ φ φ), where i indexes elements of and z; G represents a deep neural network (generator) parameterized by φ φ φ.

The autoregressive flow (AF) parameterizes a Gaussian conditional distribution for each z i , e.g., p(z i | z 1:i-1 ) = N (z i |µ i , exp(α i )), where µ i = g µi (z 1:i-1 ) and α i = g αi (z 1:i-1 ) are outputs of two neural networks g µi and g αi . The sample generation process is: z i = i exp(α i ) + µ i , with µ i = g µi (z 1:i-1 ), α i = g αi (z 1:i-1 ) and i ∼ N (0, 1). Instances of autoregressive flows include the where

z ∈ R 2 , µ 1 = -µ 2 = [3, 0], µ 3 = -µ 4 = [0, 3], σ = [1, 1]. Mog6: p(z) = 6 i=1 N (z |µ i , σ 2 )/6 where z ∈ R 2 , µ i = [4 cos(iπ/3), 4 sin(iπ/3)], σ = [1, 1]. Mog20: p(z) = 20 i=1 N (z |µ i , σ 2 )/20

where z ∈ R 2 , µ = meshgrid(x, y), x = linspace(-6, 6, 2), y = linspace (-8, 8, 4) We apply our DAMS for fast adaptive sampling on regression tasks. We follow Huang et al. (2018) and apply DAMS to sample the posterior distribution of the frequency parameter of a sine-wave function, given only three data points.

The sinewave function is defined as y(t) = sin(2πf t), with a uniform prior U ([0, 2]) and a Gaussian likelihood N (y i ; y f (t i ), 0.125).

We design a meta-sampling setting to adapt the posterior distributions, p(f |D) on the training data, to that on new test data.

Specifically, meta training data For the first setting, meta testing consists of data {(0, 0), (3/5, 0), (6/5, 0)}.

For the second setting, meta testing consists of data {(0, 0), (4/5, 0), (8/5, 0)}. Meta training data corresponds to a posterior with two modes of f ∈ {0.0, 5/4}. For the first setting, the test data corresponds to a posterior with three modes f ∈ {0.0, 5/6, 5/3}. For the second setting, the test data corresponds to four modes f ∈ {0.0, 5/8, 5/4, 15/8}. We compare our DAMS with the result of re-training from scratch with the test data. Empirical distribution with samples and kernel density estimation are plotted in Figure 13 . The first setting takes about 3.4K iterations to find the three modes in the posterior with re-training, while it takes about 0.8K iterations with meta adaptation. For the second setting, it takes more than 3.6K iterations to find the four modes with training from scratch, while it is about 0.9K iterations with meta adaptation. For both test tasks, the sampler with re-training miss at least one mode compared with the meta sampler adaptation with the same number of iterations. We can see that DAMS can adapt the training posterior to the test posterior much faster than re-training from scratch due to effective uncertainty adaption, obtaining more than 3X speedups. We show the predictive uncertainty of DAMS compared to SGHMC and NNSGHMC by exploring the posterior of neural parameters, we estimate the uncertainty for out-of-distribution data samples Lakshminarayanan et al. (2017) . We train different algorithms on the MNIST dataset, and estimate the entropy of the predictive distribution on the notMNIST dataset Bulatov (2011) . We follow Louizos & Welling (2017a) and use the empirical CDF of entropy to evaluate the uncertainty. Since the probability of observing a high confidence prediction is low, curves that are nearer to the bottom right of the figure estimates uncertainty better. The predictive distribution of the trained model is expected to be uniform over the notMNIST digits as the samples from the dataset are from unseen classes. The BNN is a CNN with 16 and 50 filters whose kernel sizes are 5 and 5 for the two convolutional layers, respectively. The hidden units of the fully connected layer is 300. The uncertainty evaluation results are shown in the right plot of Figure 6 .

:

# $ G! %

Figures and Tables:

Figure1: Overall idea of the proposed meta-sampling framework. The task is denoted with τ . The two components and specific inputs will be described in details.

Figure 2: Roll-out architecture of the proposed NIAF as a meta sampler, consisting of a generator G and an autoregressive conditional model T (•; ψ ψ ψ). The meta samples are then fed to a task network C(•; W) to encode uncertainty into specific tasks such as classification. Please see detailed descriptions in the text.

at time t = hk with k being the iteration. Theorem 1 and equation 5 indicate that the derivative of E τ w.r.t. z (i)

readily calculated with equation 7; the term ∂ z (i) k ∂ψ ψ ψ can be calculated by standard BP over the meta-sampler network (NIAF). Training in the multiplicative-parameterization setting In this case, two sets of parameters are to be learned, ψ ψ ψ and W. For ψ ψ ψ, one can optimize ψ ψ ψ by adopting the same update equation as in the sample-parameterization setting. For W, we follow Louizos & Welling (2017a) and adopt variational inference. Specifically, we first augment the task network with an auxiliary network with a conditional likelihood r θ (z | W) k N (z k |μ k , σ2 k ), parameterized by θ {μ k , σk }. Based on Ranganath et al. (2016); Louizos & Welling (2017a); Touati et al. (

Figure 3: Convergence of meta-sampling different Gaussian mixture models. SGLD-m and SVGD-m are SGLD and SVGD with the same initialization as DAMS. Other samplers are initialized randomly.

Figure 4: Adaptation efficiency in terms of testing accuracy on CIFAR10 (left) and MNIST (right).

Figure 5: Sample efficiency evaluation with 5%, 20%, 30% of training data on CIFAR10. (see text).

Figure 7: Learning curves of 5-way 1-shot classification.3.3 META SAMPLING FOR BAYESIAN FEW-SHOT LEARNINGDAMS for few-shot image classification We apply our method on two popular few-shot imageclassification tasks on the Mini-Imagenet dataset, consisting of 64, 16, and 20 classes for training, validation and testing, respectively. We compare our method with MAML and its variants with uncertainty modeling, including the Amortized Bayesian Meta-Learning (ABML)(Ravi & Beatson, 2019) and Probabilistic MAML (PMAML)(Finn et al., 2018). To get a better understanding of each component of our framework, we also conduct an ablation study with three variants of our model: MAML-SGLD, MAML-SGHMC and DAMS-SGLD. MAML-SGLD and MAML-SGHMC correspond to the variants where SGLD and SGHMC are used to sample the parameters of the classifier, respectively; and DAMS-SGLD replaces the WGF component of DAMS-NIAF with SGLD. Follow the setting in(Finn et al., 2017;2018), the network architecture includes a stacked 4-layer convolutional feature extractor, followed by a meta classifier with one single fully-connected layer using the multiplicative parameterization. Testing results are presented in Table4. With our method, we observe significant improvement of the classification accuracy at an early stage compared with MAML. The learning curves are plotted in Figure7, further demonstrating the superiority of our method, which can provide an elegant initialization for the classification network. Finally, from the ablation study, the results suggest both the NIAF and the WGF components contribute to the performance gain obtained by our method.

Figure 8: Locomotion comparison results for goal velocity task (left) and goal direction task (right).

Figure 9: Densities of three Gaussian mixtures.

Figure 13: Meta sampling Bayesian regression for two test tasks (top and bottom). Each row plots posteriors after meta training (left), meta testing (middle, 0.8K (top) and 0.9K (bottom) iterations) and re-training (right, 0.8K (top) and 0.9K (bottom) iterations).

Figure 10: Comparison among different samplers on adapting to Mixture of 4-Gaussian. Top to Bottom row: DAMS, SGLD-m, SVGD-m, SVGD, SGLD and NNSGHMC

Figure 12: Comparison among different samplers on adapting to Mixture of 6-Gaussian. Top to Bottom row: DAMS, SVGD, SGLD and NNSGHMC

Accuracy of BLR with different samplers.

Table content: <table><tr><td/><td colspan="3">Australian German Heart</td></tr><tr><td>SVGD</td><td>88.7</td><td>74.5</td><td>78.7</td></tr><tr><td>SGLD</td><td>88.4</td><td>75.2</td><td>81.5</td></tr><tr><td>SGHMC</td><td>88.7</td><td>77.0</td><td>88.6</td></tr><tr><td>DAMS with MLP</td><td>85.5</td><td>75.0</td><td>87.0</td></tr><tr><td>DAMS with IAF</td><td>87.0</td><td>75.5</td><td>88.9</td></tr><tr><td>DAMS with NIAF</td><td>89.1</td><td>77.5</td><td>90.7</td></tr></table>

Accuracy (%) for adaptation efficiency.

Table content: <table><tr><td/><td>Approach</td><td colspan="2">CIFAR MNIST</td></tr><tr><td>20 steps</td><td>NNSGHMC</td><td>23.8</td><td>28.6</td></tr><tr><td>20 steps</td><td>MNF</td><td>42.3</td><td>59.4</td></tr><tr><td>20 steps</td><td>DAMS-NIAF</td><td>42.7</td><td>60.7</td></tr><tr><td>50 steps</td><td>NNSGHMC</td><td>47.2</td><td>63.6</td></tr><tr><td>50 steps</td><td>MNF</td><td>59.8</td><td>75.5</td></tr><tr><td>50 steps</td><td>DAMS-NIAF</td><td>61.9</td><td>78.5</td></tr><tr><td colspan="2">100 steps NNSGHMC</td><td>58.1</td><td>80.1</td></tr><tr><td>100 steps</td><td>MNF</td><td>67.2</td><td>85.8</td></tr><tr><td colspan="2">100 steps DAMS-NIAF</td><td>69.5</td><td>88.0</td></tr><tr><td colspan="2">200 steps NNSGHMC</td><td>71.7</td><td>88.6</td></tr><tr><td>200 steps</td><td>MNF</td><td>72.7</td><td>90.8</td></tr><tr><td colspan="2">200 steps DAMS-NIAF</td><td>75.2</td><td>92.6</td></tr></table>

Uncertainty adaptation evaluation of Bayesian meta sampling for BNNs in terms of negative log-likelihood.

Table content: <table><tr><td/><td>Approach</td><td colspan="2">accuracy NLL/100</td></tr><tr><td>20 steps</td><td>SGHMC</td><td>32.3</td><td>79.22</td></tr><tr><td>20 steps</td><td>DAMS-NIAF</td><td>42.7</td><td>68.98</td></tr><tr><td>50 steps</td><td>SGHMC</td><td>39.2</td><td>75.34</td></tr><tr><td>50 steps</td><td>DAMS-NIAF</td><td>61.9</td><td>57.82</td></tr><tr><td>100 steps</td><td>SGHMC</td><td>46.5</td><td>67.95</td></tr><tr><td colspan="2">100 steps DAMS-NIAF</td><td>69.5</td><td>46.49</td></tr><tr><td>200 steps</td><td>SGHMC</td><td>59.1</td><td>51.79</td></tr><tr><td colspan="2">200 steps DAMS-NIAF</td><td>75.2</td><td>36.53</td></tr></table>

Testing accuracy of Mini-ImageNet Few Shot Classification. "N/A" means results not reported. SGHMC 50.18 ± 0.50 % 65.94 ± 0.48 % MAML-SGLD 50.14 ± 0.51 % 65.12 ± 0.47 % DAMS-SGLD 50.86 ± 0.48 % 66.23 ± 0.45 % DAMS-NIAF 51.43 ± 0.49 % 66.84 ± 0.45 %

Table content: <table><tr><td colspan="4">Mini-Imagenet 5-way Few Shot Classification</td></tr><tr><td>Approach</td><td>1-shot</td><td>Accuracy</td><td>5-shot</td></tr><tr><td>MAML</td><td colspan="3">48.70± 1.84 % 63.11± 0.92 %</td></tr><tr><td>PMAML</td><td colspan="2">50.13 ± 1.86 %</td><td>N/A</td></tr><tr><td>ABML</td><td colspan="2">45.0 ± 0.60 %</td><td>N/A</td></tr><tr><td>MAML-</td><td/><td/><td/></tr></table>

Peer review (iclr2020):
Thank you for an interesting read. As far as I understand, this paper presents DAMS which is a MAML-like algorithm but applied to posterior sampling. The idea is the following: 1. Construct a meta-sampler that generates good proposals/initial samples for task-specific samplers; 2. Train the meta-sampler so that the task-specific posterior sampling converges faster to the target distribution. The meta-sampler is designed as an inverse version of the neural auto-regressive flow (NIAF), and the task-specific sampler is based on the Wasserstein gradient flow (WGF). ======= novelty ====== The method is novel: 1. Probabilistic/Bayesian understanding of few-shot learning/meta learning has been proposed, but to date variational inference is the main inference engine used in the literature. This paper provides a nice complement by considering fast adaptation of posterior sampling. 2. The meta-level parameterisation method is indeed different from (probabilistic) MAML, in the sense that the initial parameters/samples z is generated from a neural network conditioned on a task, instead of using a shared initialisation across tasks. This is more inline with approaches such as hyper-networks, and I haven*t seen many MAML-like approaches doing that. The meta-sampler architecture is new but improved upon NAF so I consider the architectural novelty to be minor. ======= significance ====== The experimental section contains many results, with the two main category as (a) comparisons between DAMS and existing sampling methods on Bayesian inference tasks; and (b) comparisons to MAML on few-shot learning tasks. Compared with the baselines, DAMS achieves significantly better results, which is a good sign. However, DAMS as a whole pipeline has a lot of components, and it is not clear to me which part is the main driving force. So I think the following ablation studies will be very helpful: 1. To see whether it is necessary to use NIAF, one can replace NIAF with a set of learned particles shared across tasks (i.e. make psi = {z^(1), ..., z^(N)} which is in similar spirit as MAML). 2. To see whether WGF brings in significant improvement in DAMS, one can replace WGF in sample adapter with SVGD or SG-MCMC method such as SGHMC/preconditioned SGLD. A comparison between e.g. SGHMC vs DAMS-SGHMC vs DAMS-WGF will be helpful for this ablation. 3. It seems to me the meta-sampler requires running a few step of NAIF updates (Gamma info contains previous sample and the gradients). How many steps are required here? A detailed analysis will be useful. If a lot of steps are required, the *fast adaptation* in meta-testing is not really the case, as both meta-sampling and sample adaptation requires evaluating gradients. 4. The multiplicative normalising flow for BNN approximate posterior adds in another layer of complexity, as the method also perform meta-learning on this variational distribution. So a baseline which removes the NIAF part (i.e. also learn particles for masks z, see point 1) on few-shot learning tasks would be useful. As far as I understand this baseline approach is different from ABML/PMAML that has been reported in the paper. ======= clarity ======= The presentation needs to be improved. To me, section 3 seems to be overwhelmed by details, e.g. the long discussion of task networks and WGF. For example, to what extent do the WGF construction details matter for understanding the whole DAMS pipeline? I think the general idea works for any valid posterior sampler in the fast adaptation step. I would suggest the following structure for section 3 instead: 1. write down the whole pipeline in a more abstract way, e.g. say the meta-sampler is any generator conditioned on task information, and the sample adaptation as generic posterior sampling; 2. discuss the training algorithm, what are the loss functions, etc; 3. discuss in meta-testing how DAMS is deployed; 3. discuss the detail implementation of the meta-sampler and sample adaptation method.