Continuous-Time Meta-Learning with Forward Mode Differentiation

1 INTRODUCTION . Among the existing meta-learning algorithms , gradient-based methods as popularized by ModelAgnostic Meta-Learning ( MAML , Finn et al. , 2017 ) have received a lot of attention over the past few years . They formulate the problem of learning a new task as an inner optimization problem , typically based on a few steps of gradient descent . An outer meta-optimization problem is then responsible for updating the meta-parameters of this learning process , such as the initialization of the gradient descent procedure . However since the updates at the outer level typically require backpropagating through the learning process , this class of methods has often been limited to only a few gradient steps of adaptation , due to memory constraints . Although solutions have been proposed to alleviate the memory requirements of these algorithms , including checkpointing ( Baranchuk , 2019 ) , using implicit differentiation ( Rajeswaran et al. , 2019 ) , or reformulating the meta-learning objective ( Flennerhag et al. , 2018 ) , they are generally either more computationally demanding , or only approximate the gradients of the meta-learning objective ( Nichol et al. , 2018 ; Flennerhag et al. , 2020 ) . In this work , we propose a continuous-time formulation of gradient-based meta-learning , called Continuous-Time Meta-Learning ( COMLN ) , where the adaptation is the solution of a differential equation ( see Figure 1 ) . Moving to continuous time allows us to devise a novel algorithm , based on forward mode differentiation , to efficiently compute the exact gradients for meta-optimization , no matter how long the adaptation to a new task might be . We show that using forward mode differentiation leads to a stable algorithm , unlike the counterpart of backpropagation in continuous time called the adjoint method ( frequently used in the Neural ODE literature ) which tends to be unstable with gradient vector fields . Moreover as the length of the adaptation trajectory is a continuous quantity , as opposed to a discrete number of gradient steps fixed ahead of time , we can treat the amount of adaptation in COMLN as a meta-parameter—on par with the initialization—which we can meta-optimize using stochastic gradient descent . We verify empirically that our method is both computationally and memory efficient , and we show that COMLN outperforms other standard meta-learning algorithms on few-shot image classification datasets . 2 BACKGROUND . In this work , we consider the problem of few-shot classification , that is the problem of learning a classification model with only a small number of training examples . More precisely for a classification task τ , we assume that we have access to a ( small ) training dataset Dtrainτ = { ( xm , ym ) } Mm=1 to fit a model on task τ , and a distinct test dataset Dtestτ to evaluate how well this adapted model generalizes on that task . In the few-shot learning literature , it is standard to consider the problem of k-shot N -way classification , meaning that the model has to classify among N possible classes , and there are only k examples of each class in Dtrainτ , so that overall the number of training examples is M = kN . We use the convention that the target labels ym ∈ { 0 , 1 } N are one-hot vectors . 2.1 GRADIENT-BASED META-LEARNING . Gradient-based meta-learning methods aim to learn an initialization such that the model is able to adapt to a new task via gradient descent . Such methods are often cast as a bi-level optimization process : adapting the task-specific parameters θ in the inner loop , and training the ( task-agnostic ) meta-parameters Φ and initialization θ0 in the outer loop . The meta-learning objective is : min θ0 , Φ Eτ [ L ( θτT , Φ ; Dtestτ ) ] ( 1 ) s.t . θτt+1 = θ τ t − α∇θL ( θτt , Φ ; Dtrainτ ) θτ0 = θ0 ∀τ ∼ p ( τ ) , ( 2 ) where T is the number of inner loop updates . For example , in the case of MAML ( Finn et al. , 2017 ) , there is no additional meta-parameter other than the initialization ( Φ ≡ ∅ ) ; in ANIL ( Raghu et al. , 2019 ) , θ are the parameters of the last layer , and Φ are the parameters of the shared embedding network ; in CAVIA ( Zintgraf et al. , 2019 ) , θ are referred to as context parameters . During meta-training , the model is trained over many tasks τ . The task-specific parameters θ are learned via gradient descent on Dtrainτ . The meta-parameters are then updated by evaluating the error of the trained model on the test dataset Dtestτ . At meta-testing time , the meta-trained model is adapted on Dtrainτ —i.e . applying Equation 2 with the learned meta-parameters θ0 and Φ . 2.2 LOCAL SENSITIVITY ANALYSIS OF ORDINARY DIFFERENTIAL EQUATIONS . Consider the following ( autonomous ) Ordinary Differential Equation ( ODE ) : dz dt = g ( z ( t ) ; θ ) z ( 0 ) = z0 ( θ ) , ( 3 ) where the dynamics g and initial value z0 may depend on some external parameters θ , and integration is carried out from 0 to some time T . Local sensitivity analysis is the study of how the solution of this dynamical system responds to local changes in θ ; this effectively corresponds to calculating the derivative dz ( t ) /dθ . We present here two methods to compute this derivative , with a special focus on their memory efficiency . Adjoint sensitivity method Based on the adjoint state ( Pontryagin , 2018 ) , and taking its root in control theory ( Lions & Magenes , 2012 ) , the adjoint sensitivity method ( Bryson & Ho , 1969 ; Chavent et al. , 1974 ) provides an efficient approach for evaluating derivatives of L ( z ( T ) ; θ ) , a function of z ( T ) the solution of the ODE in Equation 3 . This method , popularized lately by the literature on Neural ODEs ( Chen et al. , 2018 ) , requires the integration of the adjoint equation da dt = −a ( t ) ∂g ( z ( t ) ; θ ) ∂z ( t ) a ( T ) = dL ( z ( T ) ; θ ) dz ( T ) , ( 4 ) backward in time . The adjoint sensitivity method can be viewed as a continuous-time counterpart to backpropagation , where the forward pass would correspond to integrating Equation 3 forward in time from 0 to T , and the backward pass to integrating Equation 4 backward in time from T to 0 . One possible implementation , reminiscent of backpropagation through time ( BPTT ) , is to store the intermediate values of z ( t ) during the forward pass , and reuse them to compute the adjoint state during the backward pass . While several sophisticated checkpointing schemes have been proposed ( Serban & Hindmarsh , 2003 ; Gholami et al. , 2019 ) , with different compute/memory trade-offs , the memory requirements of this approach typically grow with T ; this is similar to the memory limitations standard gradient-based meta-learning methods suffer from as the number of gradient steps increases . An alternative is to augment the adjoint state a ( t ) with the original state z ( t ) , and to solve this augmented dynamical system backward in time ( Chen et al. , 2018 ) . This has the notable advantage that the memory requirements are now independent of T , since z ( t ) are no longer stored during the forward pass , but they are recomputed on the fly during the backward pass . Forward sensitivity method While the adjoint method is related to reverse-mode automatic differentiation ( backpropagation ) , the forward sensitivity method ( Feehery et al. , 1997 ; Leis & Kramer , 1988 ; Maly & Petzold , 1996 ; Caracotsios & Stewart , 1985 ) , on the other hand , can be viewed as the continuous-time counterpart to forward ( tangent-linear ) mode differentiation ( Griewank & Walther , 2008 ) . This method is based on the fact that the derivative S ( t ) , dz ( t ) /dθ is the solution of the so-called forward sensitivity equation dS dt = ∂g ( z ( t ) ; θ ) ∂z ( t ) S ( t ) + ∂g ( z ( t ) ; θ ) ∂θ S ( 0 ) = ∂z0 ∂θ . ( 5 ) This equation can be found throughout the literature in optimal control and system identification ( Betts , 2010 ; Biegler , 2010 ) . Unlike the adjoint method , which requires an explicit forward and backward pass , the forward sensitivity method only requires the integration forward in time of the original ODE in Equation 3 , augmented by the sensitivity state S ( t ) with the dynamics above . The memory requirements of the forward sensitivity method do not scale with T either , but it now requires storing S ( t ) , which may be very large ; we will come back to this problem in Section 3.2 . We will simply note here that in discrete-time , this is the same issue afflicting forward-mode training of RNNs with real-time recurrent learning ( RTRL ; Williams & Zipser , 1989 ) , or other meta-learning algorithms ( Sutton , 1992 ; Franceschi et al. , 2017 ; Xu et al. , 2018 ) . 3 CONTINUOUS-TIME ADAPTATION . In the limit of infinitely small steps , some optimization algorithms can be viewed as the solution trajectory of a differential equation . This point of view has often been taken to analyze their behavior ( Platt & Barr , 1988 ; Wilson et al. , 2016 ; Su et al. , 2014 ; Orvieto & Lucchi , 2019 ) . In fact , some optimization algorithms such as gradient descent with momentum have even been introduced initially from the perspective of dynamical systems ( Polyak , 1964 ) . As the simplest example , gradient descent with a constant step size α→ 0+ ( i.e . α tends to 0 by positive values ) corresponds to following the dynamics of an autonomous ODE called a gradient vector field zt+1 = zt − α∇f ( zt ) −→ α→0+ dz dt = −∇f ( z ( t ) ) , ( 6 ) where the iterate z ( t ) is now a continuous function of time t. The solution of this dynamical system is uniquely defined by the choice of the initial condition z ( 0 ) = z0 . W ( 0 ) W ( T ) W * Forward Backward 3.1 CONTINUOUS-TIME META-LEARNING . In gradient-based meta-learning , the task-specific adaptation with gradient descent may also be replaced by a gradient vector field in the limit of infinitely small steps . Inspired by prior work in meta-learning ( Raghu et al. , 2019 ; Javed & White , 2019 ) , we assume that an embedding network fΦ with meta-parameters Φ is shared across tasks , and only the parametersW of a task-specific linear classifier are adapted , starting at some initializationW0 . Instead of being the result of a few steps of gradient descent though , the final parametersW ( T ) now correspond to integrating an ODE similar to Equation 6 up to a certain horizon T , with the initial conditionsW ( 0 ) = W0 . We call this new meta-learning algorithm Continuous-Time Meta-Learning1 ( COMLN ) . Treating the learning algorithm as a continuous-time process has the notable advantage that the adapted parameterW ( T ) is now differentiable wrt . the time horizon T ( Wiggins , 2003 , Chap . 7 ) , in addition to being differentiable wrt . the initial conditionsW0—which plays a central role in gradientbased meta-learning , as described in Section 2.1 . This allows us to view T as a meta-parameter on par with Φ andW0 , and to effectively optimize the amount of adaptation using stochastic gradient descent ( SGD ) . The meta-learning objective of COMLN can be written as min Φ , W0 , T Eτ [ L ( Wτ ( T ) ; fΦ ( Dtestτ ) ) ] ( 7 ) s.t . dWτ dt = −∇L ( Wτ ( t ) ; fΦ ( Dtrainτ ) ) Wτ ( 0 ) = W0 ∀τ ∼ p ( τ ) , ( 8 ) where fΦ ( Dtrainτ ) = { ( fΦ ( xm ) , ym ) | ( xm , ym ) ∈ Dtrainτ } is the embedded training dataset , and fΦ ( Dtestτ ) is defined similarly for Dtestτ . In practice , adaptation is implemented using a numerical integration scheme based on an iterative discretization of the problem , such as Runge-Kutta methods . Although a complete discussion of numerical solvers is outside of the scope of this paper , we recommend ( Butcher , 2008 ) for a comprehensive overview of numerical methods for solving ODEs .

Summary: The authors propose COMLN -- a gradient-based meta-learning algorithm with continuous-time inner loop adaptation. The authors discusses possible options for computing the meta-gradient (backpropagation through time, solve ODE backwards in time, and forward-mode differentiation), and identifies the forward-mode differentiation as the only method that is both numerically stable and has a memory cost that is constant w.r.t. the inner-loop length. The authors then derive a way to compute the vector-Jacobian products without explicitly constructing the Jacobians in forward-mode differentiation (specifically for the few-shot learning setting), which further reduces the memory cost. Also, by using a continuous formulation for the inner-optimization, the authors are able to compute meta-gradient w.r.t. horizon T and adapt that as well. Finally, the authors conduct experiments on standard few-shot image classification benchmarks, and show that COMLN outperforms other gradient-based meta-optimization methods in accuracy with a lot less memory cost.