Theoretical bounds on estimation error for meta-learning

Many practical machine learning applications deal with distributional shift from training to testing. One example is few-shot classification (Ravi and Larochelle, 2016; Vinyals et al., 2016), where new classes need to be learned at test time based on only a few examples for each novel class. Recently, few-shot classification has seen increased success; however, theoretical properties of this problem remain poorly understood.

In this paper we analyze the meta-learning setting, where the learner is given access to samples from a set of meta-training distributions, or tasks. At test-time, the learner is exposed to only a small number of samples from some novel task. The meta-learner aims to uncover a useful inductive bias from the original samples, which allows them to learn a new task more efficiently.¹¹1Note that this definition encompasses few-shot learning. While some progress has been made towards understanding the generalization performance of specific meta-learning algorithms (Amit and Meir, 2017; Khodak et al., 2019; Bullins et al., 2019; Denevi et al., 2019; Cao et al., 2019), little is known about the difficulty of the meta-learning problem in general. Existing work has studied generalization upper-bounds for novel data distributions (Ben-David et al., 2010; Amit and Meir, 2017), yet to our knowledge, the inherent difficulty of these tasks relative to the i.i.d case has not been characterized.

In this work, we derive novel bounds for meta learners. We first present a general information theoretic lower bound, Theorem 1, that we use to derive bounds in particular settings. Using this result, we derive lower bounds in terms of the number of training tasks, data per training task, and data available in a novel target task. Additionally, we provide a specialized analysis for the case where the space of learning tasks is only partially observed, proving that infinite training tasks or data per training task are insufficient to achieve zero minimax risk (Corollary 2).

We then derive upper and lower bounds for a particular meta-learning setting. In recent work, Grant et al. (2018) recast the popular meta-learning algorithm MAML (Finn et al., 2017) in terms of inference in a Bayesian hierarchical model. Following this, we provide a theoretical analysis of a hierarchical Bayesian model for meta-linear-regression. We compute sample complexity bounds for posterior inference under Empirical Bayes (Robbins, 1956) in this model and compare them to our predicted lower-bounds in the minimax framework. Furthermore, through asymptotic analysis of the error rate of the MAP estimator, we identify crucial features of the meta-learning environment which are necessary for novel task generalization.

Our primary contributions can be summarized as follows:

• We introduce novel lower bounds on minimax risk of parameter estimation in meta-learning.

• Through these bounds, we compare the relative utility of samples from meta-training tasks and the novel task and emphasize the importance of the relationship between the tasks.

• We provide novel upper bounds on the error rate for estimation in a hierarchical meta-linear-regression problem, which we verify through an empirical evaluation.

An early version of this work (Lucas et al., 2019) presented a restricted version of Theorem 1. The current version includes significantly more content, including more general lower bounds and corresponding upper bounds in a hierarchical Bayesian model of meta-learning (Section 5).

Baxter (2000) introduced a formulation for inductive bias learning where the learner is embedded in an environment of multiple tasks. The learner must find a hypothesis space which enables good generalization on average tasks within the environment, using finite samples. In our setting, the learner is not explicitly tasked with finding a reduced hypothesis space but instead learns using a general two-stage approach, which matches the standard meta-learning paradigm (Vilalta and Drissi, 2002). In the first stage an inductive bias is extracted from the data, and in the second stage the learner estimates using data from a novel task distribution. Further, we focus on bounding minimax risk of meta learners. Under minimax risk, an optimal learner achieves minimum error on the hardest learning problem in the environment. While average case risk of meta learners is more commonly studied, recent work has turned attention towards the minimax setting (Kpotufe and Martinet, 2018; Hanneke and Kpotufe, 2019, 2020; Mousavi Kalan et al., 2020; Mehta et al., 2012). The worst-case error in meta-learning is particularly important in safety-critical systems, for example in medical diagnosis.

Mousavi Kalan et al. (2020) study the minimax risk of transfer learning. In their setting, the learner is provided with a large amount of data from a single source task and is tasked with generalizing to a target task with a limited amount of data. They assume relatedness between tasks by imposing closeness in parameter-space (whereas in our setting, we assume closeness in distribution via KL divergence). They prove only lower bounds, but notably generalize beyond the linear setting towards single layer neural networks.

