Generalization Bounds for Meta-Learning via PAC-Bayes and Uniform Stability

1 Introduction . A major challenge with current machine learning systems is the need to acquire large amounts of training data in order to learn a new task . Over the past few decades , meta-learning [ 62 , 70 ] has emerged as a promising avenue for addressing this challenge . Meta-learning relies on the intuition that a new task often bears significant similarity to previous tasks ; hence , a learner can learn to perform a new task very quickly by exploiting data from previously-encountered related tasks . The meta-learning problem formulation thus assumes access to datasets from a variety of tasks during meta-training . The goal of the meta learner is then to learn inductive biases from these tasks in order to train a base learner to achieve few-shot generalization on a new task . Over the past few years , there has been tremendous progress in practical algorithms for metalearning ( see , e.g. , [ 61 , 55 , 25 , 32 ] ) . Techniques such as model-agnostic meta-learning ( MAML ) [ 25 ] have demonstrated the ability to perform few-shot learning in a variety of supervised learning and reinforcement learning domains . However , our theoretical understanding of these techniques lags significantly behind successes on the empirical front . In particular , the problem of deriving generalization bounds for meta-learning techniques remains an outstanding challenge . Current methods for obtaining generalization guarantees for meta-learning [ 5 , 34 , 77 ] either ( i ) produce bounds that are extremely challenging to compute or ( ii ) produce vacuous or near-vacuous bounds in even highly simplified settings ( see Section 5 for numerical examples ) . Indeed , we note that existing work on generalization theory for meta-learning techniques do not explicitly report numerical values for generalization bounds . This is in contrast to the state of generalization theory in the supervised learning setting , where recent techniques demonstrate the ability to obtain non-vacuous generalization guarantees on benchmark problems ( e.g . visual classification problems [ 24 , 78 , 54 ] ) . The generalization challenge in meta-learning is similar to , but distinct from , the supervised learning case . In particular , any generalization bound for meta-learning must account for two levels of 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . generalization . First , one must account for generalization at the base level , i.e. , the ability of the base learner to perform well on new data from a given task . This is particularly important in the few-shot learning setting . Second , one must account for generalization at the meta level , i.e. , the ability of the meta learner to generalize to new tasks not encountered during meta-training . Moreover , the generalization performance at the two levels is coupled since the meta learner is responsible for learning inductive biases that the base learner can exploit for future tasks . The key technical insight of this work is to bound the generalization error at the two levels ( base and meta ) using two different generalization theory frameworks that each are particularly well-suited for addressing the specific challenges of generalization . At the base level , we utilize the fact that a learning algorithm that exhibits uniform stability [ 14 , 15 ] also generalizes well in expectation ( see Section 4.1 for a formal statement ) . Intuitively , uniform stability quantifies the sensitivity of the output of a learning algorithm to changes in the training dataset . As demonstrated by [ 29 ] , limiting the number of training epochs of a gradient-based learning algorithm leads to uniform stability . In other words , a gradient-based algorithm that learns quickly is stable . Since the goal of meta-learning is precisely to train the base learner to learn quickly , we posit that generalization bounds based on stability are particularly well-suited to bounding the generalization error at the base level . At the meta level , we employ a generalization bound based on Probably Approximately Correct ( PAC ) -Bayes theory . Originally developed two decades ago [ 43 , 38 ] , there has been a recent resurgence of interest in PAC-Bayes due to its ability to provide strong generalization guarantees for neural networks [ 24 , 8 , 54 ] . Intuitively , the challenge of generalization at the meta level ( i.e. , generalizing to new tasks ) is similar to the challenge of generalizing to new data in the standard supervised learning setting . In both cases , one must prevent over-fitting to the particular tasks/data that have been seen during meta-training/training . Thus , the strong empirical performance of PAC-Bayes theory in supervised learning problems makes it a promising candidate for bounding the generalization error at the meta level . Contributions . The primary contributions of this work are the following . First , we leverage the insights above in order to develop a novel generalization bound for gradient-based meta-learning using uniform stability and PAC-Bayes theory ( Theorem 3 ) . Second , we develop a regularization scheme for MAML [ 25 ] that explicitly minimizes the derived bound ( Algorithm 1 ) . We refer to the resulting approach as PAC-BUS since it combines PAC-Bayes and Uniform Stability to derive generalization guarantees for meta-learning . Third , we demonstrate our approach on two metalearning problems : ( i ) a toy non-convex classification problem on the unit-ball ( Section 5.1 ) , and ( ii ) the Mini-Wiki benchmark introduced in [ 34 ] ( Section 5.2 ) . Even in these relatively small-scale settings , we demonstrate that recently-developed generalization frameworks for meta-learning provide either near-vacuous or loose bounds , while PAC-BUS provides significantly stronger bounds . Fourth , we demonstrate our approach in larger-scale settings where it remains challenging to obtain nonvacuous bounds ( for our approach as well as others ) . Here , we propose a practical regularization scheme which re-weights the terms in the rigorously-derived PAC-BUS upper bound ( PAC-BUS ( H ) ; Algorithm 3 in the appendix ) . Recent work [ 77 ] introduces a challenging variant of the Omniglot benchmark [ 35 ] which highlights and tackles challenges with memorization in meta-learning . We show that PAC-BUS ( H ) is able to prevent memorization on this variant ( Section 5.3 ) . 2 Problem formulation . Samples , tasks , and datasets . Formally , consider the setting where we have an unknown meta distribution Pt over tasks ( roughly , “ tasks ” correspond to different , but potentially related , learning problems ) . A sampled task t ⇠ Pt induces an ( unknown ) distribution Pz|t over sample space Z . We assume that all sampling is independent and identically distributed ( i.i.d. ) . Note that the sample space Z is shared between tasks , but the distribution Pz|t may be different . We then sample within-task samples z ⇠ Pz|t and within-task datasets S = { z1 , z2 , . . . , zm } ⇠ Pmz|t . We assume that each sample z has a single corresponding label o ( z ) , where the function o is an oracle which outputs the correct label of z . At meta-training time , we assume access to l datasets , which we call S = { S1 , S2 , . . . , Sl } . Each dataset Si in S is drawn by first selecting a task ti from Pt , and then drawing Si ⇠ Pmz|ti . Hypotheses and losses . Let h denote a hypothesis and L ( h , z ) be the loss incurred by hypothesis h on sample z . The loss is computed by comparing h ( z ) with the true label o ( z ) . For simplicity , we assume that there is no noise on the labels ; we can thus assume that all loss functions have access to the label oracle function o and thus the loss depends only on hypothesis h and sample z . We note that this assumption is not required for our analysis and is made for the ease of exposition . Overloading the notation , we let L ( h , Pz|t ) : = Ez⇠Pz|tL ( h , z ) and bL ( h , S ) : = 1|S| P|S| i=1 L ( h , zi ) . Meta-learning . As with model-agnostic meta-learning ( MAML ) [ 25 ] , we let meta parameters ✓ 2 Rn✓ correspond to an initialization of the base learner ’ s hypothesis . Let h✓ be the ✓-initialized hypothesis . Generally , the initialization ✓ is learned from the multiple datasets we have access to at meta-training time . In this work , we will learn a distribution P✓ over initializations so that we can use bounds from the PAC-Bayes framework . At test time , a new task t ⇠ Pt is sampled and we are provided with a new dataset S ⇠ Pmz|t . The base learner uses an algorithm A ( e.g. , gradient descent ) , the dataset S , and the initialization ✓ ⇠ P✓ in order to fine-tune the hypothesis and perform well on future samples drawn from Pz|t . We denote the base learner ’ s updated hypothesis by hA ( ✓ , S ) . More formally , our goal is to learn a distribution P✓ with the following objective : min P✓ L ( P✓ , Pt ) : = min P✓ E t⇠Pt E S⇠Pm z|t E ✓⇠P✓ L ( hA ( ✓ , S ) , Pz|t ) . ( 1 ) We are particularly interested in the the few-shot learning case , where the number of samples which the base learner can use to adapt is small . A common technique to improve test performance in the few-shot learning case is to allow for validation data at meta-training time . Thus , in addition to a generalization guarantee on meta-learning without validation data , we will derive a bound when allowing for the use of validation data Sva ⇠ Pnz|t during meta-training . 3 Related work . Meta-learning . Meta-learning is a well-studied technique for exploiting similarities between learning tasks [ 62 , 70 ] . Often used to reduce the need for large amounts of training data , a number of approaches for meta-learning have been explored over decades [ 11 , 13 , 16 , 31 , 72 , 61 , 55 , 32 ] . Recently , methods based on model-agnostic meta-learning ( MAML ) [ 25 ] have demonstrated strong performance across different application domains and benchmarks such as Omniglot [ 35 ] and MiniImageNet [ 74 ] . These methods operate by optimizing a set of initial parameters that can be quickly fine-tuned via gradient descent on a new task . The approaches mentioned above typically do not provide any generalization guarantees , and none of them compute explicit numerical bounds on generalization performance . Our approach has the structure of gradient-based meta-learning while providing guarantees on generalization . Generalization bounds for supervised learning . Multiple frameworks have been developed for providing generalization guarantees in the classical supervised learning setting . Early breakthroughs include Vapnik-Chervonenkis ( VC ) theory [ 71 , 6 ] , Rademacher complexity [ 65 ] , and the minimum description length principle [ 12 , 56 , 36 ] . More recent frameworks include algorithmic stability bounds [ 14 , 19 , 29 , 57 , 1 ] and PAC-Bayes theory [ 67 , 43 , 64 ] . The connection between stability and learnability has been established in [ 66 , 73 , 29 ] , and suggests that algorithmic stability bounds are a strong choice of generalization framework . PAC-Bayes theory in particular provides some of the tightest known generalization bounds for classical supervised learning approaches such as support vector machines [ 64 , 38 , 26 , 57 , 4 , 48 ] . Since its development , researchers have continued to tighten [ 38 , 42 , 54 ] and generalize the framework [ 17 , 18 , 59 ] . Exciting recent results [ 24 , 45 , 46 , 10 , 8 , 54 ] have demonstrated the promise of PAC-Bayes to provide strong generalization bounds for neural networks on supervised learning problems ( see [ 33 ] for a recent review of generalization bounds for neural networks ) . It is also possible to combine frameworks such as PAC-Bayes and uniform stability to derive bounds for supervised learning [ 39 ] . We will use these two frameworks to bound generalization in the two levels of meta-learning . In contrast to the standard supervised learning setting , generalization bounds for meta-learning are less common and remain loose . Generalization bounds for meta-learning . As described in Section 1 , meta-learning bounds must account for two “ levels ” of generalization ( base level and meta level ) . The approach presented in [ 41 ] utilizes algorithmic stability bounds at both levels . However , this requires both meta and base learners to be uniformly stable . This is a strong requirement that is challenging to ensure at the meta level . Another recent method , known as follow-the-meta-regularized-leader ( FMRL ) [ 34 ] , provides guarantees for a regularized meta-learning version of the follow-the-leader ( FTL ) method for online learning , see e.g . [ 30 ] . The generalization bounds provided are derived from the application of online-to-batch techniques [ 3 , 22 ] . A regret bound for meta-learning using an aggregation technique at the meta-level and an algorithm with a uniform generalization bound at the base level is provided in [ 3 ] . The techniques mentioned do not present an algorithm which makes use of validation data ( in contrast to our approach ) . Using validation data ( i.e. , held-out data ) is a common technique for improving performance in meta-learning and is particularly important for the few-shot learning case . Another method for deriving a generalization bound on meta-learning is to use PAC-Bayes bounds at both the base and meta levels [ 52 , 53 ] . In [ 5 ] , generalization bounds based on such a framework are provided along with practical optimization techniques . However , the method requires one to maintain distributions over distributions of initializations , which can result in large computation times during training and makes it extremely challenging to numerically compute the bound . Moreover , the approach also does not allow one to incorporate validation data to improve the bound . Recent work has made progress on some of these challenges . In [ 60 ] , the computational efficiency of training is improved but the challenges associated with numerically computing the generalization bound or incorporating validation data are not addressed . State-of-the-art work tightens the two-level PAC-Bayes guarantee , addresses computation times for training and evaluation of the bound , and allows for validation data [ 77 ] . However , all of the two-level PAC-Bayes bounds require a separate PAC-Bayes bound for each task , and thus a potentially loose union bound . We present a framework which , to our knowledge , is the first to combine algorithmic stability and PAC-Bayes bounds ( at the base- and meta- levels respectively ) in order to derive a meta-learning algorithm with associated generalization guarantees . As outlined in Section 1 , we believe that the algorithmic stability and PAC-Bayes frameworks are particularly well-suited to tackling the specific challenges of generalization at the different levels . We also highlight that none of the approaches mentioned above report numerical values for generalization bounds , even for relatively simple problems . Here , we empirically demonstrate that prior approaches tend to provide either near-vacuous or loose bounds even in relatively small-scale settings while our proposed method provides significantly stronger bounds .

