Generalization of Model-Agnostic Meta-Learning Algorithms: Recurring and Unseen Tasks

1 Introduction . In several machine learning problems , it is of interest to design algorithms that can be adjusted based on previous experiences and tasks to perform better on a new task . In particular , metalearning algorithms achieve such a goal through various approaches , including finding a proper meta-initialization for the new task [ 1–3 ] , updating the model architecture [ 4–6 ] , or learning the parameters of optimization algorithms [ 7 , 8 ] . A popular meta-learning framework that has shown promise in practice is Model-Agnostic MetaLearning ( MAML ) , which was first introduced in [ 1 ] . MAML algorithm uses available training data on a number of tasks to come up with a meta-initialization that performs well after it is slightly updated at test time with respect to the new task . In other words , unlike standard supervised learning , in which we aim to find a model that generalize well to a new task without any adaptation step , in MAML our goal is to find an initial model for learning a new task when we have access to limited labeled data for that task to run one ( or a few ) step ( s ) of stochastic gradient descent ( SGD ) . As shown in Fig . 1 , in MAML we are given m tasks with m corresponding datasets { Si } mi=1 in the training phase . Once the model is trained ( w∗train ) , a new task is revealed at test time for which we have access to K labeled samples drawn from Dtest . We use these labeled samples of the new task to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . update the trained model by running a step of SGD leading to a new model for the test task ( w∗new ) . We finally evaluate the performance of the updated model over the test task , denoted by Ltest ( w∗new ) . MAML and its variants have been extensively studied over the past few years from both empirical and theoretical point of view [ 2 , 9–16 ] . In particular , [ 13 ] provided convergence guarantees for MAML algorithm under the assumption that access to fresh samples at any round of the training stage is possible , and [ 15 ] extended this results to the case that multiple gradient steps can be performed at test time . However , one shortcoming of such analysis is that , at training stage , we often do not have access to fresh samples at every iteration . Instead , we have access to a large set of realized samples and we typically do multiple passes over the data points during the training stage . Hence , it is essential to come up with a novel analysis that addresses this issue by characterizing the training error and generalization error of MAML separately . In this paper , we accomplish this goal and showcase the role of different problem parameters in the generalization error of MAML . Specifically , we assume that we are given m supervised learning tasks , with ( possibly different ) underlying distributions p1 , . . . pm , where for each task we have access to n samples1 . As we measure the performance of a model by its loss after one step of SGD adaptation with K samples , the problem that one can solve in the training phase is minimizing the average loss , over all given m tasks and their n samples , after one step of SGD with K samples . This empirical loss can be considered as a surrogate for the desired expected loss ( with respect to tasks data ) over all m tasks . Here , we focus on the case that MAML is used to solve this empirical minimization problem , and our goal is to quantify the test error of MAML output . To tackle this problem , we first briefly revisit the results from the optimization literature to bound the training error of MAML , assuming that the loss functions are strongly convex . We next turn to the main focus of our paper which is the generalization properties of MAML . More specifically , we address the following questions : • If one of the m given tasks recurs uniformly at random at test time , then how well ( in expectation ) would the trained model perform after adaptation with SGD over the fresh samples of that task ? In other words , having training error minimized , what would be the generalization error and our guarantee on test error ? Here , we show that for strongly convex objective functions , we could achieve a generalization error that decays at O ( 1/mn ) . Our analysis builds on the connections between algorithmic stability and generalization of the output of algorithms . While this relation is well-understood in classic statistical learning [ 17 , 18 ] , here we propose a novel stability definition for meta-learning algorithms which allows us to restore such connection for our setting . • Assuming that the task at test time is NOT one of the m tasks at training , how would the model perform on that task after the adaptation step ? We answer this question by focusing on the case that the revealed task at the test time is a new unseen task with underlying data distribution pm+1 , and formally characterizing the generalization error of MAML in this case . We show that when the task at test time is new , the generalization error also depends on the total variation distance between pm+1 and p1 , . . . , pm . Related work : Recently , there has been significant progress in studying theoretical aspects of meta-learning , in particular , MAML . Authors in [ 19 ] proposed iMAML which updates the model using an approximation of one step of proximal point method and studied its convergence . In [ 20 ] , authors introduced the task-robust MAML by considering a minimax formulation rather than minimization . Several papers have also studied MAML through more general frameworks such as bilevel optimization [ 21 ] , stochastic compositional optimization [ 22 ] , and conditional stochastic 1More precisely , in our analysis we take 2n samples per each task to simplify derivations . optimization [ 23 ] . Also , several works have studied the extension of meta-learning theory to online learning [ 24 , 3 ] , federated learning [ 25 ] , and reinforcement learning [ 26 , 27 ] . The most relevant paper to our work is [ 28 ] that studies generalization of meta-learning algorithms using stability techniques and shows a O ( 1/ √ m ) bound for nonconvex loss functions . Here we focus on strongly convex objective functions and present an analysis that differs from this work in two fundamental aspects . First , we present a different notion of stability that allows us to capture the number of data points per task in our bound . In particular , our stability notion measures sensitivity of the algorithm to perturbations that involve changing K data points which is the data unit involved in the adaptation step of the MAML algorithm . This enables us to obtain a much tighter bound O ( 1/mn ) ( compared to O ( 1/m ) achieved in [ 28 ] for strongly convex functions ) , highlighting the dependence on the number of the data samples available for each task . Second , we also consider the generalization of MAML for the case that the task at test time is not one of the available tasks during the training stage . The generalization of MAML has also been studied in [ 29 ] from an empirical point of view . In particular , they show that the generalization of MAML to new tasks is correlated with the coherence between their adaptation trajectories in parameter space . This is aligned with the connection of generalization and closeness of underlying distributions that we observe in our results . 2 Problem formulation . In this paper , we consider the supervised learning setting , where each data point is denoted by z = ( x , y ) ∈ Z with x ∈ X being the input ( feature vector ) and y ∈ Y being its corresponding label . We use the loss function l : Rd × Z → R+ to evaluate the performance of a model parameterized by w ∈ W , where W is a convex and closed subset of Rd . In other words , for a data point z = ( x , y ) ∈ Z , the loss ` ( w , z ) denotes the error of model w in predicting the label y given input x . We consider access to m tasks denoted by T1 , . . . , Tm , where the data corresponding to each task Ti is generated from a distinct distribution pi . The population loss corresponding to task Ti for model w is defined as Li ( w ) : = Ez∼pi [ ` ( w , z ) ] . We further use the notation L̂ ( w ; D ) to denote the empirical loss corresponding to datasetD , which is defined as the average loss of w over the samples of dataset D , i.e. , L̂ ( w ; D ) : = 1|D| ∑ z∈D ` ( w , z ) , where |D| is the size of dataset D. In general , and throughout the paper , we use the hat notation to distinguish empirical losses from population losses . Our goal is to find w ∈ W that performs well on average2 over all tasks , after it is updated with respect to the new task and by using one step of stochastic gradient descent ( SGD ) with a batch of size K. To formally introduce this problem we first define the function Fi ( w ) which captures the performance of model w over task Ti once it is updated by a single step of SGD , Fi ( w ) : = EDtesti [ Li ( w − α∇L̂ ( w , Dtesti ) ) ] = EDtesti Ez∼pi [ ` ( w − α K ∑ z′∈Dtesti ∇ ` ( w , z′ ) , z ) ] ( 1 ) where Dtesti is a batch with K different samples , drawn from the probability distribution pi . Note that the outer expectation is taken with respect to the choice of elements of Dtesti while the inner one is taken with respect to the data of task i . As our goal is to find a model that performs well after one step of adaptation over all m tasks , we minimize the average expected loss over all given tasks , which can be written as min w∈W F ( w ) : = 1 m m∑ i=1 Fi ( w ) . ( 2 ) As the underlying distribution of tasks are often unknown in most applications , we are often unable to directly solve the problem in ( 2 ) . On the other hand , for each task , we often have access to data points that are drawn according to their data distributions . Therefore , instead of solving ( 2 ) , we solve its sample average surrogate problem in which each Fi is approximated by its empirical loss . 2Our analysis can be extended to the case that the distribution over tasks is not uniform . To formally define the empirical loss for each task , suppose for each task Ti we have access to a training set Si , where its elements are drawn independently according to the probability distribution pi . We further divide the set Si into two disjoint sets of size n defined as S ini and Souti , i.e. , Si : = { S ini , Souti } and |S ini | = |Souti | = n. Here , we use the elements of the S ini to estimate the inner gradient ∇L̂ ( w , Dtesti ) and use the samples in the set Souti to estimate the outer function Li ( . ) . Specifically , we define the sample average of Fi using data sets S ini and Souti as F̂i ( w , Si ) : = 1 ( n K ) ∑ Dini ⊂S in i |D in i |=K L̂ ( w − α∇L̂ ( w , Dini ) , Souti ) ( 3 ) = 1 ( n K ) ∑ Dini ⊂S in i |D in i |=K 1 n ∑ z∈Souti ` w − α K ∑ z′∈Dini ∇ ` ( w , z′ ) , z . This expression shows that we use all n elements of Souti to approximate the expectation required for the computation of Li , and we approximate the expectation with respect to the test set by averaging over all subsets of S ini that haveK elements . Given this expression , the sample average approximation ( empirical loss ) of Problem ( 1 ) is given by arg min w∈W F̂ ( w , S ) : = 1 m m∑ i=1 F̂i ( w , Si ) , ( 4 ) where S : = { Si } mi=1 is defined as the concatenation of all tasks data sets . Having the dataset S , a ( possibly randomized ) optimization algorithm A with output A ( S ) can be used to find an approximate solution to the problem in ( 4 ) . The error of this solution with respect to the MAML empirical loss , i.e. , F̂ ( A ( S ) , S ) − minW F̂ ( . , S ) , is called training error . In this paper , we are mainly interested to bound the test error which is the error of A ( S ) with respect to the population loss , i.e. , F ( A ( S ) ) −minW F . The test error is also sometimes called excess ( population ) loss . Note that the expected test error can be decomposed into three terms : EA , S [ F ( A ( S ) ) −min W F ] ( test error ) = EA , S [ F ( A ( S ) ) −F̂ ( A ( S ) , S ) ] ︸ ︷︷ ︸ generalization error +EA , S [ F̂ ( A ( S ) , S ) −min W F̂ ( . , S ) ] ︸ ︷︷ ︸ training error +ES [ min W F̂ ( . , S ) ] −min W F︸ ︷︷ ︸ ≤0 . It can be verified that the expectation of the third term ( over A and S ) is non-positive since ES [ minW F̂ ( . , S ) ] ≤ minW ES [ F̂ ( . , S ) ] and ES [ F̂ ( . , S ) ] = F. Hence , to bound the expected test error , we should bound the expectation of training and generalization errors . The Model-Agnostic Meta-Learning ( MAML ) method proposed in [ 1 ] is designed to solve the empirical minimization problem defined in ( 4 ) . The steps of MAML are outlined in Algorithm 1 . MAML solves Problem ( 4 ) by using SGD update for the average loss function F̂ ( w , S ) . To better highlight this point , note that the gradient of∇F̂ ( w , S ) can be written as 1m ∑m i=1∇F̂i ( w , Si ) , where the i-th term corresponding to task Ti is given by ∇F̂i ( w , Si ) = 1 ( n K ) ∑ Dini ⊂S in i |Dini |=K [ ( Id − α∇2L̂ ( w , Dini ) ) ×∇L̂ ( w − α∇L̂ ( w , Dini ) , Souti ) ] , ( 5 ) which involves the second-order information of the loss function . Therefore , to compute a mini-batch approximation for the above gradient , we consider the batchesDini ⊂ S ini with sizeK andDouti ⊂ Souti with b elements . Replacing the above sums with their batch approximations leads to the following stochastic gradient approximation gi ( w ; Dini , Douti ) : = ( Id − α∇2L̂ ( w , Dini ) ) ∇L̂ ( w − α∇L̂ ( w , Dini ) , Douti ) , ( 6 ) which is indeed an unbiased estimator of the gradient ∇F̂i ( w , Si ) in ( 5 ) . If for each task we perform the update of SGD with gi and then compute their average it would be similar to running SGD for Algorithm 1 : MAML [ 1 ] Input : The set of datasets S = { Si } mi=1 with Si = { S ini , Souti } ; test time batch size K ; # of tasks summoned at each round r ; # of iterations T . Choose arbitrary initial point w0 ∈ W ; for t = 0 to T − 1 do Choose r tasks uniformly at random ( out of m tasks ) and store their indices in Bt ; for all Ti with i ∈ Bt do Sample a batch Dt , ini of K different elements from S ini with replacement ; Sample a batch Dt , outi of size b from Souti and with replacement ; wt+1i : = w t − βt ( Id − α∇2L̂ ( wt , Dt , ini ) ) ∇L̂ ( wt − α∇L̂ ( wt , Dt , ini ) , D t , out i ) ; end for wt+1 : = rW ( 1 r ∑ i∈Bt w t+1 i ) ; end for Return : wT and w̄T : = 1T+1 ∑T t=0 w t the average loss ∇F̂ ( w , S ) . This is exactly how MAML is implemented in practice as outlined in Algorithm 1 . In this paper , we consider a constrained problem , and as a result , we also need an extra projection step in the last step to ensure the feasibility of iterates . Finally , the output of MAML could be the last iterate wT or the time-average of all iterates w̄T : = 1T+1 ∑T t=0 w t. As stated earlier , the convergence properties of MAML-type methods from an optimization point of view have been studied recently under different set of assumptions . In this paper , as we characterize the sum of training error and generalization error , we briefly discuss the optimization error of MAML when it is used to solve the empirical problem in ( 4 ) . However , the main focus of this paper is on studying the generalization error of MAML with respect to new samples and new tasks . Specifically , we aim to address the following questions : ( i ) How well does the solution of ( 4 ) generalize to the main problem of interest in ( 2 ) ? This could be seen as the generalization error of the MAML algorithm over new samples for recurring tasks . ( ii ) How well does the solution of ( 4 ) generalize to samples from new unseen tasks ? To be more precise , how would the obtained model preform if the new task is not one of the m tasks T1 , . . . , Tm observed at training , and it is rather a new , unseen task Tm+1 with an unknown underlying distribution pm+1 ? In the upcoming sections , we answer these questions on the generalization properties of MAML in detail and characterize the role of number of tasks m , number of samples per task n , and number of labeled samples revealed at test time K .

Summary: In this paper, the authors provide a generalization and uniform-stability-type analysis for MAML training. In specific, the authors first present a generalization error bound (based on the uniform stability bound) of $O(1/mn)$ for strongly convex objective functions when the new task at the test time belong to one of the training tasks. Then, for a more interesting setting where the new task does not belong to the training tasks, the authors provide a generalization bound consisting of a total variation distance between the data distributions of this unseen task and the training tasks. One toy experiment is provided in the supplementary materials to verify the impact of the sizes $m,n$ on the generalization performance.