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

Abstract:
In this paper, we study the generalization properties of Model-Agnostic MetaLearning (MAML) algorithms for supervised learning problems. We focus on the setting in which we train the MAML model over m tasks, each with n data points, and characterize its generalization error from two points of view: First, we assume the new task at test time is one of the training tasks, and we show that, for strongly convex objective functions, the expected excess population loss is bounded byO(1/mn). Second, we consider the MAML algorithm’s generalization to an unseen task and show that the resulting generalization error depends on the total variation distance between the underlying distributions of the new task and the tasks observed during the training process. Our proof techniques rely on the connections between algorithmic stability and generalization bounds of algorithms. In particular, we propose a new definition of stability for meta-learning algorithms, which allows us to capture the role of both the number of tasks m and number of samples per task n on the generalization error of MAML.

heading: 1 Introduction
text: 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.

heading: 2 Problem formulation
text: 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.

heading: 3 Theoretical results
text: In this section, we formally characterize the excess population loss (test error) of the MAML solution, when we measure the performance of a model after one step of SGD adaptation. In particular, we first discuss the training error of MAML in detail. Then, we establish a generalization error bound for the case that the solution of MAML is evaluated over new samples of a recurring task. Finally, we state the generalization error of MAML once its solution is applied to a new unseen task. Before stating our results, we mention our required assumptions.
Assumption 1. For any z ∈ Z , the function `(., z) is twice continuously differentiable. Furthermore, we assume it satisfies the following properties for any w, u ∈ Rd: (i) For any z ∈ Z , the function `(., z) is µ-strongly convex, i.e., ‖∇`(w, z)−∇`(u, z)‖ ≥ µ‖w−u‖; (ii) The gradient norm is uniformly bounded by G overW , i.e., ‖∇`(w, z)‖ ≤ G;
(iii) The loss is L-smooth over Rd, i.e., ‖∇`(w, z)−∇`(u, z)‖ ≤ L‖w − u‖;
(iv) Hessian is ρ-Lipschitz continuous over Rd, i.e., ‖∇2`(w, z)−∇2`(u, z)‖ ≤ ρ‖w − u‖.
We also require the following assumption on the tasks distribution. This assumption implies that, with probability one, a set of finite samples generated from a distribution pi are all different.
Assumption 2. We assume Z is a Polish space (i.e., complete, separable, and metric) and FZ is the Borel σ-algebra over Z . Moreover, for any i, pi is a non-atomic probability distribution over (Z,FZ), i.e., pi(z) = 0 for every z ∈ Z .

heading: 3.1 Training error
text: While the main focus of this paper is on studying the population error of MAML algorithm, we first study its training error which is required to provide characterization of the excess loss of MAML. To do so, we first state the following result from [24] and [13] on the strong convexity and smoothness of `(w − α∇L̂(w,D), z) for any batch D and any z ∈ Z . Lemma 1 ([13] & [24]). If Assumption 1 holds, then for an arbitrary batch D and z ∈ Z , and with α ≤ 1L , the function `(w − α∇L̂(w,D), z) is 4L + 2αρG smooth over W . Furthermore, `(w − α∇L̂(w,D), z) is µ8 -strongly convex, if α ≤ min{ 1 2L , µ 8ρG}.
An immediate consequence of this Lemma is that the MAML empirical loss F̂ defined in (4) is also µ/8-strongly convex and 4L + 2αρG smooth over W . In addition, it can be shown that the norm of gi(w;Dini ,Douti ) defined in (6), which is the unbiased gradient estimate used in MAML, is uniformly bounded above; for more details check Lemma 5 in Appendix A. Having these properties of the MAML empirical loss established, we next state the following proposition on the training error of MAML. This result is obtained by slightly modifying the well-known results on the convergence of SGD in [30–32] in order to take into account the stepsize constraints that are imposed by generalization analysis. For completeness, the proof of this result is provided in Appendix B.
Proposition 1. Consider F̂ (.,S) defined in (4) with α ≤ min{ 12L , µ 8ρG}. If Assumption 1 holds, then for MAML with βt = min(β, 8µ(t+1) ) for β ≤ 8/µ, and for any set S, the last iterate w T satisfies
E [ F̂ (wT ,S)− F̂ (w∗S ,S) ] ≤ O(1) G2(1 + 1βµ )
µ2
( L+ ραG
T + G√ T
) , (7)
and the time-average of iterates w̄T satisfies
E [ F̂ (w̄T ,S)− F̂ (w∗S ,S) ] ≤ O(1) G2(log(T ) + 1βµ )
µT ,
where w∗S := arg minw∈W F̂ (.,S) and the expectations are taken over the randomness of algorithm.
In the above expressions, the notation O(1) only hides absolute constants. It is worth noting that the termG/ √ T in (7) vanishes, ifw∗S be a minimizer of the unconstrained problem, i.e.,∇F̂ (w∗S ,S) = 0.

