Meta-Learning with Warped Gradient Descent

1 INTRODUCTION . 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 defines 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 beneficial 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 . 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 defining 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 defined point-wise in parameter space and can be seen as a recurrent operator of order 1 . We use this insight to define a trajectory agnostic meta-objective over a joint parameter search space where knowledge transfer is encoded in gradient preconditioning . To achieve a scalable and flexible 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 ( miniImageNet , tieredImageNet ; 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 tieredImageNet . We further find 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 . 2 WARPED GRADIENT DESCENT . 2.1 GRADIENT-BASED META-LEARNING . WarpGrad belongs to the family of optimisation-based meta-learners that parameterise an update rule θ ← U ( θ ; ξ ) with some meta-parameters ξ . Specifically , gradient-based meta-learners define an update rule by relying on the gradient descent , U ( θ ; ξ ) : = θ − α∇L ( θ ) for some objective L and learning rate α . A task is defined by a training set Dτtrain and a test set D τ test , which defines learning objectives LDτ ( θ ) : = E ( x , y ) ∼Dτ [ ` ( f ( x , θ ) , y ) ] over the task-learner f for some loss ` . MAML ( Finn 1Open-source implementation available at https : //github.com/flennerhag/warpgrad . et al. , 2017 ) meta-learns a shared initialisation θ0 by backpropagating through K steps of gradient descent across a given task distribution p ( τ ) , CMAML ( ξ ) : = ∑ τ∼p ( τ ) LDτtest ( θ0 − α K−1∑ k=0 UDτtrain ( θ τ k ; ξ ) ) . ( 1 ) 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 ) defines a block-diagonal preconditioning matrixB , and T-Nets ( Lee & Choi , 2018 ) embed block-diagonal preconditioning in feed-forward learners via linear projections , U ( θk ; θ0 ) : = θk − α∇L ( θk ) MAML ( 2 ) U ( θk ; θ0 , φ ) : = θk − α diag ( φ ) ∇L ( θk ) MSGD ( 3 ) U ( θk ; θ0 , φ ) : = θk − αB ( θk ; φ ) ∇L ( θk ) MC ( 4 ) U ( θk ; θ0 , φ ) : = θk − α∇L ( θk ; φ ) T-Nets . ( 5 ) 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 first 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.2 GENERAL-PURPOSE PRECONDITIONING . A preconditioned gradient descent rule , U ( θ ; φ ) : = θ − αP ( θ ; φ ) ∇L ( θ ) , defines 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 sufficient 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 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 beneficial ( 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 : ∂ L ∂θ ( i ) = E ∇ ` T L− ( i+1 ) ∏ j=0 Dxω ( L−j ) Dxh ( L−j ) Dxω ( i ) Dθh ( i ) , ( 6 ) where Dx 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 Dxω = 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 defining warp-layers as Residual Networks ( He et al. , 2017 ) improves performance on classification tasks ( Section 4.1 ) . We also introduce recurrent warp-layers for agents in a gradient-based meta-learner that is the first , 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 : R2 → 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 τ ′ ) that are smoother and more well-behaved than their respective native loss-surfaces . 2.3 THE GEOMETRY OF WARPED GRADIENT DESCENT . If the preconditioning matrix P is invertible , it defines 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 sufficient to require T to be full rank , since T explicitly defines P as a block-diagonal matrix with block entries TTT . In contrast , non-linearity in warp-layers precludes such an explicit identification . Instead , we must consider the geometry that warp-layers represent . For this , we need a metric tensor , G , which is a positive-definite , smoothly varying matrix that measures curvature on a manifoldW . The metric tensor defines 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 manifoldW with γ = Ω ( θ ; φ ) . We induce a metric G onW by push-forward ( Figure 2 ) : ∆θ : = ∇ ( L◦Ω ) ( θ ; φ ) = [ DxΩ ( θ ; φ ) ] T∇L ( γ ) P -space ( 7 ) ∆γ : = DxΩ ( θ ; φ ) ∆θ = G ( γ ; φ ) −1∇L ( γ ) W -space , ( 8 ) where G−1 : = [ DxΩ ] [ DxΩ ] T . Provided Ω is not degenerate ( G is non-singular ) , G−1 is positivedefinite , hence a valid Riemann metric . While this is the metric induced onW 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 first-order Taylor series expansion to establish first-order equivalence between the WarpGrad update in P ( Eq . 7 ) and the ideal update inW ( Eq . 8 ) , ( L◦Ω ) ( θ − α∆θ ) = L ( γ − α∆γ ) +O ( α2 ) . ( 9 ) Consequently , gradient descent under warp-layers ( in P-space ) is first-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 ( inW-space ) should yield preconditioning that points in the direction of steepest descent , accounting for global information across tasks , min φ EL , γ∼p ( L , γ ) [ L ( γ − αG ( γ ; φ ) −1∇L ( γ ) ) ] . ( 10 ) In contrast to MAML-based approaches ( Eq . 1 ) , this objective avoids backpropagation through learning processes . Instead , it defines task learning abstractly by introducing a joint distribution over objectives and parameterisations , opening up for general-purpose meta-learning at scale .

Summary: The authors propose warped gradient descent (WarpGrad) an optimisation framework for facilitating gradient-based meta-learning. WarpGrad interleaves within the learner meta-learned warp-layers that implicitly precondition the gradients of the task-specific parameters during backpropagation. In contrast to the linear projection layers employed in T-Nets, warp-layers are unrestricted in form and induce a full Jacobian preconditioning matrix. The warp layers are meta-learned in a trajectory-agnostic fashion, thus obviating the need to backpropagate through the gradient steps to compute the updates of their parameters. The framework is readily applicable to standard gradient-based meta-learners, and is shown to yield a significant boost in performance on both few-shot and multi-shot learning tasks, as well as to have promising applications to continual learning.