M ODEL -A GNOSTIC M ETA -L EARNING FOR M ULTIMODAL T ASK D ISTRIBUTIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
Gradient-based meta-learners such as MAML (Finn et al., 2017) are able to learn a meta-prior from similar tasks to adapt to novel tasks from the same distribution with few gradient updates.
One important limitation of such frameworks is that they seek a common initialization shared across the entire task distribution, substantially limiting the diversity of the task distributions that they are able to learn from.
In this paper, we augment MAML with the capability to identify tasks sampled from a multimodal task distribution and adapt quickly through gradient updates.
Speciﬁcally, we propose a multimodal MAML algorithm that is able to modulate its meta-learned prior according to the identiﬁed task, allowing faster adaptation.
We evaluate the proposed model on a diverse set of problems including regression, few-shot image classiﬁcation, and reinforcement learning.
The results demonstrate the effectiveness of our model in modulating the meta-learned prior in response to the characteristics of tasks sampled from a multimodal distribution.
I NTRODUCTION
Humans are capable of effectively utilizing prior knowledge to acquire new skills as well as adapt to new environments.
For example, imagine that we are learning how to carve on a snowboard on a mountain that we just arrived at.
While we only know the basics of snowboarding and barely know the terrain, we can still accomplish this feat quickly.
Even when the skill that we are interested in learning is only vaguely related to the set of skills that we have acquired in the past, we are usually able to exploit the relationship among skills and generalize.
For the snowboarding example, our knowledge about skiing and skateboarding can play the role that prepares us for learning snowboarding skills more efﬁciently.
Similarly, can machines rapidly master a novel skill based on a variety of related skills they have already acquired?
Recent advances in meta-learning offer machines a way to rapidly adapt to a new task using few samples by ﬁrst learning an internal representation that matches similar tasks.
Such representations can be learned by considering a distribution over similar tasks as the training data distribution.
Model-based meta-learning approaches (Duan et al., 2016; Wang et al., 2016;
Munkhdalai & Yu, 2017; Mishra et al., 2017) propose to recognize the task identity from a few sample data, use the task identity to adjust a model’s state ( e .
g .
RNN’s internal states) and make the appropriate predictions with the adjusted model.
Those methods demonstrate good performance at the expense of having to hand-design architectures, yet the optimal strategy of designing a meta-learner for arbitrary tasks may not be obvious to humans.
On the other hand, model-agnostic gradient-based meta-learning frameworks (Finn et al., 2017; 2018; Kim et al., 2018; Lee & Choi, 2018; Grant et al., 2018) seek an initialization of model parameters such that a small number of gradient updates will lead to high performance on a new task, while having ﬂexibility in the choice of models.
While most of the existing gradient-based meta-learners rely on a single initialization, different tasks from the task distribution can require substantially different parameters, making ﬁnding an initialization that is a short distance away from all of them infeasible.
If the task distribution is multimodal with disjoint and far apart modes, one can imagine that a set of separate meta-learners with each covering one mode could better master the full distribution.
However, associating each task with one of the meta-learners not only requires additional identity information about the task, which is not always available or could be ambiguous when the modes are not clearly disjoint, but also eliminates the possibility of transferring knowledge across different modes of the task distribution.
To overcome this issue, we aim to develop a meta-learner that is able to acquire a meta-learned prior and adapt quickly given tasks sampled from a multimodal task distribution.
To this end, we leverage the strengths of the two main lines of existing meta-learning techniques: model-based and gradient-based meta-learning.
Speciﬁcally, we propose to augment gradient-based meta-learners with the capability of generalizing across a multimodal task distribution.
Instead of learning a single initialization point in the parameter space, we propose to ﬁrst compute the task identity of a sampled task by examining task related samples.
Given the estimated task identity, our model then performs a step of model-based adaptation to condition the meta-learned initialization on the sampled task in a step we call modulation.
Then, from this modulated meta-prior, a few steps of gradient-based adaptation are performed towards the target task to progressively improve the performance on the task.
The main contributions of this paper are as follows: • We identify and empirically demonstrate a limitation in a family of widely used gradientbased meta-learners, including Finn et al. ; Kim et al. (2018); Lee & Choi (2018); Grant et al.
; Nichol & Schulman (2018). • We propose a model that is able to acquire a meta-learned prior and adapt quickly given tasks sampled from a multimodal task distribution by taking advantage of both model-based and gradient-based meta-learning. • We design a set of multimodal meta-learning problems and demonstrate that our model offers a better generalization ability on tasks including regression, classiﬁcation, and reinforcement learning.
R ELATED W ORK
[Figure]
learner. The former strives to identify the task from a few samples and modulate the prior accordingly; the latter performs gradient updates to effectively adapt to the speciﬁc task.
We review the several families of meta-learning, which our work builds upon:
Optimization-based meta-learning.
The approaches that learn to optimize a learner model are known as optimization-based meta-learning.
Pioneered by Schmidhuber (1987); Bengio et al.
(1992), optimization algorithms with learned parameters have been proposed, enabling the automatic exploitation of the structure of learning problems.
From a reinforcement learning perspective, Li &
Malik (2016) represent an optimization algorithm as a learning policy.
Andrychowicz et al.
(2016) train
LSTM optimizers to learn update rules from the history of gradients, and Ravi & Larochelle (2016) train a meta-learner LSTM to update a learner’s parameters for few-shot image classiﬁcation.
Unlike these methods, our algorithm does not require additional parameters to learn update rules; instead, we augment our meta-learner with the ability to compute the identity of tasks.
Model-based meta-learning.
Model-based metalearning frameworks learn to recognize the identities of tasks and adjust the model state ( e . g .
the internal state of an RNN) to ﬁt the task.
Santoro et al.
(2016) train a network with an external memory that is able to assimilate new samples and leverage this data to make accurate predictions.
Duan et al.
(2016) represent a fast RL algorithm as an
RNN and learn it from data.
Munkhdalai & Yu (2017) propose to learn task-agnostic knowledge and use it to shift its fast parameters for rapid generalization.
Though our model has a model-based meta-learning component, the main adaptation mechanism is gradient-based, which has preferable behavior on tasks outside the task distribution (Finn & Levine, 2017).
Gradient-based meta-learning.
Finn et al.
(2017) and its extensions Finn et al.
; Kim et al.
(2018); Lee & Choi (2018); Grant et al.
(2018), known as gradient-based meta-learning, aim to estimate a parameter initialization among the task-speciﬁc models, that provides a favorable inductive bias for fast adaptation.
With the model agnostic nature, appealing results have been shown on a variety of learning problems.
However, assuming tasks are sampled from a concentrated distribution and pursuing a common initialization to all tasks can substantially limit the performance of such methods on multimodal task distributions where the center in the task space becomes ambiguous.
In this paper, we propose to ﬁrst identify the mode of a sampled task, in a procedure which is similar to model-based meta-learning approaches Santoro et al.
; Mishra et al.
Then, we modulate the meta-learned prior in the parameter space to make the model better ﬁt to the mode and take gradient steps to rapidly improve the performance on the task afterwards.
P RELIMINARIES
We aim to rapidly adapt to a novel task sampled from a multimodal task distribution.
We consider a target dataset D consisting of tasks sampled from a multimodal distribution.
The dataset is split into meta-training and meta-testing sets, which are further divided into task-speciﬁc training D train T and validation D val T sets.
A meta-learner learns about the underlying structure of the task distribution through training on the meta-training set and is evaluated on meta-testing set.
Our work builds upon Model-Agnostic Meta-Learning (MAML) algorithm
(Finn et al., 2017).
MAML seeks an initialization of parameters θ of a meta-learner such that after a small number of gradient steps to minimize the task-speciﬁc loss on the training set of a new task D train T , the adapted parameters produce good results on the validation set of that task D val T .
The initialization of the parameters is trained by sampling minibatches of tasks from D , computing the adapted parameters for all D train T in the batch, evaluating the validation losses on the corresponding D val T using the adapted parameters and ﬁnally backpropagating the sum of the validation losses through the adaptation step to the initial parameters θ .
M ETHOD
We aim to develop a Multi-Modal Model-Agnostic Meta-Learner (M U M O MAML) that is able to quickly master a novel task sampled from a multimodal task distribution .
To this end, we propose to leverage the ability of model-based meta-learners to identify tasks sampled from a task distribution as well as the ability of gradient-based meta-learners to consistently improve the performance with a few gradient steps.
Speciﬁcally, we propose to learn a model-based meta-learner that produces a set of task-speciﬁc parameters to modulate the meta-learned prior parameters.
Then, this modulated prior learns to adapt to a target task rapidly through gradient-based optimization.
An illustration of our model is shown in Figure 1.
In this section, we ﬁrst introduce our model-based meta-learner and a variety of modulation operators in section 4.1.
Then we describe the training details for M U M O MAML in section 4.2 and present the implementation details in section 4.3.
4. M ODEL - BASED M ETA - LEARNER
Given K data points and labels { x k , y k } k =1 ,...,K , the task embedding network f produces an embed- ding vector υ that encodes the characteristics of a task according to υ
= f
( { x k , y k } k =1 ,...,K ; ω f ) .
The modulation network g modulates the meta-learned prior of the gradient-based meta-learner in the parameter space based on the task embedding vector υ .
To enable specialization of each block of the gradient-based meta-learner to the task, we apply the modulation block-wise to activate or deactivate the units of a block ( i .
e .
a channel of a convolutional layer or a neuron of a fullyconnected layer).
Speciﬁcally, modulation network produces the modulation vectors for each block i by τ 1 , ..., τ N = g ( υ ; ω g ) .
We formalize the procedure of applying modulation as: φ i
=
θ i (cid:12)
τ i , where φ i represents the modulated prior parameters for the gradient-based meta-learner, and (cid:12) represents a general modulation function.
In the experiments, we investigate some representative modulation operations including attention-based modulation (Mnih et al., 2014; Vaswani et al., 2017) and feature-wise linear modulation (FiLM)
(Perez et al., 2017).
(Mnih et al., 2014; Vaswani et al., 2017) has been widely used in modern deep learning models and has proved its effectiveness across various tasks (Yang et al., 2016; Mnih et al., 2014; Zhang et al., 2018; Xu et al., 2015).
Inspired by the previous works, we employed attention to modulate the prior model.
In concrete terms, attention over the outputs of all neurons (Softmax) or a binary gating value (Sigmoid) on each neuron’s output is computed by the model-based meta-learner.
These modulation vectors τ are then used to scale the pre-activation of each neural network layer F θ , such that F φ = F θ ⊗ τ .
Note that here ⊗ represents a channel-wise multiplication.
Feature-wise linear modulation (FiLM) (Perez et al., 2017) proposed to modulate neural networks for achieving the conditioning effects of data from different modalities.
We adopt FiLM as an option for modulating our gradient-based meta-learner.
Speciﬁcally, the modulation vectors τ are divided into two components τ γ and τ β such that for a certain layer of the neural network with its pre-activation F θ , we would have F φ = F θ ⊗ τ γ + τ β .
It can be viewed as a more generic form of attention mechanism.
Please refer to Perez et al.
(2017) for the complete details.
Interestingly, in a recent few-shot learning paper FiLM modulation is used in a metric learning model and high performance on few-shot learning tasks is attained (Oreshkin et al., 2018).
4. G RADIENT - BASED F AST A DAPTATION
We perform a gradient-based meta-learning step to ﬁne-tune the modulated initialization to ﬁt the target task T i .
The concrete training procedure is described by the following pseudo-code:
Algorithm 1 Meta-Training Procedure for M U M O MAML.
Input: Task distribution P ( T ) , Hyper-parameters α and β
Randomly initialize θ and ω .
while not DONE do
Sample batches of tasks T j ∼ P ( T ) for all j do
Compute the modulation vector τ = g
( { x, y } ; ω ) with K samples from D train T j
Evaluate ∇ θ L T j (cid:16) f
( x ; θ, τ ); D train T j (cid:17) with respect to the these K samples
Compute adapted parameter with gradient descent: θ (cid:48)T j = θ − α ∇ θ L T j (cid:0) f
( x ; θ, τ ); D train T j (cid:1) end for Update θ ← θ − β ∇ θ (cid:80)
T j ∼ P ( T ) L T j (cid:0) f ( x ; θ (cid:48) , τ ); D val T j (cid:1) Update ω ← ω − β ∇ ω (cid:80)
T j ∼ P ( T ) L T j (cid:0) f ( x ; θ (cid:48) , τ ); D val T j (cid:1) end while
Note that the modulation vectors τ are not updated in the inner loop.
The model-based learner is responsible for ﬁnding a good task-speciﬁc initialization through modulation.
The gradient-based learning phase shown in the inner optimization loop is responsible for ﬁtting the target task T with one or few gradient updates.
During the meta-testing phase, with a few data samples D train T i from task T i , we perform an adaptation step that corresponds to the inner loop of the meta-training procedure.
Our model-based metalearner ﬁrst computes the task embedding based on the input data samples and then computes the modulation vector τ to modulate the prior model.
Then, gradient updates are performed on the prior model f
( x ; θ, τ ) to compute the adapted parameters θ (cid:48) .
Finally, we evaluate our adapted model f
( x ; θ (cid:48) , τ ) on the meta-test.
4. I MPLEMENTATION D ETAILS
For the model-based meta-learner, we used S EQ 2S EQ (Sutskever et al., 2014) encoder structure to encode the sequence of { x, y } k =1 ,...,K with a bidirectional GRU (Chung et al., 2014) and use the last hidden state of the recurrent model as the task embedding υ .
Then, the modulation vectors τ
i are computed from υ using a separate one hidden-layer multi-layer perceptron (MLP) for each layer of the gradient-based meta-learner.
We experiment with our model in three representative learning [Table] [Figure] scenarios namely regression, few-shot classiﬁcation and reinforcement learning.
The architectures used for each task differ from each other due to the differences in the task nature and data format.
We discuss the tasks and the corresponding architectures in detail in section 5.
E XPERIMENTS
In this section, we evaluate M U M O MAML in a variety of tasks including regression, few-shot image classiﬁcation and reinforcement learning.
We design experiments with multimodal task distributions for regression and reinforcement learning tasks.
To shed some light on the behavior of the proposed model on task distributions whose structure is not clearly multimodal or unimodal, we experiment with few-shot image classiﬁcation.
5. R EGRESSION
We investigate the capability of our model to learn few-shot regression tasks sampled from multimodal task distributions.
In these tasks, a few input/output pairs { x k , y k } k =1 ,...,K sampled from a one dimensional function are given and the model is asked to predict L output values y q 1 , ..., y qL for input queries x q 1 , ..., x qL .
We set up two regression settings with two task modes (sinusoidal and linear functions) or three modes (quadratic functions added).
As a baseline beside MAML, we propose Multi-MAML, which consists of M (the number of modes) separate MAML models which are chosen for each query based on ground-truth task-mode labels.
This baseline serves as an upper-bound for the performance of MAML when the task-mode labels are available. [Figure]
The quantitative results are shown in Table 1.
We observe that Multi-MAML outperforms MAML, showing that the performance of MAML degrades on multimodal task distributions.
The marginal gap between the performance of our model in two and three mode settings indicates that M UM O MAML is able to clearly identify the task modes and has sufﬁcient capacity for all modes.
Also, M U M O MAML consistently achieves better results than Multi-MAML demonstrating that our model is able to discover and exploit transferable knowledge across different modes to improve its performance.
The qualitative results are shown in Figure 2, visualizing the predicted functions.
We observe that
M U M O MAML is able to effectively modulate its meta-learned prior to ﬁt a sampled task (see Figure 3 (a)), which greatly eases the optimization procedure of our gradient-based learner (see Figure(b)).
To gain some insights into what kind of task embeddings are produced by our model, we show tSNE (Maaten & Hinton, 2008) visualization of the predicted embedding vectors υ in Figure 4 (a).
The tSNE plot shows that our model is able to capture the mode structure in the embeddings, enabling performance gain from modulation.
We compared attention modulation with Sigmoid or Softmax and FiLM modulation and found that
FiLM achieves better results.
We therefore use FiLM in the subsequent experiments.
Please refer to
Supplementary Material for additional results.
5. F EW - SHOT I MAGE C LASSIFICATION
The task of few-shot image classiﬁcation considers a problem of classifying images into N classes with a small number (K) of labeled samples available.
To evaluate our model on this task, we conduct experiments on O MNIGLOT , a widely used handwritten character dataset of binary images.
The results are shown in Table 2, demonstrating that our method achieves comparable or better results against state-of-the-art algorithms.
Please refer to Supplementary Material for details.
To gain insights to the task embeddings υ produced by our model, we again randomly sampledtasks and employ tSNE to visualize the υ in Figure 4 (b).
While we are not able to clearly distinguish the modes of task distributions, we observe that the distribution of the produced embeddings is not uniformly distributed or unimodal, potentially indicating the model-based meta-learner ﬁnding some exploitable structure in this task distribution.
5. R EINFORCEMENT L EARNING
We experiment with M U M O MAML in reinforcement learning (RL) to verify its ability to learn to rapidly adapt to tasks sampled from multimodal task distributions given a minimum amount of interaction with an environment.
Speciﬁcally, we consider a Markov decision process (MDP) with horizon H formed by a distribution of tasks T i .
Each task has an initial state distribution p
i ( s ) , a transition distribution p i ( s t +1 | s t , a t ) , and a reward function R i ( s t , a t ) .
Each rollout contains ( s 1 , a 1 , ..., s H ) and its corresponding rewards R ( s t , a t ) and K rollouts are available for K -shot RL. [Figure] [Table]
To analyze our model’s performance on an MDP formed by a multimodal distribution of tasks, we designed and conducted experiments on a 2D navigation environment similar to Finn et al.
In the environment, an agent is trained to navigate to a goal location sampled from a bimodal distribution.
The agent takes its current location as input and has an action space limited to a 2D vector representing the moving orientation.
At each time step, it receives a reward deﬁned as the negative distance to the goal.
Instead of uniformly sampling goals from a square like Finn et al.
, we sample goals from M Gaussian distributions that are far away from each others, which makes our environment suitable for our analysis but more challenging.
To identify the mode of a task distribution in reinforcement learning, we run the meta-learned prior model without modulation or adaptation to interact with the environment and collect a single trajectory and obtained rewards.
Then the collected trajectory and rewards are fed to our model-based meta-learner to compute the task embedding υ and the modulation vectors τ .
Next, the gradient-based meta-learner operates on the modulated policy and seeks to rapidly adapt to the environment – approach the goal.
We refer the reader to A.3 for more details.
We compare M U M O MAML against MAML.
In all experiments, we set H = 100 , and M = 2 .
The average returns are shown in Figure 5 (a), demonstrating that our model is able to achieve a signiﬁcant performance gain from modulation and can consistently outperform MAML.
Qualitative comparisons are shown in Figure 6.
After taking a few gradient updates, both the models are able [Figure] to reach to goals while M U M O MAML converges in fewer steps.
Please refer to Supplementary
Material for additional results.
To investigate the correlation between sampled goals and produced embeddings υ , we samplegoals, collect the trajectories generated by the prior model and compute the task embeddings for the goals.
As shown in Figure 4 (c), a clear separation of goals sampled from two goal modes indicates that our meta-learned initial policy provides useful trajectories and our model-based learner is able to produce meaningful embeddings.
To further examine this hypothesis, we uniformly sample goals along a straight line from the center of a goal mode to the center of the other goal mode and compute the embeddings.
As shown in Figure 5, the clean curved structure demonstrate the competence of our model-based meta-learner.
C ONCLUSION
We present a novel approach that is able to leverage the strengths of both model-based and gradientbased meta-learners to discover and exploit the structure of multimodal task distributions.
Given a few samples from a target task, our model-based learner ﬁrst identiﬁes the mode of the task distribution and then modulates the meta-learned prior in a parameter space.
Next, the gradient based meta-learner efﬁciently adapts to the target task through gradient updates.
We empirically observe that our model-based learner is capable of effectively recognizing the task modes and producing embeddings that captures the structure of a multimodal task distribution.
We evaluated our proposed model in few-shot regression, image classiﬁcation and reinforcement learning, and achieved superior generalization performance on tasks sampled from multimodal task distributions.
A A DDITIONAL E XPERIMENTAL D ETAILS A. R EGRESSION
Setups.
To form multimodal task distributions, we consider a family of functions including sinusoidal functions (in forms of A · sin w · x + b + (cid:15) , with A ∈ [0 . 1 , 5 . 0] , w ∈ [0 . 5 , 2 . 0] and b ∈ [0 , 2 π ] ), linear functions (in forms of A · x + b , with A ∈ [ − 3 , 3] and b ∈ [ − 3 , 3] ) and quadratic functions (in forms of A · ( x − c ) 2 + b , with A ∈ [ − 0 .
15 , − 0 .
02] ∪ [0 .
15] , c ∈ [ − 3 . 0 , 3 . 0] and b ∈ [ − 3 . 0 , 3 . 0] ).
A Gaussian observation noise with µ = 0 and (cid:15) = 0 .
3 is added to each data point sampled from the target task.
In all the experiments, K is set to 5 and L is set to 10 .
We report the mean squared error (MSE) as the evaluation criterion.
Due to the multimodality and uncertainty, this setting is more challenging comparing to (Finn et al., 2017).
Models and Optimization.
In the regression task, we trained a 4-layer fully connected neural networks with the hidden dimensions of 100 and ReLU non-linearity for each layer, as the base model for both MAML and M U M O MAML.
In M U M O MAML, an additional model with a Bidirectional
GRU of hidden size 40 associated with multiple linear layers is trained to generate τ and to modulate each layer of the base model.
We used the same hyper-parameter settings as the regression experiments presented in Finn et al.
and used Adam Kingma & Ba (2014) as the meta-optimizer.
For all our models, we train on 5 meta-train examples and evaluate on 10 meta-val examples to compute the loss.
A. F EW - SHOT I MAGE C LASSIFICATION
Setups.
In the few-shot learning experiments, we used O MNIGLOT , a dataset consists of 50 languages, with a total of 1632 different classes with 20 instances per class.
Following Santoro et al.
(2016), we downsampled the images to 28 × 28 and perform data augmentation by rotating each member of an existing class by a multiple of 90 degrees to form new data points of a given class.
Models and Optimization.
Following prior works (Vinyals et al., 2016; Finn et al., 2017), we used the same 4-layer convolutional neural network and applied the same training and testing splits from Finn et al.
and compare our model against baselines for 5-way and 20-way, 1-shot and 5-shot classiﬁcation.
A. R EINFORCEMENT L EARNING
Environment.
In the bimodal 2D navigation environment the goals are sampled from one of two bivariate Gaussians with means of [0 . 5 , 0 . 5] and [ − 0 .
5 , − 0 .
5] and standard deviation of 0 . 5 .
Each mode is selected with equal probability.
In the beginning of each episode, the agent starts from the origin.
The agent’s observation is its 2D-location and the reward is the negative squared distance to the goal.
The agent does not observe the goal directly, instead it must learn to navigate there based on the reward function alone.
The episodes terminate after 100 steps or when the agent comes to the distance of 0 .
01 from the goal.
Models and Optimization.
The ﬁrst trajectory from the environment is sampled using the unmodulated, unadapted model and used for computing the task embedding.
We experimented with sampling more trajectories for this purpose, but found no improvement over using only one.
The batch of trajectories used for computing the ﬁrst gradient-based adaptation step is sampled using the modulated model and the batches after that use the modulated and adapted parameters from the previous update step.
For the gradient adaptation steps, we use the vanilla policy gradient algorithm Williams (1992).
As the meta-optimizer we use the trust region policy optimization algorithm
Schulman et al.
With respect to the gradient-based adaptation we follow the meta-learning procedure described in Finn et al.
The models are trained for one gradient adaptation step with a batch size of 20 trajectories.
We use 20 tasks for the meta-batch.
We chose the hyperparameters of M U M O MAML and MAML through random search.
For MAML we use the inner loop update step size 0.03 and the discounting parameter γ of 0.95, for M U M O MAML we use step size 0.01 and γ 0.99.
MAML uses a two-layer
MLP model with hidden size 32 and ReLU activations.
For M U M O MAML, we use the same size
MLP in addition to the model-based meta-learner, which consists of an RNN for the embedding network and an MLP for the modulation network.
The RNN model used is the GRU with hidden size 8.
The number of hidden units in the embedding network is 8.
B A DDITIONAL E XPERIMENTAL R ESULTS
For the better understanding of our paper, we provide additional results in this section.
B. A DDITIONAL Q UALITATIVE R ESULTS FOR R EGRESSION
Additional qualitative results for M U M O MAML after modulation are shown in Figure 7 and additional qualitative results for M U M O MAML after adaptation are shown in Figure 8.
B. A DDITIONAL Q UALITATIVE R ESULTS FOR R EINFORCEMENT L EARNING
Additional trajectories sampled from the 2D navigation environment are presented in Figure 9.
[Figure] [Figure]