Learning Expressive Meta-Representations with Mixture of Expert Neural Processes
Abstract
Neural processes (NPs) formulate exchangeable stochastic processes and are promising models for meta learning that do not require gradient updates during the testing phase.
However, most NP variants place a strong emphasis on a global latent variable.
This weakens the approximation power and restricts the scope of applications using NP variants, especially when data generative processes are complicated.
To resolve these issues, we propose to combine the M ixture o f E xpert models with N eural P rocesse s to develop more expressive exchangeable stochastic processes, referred to as Mixture of Expert Neural Processes (MoE-NPs).
Then we apply MoE-NPs to both few-shot supervised learning and meta reinforcement learning tasks.
Empirical results demonstrate MoE-NPs’ strong generalization capability to unseen tasks in these benchmarks.
Introduction
Humans can naturally accommodate themselves to new environments after developing related skills and this kind of adaptability relies on the good abstraction of environments.
Similarly, meta learning or learning to learn tries to leverage past experiences, and with help of the incorporated meta learned knowledge, a new skill can be mastered rapidly with a few instances.
During the past decade, an increasing number of methods have emerged in meta learning domains.
In this paper, we concentrate on a special branch of meta learning methods, referred to as context-based meta learning [1; 2].
A representative one is a neural process (NP) [1], which was initially proposed to approximate Gaussian processes with lower computational cost.
The core purpose of NPs is to learn meta-representations [3], which encode context points into latent variables and represent the task in a functional form.
In comparison to gradient-based meta learning algorithms, e.g. model-agnostic meta learning (MAML) [4], the NP directly learns a functional representation and does not require additional gradient updates in fast adaptation.
Research Motivations.
Fundamentally, vanilla NPs employ global Gaussian latent variables to specify different tasks.
This setting raises several concerns in some scenarios.
(i)
When observations originate from a mixture of stochastic processes [5], a single
Gaussian latent variable is faced with deficiencies in formulating complex functional forms.
(ii) In context-based meta reinforcement learning, the uncertainty of value functions revealed from latent variables encourages effective exploration in environments [6].
But when tasks are governed by multiple variate, e.g. velocities, masses or goals in Mujoco robots [7], the use of a unimodal Gaussian latent variable restricts the flexibility of randomized value functions [8], leading to sub-optimality in performance.
Developed Methods.
Instead of using a global latent variable in modeling, we employ multiple latent variables to induce a mixture of expert NPs to specify diverse functional priors.
Meanwhile, the discrete latent variables as assignment variables are introduced to establish connections between a single data point and expert NPs.
To optimize this model with hybrid types of latent variables, we utilize variational inference to formulate the evidence lower bound.
Additionally, special modules are designed to accommodate few-shot supervised learning and meta reinforcement learning tasks.
Outline & Contributions.
We overview the properties of NPs and general notations for context-based meta learning tasks in Section 3.
Section 2 summarizes related work in meta learning and the NP family.
In Section 4, M ixture o f E xpert N eural P rocesse s (MoE-NPs) are elaborated, together with required modules for meta learning tasks.
Experimental results and analysis are reported in Section 5.
Discussions and limitations are included in Conclusion Section.
In principle, our contribution is two-fold:
• We introduce a new exchangeable stochastic process, referred to as MoE-NPs, to enrich the family of NPs.
Our model inherits the advantages of both mixtures of experts models and NPs, enabling multi-modal meta representations for functions.
• We specify inference modules in MoE-NPs for few shot supervised learning and meta reinforcement learning tasks.
Extensive meta learning experiments show that MoE-NPs can achieve competitive performance in comparison to most existing methods.
Literature Review
Meta Learning.
Meta learning is a paradigm to enable fast learning (fast adaptation to new tasks) via slow learning (meta training in the distribution of tasks).
There exist several branches of meta learning algorithms.
Gradient-based meta learning algorithms, e.g. MAML [4] and its variants [9; 10; 11], perform gradient updates over model parameters to achieve fast adaptation with a few instances.
Metrics-based meta learning algorithms try to learn representations of tasks in distance space, and models e.g. prototypical networks [12; 13] are popular in computer vision domains.
As for context-based meta learning methods of our interest, latent variable models, e.g. NPs [1], are designed to learn task representations in a functional space.
This family does not require gradient updates in fast adaptation.
Neural Processes Family.
Apart from vanilla NPs or CNPs [2; 1], other variants are developed and these are built on various inductive biases.
To address underfitting issues, attention networks [14; 15] are introduced to augment NPs.
[16; 17] improve the generalization capability of NPs with help of convolutional operations.
To learn distributions of group equivariant functions, [18] has proposed
EquiCNP.
Similarly, SteerCNPs also incorporate equivariance to approximate stochastic fields [19].
Our work is to get NPs married with Mixture of Experts (MoEs) models [20; 21], which model datasets with a collection of expert NPs.
We provide more technical summary of the NP family together with other probabilistic meta learning methods [22; 23; 24; 25] in the Appendix (F).
Preliminaries
Notations.
The paradigm of meta learning is considered in the distribution of tasks p ( T ) , and a task sampled from p ( T ) is denoted by τ in this paper.
The form of a task depends on applications.
For example, a task of our interest in regressions can be a function f to fit, which is a realization from unknown stochastic processes [26].
Let D Cτ refer to a set of context points used to specify the underlying task, and D Tτ =
[ x T , y T ] is a set of target points to predict, e.g. f ( x T ) =
y T .
In the context-based meta learning with latent variables, we write the probabilistic dependencies in distributions of target functions as follows,
[Equation]
where the functional prior p ( z |D Cτ ) is injected in modeling via latent variables z .
Neural Processes.
The family of NPs [1] belongs to exchangeable stochastic processes [27].
A generative process is written as Eq.
(2) with a global Gaussian latent variable z placed in Eq.
[Equation]
where f θ is a mean function and σ 2 i is the corresponding variance.
In our settings, we treat the conditional neural process [2] as a special case in NPs, when the distribution of z is collapsed into a
Dirac delta distribution p ( z ) =
δ ( z − ˆ z ) with ˆ z a fixed real valued vector.
3.1 Few-Shot Supervised Learning
In the context-based meta learning, we formulate the few-shot supervised learning objective within the expected risk minimization principle as follows.
[Equation]
[Equation] The risk function L , e.g. negative log-likelihoods, measures performance of meta learning structure on the task-specific dataset D Cτ and D Tτ , and Θ means parameters of common knowledge shared across tasks and parameters for fast adaptation ( e.g , Θ denotes parameters of the encoder ϕ and decoder θ , and L is the approximate objective in NPs).
With the set of context points D Cτ = { ( x 1 , y 1 ) , . . .
, ( x m , y m ) } and the target points D Tτ , the posterior of a global latent variable z in Eq.
(1) is approximated with a variational distribution q ϕ ( z |D Tτ ) and an evidence lower bound (ELBO) is derived to optimize in practice.
A general meta training objective of NPs in few-shot supervised learning is L ( θ, ϕ ) in Eq.
(4), where p θ ( D Tτ | z ) =
(cid:81) ni
=1 p θ ( y i |
[ x i , z ]) .
E
[Equation]
3.2 Meta Reinforcement Learning
For the context-based meta reinforcement learning, the context points D Cτ are a set of random transition samples from an environment as D Cτ
= { ( s 1 , a 1 , s 2 , r 1 ) , . .
, ( s H , a H , s H +1 , r H ) } , where r t is the one-step reward after performing action a t at state s t .
Here D Cτ plays a role in task inference [28] to obtain the information bottleneck q ϕ ( z |D Cτ ) and D Tτ is the dataset of state action values to fit.
For example, in an off-policy meta reinforcement learning algorithm, e.g. PEARL [6] or FCRL [3], the general optimization objective consists of two parts: (i) to approximate distributions of task-specific optimal value functions in Eq.
(5), where Q θ is optimal Q -value with the state value ˆ V (ii) to maximize the cumulative rewards E τ (cid:2)
E q ϕ ( z |D Cτ )
[ R ( τ, z ; φ )] (cid:3) , where R is the expected cumulative rewards in the environment τ given policies π φ ( a | [ s, z ]) .
[Equation]
Different from the few-shot supervised learning, here the context points are not part of fitting dataset, which means D Cτ ̸⊂ D Tτ .
As implemented in [6], the prior distribution p ( z ) is typically selected as a fixed one, e.g. N (0 , I ) .
The induced distribution of task-specific value functions p ( Q θ ([ s, z ] , a )) enables posterior sampling [29] in meta learning scenarios, which brings additional benefits of exploration for continuous control problems.
Model & Algorithm
In this section, we present our developed MoE-NPs and connect them to the hierarchical Bayes framework.
Then approximate objectives are derived and stochastic gradient variational Bayes [30] is used to optimize the developed model.
Finally, specialized neural modules are described for
MoE-NPs application to different meta learning tasks.
We have attached concepts of variational priors and posteriors and detailed computational diagrams in training and testing in Appendix (C).
For the sake of simplicity, we derive equations w.r.t.
a task τ in the following section, but a batch of tasks are considered in training in implementation.
4.1 Mixture of Expert Neural Processes
Vanilla NPs often suffer underfitting in experiments [1; 14].
This can be attributed to expressiveness bottlenecks when employing a global latent variable in learning functional priors of tasks from unknown distributions [31].
To alleviate mentioned deficiencies, we make two modifications for the NP family.
(i)
Multiple functional priors are encoded in modeling with help of K expert latent variables, which can capture statistical traits, e.g. distributional multi-modality, in data points.
This setting is also an extension of
Mixture of Experts (MoEs) models [20; 32; 33; 34] to meta learning scenarios.
(ii) Like the gating mechanism in [5], assignment latent variables are included in modeling to select functional forms for each data point in prediction.
The resulting MoE-NPs can learn more expressive functional priors and exhibit the approximation power for local properties of the dataset.
Generative Process.
As displayed in Fig.
(1), the graphical model involves two types of latent variables, respectively the continuous expert latent variables z 1: K and the discrete assignment latent variable e .
Further, we can translate the generative process into equations as follows,
[Equation]
where the sampled assignment variable e is in the form of K -dimensional one-hot encoding, and e k = 1 in Eq.
(6) indicates the k -th expert z k is selected from z 1: K for prediction.
A more detailed probabilistic generative process can also be found in Appendix (E.1).
In this way, our developed model constitutes an exchangeable stochastic process .
And we demonstrate this claim with help of
Kolmogorov Extension Theorem [35] in Appendix (E.2).
Link to Hierarchical Bayes.
Note that latent variables in Eq.
(6) are of hybrid types.
K functional priors are incorporated in expert latent variables z 1: K , while the assignment latent variable e is input dependent.
The dependencies between z 1: K and e are reflected in modeling, and this connects our work to Hierarchical Bayes [36; 37] in a latent variable sense.
Also when only one expert latent variable is used here, the hierarchical model degenerates to the vanilla (C)NPs [1; 2].
4.2 Scalable Training & Prediction
Inference Process.
Given a task τ , due to existence of unknown latent variables, it is intractable to perform exact inference w.r.t.
p ( D T τ |D C τ ) .
As an alternative, we apply variational inference to our developed model.
Here we use [ x, y ] to denote a single data point from a set of target points D Tτ .
[Equation]
An example for the k -th expert latent variable z k can be in the form of a Gaussian distribution
N ( z k ; µ k , Σ k ) .
And in meta training processes, the assignment variable e is assumed to be drawn from a categorical distribution Cat ( e ; K, α ( x, y, z 1: K )) with parameters
α ( x, y, z 1: K ) .
The existence of discrete latent variables e makes it tough to optimize using traditional methods.
This is because either sampling algorithms or expectation maximization algorithms are computationally intensive when utilized here (we have discussed this point in Appendix (G)) for expert latent variables.
To reduce computational cost, we again utilize variational inference and the decoder directly formulates the output as a mixture of log-likelihoods E q ϕ 2 , 1 [ln p θ ( y | x, z 1: K , e )]
= (cid:80)
Kk
=1 α k ( x, y, z 1: K ; ϕ 2 , 1 ) ln p θ ( y | x, z k ) .
This results in a general ELBO as Eq.
(7) for few-shot supervised learning in meta training, where q ϕ 1 denotes a collection of K independent variational distribution { q ϕ 1 , 1 , . . .
, q ϕ 1 ,K } .
q ϕ 2 , 1 and p ϕ 2 , 2 respectively define the variational posterior and prior for assignment latent variables.
Please refer to
Appendix (C)/(G) for definitions and more detailed derivations.
Monte Carlo Estimates & Predictions.
Meta-training processes consider a batch of tasks to optimize in iterations, and we apply Monte Carlo methods to the obtained negative ELBO L ( θ, ϕ 1 , ϕ 2 ) as follows.
[Equation]
With the number of tasks B and the number of data points N in mini-batches, the Monte Carlo estimate with one stochastic forward pass is Eq.
(8) for meta training objectives.
Like that in NPs [1], we derive the predictive distribution as Eq.
(9) with one stochastic forward pass and parameters of discrete latent variables
[Equation]
1: K ; ϕ 2 , 2 ) .
p ( y ∗ | x ∗ , D Cτ ) ≈
K (cid:88) k = α k ( x, z 1: K ; ϕ 2 , 2 ) p θ ( y | x ∗ , z k ) with z 1: K ∼ q ϕ 1 ( z 1: K |D Cτ ) (9)
And the point estimate in prediction E
[ Y | X = x, D Cτ ] can also be obtained in Appendix (G.4).
4.3 Module Details for Meta Learning 4.3.1 Inference Modules in MoE-NPs
[Figure]
The equations so far define a framework that can be implemented in different meta learning tasks.
Two examples are given in Fig.
Note that two types of latent variables are involved in modelling, we need different structures of encoders for latent variables.
Inference Modules are required for them, satisfying different conditions.
In variational inference, distribution parameters of these latent variables are approximated with the output of these specialized encoders.
Inference Modules for Continuous Latent Variables.
For neural networks to parameterize the encoder of continuous latent variables q ϕ 1 , we use the same architectures in (C)NPs [1; 2], which are permutation invariant to the order of context points [ x C , y C ]
= { [ x 1 , y 1 ] , . . .
, [ x N , y N ] } .
That is, for any permutation operator σ over the set of context points, the neural network ( NN ) parameters of an output distribution for each expert z k should satisfy [ µ k , Σ k ] =
NN ϕ 1 ,k
([ x σ (1: N ) , y σ (1: N ) ]) .
[Equation]
,i
[Equation]
Eq.
(10) is an example, where h is the embedding function, (cid:76) denotes a mean pooling operation, and g is part of encoder networks.
Inference Modules for Categorical Latent Variables.
For neural networks to parameterize the encoder of discrete latent variables q ϕ 2 , 1 and p ϕ 2 , 2 , we need the categorical distribution parameters α to be permutation equivariant [38] with respect to the order of z 1: K .
This means for any order permutation operation σ , the condition is satisfied as [ α σ (1) , α σ (2) , . . .
, α σ ( K ) ]
=
NN ϕ 2 , 1 ( x, y, z σ (1: K ) ) .
[Equation]
An example implementation for the variational posterior Cat ( e ; K, α ( x, y, z 1: K )) can be Eq.
(11), where the vector of logits is b =
[ b 1 , b 2 , . . .
, b K ] with t a temperature parameter.
And this implementation applies to prior networks NN ϕ 2 , 2 ( x, z σ (1: K ) ) to learn distribution parameters of
Cat ( e ; K, α ( x, z 1: K )) in the same way.
4.3.2 Meta RL Modules in MoE-NPs.
When extending MoE-NPs to meta RL tasks, optimization objectives in Eq.
(7) need to be modified for Actor-Critic methods, which are employed in our settings.
Like that in PEARL [6] and FCRL [3], the soft actor critic (SAC) algorithm [39] is used to learn policies due to good sample efficiency.
Given a specific MDP τ , posterior distributions of optimal value functions are formulated via latent variables z in context-based meta RL.
That is, p ( Q θ ( s, a ; M )) is approximated in the form p ( Q θ ([ s, z ] , a )) .
The resulting objectives for the Actor and Critic functions are respectively in Eq. (12) and Eq.
(13), where Z θ is a normalization factor.
[Equation]
The variational posterior q ϕ ( z | s, D Cτ ) in Eq.
(13) is a state dependent distribution with ˆ V a state value function, and sampling processes refer to steps in Algorithm (3).
[Equation]
A key difference from PEARL [6] lies in that several expert latent variables and assignment latent variables are involved in modeling.
So we refer the Kullback–Leibler divergence term to Eq.
(14) in
MoE-NPs with coefficient β 0 and β 1 .
[Equation]
The Monte Carlo estimates w.r.t.
Eq.
(12/13/14) are used in meta training, and Pseudo code to optimize these functions is listed in Appendix (A).
Experiments and Analysis 5.1 General Settings
The implementation of MoE-NPs in meta training can be found in Appendix Algorithms (1)/(3), and also please refer to Appendix Algorithms (2)/(4) for the corresponding meta-testing processes.
We leave the details of experimental implementations ( e.g. parameters, neural architectures, corresponding PyTorch modules and example codes) in Appendix (H).
Baselines for Learning Tasks.
Apart from MoE-NPs, methods involved in comparisons are contextbased methods such as CNPs [2], NPs [1] and FCRL [3], and gradient-based methods such as MAML [4] and CAVIA [9].
For FCRL, contrastive terms from SimCLR [40] are included in the objective.
In meta RL, the modified NP model corresponds to PEARL [6].
Meanwhile, in Appendix (I), we include additional experimental results compared with other NPs models augmented by attentive modules [14] or convolutional modules [17].
5.2 Illustration in Toy Regression
To see different roles of latent variables, we visually show effects of stochastic function fitting and quantified uncertainty in toy dataset.
Our goal is to discover potential components of distributions from limited observed data points.
[Figure]
The learning data points are sampled in x -domain
[ − π, 3 π ] and merged from a mixture of randomized functions f 1 ( x ) = sin( x ) + ϵ 1 and f 2 ( x ) =
cos( x ) + ϵ 2 with equal probability for mixture components, where ϵ 1 ∼ N (0 , 0 .
03 2 ) and ϵ 2 ∼ N (0 , 0 .
In each training iteration, we sample a batch of data points and randomly partition context points and target points for learning.
In testing phase, we draw up 100 data points from this mixture of distributions with 15 random data points selected as the context.
The fitting results for one sampled mixture curve are shown in Fig.
It can be seen that both CNPs and FCRL display similar patterns, overestimate the uncertain in the mixture curve of the second component.
NPs show intermediate performance and still fails to match context points well.
As for MoE-NPs, with help of predicted assignment variables parameters e ∗ = one_hot [arg k ∈{ 1 , 2 } max α k ] , we set the number of experts as two and partition data points to visualize predictive distributions p θ ( y ∗ | x ∗ , z ∗ ) .
The MoE-NP is able to precisely separate mixture components inside the dataset and provides more reliable uncertainty.
5.3 Few-Shot Supervised Learning
We evaluate the performance of models on a system identification task in Acrobot [41] and image completion task in CIFAR10 [42].
Both tasks are common benchmarks in the meta learning or NPs literature [41; 43; 1; 2; 14].
[Table]
System Identification.
For Acrobot systems, different tasks are generated by varying masses of two pendulums.
A dataset of state transitions is collected by using a complete random policy to interact with sampled environments.
The state consists of continuous as angles and angular velocities [ θ 1 , θ ′ 1 , θ 2 , θ ′ 2 ] .
The objective is to predict the resulting state after a selected Torque action from {− 1 , 0 , +1 } .
For more details about meta training dataset and environment information, refer to
Appendix (H.2).
In the meta testing phase, 15 episodes with the length of horizon 200 are collected for each task and we report the average predictive errors and standard deviations for all transitions.
Here we usetransitions as the context points to identify the task.
As exhibited in Table (1), gradient-based methods, e.g. CAVIA and MAML, beat NP in terms of predictive accuracy but show higher variances than all other models.
With three experts in modeling, MoE-NPs significantly outperform other baselines in terms of dynamics prediction.
Our finding is consistent with observations in [44], where multi-modal distributions are necessary for Acrobot systems.
We also illustrate the asymptotic performance of
MoE-NPs with the increase of the context points in the following Section (5.5) Ablation part.
Image Completion.
We use CIFAR10 dataset [42] in this experiment, which is formulated with 32x32 RGB images.
In the meta training process, a random number of pixels are masked to complete
[Figure] -NP in images. That is, given the context pixel locations and values [ x C , y C ] , we need to learn a map from each 2-D pixel location x ∈ [0 , 1] 2 to pixel values y ∈ R 3 .
Here two expert latent variables are used in MoE-NPs.
In Fig.
(4), we evaluate image completion performance on the test dataset and the number of context pixels is varied in three levels.
It can be found that CAVIA works best in cases withrandom context pixels or less than 500 ordered context pixels.
In other cases, MoE-NP surpasses all baselines.
With more observed pixels, the predictive errors of MoE-NPs can be decreased in both random and ordered context cases.
An example of image completion results is displayed in Fig.
[Figure]
For gradient-based methods, CAVIA and MAML are sensitive to the number of context points and do not exhibit asymptotic performance like that in MoE-NP.
NP still suffers underfitting in performance.
5.4 Meta Reinforcement Learning
To evaluate the meta RL implementation of our model, we conduct the experiments in a 2-D point robot and Mujoco environments [7].
Fig.
(6) exhibits the environments used in this paper, and we leave more details in Appendix (H.1). 2D Navigation.
For 2-D point robot tasks, the agent needs to navigate with sparse rewards.
The navigation goal of a task is sampled from a mixture of arcs in a semi-circle in Fig.
(6.a) during meta training processes.
From Fig.
(7.a-c), we can observe the evaluation performance of agents over iterations.
For gradient-based methods, CAVIA shows better performance than MAML in exploration but both are weaker than context-based baselines.
MoE-NPs can converge earlier with less training samples and show slight advantage over vanilla PEARL.
In particular, we test the asymptotic performance in out of distributions (O.O.D.) tasks and show results in Fig. (7.d).
We notice O.O.D. tasks are challenging for all algorithms to generalize but average returns are gradually increased with more trials.
PEARL and FCRL achieve comparable rewards, while MoE-NP behaves better in this case.
[Figure]
Locomotion.
As for Half Cheetah-Complex-Direction (HCheetah-CD) and Slim Humanoid-Complex-Goal (S-HumanoidCG) tasks, these correspond to locomotion in complicated environments.
Note that multiple directions and goals are involved in tasks.
Fig.
(7) illustrates the performance of learned policies in meta learning tasks.
In H-Cheetah-CD, MoE-NP shows a slight advantage over FCRL, and it exhibits comparable performance in S-Humanoid-CG.
In both environments, MoE-NP and FCRL outperform other baselines.
This implies the importance of func- [Figure] tional representations for task-specific value functions.
Either contrastive or multiple functional priors lead to better exploration and have a potential to boost performance in continuous control.
For gradient-based methods, observations show that they can easily get stuck in the local optimal [6; 45].
5.5 Ablation Studies
Number of Experts.
We examine the influence of the number of experts in meta-trained MoE-NPs, and system identification in the Acrobot system is selected as an example here.
As displayed in Fig.
(8), the number of experts are 3, 5,and 9 in different MoE-NPs.
Here we test the predictive performance of meta-trained MoE-NP by varying the number of transitions.
We set respectively 15, 25, 50, 100 transition samples as the number of context points to identify the system.
All settings for meta testing processes are already described in Section (5.3).
It can be seen when the expert number is 5, the predictive performance is largely enhanced with the increase of context points and the variance shrinks accordingly.
But with more experts, e.g. greater than 5, MoE-NPs exhibit higher predictive errors, and show no significant performance improvement with the increase of the number of context points.
These suggest increasing the number of experts beyond a certain point can deteriorate the predictive performance of the MoE-NPs.
[Figure]
Latent Variables in Meta RL.
As mentioned in Preliminaries
Section, the use of latent variable is able to induce task-specific optimal value functions.
Here we take the S-Humanoid-CG as the example, and the expert encoder of MoE-NPs (Deterministic) is set to the deterministic.
In Fig.
(9), we observe the performance degrades a lot using deterministic expert latent variables.
These further verify findings in PEARL [6].
The randomness of value function distributions captures task uncertainty and encourages more efficient exploration.
Conclusions
Technical Discussions.
In this work, we have developed a new variant of NP models by introducing multiple expert latent variables.
Our work illustrates the roles of different latent variables in MoE-NPs for meta learning tasks.
MoE-NPs are able to separate data points from different clusters of stochastic processes and exhibit superior performance in few-shot supervised learning tasks.
Also, MoE-NPs are consistently among the best methods in meta learning continuous control.
Existing Limitations.
Though the developed model provides more expressive functional priors, the appropriate number of experts is still hard to determine.
Also, the mechanism of gradually incorporating new expert latent variables has not been explored and this raises concerns in additional computational cost and more effective inference.
Future Extensions.
Here we provide a couple of heuristics to determine optimal number of experts for
MoE-NPs in the future.
Information metrics, e g. Bayesian information criterion, can be incorporated in modeling.
Another way is to place priors over distributions of discrete latent variables like that in hierarchical Dirichlet processes [46] and select the optimal number of experts in a Bayesian way. Acknowledgement
We thank NeurIPS anonymous reviewers and meta reviewers for their comments and constructive suggestions for our manuscript.
Their engagement in the review/rebuttal period helps improve our manuscript a lot
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] or [N/A] .
You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[Yes] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
[-Contents
-]Contents 1 Introduction 2 Literature Review 3 Preliminaries 3.
Few-Shot Supervised Learning . .
Meta Reinforcement Learning . .
4 Model & Algorithm 4.
Mixture of Expert Neural Processes . . . . . . . . . . . . . . . . . . . . . . . . .
Scalable Training & Prediction . .
Module Details for Meta Learning . . . . . . . . . . . . . . . . . . . . . . . . . .
Inference Modules in MoE-NPs . . . . . . . . . . . . . . . . . . . . . . .
Meta RL Modules in MoE-NPs. . . . . . . . . . . . . . . . . . . . . . . .
5 Experiments and Analysis 5.
General Settings . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Illustration in Toy Regression . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Few-Shot Supervised Learning . .
Meta Reinforcement Learning . .
Ablation Studies . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
6 Conclusions[-
-]
A Pseudo Code of Algorithms in Meta Learning[-
-]
B Frequently Asked Questions[-
-]
C Probabilistic Graphs in Meta Training/Testing[-
-]
D More Descriptions of NP Family Models and Meta RL[-
-]
E MoE-NPs as Exchangeable SP s[-
-]
E.1 Generative Processes . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.2 Consistency Properties . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F Summary of Existing NP Related Models[-
-]
F.
Comparison in Technical Details . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Time Complexity . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Additional Literature Review . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G Formulation of Evidence Lower Bounds[-
-]
G.1 Variational Distributions . .
G
.2 Lower Bound on the Evidence for Few-Shot Supervised Learning
G.3 Selection of Categorical Approximate Posteriors/Priors . . . . . . . . . . . . . . .
G.4 Stochastic Gradient Estimates . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.5 Estimates of Statistics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H Experimental Settings and Neural Architectures
H.1 Dataset & Environments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.1.1 Dataset in Few-shot Supervised Learning . . . . . . . . . . . . . . . . . .
H.1.2 Environments in Meta Reinforcement Learning . . . . . . . . . . . . . . .
H.2 Implementations in Meta Learning Tasks . . . . . . . . . . . . . . . . . . . . . . .
H.2.1 Toy Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.2.2 Few-Shot Supervised Learning . .
H.2.3 Meta Reinforcement Learning . . . . . . . . . . . . . . . . . . . . . . . .
H.3 Neural Architectures . .
I Additional Experimental Results
[-Results 
-]
I.
Additional Analysis of Learned Latent Variables . . . . . . . . . . . . . . . . . . .
I.
Additional Results of NP Variants in Toy Regression . . . . . . . . . . . . . . . .
I
Additional Results of NP Variants in System Identification . . . . . . . . . . . . .
I
Comparison with Attentive Neural Processes . . . . . . . . . . . . . . . . . . . .
I.
Augmenting MoE-NPs with Convolutional Modules . .
I.
More Visualization Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J Computational Devices
parameters θ .
1 Initialize parameters ϕ θ ; 2 for i 1 to m do
Sample a random value
N C ∼ U [1 , N max ] as the number of context points;
Sample mini-batch instances D to split dataset { ( x C , y C , x T , y T ) bs } Bbs =1 ; // generative process
Sample expert latent variables z 1: K ∼ q ϕ 1 ( z 1: K |D T ) ;
Compute distribution parameters α with Eq.
Compute negative ELBOs L MC ( θ, ϕ ) in Eq.
; // amortized inference process ϕ ← ϕ − λ 1 ∇ ϕ L MC ( θ ϕ ) in Eq.
; θ ← θ − λ 2 ∇ θ L MC ( θ ϕ ) in Eq.
; 12 end
: Predictive distributions.
parameters
; 2 Set the number context points N C ; 3 Split test dataset into the /target 4 { ( x C , y C , x T , y T ) bs } Bbs =1 ∼ D ; // generative process 5 Sample expert latent variables of 6 the mini-batch z 1: K ∼ q ϕ 1 ( z 1: K |D Cτ ) ; 7 if discrete l.v.s for hard assignment then
Sample assignment latent variables of the mini-batch e ∼ p ϕ 2 , 2 ( e | x T , z 1: K ) ;
Output the distribution p θ
; 11 else
Compute the distribution parameters α of assignment latent variables via Output the predictive distribution p
; 16 end do
Sample a batch of tasks { τ } B ∼ p ( T ) ; // collect context transitions for each τ ∈ { τ } B do
Initialize the context D Cτ = {} ;
Execute Algorithm (4) in Appendix to update D Cτ end // actor critic learning in batchs for i = 1 to m do for each τ ∈ { τ } B do
Sample context points D Cτ ∼ S c ( M Cτ ) & batch of transitions b τ ∼ M τ ;
Sample z 1: K ∼ q ϕ 1 ( z 1: K |D Cτ ) ; for each s
∈ b τ do
Sample e ∼ q ϕ 2 ( e | z 1: K , s ) to select z and augment the state as [ s, z ] ∈ b τ ; end // run forward propagation
L τA =
L τA ( b τ ) in Eq.
L τC =
L τC ( b τ ) in Eq.
L τKL =
L τKL ( D Cτ , b τ ) in Eq.
(14) end // run back propagation ϕ ← ϕ − λ 1 ∇ ϕ (cid:80) τ
( L τC + L τKL ) in Eq.
(13/14); φ ← φ − λ 2 ∇ φ (cid:80) τ
L τA in Eq.
(12); θ ← θ − λ 3 ∇ θ (cid:80) τ
L τC in Eq.
(13); end 26 end
Input : MDP distribution p
B Frequently Asked Questions
Here we collect some frequently asked questions from reviewers and other literature researchers.
We thank these reviewers for these precious questions and add more explanation.
Selection of Benchmarks.
Admittedly, NP variants can be applied a series of downstream tasks.
Our selection of benchmark missions is based on existing literature for NP models.
The system identification task was previously investigated with NP variants in work [43; 47], in which learning transferable physics dynamics with NPs is a crucial application.
The image completion task is more commonly used in work [9; 2; 1].
The meta reinforcement learning task can also be studied within
NP framework [3; 48].
NP Family in Classification Tasks.
We have tried to search neural architectures for few-shot image classification tasks, but the performance is not ideal in comparison to other metrics based methods.
Meanwhile, we have gone through most NP related work, and it is challenging to achieve SOTA few-shot image classification results with standard neural architectures in NPs, e.g. multi-layer perceptrons (MLPs).
Maybe this is due to the nature of stochastic processes, which can address regression problems more efficiently.
Unless specialize modules instead of MLPs are used, we do not expect NP variants with MLPs can achieve SOTA performance.
The aim of this paper focuses more on mixture expert inductive biases and place less attention on neural architecture search.
So MLPs are shared across all baselines to enable fair comparison.
Expressiveness of Mixture of Expert Inductive Biases.
A natural question about learning diverse functional representation is whether these multiple expert latent variables will collapse into one.
We refer the collapsed representation to vanilla (C)NPs, and the previous empirical results show the collapsed ones work poorer than mixture of expert ones in both few-shot regression tasks and meta reinforcement learning tasks.
Also, from the multimodal simulation result, we discover both latent assignment variables and expert variables are meaningful, which reflects the effectiveness of mixture expert inductive biases.
That means, we have not encountered meta representation collapse issues in experiments.
Extension with other Mixture of Experts Models.
Our work is the first time to examine MoEs inductive biases in NPs family, and the used MoE module is an amortized inference one.
We have not found a trivial implementation of MoEs in meta learning domain.
But in MoEs literature, there exist other more effective MoE models, which can better trade off communication/memory and performance.
So NPs family can also be combined with these models, such as GShard [49],
Deepsepeed MoEs [50] and etc.
Potential Applications in Industry.
Here we provide two available applications with MoE-NPs in the industry.
One is in multilingual machine translation or multilingual language auto-completion.
In this case, a mixture of experts corresponds to multilingual functional priors for multi-modal signals [51] and enables the prediction with partial observations.
Another application lies in modelling irregular time series [52; 53].
In this case, diverse experts can handle discontinuous components in a rich family of stochastic functions.
Meanwhile, the entropy of learned assignment latent variables can tell us the regions likely to be discontinuous, which is quite helpful in anomaly detection in a black-box system.
C Probabilistic Graphs in Meta Training/Testing
As exhibited in Fig.
(10)/(11), shown are computational diagrams when implementing MoE-NPs in meta learning tasks.
Variational Posteriors.
Since the real posteriors for both { z k } Kk =1 and e T are computationally intractable, the approximate ones are used in practice.
These are called variational posteriors, e.g. q ϕ 1 ,k for an expert latent variable z k and q ϕ 2 , 1 ( e | x, y, z 1: K ) for assignment latent variables e .
For the sake of simplicity, we denote { q ϕ 1 ,k } Kk =1 by q ϕ 1 in some time.
Importantly, the Gumbel-softmax trick [54] is used to sample assignment latent variable e from categorical approximate distributions.
Variational Priors.
In some cases, the prior distributions for { z k } Kk =1 and e T are set to constrain the scope of prior distributions.
For example, in few-shot supervised learning, since the context and the target have the same form, the variational prior is selected to be q ϕ 1 as well to ensure the consistency and this works in meta testing phases.
For the assignment latent variable e , this uses the same form in conditional VAE [55] as p ϕ 2 , 2 .
[Figure]
D More Descriptions of NP Family Models and Meta RL
In the main paper, we unify the description of NP family models in both few-shot supervised learning and meta reinforcement learning.
This is the same with that in FCRL [3].
Meta learning in NP related models is to learn functional representations of different tasks and formulate the fast adaptation via inferring the task specific conditional distribution p ( D Tτ |D
Cτ ) =
(cid:82) p ( D Tτ | z ) p ( z |D Cτ ) dz (equivalent to Eq.
To make the downstream reinforcement learning task using NP family models clearer, we add the following explanations.
In few-shot supervised learning, D Cτ and D Tτ are of the same form.
However, in context-based meta reinforcement learning, D Cτ is a set of task specific transitions and D Tτ is a set of state (action) values.
As the result, the approximate posteriors and the selected priors to resolve
Eq.
(1) are distinguished in separate meta learning cases.
In context-based meta reinforcement learning, we can translate our problem into finding the distribution of optimal value functions in Eq.
(5), this corresponds to learning meta critic modules with NP family models.
Given a transition sample [ s, a, r ( s, a ) , s ′ ] , the target input is the state x T
= s and the target output is the temporal difference target y T
= ˆ Q ( s, a ) =
r ( s, a ) + γV ([ s ′ , z ′ ]) .
The standard
Gaussian distribution is used as the prior p ( z k ) =
N (0 , I ) in Eq.
(14), while the approximate posterior is learned from D Cτ with permutation invariant functions.
In total, sampling from the state dependent approximate posterior z ∼ q ϕ
( z | s, D C τ ) corresponds to Eq.
(15), where the operator ⊙ denotes the selection process with help of Hadamard products.
[Equation]
E MoE-NPs as Exchangeable SP s E.1 Generative Processes
To better understand our developed model in meta learning scenarios, we translate Eq.
(6) into a step-wise generative process.
The same with that in the main paper, the task distribution is denoted by p ( T ) and we presume K -experts to summarize the stochastic traits of a task.
[Equation]
Here a MoE-NP for the task τ is specified with K -expert latent variables z 1: K in Eq.
(16.a).
The probability mass function for a data point related categorical distribution Cat ( K, α ( x, z 1: K )) is denoted by p ( e | x, z 1: K ) =
(cid:81)
Kk =1 α k ( x, z 1: K ) I [ e k =1] in Eq.
(16.b), and e is an assignment latent variable to select an expert for the generative process.
After that, the distributional parameters for the output of a data point are learned via a function g θ in Eq.
(16.c), followed by the output distribution
N ( µ x , Σ x + ϵ 2 I ) .
Note that a collection of sampled functional experts are represented in a vector of variables z 1: K =
[ z 1 , z 2 , . . .
, z K ] T and e = [0 , · · · , 1 (cid:124)(cid:123)(cid:122)(cid:125)
k -th position , · · · , 0] T ⇔ e k = 1 is a one-hot vector in Eq.
(16.b).
In
Eq.
(16.c), the expert is selected in a way z
= z 1: K ⊙ e .
For the sake of generality, irreducible noise
N (0 , ϵ 2 I ) is injected in the output.
In experiments, K -expert latent variables z 1: K as well as discrete assignment latent variables e are non-observable.
E.2 Consistency Properties
Definition 1.
(Exchangeable Stochastic Process)
Given a probability space be (Ω , F , P ) , let µ x 1 ,...,x N be a probability measure on R d with { x 1 , . . .
, x N } as a finite index set.
The defined process is called an exchangeable stochastic process ( SP ), S : X × Ω → R d such that µ x 1 ,...,x N ( F 1 × · · · × F N ) =
P ( S x 1 ∈ F 1 , . .
S x N ∈ F N ) when it satisfies the exchangeable consistency and marginalization consistency.
Remember that the generative model is induced in the main paper as follows.
And we claim that our designed generative model MoE-NP formulates a family of exchangeable SP in Definition 1 .
[Equation]
So it is necessary to verify two formerly mentioned consistencies according to Kolmogorov Extension
Theorem [35] and de Finetti’s Theorem [56].
This is to show the existence of SP s in Eq.
Exchangeability Consistency.
For N data points from Eq.
(17), we impose any permutation operation σ over their indices, and this results in σ : [1 , 2 , . . .
, N ] → [ σ 1 , σ 2 , . .
, σ N ] .
Then we can check the following equation is satisfied since the element-wise product of probabilities can be swapped.
[Equation]
Marginalization Consistency.
Given the assumption that the integral in Eq.
(17) is finite, we pick up a subset of indices [ M + 1 , M + 2 , . . .
, N ] and make M < N without difference in orders.
And the result after marginalization over y -variable in the selected indices can be verified based on the following equation.
[Equation]
: N [Equation] Built on these two sufficient conditions, our developed MoE-NP is a well defined exchangeable SP .
F Summary of Existing NP Related Models F.1 Comparison in Technical Details
Here we give a brief summary on difference between MoE-NPs and existing typical Neural Process models in Table (2).
Some crucial traits include forms of encoders and decoders structures, types of latent variable and inductive biases injected in modelling.
Especially, the inductive bias for MoE-NP is reduced to be multiple functional priors, which means a collection of expert neural processes to induce the generated dataset.
Since a general inductive bias behind NPs related models is the modelling of exchangeable stochastic processes with cheap computations.
This corresponds to a distribution of functions, termed as functional in the Table.
Note that the recognition model of
NP models in Meta Training Scenarios is replaced with q ϕ ( z | [ x T , y T ]) since all target points can be available, but in Meta Testing Scenarios, only [ x C , y C ] are accessible.
F.2 Time Complexity
As for running time complexity, the vanilla NPs and CNPs are with O ( N + M ) , while MoE-NPs are with O ( K ∗ ( N + M )) (making M predictions with N observations).
In practice, the number of experts is small, so the increase of running time complexity can be ignored in practice.
In contrast, traditional Gaussian processes are O (( N + M ) 3 ) in terms of running time complexity.
F.3 Additional Literature Review
Due to page limit in the main paper, we include other related works in this subsection.
In unsupervised learning, the Neural Statistician Model [57] is introduced to compute summary statistics inside the dataset.
The Generative Query Network [48], a variant of NPs for visual sensory dataset, makes use of a latent variable to abstract scenes in high dimensions.
To capture heteroscedastic noise inside the stochastic process, DSVNP [58] induces latent variables at different levels.
The functional neural
[Table]
processes infer the directed acyclic graph in the latent space and formulate flexible exchangeable stochastic processes for single task problems [59].
Inspired by self-supervised learning, [3; 60] propose to augment the neural process with contrastive losses.
[61] combines context memories and recurrent memories to formulate sequential neural processes (SNPs).
Though there exist a number of
NP variants, none of them consider to inject multiple functional inductive bias in modeling.
G Formulation of Evidence Lower Bounds
Since functional priors reflected in the K -expert latent variables z 1: K are learned via approximate distributions, this can be directly optimized within the framework of variational inference.
So we leave these out in this discussion.
The difficulty of optimization principally comes from the involvement of discrete latent variables.
We therefore discuss chance of using another traditional optimization algorithm, called Expectation Maximization (EM) [62], in our settings.
Omitting the K -expert latent variables z 1: K and corresponding variational distributions, we take a closer look at the assignment latent variable e in the logarithm likelihood as ln (cid:16)(cid:80)
Kk =1 p ( y | x, z 1: K , e k = 1) p ( e k = 1 | x, z 1: K ) (cid:17) and derive the corresponding EM algorithm.
Expectation(E)-Step : Note that the assignment variable e is discrete with the categorical probability function p ( e | x, z 1: K ) =
Cat ( e ; K, α ( x, z 1: K )) .
This step is to update the posterior of the proportional coefficients α ( x, z 1: K ) based on the last time step model parameters θ ( t ) .
[Equation]
Here the prior distribution p ( e ) can be a commonly used one Cat ( K, [ 1 K , 1 K , . .
, 1 K ]) or the last time updated one p ( t ) ( e ) .
As a result, updated categorical distribution parameters are:
[Equation]

[Equation]

 (21) where τ is the temperature parameter.
M )-Step
Once the distributional parameter of assignment latent variables are updated, the next step is to maximize the logarithm likelihood as θ ( t +1)
=
arg max θ
[Equation]
x,y ) ∈D ln
[ p θ ( t ) ( y | x, z 1: K , e )] given the last time updated model parameter θ ( t ) .
With help of gradient ascent, this can be written as follows,
[Equation]
where λ is the learning rate.
Note that the coefficient α is data point dependent and the derivation of EM algorithms considers a subset of data points D .
However, in meta learning scenarios, we handle large-scale dataset and the above-mentioned EM framework is computationally expensive and impractical.
Due to these considerations, we do not apply EM algorithms to estimate the discrete distribution and instead variational inference is employed for the assignment latent variable in optimization.
G.1 Variational Distributions
For continuous latent variables, diagonal Gaussians are commonly used as variational distributions.
With Gaussian variational posteriors N ( z ; µ, Σ) and corresponding priors N ( z ; µ p , Σ p ) , the
Kullback–Leibler Divergence can be analytically computed as follows.
[Equation]
Meanwhile, when it comes to categorical distributions, the corresponding prior distribution is selected as Cat ( K, α 0 ) with distribution parameters α
=
[ α 0 , 1 , α 0 , 2 , . .
, α 0 ,K ] .
And the Kullback–Leibler
Divergence is computed as follows.
[Equation]
When α 0 =
[ 1 K , 1 K , . .
, 1 K ] , the divergence is further simplified as follows.
[Equation]
G.2 Lower Bound on the Evidence for Few-Shot Supervised Learning
Since K -expert latent variables are independent in settings, we denote the corresponding variational parameters by q ϕ 1 = { q ϕ 1 , 1 , q ϕ 1 , 2 , . . .
, q ϕ 1 ,K } .
ϕ 1 ,k denotes parameters of encoders for k -th expert model.
Hence, the distribution follows that q ϕ 1 ( z 1: K |D Cτ ) =
(cid:81)
Kk =
1 q ϕ 1 ,k
( z k |D Cτ ) and q ϕ 1 ( z 1: K |D Tτ ) =
(cid:81)
Kk =
1 q ϕ 1 ,k
( z k |D Tτ ) .
Note that ( x, y ) ∈ D Tτ and the variational posterior for expert latent variables are q ϕ 1 ( z 1: K |D Tτ ) in the general NPs.
As for the variational posterior for assignment latent variables, we choose q ϕ 2 , 1 ( e | x, y, z 1: K ) as default.
We will use these notations to formulate the evidence lower bound (ELBO) as follows.
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
By introducing the variational distribution q ϕ 2 for the discrete assignment latent variable e , Eq.(26.d) is further bounded by Eq.
(26.e-g).
Recall that when vanilla NP modules are used here, the approximate posterior in Eq.
(26) in meta training should be substituted with q ϕ 1 ( z 1: K |D Tτ ) with the corresponding approximate prior p ( z k ) =
q ϕ 1 ,k
( z k |D Cτ ) .
And this matches the general form in the main paper for −L ( θ, ϕ 1 , ϕ 2 ) in Eq.
When Dirac delta distributions are used in MoE-NPs, the divergence term about the continuous latent variable is removed as default.
Denoting the approximate posterior by q ϕ 2 , 1 ( e | x, y, z 1: K ) =
Cat ( e ; [ α 1 ( x, y, z 1: K ) , α 2 ( x, y, z 1: K ) , . .
. , α K ( x, y, z 1: K )]) , we rewrite the log-likelihood inside the ELBO as Eq.
[Equation]
As for the approximate posterior of the assignment latent variable q ϕ 2 , 1 ( e | x, y, z 1: K ) , we provide two ways of implementations in our experiments: (i) use the target input y as the additional input to formulate q ϕ 2 , 1 ( e | x, y, z 1: K ) (ii) use the same form as the conditional prior q ϕ 2 , 1 ( e | x, y, z 1: K )
=
p ϕ 2 , 2 ( e | x, z 1: K ) .
G.3 Selection of Categorical Approximate Posteriors/Priors
As previously observed in Acrobot system identification results, increasing the number of expert latent variables tends to weaken the generalization capability.
This also happens in image completion, so we set the number of experts used is 2 in the task.
We can attribute this to inference sub-optimality in categorical approximate posteriors/priors.
Remember that in image completion and Acrobot system identification, the used approximate posterior for the categorical latent variable is q ϕ 2 , 1 ( e | x, y, z 1: K ) with the target information y for the input.
Since the developed MoE-NP is a VAE-like models [30], the number of expert latent variables
K decides the dimension of the assignment latent variable e .
In auto-encoder models, when the dimension of latent variables in all hidden layers is higher than that of the input, the model tends to copy the input to the output and fails to learn effective representations.
This is the direct source of overfitting and applies to conditional VAE methods [55].
For example, the dimension of output in
Acrobot is 6, which implies the bottleneck constraint is weaker when the number of experts is greater than 6.
It is reasonable to alleviate such sub-optimality by directly using the conditional prior as the approximate posterior q ϕ 2 , 1 ( e | x, y, z 1: K ) =
p ϕ 2 , 2 ( e | x, z 1: K ) .
You can find more clues from the following stochastic gradient estimates for the assignment latent variables in Eq.
Meanwhile, we report empirical results in image completion when q ϕ 2 , 1 ( e | x, y, z 1: K ) =
p ϕ 2 , 2 ( e | x, z 1: K ) in Sec.
(I.1).
And you can see inference in this way does not suffer the overfitting issue caused by more experts
[Equation]
G.4 Stochastic Gradient Estimates
Here the stochastic gradient estimates with respect to parameters in the negative ELBO L ( θ, ϕ 1 , ϕ 2 ) in Eq.
(7) are provided as follows.
[Equation]
The reparameterization trick [30] is used to sample values from variational distributions of expert latent variables throughout the inference process and stochastic gradient estimates in Eq.
In prediction processes, the way to get values of assignment latent variables follows that in [63].
Besides, we provide the stochastic gradient estimate for another case when the variational posterior for the assignment latent variable is selected as the variational prior, which means q ϕ 2 , 1 ( e | x, y, z 1: K ) =
p ϕ 2 , 2 ( e | x, z 1: K ) .
This case can drop off the divergence term for the discrete variable.
Let the conditional prior for the discrete variable be p ϕ 2 , 2 ( e | x, z 1: K ) =
Cat ( e ; [ α 1 ( x, z 1: K ) , α 2 ( x, z 1: K ) , . . .
, α K ( x, z 1: K )]) , we apply the log-derivative trick in a REINFORCE estimator [64] to Eq.
(30) and can obtain the following equation as the gradient estimator 1 .
[Equation]
As can be seen from Eq.
[Equation] the posterior update implicitly exploits supervision information.
G.5 Estimates of Statistics
Momentum.
Given the pre-trained MoE-NPs, we can formulate the statistical momentum in predictive distributions.
For the first order momentum, equivalently mean of the predictive distribution, we need to compute the conditional version E
[ Y | X
= x, D Cτ ] in meta learning scenarios.
Here the predictive distribution of one expert is parameterized in the form p
[Equation] where m is the learned mean function using a neural network and σ 2 is a variance parameter.
Using a single stochastic forward pass in expert latent variable z 1: K , we can derive the estimate of the predictive mean ˆ m =
E
[ Y | X
= x, D Cτ ] .
[Equation]
The second order moment can be estimated accordingly.
Here we consider the case when the output is one dimensional and Σ k
= σ 2 k .
[Equation]
Entropy.
Note that our developed MoE-NPs can also be applied to out of detection (O.O.D) tasks.
In this case, the entropy of predictive distribution plays a crucial role.
Though the exact estimate of the predictive entropy for MoE-NPs is intractable due to the complexity inside the mixture components, we can measure the expected result of the entropy E [ H ( Y ) | X
= x, D Cτ ] in prediction.
We still use a single stochastic forward pass in expert latent variable z 1: K in estimation.
If H ( Y k | X = x, z k ) =
− (cid:82) p ( Y =
y
| X = x, z k ) ln p ( Y = y
| X = x, z k ) dY is bounded ∀ k ∈ { 1 , . .
, K } , the estimate of entropy term is as follows.
[Equation]
[Equation]
H Experimental Settings and Neural Architectures
In this section, we provide with more experimental details.
Importantly, neural modules of MoE-NPs in the PyTorch version are listed.
For the few-shot regression, we provide an example of our implementation of MoE-NPs from the anonymous Github link https://github.com/codeanonymous233/ICMoENP .
For the context-based meta RL algorithms, you can find the implementation of MoE-NPs from the anonymous Github link https://github.com/codeanonymous233/MoENP .
H.1 Dataset & Environments H.1.1 Dataset in Few-shot Supervised Learning
System Identification.
Note that in the used Acrobot simulator 2 , the observation is the preprocessed state as a 6 dimensional vector [sin( θ 1 ) , cos( θ 1 ) , sin( θ 2 ) , cos( θ 2 ) , θ ′ 1 , θ ′ 2 ] .
The input of the Acrobot system is the concatenation of the observation and the executed action [sin( θ 1 ) , cos( θ 1 ) , sin( θ 2 ) , cos( θ 2 ) , θ ′ 1 , θ ′ 2 , a ] .
The output of Acrobot system is the predicted transited state.
We generate 16 meta training tasks by varying the masses of two pendulums m 1 and m 2 , which means the hyper-parameters of the system come from the Cartesian combination of the set m 1 ∈ { 0 . 75 , 0 . 85 , 0 .
95 , 1 . 15 } and m 2 ∈ { 0 . 75 , 0 . 85 , 0 .
In meta training processes, a complete random policy interacts with batch of sampled MDPs to formulate transition dataset.
As for meta testing tasks, we follow the same way to generate tasks by setting m 1 ∈ { 0 . 85 , 1 . 05 , 1 . 25 } and m 2 ∈ { 0 . 85 , 1 . 05 , 1 . 25 } .
Image Completion.
CIFAR10 dataset consists of 60000 32x32 color images in 10 categories.
Among these images, 50000 are for meta training with the rest for meta testing as the default in image completion tasks.
CIFAR10 images are processed via torchvision modules to normalize the pixel values between [0 , 1] .
H.1.2 Environments in Meta Reinforcement Learning
Note that 2-D point robot navigation tasks, the distribution for meta training is a mixture of uniform distributions [0 , 2 π/ 12] ∪ [5 π/ 12 , 7 π/ 12] ∪ [10 π/ 12 , π ] .
The rest of regions along the arc is for out of distribution tasks.
The tasks in Mujoco [7] follows adaptations from [6; 65], where goals/velocities or multiple hyper parameters of simulation systems are sampled from mixture distributions.
The horizon of an episode for mixture of point robots is 20, while that for Mujoco environments is 500.
The required numbers of environment steps in meta training processes are respectively 2 .
5 ∗ 1 e 6 for point robots, 7 . 5 ∗ 1 e 6 for Half-Cheetah-CD and 6 .
5 ∗ 1 e 6 for Slim-Humanoid-CG.
We leave more details and settings of each environment in the above github code link. 2-D Point Robot.
2-D robots attempt to reach goals located in specified regions of the arc.
Cheetah-Complex-Direction.
2-D Cheetah robots aim at running in given directions.
The task includes multiple target directions and these change with split steps of episodes.
Humanoid-Complex-Goals.
3-D Humanoid robots aim at running towards goals.
The task includes multiple goals and these change with split steps of episodes.
H.2 Implementations in Meta Learning Tasks H.2.1 Toy Experiments
The input of functions is in a range [ − π, π ] ∪ [ π, 2 π ] .
The general implementations of MoE-NPs are as follows.
The x -domain for the function f 1 ( x ) = sin( x )
+ ϵ 1 with ϵ 1 ∼ N (0 , 0 .
03 2 ) is [ − π, π ] , while that for f 2 ( x ) = cos( x ) + ϵ 2 with ϵ 2 ∼ N (0 , 0 .
01 2 ) is [ π, 2 π ] .
Sampling from these two components leads to a mixture dataset.
The Encoder for all continuous latent variables is with two hidden layers (32 neuron units each). The Gaussian distribution is used for continuous latent variables in MoE-NPs. The Encoder for the discrete assignment latent variable in MoE-NPs corresponds to a softmax-output neural network with two hidden layers (16 neuron units each). The Decoder is with three hidden
H.2.2 Few-Shot Supervised Learning
For gradient based methods, we use implementations of MAML 3 and CAVIA 4 .
System Identification.
The general implementations are as follows.
For all NP related models, the dimension of a latent variable is set to be 16.
The Encoder for all continuous latent variables is with two hidden layers (32 neuron units each).
Note that Dirac delta distributions are used for continuous latent variables in MoE-NPs since this case works best in few-shot supervised learning.
For MoE-NPs, we use three expert latent variables as the default, and the Encoder for the discrete assignment latent variable in MoE-NPs corresponds to a softmax-output neural network with two hidden layers (32 neuron units each).
The Decoder is with four hidden layers (200 neuron units each) as well.
The number of tasks in batch training is 4, batch size in training is 200 transition steps (for each task).
The horizon for each episode of transitions is 200 time steps.
In each iteration of meta training, 4 different environments are randomly selected, and the uniform random controller is used to interact for the collection of 5 episodes (for each task).
In training dynamics systems, the training batch size in transition buffer is 200, the training epoch is 5, and the whole process iterates until convergence (the iteration number is 25).
The learning rate for Adam optimizer is 1 e − 3 as the default.
Image Completion.
The general implementations follow that in [9; 2] and are applied to all baselines and MoE-NPs.
The dimension of a latent variable is set to be 128.
The Encoder for all NP variants is with three hidden layers (128 neuron units each).
Note that Dirac delta distributions are used for continuous latent variables in MoE-NPs since this case works best in few-shot supervised learning.
The Decoder is with five hidden layers (128 neuron units each) as well.
For MoE-NPs, we use three expert latent variables as the default, and the Encoder for the discrete assignment latent variable corresponds to a softmax-output neural network with two hidden layers (32 neuron units each).
Adam [30] is used as the optimizer, where the learning rate is set to be 5 e − 4 .
The batch size in training is 8 images, and we meta train the model until convergence (the maximum epoch number is 50).
Also note that the number of context pixels in CAVIA is 10 for fast adaptation in default implementations, and this leads to the best testing result of CAVIA in 10 pixel cases in Fig.
Note that, to train
NP models, including CNP/NP/FCRL/MoE-NP, we set the form of negative log-likelihood objective consistent based on that in [66].
But in evaluation, to keep results of all methods consistent, we follow that in [2; 9; 3] and report the MSEs in Fig.
(4) in the testing phase.
H.2.3 Meta Reinforcement Learning
In terms of implementations of baselines, we directly use the following open sourced code: PEARL 5 ,
MAML 6 and CAVIA 7 .
Note that the TRPO algorithm [67] is used for MAML/CAVIA as the default.
We do not change too much except the replacement of our environments when running experiments.
Further we provide more details on how to modify MoE-NPs in meta RL domains.
Notice that
MoE-NP can also be seen as a latent variable model, and there exists a close relationship with PEARL algorithms [6] when it comes into meta RL.
Implementations of MoE-NPs are the same as that in
PEARL [6] except latent variable distributions and the inference way.
Note that Soft Actor Critic (SAC) algorithm [39] is used in policy optimization, which requires parameterization of both actor and critic functions.
As for the number of experts in MoE-NPs, we use 3 for all environments as default.
You can find more details about neural architectures/optimizers for each modules from the above mentioned link.
As mentioned before, we use p ( z k ) =
N (0 , I ) as the prior distribution for expert latent variables.
The approximate posterior is parameterized with a diagonal Gaussian distribution.
The coefficient for KL divergence terms in Eq.
(14) are β 0 = 1 . 0 , β 1 = 1 . 0 .
The meta-training processes for reinforcement learning can be found in Algorithm (3) in the main paper.
In terms of meta-testing processes, we report the pseudo code in Algorithm (4)
H.3 Neural Architectures
To help readers better understand our models [Table] we attach the python code of the Encoder for separate latent variables as follows.
As for the Decoder , the structure is the same with that in a vanilla
NP/CNP [1; 2].
[Table] [Table]
variables #####################################################################
[Table] [Table]
if
[Table]
in_
[Table]
: self.
[Table]
_modules.
[Table]
(nn.
[Table]
(self.dim_xz,self.dim_logit_h))
foriinrange(self.num_logit_layers):self.logit_modules.append(nn
[Table]
ReLU()) self.logit_modules.append(nn.
Linear(self.dim_logit_h,self.dim_logit_
[Table]
self.logit_
[Table]
.append(nn.
ReLU())self.logit_
[Table] [Table]
[Table] [Table]
.dim
[Table] [Table]
[Table]
) else: self.
[Table] [Table]
[Table] [Table]
[Table]
foriinrange(
[Table]
num
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
def
[Table]
[Table] [Table]
if self
[Table]
_in_gates:logit_output=self.logit_net(x_z
[Table]
else:
[Table]
_z
[Table]
torch
[Table]
mean(
[Table]
_z
[Table]
dim=
[Table] [Table]
logit_output =
self.logit_net(
[Table] [Table]
if not
[Table]
logit
[Table]
(-1) if gumbel_max:logit_output = logit_output \ + sample_gumbel(logit_output.size())
softmax_y = F.softmax(logit_output/temperature, dim=-2) softmax_y = softmax_y.squeeze(-1) shape = softmax_y.size() _,ind = softmax_y.max(dim=-1) y_hard
= torch.zeros_like(softmax_y).view(-1, shape[-1]) y_hard.scatter_(1, ind.view(-1, 1), 1) y_hard
= y_hard.view(*shape) y_hard
= (y_hard-softmax_y).detach() \ + softmax_y softmax_y, y_hard = softmax_y.unsqueeze(-1), y_hard.unsqueeze(-1) return softmax_y, y_hard
Putting the mentioned structures together, we include the MoE-NPs as follows.
[Table]
for i in range(self.num_h): self.dec_modules.append(get_act(self.act_type))
self.dec_modules.append(nn.
Linear(self.dim_h,self.dim_h))
if self.amort_y:self.dec_modules.append(get_act(self.act_type)) self.dec_modules.append(nn.
Linear(self.dim_h,2*self.dim_y))
else:self.dec_modules.append(get_act(self.act_type))
self.dec_modules.append(nn.
Linear(self.dim_h,self.dim_y)) self.dec_net
=nn.Sequential(*self.dec_modules) def get_context_idx(self,M): # generate the indeces of the N context points from M points
N = random.randint(1,M) idx =
random.sample(range(0, M), N) idx =
torch.tensor(idx).cuda() return idx def idx_to_data(self,data,sample_dim,idx): # get subset of an array ind_data= torch.index_select(data, dim=sample_dim, index=idx) return ind_data def reparameterization(self,mu,logvar):std=torch.exp(0.5*logvar) eps=torch.randn_like(std) return mu+eps*std def encoder(self,x_c,y_c,x_t,y_t):ifself.training: memo_c,memo_t=torch.cat((x_c,y_c),dim=-1),torch.cat((x_t,y_t),dim=-1)emb_c_list_mu=torch.cat([expert_module(memo_c)[0].unsqueeze(0) forexpert_moduleinself.expert_modules]) emb_c_list_logvar=torch.cat([expert_module(memo_c)[1].unsqueeze(0) for expert_module in self.expert_modules]) emb_t_list_mu
=torch.cat([expert_module(memo_t)[0].unsqueeze(0) forexpert_moduleinself.expert_modules]) emb_t_list_logvar=
torch.cat([expert_module(memo_t)[1].unsqueeze(0)forexpert_module inself.expert_modules])
emb_c_list_mu,emb_c_list_logvar=
emb_c_list_mu.permute(1,0,2),\emb_c_list_logvar.permute(1,
[Table]
,2) emb_
[Table]
_list_mu,
[Table]
_t_list_
[Table]
=emb_t_list_mu.
[Table] [Table]
,2),\emb_
[Table]
_list_
[Table] [Table]
[Table] [Table]
,2) else:memo_c=torch.cat((x_c,y_c),dim=-1)
emb_c_list_mu=
torch.cat([expert_module(memo_c)[0].unsqueeze(0)forexpert_module inself.expert_modules])
) if self.training:ifself.info_bottleneck: z_experts=self.reparameterization(mu_t,logvar_t)else: z_experts=
mu_c else: assert y_pred==None if self.info_bottleneck: z_experts=self.reparameterization(mu_c,logvar_c) else: z_experts=
mu_c z_experts_unsq=z_experts.unsqueeze(1).expand(-1,x_pred.size()[1],-1,-1)
x_exp=
x_pred.unsqueeze(2).expand(-1,-1,z_experts_unsq.size()[2],-1)
if self.training:y_exp=y_pred.unsqueeze(2).expand(-1,-1,z_experts_unsq.size()[2],-1) ifself.experts_in_gates:
xz_exp=torch.cat((x_exp,z_experts_unsq),dim=-1)
xy_exp=torch.cat((x_exp,y_exp),dim=-1)xyz_exp=torch.cat((xy_exp,z_experts_unsq),dim=-1)
alpha_post,y_hard_post=self.logit_net_post(x_z=xyz_exp,temperature=self.temperature, gumbel_max=
self.gumbel_max)alpha_prior,y_hard_prior=
self.logit_net_prior(x_z=xz_exp, temperature=self.temperature,gumbel_max=self.gumbel_max) else:xy_exp=torch.cat((x_exp,y_exp),dim=-1)
alpha_post,y_hard_post=self.logit_net_post(x_z
=
xy_exp, temperature=self.temperature, gumbel_max=self.gumbel_max) alpha_prior,y_hard_prior=
self.logit_net_prior(x_z=x_exp, temperature=self.temperature,gumbel_max=self.gumbel_max) else:if self.experts_in_gates:
xz_exp=torch.cat((x_exp,z_experts_unsq),dim=-1)alpha_post,y_hard_post=0, alpha_prior,y_hard_prior=
self.logit_net_prior(x_z=xz_exp,temperature=self.temperature, gumbel_max=self.gumbel_max) else: alpha_post,y_hard_post=0,0 alpha_prior,y_hard_prior=
self.logit_net_prior(x_z=x_exp, temperature=self.temperature,gumbel_max=self.gumbel_max) output=self.dec_net(torch.cat((x_exp,z_experts_unsq),dim=-1))
if whether_acrobot:ifself.amort_y: y_mean,y_std=output[...
,:self.dim_y],\F.softplus(output[...,self.dim_y:]) returnmu_c,logvar_c,mu_t,logvar_t,\y_mean,
y_std,alpha_post,alpha_prior else:y_pred=torch.cat((torch.cos(output[...,0:1]), torch.sin(output[...,0:1]),torch.cos(output[...,1:2]), torch.sin(output[...,1:2]),4*pi*torch.tanh(output[...,2:3]), 9*pi*torch.tanh(output[...,3:4])),axis=-1)returnmu_c,logvar_c,mu_t,logvar_t,\ y_pred,alpha_post,alpha_prior elif whether_image: ifself.amort_y:y_mean,y_std=F.sigmoid(output[...,:self.dim_y]),\
F.softplus(output[...,self.dim_y:])returnmu_c,logvar_c,mu_t,logvar_t,\ y_mean,y_std,alpha_post,alpha_prior else: y_pred=F.sigmoid(output)returnmu_c,logvar_c,mu_t,logvar_t,\ y_pred,alpha_post,alpha_prior else: if self.amort_y:y_mean,y_std=output[...,:self.dim_y],\
F.softplus(output[...,self.dim_y:])returnmu_c,logvar_c,mu_t,logvar_t,\ y_mean,y_std,alpha_post,alpha_prior else: y
I Additional Experimental Results I.1 Additional Analysis of Learned Latent Variables
Here we give more detailed analysis w.r.t.
learned latent variables in MoE-NPs.
Entropy of Assignment Latent Variables.
Still we take the 1-dimensional toy stochastic function as example because it is intuitive to understand the latent variable meanings of different levels.
Remember that the role of the discrete latent variable e is to assign the diverse functional prior z 1: K to a given data point.
With the learned conditional prior p ϕ 2 , 2 ( e | z 1: K , x i ) for a data point x i , we can quantify the uncertainty of assignment via the entropy of such a Bernoulli latent variable H [ e ] .
[Equation]
This has a practical significance in discontinuous functions.
For example, in regions close to demarcation points, it should be difficult to judge the best expert z 1: K to handle these data points, which means the set of H [ e ] theoretically exhibits higher uncertainty.
Similarly, in regions without context points, it is hard to determine the function as well.
Interestingly, we observe that MoE-NPs are able to exhibit the above effect on the right side of Fig.
The sampled function consists of two components respectively in the interval [ − π, π ] and [ π, 3 π ] .
The entropy values of our interest are computed via Eq.
Here K = 2 and the learned conditional prior p ϕ 2 , 2 ( e | z 1: K , x i ) has highest entropy around the demarcation data point π and the
[Figure]
data point − 2 . [-
-]0 with no context points nearby.
This finding further verifies the role of the assignment latent variable in MoE-NPs.
Number of Expert Latent Variables.
By setting the approximate posterior of assignment latent variables as q ϕ 2 , 1 ( e | x, y, z 1: K ) =
p ϕ 2 , 2 ( e | x, z 1: K ) , we further investigate the scalability issue of
MoE-NPs in CIFAR10 image completion.
As reported in Table (3), we can find with more experts ( > = 3 ), the performance can be further improved and no overfitting issue occurs.
It can also be inferred when the number of experts reaches a certain level, the improvement from the increase of expert numbers is quite limited.
So in general, when the output dimension is lower, the best choice of the approximate posterior for assignment latent variables is in a form without y as the input.
[Table]
I.2 Additional Results of NP Variants in Toy Regression
Note that the variation of tasks in the previous toy regression is quite limited and its goal is to show the role of latent variables.
To further examine the performance, we construct the mixture of sinusoidal functions by varying the amplitude and the phase as follows.
The learning data points are randomly sampled in x -domain and merged from a mixture of randomized functions f 1 ( x ) =
A sin( x − B )+ ϵ in x -domain [ − π, π ] and f 2 ( x ) =
A cos( x − B )+ ϵ in x -domain [ π, 3 π ] with equal probability, where ϵ ∼ N (0 , 0 . 03 2 ) .
The range of the amplitude is A ∈ [0 .
1 , 5 . 0] while that of the phase is B ∈ [0 , π ] .
In each training iteration, we sample a batch of data points and randomly partition context points and target points for learning.
Each task consists of 100 randomly sampled data points from the mixture of sinusoidal functions with the random number of context points between [5 , 50] .
The default number of tasks in a meta batch is {+
+} and we set the number of iteration steps at most 50000.
As for neural architectures of all baselines, we retain that in the previous toy regression in Sec.
(H.2.1).
Still we use two experts for MoE-NPs as default to fit mixture sinusoidal functions.
[Table]
In meta testing phase, we draw up 500 tasks with 15 random data points selected as the context.
These testing tasks are generated in the way: the couple of the amplitude and the phase [ A, B ] are orderly set from the amplitude list numpy.linspace (0 . 1 , 5 . 0 , num = 500) and phase list numpy.linspace (0 . 0 , π, num = 500) .
The sampled x -values for these tasks are a list torch.linspace ( − π, 3 π, steps = 500) .
As shown in Table (4), ANP achieves best performance in meta testing, followed by MoE-NP.
I.3 Additional Results of NP Variants in System Identification
To understand how performance evolves with more context transitions in Acrobot system, we extend the result of 50 context points in the main paper to Fig.
As can be seen, MoE-NP still outperforms all NP baselines in all cases.
And with the increase of context transitions, we can find the predictive MSEs degrade accordingly.
[Figure]
I.4 Comparison with Attentive Neural Processes
Since neural architectures for attentive neural processes (ANPs) [14] are bit different from used baselines and cannot be trivially modified to meta RL cases, we report additional results in this subsection.
We implement ANPs with dot attention networks (since ANPs have more model complexity and can easily lead to cuda out of memory in practice, we choose the basic version of ANPs), the input embedding dimension of to compute attention weights is 32 and 4 layers are used to transform the deterministic embedding z attn .
The local embedding is concatenated with the input x and the global latent variable z for the decoder.
We use three heads for system identification tasks and one head for image completion tasks.
The related results are reported as follows.
It can be seen in Table (5),
MoE-NPs still outperform ANPs, while ANPs beat NPs a lot in predicting Acrobot system dynamics.
As for CIFAR10 image completion, we can draw the same conclusion in Table (6) that Mixture
Expert inductive biases are more effective than local latent variables embedded in attention modules.
[Table]
I.5 Augmenting MoE-NPs with Convolutional Modules
In this section, we examine the chance of Since neural architectures of encoders are quite different between ConvCNPs [17] and previous mentioned NP baselines, we only report the results of NP related models with the same functional encoder structures in the main paper.
Note that the translation equivariance is injected in ConvCNPs, which is a strong inductive bias for image dataset.
Naturally,
[Table]
we also develop the convolutional version of MoE-NPs.
And we report the separate results in image completion here.
In our settings, MoE-ConvCNPs use 3 experts in convolutional modules.
The implementation of ConvCNPs can be found in [17].
It can be seen that in Table.
(7), in comparison to ConvCNP, the image completion performance is further improved with help of multiple experts
[Table]
I.6 More Visualization Results
Here we include more visualized CelebA images by varying the number of observed pixels.
These are produced using CNN Augmented MoE-NPs (MoE-ConvCNPs).
Fig.s (14)/(15)/(16)/(17) are image completion results given the fixed number of random context pixels.
Fig.s (18)/(19)/(20)/(21) are image completion results given the fixed number of ordered context pixels.
[Figure]
J Computational Devices
Throughout the research process, we use NVIDIA 1080-Ti GPUs and Pytorch is used as the deep learning toolkit.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]