heading: 3.2 Generalization error
text: We derive our generalization bounds for MAML by establishing its algorithmic stability properties. The stability approach has been used widely to characterize the generalization properties for optimization algorithms such as stochastic gradient descent [18] or differentially private methods [33]. These arguments are based on showing the uniform stability of algorithms [17] which we restate it here.
Definition 1 ([17]). Consider the problem of minimizing the empirical function L̂(w,H) for some datasetH. A randomized algorithm A with output wH given datasetH is called γ-uniformly stable if the following condition holds: Take the dataset H̃ which is the same as H, except at one data points. Then, we have supz̃∈Z EA [|`(wH, z̃)− `(wH̃, z̃)|] ≤ γ, where the expectation is taken over the randomness of A.
The above definition captures the stability of an algorithm. Specifically, it states that Algorithm A is γ-stable, if the resulting loss of its outputs, when it is run using to two different datasets that only differ in one data point, are at most γ away from each other. Note that the above definition holds if the difference between the losses evaluated at any point z̃ is bounded by γ. The main importance of this definition is its connection with generalization error. In particular, it can be shown that if an algorithm is γ-uniformly stable and “symmetric", then its generalization error is bounded above by γ; see, e.g., [17]. Next, we formally state the definition of a symmetric algorithm.
Definition 2. An algorithm A : Zn → Rd is called symmetric, if for any S ⊂ Zn, the distribution of its output, i.e., A(S), does not depend on the ordering of elements of S, i.e., if we take S ′ as a permutation of S, the distribution of A(S) and A(S ′) would be similar.
Note that Definition 1 is useful for the case where we measure the performance of a model w by its loss function over a sample, i.e., `(w, z̃). However, in this paper we measure the performance of a model by looking at its loss after one step of SGD which involves K data points, as defined in (6). Therefore, we cannot directly use Definition 1 for characterizing the generalization error of MAML. In fact, in what follows, we first propose a modified version of the uniform stability definition, which is compatible with our setting, and then show how such stability could lead to generalization bounds for MAML-type algorithms. Definition 3. Consider the problem in (4). A randomized algorithm A with output wS given dataset S is called (γ,K)-uniformly stable if the following condition holds for any i ∈ {1, . . .m}: Take the dataset S̃ which is the same as S , except that S̃ ini and S̃outi differ from S ini and Souti in at most K and one data points, respectively. Then, for any z̃ ∈ Z and any K distinct points {z1, ..., zK} in Z ,
EA [∣∣∣`(wS − α∇L̂(wS , {zj}Kj=1), z̃)− `(wS̃ − α∇L̂(wS̃ , {zj}Kj=1), z̃)∣∣∣] ≤ γ,
where the expectation is taken over the randomness of A.
A few remarks about the above definition follow. First, one might wonder, why it is needed to change K points of the set S ini , while we change only one point of the set Souti . Note that, going from (1) to (3), the expectation EDtesti [.] is replaced by the sum over all ( n K ) possible batches Dini of size K from S ini . In other words, for the empirical sum in (3), each batch Dini can be seen as a data unit. That said, and similar to Definition 1, to characterize the stability, we need to change one data unit which is one batch of size K. That is why we change K data points of S ini in the definition of (γ,K)-uniformly stability. On the other hand, we replace Li(.) = Ez∼pi [`(., z)] in (1) with a sum over n points of Souti in (3), and thus, for this one, each data unit is just a single data point. So, similar to Definition 1, we just change one data point for the set Souti . Second, it is worth comparing this definition with the other definition given for stability of metalearning algorithms in [28]. In that paper, the definition of stability is based on modifying the whole dataset Si rather than what we do here which is changing just K + 1 points. While taking such a definition makes the analysis relatively simpler, it prohibits us from characterizing the dependence of generalization error on n, and hence the resulting upper bound for generalization error would be larger. We will come back to this point later when we derive the stability of MAML with respect to Definition 3 and compare it with the one obtained in [28].
As we discussed, the main reason that we are interested in the uniform stability of an algorithm is its connection with generalization error. In the next theorem, we formalize this connection for MAML formulation and show that if an Algorithm A is (γ,K)-uniformly stable and symmetric, then its output generalization error is bounded above by γ. The proof of this result is available in Appendix C. Theorem 1. Consider the population and empirical losses defined in (2) and (4), respectively. If Assumption 2 holds and A is a (possibly randomized) symmetric and (γ,K)-uniformly stable algorithm with output wS ∈ W , then EA,S [ F (wS)− F̂ (wS ,S) ] ≤ γ.
This result shows that if we prove a symmetric algorithm is (γ,K)-uniformly stable as defined in Definition 3, then we can bound its output model generalization error by γ. Hence, to characterize the generalization error of the model trained by MAML algorithm, we only need to capture the uniform stability parameter of MAML. Before stating this result, it is worth noting that while we limit our focus to MAML in this paper, Definition 3 and Theorem 1 could provide a framework for studying the generalization properties of a broader class of gradient-based meta-learning algorithms such as Reptile [34], First-order MAML [1], and Hessian-Free MAML [13]. Theorem 2. If Assumption 1 holds, then MAML (Algorithm 1) with both last iterate and average iterate outputs and with α ≤ min{ 12L , µ 8ρG} and βt ≤ 1 4L+2αρG is (γ,K)-uniformly stable, where γ := O(1)G 2(1+αLK) mnµ .
According to the above discussion, the result of Theorem 2 guarantees that the generalization error of MAML solution decays by a factor of O(K/mn), where m is the number of tasks in the training set
and n is the number of available samples per task. The classic lower bound for SGD over strongly convex functions translates to a O(1/mn) lower bound in our setting. Hence, our bound is tight in the small K regime, which is generally the case in few-shot learning problems. However, one shortcoming of this result is that it is not tight in the large K regime. In Appendix E we show how we could improve this result for the large K regime. However, throughout the paper, we keep our discussion limited to the small K regime. Remark 1. If instead of using our uniform stability definition (i.e., Definition 3), one uses the stability definition given in [28], the resulted stability constant γ would be proportional to (1/m) rather than (1/mn). In fact, our proposed uniform-stability definition empowers us to obtain a better bound and indicates the role of number of samples per task n in the generalization error. Remark 2. The algorithmic stability technique is mainly limited to the convex setting, since, in the nonconvex case, we need to keep learning rate very small to obtain meaningful generalization results which makes it impractical (Check Appendix G for further discussions on this matter). In fact, the main reason that we assume ` is strongly convex and α ≤ Ω(µ) is to ensure that the meta-objective is convex, as, in general, relaxing any of these two could lead to a nonconvex meta-objective function. However, these two assumptions together make the objective function strongly convex, which is not necessarily needed in our analysis. In fact, if we assume that ` and the meta-function are convex (but not necessarily strongly convex), we could still use Definition 3 to derive similar generalization bounds.
Putting Proposition 1 and Theorem 2 together, we obtain the following result on the excess population loss of MAML algorithm. We only report the result for the averaged iterates here, but one can obtain the result for the last iterate similarly by using Proposition 1. Proposition 2. Consider the function F defined in (2) with α ≤ min{ 12L , µ 8ρG}. If Assumptions 1 and 2 hold, then the average of iterates generated by MAML (Algorithm 1) with βt = min( 14L+2αρG , 8 µ(t+1) ) after T iterations satisfies
EA,S [ F (w̄T )−min W F ] ≤ O(1)G 2 µ ( log(T ) + L/µ T + 1 + αLK mn ) ,
where the expectation is taken over the sampling of S and the randomness of MAML algorithm.
As an immediate application, the following corollary characterizes MAML test error.
Corollary 1. Under the premise of Proposition 2, MAML algorithm after T = Õ(mnL/µ) iterations returns w̄T such that EA,S [ F (w̄T )−minW F ] ≤ O ( G2(1 + αLK)/(mnµ) ) .