There is a large volume of prior work studying upper-bounds on generalization error in multi-task environments (Ben-David and Borbely, 2008; Ben-David et al., 2010; Pentina and Lampert, 2014; Amit and Meir, 2017; Mehta et al., 2012). While the approaches in these works vary, one common factor is the need to characterize task-relatedness. Broadly, these approaches either assume a shared distribution for sampling tasks (Baxter, 2000; Pentina and Lampert, 2014; Amit and Meir, 2017), or a measure of distance between distributions (Ben-David and Borbely, 2008; Ben-David et al., 2010; Mohri and Medina, 2012). Our lower-bounds utilize a weak form of task relatedness, assuming that the environment contains a finite set that is suitably separated in parameter space but close in KL divergence—this set of assumptions also arises often when computing i.i.d minimax lower bounds (Loh, 2017).

One practical approach to meta-learning is learning a linear mapping on top of a learned feature space. Prototypical Networks (Snell et al., 2017) effectively learn a discriminative embedding function and performs linear classification on top using the novel task data. Analyzing these approaches is challenging due to metric-learning inspired objectives (that require non-i.i.d sampling) and the simultaneous learning of feature mappings and top-level linear functions. Though some progress has been made (Jin et al., 2009; Saunshi et al., 2019; Wang et al., 2019; Du et al., 2020). Maurer (2009), for example, explores linear models fitted over a shared linear feature map in a Hilbert space. Our results can be applied in these settings if a suitable packing of the representation space is defined.

Other approaches to meta-learning aim to parameterize learning algorithms themselves. Traditionally, this has been achieved by hyper-parameter tuning (Rasmussen and Nickisch, 2010; MacKay et al., 2019) but recent fully parameterized optimizers also show promising performance in deep neural network optimization (Andrychowicz et al., 2016), few-shot learning (Ravi and Larochelle, 2016), unsupervised learning (Metz et al., 2019), and reinforcement learning (Duan et al., 2016). Yet another approach learns the initialization of task-specific parameters, that are further adapted through regular gradient descent. Model-Agnostic Meta-Learning (Finn et al., 2017), or MAML, augments the global parameters with a meta-initialization of the weight parameters. Grant et al. (2018) recast MAML in terms of inference in a Bayesian hierarchical model. In Section 5, we consider learning in a hierarchical environment of linear models and provide both lower and upper bounds on the error of estimating the parameters of a novel linear regression problem.

Lower bounding estimation error is a critical component of understanding learning problems (and algorithms). Accordingly, there is a large body of literature producing such lower bounds (Khas’ minskii, 1979; Yang and Barron, 1999; Loh, 2017). We focus on producing lower-bounds for parameter estimation using local packing sets, but expect that extending these results to density estimation or non-parametric estimation is feasible.

Most existing theoretical work studying out-of-distribution generalization focuses on providing upper-bounds on generalization performance (Ben-David et al., 2010; Pentina and Lampert, 2014; Amit and Meir, 2017). We begin by instead exploring the converse: what is the best performance we can hope to achieve on any given task in the environment? After introducing notation and minimax risks, we then show how these ideas can be applied, using meta linear regression as an example.

A full reference table for notation can be found in Appendix A and a short summary is given here. We consider algorithms that learn in an environment $(Z,P)$, with data domain $Z=X\times Y$ and $P$ a space of distributions with support $Z$. In the typical i.i.d setting, the algorithm is provided training data $S\in Z^k$, consisting of $k$ i.i.d samples from $P\in P$.

In the standard multi-task setting, we sample training data from a set of training tasks $\{P_1,\dots ,P_{M+1}\}\subset P$. We extend this to a meta-learning, or novel-task setting by first drawing $S_{1:M}$: $n$ training data points from the first $M$ distributions, for a total of $nM$ samples. We call this the meta-training set. We then draw a small sample of novel data, called a support set, $S_{M+1}\in Z^k$, from $P_{M+1}$.

Consider a symmetric loss function $\ell (a,b)=\psi \left(\rho \right(a,b\left)\right)$ for non-decreasing $\psi$ and arbitrary metric $\rho$. We seek to estimate the output of $\theta :P\to \Omega$, a functional that maps distributions to a metric space $\Omega$. For example, $\theta \left(P\right)$ may describe the coefficient vector of a high-dimensional hyperplane when $P$ is a space of linear models, and $\rho$ may be the Euclidean distance.

In this section, we first present our most general result: Theorem 1. Using this, we derive Corollary 1 that gives a lower bound in terms of the sample size in the training and novel tasks. Corollary 1 recovers a well-known i.i.d lower bound (Theorem 2) when $Mn=0$, and, importantly, highlights that the novel task data is significantly more valuable than the training task data. Additionally, we provide a specialized bound that applies when the environment is partially observed — proving that in this setting training task data is insufficient to drive the minimax risk to zero.

