M ETA -L EARNING WITH W ARPED G RADIENT D ESCENT
A BSTRACT
Learning an efﬁcient update rule from data that promotes rapid learning of new tasks from the same distribution remains an open problem in meta-learning.
Typically, previous works have approached this issue either by attempting to train a neural network that directly produces updates or by attempting to learn better initialisations or scaling factors for a gradient-based update rule.
Both of these approaches pose challenges.
On one hand, directly producing an update forgoes a useful inductive bias and can easily lead to non-converging behaviour.
On the other hand, approaches that try to control a gradient-based update rule typically resort to computing gradients through the learning process to obtain their metagradients, leading to methods that can not scale beyond few-shot task adaptation.
In this work, we propose Warped Gradient Descent (WarpGrad), a method that intersects these approaches to mitigate their limitations.
WarpGrad meta-learns an efﬁciently parameterised preconditioning matrix that facilitates gradient descent across the task distribution.
Preconditioning arises by interleaving non-linear layers, referred to as warp-layers , between the layers of a task-learner.
Warp-layers are meta-learned without backpropagating through the task training process in a manner similar to methods that learn to directly produce updates.
WarpGrad is computationally efﬁcient, easy to implement, and can scale to arbitrarily large meta-learning problems.
We provide a geometrical interpretation of the approach and evaluate its effectiveness in a variety of settings, including few-shot, standard supervised, continual and reinforcement learning.
I NTRODUCTION
Learning (how) to learn implies inferring a learning strategy from some set of past experiences via a meta-learner that a task-learner can leverage when learning a new task.
One approach is to directly parameterise an update rule via the memory of a recurrent neural network (Andrychowicz et al., 2016; Ravi & Larochelle, 2016; Li & Malik, 2016;
Chen et al., 2017).
Such memory-based methods can, in principle, represent any learning rule by virtue of being universal function approximators (Cybenko, 1989; Hornik, 1991; Schäfer & Zimmermann, 2007).
They can also scale to long learning processes by using truncated backpropagation through time, but they lack an inductive bias as to what constitutes a reasonable learning rule.
This renders them hard to train and brittle to generalisation as their parameter updates have no guarantees of convergence.
An alternative family of approaches deﬁnes a gradient-based update rule and meta-learns a shared initialisation that facilitates task adaptation across a distribution of tasks (Finn et al., 2017; Nichol et al., 2018;
Flennerhag et al., 2019).
Such methods are imbued with a strong inductive bias— gradient descent—but restrict knowledge transfer to the initialisation.
Recent work has shown that it is beneﬁcial to more directly control gradient descent by meta-learning an approximation of a parameterised matrix (Li et al., 2017; Lee & Choi, 2018; Park & Oliva, 2019) that preconditions gradients during task training, similarly to second-order and Natural Gradient Descent methods (Nocedal & Wright, 2006; Amari & Nagaoka, 2007).
To meta-learn preconditioning, these methods backpropagate through the gradient descent process, limiting them to few-shot learning.
x f ( x ) h (1) ω (1) h (2) ω (2) θ θ (1) φ (1) θ (2) φ (2) ∇ f ( x ) L ( θ ; φ )
P (1) ∇ θ (1) L
Dω (1)
P (2) ∇ θ
(2) L
Dω (2) θ θ (cid:48) θ (cid:48)(cid:48) ∇ L
P ∇ L min E θ
[ J ( φ )] φ
Task Adaptation
Meta
[Figure]
Learning
Task-learners
Shared
[Figure]
In this paper, we propose a novel framework called Warped Gradient Descent (WarpGrad) 1 , that relies on the inductive bias of gradient-based meta-learners by deﬁning an update rule that preconditions gradients, but that is meta-learned using insights from memory-based methods.
In particular, we leverage that gradient preconditioning is deﬁned point-wise in parameter space and can be seen as a recurrent operator of order 1.
We use this insight to deﬁne a trajectory agnostic meta-objective over a joint parameter search space where knowledge transfer is encoded in gradient preconditioning.
To achieve a scalable and ﬂexible form of preconditioning, we take inspiration from works that embed preconditioning in task-learners (Desjardins et al., 2015; Lee & Choi, 2018), but we relax the assumption that task-learners are feed-forward and replace their linear projection with a generic neural network ω , referred to as a warp layer .
By introducing non-linearity, preconditioning is rendered data-dependent.
This allows WarpGrad to model preconditioning beyond the block-diagonal structure of prior works and enables it to meta-learn over arbitrary adaptation processes.
We empirically validate WarpGrad and show it surpasses baseline gradient-based meta-learners on standard few-shot learning tasks ( mini ImageNet, tiered ImageNet; Vinyals et al., 2016; Ravi &
Larochelle, 2016; Ren et al., 2018), while scaling beyond few-shot learning to standard supervised settings on the “multi”-shot Omniglot benchmark (Flennerhag et al., 2019) and a multi-shot version of tiered ImageNet.
We further ﬁnd that WarpGrad outperforms competing methods in a reinforcement learning (RL) setting where previous gradient-based meta-learners fail (maze navigation with recurrent neural networks (Miconi et al., 2019)) and can be used to meta-learn an optimiser that prevents catastrophic forgetting in a continual learning setting.
W ARPED G RADIENT D ESCENT 2. G RADIENT - BASED M ETA -L EARNING
WarpGrad belongs to the family of optimisation-based meta-learners that parameterise an update rule θ ← U ( θ ; ξ ) with some meta-parameters ξ .
Speciﬁcally, gradient-based meta-learners deﬁne an update rule by relying on the gradient descent, U ( θ ; ξ ) :
= θ − α ∇ L ( θ ) for some objective L and learning rate α .
A task is deﬁned by a training set D τ train and a test set D τ test , which deﬁnes learning objectives L D τ ( θ ) :
=
E ( x,y )
∼D τ [ (cid:96) ( f ( x, θ ) , y )] over the task-learner f for some loss (cid:96) .
MAML (Finn
[Figure]
et al., 2017) meta-learns a shared initialisation θ by
[Equation]
Subsequent works on gradient-based meta-learning primarily differ in the parameterisation of U .
Meta-SGD (MSGD; Li & Malik, 2016) learns a vector of learning rates, whereas Meta-Curvature (MC; Park & Oliva, 2019) deﬁnes a block-diagonal preconditioning matrix B , and T-Nets (Lee & Choi, 2018) embed block-diagonal preconditioning in feed-forward learners via linear projections,
[Equation]
k
[Equation]
φ )
T-Nets .
These methods optimise for meta-parameters ξ = { θ 0 , φ } by backpropagating through the gradient descent process (Eq. 1).
This trajectory dependence limits them to few-shot learning as they become (1) computationally expensive, (2) susceptible to exploding/vanishing gradients, and (3) susceptible to a credit assignment problem (Wu et al., 2018; Antoniou et al., 2019; Liu et al., 2019).
Our goal is to develop a meta-learner that overcomes all three limitations.
To do so, we depart from the paradigm of backpropagating to the initialisation and exploit the fact that learning to precondition gradients can be seen as a Markov Process of order 1 that depends on the state but not the trajectory (Li et al., 2017).
To develop this notion, we ﬁrst establish a general-purpose form of preconditioning (Section 2.2).
Based on this, we obtain a canonical meta-objective from a geometrical point of view (Section 2.3), from which we derive a trajectory-agnostic meta-objective (Section 2.4).
2. G ENERAL -P URPOSE P RECONDITIONING
A preconditioned gradient descent rule, U ( θ ; φ ) :
= θ − αP ( θ ; φ ) ∇ L ( θ ) , deﬁnes a geometry via
P .
To disentangle the expressive capacity of this geometry from the expressive capacity of the tasklearner f , we take inspiration from T-Nets that embed linear projections T in feed-forward layers, h
= σ ( TWx + b ) .
This in itself is not sufﬁcient to achieve disentanglement since the parameterisation of T is directly linked to that of W , but it can be achieved under non-linear preconditioning.
To this end, we relax the assumption that the task-learner is feed-forward and consider an arbitrary neural network, f = h ( L ) ◦· · ·◦ h (1) .
We insert warp-layers that are universal function approximators parameterised by neural networks into the task-learner without restricting their form or how they
P τ
W τ
P τ (cid:48)
W τ (cid:48) θ θ (cid:48) γ γ (cid:48)
ω ω ∇ ( L ◦ ω )
( θ )
G − 1 ∇ L ( γ ) interact with f .
In the simplest case, we interleave warp-layers between layers of the task-learner to obtain ˆ f = ω ( L ) ◦ h ( L ) ◦ · · · ◦ ω (1) ◦ h (1) , but other forms of interaction can be beneﬁcial (see Appendix A for practical guidelines).
Backpropagation automatically induces gradient preconditioning, as in T-Nets, but in our case via the Jacobians of the warp-layers:
[Equation]
where D x and D θ denote the Jacobian with respect to input and parameters, respectively.
In the special case where f is feed-forward and each ω a linear projection, we obtain an instance of
WarpGrad that is akin to T-Nets since preconditioning is given by D x ω = T .
Conversely, by making warp-layers non-linear, we can induce interdependence between warp-layers, allowing WarpGrad to model preconditioning beyond the block-diagonal structure imposed by prior works.
Further, this enables a form of task-conditioning by making Jacobians of warp-layers data dependent.
As we have made no assumptions on the form of the task-learner or warp-layers, WarpGrad methods can act on any neural network through any form of warping, including recurrence.
We show that increasing the capacity of the meta-learner by deﬁning warp-layers as Residual Networks (He et al., 2017) improves performance on classiﬁcation tasks (Section 4.1).
We also introduce recurrent warp-layers for agents in a gradient-based meta-learner that is the ﬁrst, to the best of our knowledge, to outperform memory-based meta-learners on a maze navigation task that requires memory (Section 4.3).
Warp-layers imbue WarpGrad with three powerful properties.
First, due to preconditioned gradients,
WarpGrad inherits gradient descent properties, importantly guarantees of convergence.
Second, warp-layers form a distributed representation of preconditioning that disentangles the expressiveness of the geometry it encodes from the expressive capacity of the task-learner.
Third, warp-layers are meta-learned across tasks and trajectories and can therefore capture properties of the task-distribution beyond local information.
Figure 3 illustrates these properties in a synthetic scenario, where we construct a family of tasks f : R 2 → R (see Appendix D for details) and meta-learn across the task distribution.
WarpGrad learns to produce warped loss surfaces (illustrated on two tasks τ and τ (cid:48) ) that are smoother and more well-behaved than their respective native loss-surfaces.
2. T HE G EOMETRY OF W ARPED G RADIENT D ESCENT
If the preconditioning matrix P is invertible, it deﬁnes a valid Riemann metric (Amari, 1998) and therefore enjoys similar convergence guarantees to gradient descent.
Thus, if warp-layers represent a valid (meta-learned) Riemann metric, WarpGrad is well-behaved.
For T-Nets, it is sufﬁcient to require T to be full rank, since T explicitly deﬁnes P as a block-diagonal matrix with block entries
TT T .
In contrast, non-linearity in warp-layers precludes such an explicit identiﬁcation.
Instead, we must consider the geometry that warp-layers represent.
For this, we need a metric tensor,
G , which is a positive-deﬁnite, smoothly varying matrix that measures curvature on a manifold W .
The metric tensor deﬁnes the steepest direction of descent by − G − 1 ∇ L (Lee, 2003), hence our goal is to establish that warp-layers approximate some G − 1 .
Let Ω represent the effect of warp-layers by a reparameterisation h ( i ) (
x ; Ω( θ ; φ ) ( i ) )
=
ω ( i ) ( h ( i ) (
x ; θ ( i ) ); φ ) ∀ x, i that maps from a space P onto the manifold W with γ = Ω( θ ; φ ) .
We induce a metric G on W by push-forward (Figure 2):
[Equation]
where G − 1 :
=
[ D x Ω][ D x Ω] T .
Provided Ω is not degenerate ( G is non-singular), G − 1 is positivedeﬁnite, hence a valid Riemann metric.
While this is the metric induced on W by warp-layers, it is not the metric used to precondition gradients since we take gradient steps in P which introduces an error term (Figure 2).
We can bound the error by ﬁrst-order Taylor series expansion to establish ﬁrst-order equivalence between the WarpGrad update in P (Eq. 7) and the ideal update in W (Eq.
[Equation]
Consequently, gradient descent under warp-layers (in P -space) is ﬁrst-order equivalent to warping the native loss surface under a metric G to facilitate task adaptation.
Warp parameters φ control the geometry induced by warping, and therefore what task-learners converge to.
By meta-learning φ we can accumulate information that is conducive to task adaptation but that may not be available during that process.
This suggests that an ideal geometry (in W -space) should yield preconditioning that points in the direction of steepest descent, accounting for global information across tasks,
[Equation]
In contrast to MAML-based approaches (Eq. 1), this objective avoids backpropagation through learning processes.
Instead, it deﬁnes task learning abstractly by introducing a joint distribution over objectives and parameterisations, opening up for general-purpose meta-learning at scale.
2. M ETA -L EARNING W ARP P ARAMETERS
The canonical objective in Eq.
10 describes a meta-objective for learning a geometry on ﬁrst principles that we can render into a trajectory-agnostic update rule for warp-layers.
To do so, we deﬁne a task τ
= ( h τ , L τ meta , L τ task ) by a task-learner ˆ f that is embedded with a shared WarpGrad optimiser, a meta-training objective L τ meta , and a task adaptation objective L τ task .
We use L τ task to adapt task parameters θ and L τ meta to adapt warp parameters φ .
Note that we allow meta and task objectives to differ in arbitrary ways, but both are expectations over some data, as above.
In the simplest case, they differ in terms of validation versus training data, but they may differ in terms of learning paradigm as well, as we demonstrate in continual learning experiment (Section 4.3).
To obtain our meta-objective, we recast the canonical objective (Eq. 10) in terms of θ using ﬁrst-order equivalence of gradient steps (Eq. 9).
Next, we factorise p ( τ, θ ) into p ( θ | τ ) p ( τ ) .
Since p ( τ ) is given, it remains to consider a sampling strategy for p ( θ | τ ) .
For meta-learning of warp-layers, we assume this distribution is given.
We later show how to incorporate meta-learning of a prior p ( θ 0 | τ ) .
While any sampling strategy is valid, in this paper we exploit that task learning under stochastic gradient descent can be seen as sampling from an empirical prior p ( θ | τ ) (Grant et al., 2018); in particular, each iterate θ τk can be seen as a sample from p ( θ τk | θ τk − 1 , φ ) .
Thus, K -steps of gradient descent forms a Monte-Carlo chain θ τ 0 , . . .
, θ τK and sampling such chains deﬁne an empirical distribution p ( θ | τ ) around some prior p ( θ 0 | τ ) , which we will discuss in Section 2.5.
The joint distribution p ( τ, θ ) deﬁnes a joint search space across tasks.
Meta-learning therefore learns a geometry over this space with the steepest expected direction of descent.
This direction is however not with respect to the objective that produced the gradient, L τ task , but with respect to L τ meta ,
[Equation]
Decoupling the task gradient operator ∇L τ task from the geometry learned by L τ meta lets us infuse global knowledge in warp-layers, a promising avenue for future research (Metz et al., 2019; Mendonca et al., 2019).
For example, in Section 4.3, we meta-learn an update-rule that mitigates catastrophic forgetting by deﬁning L τ meta over current and previous tasks.
In contrast to other gradient-based meta-learners, the WarpGrad meta-objective is an expectation over gradient update steps sampled from the search space induced by task adaptation (for example, K steps of stochastic gradient descent; Figure 2).
It is therefore trajectory agnostic and hence compatible with arbitrary task learning processes.
Because the meta-gradient is independent of the number of task gradient steps, it avoids vanishing/exploding gradients and the credit assignment problem by design.
It does rely on second-order gradients, a requirement we can relax by detaching task parameter gradients ( ∇L τ task ) in Eq.
[Equation]
[Equation]
where sg is the stop-gradient operator.
In contrast to the ﬁrst-order approximation of MAML (Finn et al., 2017), which ignores the entire trajectory except for the ﬁnal gradient, this approximation retains all gradient terms and only discards local second-order effects, which are typically dominated by ﬁrst-order effect in long parameter trajectories (Flennerhag et al., 2019).
Empirically, we ﬁnd that our approximation only incurs a minor loss of performance in an ablation study on Omniglot (Appendix F).
Interestingly, this approximation is a form of multitask learning with respect to φ (Li & Hoiem, 2016;
Bilen & Vedaldi, 2017; Rebufﬁ et al., 2017) that marginalises over task parameters θ τ .
Require: p ( τ ) : distribution over tasks
Require: α, β, λ : hyper-parameters 1: initialise φ and p ( θ 0 | τ ) 2: while not done do 3: sample mini-batch of tasks T from p ( τ ) 4: g φ , g θ 0 ←5: for all τ ∈ T do 6: θ τ 0 ∼ p ( θ 0 | τ ) 7: for all k in 0 , . . .
, K τ − 1 do 8: θ
τk +1 ← θ τk − α ∇L τ task ( θ τk ; φ ) 9: g φ ← g φ + ∇ L ( φ ; θ τk ) 10: g θ 0 ← g θ 0 + ∇ C ( θ 0 ; θ τ 0: k ) 11: end for 12: end for 13: φ ← φ − βg φ 14: θ 0 ← θ 0 − λβg θ15: end while
Require: p ( τ ) : distribution over tasks
Require: α, β, λ, η : hyper-parameters 1: initialise φ , p ( θ 0 | τ ) 2: while not done do 3: initialise buffer B = {} 4: sample mini-batch of tasks T from p ( τ ) 5: for all τ ∈ T do 6: θ τ 0 ∼ p ( θ
| τ ) 7:
B [ τ ] =
[ θ τ 0 ] 8: for all k in 0 , . .
, K τ − 1 do 9: θ τk +1 ← θ τk − α ∇L τ task ( θ τk ; φ ) 10:
B [ τ ] .append ( θ τk +1 ) 11: end for 12: end for 13: i, g φ , g θ 0 ←14: for all
( τ, k ) ∈ B do 15: g φ ← g φ + ∇ L ( φ ; θ τk ) 16: g θ 0 ← g θ 0 + ∇ C
( θ τ 0 ; θ τ 0: k ) 17: i ← i +18: if i =
η then 19: φ ← φ − βg φ 20: θ 0 ← θ 0 − λβg θ21: i, g φ , g θ 0 ←22: end if 23: end for 24: end while
2. I NTEGRATION WITH L EARNED I NITIALISATIONS
WarpGrad is a method for learning warp layer parameters φ over a joint search space deﬁned by p ( τ, θ ) .
Because WarpGrad takes this distribution as given, we can integrate WarpGrad with methods that deﬁne or learn some form of “prior” p ( θ 0 | τ ) over θ τ 0 .
For instance, (a) Multi-task solution : in online learning, we can alternate between updating a multi-task solution and tuning warp parameters.
We use this approach in our Reinforcement Learning experiment (Section 4.3); (b) Meta-learned point-estimate : when task adaptation occurs in batch mode, we can meta-learn a shared initialisation θ 0 .
Our few-shot and supervised learning experiments take this approach (Section 4.1); (c) Metalearned prior : WarpGrad can be combined with Bayesian methods that deﬁne a full prior (Rusu et al., 2019; Oreshkin et al., 2018; Lacoste et al., 2018;
Kim et al., 2018).
We incorporate such methods by some objective C (potentially vacuous) over θ 0 that we optimise jointly with WarpGrad,
[Equation]
) where L can be substituted for by ˆ L and λ ∈ [0 , ∞ ) is a hyper-parameter.
We train the WarpGrad optimiser via stochastic gradient descent and solve Eq. 13 by alternating between sampling task parameters from p ( τ, θ ) given the current parameter values for φ and taking meta-gradient steps over these samples to update φ .
As such, our method can also be seen as a generalised form of gradient descent in the form of Mirror Descent with a meta-learned dual space (Desjardins et al., 2015; Beck & Teboulle, 2003).
The details of the sampling procedure may vary depending on the speciﬁcs of the tasks (static, sequential), the design of the task-learner (feed-forward, recurrent), and the learning objective (supervised, self-supervised, reinforcement learning).
In Algorithm 1 we illustrate a simple online algorithm with constant memory and linear complexity in K , assuming the same holds for C .
A drawback of this approach is that it is relatively data inefﬁcient; in Appendix B we detail a more complex ofﬂine training algorithm that stores task parameters in a replay buffer for mini-batched training of φ .
The gains of the ofﬂine variant can be dramatic: in our Omniglot experiment (Section 4.1), ofﬂine meta-training allows us to update warp parameters 2000 times with each meta-batch, improving ﬁnal test accuracy from 76 .
3% to 84 .
3 % (Appendix F).
R ELATED W ORK
Learning to learn, or meta-learning, has previously been explored in a variety of settings.
Early work focused on evolutionary approaches (Schmidhuber, 1987; Bengio et al., 1991; Thrun & Pratt, 1998).
Hochreiter et al.
(2001) introduced gradient descent methods to meta-learning, speciﬁcally for recurrent meta-learning algorithms, extended to RL by Wang et al.
(2016) and Duan et al.
A similar approach was taken by Andrychowicz et al.
and Ravi & Larochelle (2016) to meta-learn a parameterised update rule in the form of a Recurrent Neural Network (RNN).
A related idea is to separate parameters into “slow” and “fast” weights, where the former captures meta-information and the latter encapsulates rapid adaptation (Hinton & Plaut, 1987; Schmidhuber, 1992; Ba et al., 2016).
This can be implemented by embedding a neural network that dynamically adapts the parameters of a main architecture (Ha et al., 2016).
WarpGrad can be seen as learning slow warp-parameters that precondition adaptation of fast weights.
Recent meta-learning focuses almost exclusively on few-shot learning, where tasks are characterised by severe data scarcity.
In this setting, tasks must be sufﬁciently similar that a new task can be learned from a single or handful of examples (Lake et al., 2015; Vinyals et al., 2016;
Snell et al., 2017; Ren et al., 2018).
Several meta-learners have been proposed that directly predict the parameters of the tasklearner (Bertinetto et al., 2016; Munkhdalai et al., 2018; Gidaris & Komodakis, 2018; Qiao et al., 2018).
To scale, such methods typically pretrain a feature extractor and predict a small subset of the parameters.
Closely related to our work are gradient-based few-shot learning methods that extend
MAML by sharing some subset of parameters between task-learners that is ﬁxed during task training but meta-learner across tasks, which may reduce overﬁtting (Mishra et al., 2018; Lee & Choi, 2018;
Munkhdalai et al., 2018) or induce more robust convergence (Zintgraf et al., 2019).
It can also be used to model latent variables for concept or task inference, which implicitly induce gradient modulation (Zhou et al., 2018;
Oreshkin et al., 2018; Rusu et al., 2019; Lee et al., 2019).
Our work is also related to gradient-based meta-learning of a shared initialisation that scales beyond few-shot learning (Nichol et al., 2018; Flennerhag et al., 2019).
Meta-learned preconditioning is closely related to parallel work on second-order optimisation methods for high dimensional non-convex loss surfaces (Nocedal & Wright, 2006; Saxe et al., 2013; Kingma & Ba, 2015; Arora et al., 2018).
In this setting, second-order optimisers typically struggle to improve upon ﬁrst-order baselines (Sutskever et al., 2013).
As second-order curvature is typically intractable to compute, such methods resort to low-rank approximations (Nocedal & Wright, 2006; Martens, 2010; Martens & Grosse, 2015) and suffer from instability (Byrd et al., 2016).
In particular, Natural
Gradient Descent (Amari, 1998) is a method that uses the Fisher Information Matrix as curvature metric (Amari & Nagaoka, 2007).
Several proposed methods for amortising the cost of estimating the metric (Pascanu & Bengio, 2014; Martens & Grosse, 2015; Desjardins et al., 2015).
As noted by Desjardins et al. (2015), expressing preconditioning through interleaved projections can be seen as a form of Mirror Descent (Beck & Teboulle, 2003).
WarpGrad offers a new perspective on gradient preconditioning by introducing a generic form of model-embedded preconditioning that exploits global information beyond the task at hand.
E XPERIMENTS
[Table]
We evaluate WarpGrad in a set of experiments designed to answer three questions: (1) do WarpGrad methods retain the inductive bias of MAML-based fewshot learners?
(2) Can WarpGrad methods scale to problems beyond the reach of such methods?
(3) Can WarpGrad generalise to complex meta-learning problems?
4. F EW -S HOT L EARNING
For few-shot learning, we test whether
WarpGrad retains the inductive bias of gradient-based meta-learners while avoiding backpropagation through the gradient descent process.
To isolate the effect of the WarpGrad objective, we use linear warp-layers that we train using online meta-training (Algorithm 1) to make WarpGrad as close to T-Nets as possible.
For a fair comparison, we meta-learn the initialisation using MAML (Warp-MAML) with J ( θ 0 , φ ) :
= L ( φ ) + λC MAML ( θ 0 ) .
We evaluate the importance of metalearning the initialisation in Appendix G and ﬁnd that WarpGrad achieves similar performance under random task parameter initialisation.
All task-learners use a convolutional architecture that stacks 4 blocks made up of a 3 × 3 convolution, max-pooling, batch-norm, and ReLU activation.
We deﬁne Warp-MAML by inserting warp-layers in the form of 3 × 3 convolutions after each block in the baseline task-learner.
All baselines are tuned with identical and independent hyper-parameter searches (including ﬁlter sizes – full experimental settings in Appendix H), and we report best results from our experiments or the literature.
WarpMAML outperforms all baselines (Table 1), improving 1- and 5-shot accuracy by 3 . 6 and 5 .percentage points on mini ImageNet (Vinyals et al., 2016; Ravi & Larochelle, 2016) and by 5 . 2 and 3 .
8 percentage points on tiered ImageNet (Ren et al., 2018), which indicates that WarpGrad retains the inductive bias of MAML-based meta-learners.
[Figure]
T e s t a cc u r a c y o n h e l d - o u t t a s k s
[Figure]
4. M ULTI -S HOT L EARNING
Next, we evaluate whether WarpGrad can scale beyond few-shot adaptation on similar supervised problems.
We propose a new protocol for tiered ImageNet that increases the number of adaptation steps to 640 and use 6 convolutional blocks in task-learners, which are otherwise deﬁned as above.
Since MAML-based approaches cannot backpropagate through 640 adaptation steps for models of this size, we evaluate WarpGrad against two gradient-based meta-learners that meta-learn an initialisation without such backpropagation, Reptile (Nichol et al., 2018) and Leap (Flennerhag et al., 2019), and we deﬁne a Warp-Leap meta-learner by J ( θ 0 , φ ) :
= L ( φ ) + λC Leap ( θ 0 ) .
Leap is an attractive complement as it minimises the expected gradient descent trajectory length across tasks.
Under
WarpGrad, this becomes a joint search for a geometry in which task adaptation deﬁnes geodesics (shortest paths, see Appendix C for details).
While Reptile outperforms Leap by 2 . 6 percentage points on this benchmark, Warp-Leap surpasses both, with a margin of 3 . 88 to Reptile (Table 1).
We further evaluate Warp-Leap on the multi-shot Omniglot
(Lake et al., 2011) protocol proposed by Flennerhag et al.
(2019), where each of the 50 alphabets is a 20-way classiﬁcation task.
Task adaptation involves 100 gradient steps on random samples that are preprocessed by random afﬁne transformations.
We report results for Warp-Leap under ofﬂine meta-training (Algorithm 2), which updates warp parameters 2000 times per meta step (see Appendix E for experimental details).
WarpLeap enjoys similar performance on this task as well, improving over Leap and Reptile by 8 . 1 and 12 . 8 points respectively (Table 1).
We also perform an extensive ablation study varying the number of tasks in the meta-training set.
Except for the case of a single task, Warp-Leap substantially outperforms all baselines (Figure 4), achieving a higher rate of convergence and reducing the ﬁnal test error from ~30% to ~15% .
Non-linear warps, which go beyond block-diagonal preconditioning, reach ~11% test error (refer to Appendix F and Table 2 for the full results).
Finally, we ﬁnd that
WarpGrad methods behave distinctly different from Natural Gradient Descent methods in an ablation study (Appendix G).
It reduces ﬁnal test error from ~42% to ~19%, controlling for initialisation, while its preconditioning matrices differ from what the literature suggests (Desjardins et al., 2015).
4. C OMPLEX M ETA -L EARNING : R EINFORCEMENT AND C ONTINUAL L EARNING (c.1) Reinforcement
Learning To illustrate how WarpGrad may be used both with recurrent neural networks and in meta-reinforcement learning, we evaluate it in a maze navigation task proposed by Miconi et al.
The environment is a ﬁxed maze and a task is deﬁned by randomly choosing a goal location.
The agent’s objective is to ﬁnd the location as many times as possible, being teleported to a random location each time it ﬁnds it.
We use advantage actor-critic with a basic recurrent neural network (Wang et al., 2016) as the task-learner, and we design a Warp-RNN as a HyperNetwork (Ha et al., 2016) that uses an LSTM that is ﬁxed during task training.
This LSTM modulates the weights of the task-learning RNN (deﬁned in Appendix I), which in turn is trained on mini-batches of 30 episodes for 200 000 steps.
We accumulate the gradient of ﬁxed warp-parameters continually (Algorithm 3, Appendix B) at each task parameter update.
Warp parameters are updated on every 30 th step on task parameters (we control for meta-LSTM capacity in Appendix I).
We compare against
Learning to Reinforcement Learn (Wang et al., 2016) and Hebbian meta-learning (Miconi et al., 2018; 2019); see Appendix I for details.
Notably, linear warps (T-Nets) do worse than the baseline
RNN on this task while the Warp-RNN converges to a mean cumulative reward of ~160 in 60episodes, compared to baselines that reach at most a mean cumulative reward of ~125 after 100episodes (Figure 4)
~
200 000 episodes (I).
c.
Continual Learning We test if WarpGrad can prevent catastrophic forgetting (French, 1999) in a continual learning scenario.
To this end, we design a continual learning version of the sine regression meta-learning experiment in Finn et al.
(2017) by splitting the input interval [ − 5 , 5] ⊂ R into 5 consecutive sub-tasks (an alternative protocol was recently proposed independently by Javed &
White, 2019).
Each sub-task is a regression problem with the target being a mixture of two random sine waves.
We train 4-layer feed-forward task-learner with interleaved warp-layers incrementally on one sub-task at a time (see Appendix J for details).
To isolate the behaviour of WarpGrad parameters, we use a ﬁxed random initialisation for each task sequence.
Warp parameters are meta-learned to prevent catastrophic forgetting by deﬁning L τ meta to be the average task loss over current and previous sub-tasks, for each sub-task in a task sequence.
This forces warp-parameters to disentangle the adaptation process of current and previous sub-tasks.
We train on each sub-task for 20 steps, for a total of 100 task adaptation steps.
We evaluate WarpGrad on 100 random tasks and ﬁnd that it learns new sub-tasks well, with mean losses on an order of magnitude 10 − 3 .
When switching sub-task, performance immediately deteriorates to ~ 10 − 2 but is stable for the remainder of training (Figure 5).
Our results indicate that WarpGrad can be an effective mechanism against catastrophic forgetting, a promising avenue for further research.
For detailed results, see Appendix J.
C ONCLUSION
We propose WarpGrad, a novel meta-learner that combines the expressive capacity and ﬂexibility of memory-based meta-learners with the inductive bias of gradient-based meta-learners.
WarpGrad meta-learns to precondition gradients during task adaptation without backpropagating through the adaptation process and we ﬁnd empirically that it retains the inductive bias of MAML-based few-shot learners while being able to scale to complex problems and architectures.
Further, by expressing preconditioning through warp-layers that are universal function approximators, WarpGrad can express geometries beyond the block-diagonal structure of prior works.
WarpGrad provides a principled framework for general-purpose meta-learning that integrates learning paradigms, such as continual learning, an exciting avenue for future research.
We introduce novel means for preconditioning, for instance with residual and recurrent warp-layers.
Understanding how
WarpGrad manifolds relate to second-order optimisation methods will further our understanding of gradient-based meta-learning and aid us in designing warp-layers with stronger inductive bias.
In their current form, WarpGradshare some of the limitations of many popular meta-learning approaches.
While WarpGrad avoids backpropagating through the task training process, as in Warp-Leap , the WarpGrad objective samples from parameter trajectories and has therefore linear computational complexity in the number of adaptation steps, currently an unresolved limitation of gradient-based meta-learning.
Algorithm 2 hints at exciting possibilities for overcoming this limitation.
A CKNOWLEDGEMENTS
The authors would like to thank Guillaume Desjardins for helpful discussions on an early draft as well as anonymous reviewers for their comments.
SF gratefully acknowledges support from North
West Doctoral Training Centre under ESRC grant ES/J500094/1 and by The Alan Turing Institute under EPSRC grant EP/N510129/
B W ARP G RAD M ETA -T RAINING A LGORITHMS
In this section, we provide the variants of WarpGrad training algorithms used in this paper.
Algorithm 1 describes a simple online algorithm, which accumulates meta-gradients online during task adaptation.
This algorithm has constant memory and scales linearly in the length of task trajectories.
In Algorithm 2, we describe an ofﬂine meta-training algorithm.
This algorithm is similar to Algorithm 1 in many respects, but differs in that we do not compute meta-gradients online during task adaptation.
Instead, we accumulate them into a replay buffer of sampled task parameterisations.
This buffer is a Monte-Carlo sample of the expectation in the meta objective (Eq. 13) that can be thought of as a dataset in its own right.
Hence, we can apply standard mini-batching with respect to the buffer and perform mini-batch gradient descent on warp parameters.
This allows us to update warp parameters several times for a given sample of task parameter trajectories, which can greatly improve data efﬁciency.
In our Omniglot experiment, we found ofﬂine meta-training to converge faster: in fact, a mini-batch size of 1 (i.e. η = 1 in Algorithm 2 converges rapidly without any instability.
Finally, in Algorithm 3, we present a continual meta-training process where meta-training occurs throughout a stream of learning experiences.
Here, C represents a multi-task objective, such as the average task loss, C multi = (cid:80) τ ∼ p ( τ ) L τ task .
Meta-learning arises by collecting experiences continuously (across different tasks) and using these to accumulate the meta-gradient online.
Warp parameters are updated intermittently with the accumulated meta-gradient.
We use this algorithm in our maze navigation experiment, where task adaptation is internalised within the RNN tasklearner.
C W ARP G RAD O PTIMISERS
In this section, we detail WarpGrad methods used in our experiments.
Warp-MAML We use this model for few-shot learning (Section 4.1).
We use the full warpobjective in Eq.
11 together with the MAML objective (Eq. 1),
[Equation]
14) where C MAML =
L MAML under the constraint P = I .
In our experiments, we trained Warp-MAML using the online training algorithm
(Algorithm 1).
Warp-Leap We use this model for multi-shot meta-learning.
It is deﬁned by applying Leap (Flennerhag et al., 2019) to θ
(Eq. 16),
[Equation]
15) where the Leap objective is deﬁned by minimising the expected cumulative chordal distance,
[Equation]
Note that the Leap meta-gradient makes a ﬁrst-order approximation to avoid backpropagating through the adaptation process.
It is given by
[Equation]
where ∆ L τ task ( θ τk ; φ ) :
= L τ task ( θ τk ; φ ) − L τ task (cid:0) θ τk − 1 ; φ (cid:1) and ∆ θ τk := θ τk − θ τk − 1 .
In our experiments, we train Warp-Leap using Algorithm 1 in the multi-shot tiered ImageNet experiment and Algorithm in the Omniglot experiment.
We perform an ablation study for training algorithms, comparing exact (Eq. 11) versus approximate (Eq. 12) meta-objectives, and several implementations of the warp-layers on Omniglot in Appendix F.
Warp-RNN For our Reinforcement Learning experiment, we deﬁne a WarpGrad optimiser by meta-learning an LSTM that modulates the weights of the task-learner (see Appendix I for details).
For this algorithm, we face a continuous stream of experiences (episodes) that we meta-learn using our continual meta-training algorithm (Algorithm 3).
In our experiment, both L τ task and L τ meta are the advantage actor-critic objective (Wang et al., 2016); C is computed on one batch of 30 episodes, whereas L is accumulated over η = 30 such batches, for a total of 900 episodes.
As each episode involves 300 steps in the environment, we cannot apply the exact meta objective, but use the approximate meta objective (Eq. 12).
Speciﬁcally, let E τ = { s 0 , a 1 , r 1 , s 1 , . . .
, s T , a T , r T , s T +1 } denote an episode on task τ , where s denotes state, a action, and r instantaneous reward.
Denote a minibatch of randomly sampled task episodes by E =
{ E τ } τ ∼ p ( τ ) and an ordered set of k consecutive mini-batches by E k =
E k − i } k − 1 i =0 .
Then ˆ L ( φ ; E k )
= 1 /n (cid:80) E i ∈E k
(cid:80) E τi,j ∈ E i
L τ meta ( φ ; θ, E τi,j ) and
Require: p ( τ ) : distribution over tasks
Require:
1: initialise
2: while not done do 3:
Sample mini-batch tasks
) 4: g
5: for all do 6: θ
←
7: for all k in 0 , . . .
, K τ − 1 do 8: θ
τk +1 ← θ τk − α ∇L τ task ( θ τk ; φ ) 9: g φ ← g φ + ∇ L ( φ ; θ τk ) 10: g θ 0 ← g θ 0 + ∇ C ( θ τ 0 ; θ τ 0: k ) 11: end for 12: end for 13: φ ← φ − βg φ 14: θ 0 ← θ 0 − λβg θ15: end while
Require: p ( τ ) : distribution over tasks
Require: α, β, λ, η : hyper-parameters 1: initialise φ and θ 2: i, g φ , g θ ←3: while not done do 4:
Sample mini-batch of tasks B from p ( τ ) 5: for all τ ∈ B do 6: g φ ← g φ + ∇ L ( φ ; θ ) 7: g θ ← g θ + ∇ C ( θ ; φ ) 8: end for 9: θ ← θ − λβg θ 10: g θ , i ← 0 , i +11: if i = η then 12: φ ← φ − βg φ 13: i, g θ ←14: end if 15: end while
Require: p ( τ ) : distribution over tasks
Require: α, β, λ, η : hyper-parameters 1: initialise φ and θ2: while not done do 3:
Sample mini-batch of tasks B from p ( τ ) 4:
T ← { τ : [ θ 0 ] for τ in B} 5: for all τ ∈ B do 6: θ τ 0 ← θ7: for all k in 0 , . . .
, K τ − 1 do 8: θ τk +1 ← θ τk − α ∇L τ task ( θ τk ; φ ) 9:
T [ τ ] .append ( θ τk +1 ) 10: end for 11: end for 12: i, g φ , g θ 0 ←13: while T not empty do 14: sample τ, k without replacement 15: g φ ← g φ + ∇ L ( φ ; θ τk ) 16: g θ 0 ← g θ 0 + ∇ C ( θ τ 0 ; θ τ 0: k ) 17: i ← i +18: if i =
η then 19: φ ← φ − βg φ 20: θ 0 ← θ 0 − λβg θ21: i, g φ , g θ 0 ←22: end if 23: end while 24: end while
C multi ( θ ; E k ) =
1 /n (cid:48)
(cid:80) E τk,j ∈ E k L τ task ( θ ; φ, E τk,j ) , where n and n (cid:48) are normalising constants.
The
Warp-RNN objective is deﬁned by
[Equation]
WarpGrad for Continual Learning For this experiment, we focus on meta-learning warpparameters.
Hence, the initialisation for each task sequence is a ﬁxed random initialisation, (i.e. λC ( θ 0 ) = 0 ).
For the warp meta-objective, we take expectations over N task sequences, where each task sequence is a sequence of T = 5 sub-tasks that the task-learner observes one at a time; thus while the task loss is deﬁned over the current sub-task, the meta-loss averages of the current and all prior sub-tasks, for each sub-task in the sequence.
See Appendix J for detailed deﬁnitions.
Importantly, because WarpGrad deﬁnes task adaptation abstractly by a probability distribution, we can readily implement a continual learning objective by modifying the joint task parameter distribution p ( τ, θ ) that we use in the meta-objective (Eq. 11).
A task deﬁnes a sequence of sub-tasks over which we generate parameter trajectories θ τ .
Thus, the only difference from multi-task meta-learning is that parameter trajectories are not generated under a ﬁxed task, but arise as a function of the continual learning algorithm used for adaptation.
We deﬁne the conditional distribution p ( θ | τ ) as before by sampling sub-task parameters θ τ t from a mini-batch of such task trajectories, keeping track of which sub-task t it belongs to and which sub-tasks came before it in the given task sequence τ .
The meta-objective is constructed, for any sub-task parameterisation θ τ t , as L τ meta ( θ τ t )
= 1 /t (cid:80)
[Equation]
where D j is data from sub-task j (Appendix J).
The meta-objective is an expectation over task parameterisations,
[Equation]
D S YNTHETIC E XPERIMENT
To build intuition for what it means to warp space, we construct a simple 2-D problem over loss surfaces.
A learner is faced with the task of minimising an objective function of the form f τ ( x 1 , x 2 ) =
g τ 1 ( x 1 ) exp( g τ 2 ( x 2 ))
− g τ 3 ( x 1 ) exp( g τ 4 ( x 1 , x 2 )) − g τ 5 exp( g τ 6 ( x 1 )) , where each task f τ is deﬁned by scale and rotation functions g τ that are randomly sampled from a predeﬁned distribution.
Speciﬁcally, each task is deﬁned by the objective function
[Equation]
where each a, b and s are randomly sampled parameters from
[Equation]
The task is to minimise the given objective from a randomly sampled initialisation, x { i =1 , 2 } ∼ U ( − 3 , 3) .
During meta-training, we train on a task for 100 steps using a learning rate of 0 . 1 .
Each task has a unique loss-surface that the learner traverses from the randomly sampled initialisation.
While each loss-surface is unique, they share an underlying structure.
Thus, by metalearning a warp over trajectories on randomly sampled loss surfaces, we expect WarpGrad to learn a warp that is close to invariant to spurious descent directions.
In particular, WarpGrad should produce a smooth warped space that is quasi-convex for any given task to ensure that the task-learner ﬁnds a minimum as fast as possible regardless of initialisation.
To visualise the geometry, we use an explicit warp Ω deﬁned by a 2-layer feed-forward network with a hidden-state size of 30 and tanh non-linearities.
We train warp parameters for 100 meta-training steps; in each meta-step we sample a new task surface and a mini-batch of 10 random initialisations that we train separately.
We train to convergence and accumulate the warp meta-gradient online (Algorithm 1).
We evaluate against gradient descent on randomly sampled loss surfaces (Figure 7).
Both optimisers start from the same initialisation, chosen such that standard gradient descent struggles; we expect the WarpGrad optimisers to learn a geometry that is robust to the initialisation (top row).
This is indeed what we ﬁnd; the geometry learned by WarpGrad smoothly warps the native loss surface into a well-behaved space where gradient descent converges to a local minimum.
E O MNIGLOT
We follow the protocol of Flennerhag et al.
(2019), including the choice of hyper-parameters unless otherwise noted.
In this setup, each of the 50 alphabets that comprise the dataset constitutes a distinct task.
Each task is treated as a 20-way classiﬁcation problem.
Four alphabets have fewer thancharacters in the alphabet and are discarded, leaving us with 46 alphabets in total.
10 alphabets are held-out for ﬁnal meta-testing; which alphabets are held out depend on the seed to account for variations across alphabets; we train and evaluate all baselines on 10 seeds.
For each character in an alphabet, there are 20 raw samples.
Of these, 5 are held out for ﬁnal evaluation on the task while the remainder is used to construct a training set.
Raw samples are pre-processed by random afﬁne transformations in the form of (a) scaling between [0 . 8 , 1 . 2] , (b) rotation [0 , 360) , and (c) cropping height and width by a factor of [ − 0 . 2 , 0 .
2] in each dimension.
This ensures tasks are too hard for few-shot learning.
During task adaptation, mini-batches are sampled at random without ensuring class-balance (in contrast to few-shot classiﬁcation protocols (Vinyals et al., 2016)).
Note that benchmarks under this protocol are not compatible with few-shot learning benchmarks.
We use the same convolutional neural network architecture and hyper-parameters as in Flennerhag et al.
This learner stacks a convolutional block comprised of a 3 × 3 convolution withﬁlters, followed by 2 × 2 max-pooling, batch-normalisation, and ReLU activation, four times.
All images are down-sampled to 28 × 28 , resulting in a 1 × 1 × 64 feature map that is passed on to a ﬁnal linear layer.
We create a Warp Leap meta-learner that inserts warp-layers between each convolutional block, W ◦ ω 4 ◦ h 4 ◦ · · · ◦ ω 1 ◦ h 1 , where each h is deﬁned as above.
In our main experiment, each ω i is simply a 3 × 3 convolutional layer with zero padding; in Appendix F we consider both simpler and more sophisticated versions.
We ﬁnd that relatively simple warp-layers do quite well.
However, adding capacity does improve generalisation performance.
We meta-learn the initialisation of task parameters using the Leap objective (Eq. 16), detailed in Appendix C.
Both L τ meta and L τ task are deﬁned as the negative log-likelihood loss; importantly, we evaluate them on different batches of task data to ensure warp-layers encourage generalisation.
We found no additional beneﬁt in this experiment from using held-out data to evaluate L τ meta .
We use the ofﬂine meta-training algorithm (Appendix B, Algorithm 2); in particular, during meta-training, we sample mini-batches
Number of tasks in meta-training set 0 .0 .0 .0 .0 .
T e s t a cc u r a c y o n h e l do u t t a s k s
Warp-Leap
Leap
Reptile
FT †
SGD ‡
KFAC ‡
Number of tasks in meta-training set 0 .0 .0 .0 .0 .
T r a i n
A cc
AU C o n h e l do u t t a s k s
Warp-Leap
Leap
Reptile
FT †
SGD ‡
KFAC ‡ of 20 tasks and train task-learners for 100 steps to collect 2000 task parameterisations into a replay buffer.
Task-learners share a common initialisation and warp parameters that are held ﬁxed during task adaptation.
Once collected, we iterate over the buffer by randomly sampling mini-batches of task parameterisations without replacement.
Unless otherwise noted, we used a batch size of η = 1 .
For each mini-batch, we update φ by applying gradient descent under the canonical meta-objective (Eq. 11), where we evaluate L τ meta on a randomly sampled mini-batch of data from the corresponding task.
Consequently, for each meta-batch, we take (up to) 2000 meta-gradient steps on warp parameters φ .
We ﬁnd that this form of mini-batching causes the meta-training loop to converge much faster and induces no discernible instability.
We compare Warp-Leap against no meta-learning with standard gradient descent (SGD) or
KFAC (Martens & Grosse, 2015).
We also benchmark against baselines provided in Flennerhag et al.
; Leap, Reptile (Nichol et al., 2018), MAML, and multi-headed ﬁne-tuning.
All learners beneﬁt substantially from large batch sizes as this enables higher learning rates.
To render no-pretraining a competitive option within a fair computational budget, we allow SGD and KFAC to use 4x larger batch sizes, enabling 10x larger learning rates.
[Table]
F A BLATION STUDY : W ARP L AYERS , M ETA -O BJECTIVE , AND M ETA -T RAINING
WarpGrad provides a principled approach for model-informed meta-learning and offers several degrees of freedom.
To evaluate these design choices, we conduct an ablation study on Warp-Leap where we vary the design of warp-layers as well as meta-training approaches.
For the ablation study, we ﬁxed the number of pretraining tasks to 25 and report ﬁnal test accuracy over 4 independent runs.
All ablations use the same hyper-parameters, except for online meta-training which uses a learning rate of 0 . 001 .
First, we vary the meta-training protocol by (a) using the approximate objective (Eq. 12), (b) using online meta-training (Algorithm 1), and (c) whether meta-learning the learning rate used for task adaptation is beneﬁcial in this experiment.
We meta-learn a single scalar learning rate (as warp parameters can learn layer-wise scaling).
Meta-gradients for the learning rate are clipped at 0 .and
we use a learning rate of 0 . 001 .
Note that when using ofﬂine meta-training, we store both task parameterisations and the momentum buffer in that phase and use them in the update rule when computing the canonical objective (Eq. 11).
Further, we vary the architecture used for warp-layers.
We study simpler versions that use channelwise scaling and more complex versions that use non-linearities and residual connections.
We also evaluate a version where each warp-layer has two stacked convolutions, where the ﬁrst warp convolution outputs 128 ﬁlters and the second warp convolution outputs 64 ﬁlters.
Finally, in the two-layer warp-architecture, we evaluate a version that inserts a FiLM layer between the two warp convolutions.
These are adapted during task training from a 0 initialisation; they amount to task embeddings that condition gradient warping on task statistics.
Full results are reported in Table 3.
G A BLATION STUDY : W ARP G RAD AND N ATURAL G RADIENT D ESCENT
[Table]
Here, we perform ablation studies to compare the geometry that a WarpGrad optimiser learns to the geometry that
Natural Gradient Descent (NGD) methods represent (approximately).
For consistency, we run the ablation on Omniglot.
As computing the true Fisher Information Matrix is intractable, we can compare WarpGrad against two common block-diagonal approximations,
KFAC (Martens & Grosse, 2015) and Natural Neural Nets (Desjardins et al., 2015).
First, we isolate the effect of warping task loss surfaces by ﬁxing a random initialisation and only meta-learning warp parameters.
That is, in this experiment, we set λC ( θ 0 )
= 0 .
We compare against two baselines, stochastic gradient descent (SGD) and KFAC, both trained from a random initialisation.
We use task mini-batch sizes of 200 and task learning rates of 1.0, otherwise we use
[Table]
[Figure]
x p e c t e
[Figure]
c t i v a t
i
[Figure]
the same hyper-parameters as in the main experiment.
For WarpGrad, we meta-train with these hyper-parameters as well.
We evaluate two WarpGrad architectures, in one, we use linear warp-layers, which gives a block-diagonal preconditioning, as in KFAC.
In the other, we use our most expressive warp conﬁguration from the ablation experiment in appendix F, where warp-layers are two-layer convolutional block with residual connections, batch normalisation, and ReLU activation.
We ﬁnd that warped geometries facilitate task adaptation on held-out tasks to a greater degree than either
SGD or KFAC by a signiﬁcant margin (table 4).
We further ﬁnd that going beyond block-diagonal preconditioning yields a signiﬁcant improvement in performance.
Second, we explore whether the geometry that we meta-learn under in the full Warp-Leap algorithm is approximately Fisher.
In this experiment, we use the main Warp-Leap architecture.
We use a meta-learner trained on 25 tasks and that we evaluate on 10 held-out tasks.
Because warp-layers are linear in this conﬁguration, if the learned geometry is approximately Fisher, post-warp activations should be zero-centred and the layer-wise covariance matrix should satisfy
Cov( ω ( i ) ( h ( i ) ( x )) , ω ( i ) ( h ( i ) ( x )))
=
I , where I is the identity matrix (Desjardins et al., 2015).
If true, Warp-Leap would learn a block-diagonal approximation to the Inverse Fisher Matrix, as Natural
Neural Nets.
To test this, during task adaptation on held-out tasks, we compute the mean activation in each convolutional layer pre- and post-warping.
We also compute the Shatten-1 norm of the difference between layer activation covariance and the identity matrix pre- and post-warping, as described above.
We average statistics over task and adaptation step (we found no signiﬁcant variation in these dimensions).
Figure 9 summarise our results.
We ﬁnd that, in general, WarpGrad-Leap has zero-centered post-warp activations.
That pre-warp activations are positive is an artefact of the ReLU activation function.
However, we ﬁnd that the correlation structure is signiﬁcantly different from what we would expect if
Warp-Leap were to represent the Fisher matrix; post-warp covariances are signiﬁcantly dissimilar from the identity matrix and varies across layers.
These results indicate that WarpGrad methods behave distinctly different from Natural Gradient
Descent methods.
One possibility is that WarpGrad methods do approximate the Fisher Information
Matrix, but with higher accuracy than other methods.
A more likely explanation is that WarpGrad methods encode a different geometry since they can learn to leverage global information beyond the task at hand, which enables them to express geometries that standard Natural Gradient Descent cannot.
H mini I MAGE N ET AND tiered I MAGE N ET
mini ImageNet This dataset is a subset of 100 classes sampled randomly from the 1000 base classes in the ILSVRC-12 training set, with 600 images for each class.
Following (Ravi & Larochelle, 2016), classes are split into non-overlapping meta-training, meta-validation and meta-tests sets with 64, 16, and 20 classes in each respectively.
tiered ImageNet As described in (Ren et al., 2018), this dataset is a subset of ILSVRC-12 that stratiﬁes 608 classes into 34 higher-level categories in the ImageNet human-curated hierarchy (Deng et al., 2009).
In order to increase the separation between meta-train and meta-evaluation splits,of these categories are used for meta-training, while 6 and 8 are used for meta-validation and metatesting respectively.
Slicing the class hierarchy closer to the root creates more similarity within each split, and correspondingly more diversity between splits, rendering the meta-learning problem more challenging.
High-level categories are further divided into 351 classes used for meta-training, 97 for meta-validation and 160 for meta-testing, for a total of 608 base categories.
All the training images in ILSVRC-12 for these base classes are used to generate problem instances for tiered ImageNet, of which there are a minimum of 732 and a maximum of 1300 images per class.
For all experiments, N -way K -shot classiﬁcation problem instances were sampled following the standard image classiﬁcation methodology for meta-learning proposed in Vinyals et al.
A subset of N classes was sampled at random from the corresponding split.
For each class, K arbitrary images were chosen without replacement to form the training dataset of that problem instance.
As usual, a disjoint set of L images per class were selected for the validation set.
Few-shot classiﬁcation In these experiments we used the established experimental protocol for evaluation in meta-validation and meta-testing: 600 task instances were selected, all using N = 5 ,
K = 1 or K = 5 , as speciﬁed, and L = 15 .
During meta-training we used N = 5 , K = 5 or
K = 15 respectively, and L = 15 .
Task-learners used 4 convolutional blocks deﬁned by with 128 ﬁlters (or less, chosen by hyperparameter tuning), 3 × 3 kernels and strides set to 1 , followed by batch normalisation with learned scales and offsets, a ReLU non-linearity and 2 × 2 max-pooling.
The output of the convolutional stack ( 5 × 5 × 128 ) was ﬂattened and mapped, using a linear layer, to the 5 output units.
The last 3 convolutional layers were followed by warp-layers with 128 ﬁlters each.
Only the ﬁnaltask-layer parameters and their corresponding scale and offset batch-norm parameters were adapted during task-training, with the corresponding warp-layers and the initial convolutional layer kept ﬁxed and meta-learned using the WarpGrad objective.
Note that, with the exception of CAVIA, other baselines do worse with 128 ﬁlters as they overﬁt; MAML and T-Nets achieve 46% and% 5-way-1-shot test accuracy with 128 ﬁlters, compared to their best reported results (48.7% and 51.7%, respectively).
Hyper-parameters were tuned independently for each condition using random grid search for highest test accuracy on meta-validation left-out tasks.
Grid sizes were 50 for all experiments.
We choose the optimal hyper-parameters (using early stopping at the meta-level) in terms of meta-validation test set accuracy for each condition and we report test accuracy on the meta-test set of tasks.
60000 metatraining steps were performed using meta-gradients over a single randomly selected task instances and their entire trajectories of 5 adaptation steps.
Task-speciﬁc adaptation was done using stochastic gradient descent without momentum.
We use Adam (Kingma & Ba, 2015) for meta-updates.
Multi-shot classiﬁcation
For these experiments we used N
=
, K = and
L =
Task learners are deﬁned similarly, but stacking 6 convolutional blocks deﬁned by
× kernels and strides set to
, followed by batch normalisation with learned scales and offsets, a ReLU non-linearity and 2 × 2 max-pooling ( ﬁrst
5 layers).
The sizes of convolutional layers were chosen by hyper-parameter tuning to
The output of the convolutional stack
2 × 2 × was ﬂattened and mapped, using a linear layer, to the 10 output units.
Hyper-parameters were tuned independently for each algorithm, version, and baseline using random grid search for highest test accuracy on meta-validation left-out tasks.
Grid sizes were 200 for all multi-shot experiments.
We choose the optimal hyper-parameters in terms of mean meta-validation test set accuracy AUC (using early stopping at the meta-level for each condition and we report test accuracy on the meta-test set of tasks.
2000 meta-training steps were performed using averaged meta-gradients over 5 random task instances and their entire trajectories of 100 adaptation steps with batch size 64 , or inner-loops.
Task-speciﬁc adaptation was done using stochastic gradient descent with momentum ( 0 . 9 ).
Meta-gradients were passed to Adam in the outer loop.
We test WarpGrad against
Leap
Reptile, and training from scratch with large batches and tuned momentum.
We tune all meta-learners for optimal performance on the validation set.
WarpGrad outperforms all baselines both in terms of rate of convergence and ﬁnal test performance (Figure 10).
I M AZE N AVIGATION
To illustrate both how WarpGrad may be used with Recurrent Neural Networks in an online metalearning setting, as well as in a Reinforcement Learning environment, we evaluate it in a maze navigation task proposed by Miconi et al.
The environment is a ﬁxed maze and a task is deﬁned by randomly choosing a goal location in the maze.
During a task episode of length 200 , the goal location is ﬁxed but the agent gets teleported once it ﬁnds it.
Thus, during an episode the agent must ﬁrst locate the goal, then return to it as many times as possible, each time being randomly teleported to a new starting location.
We use an identical setup as Miconi et al.
(2019), except our grid is of size 11 × 11 as opposed to 9 × 9 .
We compare our Warp-RNN to Learning to Reinforcement
Learn (Wang et al., 2016) and Hebbian meta-learners (Miconi et al., 2018; 2019).
The task-learner in all cases is an advantage actor-critic (Wang et al., 2016), where the actor and critic share an underlying basic RNN, whose hidden state is projected into a policy and value function by two separate linear layers.
The RNN has a hidden state size of 100 and tanh non-linearities.
Following (Miconi et al., 2019), for all benchmarks, we train the task-learner using Adam with a learning rate of 1 e − 3 for 200 000 steps using batches of 30 episodes, each of length 200 .
Metalearning arises in this setting as each episode encodes a different task, as the goal location moves, and by learning across episodes the RNN is encoding meta-information in its parameters that it can leverage during task adaptation (via its hidden state (Hochreiter & Schmidhuber, 1997; Wang et al., 2016)).
See Miconi et al.
for further details.
We design a Warp-RNN by introducing a warp-layer in the form of an LSTM that is frozen for most of the training process.
Following Flennerhag et al.
(2018), we use this meta-LSTM to modulate the task RNN.
Given an episode with input vector x t , the task RNN is deﬁned by
[Equation]
where W, V, b are task-adaptable parameters; each U ( i ) j,t is a diagonal warp matrix produced by projecting from the hidden state of the meta-LSTM, U ( i ) j,t = diag(tanh( P ( i ) j z t )) , where z is the hidden-state of the meta-LSTM.
See Flennerhag et al.
(2018) for details.
Thus, our Warp-RNN is a form of HyperNetwork (see Figure 6, Appendix A).
Because the meta-LSTM is frozen for most of the training process, task adaptable parameters correspond to those of the baseline RNN.
To control for the capacity of the meta-LSTM, we also train a HyperRNN where the LSTM is updated with every task adaptation; we ﬁnd this model does worse than the WarpGrad-RNN.
We also compare the non-linear preconditioning that we obtain in our Warp-RNN to linear forms of preconditioning deﬁned in prior works.
We implement a T-Nets-RNN meta-learner, deﬁned by embedding linear projections T h , T x and T b that are meta-learned in the task RNN, h t = tanh( T h V
h t
+ T x
Wx t + b ) .
Note that we cannot backpropagate to these meta-parameters as per the T-Nets (MAML) framework.
Instead, we train T h , T x , T b with the meta-objective and meta-training algorithm we use for the
Warp-RNN.
The T-Nets-RNN does worse than the baseline RNN and generally fails to learn.
We meta-train the Warp-RNN using the continual meta-training algorithm (Algorithm 3, see Appendix B for details), which accumulates meta-gradients continuously during training.
Because task training is a continuous stream of batches of episodes, we accumulating the meta-gradient using the approximate objective (Eq. 12, where L τ task and L τ meta are both the same advantage actor-critic objective) and update warp-parameters on every 30th task parameter update.
We detail the meta-objective in Appendix C (see Eq. 18).
Our implementation of a Warp-RNN can be seen as meta-learning “slow” weights to facilitate learning of “fast” weights (Schmidhuber, 1992; Mujika et al., 2017).
Implementing Warp-RNN requires four lines of code on top of the standard training script.
The task-learner is the same in all experiments with the same number of learnable parameters and hidden state size.
Compared to all baselines, we ﬁnd that the Warp-RNN converges faster and achieves a higher cumulative reward (Figure 4 and Figure 11).
J M ETA -L EARNING FOR C ONTINUAL L EARNING
Online SGD and related optimisation methods tend to adapt neural network models to the data distribution encountered last during training, usually leading to what has been termed “catastrophic
[Figure]
forgetting” (French, 1999).
In this experiment, we investigate whether WarpGrad optimisers can meta-learn to avoid this problem altogether and directly minimise the joint objective over all tasks with every update in the fully online learning setting where no past data is retained.
Continual Sine Regression We propose a continual learning version of the sine regression metalearning experiment in Finn et al.
We split the input interval [ − 5 , 5] ⊂ R evenly intoconsecutive sub-intervals, corresponding to 5 regression tasks.
These are presented one at a time to a task-learner, which adapts to each sub-task using 20 gradient steps on data from the given sub-task only.
Batch sizes were set to 5 samples.
Sub-tasks thus differ in their input domain.
A task sequence is deﬁned by a target function composed of two randomly mixed sine functions of the form f a i ,b i ( x )
=
a i sin( x − b i )
each with randomly sampled amplitudes a i ∈ [0 . 1 , 5] and phases b i ∈ [0 , π ] .
A task τ = ( a 1 , b 1 , a 2 , b 2 , o ) is therefore deﬁned by sampling the parameters that specify this mixture; a task speciﬁes a target function g τ by
[Equation]
[Equation]
where α o ( x ) = σ ( x + o ) for a randomly sampled offset o ∈ [ − 5 , 5] , with σ being the sigmoid activation function.
Model We deﬁne a task-learner as 4 -layer feed-forward networks with hidden layer size 200 and
ReLU non-linearities to learn the mapping between inputs and regression targets, f ( · , θ, φ ) .
For each task sequence τ , a task-learner is initialised from a ﬁxed random initialisation θ 0 (that is not meta-learned).
Each non-linearity is followed by a residual warping block consisting of 2 -layer feedforward networks with 100 hidden units and tanh non-linearities, with meta-learned parameters φ which are ﬁxed during the task adaptation process.
Continual learning as task adaptation The task target function g τ is partitioned into 5 sets of sub-tasks.
The task-learner sees one partition at a time and is given n = 20 gradient steps to adapt, for a total of K = 100 steps of online gradient descent updates for the full task sequence; recall that every such sequence starts from a ﬁxed random initialisation θ 0 .
The adaptation is completely online since at step k = 1 , . .
, K we sample a new mini-batch D k task of 5 samples from a single sub-task (sub-interval).
The data distribution changes after each n = 20 steps with inputs x coming from the next sub-interval and targets form the same function g τ ( x ) .
During meta-training we always present tasks in the same order, presenting intervals from left to right.
The online (sub-)task loss is deﬁned on the current mini-batch D k task at step k :
[Equation]
[Equation]
Adaptation to each sub-task uses sub-task data only to form task parameter updates θ τk +1 ← θ τk − α ∇L τ task (cid:0) θ τk , D k task ; φ (cid:1) .
We used a constant learning rate α = 0 . 001 .
Warp-parameters φ are ﬁxed across the full task sequence during adaptation and are meta-learned across random samples of task sequences, which we describe next.
Meta-learning an optimiser for continual learning To investigate the ability of WarpGrad to learn an optimiser for continual learning that mitigates catastrophic forgetting, we ﬁx a random initialisation prior to meta-training that is not meta-learned; every task-learner is initialised with these parameters.
To meta-learn an optimiser for continual learning, we need a meta-objective that encourages such behaviour.
Here, we take a ﬁrst step towards a framework for meta-learned continual learning.
We deﬁne the meta-objective L τ meta as an incremental multitask objective that, for each sub-task!
τ t in a given task sequence τ , averages the validation sub-task losses (Eq. 22) for the current and every preceding loss in the task sequence.
The task meta-objective is deﬁned by summing over all sub-tasks in the task sequence.
For some sub-task parameterisation θ τ t , we have
[Equation]
As before, the full meta-objective is an expectation over the joint task parameter distribution (Eq. 11); for further details on the meta-objective, see Appendix C, Eq. 19.
This meta-objective gives equal weight to all the tasks in the sequence by averaging the regression step loss over all sub-tasks where a prior sub-task should be learned or remembered.
For example, losses from the ﬁrst sub-task, deﬁned using the interval [ − 5 , − 3] , will appear nT times in the meta-objective.
Conversely, the last sub-task in a sequence, deﬁned on the interval [3 , 5] , is learned only in the last n = 20 steps of task adaptation, and hence appears n times in the meta-objective.
Normalising on number of appearances corrects for this bias.
We trained warp-parameters using Adam and a meta-learning rate of 0 .
001 , samplingrandom tasks to form a meta-batch and repeating the process for 20 000 steps of meta-training.
Results Figure 12 shows a breakdown of the validation loss across the 5 sequentially learned tasks over the 100 steps of online learning during task adaptation.
Results are averaged over 100 random regression problem instances.
The meta-learned WarpGrad optimiser reduces the loss of the task currently being learned in each interval while also largely retaining performance on previous tasks.
There is an immediate relatively minor loss of performance, after which performance on previous tasks is retained.
We hypothesise that this is because the meta-objectives averages over the full learning curve, as opposed to only the performance once a task has been adapted to.
As such, the
WarpGrad optimiser may allow for some degree of performance loss.
Intriguingly, in all cases, after an initial spike in previous sub-task losses when switching to a new task, the spike starts to revert back some way towards optimal performance, suggesting that the WarpGrad optimiser facilitates positive backward transfer, without this being explicitly enforced in the meta-objective.
Deriving a principled meta-objective for continual learning is an exciting area for future research.
[Figure] [Figure]