heading: 3.3 Generalization to an unseen task
text: As we discussed in Section 2, another generalization measure is how the model trained with respect to the empirical problem in (4) performs on a new and unseen task Tm+1 with corresponding distribution pm+1. To state our result for this case, we first need to introduce the following distance notion between probability distributions. Definition 4. For two distributions P and Q, defined over the sample space Ω and σ-field F , the total variation distance is defined as ‖P −Q‖TV := supA∈F |P (A)−Q(A)|.
It is well-known that the total variation distance admits the following characterization
‖P −Q‖TV = sup f :0≤f≤1 Ex∼P [f(x)]− Ex∼Q[f(x)]. (8)
Also, we require the following boundedness assumption for our result. Assumption 3. For any z ∈ Z , the function `(., z) is M -bounded overW .
Considering these assumptions, we are ready to state our result for the case when the task at test time is a new task and is not observed during training. Theorem 3. Consider the population losses defined in (1) and (2). Suppose Assumptions 1, 2 and 3 hold. Then, for any w ∈ W , we have
|Fm+1(w)− F (w)| ≤ D(pm+1, {pi}mi=1), (9)
where
D(pm+1, {pi}mi=1) := 4αG2
m m∑ i=1 ‖pm+1 − pi‖TV + (M + 2αG2)‖pm+1 − 1 m m∑ i=1 pi‖TV . (10)
While the proof is provided in detail in Appendix F, here we discuss a sketch of it to highlight the main technical contributions. To simplify the notation here, let us assume m = 1, meaning that p1 is the distribution used for training and p2 is the distribution corresponding to the new task. Note that we aim to bound |F2(w)− F1(w)|. Recalling the definition of population loss (2), we need to bound the following expression (we drop the absolute value due to symmetry)
E{z2j∼p2}Kj=1,z̃2∼p2 [ l ( w − α∇L̂(w, {z2j }j), z̃2 )] − E{z1j∼p1}Kj=1,z̃1∼p1 [ l ( w − α∇L̂(w, {z1j }j), z̃1 )] .
(11)
Notice that this difference can be cast as E({zj}Kj=1,z̃)∼pK+12 [X] − E({zj}Kj=1,z̃)∼pK+11 [X], with X := l ( w − α∇L̂(w, {zj}Kj=1), z̃ ) . As a result, a naive approach would be using Lipschitz and
boundedness properties of l (Assumptions 1 and 3) along with (8) to obtain a bound depending on ‖pK+11 − p K+1 2 ‖TV = O(K)‖p1 − p2‖TV . However, this bound is not tight as it grows with K.
To address this issue, we exploit a coupling technique. Note that the expression in (11) does not depend on the joint distribution of z1j and z 2 j , and instead, it only depends on the marginal distribution of z1j and z 2 j . That said, for each j, we assume that z 1 j and z 2 j are sampled from a distribution µ on Z × Z such that z1j ∼ p1, z2j ∼ p2, and µ(z1j 6= z2j ) = ‖p1 − p2‖TV . Such a coupling exists and is called maximal coupling of p1 and p2 [35]. Using this idea, as we show in Appendix F, we can eliminate the dependence on K, and as a result, the upper bound in (9) is independent of number of available labeled samples at test time denoted by K. Remark 3. Note that the terms 1m ∑m i=1 ‖pm+1 − pi‖TV and ‖pm+1 − 1 m ∑m i=1 pi‖TV in D(pm+1, {pi}mi=1) come from the fact that we consider uniform distribution over tasks in the empirical problem (4). In particular, if we instead consider the empirical problem arg minw∈W ∑m i=1 qiF̂i(w,Si), for some non-negative weights qi with ∑m i=1 qi = 1, then D(pm+1, {pi}mi=1) on the right hand side of (9) would change to
(M + 2αG2)‖pm+1 − m∑ i=1 qipi‖TV + 12αG2 m∑ i=1 qi‖pm+1 − pi‖TV .
This result shows that by changing the training problem we can achieve a lower generalization error for MAML, if we have some information about the distribution pm+1 at training time. For instance, if we know pm+1 will be much closer to p1 compared to p2, making the weight of p1 larger than p2 would decrease the generalization error of MAML. Corollary 2. Recall the population loss Fm+1 defined in (2) and D(pm+1, {pi}mi=1) defined in Theorem 3. Let A be an algorithm for solving the empirical problem (4) which achieves excess risk, i.e., EA,S [F (A(S))]−minW F ≤ . If Assumptions 1, 2 and 3 hold, then algorithmA finds a model wS which achieves +D(pm+1, {pi}mi=1) excess loss with respect to Fm+1,
EA,S [Fm+1(wS)]−min W Fm+1 ≤ + 2D(pm+1, {pi}mi=1).
This corollary and Proposition 2 together imply that the MAML algorithm’s test error with respect to the new task Tm+1 is O(1) ( 1 mn +D(pm+1, {pi} m i=1) ) . As a result, if the new task’s distribution pm+1 is sufficiently close to the other tasks’ distributions, MAML will have a low test error on the new unseen task. On the other hand, if pm+1 is far from p1, . . . , pm in TV distance, then test error of the model trained {T mi=1} over Tm+1 could be potentially large. In Appendix F.2 we show how this result can be extended to the case that the task at test time is generated from a distribution over both recurring tasks {Ti}mi=1 and the unseen task Tm+1.

heading: 4 Conclusion and future work
text: In this work, we studied the generalization of MAML algorithm in two key cases: a) when the test time task is a recurring task from the ones observed during the training stage, b) when it is a
new and unseen one. For the first one, and under strong convexity assumption, we showed that the generalization error improves as the number of tasks or the number of samples per task increases. For the second case, we showed that when the distance between the unseen task’s distribution and the distributions of training tasks is sufficiently small, the MAML output generalizes well to the new task revealed at test time.
While we focused on the convex case in this paper, deriving generalization bounds when the metafunction is nonconvex is a natural future direction to explore. However, this could be challenging since the generalization of gradient methods is not well understood in the nonconvex setting even for the classic supervised learning problem.

heading: 5 Acknowledgment
text: Alireza Fallah acknowledges support from the Apple Scholars in AI/ML PhD fellowship and the MathWorks Engineering Fellowship. This research is sponsored by the United States Air Force Research Laboratory and the United States Air Force Artificial Intelligence Accelerator and was accomplished under Cooperative Agreement Number FA8750-19-2-1000. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the United States Air Force or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein. This research of Aryan Mokhtari is supported in part by NSF Grant 2007668, ARO Grant W911NF2110226, the Machine Learning Laboratory at UT Austin, and the NSF AI Institute for Foundations of Machine Learning.

Review summaries:
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. 

This paper analyzes the generalization bound of MAML-type algorithms from two aspects: (1) when the test task is a sample from the training tasks and (2) when the test task is unseen during training stage. A generalization bound on strongly convex case is given in case (1) showing the generalization error is deceasing at a rate of $O(1/mn)$. On the other hand, a bound under the total variation assumption is given in case (2), showing the bound depends on the average total variation between all the training tasks and the unseen task. The paper focuses on theory and thus no experiments are included.

The paper analyzes the generalization error of MAML in 2 scenarios that are different in the new tasks at test time which are (i) the training tasks and (ii) unseen tasks, under the assumptions that the loss function is (strongly) convex and the number of steps of SGD in the adaptation step of MAML is 1. In the former case, this work proposes a new notion of stability that is capable of capturing the dependence on the number of data points per task in the bound. In the latter case, this work formulates a generalization error bound based on the total variation distance between the distribution of the new task and those of the training task.

This paper studies the generalization properties of MAML in the strongly convex setting. It's assumed that there are $m$ training tasks, each with $n$ samples, and the inner loop adaptation is one step of SGD using $K$ samples. It's proved that if the new task is a uniformly random task from one of the $m$ training tasks, MAML incurs a generalization error of $O(\frac{K}{mn})$; if the new task is an unseen task, the error depends on the TV distance between the new task distribution and all the training task distributions as well as their average. The analysis is based on a modified version of algorithmic stability tailored to the MAML setup.