In Theorem 1, we assume that $P$ contains $J$ distinct $2\delta$-separated distributions but only $M+1\le J$ tasks are visible to the learner. Intuitively, the error rate lower-bound shrinks as the amount of information shared between the training tasks and the novel task grows. All proofs are given in Appendix B.1. Recall $\ell (a,b)=\psi \left(\rho \right(a,b\left)\right)$ for non-decreasing $\psi$ and arbitrary metric $\rho$.

Note that $\delta$ is a property of the so-called packing set, $J$, and may depend on the sample size, $\beta$, and other properties of $P$. For example, practical instances of this bound typically require $\psi \left(\delta \right)=O\left(1/k\right)$ or similar, as in Theorem 3 below. To derive this result, we bound the statistical estimation error by the error on a corresponding decoding problem where we must predict the novel task index, given the meta-training set $S_{1:M}$ and $S_{M+1}$. Fano’s inequality provides best-case error rates for this problem.

Using Theorem 1, we derive our first bound on the novel-task minimax risk that depends on the number of meta-training tasks ($M$) and datapoints per training task ($n$, $k$), via a local-packing argument. The following corollary implies that if we have $J$ meta-training tasks in our $2\delta$-packing that are close (in terms of their pairwise KL distance), then learning a novel task from training samples drawn from the meta-training tasks requires significantly more examples; in particular, learning the novel task from samples drawn from the meta training set requires $\Omega \left(J\right)$ times the sample complexity of the novel task. This matches our intuition that learning the novel task implies the ability to distinguish it from all $J$ well-separated meta-training tasks.

Our goal is to analyze the sample complexity of meta-learning for linear regression, where samples are drawn from multiple meta-training tasks and we want to generalize to a new task with only a few data points. After introducing the setting, we will compute lower-bounds on the minimax risk using our results from Section 4, revealing a $2^d$ scaling on the meta-training sample complexity. Following the lower bound, we derive an accompanying upper-bound on the risk of a MAP estimator, derived from an empirical Bayes estimate over a hierarchical Bayesian model. Asymptotic analysis of this bound reveals that if the observed samples from the novel task vary considerably more than the task parameters, then observing more meta-training samples may significantly improve convergence in the small $k$ regime. This is validated empirically in Section 6.

For $i=\mathrm{1...}M+1$, where $M+1$ is the total number of tasks, we define,

Each task has some design matrix $X_i$ and unknown parameters ${\theta }_i$. For simplicity, we assume known isotropic noise models and that $n_i=n$ for all $i\le M$, with $n_{M+1}=k$.

Our meta learner will fit the data using an empirical Bayes estimate in a hierarchical Bayesian model:

We will consider the Maximum a Posterior estimator,

and will characterize its risk, $E[\parallel {\hat{\theta }}_{M+1}-{\theta }_{M+1}{\parallel }_2^2]$, where the expectation is with respect to sampled data only. The posterior distribution under the Empirical Bayes estimate for $\tau$ is given in Appendix C.2. The derivation is standard but dense and we recommend dedicated readers to consult Gelman et al. (2013), or an equivalent text, for more details.

In this section, we provide additional quantitative exploration of the upper bound studied in Section 5. The aim is to take steps towards relating the bounds to experimental results; we know of little theoretical work in meta-learning that attempt to relate their results to practical empirical datasets. Full details of the experiments in this section can be found in Appendix D.

Meta-learning algorithms identify the inductive bias from source tasks and make models more adaptive towards unseen novel distribution. In this paper, we take initial steps towards characterizing the difficulty of meta-learning and understanding how these limitations present in practice. We have derived both lower bounds and upper bounds on the error of meta-learners, which are particularly relevant in the few-shot learning setting where $k$ is small. Our bounds capture key features of the meta-learning problem, such as the effect of increasing the number of shots or training tasks. We have also identified a gap between our lower and upper bounds when there are a large number of training tasks, which we hypothesize is a limitation of the proof technique that we applied to derive the lower bounds — suggesting an exciting direction for future research.

This work benefited greatly from the input of many other researchers. In particular, we extend our thanks to Shai Ben-David, Karolina Dziugaite, Samory Kpotufe, and Daniel Roy for discussions and feedback on the results presented in this work. We thank Ahmad Beirami and anonymous reviewers for their valuable feedback that led to significant improvements to this paper. We also thank Elliot Creager, Will Grathwohl, Mufan Li, and many of our other colleagues at the Vector Institute for feedback that greatly improved the presentation of this work. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute (www.vectorinstitute.ai/partners).