Covariate Distribution Aware Meta-learning
Abstract
Meta-learning has proven to be successful at fewshot learning across the regression, classiﬁcation and reinforcement learning paradigms.
Recent approaches have adopted Bayesian interpretations to improve gradient based meta-learners by quantifying the uncertainty of the post-adaptation estimates.
Most of these works almost completely ignore the latent relationship between the covariate distribution ( p ( x )) of a task and the corresponding conditional distribution p ( y | x ) .
In this paper, we identify the need to explicitly model the meta-distribution over the task covariates in a hierarchical Bayesian framework.
We begin by introducing a graphical model that explicitly leverages very few samples drawn from p ( x ) to better infer the posterior over the optimal parameters of the conditional distribution ( p ( y | x )) for each task.
Based on this model we provide an inference strategy and a corresponding meta-algorithm that explicitly accounts for the meta-distribution over task covariates.
Finally, we demonstrate the signiﬁcant gains of our proposed algorithm on a synthetic regression dataset.
1. Introduction
Learning quickly or with very few samples has been a longterm goal of the machine learning community.
The ﬁeld of meta-learning has recently made signiﬁcant strides towards achieving that goal.
Meta-learning (Nichol et al., 2018;
Ravi & Larochelle, 2016;
Finn et al., 2017) comprises of a set of algorithms designed to exploit prior experiences from multiple tasks (drawn from a task distribution) for improving sample-efﬁciency on a new but related task from the same distribution.
Given the increasing cost of getting annotated samples on an ever-increasing variety of related tasks, the practical scope of these algorithms is immense.
Most meta-learning methods can be classiﬁed into two broad categories (i) gradient-based (Ravi & Beatson, 2018;
> . Denevi et al., 2019; Finn et al., 2017) approaches that metalearn parameters of optimization algorithms (like initialization and learning rate) in a way that the meta-learner (optimizer) is amenable to quickly adapt on a new task by performing gradient descent on a very small number of labeled samples, and (ii) amortized-inference (Snell et al., 2017;
Lee et al., 2019; Bertinetto et al., 2018)
based approaches that directly infer the optimal parameters of a new task without performing any gradient based optimization.
In general, such algorithms learn to adapt the parameters of a complex neural network using only a few samples from an unseen task in a way that the adapted network generalizes well (Wang et al., 2019).
In this work, although we focus on improving gradient-based methods, we believe that our core idea can be adapted to the latter as well.
Recent works (Finn et al., 2018), (Ravi & Beatson, 2018), (Kim et al., 2018) have used a Bayesian framework to learn a suitable prior over the network parameters by leveraging the inherent structure of the task distribution.
By viewing the parameters of a meta-learner through a Bayesian lens we can use the predictive posterior (Gal & Ghahramani, 2016) to estimate the uncertainty of the adapted parameters for each task (Ravi & Beatson, 2018).
In a Bayesian meta-learner (Kim et al., 2018; Finn et al., 2018; Ravi & Beatson, 2018), the posterior over the adapted network parameters for a new task is typically inferred using a few samples from the task along with a meta-learned prior.
In this work, we hypothesize that the covariate distribution of a task can also inﬂuence the posterior over the adapted network parameters.
To the best of our knowledge none of the existing meta-learning algorithms like Bertinetto et al.
(2018); Rajeswaran et al.
; Ravi & Beatson (2018);
Finn et al.
(2018) explicitly utilize the information present in the covariates to improve the estimate of the adapted parameters.
This is done by modeling the latent factors of the covariate distribution.
We deﬁne a prior not only on the network parameters (which determine the conditional p ( y | x ) ) but also on the covariate distribution p ( x ) .
Our meta-learning objective involves maximizing the joint likelihood p ( x, y ) as opposed to just p ( y | x ) which leads to meta-parameters sharing information about the covariates across tasks, in addition to the optimal network parameters.
This way the latent factors of the covariate distribution p ( x ) of a new task can be quickly inferred from very few covariates.
Finally, the inferred latent covariate factors are used to infer the posterior over the adapted network parameters.
For simplicity, we motivate the need to model covariates via a synthetic example in Fig.
The metadistribution consists of tasks with optimal hypothesis h ∈
H =
(cid:83) {H s , H l , H t , H q } that can be classiﬁed into four hypothesis classes: sinusoidal ( H s ), linear
( H l ), tanh ( H t ) and quadratic ( H q ).
We note that the support of the input distribution is vastly different for each of the four hypothesis classes.
Thus, intuitively we can see that inferring properties of the covariate distribution of a task can be helpful in adjusting the posterior over the network parameters.
For example, for a given task if we only observe negative covariates in the range [ − 10 , − 5] we can adapt the posterior to have a higher measure for H s (sinusoidal hypothesis).
We can generalize the above intuition to real-world metalearning problems as well, especially for high-dimensional data like images.
Particularly, in the limited availability of labeled data, semi-supervised methods (Chapelle et al., 2009) tend to leverage unlabeled data (and hence covariate distribution) to attain generalizable models (Berthelot et al., 2019)
On similar lines, modeling the meta-distribution over the covariates can help inferring the task-speciﬁc covariate distribution which can then better inform task-speciﬁc features for image classiﬁcation.
In few-shot classiﬁcation, images for different tasks can lie on different manifolds (Saul & Roweis, 2003) of varying complexities.
Information about the covariates (by modeling the unlabeled data) can help us better estimate the required complexity of the discriminative features for a given task.
For example, distinguishing species of plants can be considered harder than classifying mammals since the features of plant images may be cluttered on a low-dimensional manifold as opposed to the possibly well separated features in the case of mammals.
Thus the former may require having complex (non-linear) decision boundaries as opposed to simpler linear classiﬁers in the case of the latter.
The main contributions of our work are as follows: (1) we identify the need to model the latent structure present in the covariate distributions ( p ( x )) for a sequence of tasks (2) to the best of our knowledge we are the ﬁrst to propose a Bayesian framework which exploits this latent information to better infer the posterior over the adapted network parameters (that deﬁne p ( y | x ) ) (3) we propose a gradient based model-agnostic meta-learning algorithm that is an instantiation of our probabilistic theory and demonstrate its beneﬁts on synthetic regression datasets.
2. Related Work
Our methodology is complementary to most existing works in the probabilistic meta-learning literature.
We borrow the basic hierarchical Bayes framework from Ravi & Beatson (2018); Finn et al.
and extend it to model Bayesian variables that generate the covariate distribution for a task.
This enables our method to be model-agnostic while having the ability to beneﬁt from the latent relationship between the task covariates and the optimal parameters as mentioned in
Sec. 1.
In the non-Bayesian setting, the M - MAML algorithm proposed by Vuorio et al.
(2019) is mildly similar to our approach in the sense that they learn task speciﬁc initializations instead of a single one as originally introduced by
Finn et al.
M - MAML uses the labeled samples to infer an initialization for a given task and hence one can view the covariate distribution as being used indirectly.
But they fail to explicitly model the mutual information between the covariates and adapted parameters.
On the other hand, our approach is more direct since it ﬁrst infers the posterior over the latent factors of the covariate distribution via a maximum likelihood objective and then uses the inferred posterior to improve the adaptation of network parameters.
Additionally, our framework is capable of modeling the uncertainty of the adaptation which can prove to be be critical in the few shot scenario.
3. Methodology
We begin by introducing some notations for the metalearning setup used in the rest of the paper followed by the proposed probabilistic framework which explicitly exploits (i) the structure of the covariate distributions across tasks (ii) the relation between the covariate distribution and optimal hypothesis for a given task.
We then derive the
Maximum Likelihood Estimation ( MLE ) objectives for the observed variables in our model and show how the MLE derivations can inform a novel meta-learning objective.
Finally, we discuss a speciﬁc meta-learning algorithm that can efﬁciently optimize the proposed objective.
We do this via an instantiation of the generic approach obtained by making certain simplifying assumptions in the original framework.
Notations We are given a sequence of n tasks {T i } ni =1 with each task T i having m labeled samples given by the dataset
D i
= { x ( i ) j , y ( i ) j } mj =1 where x ( i ) j ∈ X ⊂ R k
and y ( i )
j ∈
Y ⊂ R .
Following the deﬁnitions introduced by Finn et al.
we split the dataset D i :=
{D Si , D Qi } into support ( D Si ) and query ( D Qi ) sets respectively with |D Si | =
m (cid:48) , |D Qi | = m − m (cid:48) .
Each sample in D i is drawn from the joint distribution p
i ( x, y ) over X × Y with the marginals given by p i ( x ) and p i ( y ) .
The probabilistic model we consider in our work is summarized in Fig.
Without making any assumptions on the nature of p i ( x, y ) , we assume the existence of meta-parameters θ that govern the common structure shared across the set of joint distributions { p i ( x, y ) } ni =1 .
Within each task the generative model for the input x ( i ) involves a random variable z xi which we shall refer to as the latent factors of the covariate distribution.
Also, each task has an additional latent variable z y i which plays a role in the generative model for the response variable y
( i )
given the input x ( i ) .
In most settings, a naive assumption of independence is made over the latent factors z xi and z yi .
On the other hand, we refrain from making such assumptions and instead exploit the information present in the covariates { x ( i ) j } mj =1 to better infer the posterior over the latent variable z yi which inﬂuences the conditional p
i ( y | x ) .
3.1. Formal Derivations
In this section, we derive a lower bound for the likelihood of the observed data {D Si } ni
=1 using hierarchical variational inference.
This gives us the meta-learning objective that can be optimized using standard gradient-based approaches.
[Equation]
In the above equation, the distribution q ( θ ; β ) is a variational approximation (with parameters β ) for the true posterior over the meta-parameter θ .
For the derivations henceforth we shall drop the notations i and S when understood from context.
The log-likelihood of the dataset D i given by p ( D i | θ ) , can be written as an integral over the factors p
[Equation]
D i | z xi , z yi ) , p ( z yi | z xi , θ ) and p ( z xi | θ ) .
[Equation]
To lower bound the log of the above objective we introduce two variational approximations (i) q ( z xi ; κ i ) with parameters κ i for the true posterior p ( z xi |{ x ( i ) j } m (cid:48)
j =1 ) and (ii) q ( z yi ; λ i ) with parameters
λ
i for the true posterior p ( z yi |
z xi , D Si ) .
[Equation]
Since z xi is the latent factor in the generative model for x ( i ) and z yi is the corresponding latent variable for y ( i ) , we arrive at the independence: y ( i )
j ⊥ z xi | z yi , x ( i ) j and x ( i )
j ⊥ z yi | z xi .
Based on this, we ﬁnally arrive at the following Evidence
Lower Bound ( ELBO ) for log p
( D i | θ ) which we shall refer to as L D i ( κ i , λ i , θ ) .
[Equation]
Therefore, the ELBO on the likelihood of the dataset for i th task is a function of the task-speciﬁc variational parameters κ i , λ i and the variational meta-parameter β .
For each task, the optimal variational parameters that approximate the true posteriors are distinct.
Hence, κ i , λ i need to be adapted for each task individually.
Given L D i ( κ i , λ i , θ ) we can re-write the lower bound in Eq.
1 as:
[Equation]
3.2. Algorithm
The primary aim of any meta-learning algorithm is to optimize for the meta-parameter θ given the datasets {D i } ni = from the corresponding sequence of tasks.
This is generally a two step process where step-I involves identifying the optimal task-speciﬁc parameters using θ and the support set
D Si .
In step-II, based on the task-speciﬁc adapted parameters from step-I the meta-parameter θ is optimized over the query set D Qi .
Within our framework, since both the meta-parameter θ and the task-speciﬁc parameters z xi , z yi are Bayesian random variables with variational parameters given by β, λ i , κ i respectively, we instead deﬁne an algorithm to optimize the ELBO in Eq.
[Equation]
In order to optimize the objectives in Eqs.
6, 7 we introduce certain simplifying assumptions over each of the variational approximations in Sec. 3.1.
Assumption 1.
The variational approximation q ( θ ; β ) follows a δ − distribution given by δ β and the prior p ( θ ) is given by N ( 0 , I ) .
We note that the assumptions taken are effective in arriving at a computationally feasible algorithm and are
Assumption 2.
The variational approximation q ( z yi ; κ i ) follows a normal distribution with mean, diagonal covariance matrix given by κ ∗ i =
( γ µ ( D Si ; β ) , γ 2 σ ( D Si ; β )) .
The prior p ( z xi | θ ) is taken as N ( 0 , I ) .
On the other hand, the distribution q ( z xi ; λ i ) is given by a δ − distribution: δ λ i .
Using Assms.
1, 2 the Eqs.
6, 7 can be re-written as a two layer log-likelihood objective with an l 2 regularization term.
[Equation]
To be concrete, the task-speciﬁc random variable z yi represents the parameters of the neural network for the i th task, that takes as input x
( i ) j and outputs a prediction
ˆy ( i ) j .
On the other hand, as stated previously z xi represents the latent variable for the covariate distribution p
i ( x ) .
Furthermore, the optimal variational parameters for the distribution q ( z ix ; κ i ) is given by κ ∗ i which consists of the mean γ µ ( D Si ; β ) and std.
deviation γ σ ( D Si ; β ) of a normal distribution.
Here, γ µ ( · , β ) , γ σ ( · , β ) represent neural networks which take as input the support set D S i and output
κ ∗ i .
It is important to note that even though the parameters of γ µ , γ σ are taskagnostic, the variational parameter κ ∗ i is still different for each task and is determined using the covariates in D Si .
Having identiﬁed κ ∗
i we now describe the optimization algorithm for λ i in Eq.
Notice that to obtain λ ∗
i it is sufﬁcient to only minimize the objective − E q ( z xi ; κ ∗ ) L (cid:48)D i ( z xi , λ i , β ) .
In Eq.
4 the KL term acts as a regularizer in the optimization objective for λ i .
Since the most common algorithm for optimization is Stochastic Gradient Descent ( SGD ) many meta-learning algorithms avoid the KL term by choosing a regularization speciﬁc to SGD .
In most works (Ravi &
Beatson, 2018; Finn et al., 2018; Kim et al., 2018), the KL term is a function of only the meta-parameter θ (or β given
Assm. 1).
Hence the regularization is induced by letting the initialization for the optimization of λ i (given by λ (0) i ) be determined by β .
In our framework, we realize that the KL term is a function of both β and the latent variable z xi for the task speciﬁc covariate distribution.
Hence we model the initialization λ (0)
i using a neural-network whose parameters are subsumed in β and thus without loss of expressivity λ (0) i
= f β ( z x i ) .
Thus, the optimal parameters of the variational approximation ( λ ∗ i ) would be given by performing
K steps of SGD on the MLE objective in Eq.
4 with the initialization given by
[Equation]
The expectation in Eq 11 is computed using monte-carlo approximation.
We ﬁnd that sampling a single value of z xi ∼ q ( z xi ; κ ∗ i ) is sufﬁcient to optimize for λ i .
Algorithm
Given: n datasets: {D Si , D Qi } ni =1 , learning rates: η 0 , η 1 , number of update steps: K .
p ( θ ) , p ( z xi | θ ) ← N ( 0 , I ) for i=1 to n do κ ∗ i ( µ ) , κ ∗ i ( σ ) ← γ µ ( D Si ; β ) , γ σ ( D Si ; β ) λ (0) i
= f β ( κ ∗ i ( µ )
+ (cid:15) ◦ κ ∗ i ( σ )) ; (cid:15) ∼ N ( 0 , I ) for t=0 to K-1 do λ ( t +1)
i ← λ ( t ) i − η 0 ∇ λ
( t ) i m (cid:48) (cid:80) j =1 log p ( y ( i ) j | λ ( t ) i , x ( i ) j ) end for λ ∗
i ← λ ( K ) i β ← β − η 1 ∇ β (cid:20) n (cid:80) i =
1 L D Qi ( κ ∗ i , λ ∗ i , β ) − 12 (cid:107) β (cid:107)(cid:21) end for
Finally, we note that the meta-parameter β constitutes the parameters of the network f β which determines the initialization λ (0) i as well as the parameters of γ µ ( · ; · ) , γ σ ( · ; · ) which output κ ∗ i .
Thus, β is optimized to jointly maximize the likelihood of the covariates of a sequence of tasks as well as for learning to choose covariate dependent initializations suitable for few-shot adaptation.
For the optimization objective in Eq.
8, we use the standard re-parameterization trick ( 2 nd step of the outer for-loop in Algorithm 1) commonly used in Variational Auto-Encoders ( VAE s).
This is done so as to be able to differentiate through the expectation over q ( z xi ; κ ∗ i ) in Eqs.
The step-by-step procedure for the meta-training and meta-testing phases are given by
Algorithm 1 and Algorithm 2 respectively.
Given: dataset D T = { x ( T ) j , y ( T ) j } mj =1 , parameter: β ∗ , learning rate η 0 , number of update steps: K. κ ∗T ( µ ) , κ ∗T ( σ ) ← γ µ ( D T ; β ∗ ) , γ σ ( D T ; β ∗ ) λ (0) T = f β ∗ ( κ ∗T ( µ ) + (cid:15) ◦ κ ∗T ( σ )) ; (cid:15) ∼ N ( 0 , I ) for t=0 to K-1 do λ
( t +1) T ← λ ( t ) T − η 0 ∇ λ ( t ) T m (cid:80)
j =1 log p ( y ( T ) j | λ ( t ) T , x ( T ) j ) end for λ ∗T ← λ ( K ) T
4. Experiments and Results
In order to ﬁrst litmus test our approach on a simpler task we begin by evaluating it on a synthetic regression dataset borrowed from Vuorio et al.
and defer further experimentation on more complicated real world datasets to future work.
Most meta-learning algorithms have been tested on regression datasets where the covariate distribution is same across all tasks.
We thus describe how we suitably modify the original dataset so that there exists a structure over the set of covariate distributions across tasks.
This enables us to fairly evaluate our method against the baselines in the proposed setting.
We compare our algorithm with the popular gradient based meta-learning approach MAML introduced by Finn et al.
Additionally, we also chose as baselines Amortized
MAML (Ravi & Beatson, 2018) and M - MAML (Vuorio et al., 2019) which are exemplars of the Bayesian and task-speciﬁc initialization type approaches respectively.
These methods either model the task-parameters as Bayesian random variables (former) or adapt the parameters of the optimizer based on the input dataset (latter) and hence warrant a close comparison with our proposed methodology.
[Table] [Table]
Dataset The covariate distribution for each task T is given by a normal N ( µ T , σ T ) whose parameters are sampled from a discrete distribution over P pairs { ( µ p , σ p ) }
Pp =1 .
The pairs p 1 , . . .
p P are ﬁxed in the beginning once they are sampled from a pair of independent uniform priors, p i ∼ ( U ( − 10 , 10) , U (0 , 10)) .
The parameters of the discrete distribution are sampled from a Dirichlet prior ( Dir ( α p = 1)) .
Following (Vuorio et al., 2019), the optimal hypothesis for each task is sampled from one of many modalities or hypothesis classes.
The hypothesis classes considered are : sine , linear , quad , transformed-L 1 and tanh .
For each task, having chosen a hypothesis class, the parameters of the optimal hypothesis (like slope of a linear functions) is chosen based on uniform distributions over the following 3 :
sine : f ( x ) =
A · ( sin ( w · x ) + b ) , with A ∈ [0 . 1 , 5 . 0] , w ∈ [0 . 5 , 2 . 0] and b ∈ [0 , 2 π ] .
quad : f ( x ) =
A · ( x − c ) 2 + b with A ∈ [ − 0 .
15 , − 0 .
02] ∪ [0 .
15] , c ∈ [ − 3 . 0 , 3 . 0] and b ∈ [ − 3 . 0 , 3 . 0] .
linear : f ( x ) =
A · x + b , A ∈
[ − 3 , 3] , b ∈ [ − 3 , 3] .
transformed-L 1 : f ( x ) =
A · | x − c | + b , with A ∈ [ − 0 .
15 , − 0 .
02] ∪ [0 .
15] , c ∈ [ − 3 . 0 , 3 . 0] and b ∈ [ − 3 . 0 , 3 . 0] .
tanh : f ( x ) =
A · tanh ( x − c )
+ b with A ∈ [ − 3 . 0 , 3 . 0] , c ∈ [ − 3 . 0 , 3 . 0] , b ∈ [ − 3 . 0 , 3 . 0] .
For each of the true hypothesis classes above the ﬁnal value of y ( i ) j is generated by adding an independent error term (cid:15) sampled from a normal distribution with mean 0 and standard deviation of 0 .
3 i.e y ( i )
j = f ( x ( i ) j ) + (cid:15) .
We consider two cases, ﬁrst where there exists a relation between the parameters (mean, variance) of the covariate distribution and the optimal hypothesis class chosen for a task and second when the optimal hypothesis class is chosen independent of the mean, variance of the covariate distribution.
We highlight the results in each case separately.
Case-I: With a speciﬁc relation This setting conforms to the case when z y (cid:54)⊥⊥ z x | θ in Fig.
We experiment with three different meta-distributions which are of different complexities owing to the variety (number) of hypothesis classes each of them span.
The datasets are listed as follows:
sine : true hypothesis for each class is given by a sinusoidal function with parameters sampled from distributions mentioned previously.
The range for each of the parameters is split into P = 3 disjoint sets, each one corresponding to a speciﬁc covariate distribution.
sine-quad-linear : each of the three hypothesis classes are mapped to a speciﬁc covariate distribution.
Thus, once the parameters of the covariate distribution are sampled based on the discrete prior, the corresponding hypothesis class is also chosen and a task is sampled from it.
ﬁve : similar to the previous case with the distinction that all ﬁve hypothesis classes are considered in this dataset i.e P = 5 .
Table 1 highlights the Mean Squared Errors ( MSE ) achieved by our method and the baselines on the three regression datasets described above.
We can see that when there exists a relation between the true hypothesis class and the covariate distribution our approach performs signiﬁcantly better than other state-of-the-art approaches for regression.
Interestingly, improvements are also observed for the ﬁve dataset which was speciﬁcally introduced by Vuorio et al.
for evaluating M - MAML .
The Bayesian model Amortized
MAML performs poorly since unlike our approach it fails to acknowledge the latent relationship between the covariate distribution and the posterior over the ﬁve hypothesis classes.
Case-II: Independent
This setting conforms to the case when z y ⊥⊥ z x | θ in Fig.
In Table 2 we demonstrate that the performance of our approach is no worse (if not better) than other methods which assume the independence by default.
Once again we experiment with the same three types of meta-distribution mentioned in the previous case.
Implementation The conditional p ( y | x ) is modeled using a three-layered neural network with hidden sizes 100 , 100 , 100 .
The networks γ µ ( · ; β ) , γ σ ( · ; β ) which take as input the sequence of labeled samples in a dataset D Si are modeled using a common Recurrent Neural Network
RNN backbone with output of dimension 28 .
This can be referred to as the task-embedding (Vuorio et al., 2019).
The mapping of task-embedding to λ (0) i , given by f β ( · ) is modeled similar to M - MAML as well as the modulation over λ (0) i which gives us the ﬁnal initialization of the network parameters 4 .
Following prior work (Finn & Levine, 2017), a bias-transformation of size 20 was appended to the inputs.
The total number of training tasks used were 10 , 000 with m (cid:48)
= m − m (cid:48)
= 5 samples for the support and query sets for each task.
Adam optimizer with an initial learning rate of 0 .
001 was used to train the meta-learner.
Additionally, all KL terms were re-weighted with a weight of 0 . 01 .
5. Conclusion
Cognizant of the fact that the generalization performance of few-shot algorithms depends on a varying number of factors ranging from sample size, hypothesis class complexity to the optimization algorithm, input distribution; in this work we focus our efforts on improving meta-learning algorithms by using the covariate distribution to infer the adapted parameters via a principled Bayesian approach.
We begin by deriving ELBO bounds for the hierarchical Bayes formulation and follow it up with a meta-learning algorithm to infer the posterior over the network parameters.
Finally, we show some preliminary results on a synthetic regression dataset designed to test the usefulness of our method.
Motivated by the empirical gains observed, we plan to extend our work to more challenging few-shot image classiﬁcation benchmarks like mini-imagenet (Ravi & Larochelle, 2016) and FC100 (Oreshkin et al., 2018).
We also acknowledge that the proposed framework warrants a more rigorous theoretical analysis to understand exactly how inferring the covariate distribution can impact regret bounds (Khodak et al., 2019b;a) in the online few-shot setting or even generalization error one