SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sang Keun Choe $^{1*}$  Sanket Vaibhav Mehta $^{1}$  Hwijeen Ahn $^{1}$  Willie Neiswanger $^{2}$  Pengtao Xie $^{3,5}$  Emma Strubell $^{1,4}$  Eric Xing $^{1,5}$
$^{1}$ Carnegie Mellon University  $^{2}$ Stanford University  $^{3}$ UCSD  $^{4}$ Allen Institute for AI  $^{5}$ MBZUAI
Despite its flexibility to learn diverse inductive biases in machine learning programs, meta learning (i.e., learning to learn) has long been recognized to suffer from poor scalability due to its tremendous compute/memory costs, training instability, and a lack of efficient distributed training support. In this work, we focus on making scalable meta learning practical by introducing SAMA, which combines advances in both implicit differentiation algorithms and systems. Specifically, SAMA is designed to flexibly support a broad range of adaptive optimizers in the base level of meta learning programs, while reducing computational burden by avoiding explicit computation of second-order gradient information, and exploiting efficient distributed training techniques implemented for first-order gradients. Evaluated on multiple large-scale meta learning benchmarks, SAMA showcases up to  $1.7/4.8 \times$  increase in throughput and  $2.0/3.8 \times$  decrease in memory consumption respectively on single-/multi-GPU setups compared to other baseline meta learning algorithms. Furthermore, we show that SAMA-based data optimization leads to consistent improvements in text classification accuracy with BERT and RoBERTa large language models, and achieves state-of-the-art results in both small- and large-scale data pruning on image classification tasks, demonstrating the practical applicability of scalable meta learning across language and vision domains.
Meta learning aims to learn the inductive biases (e.g. training data, neural architecture) of a machine learning program in such a way that a model trained with these inductive biases achieves optimal performance on user-specified objectives (e.g. fairness, quick generalization). This concept of meta learning can naturally be formulated as bilevel optimization, where the upper (meta) level problem encodes inductive biases and objectives, and the lower (base) level optimization problem represents the main machine learning program of interest, such as image classification or language modeling. Depending on the design of inductive biases and objectives, meta learning has found many applications in machine learning, including hyperparameter optimization [16], data optimization [21, 58], neural architecture search [38, 69], learned optimizers [43, 44], and few-shot learning [14, 51].
Following its versatility, numerous algorithms have been proposed to solve meta learning. Among them, gradient-based meta learning (GBML) has in particular gained considerable attention, due to its capability to optimize a wide range of high-dimensional inductive biases in an efficient manner. For example, MAML [14] finds optimal initialization weights (inductive bias) that achieve quick generalization to new tasks (objective), and L2RW [54] optimizes training sample weights (inductive bias) to achieve robustness against label noise (objective). However, the above benefits of GBML oftentimes get overshadowed by its poor scalability in practice, especially under the recent trend of large models [5, 33, 50], which arises due to several factors. First, many GBML algorithms [14, 40, 51] require inversion of a large Jacobian matrix, which suffers from both algorithmic instability as well
*Correspondence: sangkeuc@andrew.cmu.edu
Constant Memory	Jacobian Inverse Free	Adaptive Optimizer Support	(Efficient) Distributed Training Support	Overall Scalability
Iterative Differentiation [14, 15, 42]	X	X	✓	X	X
Recurrent Backpropagation [36]	✓	X	✓	X	X
T1-T2[38, 41]	✓	✓	X	X	✓
Neumann Series [40]	✓	X	✓	X	✓
Conjugate Gradient [51]	✓	X	✓	X	✓
SAMA (ours)	✓	✓	✓	✓	✓
Figure 1: Top: Table showing a scalability comparison. Bottom left: Plot of throughput vs memory of different GBML algorithms on the noisy finetuning of BERT-base experiment. SAMA achieves better memory/compute efficiency overall given a fixed model, and the gap further widens by distributing compute across multiple GPUs with our efficient distributed training strategy. Bottom right: Plot of memory vs model size (i.e., # of parameters) of different GBML algorithms on the continued pretraining of RoBERTa experiment. SAMA demonstrates the least significant increase in GPU memory usage with the increasing model size compared to baseline methods.
as exorbitant compute/memory costs. Second, most GBML research assumes that lower (base) level optimization is performed with SGD, whereas most large models, exemplified by Transformers [62], are by default optimized with adaptive optimizers like Adam [32]; consequently, the applicability of SGD-based GBML methods to large models trained with adaptive optimizers remains unclear. Finally, most GBML research to date has been limited to the single-GPU setup due to the lack of distributed training support [9, 20], which is essential in large-scale learning.
In this work, we endeavor to resolve the aforementioned scalability issues in GBML by co-developing algorithms and systems, and explore the initial potential of scalable meta learning in diverse applications. Our main contributions can be summarized as follows:
1. We investigate the root causes of substantial memory/compute costs, algorithmic instability, and a lack of distributed training support in GBML, all of which significantly limit its scalability, through a technical analysis of implicit differentiation. In doing so, we identify three major factors, namely (i) base Jacobian inversion, (ii) a lack of algorithmic adaptation for adaptive optimizers, and (iii) a need for the custom implementation of the backward pass of meta gradient computation, and discuss how each of them contributes negatively to the above limitations in depth.
2. Taking one step further, we propose an initial solution to each of the aforementioned issues by respectively (i) approximating the base Jacobian with an identity matrix, (ii) additionally expanding the meta Jacobian via the chain rule, and (iii) devising a novel communication strategy that efficiently uses the communication-computation overlap trick [35]. Combining all these solutions, we develop SAMA, a holistic and practically ScalAble Meta learning Algorithm.
3. We evaluate the scalability and the overall performance of SAMA on a multitude of large-scale meta learning benchmarks involving large language models (e.g., BERT [31] and RoBERTa [39]) or large datasets (e.g., ImageNet-1k [10]). Notably, SAMA showcases up to  $1.7 / 4.8 \times$  increase in throughput and  $2.0 / 3.8 \times$  decrease in memory consumption respectively on single-/multi-GPU setups compared to other baseline meta learning algorithms. In addition, we observe that SAMA-based data optimization consistently leads to improvements in text classification accuracy with large language models, and achieves state-of-the-art results in both small-/large-scale data pruning, demonstrating the initial potential of scalable meta learning.
2 Background: Gradient-Based Meta Learning
We begin by reviewing the basics of (gradient-based) meta learning in order to establish the key aspects that have limited its scalability. Mathematically, meta learning is commonly formulated as bilevel optimization as follows:
$$
\lambda^ {*} = \operatorname * {a r g m i n} _ {\lambda} L _ {m e t a} (D _ {m e t a}; \theta^ {*} (\lambda))
$$
$$
s. t. \theta^ {*} (\lambda) = \underset {\theta} {\operatorname {a r g m i n}} L _ {b a s e} (D _ {b a s e}; \theta , \lambda)
$$
where  $\lambda$  (respectively,  $\theta$ ) are the parameters of meta (base) learners,  $D_{meta}$  ( $D_{base}$ ) are meta (base) datasets, and  $L_{meta}$  ( $L_{base}$ ) are meta (base) loss functions. An important implication of the above formulation is that meta learning changes the task of finding the optimal inductive biases from designing heuristics to designing meta optimization problems. As an example, consider the problem of finding the optimal inductive bias for fair classification given class-imbalanced training data. A traditional approach to this problem is to use a heuristic that reweights training samples inversely proportional to class frequencies. On the contrary, L2RW [54] designs a meta optimization problem by curating a small number of class-balanced data for the meta dataset  $D_{meta}$  and setting the meta learner  $\lambda$  to be importance weights for all training data. In short, unlike heuristic-based methods that explicitly specify "how to learn," meta learning methods only specify "what to learn" and let the meta learner automatically determine "how." From a programming paradigm perspective, such a difference can be understood as a transition from imperative to declarative programming.
While there are multiple approaches to solving meta learning, in this work we focus on gradient-based approaches due to their ability to efficiently solve high-dimensional meta optimization (i.e.  $\dim(\lambda) \gg 1$ ) problems. Such an ability is essential given that the search space for inductive biases (e.g. importance weights for training data) can increase exponentially with recent large models and datasets. Concretely, GBML computes a meta gradient composed of two terms—the best-response Jacobian and direct gradient—with the chain rule, as follows:
$$
\frac {\partial L _ {m e t a}}{\partial \lambda} = \underbrace {\frac {\partial \theta^ {*}}{\partial \lambda}} _ {\text {b e s t - r e s p o n s e J a c o b i a n}} \cdot \underbrace {\frac {\partial L _ {m e t a}}{\partial \theta^ {*}}} _ {\text {d i r e c t g r a d i e n t}} \tag {1}
$$
Since the direct gradient (teal) computation is straightforward with the underlying automatic differentiation library, the major challenge in GBML lies in computing the best-response Jacobian (purple), of which two common solutions are iterative differentiation [14, 15, 16, 42] and implicit differentiation [25, 40, 49, 51]. Between these two, in this paper we adopt implicit differentiation as our baseline solution to GBML, as it achieves better computation and memory efficiency than iterative differentiation [19], both of which are vital in accomplishing our goal of scalable GBML.
The gist of implicit differentiation is that it calculates the best-response Jacobian by leveraging Cauchy's Implicit Function Theorem (IFT) and re-interpreting the base optimization problem from the perspective of fixed-point iteration given an iterative solver  $u$ , as follows:
$$
\frac {\partial \theta^ {*}}{\partial \lambda} = - \underbrace {\frac {\partial u}{\partial \lambda}} _ {\text {m e t a J a c o b i a n}} \cdot \left(\underbrace {\frac {\partial u}{\partial \theta^ {*}}} _ {\text {b a s e J a c o b i a n}}\right) ^ {- 1} \quad \text {w h e r e} \left\{\begin{array}{l}\theta^ {*} = \lim  _ {t \rightarrow \infty} \theta_ {t}\\\theta_ {t} = \theta_ {t - 1} - u \left(\theta_ {t - 1}; \lambda\right)\end{array}\right. \tag {2}
$$
While exact implicit differentiation requires solving the base optimization problem to convergence  $\theta^{*}$  by repeatedly applying an iterative solver  $u$  (e.g., SGD or Adam) to calculate base (blue) and meta (red) Jacobians, this is computationally impractical, especially in most large-scale learning settings. Therefore, researchers oftentimes approximate  $\theta^{*}$  with a small number of unrolled update steps of  $u$ . This results in a solution that alternates gradient descent between base and meta optimization problems, where the base gradient is calculated with standard backpropagation and the meta gradient with Eqs. (1) & (2). Noting that many techniques have been developed to efficiently perform and scale up standard backpropagation, we deduce that the major challenges in scaling GBML lie in meta gradient computation, which will be discussed in depth in the next section.
3 Scaling Meta Learning
It has long been recognized that meta gradient computation in GBML suffers from a substantial compute/memory cost [40, 51], algorithmic instability [2, 12], and a lack of efficient distributed
training support [3, 6, 9], all of which can significantly limit its scalability. In this section, we first attempt to understand the above limitations at a technical level. Toward this end, we investigate three aspects in Eqs. (1) & (2), namely (i) base Jacobian inversion, (ii) algorithmic adaptation for adaptive optimizers, and (iii) a need for the custom implementation of meta gradient computation, and discuss how they lead to the aforementioned limitations. Next, we propose initial solutions for each of these issues, based on which we build a holistic and ScalAble Meta Learning Algorithm, SAMA.
3.1 Base Jacobian Inverse
Problem Denoting the size of the base learner (i.e.,  $\dim(\theta)$ ) as  $n_b$ , the computational complexity of the naive base Jacobian (blue) inversion in Eq. (2) is  $\mathcal{O}(n_b^3)$ . Since such cubic complexity is impractical even for small-sized base learners, practitioners typically utilize various linear systems techniques such as Neumann series [40] or conjugate gradient [51], and directly approximate  $\left(\frac{\partial u}{\partial \theta^*}\right)^{-1} \frac{\partial L_{meta}}{\partial \theta^*}$ . Given that the base Jacobian is a function of the Hessian matrix of the base optimization problem in GBML, these algorithms solve linear systems by iteratively performing Hessian-vector products. However, this Hessian-vector product computation contributes negatively to all three above limitations. First, while many automatic differentiation engines provide an efficient Hessian-vector product implementation like Pearlmutter's algorithm [48], its memory/compute cost is still prohibitive with larger models, as demonstrated in Fig. 1. Second, in most cases, we only have access to a stochastic estimation of the Hessian due to mini-batch sampling [34]. Hence, meta gradients obtained with noisy Hessian-vector products can be biased, which may result in training instability. Finally, the most efficient distributed training features, such as communication-computation overlap [35], are designed for first-order gradients, rather than the higher-order gradients involved in Hessian-vector products.
Solution A simple solution for avoiding the aforementioned issues stemming from base Jacobian inversion is to approximate the base Jacobian with an identity matrix as follows:
$$
\frac {\partial L _ {m e t a}}{\partial \lambda} = - \frac {\partial u}{\partial \lambda} \cdot \left(\frac {\partial u}{\partial \theta^ {*}}\right) ^ {- 1} \cdot \frac {\partial L _ {m e t a}}{\partial \theta^ {*}} \approx - \frac {\partial u}{\partial \lambda} \cdot \frac {\partial L _ {m e t a}}{\partial \theta^ {*}} \tag {3}
$$
Under the deep equilibrium model setting, Jacobian-free backpropagation [17] shows that such an approximation can be understood as preconditioning the original meta gradient. Our approximation also resembles approximating the Hessian as an identity matrix in one-step unrolling techniques (i.e.  $T_{1} - T_{2}$ ) from [38, 41]. While this approximation is exact when the iterative solver  $u$  is naive SGD where  $u = \frac{\partial L_{base}}{\partial\theta}$ , we note that the base Jacobian does not necessarily equate with the Hessian when an adaptive optimizer is used in base optimization (more detailed discussion on this issue is deferred to Sec. 3.2). Furthermore, their methods calculate the meta Jacobian at initialization  $\theta$  instead of at convergence  $\theta^{*}$  due to their close connection to iterative differentiation [42], and thereby inhibit unroll steps larger than 1, unlike our approach. Considering that a larger number of unroll steps allows for less frequent computations of expensive meta gradients, our implicit-differentiation-based derivation can lead to a further computational gain in large-scale meta learning. In Appendix E, we investigate the effect of this identity approximation in the "biased regression" setting where the closed-form solution can be analytically calculated, and empirically show that the identity approximation still allows for accurate estimation of the meta gradient  $\frac{\partial L_{meta}}{\partial\lambda}$  and the optimal meta solution  $\lambda^{*}$ , even when the true base Jacobian is not an identity matrix.
3.2 Algorithmic Adaptation for Adaptive Optimizers
Problem Most existing implicit differentiation algorithms [21, 26, 40] compute the best-response Jacobian in Eq. (1) based on the assumption that the iterative solver  $u$  for base optimization is vanilla SGD, whereas recent large models, exemplified by Transformers [5, 62], are by default trained with adaptive optimizers, such as Adam [32]. While the fixed point condition can be theoretically identical for any gradient-based optimizer at convergence (i.e.,  $\frac{\partial\mathcal{L}_{base}}{\partial\theta^{*}} = 0$ ), researchers in practice approximate  $\theta^{*}$  with a small number of gradient steps, at which the above fixed point condition would unlikely hold. Thus, the inconsistency between assumed and actual optimizers results in an incorrect meta gradient, which is a source of training instabilities and reduced performance in meta learning (as we demonstrate in Table 1).
Figure 2: The overall workflow of meta gradient computation with SAMA in the distributed data parallel setting. In detail, SAMA consists of three first-order backward passes performed with the underlying automatic differentiation engine, and one manual backward pass for algorithmic adaptation for the adaptive optimizer. Gradient synchronization is performed only once in the last backward pass with communication-computation overlap to minimize the communication bottleneck.
Solution To take adaptive update rules into account, we propose to further expand a meta Jacobian term in Eq. (3) with the chain rule as follows:
$$
\begin{array}{l} \frac {\partial u}{\partial \lambda} = \frac {\partial g _ {b a s e}}{\partial \lambda} \cdot \frac {\partial u}{\partial g _ {b a s e}} = \frac {\partial^ {2} L _ {b a s e}}{\partial \lambda \partial \theta^ {*}} \cdot \frac {\partial u}{\partial g _ {b a s e}} \\ \Rightarrow \frac {\partial L _ {m e t a}}{\partial \lambda} \approx - \frac {\partial u}{\partial \lambda} \cdot \frac {\partial L _ {m e t a}}{\partial \theta^ {*}} = - \frac {\partial^ {2} L _ {b a s e}}{\partial \lambda \partial \theta^ {*}} \cdot \frac {\partial u}{\partial g _ {b a s e}} \cdot \frac {\partial L _ {m e t a}}{\partial \theta^ {*}} \tag {4} \\ \end{array}
$$
where  $g_{base}$  is base-gradient computed at convergence (i.e.  $\frac{\partial L_{base}}{\partial\theta^{*}}$ ). In short, we accomplish the algorithmic adaptation for any adaptive optimizer with the update rule  $u$  through the middle term  $\frac{\partial u}{\partial g_{base}}$  in Eq. (4), which reduces to an identity matrix in the case of SGD. To analyze the adaptation cost, we note that parameter updates  $u$  in most optimizers are performed only with element-wise operations. Thus, the adaptation matrix is diagonal, for which computation/memory complexities are only  $\mathcal{O}(n_b)$ . As a concrete example, we provide an adaptation matrix for the most popular Adam optimizer [32] in Appendix C.
Furthermore, to reduce computation/memory complexities involving the costly second-order derivative  $\frac{\partial^2L_{base}}{\partial\lambda\partial\theta^*}$ , we instead perform the matrix-vector product with the central-difference method from DARTS [38] and the associative rule of matrix multiplication. All combined, we propose SAMA, a highly compute/memory efficient meta gradient algorithm for scalable meta learning, the formulation of which is as follows:
$$
\frac {\partial L _ {m e t a}}{\partial \lambda} \approx - \frac {\partial^ {2} L _ {b a s e}}{\partial \lambda \partial \theta^ {*}} \cdot \left(\frac {\partial u}{\partial g _ {b a s e}} \cdot \frac {\partial L _ {m e t a}}{\partial \theta^ {*}}\right) \approx - \frac {\frac {\partial L _ {b a s e} \left(\theta^ {+} , \lambda\right)}{\partial \lambda} - \frac {\partial L _ {b a s e} \left(\theta^ {-} , \lambda\right)}{\partial \lambda}}{2 \epsilon} \tag {5}
$$
where  $\theta^{\pm} = \theta^{*}\pm \epsilon v$  with the perturbation vector  $v = \frac{\partial u}{\partial g_{base}}\cdot \frac{\partial L_{meta}}{\partial\theta^{*}}$  and the step size  $\epsilon = \frac{\alpha}{\|v\|_2}$ . We empirically observe that  $\alpha = 1.0$  works well across a multitude of tasks without further tuning. Finally, we notice that previous work in penalty-based bilevel optimization (e.g. F2SA [34] and BOME [37]) further uses the direct gradient  $\frac{\partial L_{meta}}{\partial\theta^{*}}$  explicitly in the base level optimization to maximize the performance of the final base parameter  $\theta^{*}$  on the meta objective. Given that our perturbation vector  $v$  includes the direct gradient term, we also follow a similar strategy and update the base parameter  $\theta$  in the direction of  $v$  (i.e.  $\theta_{t + 1} = \theta_t - \epsilon v$ ) every time the meta update is performed.
3.3 Efficient Distributed Training & Implementation
Problem In large-scale learning, distributed data parallelism (DDP), which communicates and synchronizes local gradients from each device before the parameter update, is necessary to improve both compute and memory efficiency. However, most automatic differentiation libraries like PyTorch [46] only have native DDP support for their basic backward function, whereas SAMA, similar to other meta gradient algorithms, requires a custom implementation of the backward function, which consists of three basic backward passes as shown in Eq. (5). In addition, while a few meta learning
$^2$ Eq. (5) technically involves four derivatives in total, but the adaptation matrix  $\frac{\partial u}{\partial g_{base}}$  can be calculated analytically without backpropagation.
libraries, such as Betty [6], have preliminary DDP support, they do not offer further communication cost optimization. As a result, meta learning research to date has either been limited to a single-GPU setup or suffered from communication inefficiency.
Solution Since we avoid any explicit computations of second-order gradient information in SAMA, we can utilize various efficient distributed training tricks that have been implemented for first-order gradients. More specifically, to enable efficient DDP in SAMA, we develop a novel communication strategy that performs first two backward passes locally on each device, and then overlaps computation in the final backward pass with communication. In PyTorch, this can be neatly achieved by implementing the first two backward passes with torch autograd grad and the last one with torch autograd backward. The overall workflow diagram of SAMA is presented in Figure 2. To facilitate research in scalable meta learning, we provide our implementation of SAMA with the above communication optimization in  $\mathsf{Betty}^3$  that only requires a one-line change in the configuration.
While few-shot learning has traditionally been the most popular application of meta learning, most recent large models such as GPT [5], ViT [13], and Whisper [50], provide few-shot generalization capability out of the box. Therefore, we in this work focus on another promising application of meta learning, data optimization, where we transform (e.g., reweight, label correct) downstream/pretraining data given the specific objectives in the meta optimization problem. Indeed, there is an increasing number of works originating in data-centric AI that empirically show that the quality of training data significantly affects the final performance of large models [1, 5, 18, 50, 56, 65]. Nevertheless, solutions proposed in these works to improve training data quality mostly rely on hand-designed heuristics, which typically result in suboptimal performance. Given that training data of large models serves as extremely high-dimensional inductive biases in machine learning, we expect that GBML's ability to efficiently optimize high-dimensional inductive biases can be fully unlocked in the data optimization application.