Summary: This submission is about generalization bounds for gradient-based meta-learning algorithms. It is motivated by the observations that meta-learning algorithms have demonstrated promising empirical results in a variety of tasks, while our theoretical understanding of these algorithms is lagging behind and in particular the existing generalization bounds for meta-learning are challenging to compute or give vacuous values even in simple settings.  I like how this work reasons about the challenges in generalization for meta-learning. The reasoning is by separately considering generalization at the base level and generalization at the meta level, and how the two levels must be coupled. The work proposes to rely on uniform stability to reason about generalization at the base level, and to use PAC-Bayes bounds to reason about generalization at the meta level. The claimed main contribution is that this approach leads to a novel generalization bound (Theorem 3) for gradient-based meta-learning, which is then leveraged to develop a learning method (PAC-BUS) by using the bound as optimization objective, and the method is evaluated in two meta-learning problems. It is not too clear to me how the two levels are coupled to learn the inductive biases that can be used for future tasks, but I suspect that the PAC-Bayes prior must play a role in connecting the two levels.  The argument behind the proof of the main result (Theorem 3) appears to consist of three blocks: (1) writing the distribution that generates data sets of a given size $m$ as a product distribution between the random tasks and the random samples for given tasks; (2) using the classical PAC-Bayes bound of McAllester for bounding the expected loss (expectation over data sets and over parameters) in terms of its empirical counterpart (that sums over the training tasks) plus a KL regularization term; and (3) using Uniform Stability for bounding the expected loss on a given task in terms of its empirical counterpart and stability. However, I missed an outline of the proof or indication of where the reader can see the full details of the proof.