Meta-Learning with Neural Tangent Kernels

1 INTRODUCTION . Meta-learning ( Schmidhuber , 1987 ) has made tremendous progresses in the last few years . It aims to learn abstract knowledge from many related tasks so that fast adaption to new and unseen tasks becomes possible . For example , in few-shot learning , meta-learning corresponds to learning a meta-model or meta-parameters so that they can fast adapt to new tasks with a limited number of data samples . Among all existing meta-learning methods , Model Agnostic Meta-Learning ( MAML ) ( Finn et al. , 2017 ) is perhaps one of the most popular and flexible ones , with a number of follow-up works such as ( Nichol et al. , 2018 ; Finn et al. , 2018 ; Yao et al. , 2019 ; Khodak et al. , 2019a ; b ; Denevi et al. , 2019 ; Fallah et al. , 2020 ; Lee et al. , 2020 ; Tripuraneni et al. , 2020 ) . MAML adopts a double-looped optimization framework , where adaptation is achieved by one or several gradientdescent steps in the inner-loop optimization . Such a framework could lead to some undesirable issues related to computational inefficiency and sub-optimal solutions . The main reasons are that 1 ) it is computationally expensive to back-propagate through a stochastic-gradient-descent chain , and 2 ) it is hard to tune the number of adaptation steps in the inner-loop as it can be different for both training and testing . Several previous works tried to address these issues , but they can only alleviate them to certain extents . For example , first order MAML ( FOMAML ) ( Finn et al. , 2017 ) ignores the high-order terms of the standard MAML , which can speed up the training but may lead to deteriorated performance ; MAML with Implicit Gradient ( iMAML ) ( Rajeswaran et al. , 2019 ) directly minimizes the objective of the outer-loop without performing the inner-loop optimization . But it still needs an iterative solver to estimate the meta-gradient . To better address these issues , we propose two algorithms that generalize meta-learning to the Reproducing Kernel Hilbert Space ( RKHS ) induced by the meta-model ’ s Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ) . In this RKHS , instead of using parameter adaptation , we propose to perform an implicit function adaptation . To this end , we introduce two algorithms to avoid explicit ∗The first two authors contribute equally . Correspondence to Changyou Chen ( changyou @ buffalo.edu ) . †The research of the first and fifth authors was supported in part by NSF through grants CCF-1716400 and IIS-1910492 . function adaptation : one replaces the function adaptation step in the inner-loop with a new metaobjective with a fast-adaptive regularizer inspired by MAML ; the other solves the adaptation problem analytically based on tools from NTK so that the meta-objective can be directly evaluated on samples in a closed-form . When restricting the function space to be RKHS , the solutions to the proposed two algorithms become conveniently solvable . In addition , we provide theoretical analysis on our proposed algorithms in the cases of using fully-connected neural networks and convolutional neural networks as the meta-model . Our analysis shows close connections between our methods and the existing ones . Particularly , we prove that one of our algorithms is closely related to MAML with some high-order terms ignored in the meta-objective function , thus endowing effective optimization . In summary , our main contributions are : • We re-analyze the meta-learning problem and introduce two new algorithms for metalearning in RKHS . Different from all existing meta-learning algorithms , our proposed methods can be solved efficiently without cumbersome chain-based adaptations . • We conduct theoretically analysis on the proposed algorithms , which suggests that our pro- posed algorithms are closely related to the existing MAML methods when fully-connected neural networks and convolutional neural networks are used as the meta-model . • We conduct extensive experiments to validate our algorithms . Experimental results indicate the effectiveness of our proposed methods , through standard few-shot learning , robustness to adversarial attacks and out-of-distribution adaptation . 2 PRELIMINARIES . 2.1 META-LEARNING . Meta-learning can be roughly categorized as black-box adaptation methods ( Andrychowicz et al. , 2016 ; Graves et al. , 2014 ; Mishra et al. , 2018 ) , optimization-based methods ( Finn et al. , 2017 ) , non-parametric methods ( Vinyals et al. , 2016 ; Snell et al. , 2017 ; Triantafillou et al. , 2020 ) and Bayesian meta-learning methods ( Finn et al. , 2018 ; Yoon et al. , 2018 ; Ravi & Beatson , 2019 ) . In this paper , we focus on the framework of Model Agnostic Meta-Learning ( MAML ) ( Finn et al. , 2017 ) , which has two key components , meta initialization and fast adaptation . Specifically , MAML solves the meta-learning problem through a double-looped optimization procedure . In the inner-loop , MAML runs a task-specific adaptation procedure to transform a meta-parameter , θ , to a task-specific parameter , { φm } Bm=1 , for a total of B different tasks . In the outer-loop , MAML minimizes a total loss of ∑B m=1 L ( fφm ) with respect to meta-parameter θ , where fφm is the model adapted on task m that is typically represented by a deep neural network . It is worth noting that in MAML , one potential problem is to compute the meta-gradient∇θ ∑B m=1 L ( fφm ) . It requires one to differentiate through the whole inner-loop optimization path , which could be very inefficient . 2.2 GRADIENT FLOW . Our proposed method relies on the concept of gradient flow . Generally speaking , gradient flow is a continuous-time version of gradient descent . In the finite-dimensional parameter space , a gradient flow is defined by an ordinary differential equation ( ODE ) , dθt/dt = −∇θtF ( θt ) , with a starting point θ0 and function F : Rd → R. Gradient flow is also known as steepest descent curve . One can generalize gradient flows to infinite-dimensional function spaces . Specifically , given a function spaceH , a functional F : H → R , and a starting point f0 ∈ H , a gradient flow is similarly defined as the solution of df t/dt = −∇ftF ( f t ) . This is a curve in the function space H. In this paper , we use notation∇ftF ( f t ) , instead of∇HF ( f t ) , to denote the general function derivative of the energy functional F with respect to function f t ( Villani , 2008 ) . 2.3 THE NEURAL TANGENT KERNEL . Neural Tangent Kernel ( NTK ) is a recently proposed technique for characterizing the dynamics of a neural network under gradient descent ( Jacot et al. , 2018 ; Arora et al. , 2019 ; Lee et al. , 2019 ) . NTK allows one to analyze deep neural networks ( DNNs ) in RKHS induced by NTK . One immediate benefit of this is that the loss functional in the function space is often convex , even when it is highly non-convex in the parameter space ( Jacot et al. , 2018 ) ∗ . This property allows one to better understand the property of DNNs . Specifically , let fθ be a DNN parameterized by θ . The corresponding NTK Θ ∗LetH be the function space , F be the realization function for neural network defined in Section 3.2 . Note even if a functional loss ( e.g. , L2 loss ) E : H → R is convex onH , the composition E ◦ F is in general not . is defined as : Θ ( x1 , x2 ) = ∂fθ ( x1 ) ∂θ ∂fθ ( x2 ) ∂θ ᵀ , where x1 , x2 are two data points . In our paper , we will define meta-learning on an RKHS induced by such a kernel . 3 META-LEARNING IN RKHS . We first define the meta-learning problem in a general function space , and then restrict the function space to be an RKHS , where two frameworks will be proposed to make meta-learning feasible in RKHS , along with some theoretical analysis . For simplicity , in the following we will hide the superscript time t unless necessary , e.g. , when the analysis involves time-changing . 3.1 META-LEARNING IN FUNCTION SPACE . Given a function space H , a distribution of tasks P ( T ) , and a loss function L , the goal of metalearning is to find a meta function f∗ ∈ H , so that it performs well after simple adaptation on a specific task . Let Dtrm and Dtestm be the training and testing sets , respectively , sampled from a data distribution of task Tm . The meta-learning problem on function spaceH is defined as : f∗ = arg min f∈H E ( f ) , with E ( f ) = ETm [ L ( Adapt ( f , Dtrm ) , Dtestm ) ] ( 1 ) where Adapt denotes some adaptation algorithms , e.g. , several steps of gradient descent ; E : H → R is called energy functional , which is used to evaluate the model represented by the function f . In theory , solving equation 1 is equivalent to solving the gradient flow equation df t/dt = −∇ftE ( f t ) . However , solving the gradient flow equation is generally infeasible , since i ) it is hard to directly apply optimization methods in function space and ii ) the energy functional E contains an adaptation algorithm Adapt , making the functional gradient infeasible . Thus , a better way is to design a special energy functional so that it can be directly optimized without running the specific adaptation algorithm . In the following , we first specify the functional meta-learning problem in RKHS , and then propose two methods to derive efficient solutions for the problem . 3.2 META-LEARNING IN RKHS . We consider a function f that is parameterized by θ ∈ RP , denoted as fθ , with P being the number of parameters . Define a realization function F : RP → H that maps parameters to a function . With these , we can then define an energy function in the parameter space as E , E ◦ F : RP → R with ◦ being the composition operator . Consequently , with an initialized θ0 , we can define the gradient flow of E ( θt ) in parameter space as : dθt/dt = −∇θtE ( θt ) . In the following , we first establish an equivalence between the gradient flow in RKHS and the gradient flow in the parameter space . We then propose two algorithms for meta-learning in the RKHS induced by NTK . Theorem 1 Let H be an RKHS induced by the NTK Θ of fθ . With f0 = fθ0 , the gradient flow of E ( f t ) coincides with the function evolution of fθt driven by the gradient flow of E ( θt ) . The proof of Theorem 1 relies on the property of NTK ( Jacot et al. , 2018 ) , and is provided in the Appendix . Theorem 1 serves as a foundation of our proposed methods , which indicates that solving the meta-learning problem in RKHS can be done by some appropriate manipulations . In the following , we describe two different approaches termed Meta-RKHS-I and Meta-RKHS-II , respectively . 3.3 META-RKHS-I : META-LEARNING IN RKHS WITHOUT ADAPTATION . Our goal is to design an energy functional that has no adaptation component , but is capable of achieving fast adaptation . For this purpose , we first introduce two definitions : empirical loss function L ( fθ , Dm ) and expected loss function L ( fθ ) . Let Dm = { xm , i , ym , i } ni=1 be a set containing the data of a regression task Tm . The empirical loss function L ( fθ , Dm ) and the expected loss function Lm ( fθ ) can be defined as : L ( fθ , Dm ) = 1 2n n∑ i=1 ∥∥f ( xm , i ) − ym , i ∥∥2 , Lm ( fθ ) = Exm , ym [ 12∥∥f ( xm ) − ym ∥∥2 ] . Our idea is to define a regularized functional such that it endows the ability of fast adaptation in RKHS . Our solution is based on some property of the standard MAML . We start from analyzing the meta-objective of MAML with a k-step gradient-descent adaptation , i.e. , applying k gradient-descent steps in the inner-loop . The objective can be formulated as θ∗ = arg min θ ETm [ L ( fφ , Dtestm ) ] , with φ = θ − α k−1∑ i=0 ∇θiL ( fθi , Dtrm ) , where α is the learning rate of the inner-loop , θ0 = θ , and θi+1 = θi − α∇θiL ( fθi , Dtrm ) † . By Taylor expansion , we have ETm [ L ( fφ , Dtestm ) ] ≈ ETm [ L ( fθ , Dtestm ) − α k−1∑ i=0 ∇θiL ( fθi , Dtrm ) ∇θL ( fθ , Dtestm ) ᵀ ] . ( 2 ) Since Dtrm and Dtestm come from the same distribution , equation 2 is an unbiased estimator of Mk = ETm [ Lm ( fθ ) − k−1∑ i=0 βi ] , where βi = α∇θiLm ( fθi ) ∇θLm ( fθ ) ᵀ . ( 3 ) We focus on the case of k = 1 , which is M1 = ETm [ Lm ( fθ ) ] − αETm [ ‖∇θLm ( fθ ) ‖2 ] . The first term on the RHS is the traditional multi-task loss evaluated at θ for all tasks . The second term corresponds to the negative gradient norm ; minimizing it means choosing a θ with the maximum gradient norm . Intuitively , when θ is not a stationary point of a task , one should choose the steepest descent direction to reduce the loss maximally for a specific task , thus leading to fast adaptation . The above understanding suggests us to propose the following regularized energy functional , Ẽα , for meta-learning in the RKHS induced with the NTK for fast function adaptation : Ẽ ( α , fθ ) = ETm [ Lm ( fθ ) − α‖∇fθLm ( fθ ) ‖2H ] , ( 4 ) where ‖ · ‖H denotes the functional norm inH , and α is a hyper-parameter . The above objective is inspired by the Taylor expansion of the MAML objective , but is defined in the RKHS induced by the NTK . Its connection with MAML and some functional-space properties will be discussed later . Solving the Function Optimization Problem To minimize equation 4 , we first derive Theorem 2 to reduce the function optimization problem to a parameter optimization problem . Theorem 2 Let fθ be a neural network with parameter θ andH be the RKHS induced by the NTK Θ of fθ . Then , the following are equivalent Ẽ ( α , fθ ) =M1 , and ‖∇fθLm ( fθ ) ‖2H = ‖∇θLm ( fθ ) ‖2 . Theorem 2 is crucial to our approach as it indicates that solving problem equation 4 is no more difficult than the original parameter-based MAML , although it only considers one-step adaptation case . Next , we will show that multi-step adaptation in the parameter space can also be well-approximated by our objective equation 4 but with a scaled regularized parameter α . In the following , we consider the squared loss L. The case with the cross-entropy loss is discussed in the Appendix . We assume that fθ is parameterized by either fully-connected or convolutional neural networks , and only consider the impact of number of hidden layers L in our theoretical results . Theorem 3 Let fθ be a fully-connected neural network with L hidden layers and ReLU activation function , s1 , ... , sL+1 be the spectral norm of the weight matrices , s = maxh sh , andα be the learning rate of gradient descent . If α ≤ O ( qr ) with q = min ( 1/ ( LsL ) , L−1/ ( L+1 ) ) and r = min ( s−L , s ) , then the following holds |Mk − Ẽ ( kα , fθ ) | ≤ O ( 1 L ) . Theorem 4 Let fθ be a convolutional neural network with L − l convolutional layers and l fullyconnected layers and with ReLU activation function , and dx be the input dimension . Denote by Wh the parameter vector of the convolutional layer for h ≤ L− l , and the weight matrices of the fully connected layers for L− l + 1 ≤ h ≤ L+ 1 . ‖ · ‖2 means both the spectral norm of a matrix †For ease of our later notation , we write the gradient∇θiL ( thus the parameter as well ) as a row vector . and the Euclidean norm of a vector . Define sh = √ dx‖Wh‖2 if h = 1 , ... , L − l , and ‖Wh‖2 if L − l + 1 ≤ h ≤ L + 1 . Let s = maxh sh and α be the learning rate of gradient descent . If α ≤ O ( qr ) with q = min ( 1/ ( LsL ) , L−1/ ( L+1 ) ) and r = min ( s−L , s ) , the following holds |Mk − Ẽ ( kα , fθ ) | ≤ O ( 1 L ) . The above Theorems indicate that , for a meta-model with fully-connected and convolutional layers , the proposed Meta-RKHS-I can be an efficient approximation of MAML with a bounded error . Comparisons with Reptile and MAML Similar to Reptile and MAML , the testing stage of MetaRKHS-I also requires gradient-based adaptation on meta-test tasks . By Theorem 1 , we known that gradient flow of an energy functional can be approximated by gradient descent in a parameter space . Reptile with 1-step adaptation ( Nichol et al. , 2018 ) is equivalent to the approximation of the gradient flow of Ẽ ( α , fθ ) with α = 0 , which does not include the fast-adaptation regularization as in our method . For a fairer comparison on the efficiency , we will discuss the computational complexity later . From the equivalent parameter-optimization form indicated in Theorem 2 , we know that our energy functional Ẽ ( α , fθ ) is closely related to MAML . However , with this form , our method does not need the explicit adaptation steps in training ( i.e. , the inner-loop of MAML ) , leading to a simpler optimization problem . We will show that our proposed method leads to better results .

Summary: This paper mainly deals with the computational issues of Model Agnostic Meta-Learning (MAML). Specifically, it proposes two meta-learning algorithms where the hypothesis class (i.e. the mapping function set) is defined in RKHS induced by NTK. Extensive experimental studies on many tasks (i.e. regression, few-shot image classification, robustness to adversarial attack, and out-of-distribution generalization) illustrate its superiority compared with other baselines.