Learning Expressive Meta-Representations with Mixture of Expert Neural Processes
Abstract
Neural processes (NPs) formulate exchangeable stochastic processes and arepromising models for meta learning that do not require gradient updates duringthe testing phase.
However, most NP variants place a strong emphasis on a globallatent variable.
This weakens the approximation power and restricts the scopeof applications using NP variants, especially when data generative processes arecomplicated.
To resolve these issues, we propose to combine the M ixture o f E xpertmodels with N eural P rocesse s to develop more expressive exchangeable stochasticprocesses, referred to as Mixture of Expert Neural Processes (MoE-NPs).
Thenwe apply MoE-NPs to both few-shot supervised learning and meta reinforcementlearning tasks.
Empirical results demonstrate MoE-NPs’ strong generalizationcapability to unseen tasks in these benchmarks
Introduction
[Figure]
Humans can naturally accommodate themselves to new environments after developing related skills, and this kind of adaptabilityrelies on the good abstraction of environments.
Similarly, metalearning or learning to learn tries to leverage past experiences,and with help of the incorporated meta learned knowledge, anew skill can be mastered rapidly with a few instances.
During the past decade, an increasing number of methods haveemerged in meta learning domains.
In this paper, we concentrateon a special branch of meta learning methods, referred to ascontext-based meta learning [1; 2].
A representative one is aneural process (NP) [1], which was initially proposed to approximate Gaussian processes with lower computational cost.
Thecore purpose of NPs is to learn meta-representations [3], whichencode context points into latent variables and represent thetask in a functional form.
In comparison to gradient-based metalearning algorithms, e.g. model-agnostic meta learning (MAML)[4], the NP directly learns a functional representation and doesnot require additional gradient updates in fast adaptation.
Research Motivations.
Fundamentally, vanilla NPs employglobal
Gaussian latent variables to specify different tasks.
This setting raises several concerns insome 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) Incontext-based meta reinforcement learning, the uncertainty of value functions revealed from latentvariables 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 Gaussianlatent variable restricts the flexibility of randomized value functions [8], leading to sub-optimality inperformance.
Developed Methods.
Instead of using a global latent variable in modeling, we employ multiplelatent 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 betweena single data point and expert NPs.
To optimize this model with hybrid types of latent variables, weutilize variational inference to formulate the evidence lower bound.
Additionally, special modules aredesigned to accommodate few-shot supervised learning and meta reinforcement learning tasks.
Outline & Contributions.
We overview the properties of NPs and general notations for context-basedmeta learning tasks in Section 3.
Section 2 summarizes related work in meta learning and the NPfamily.
In Section 4, M ixture o f E xpert N eural P rocesse s (MoE-NPs) are elaborated, together withrequired modules for meta learning tasks.
Experimental results and analysis are reported in Section5.
Discussions and limitations are included in Conclusion Section.
In principle, our contribution istwo-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 andNPs, enabling multi-modal meta representations for functions.
• We specify inference modules in MoE-NPs for few shot supervised learning and metareinforcement 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 metalearning 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 fewinstances.
Metrics-based meta learning algorithms try to learn representations of tasks in distancespace, and models e.g. prototypical networks [12; 13] are popular in computer vision domains.
Asfor context-based meta learning methods of our interest, latent variable models, e.g. NPs [1], aredesigned to learn task representations in a functional space.
A major advantage of this family is thatno gradient updates are required in fast adaptation and computational time is reduced accordingly.
Neural Processes Family.
Apart from vanilla NPs or CNPs [2; 1], other variants are developed andthese are built on various inductive biases.
To address underfitting issues, attention networks [14]are introduced to augment NPs.
[15; 16] improve the generalization capability of NPs with help ofconvolutional operations.
To learn distributions of group equivariant functions, [17] has proposed
EquiCNP.
Similarly, SteerCNPs also incorporate equivariance to approximate stochastic fields [18].
Our work is to get NPs married with Mixture of Experts (MoEs) models [19; 20], which modeldatasets with a collection of expert NPs.
We provide more technical summary of the NP familytogether with other probabilistic meta learning methods [21; 22; 23; 24] in the Appendix (F).
Preliminaries
Notations.
The paradigm of meta learning is considered in the distribution of tasks p ( T ) , and a tasksampled from p ( T ) is denoted by τ in this paper.
The form of a task depends on applications.
Forexample, a task of our interest in regressions can be a function f to fit, which is a realization fromunknown stochastic processes [25].
Let D Cτ refer to a set of context points used to specify the underlying task, and D Tτ =
[ x T , y T ] isa set of target points to predict, e.g. f ( x T ) =
y T .
In the context-based meta learning with latentvariables, 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 [26].
Agenerative 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 theconditional 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 withinthe expected risk minimization principle as follows.
[Equation]
The risk function L , e.g. negative log-likelihoods, measures performance of meta learning structureon the task-specific dataset D Cτ and D Tτ , and Θ means parameters of common knowledge sharedacross tasks and parameters for fast adaptation ( e.g , Θ denotes parameters of the encoder ϕ anddecoder θ , 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 posteriorof a global latent variable z in Eq.
(1) is approximated with a variational distribution q ϕ ( z |D Tτ ) andan evidence lower bound (ELBO) is derived to optimize in practice.
A general meta training objectiveof NPs in few-shot supervised learning is L ( θ, ϕ ) in Eq.
(4), where p θ ( D Tτ | z ) =
(cid:81) ni
=1 p θ ( y
[Equation] [Equation]
[Equation]
3.2 Meta Reinforcement Learning
For the context-based meta reinforcement learning, the context points D Cτ are a set of randomtransition samples from an environment as D Cτ
= { ( s 1 , a 1 , s 2 , r 1 ) , . .
, ( s H , a H , s H +1 , r H ) } , wherer t is the one-step reward after performing action a t at state s t .
Here D Cτ plays a role in task inference[27] 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 oftask-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 expectedcumulative rewards in the environment τ given policies π φ ( a | [ s, z ]) .
[Equation] [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 asa fixed one, e.g. N (0 , I ) .
The induced distribution of task-specific value functions p ( Q θ ([ s, z ] , a ))enables posterior sampling [28] in meta learning scenarios, which brings additional benefits ofexploration for continuous control problems.
Model & Algorithm
In this section, we present our developed MoE-NPs and connect them to the hierarchical Bayesframework.
Then approximate objectives are derived and stochastic gradient variational Bayes [29]is used to optimize the developed model.
Finally, specialized neural modules are described for
MoE-NPs application to different meta learning tasks.
Detailed computational diagrams in trainingand testing are attached 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 expressivenessbottlenecks when employing a global latent variable in learning functional priors of tasks fromunknown distributions [30].
To alleviate mentioned deficiencies, we make two modifications for the NP family.
(i)
Multiplefunctional priors are encoded in modeling with help of K expert latent variables, which can capturestatistical traits, e.g. distributional multi-modality, in data points.
This setting is also an extension of
Mixture of Experts (MoEs) models [19; 31; 32; 33] to meta learning scenarios.
(ii) Like the gatingmechanism in [5], assignment latent variables are included in modeling to select functional forms foreach data point in prediction.
The resulting MoE-NPs can learn more expressive functional priorsand exhibit the approximation power for local properties of the dataset.
Generative Process.
As displayed in Fig.
(1), the graphical model involves two types of latentvariables, respectively the continuous expert latent variables z 1: K and the discrete assignment latentvariable 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, ande k = 1 in Eq.
(6) indicates the k -th expert z k is selected from z 1: K for prediction.
A more detailedprobabilistic generative process can also be found in Appendix (E.1).
In this way, our developedmodel constitutes an exchangeable stochastic process .
And we demonstrate this claim with help of
Kolmogorov Extension Theorem [34] in Appendix (E.2).
Link to Hierarchical Bayes.
Note that latent variables in Eq.
(6) are of hybrid types.
K functionalpriors are incorporated in expert latent variables z 1: K , while the assignment latent variable e is inputdependent.
The dependencies between z 1: K and e are reflected in modeling, and this connects ourwork to Hierarchical Bayes [35; 36] in a latent variable sense.
Also when only one expert latentvariable is used here, the hierarchical model degenerates to the vanilla NPs [1].
4.2 Scalable Training & Prediction
Inference Process.
Given a task τ , due to existence of unknown latent variables, it is intractable toperform exact inference w.r.t.
p ( D Tτ |D Cτ ) .
As an alternative, we apply variational inference to ourdeveloped model.
Here we use [ x, y ] to denote a single data point from a set of target points D Tτ .
[Equation] [Equation]
An example for 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 drawnfrom a categorical distribution Cat ( e ; K, α ( x, y, z 1: K )) with parameters
α ( x, y, z 1: K ) .
The existenceof 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 whenutilized 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 asa 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, whereq ϕ 1 denotes a collection of K independent variational distribution { q ϕ 1 , 1 , . . .
, q ϕ 1 ,K } .
q ϕ 2 , 1 and p ϕ 2 , 2 161 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 optimizein 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 Carloestimate 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 passand parameters of discrete latent variables p ϕ
[Equation]
x
[Equation]
K [Equation] =
[Equation]
x
[Equation] [Equation]
with
[Equation]
[Equation]
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
ReplayBuffer
[Figure]
The equations so far define a framework that can be implemented in different meta learning tasks.
Two examples are givenin Fig.
Note that two types of latent variables are involvedin modelling, we need different structures of encoders for latentvariables.
Inference Modules are required for them, satisfyingdifferent conditions.
In variational inference, distribution parameters of these latent variables are approximated with the outputof these specialized encoders.
Inference Modules for Continuous Latent Variables.
Forneural networks to parameterize the encoder of continuous latentvariables q ϕ 1 , we use the same architecture 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 ) output distribution parameters for each expert z k shouldsatisfy [ µ k , Σ k ] =
NN ϕ 1 ,k
([ x σ (1: N ) , y σ (1: N ) ]) .r
k,i =
h ϕ 1 ,k
([ x i , y i ]) , r k =
N (cid:77) i =
r k,i , [ µ k , Σ k ] =
g ϕ 1 ,k ( r k )
[Equation]
Eq.
(10) is an example, where h is the embedding function, 190 (cid:76) denotes a mean pooling operation, and g is part of encodernetworks.
Inference Modules for Categorical Latent Variables.
For neural networks to parameterize theencoder of discrete latent variables q ϕ 2 , 1 and p ϕ 2 , 2 , we need the categorical distribution parameters
α 194 to be permutation equivariant [37] with respect to the order of z 1: K .
This means for any order permutation operation σ , the condition is satisfied as [ α σ (1) , α σ (2) , . . .
, α σ ( K ) ] =
[Equation]
ϕ 2 , 1 ( x, y, z σ
[Equation]
[Equation]
[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 modifiedfor Actor-Critic methods, which are employed in our settings.
Like that in PEARL [6] and FCRL [3],the soft actor critic (SAC) algorithm [38] is used to learn policies due to good sample efficiency.
Given a specific MDP τ , posterior distributions of optimal value functions are formulated vialatent variables z in context-based meta RL.
That is, p ( Q θ ( s, a ; M )) is approximated in the formp ( 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]
[Equation]
The variational posterior q ϕ ( z | s, D Cτ ) in Eq.
(13) is a state dependent distribution with ˆ V a statevalue 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 latentvariables 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 tooptimize these functions is listed in Appendix (A).
Experiments and Analysis5.1 General Settings
The implementation of MoE-NPs in meta training can be found in Appendix Algorithms (1)/(3), andalso please refer to Appendix Algorithms (2)/(4) for the corresponding meta-testing processes.
Weleave 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 [39] are included in the objective.
In meta RL, the modified NP model corresponds to PEARL [6].
Meanwhile, in Appendix (I), weinclude additional experimental results compared with other NPs models augmented by attentivemodules [14] or convolutional modules [16].
5.2 Illustration in Toy Regression
To see different roles of latent variables, we visually show effects of stochastic function fitting andquantified uncertainty in toy dataset.
Our goal is to discover potential components of distributionsfrom limited observed data points.
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 mixturecomponents, where ϵ 1 ∼ N (0 , 0 .
03 2 ) and ϵ 2 ∼ N (0 , 0 .
In each training iteration, we sample abatch of data points and randomly partition context points and target points for learning.
In testing
[Figure]
phase, we draw up 100 data points from this mixture of distributions with 15 random data pointsselected 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 inthe mixture curve of the second component.
NPs show intermediate performance and still fails tomatch context points well.
As for MoE-NPs, with help of predicted assignment variables parameterse ∗ =
one_hot [arg k ∈{ 1 , 2 } max α k ] , we set the number of experts as two and partition data points tovisualize predictive distributions p θ ( y ∗ | x ∗ , z ∗ ) .
The MoE-NP is able to precisely separate mixturecomponents 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 [40] and imagecompletion task in CIFAR10 [41].
Both tasks are common benchmarks in the meta learning or NPsliterature [40; 42; 1; 2; 14].
[Table]
248 System Identification.
For Acrobot systems, different tasks are generated by varying masses oftwo pendulums.
A dataset of state transitions is collected by using a complete random policy tointeract 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 andwe report the average predictive errors and standard deviations for all transitions.
Here we use 50transitions 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 allother models.
With three experts in modeling, MoE-NPs significantly outperform other baselines interms of dynamics prediction.
Our finding is consistent with observations in [43], where multi-modaldistributions 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 [41] in this experiment, which is formulated with32x32 RGB images.
In the meta training process, a random number of pixels are masked to completein images.
That is, given the context pixel locations and values [ x C , y C ] , we need to learn a mapfrom each 2-D pixel location x ∈ [0 , 1] 2 to pixel values y ∈ R 3 .
Here two expert latent variables areused in MoE-NPs.
In Fig.
(4), we evaluate image completion performance on the test dataset and the number ofcontext pixels is varied in three levels.
It can be found that CAVIA works best in cases with 10random context pixels or less than 500 ordered context pixels.
In other cases, MoE-NP surpasses allbaselines.
With more observed pixels, the predictive errors of MoE-NPs can be decreased in bothrandom and ordered context cases.
An example of image completion results is displayed in Fig.
[Figure]
[Figure]
For gradient-based methods, CAVIA and MAML are sensitiveto the number of context points and do not exhibit asymptoticperformance like that in MoE-NP.
NP still suffers underfittingin 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 tonavigate with sparse rewards.
The navigation goal of a task issampled from a mixture of arcs in a semi-circle in Fig.
(6.a)during meta training processes.
[Figure]
From Fig.
(7.a-c), we can observe the evaluation performanceof agents over iterations.
For gradient-based methods, CAVIAshows better performance than MAML in exploration but bothare weaker than context-based baselines.
MoE-NPs can convergeearlier with less training samples and show slight advantage overvanilla PEARL.
In particular, we test the asymptotic performancein out of distributions (O.O.D.) tasks and show results in Fig.(7.d).
We notice O.O.D. tasks are challenging for all algorithmsto generalize but average returns are gradually increased withmore trials.
PEARL and FCRL achieve comparable rewards,while MoE-NP behaves better in this case.
[Figure]
Locomotion.
As for Half Cheetah-Complex-Direction (H-Cheetah-CD) and Slim HumanoidComplex-Goal (S-Humanoid-CG) 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, MoENP 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 offunctional representations for task-specific value functions.
Either contrastive or multiple functionalpriors lead to better exploration and have a potential to boost performance in continuous control.
Forgradient-based methods
, observations show that they can easily get stuck in the local optimal [6; 44].
[Figure]
Ablation
[Figure]
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, 7and 9 in different MoE-NPs.
Here we test the predictive performance of meta-trained MoE-NP by varying the number oftransitions.
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 describedin Section (5.3).
It can be seen when the expert number is5, 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 exhibithigher predictive errors, and show no significant performanceimprovement with the increase of the number of context points.
These suggest increasing the number of experts beyond a certainpoint 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-specificoptimal value functions.
Here we take the S-Humanoid-CG asthe 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 randomnessof value function distributions captures task uncertainty andencourages more efficient exploration.
Conclusions
Technical Discussions.
In this work, we have developed a newvariant of NP models by introducing multiple expert latent variables.
Our work illustrates the roles of different latent variablesin MoE-NPs for meta learning tasks.
MoE-NPs are able to separate data points from different clustersof 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 graduallyincorporating new expert latent variables has not been explored and this raises concerns in additionalcomputational 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 incorporatedin modeling.
Another way is to place priors over distributions of discrete latent variables like that inhierarchical Dirichlet processes [45] and select the optimal number of experts in a Bayesian way.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow 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 referencingthe 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 areproprietary.•
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 instructionsblock 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]
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]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental 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 theywere 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., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
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]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[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
12 Literature Review 23 Preliminaries 23.
Few-Shot Supervised Learning . .
Meta Reinforcement Learning . .
34 Model & Algorithm 34.
Mixture of Expert Neural Processes . . . . . . . . . . . . . . . . . . . . . . . . .
Scalable Training & Prediction . .
Module Details for Meta Learning . . . . . . . . . . . . . . . . . . . . . . . . . .
Inference Modules in MoE-NPs . . . . . . . . . . . . . . . . . . . . . . .
Meta RL Modules in MoE-NPs. . . . . . . . . . . . . . . . . . . . . . . .
65 Experiments and Analysis 65.
General Settings . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Illustration in Toy Regression . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Few-Shot Supervised Learning . .
Meta Reinforcement Learning . .
Ablation Studies . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Conclusions[-
-]
Code of Algorithms in Meta Learning[-
-]
B Frequently Asked Questions
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 Stochastic Gradient Estimates . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.4 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.-]
I.
Additional Results of NP Variants in System Identification . . . . . . . . . . . . .
I.
Comparison with Attentive Neural Processes . . . . . . . . . . . . . . . . . . . .
I.
I.
J Computational Devices
Learning rates
Output trained parameters ϕ =
ϕ 2 ] and θ .
1 Initialize parameters ϕ and θ ; 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.
(8); // amortized inference process ϕ ← ϕ − λ 1 ∇ ϕ L MC ( θ, ϕ ) in Eq.
(8); θ ← θ − λ 2 ∇ θ L MC ( θ, ϕ ) in Eq.
(8); 12 end trained distributions.
1 Initialize parameters ϕ and θ ; 2 Set the number of context points N C ; 3 Split test dataset into the context/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
: 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
11 else
Compute parameters
of assignment Output 16 end Input : distribution
2 while Completed 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
B Frequently Asked Questions
Here we collect some frequently asked questions from reviewers and other literature researchers.
Wethank these reviewers for these questions and add more explanation.
Selection of Benchmarks in This Paper.
Admittedly, NP variants can be applied a series ofdownstream 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 [42; 46], in whichlearning transferable physics dynamics with NPs is a crucial application.
The image completion taskis more commonly used in work [9; 2; 1].
The meta reinforcement learning task can also be studiedwithin NP framework [3; 47].
NP Family in Classification Tasks.
I have tried to search neural architectures for few-shot imageclassification 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 SOTAfew-shot image classification results with standard neural architectures in NPs, e.g. multi-layerperceptrons (MLPs).
Maybe this is due to the nature of stochastic processes, which can addressregression problems more efficiently.
Unless specialize modules instead of MLPs are used, we do notexpect NP variants with MLPs can achieve SOTA performance.
The aim of this paper focuses moreon mixture expert inductive biases and place less attention on neural architecture search.
So MLPsare shared across all baselines to enable fair comparison.
Expressiveness of Mixture of Expert Inductive Biases.
A natural question about learning diversefunctional representation is whether these multiple expert latent variables will collapse into one.
Werefer the collapsed representation to vanilla (C)NPs, and the previous empirical results show thecollapsed ones work poorer than mixture of expert ones in both few-shot regression tasks and metareinforcement learning tasks.
Also, from the multimodal simulation result, we discover both latentassignment variables and expert variables are meaningful, which reflects the effectiveness of mixtureexpert inductive biases.
That means, we have not encountered meta representation collapse issues inexperiments.
C Probabilistic Graphs in Meta Training/Testing
As exhibited in Fig.
(10)/(11), shown are computational diagrams when implementing MoE-NPs inmeta learning tasks.
Variational Posteriors.
Since the real posteriors for both { z k } Kk =1 and e T are computationallyintractable, 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 thesake of simplicity, we denote { q ϕ 1 ,k } Kk =1 by q ϕ 1 in some time.
Variational Priors.
In some cases, the prior distributions for { z k } Kk =1 and e T are set to constrain thescope of prior distributions.
For example, in few-shot supervised learning, since the context and thetarget have the same form, the variational prior is selected to be q ϕ 1 as well to ensure the consistencyand this works in meta testing phases.
For the assignment latent variable e , this uses the same form inconditional VAE [48] 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 learningand meta reinforcement learning.
This is the same with that in FCRL [3].
Meta learning in NP relatedmodels is to learn functional representations of different tasks and formulate the fast adaptation viainferring the task specific conditional distribution p ( D Tτ |D
Cτ ) =
(cid:82) p ( D Tτ | z ) p ( z |D Cτ ) dz (equivalentto Eq.
To make the downstream reinforcement learning task using NP family models clearer, we add thefollowing 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 setof 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 NPfamily models.
Given a transition sample [ s, a, r ( s, a ) , s ′ ] , the target input is the state x T
= s
and thetarget 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 approximateposterior is learned from D Cτ with permutation invariant functions.
In total, sampling from the statedependent 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 sE.1 Generative Processes
To better understand our developed model in meta learning scenarios, we translate Eq.
(6) into astep-wise generative process.
The same with that in the main paper, the task distribution is denotedby p ( T ) and we presume K -experts to summarize the stochastic traits of a task. [Equation] (16c
[Equation]
Here a MoE-NP for the task τ is specified with K -expert latent variables z 1: K in Eq.
(16.a).
Theprobability mass function for a data point related categorical distribution Cat ( K, α ( x, z 1: K )) isdenoted by p ( e | x, z 1: K ) =
(cid:81)
[Equation] ( x, z 1: K [Equation] 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 theoutput 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 discreteassignment 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 processis 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 andmarginalization consistency.
Remember that the generative model is induced in the main paper as follows.
And we claim that ourdesigned 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 [34] and de Finetti’s Theorem [49].
This is to show the existence of SP s in Eq.
Exchangeability Consistency.
For N data points from Eq.
(17), we impose any permutationoperation σ over their indices, and this results in σ : [1 , 2 , . . .
, N ] → [ σ 1 , σ 2 , . .
, σ N ] .
Then wecan check the following equation is satisfied since the element-wise product of probabilities can beswapped.
[Equation]
Marginalization Consistency.
Given the assumption that the integral in Eq.
(17) is finite, we pickup a subset of indices [ M + 1 , M + 2 , . .
. , N ] and make M < N without difference in orders.
Andthe result after marginalization over y -variable in the selected indices can be verified based on thefollowing equation.
[Equation]
Built on these two sufficient conditions, our developed MoE-NP is a well defined exchangeable SP .
F Summary of Existing NP Related ModelsF.1 Comparison in Technical Details
Here we give a brief summary on difference between MoE-NPs and existing typical Neural Processmodels in Table (2).
Some crucial traits include forms of encoders and decoders structures, types oflatent variable and inductive biases injected in modelling.
Especially, the inductive bias for MoE-NPis reduced to be multiple functional priors, which means a collection of expert neural processesto induce the generated dataset.
Since a general inductive bias behind NPs related models is themodelling of exchangeable stochastic processes with cheap computations.
This corresponds to adistribution 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 beavailable, 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-NPsare with O ( K ∗ ( N + M )) (making M predictions with N observations).
In practice, the number ofexperts 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.
[Table] CNP [2] z = f ϕ ( x C , y C ) p θ ( y | [ x, z ]) continuous [Table] functional
NP [1] q ϕ ( z | [ x C , y C ]) p θ ( y | [ x, z ]) continuous global functional
ANP [14] q
[Table]
| [
[Table]
[Table]
|
[Table]
[Table]
| [
[Table]
[Table]
[Table]
ϕ (
[Table]
∗ | [
[Table]
[Table]
[Table]
θ ( y | [
[Table]
z ])
[Table]
[Table]
[Table]
∗ | [
[Table]
y C ]
[Table]
∗ )
[Table]
| [
[Table]
z ])
[Table]
CNP
[Table]
[Table]
| [
[Table]
[Table]
[Table]
F.3 Additional Literature Review
Due to page limit in the main paper, we include other related works in this subsection.
In unsupervisedlearning, the Neural Statistician Model [50] is introduced to compute summary statistics inside thedataset.
The Generative Query Network [47], a variant of NPs for visual sensory dataset, makes useof a latent variable to abstract scenes in high dimensions.
To capture heteroscedastic noise inside thestochastic process, DSVNP [51] induces latent variables at different levels.
The functional neuralprocesses infer the directed acyclic graph in the latent space and formulate flexible exchangeablestochastic processes for single task problems [52].
Inspired by self-supervised learning, [3; 53]propose to augment the neural process with contrastive losses.
[54] combines context memories andrecurrent 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 approximatedistributions, 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 involvementof discrete latent variables.
We therefore discuss chance of using another traditional optimizationalgorithm, called Expectation Maximization (EM) [55], in our settings.
Omitting the K -expert latentvariables z 1: K and corresponding variational distributions, we take a closer look at the assignmentlatent 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 probabilityfunction p ( e | x, z 1: K ) =
Cat ( e ; K, α ( x, z 1: K )) .
This step is to update the posterior of the proportionalcoefficients α ( 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 timeupdated one p ( t ) ( e ) .
As a result, updated categorical distribution parameters are:
[Equation]
t ) ( y | [Equation] e [Equation]
[Equation]
[Equation]
) (cid:80)
[Equation]
[Equation]
e
[Equation]
[Equation]
[Equation]
 
[Equation]
21) where τ is the temperature parameter.
Maximization (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 θ (cid:80) ( x,y ) ∈D ln
[ p θ ( t ) ( y | x, z 1: K , e )] given the last time updated model parameter θ ( t ) .
Withhelp 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 considersa subset of data points D .
However, in meta learning scenarios, we handle large-scale dataset andthe above-mentioned EM framework is computationally expensive and impractical.
Due to theseconsiderations, we do not apply EM algorithms to estimate the discrete distribution and insteadvariational 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 selectedas Cat ( K, α 0 ) with distribution parameters α
=
[ α 0 , 1 , α 0 , 2 , . .
, α 0 ,K ] .
And the Kullback–Leibler
Divergence is computed as follows.
[Equation]
[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 fork -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τ ) andq ϕ 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 chooseq ϕ 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]
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, theapproximate posterior in Eq.
(26) in meta training should be substituted with q ϕ 1 ( z 1: K |D Tτ ) with thecorresponding approximate prior p ( z k ) =
q ϕ 1 ,k
( z k |D Cτ ) .
And this matches the general form in themain paper for −L ( θ, ϕ 1 , ϕ 2 ) in Eq.
When Dirac delta distributions are used in MoE-NPs, the divergence term about the continuouslatent variable is removed as default.
Denoting the approximate posterior by q ϕ 2 ( e | x
, 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]
G.3 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 [29] 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 variable follows that in [56].
G.4 Estimates of Statistics
Momentum.
Given the pre-trained MoE-NPs, we can formulate the statistical momentum inpredictive 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 thepredictive distribution of one expert is parameterized in the form p ( y | x, z k ) =
N ( y ; m ( x, z k ) , Σ k ) ,where m is the learned mean function using a neural network and σ 2 is a variance parameter.
Usinga single stochastic forward pass in expert latent variable z 1: K , we can derive the estimate of thepredictive mean ˆ m =
E
[ Y | X
= x, D Cτ ] .
[Equation]
The second order moment can be estimated accordingly.
Here we consider the case when the outputis 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.
Inthis case, the entropy of predictive distribution plays a crucial role.
Though the exact estimate of thepredictive 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 asingle 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 ofentropy term is as follows.
[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 context-based meta RL algorithms, you can find the implementation of MoE-NPs from the anonymous Github linkhttps://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 1 , 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 andm 2 , which means the hyper-parameters of the system come from the Cartesian combination of theset m 1 ∈ { 0 . 75 , 0 . 85 , 0 .
95 , 1 . 15 } and m 2 ∈ { 0 . 75 , 0 . 85 , 0 .
In meta training processes, acomplete random policy interacts with batch of sampled MDPs to formulate transition dataset.
As formeta testing tasks, we follow the same way to generate tasks by setting m 1 ∈ { 0 . 85 , 1 . 05 , 1 . 25 } andm 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 imagecompletion tasks.
CIFAR10 images are processed via torchvision modules to normalize the pixelvalues 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 uniformdistributions [0 , 2 π/ 12] ∪ [5 π/ 12 , 7 π/ 12] ∪ [10 π/ 12 , π ] .
The rest of regions along the arc is for outof distribution tasks.
The tasks in Mujoco [7] follows adaptations from [6; 57], where goals/velocitiesor multiple hyper parameters of simulation systems are sampled from mixture distributions.
Thehorizon 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 forpoint robots, 7 . 5 ∗ 1 e 6 for Half-Cheetah-CD and 6 .
5 ∗ 1 e 6 for Slim-Humanoid-CG.
We leave moredetails 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 taskincludes 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 includesmultiple goals and these change with split steps of episodes.
H.2 Implementations in Meta Learning TasksH.2.1 Toy Experiments
The input of functions is in a range [ − π, π ] ∪ [ π, 2 π ] .
The general implementations of MoE-NPs areas 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
H.2.2 Few-Shot Supervised Learning
For gradient based methods, we use implementations of MAML 2 and CAVIA 3 .
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 variablesis with two hidden layers (32 neuron units each).
Note that Dirac delta distributions are used forcontinuous 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 discreteassignment latent variable in MoE-NPs corresponds to a softmax-output neural network with twohidden layers (32 neuron units each).
The Decoder is with four hidden layers (200 neuron unitseach) 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 ofmeta training, 4 different environments are randomly selected, and the uniform random controller isused to interact for the collection of 5 episodes (for each task).
In training dynamics systems, thetraining batch size in transition buffer is 200, the training epoch is 5, and the whole process iteratesuntil convergence (the iteration number is 25).
The learning rate for Adam optimizer is 1 e − 3 as thedefault.
Image Completion.
The general implementations follow that in [9; 2] and are applied to all baselinesand MoE-NPs.
The dimension of a latent variable is set to be 128.
The Encoder for all NP variantsis with three hidden layers (128 neuron units each).
Note that Dirac delta distributions are used forcontinuous 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 threeexpert latent variables as the default, and the Encoder for the discrete assignment latent variablecorresponds to a softmax-output neural network with two hidden layers (32 neuron units each).
Adam[29] is used as the optimizer, where the learning rate is set to be 5 e − 4 .
The batch size in training is8 images, and we meta train the model until convergence (the maximum epoch number is 50).
Alsonote 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 objectiveconsistent based on that in [58].
But in evaluation, to keep results of all methods consistent, we followthat 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 4 ,
MAML 5 and CAVIA 6 .
Note that the TRPO algorithm [59] 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 PEARLalgorithms [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
[38] is used in policy optimization, which requires parameterization of both actorand critic functions.
As for the number of experts in MoE-NPs, we use 3 for all environments asdefault.
You can find more details about neural architectures/optimizers for each modules from theabove mentioned link.
As mentioned before, we use p ( z k ) =
N (0 , I ) as the prior distribution for expert latent variables.
Theapproximate posterior is parameterized with a diagonal Gaussian distribution.
The coefficient for KLdivergence terms in Eq.
(14) are β 0 = 1 . 0 , β 1 = 1 . 0 .
The meta-training processes for reinforcementlearning can be found in Algorithm (3) in the main paper.
In terms of meta-testing processes, wereport the pseudo code in Algorithm (4).
H.3 Neural Architectures
To help readers better understand our models, we attach the python code of the Encoder for separatelatent variables as follows.
As for the Decoder , the structure is the same with that in a vanilla
NP/CNP [1; 2].
[Table] [Table]
[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,0,2)
emb_t_list_mu,emb_t_list_logvar=emb_t_list_mu.permute(1,0,2),\emb_
[Table]
_list_
[Table]
.permute(1,0,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]),\
[-H [ p ϕ 2 , 2 ( e | z 1: K , x i )] data point − 2 . 
-]
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 return
,\ y post prior
I Additional Experimental Results
[Figure]
I.1 Additional Results of NP Variants in System Identification
To understand how performance evolves with more context transitions in Acrobot system, we extendthe result of 50 context points in the main paper to Fig.
As can be seen, MoE-NP stilloutperforms all NP baselines in all cases.
And with the increase of context transitions, we can findthe predictive MSEs degrade accordingly.
[Figure]
I.2 Comparison with Attentive Neural Processes
Since neural architectures for attentive neural processes (ANPs) [14] are bit different from usedbaselines and cannot be trivially modified to meta RL cases, we report additional results in thissubsection.
We implement ANPs with dot attention networks (since ANPs have more model complexity andcan easily lead to cuda out of memory in practice, we choose the basic version of ANPs), the inputembedding dimension of to compute attention weights is 32 and 4 layers are used to transform thedeterministic embedding z attn .
The local embedding is concatenated with the input x and the globallatent variable z for the decoder.
We use three heads for system identification tasks and one headfor image completion tasks.
The related results are reported as follows.
It can be seen in Table (3),
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 (4) and interestingly
ANPs fail to exhibit the phenomenon that the predictive error decrease with the increase of contextpoints.
[Table]
[Table]
I.3 Augmenting MoE-NPs with Convolutional Modules
In this section [Table] we examine the chance of Since neural architectures of encoders are quite differentbetween ConvCNPs [16] and previous mentioned NP baselines, we only report the results of NPrelated models with the same functional encoder structures in the main paper.
Note that the translationequivariance is injected in ConvCNPs, which is a strong inductive bias for image dataset.
Naturally,we also develop the convolutional version of MoE-NPs.
And we report the separate results in imagecompletion here.
In our settings, MoE-ConvCNPs use 3 experts in convolutional modules.
Theimplementation of ConvCNPs can be found in [16].
It can be seen that in Table.
(5), in comparisonto ConvCNP, the image completion performance is further improved with help of multiple experts.
[Table]
I.4 More Visualization Results
Here we include more visualized CelebA images by varying the number of observed pixels.
Theseare produced using CNN Augmented MoE-NPs (MoE-ConvCNPs).
Fig.s (13)/(14)/(15)/(16) areimage completion results given the fixed number of random context pixels.
Fig.s (17)/(18)/(19)/(20)are image completion results given the fixed number of ordered context pixels. [Figure] [Figure]
[Figure]
J Computational Devices
Throughout the research process, we use NVIDIA 1080-Ti GPUs and Pytorch is used as the deeplearning toolkit.
[Figure]
[Figure]
[Figure]
[Figure] [Figure]