Making Scalable Meta Learning Practical
Authors: Sang Keun Choe
Sanket Vaibhav Mehta
Hwijeen Ahn
Willie Neiswanger
Pengtao Xie
Emma Strubell
Eric Xing
Venue: NeurIPS 2023 poster (2023)
Keywords: meta learning
bilevel optimization
large-scale learning
implicit differentiation
TL;DR: We enable efficient scaling of meta learning to large models and datasets.

Abstract: 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.8x increase in throughput and 2.0/3.8x 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.

Making Scalable Meta Learning Practical
Sang Keun Choe
Sanket Vaibhav Mehta'
Hwijeen Ahn'
Willie Neiswanger
Pengtao X.e
Emma Strubell
Eric Xing
'Carnegie Mellon University Stanford University
UCCD
4 *Allen Institute for AI
MMBZUA
Abstract
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 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 designed flexibly support broad range of adaptive optimizers in the base level 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- meta learning benchmarks, SAMA showcases up 1.7/4.8x x increase in throughput and 2.0/3.8x decrease in memory consumption respectively single-/multi-GHie setups compared to other baseline meta learning algorithms. Furthermore, we show that SAMA-basee data optimization leads to consistent improvements in text classification accuracy with BERT and RoBERTa large language models, and achieves state-of-the-ar 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.
Introduction
Meta learning aims to learn the inductive biases (e.g. training data, neural architecture) of a machine learning program in such a way that 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 [88) neural
architecture search 38 learned optimizers [3] 44 and few-shot learning [4
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 wide range of nigh-dimensional inductive biases in an efficient manner. For example, MAML [4] finds optimal initialization weights (inductive bias) that achieve quick generalization to new tasks (objective), and L.2RW (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
Coorespondence: sangkeuc@andreu www.errwwww.......mmm edu
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Constant Memory
Jacobiam
Adapiive Optimizeer Support
(Ellicient) Distributed Training Support
Overall Scalability
Iterative Differentiation
T
15 42
K
X
Recurrett Backpropagation
36
x
Ti
T2
BS
EN
X
Noumann Scries
FD
Conjugate Gradient
5
SAMA (ours)
Noisy Finetuning w/ BERT-base
Continued Protraining w/ RoBERTa Variants
50
500000 Xux Kunvyy 168C0)
Neumann
RoBERTIa large
Throughput (samgless
350
SAMM-NA
40
SAMA
300
SAMA (4 GPUs)
SAMA CPUs)
Memory Usage
30
250
Kn-
200
better
20
10
Noumana
100
SAMA
25
50
100
150
200
250
300
350
100
GPU Memory Usage (GB)
Parameter Count (M)
Figure 1: Top: Table showing 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 ( 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 applica- tions. 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 technical analysis of implicit differentiation. In doing so, identify three major factors, namely (i) base Jacobian inversion, (ii) a lack of algorithmic adaptation for adaptive optimizers, and (ifi) 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.
Taking one step further, we propose an initial solution to each of the aforementioned issues by respectively (i) approximating the base Jacobian with identity matrix, (ii) additionally expanding the meta Jacobian via the chain rule, and (iii) devising a novel communication strategy that efficiently uses the communication-comme
overlap trick [35]. Combining all these solutions, we develop SAMA, a holistic and practically ScalAble Meta learning Algorithm.
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 [O]). Notably, SAMA showcases up to 1.7/4.8x increase in throughput and 2.0/3.8 x decrease in memory consumption respectively on single-/multi-GPP setups compared to other baseline meta learning algorithms. addition, we observe that SAMA- based data optimization consistently leads to improvements in text classification accuracy with large language models, and achieves oate-of-the-are results in both small-/large-scal 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:
X argmin Lmeea (Dmetai 8 (X))
s.l. 0" (A) argmin Lease! Dbases 0, X)
0
where X (respectively, 0) are the parameters of meta (base) learners, Dmata (Drase) are meta (base) datasets, and T-meto (Ieasee 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-imbalang traditional approach to this problem 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 dlass-balanced data for the meta dataset Metta and setting the meta learner À to be importance weights for all training data. In short, unlike heuristic-base methods that explicitly specify "how to learn," " meta learning methods only specify "what learn" and let the meta learner automatically determine "how." From a programming paradigm perspective, such difference can be understood as a transition from imperative to declarative programming.
training data. A
While there are multiple approaches to solving meta learning, in this work we focus on gradient- based approaches due their ability to efficiently solve high-dimensional meta optimization (i.e. dim(A) >> 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 meta gradient composed of two terms the best-response Jacobian and direct gradient with the chain rule, as follows:
aLmete Ox
00"
Ommet
aX
a0
best response Jacobian
grudient
Since the direct gradient (teal) computation is straightforward with the underlying automatic dif- ferentiation library, the major challenge in GBML lies in computing the best-response Jacobian (purple), of which common solutions are iterative differentiation [14 15 16 42
and implicit differentiation 255 40 49 (51 Between these two, in this paper we adopt implicit differentiation
as our baseline solution to GBML, it achieves better computation and memory efficiency than iterative differentiation 199. 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 's Implicit Function Theorem (IFT) and re- re-interpreting the base optimization problem from the perspective of fixed- weed-point iteration given iterative solver u. follows:
00 0X
Ou 0x
Ou
where {
{ 0" lim Oz
10000
30
00=00-1-- =
mota Jacobian
base Jacobian
While exact implicit differentiation requires solving the base optimization problem to convergence 6* 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 0* 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. ( & 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
has long been recognized that meta gradient computation in GBML suffers from substantial compute/memory cost [40] 51), algorithmic instability [2]. and lack of efficient distributed
3

---

training support [ 6 19, 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. (T) & 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( as ng. the computational complexity of the naive base Jacobian (blue) inversion in Eq. is (n Since such cubic complexity is impracti- cal even for small-sized base learners, practitioners typically utilize various linear systems techniques
such as Neumann series [40] or conjugate gradient [5T), and directly approximate Ou
Given that the base Jacobian is a function of the Hessian matrix of the base optimization problem in GBML, these algorithmss solve linear systems by iteratively performing Hessian-vectte products. However, this Hessian-vecto product computation contributes negatively to all three above limita- tions. First, while many automatic differentiation engines provide efficient Hessian-vector product implementation like Pearlmutter's 's algorithm [88. its memory/compute cost is still prohibitive with larger models, as demonstrated in Fig. [ Second, in most cases, we only have access stochastic estimation of the Hessian due to mini-batch sampling (34). Hence, meta gradients obtained with noisy Hessian-ved veetorr products can be biased, which may result in training instability. Finally, the most efficient distributed training features, such as communicaiion-comptaaaii designed for first-order gradients, rather than the higher-order gradients involved in Hessian-vector products.
overlap [651, are
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:
o ammeta 0>
Ou ax
Du 20-
aLmeta 00"
Du
iLmeta
ax
00
Under the deep equilíbrium 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 - T2) from [38 41). While this approximation is exact when the iterative solver u is naive SGD ahaaes
where u = 30 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 0 instead of at convergence 0* due to their close connection 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-differ dilieiittiiifeeeetaiiioo.oo barereeentttiio-bbbeedded derivation can lead to a further computational gain in large-scale meta learning. In Appendix 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 Almata and the optimal meta solution X", even when the true base Jacobian is not an identity matrix.
[E]
8x
3.2 Algorithmic Adaptation for Adaptive Optimizers
Problem Most existing implicit differentiation algorithms 21 26 compute the best-response Jacobian in Eq. (I) based on the assumption that the iterative solver for base optimization is vanilla SGD, whereas recent large models, exemplified by Transformers 5 62,, are by default trained with adaptive optimizers, such Adam [32). While the fixed point condition can be theoretically identical for any gradient-based optimizer at convergence (i.e., DCeces = 0), researchers in practice 08-
approximate 0* 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 meta learning (as we demonstrate in Table Ip.

---

Device
SLaar (0'.))
de
Diare
syschroniizaion comm-comp
overiap
Aaani
Blesn (0".)
dLic
Device
ar'
01
da
local backward
global backward
manual backward
deta
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 compuuattiono overlap to minimize the communication bottleneck.
Figure 2:
Solution To take adaptive update rules into account, we propose further expand Jacobian term in Eq. 4 with the chain rule as follows:
Ou
Oghase
Ou
02 Liaase
Ou
0A
OA
0gbase
0A00~
dghase
ammeta
Du
aLmeta
Liass
Ou
aLmeta
OX
00"
0200 Oybase 80*
where Ibase is base-gradien computed at convergence (i.e. DLrass) ). In short, accomplish the 88
algorithmic adaptation for any adaptive optimizer with the update rule u through the middle term Oghaen in Eq. ( 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 computationnmmemory complexities only O(n)). As a concrete example, we provide an adaptation matrix for the most popular Adam optimizer 32 in Appendix
Furthermore, to reduce computation/memort complexities involving the costly second-order deriva-
8 Li we instead perform the matrix-vector product with the central-difference method from tive DI
DARTS 38 and the associative rule of matrix multiplication. All combined, we propose SAMA, highly compute/memorr efficient meta gradient algorithm for scalable meta learning, the formulation of which as follows:
OLmeta 0%
8 Liase
du Jygase
aLmeta 20"
L A) 9A
aLaace (0_ .A)
83
0x00
2f
where 0 0" - eu with the perturbation vector V
Doanta and the step size e Jf
Zigaane
|e
We empirically observe that a = 1.0 works well across a multitude of tasks without further tuning. Finally, notice that prevoius work in penalty-based bilevel optimization (e.g. F2SA 34 and BOME (77) further the direct gradient aLmate explicitly in the base level optimization to maximize the performance the final base parameter 6* on the meta objective. Given that our perturbation vector includes the direct gradient term, also follow similar strategy and update the base parameter in the direction of (i.e. 0,+1 = O2 Ew) 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, necessary to improve both compute and memory efficiency. However, most automatic differentiation libraries like Py- Torch [46] only have native DDP support for their basic backward function, whereas SAMA, similar to other meta gradient algorithms, requires custom implementation of the backward function, which consists of three basic backward passes as shown Eq. 5 In addition, while a few meta learning
Pgq. (5) technically involves four derivatives total, but the adaptation matrix
can be calculated
aghase
analytically without backpropagation.
5

---

libraries, such as Betty [6], have preliminary DDP support, they do not offer further communication cost optimization. As result, meta learning research to date has either been limited to a single-GPU setup or suffered from communication inefficiency.
Since 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 com- putation 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 ad. backward. The overall workflow diagram of SAMA is presented in Figure [ To facilitate research in scalable meta learning, we provide our inplementation of SAMA with the above communication optimization Betty that only requires a one-line change in the configuration.
Solution
Experiments
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 optinization, where we transform (e.g., reweight, label correct) www.strream/prerriinne data given the specific objectives 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 [ 5 solutions proposed in these works to improve training data quality mostly rely hand bund-designed heuristics, which typically result in suboptimal performance. Given that training data of large models serves extremely nigh-dimensional inductive biases in machine learning. expect that GBML's ability to efficiently optimize high-dimensional inductive biases can be fully unlocked in the data optimization application.
18 50, 56 665. Nevertheless,
From a technical perspective, large-scale data optimization has substantial compute/memory cost and frequently involves models that are trained with adaptive optimizers. Therefore, it serves as an ideal benchmark to (a) evaluate the scalability of SAMA compared to existing approaches, (b) study the effectiveness of each component in SAMA, and (c) investigate the practical usefulness of scalable meta learning across diverse domains. Specifically, in this section, we consider three data optimization applications, namely: Noisy finetuning of large language models (Sec. 4.1 continued
pretraining of large language models (Sec. (.2). and scale-agnose agnostic data pruning (Sec. 4.3 While not an application of data optimization, we also include a preliminary analysis the effect of model size
on few-shot image classification accuracy in Appendix D Finally, we note that all experiment details including baselines, hyperparameters, and compute resources are provided in Appendix B
4.1 Noisy Finetuning of Large Language Models
Weak supervision 5
proposes to achieve a significant reduction in the data labeling cost by letting users quickly generate labels for a large amount of data by exploiting multiple weak labeling functions, such as hand-designed rules and other neural networks. While increasing the quantity of labeled data with weak supervision has led to noticeable improvements in multiple applications, a poor quality of generated labels results in a degradation in test accuracy, leaving room for further improvement. Here, we attempt to alleviate this data quality issue in weak supervision by utilizing meta learning to automatically optimize noisy training data guided by a small amount of clean data in the meta level. In detail, we use data reweighting (58) and label correction (70] as our data optimization operations, of which the bilevel optimization formulation is as follows:
X = (Ar, Ac) = argmin L(Delean; 8" (Arr Ac)) Arcde
S.L. 0" Ao) = argmin
Duaisyy (a.s)EDnooixx
D
where wn::; Xx.) and c(:; Ax) are meta learners, respectively, for data reweighting and label correction. To evaluate the scaling efficiency of SAMA, we perform text classification with a BERT-base model
htttps//ggitaab..o///epprrrr

---

with 110M parameters on multiple weak supervision datasets from the WRENCH benchmark [67]] Furthermore, to study the effectiveness of our algorithmic adaptation strategy (Sec. 3.2 we conduct experiments with a variant of SAMA (i.e., SAMA-NA) that does not include algorithmic adaptation, and present the experiment results in Table [
Algarithm
TREC
SemEvva
IMDB
ChemnPol
AGNews
Yelp
Finetune (orig)
67
66.56 (2.31)
83,93
79.73 (2.60)
56.09 (1.08)
86.27 (0.53)
82.26 (3,50)
COSINE 66
76.56 (0.08)
86.80 (0.46)
82.98 (0.05)
58.47 (008)
87.03 (0.00)
89.22 (0.05)
Finetune (ours)
67.93 (2.55)
79.78 (1.78)
78.16 (2.28)
57.35 (1.43)
85.79 (0.49)
84.32
SAMA-NA
74.33 (2.34)
87.11 (1.11)
(1.74)
60.88 (0.00)
86.83 (0.19)
80.96 (3.04)
+R & C
SAMA-NA
(2.62)
87.67 (1.36)
80:44 (0.97)
64,05 (0.52)
87.05 (0.39)
80.73 (3.11)
SAMA
85.73
89.67 (0.67)
84.31
76.89 (.339)
89. 05 (0.
93. 23.64
+R &
SAMA
87.93(1.17)
38.83.(2.32)
85.71 (0.82)
77.78 (0.59)
89.79 (0.27)
93.77 (0.08)
Table I: WRENCH results. R and C in the first column stand for data reweighting and label correction operations. The number in parentheses indicates standard deviation for each experiment over 3 runs.
In this experiment, we are able to make two observations. First, we notice that SAMA-based data reweighting and label correction both lead to noticeable improvements over finetuning and self- training (COSINE [66]) baselines in noisy text classification accuracy of the large Transformer model across all benchmarks with the help of small additional clean data D Delean in the meta level. Second, given the superior performance of SAMA compared to SAMA-NA, the importance of algorithmic adaptation when an adaptive optimizer is used to train this Transformer base learner.
empirically verify
Additionally, we compare compute/memory efficiency of SAMA to that of two other implici-diiferrnniation--ssss meta learning algorithms, namely Neumann Series [40] and
GPUs
Mennoty
Throughpul
Neumann
26.0
82.9
conjugate gradient [51]. For fair comparison, evaluate
CG
I
28.4
82.
GPU memory usage (MB) and throughput (samples/sec) on the AGNews dataset from Wrench with fixed global batch size of 48, and summarize the result in Table Three observations that can be made here are that (1) SAMA is generally more compute/memory efficient than the base- lines, (2) the cost of algorithmic adaptation for adaptive optimizers is marginal as expected, and (3) the efficiency gap further widens as we distribute compute/memory across multiple GPUs with our efficient DDP communication strategy. Since we were only able to conduct experiments with up to 4 V100 GPUs due to the limited compute resources, exploring extremely large-scale GBML with larger GPU servers remains a promising future research direction. In Appendix [ we provide more extensive ablation study for each component of SAMA, as well as compare test accuracy, GPU memory usage, and throughput of SAMA against various meta learning algorithms on IMDB and AGNews datasets from the Wrench benchmark.
SAMA-NA
I
13.7
144.1
SAMA
1
14.3
142.0
SAMA
2
10.4
241.2
SAMA
396.7
Table 2: Memory and throughput analy- sis on AGNews with 4 V100 GPUs.
4.2 Continued Pretraining of Large Language Models
DAPT/TAPT 24 empirically demonstrate that additional pretraining (i.e., continued pretraining) of the generic language model on the domain or task task-specific data can further improve downstream performance on diverse benchmarks. However, the inclusion of low-quality samples for continued pertaining tasks can potentially hinder pretraining by amplifying negative interference 164], which could lead suboptimal downstream performance. Here, we attempt to minimize such negative transfer by reweighting samples from the continued pretraining task with meta learning. To this end, we adopt the auxiliary learning technique from TARTAN [1] and simplify the two-stage pretraining- finetuning pipeline into a one-stage multitask learning pipeline with the reweighting scheme applied to the pretraining loss. The bilevel optimization formulation is as follows:
X argmin L (A))
1 Σ
s.t. 0- (2) = argmin Cr 0) + Dpm De
w(I; A) (x;
7

---

loss functions, Dr are finetunin/pprerainingg datasets, and w(: A) the data reweighting network. Following the experiment setup in TARTAN [T]. we use task-specific data and a masked language modeling loss in our auxiliary task and perform experiments with RoBERTa-base on 4 datasets from the original DAPT/TAPT paper. We compare our SAMA- based data optimization against DAPT and TARTAN-MT. We exclude TAPT and TARTAN- respectively because (1) TAPT consistently underperforms TARTAN-MT [9] and (2) TARTAN-
where Cr are
uses additional validation data in the meta level of the downstream tasks, making the comparison unfair. We report our experiment results in Table
ChemProt
HyperPartisan
ACL-ARC
SciERC
Average
Baseline
82.70 (0.45)
89.03 (2:25
68.17 (2:52)
79.83 (0.89)
79.93
DAPT PE]
84.17 (0.50)
87.23 (3x5)
71.84 (4.78)
80.42 (1.57)
80.92
TARTAN-MT
84, 18 (0.30)
94.64 (0.91)
72.41 (1.94)
80.83 (0.71)
83.02
SAMA (ours)
84.49 (0.13)
95.18 (0.00)
71.63 (1.68)
81.84 (0.05)
83.29
Table 3: Experiment results for auxiliary learning with the continued pretraining task. Following 24]. we report test miero-F1 for ChemProt and macro-F1 for the other datasets. The number in parentheses indicates the standard deviation for each experiment over 3 runs.
As shown above, SAMA-based data optimization leads to improvements in downstream performance on almost all datasets. This indirectly demonstrates that SAMA-based data reweighting can identify more/less relevant data in the auxiliary task and accordingly up-ldown-weight them, unlike TARTAN- MT which allocates equal importance weights on all auxiliary data. Therefore, we expect that our method would likely benefit from additional auxiliary data by automatically figuring out and exploiting only relevant data, whereas TARTAN-M is much more susceptible to negative transfer. While we only used task-specific data in our auxiliary task for the fair comparison with TARTAN-MT, extending auxiliary data to domain-specific or even general text data and comparing SAMA against DAPT or TARTAN-MT would be an intriguing future research direction. Finally, we analyze the GPU memory usage of different-sized RoBERTa in this experiment and present the result in Figure| []] The figure clearly shows the superior memory efficiency of SAMA with the increasing model size.
4.3 Scale-Agnostic Efficient Data Pruning
Data pruning 47 59 61 63 has recently received the limelight in the machine learning community as means to both improve training efficiency and reduce (semantic) redundancy in training data. In particular, Sorscher et al. (59) showed both theoretically and experimentally that neural scaling laws can be beaten by data pruning. Nevertheless, they point out that the optimal data pruning metric varies across different dataset scales and further research in scalable data pruning metrics is needed. Here, we propose to forgo hand-designed data pruning metrics, and rather automatically meta-learn the importance weight of each training data following Meta-Weight-No modifications. First, we replace their iterative differentiation meta gradient algorithm with SAMA to achieve improved memory/compute efficiencies. Second, we further speed up meta learning by enabling distributed training with our efficient communication strategy. Third, we use the uncertainty input to MWN to better estimate importance
(MWN) (58) with four major
of the prediction in addition to the loss value weight of each training data. Last, we use training data both in the base and the meta levels, assuming no additional validation data. A bilevel optimization formulation of our method is as follows:
X = argmin C(Derain; 8" (2))
s.t. 8" (A) = argmin [Dtram (2.3))
Σ
w(C, U; A) L(x, 0)
where w(::>) A) is MWN that takes the loss value L and the uncertainty 'i of the training sample (æ, y) as input and outputs the importance weight. Under this setup, run meta learning with SAMA for 30 / 50 epachs respectively for ImageNet- / CIFAR-10 and obtain the pruning metrics by averaging the importance weights of the last 5 epochs. We compare method to several popular static/dynamic data pruning baselines, and present the results in Figure 3
As expected, GBML-based data pruning with SAMA not only outperforms heuristics-based data pruning but also works well across different dataset scales. Surprisingly, we observe that GBML.- based data pruning even leads to improvements in test accuracy at the pruning ratio of 0.1 and 0.2
8

---

ImageNet-1ics ResNet-50
CIEAR-10 ResNet-18
1.00
1.00
Relative Aecuracy
Relalive Top-1, Accurscc
0.99
0.99
0.98
0.97
Random
0.98
0.96
EL2N
Random
GrodND
0.95
0.97
forgelling
0.94
DynaMS
SAMA
SAMA
manga:
0.1
0.1
0.2
Praning Ratio
Pruning Retio
EL2N,
EL2N
GradNd
DynaMS
SAMA (ours)
Relative Search Time
0.16
1.65
1.65
0.16
0.46
Figure 3: Top Left: ImageNet- Ik data pruning results with ResNet-50. Reported numbers are relative accuracy compared full training accuracy (i.e., pruned_ ceoedddcc/fulllaeeeeeeeeee poc/ffullaacc......... Accuracy for other baseline methods is obtained from DynaMS [63] Top Right: CIFAR-10 data pruning results with ResNet- 18. Accuracy for other baseline methods obtained from Deepcore [22]. Bottom: Relative time spent in finding data to prune compared to full ImageNet- Ik training time.
on ImageNet- 1k. The potential implication is that ImageNet-1K may have noisy labels or semantic redundancy and that GBML is able to automatically figure and filter out these samples. Further in-depth investigation of filtered data remains interesting research direction. Considering that compute/memory inefficiency has traditionally been the major bottleneck in GBML applications, we also compare the relative search time for data pruning. Our result shows that SAMA demonstrates comparable or even shorter search time than heuristics-brr -based methods. We also note that, while the original MWN 588 encounters the OOM error under our setup of batch_ natch_size=256, the throughput analysis with the reduced batch size reveals that efficient distributed training with SAMA on 4 GPUs achieves x speed up compared to the original MWN that lacks distributed training support.
Related Work
Algorithms Two major lines of research in muudient-based meta learning algorithms are iterative and implicit differentiation [19]. Iterative differentiation [14 16 42 computes meta gradients by differentiating through the optimization path and therefore requires saving all intermediate states on the path. This makes the memory/compute costs of iterative differentiation increase linearly with the number of unrolling steps. While the linear cost can be avoided with the use of techniques like truncated backpropagation [36). it is still more expensive than that of implicit diffcrentiation. in which meta gradient computation is independent of the length of the optimization path. More specifically, meta gradient computation in implicit differentiation depends only on the final state of the optimization path. To compute base Jacobian inversion in implicit differentiation, multitude of variants have been proposed, each of which uses Neumann series 7 40], conjugate gradient [51], Nystrom method (25), and more. While generally being more compute/memory efficient than iterative differentiation, most existing implicit differentiation algorithms have poor scalability due to the issues studied in Sec. 3
Applications Meta learning has found many applications in machine learning including few-shot
learning 51] neural architecture search (38 69]. hyperparameter optimization [15] 16 40 42)
data optimization [2T 54 58 [70) and reinforcement learning (23 29 52 to name a few. Notably,
most of these applications share the underlying mathematical formulation of bilevel optimization and are conceptually related to optimal design and inductive/deellaative programming paradigms.
Systems Compared to algorithms and systems research, there are relatively fewer research efforts in meta learning systems. an attempt to facilitate research few-shot image classification, higher (20), learn2learn (3) and TorchMeta [9] have been developed. However, due to their specific focus on few-shot image classification, these libraries have not been actively used in other meta learning tasks, such as data optimization neural architecture search. Recently, software libraries

---

for implicit differentiation including JaxOpt 4 and Betty [6], have been proposed. Given that Betty's software architecture is specifically designed support various systems optimization for large-scale meta learning, we chose to implement SAMA in this framework.
6 Conclusion
In this paper, we strived to make scalable meta learning practical via both algorithmic and systems advancements. Towards this goal, we investigated diverse scaling bottlenecks in meta learning at technical level and resolved them by developing SAMA. Tested on multiple benchmarks, SAMA empirically demonstrated its scaling efficiency well its capability to optimize variety of high-dimensional inductive biases large-scale learning. In future work, we plan to explore two directions. First, given that training extremely large models with 10B+ parameters require various systems techniques such as model/pipeline parallelism or optimizer sharding, extending SAMA to be compatible with these techniques would be highly important for further scalability. Second, we will also focus on large-scale meta learning application research, such as neural architecture search for Transformer-familn fammly models. Overall, we hope that work serve stepping stone for lot of interesting scalable meta learning research to come.
Limitations & Broader Impacts While SAMA has demonstrated significantly improved com- pute/memory efficiencies, and stably worked with fairly large models like BERT/RoBERTa, we could not test it larger models with 1B+ parameters due lack of computational resources. Meta learning itself is mostly value-neutral, but there a chance that practitioners amplify the bias or toxicity of the machine learning programs by enforcing them in the meta objective.
Acknowledgements
We thank all the reviewers for their invaluable comments and feedback. We like to acknowledge CMU Workhorse and TIR group for providing compute resources for this work. EX and SKC acknowledge the support of NSF IIS23 1990, NSF IIS2123952 NSF CNS2008248, NSF BCS2040381, NGA RMM4762010002, NIGMS ROIGM 40467, and Amazon 2023-5007107. WN was supported in part by NSF (# 15555)), AFOSR 01555019100244 ARO (W91 1NF-211-111101 CZ Biohub, and Sloan Fellowship. ES and SVM acknowledge the support part by DSO National Laboratories.
References
[1] Amro Kamal Mohamed Abbas, Kushal Tirumala, Daniel Simig, Surya Ganguli, and Ari S Moreos. Semdedup: Data-efficient learning web-scale through semantic deduplication. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
[2] Antreas Antoniou, Harrison Edwards, and Amos Storkey. How train your MAML. In International Conference on Learning Representations, 2019.
[3] Sébastien MR Arnold, Praateek Mahajan, Debajyoti Datta, Ian Bunner, and Konstantinos Saitas Zarkias. learn2learm: A library for meta-leaming research. arXiv preprint arXhv:2008.12284 2020.
[4] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares- López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. Advances in neural information processing systems, 35:5230-5242, 2022.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[6] Sang Keun Choe, Willie Neiswanger, Pengtao Xie, and Eric Xing. Betty: An automatic differentiation library for multilevel optimization. In The Eleventh International Conference on Learning Representationss 2023.
[7] Ross M Clarke, Elre Talea Oldewage, and José Miguel Hermández-Lobato. Scalable one-pass optimisation high-dimensional weight-update hyperparameters by implieit differentiation. In International Conference Learning Representations, 2022.
10

---

[8]
Cody Coleman, Christopher Yeh, Stephen Mussmann, Baharan Mirzasoleiman, Peter Bailis, Percy Liang, Jure Leskovec, and Matei Zaharia. Selection via proxy: Efficient data selection for deep learning. In International Conference on Learning Representations, 2020.
[9] Tristan Deleu, Tobias Würfl,, Mandana Samiei, Joseph Paul Cohen, and Yoshua Ben- gio. Torchmeta: A Meta-Learning library for PyTorch, 2019. Available at: https://iiibbbooottttaaadduuuunp
[10] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. leee, 2009.
[11] Lucio M. Dery, Paul Michel, Ameet Talwalkar, and Graham Neubig. Should we be pre-training? an argument for end-task aware training as an alternative. In International Conference on Learning Representations, 2022.
[12] Guneet Singh Dhillon, Pratik Chaudhari, Avinash Ravichandran, and Stefano Soatto. A baseline for few-shot image classification. In International Conference on Learning Representations. 2020.
[13] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[14] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In International conference on machine learning, pages 1126-110 1135.135. PMLR, 2017.
[15] Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil. Forward and reverse gradient- based hyperparameter optimization. In International Conference on Machine Learning, pages 1 1165-1173. PMLR, 2017.
[16] Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-learning. In International Conference on Machine Learning, pages 1568-1577. PMLR, 2018.
[17] Samy Wu Fung. Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley Osher, and Wotao Yin. Jfb: Jacobian-fres backpropagation for implicit networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 6648-6656, 2022.
[18] Samir Yitzhak Gadre, Gabriel Ilharco, Alex Fang, Jonathan Hayase, Georgios Smyriss, Thao Nguyen, Ryan Marten, Mitchell Wortsman, Dhruba Ghosh, Jieyu Zhang, et al. Datacomp: In search of the next generation of multimodal datasets. arXiv preprint arXiv:2304. 1404.1418,, 2023.
[19] Riccardo Grazzi, Luca Franceschi, Massimiliano Pontil, and Saverio Salzo. On the iteration complexity of hypergradient computation. In International Conference on Machine Learning, pages 3748-3758. PMLR, 2020.
[20] Edward Grefenstette, Brandon Amos, Denis Yarats, Phu Mon Htut, Artem Molchanov, Franziska Meier, Douwe Kiela, Kyunghyun Cho, and Soumith Chintala. Generalized inner loop meta- learning. arXiv preprint
2019.
[21] Denis Gudovskiy, Luca Rigazio, Shun Ishizaka, Kazuki Kozuka, and Sotaro Tsukizawa. Autodo: Robust autoaugment for biased data with label noise via scalable probabilistic implicit dif- ferentiation. In Proceedings of the JEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16601-16610, 2021.
[22] Chengeheng Guo, Bo Zhao, and Yanbing Bai. Deepcore: A comprehensive library for coreset selection in deep learning. In Database and Expert Systems Applications: 33rd International Conference, DEXA 2022, Vienna, Austria, August 22-24, 2022, Proceedings, Part I, pages 181-195. Springer, 2022.
[23] Abhishek Gupta, Russell Mendonca, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning of structured exploration strategies. Advances in neural information processing systems, 31, 2018.
[24] Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A. Smith. Don't stop pretraining: Adapt language models to domains and tasks. In Proceedings of ACL, 2020.
11

---

[25] Ryuichiro Hataya and Makoto Yamada. Nyström method for accurate and scalable implicit differentiation. In International Conference on Artificial Intelligence and Statisties, pages 1643.4654. PMLR, 2023.
[26] Ryuichiro Hataya, Jan Zdenek, Kazuki Yoshizoe, and Hideki Nakayama. Meta approach to data augmentation optimization. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 2574-2583, 2022.
[27] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[28] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe V Le. Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
[29] Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint vrXiv:1905.06424 2019.
[30] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXhv:2001.0836,, 2020.
[31] Jacob Devlin Ming-We Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACI-HILT, pages 4171-4186, 2019.
[32] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint Drriiv//422 69800, 2014.
[33] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan- Yen Lo, al. Segment anything. arXiv preprint arXiw:2304.00644. 2023.
[34] Jeongyeol Kwon, Dohyun Kwon, Stephen Wright, and Robert D Nowak. A fully first-order method for stochastic bilevel optimization. In International Conference on Machine Learning. pages 18083-18113, PMLR, 2023.
[35] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, and Soumith Chintala. Pytorch distributed: experiences accelerating data parallel training. Proceedings of the VLDB Endowment, 13:3005-3018, 08 2020.
[36] Renjie Liao, Yuwen Xiong, Ethan Fetaya, Lisa Zhang, KiJung Yoon, Xaq Pitkow, Raquel Urta- sun, and Richard Zemel. Reviving and improving recurrent pack-propagation. In International Conference on Machine Learning, pages 3082-3091. PMLR, 2018.
[37] Bo Liu, Mao Ye, Stephen Wright, Peter Stone, and qiang liu. BOME! bilevel optimization made easy: A simple first-order approach. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[38] Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search.
In International Conference on Learning Representations, 2019.
[39] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint 1907. 1692, 2019.
[40] Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference Artificial Intelligence and Statistics, pages 1550-1552. PMLR, 2020.
[41] Jelena Luketina, Mathias Berglund, Klaus Greff, and Tapani Raiko. Scalable gradient- tuning of continuous regularization hyperparameters. In International conference on machine learning, pages 2952-2960. PMLR, 2016.
[42] Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based hyperparameter opti- mization through reversible learning. In International conference on machine learning, pages 2113-2122 PMLR, 2015.
12

---

[43]
Luke Metz, James Harrison, Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury. Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, et al. Velo: Training versatile learned optimizers by scaling up. arXiv preprint arXiv:221 1eXiv:2211.0966// 1.09760, 2022.
[44] Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In International Conference on Machine Learning, pages 4556-4565. PMLR, 2019.
[45] Luke Metz, Ben Poole, David Pfau, and Jascha Sohl-Dicksteen. Unrolled generative adversarial networks. In International Conference on Learning Representations, 2017.
[46] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- nigh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[47] Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. Deep learning on a data diet: Finding important examples early in training. Advances in Neural Information Processing Systems, 34:20596-20077, 2021.
[48] Barak A Pearlmutter. Fast exact multiplication by the hessian. Neural computation, 6(1):147- 160, 1994.
[49] Fabian Pedregosa. Hyperparameter optimization with approximate gradient. In International conference on machine learning, pages 737-746. PMLR, 2016.
[50] Alec Radford, Jong Wook Kim, Tao Xu, Greg Brockman, Christine McLeavey, and Ilya Sutskever. Robust speech recognition via large-scale weak supervision. In International Conference on Machine Learning, pages 28492-28511 28518. PMLR, 2023.
[51] Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. Advances in neural information processing systems, 32, 2019.
[52] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International conference on machine learning, pages 5331-5340 PMLR, 2019.
[53] Alexander J Ratner, Christopher M De Sa, Sen Wu, Daniel Selsam, and Christopher Ré. Data programming: Creating large training sets, quickly. Advances in neural information processing systems, 29, 2016.
[54] Mengye Ren, Wenyuan Zeng, Bin Yang, and Raquel Urtasun. Learning to reweight examples for robust deep learning. In International conference on machine learning, pages 4334 4343.
PMLR, 2018
[55] David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-proy pookk-rppaggaing errors. nature, (62(888)) 1986.
[56] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade W Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa R Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmar- czyk, and Jenia Jitsev. LAION-5b: An open large-scale dataset for training next generation image-text models. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022.
[57] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint 1r0X.vv:900.0000
2019.
[58] Jun Shu, Qi Xie, Lixuan Yi, Qian Zhao, Sanping Zhou, Zongben Xu, and Deyu Meng. Meta- weight-net: Learning an explicit mapping for sample weighting. Advances in neural information processing systems, 32, 2019.
[59] Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. Advances in Neural Information Processing Systems, 35:19523-19536, 2022.
[60] Richard Sutton. The bitter lesson. Incomplete Ideas (blog), 13(1), 2019.
[61] Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Bengio, and Geoffrey J. Gordon. An empirical study of example forgetting during deep neural network learning. In International Conference Learning Representations. 2019.
13

---

[62] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[63] Jiaxing Wang, Yong Li, Jingwei Zhuo, Xupeng Shi, WEIZHONG ZHANG, Lixing Gong, Tong Tao, Pengzhang Liu, Yongjun Bao, and Weipeng Yan. DynaMS: Dyanmic margin selection for efficient deep learning, In The Eleventh International Conference on Learning Representations, 2023.
[64] Zirui Wang, Zihang Dai, Barnabás Póczos, and Jaime Carbonell. Characterizing and avoiding negative transfer. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 1 11293-11302, 2019.
[65] Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. arXiv preprint arXiv:2302.03169 2023.
[66] Yue Yu, Simiao Zuo, Haoming Jiang, Wendi Ren, Tuo Zhao, and Chao Zhang, Fine-tuning pre- trained language model with weak supervision: A contrastive-regulaaized self-training approach. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1063-1077, 2021.
[67] Jieyu Zhang, Yue Yu, Yinghao Li, Yujing Wang, Yaming Yang. Mao Yang. and Alexander Ratner. WRENCH: A comprehensive benchmark for weak supervision. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
[68] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems,
2020.
[69] Miao Zhang, Steven W Su, Shirui Pan, Xiaojun Chang. Ehsan M Abbasnejad, and Reza Haffari. idarts: Differentiable architecture search with stochastic implicit gradients. In International Conference on Muchine Learning, pages 12557-12566. PMLR, 2021.
[70] Guoqing Zheng, Ahmed Hassan Awadallah, and Susan Dumais. Meta label correction for noisy label learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 11053-11061, 2021.
[71] Nicolas Zucchet. Simon Schug, Johannes Von Oswald, Dominic Zhao, and João Sacramento. contrastive rule for meta-learning. Advances in Neural Information Processing Systems, 35:25921-25936, 2022.
14

---

A Philosophies behind SAMA & Scalable Meta Learning
Here, we additionally discuss several important design principles and philosophies behind scalable meta learning and SAMA in Q&A format.
Q. Why do we study scalable meta learning?
A. Richard Sutton points out in his article "The Bitter Lesson" [60] that machine learning algorithms that stand the test of time are ones that continue to scale gracefully with the increased computation budger (i.e., scalable algorithms). Given that meta learning is an important topic in machine learning with many applications including data optimization [54], hyperparameter optimization [16], few-shot learning [T). and adversarial learning [45), was natural call for to investigate the scalability of meta learning algorithms following the spirit of "The Bitter Lesson". Interestingly, such focus on the scalability meta learning algorithms distinguishes our work from most other meta learning works, in which the typical focus is to improve the overall performance of meta learning algorithms under limited computation budgel (usually bounded by a single GPU).
Q. What are the design principles behind scalable meta learning?
A. The increased computation budget powered by hardware advancements (e.g., Moore's 's law) has evolved new ecosystem of large models and datasets machine learning over time, which in- volves both systems and algorithms components. For example, to efficiently leverage the increased computation for large-scale learning, diverse systems techniques, such as date//odel/pipellnn par- allelism have been developed [28] 35 57.. At the same time, researchers have devised various algorithms that are highly effective for large-scale learning, such as backpropagation [55], skip con- nections (27), Adam optimizer (32), self-attention (62), etc. Accordingly, addition to guaranteeing memory/compute efficiency for scalability, our major design principle for scalable meta learning was to ensure compatibility with existing systems and algorithms in the large-scale learning ecosystem.
Systems compatibility Given that a great deal of systems support in machine learning, such as communication-comme computtaionoo overlap R5,, has been developed for first-order gradient methods, avoid- ing explicit computations of higher-order gradient information including Hessian-vector products was important design principle in SAMA. Even though we mostly explored distributed training in this work, SAMA is also compatible with other system features such as half-precision training and activation checkpointing, which could further improve memory efficiency.
Algorithms compatibility While there exist several meta learning algorithms that avoid the com- putation of higher-orde gradient information [34] 38 37), many of these algorithms either assume the use of naive SGD update rule or devise specific update rules tailored to their own algorithms at the base level, significantly hampering their algorithm compatibility. In contrast, SAMA allows for the use of arbitrary optimizers at the base level via algorithmic adaptation.
15

---

B Experiment Details
In this section, discuss various experiment details such hyperparameters, baselines, and compute resources used for our experiments in Section 4. Our experiment codes are available the supplementary material as well under the official Betty GitHub repository [
B.1 Noisy Finetuning of Large Language Models
Hyperparameters We ran training for 1000 iterations TRECSSmmEE///MDDBCCienPOOOK. AGNews datasets from the WRENCH benchmark [67]. with a batch size of 32, a weak supervision algorithm of majority voting, and the hyperparameters in Table below.
optimizer
Ir scheduler
wdecay
step
WRENCH train set (with majority voting)
Basc
BERT-base
Adam
1o-5
cosine
10
1.0
Meta
(Reweight)
2-layer MLP
Adam
lo-5
None
WRENCH
Meta
(Correct)
2-yye ML.
WRENCH set
Table 4: Hyperparameters for noisy finetuning of large language models experiments.
Baselines
We adopted naive finetuning and self-training (i.e., COSINE [66]
approaches from the
original WRENCH benchmark paper [67 as our baseline.
Compute Resources We used NVIDIA RTX 2080Ti GPU for the main experiment, and 4 NVIDIA Tesla V100 GPUs for the throughput-memory analysis in Table 2 and Figure
B.2 Continued Pretraining of Large Language Models
Hyperparameters ran training for epochs with batch size of 16, a maximum sequence length of 256, and the hyperparameters in Table 5 below.
model
optimizer
Init
Ir schocller
welecay
detaset
unnoll SAMA
Bose (Downstream)
RoBeRTh-Aase
Adam
2e-5
linear decay wammup linear (wammup propurtion
train split of ChessPotHHypsPatiian ACL-ARUSSIERC
train split of ChemPPollyyssttrtiaan ACL ARCCSSiEERRCC
linear decay wannup (wammup proportion 0.6)
(Auxiliary)
split ChesProdRypss ACL-ARCSciERC
2-layer
Adam
NIA
Table 5: Hyperparameters for continued pretraining of large language models experiments.
Baselines We adopt DAPT 24 and TARTAN-MT 11 as our baselines for this experiment. In detail, DAPT (24] performs additional masked language model pretraining on comain-specifie data on top of the pretrained RoBERTa-bar base model and then finetunes the model on the downstream text classification task. We follow [24] (see Table 14 in the original paper) for setting downstream finetuning hyperparameters. Alternatively, TARTAN-MT [I]] performs masked language modeling with task specific data and downstream text classification training simultaneously multitask fashion through two different heads.
Compute Resources We used NVIDIA Tesla V100 GPU for the main experiment, and NVIDIA RTX A6000 GPU for the "memory vs model-size analysis" in Figure 1.
B.3 Scale-Agnostic Efficient Data Pruning
Hyperparameters We ran meta learning for 30 epochs with a batch size of 256 and the configuration shown in Table beloo.. After pruning data based on the meta learning result, we ran ImageNet- lk
http::///hthhhcooooooooee/e//e// examples
16

---

training for 120 epochs with the learning rate decayed by 10 at epochs [40, 80] following the set up
DynaMS 63
model
optimizer
init_rr
Ir.scheduler
wddecay
dataset
unroll step
SAMA
Base
ResNet-50
SGD
le-1
None
le.4
Tmage
Meta
2-aayer MLP
Adam
le-5
None
0
ImageNet-Ik Ik train set
NA
NA
Table 6: Hyperparameters for ImageNet-IK data pruning experiments
For the CIFAR-10 data pruning experiment, we ran meta learning for 50 epochs with batch size of 128, and configuration in Table After pruning the data based on the meta learning result, we ran CIFAR-10 training for 200 epochs with the cosine learning rate decay schedule following the setup in DeepCore 22
model
optimizer
init_rr
Ic_schoduler
wdeeay
dataset
unroll step
SAMA
Base
ResNet-18
SGD
le-1
None
Se-4
CIFAR- 10 train set
Meta
2-aayer MI.P
Adam
1o-5
None
CIFAR 10 main
N/A
N/A
Table Hyperparameters for CIFAR-IO data pruning experiments
Baselines We adopt EL2N 47).
GraNd [47, DynaMS 63 our baselines for the ImageNet-
Ik experiments and GraNd 47). forgetting [6], margin g for the CIFAR-I0 experiments. In
detail, EL2N/GraND 47 respectively select samples with large 1.2-loss/rradient-non norm values, forgetting (47 chooses samples that are frequently forgotten during training, and margin [8 chooses samples with least confidence. While these baselines considered static pruming. DynaMS [63] falls under the category of dynamic pruning where data to be pruned change during training. Dynamic pruning may the whole training data across different epochs, making fair comparison difficult. Surprisingly, despite being static pruning algorithm, SAMA-based data pruning still achieves a better performance than DynaMS.
Compute Resources We used NVIDIA Tesla V100 GPUs for Imagenet- Ik data pruning meta learning experiments and NVIDIA RTX 2080Ti GPU for CIFAR-10 experiments.
Additional Information We measured the uncertainty via the difference between the predictions of the current model and the mow.lll--oooogaaveee model.
17

---

C Algorithmic Adaptation for Adam Optimizer
Since the Adam optimizer has been the most popular optimizer train large models, exemplified by Transformers [68], here we provide the adaptation matrix for Adam. We denote the first and second moments of the gradient Adam as m and respectively, and the learning rate as 7.
Wadam
8u
3im + (1 31)g
7
Ög
dg
1/10 + (1 (1-B11) B1)g² +
(1 B1) (1 -22) + (1 31) //0 // (1- 32)g²
7
V + +1 (1 31)g² 3)) BI (1- 31)g² +e)² +e)²
(1 - 31) (1 32m
2 7
(because <
V B10 (1 81)9² V 310 (1 81)g² +c)
Adaptation matrices can be similarly derived for other adaptive optimizers.
D
The Effect of Scaling in Model-Agnostic Meta Learning
myriad of algorithms have been proposed to improve few-shot image classification while assuming fixed network architecture. In contrast, here we shift our focus from the algorithm the scale, and propose to study the following question: "Leveraging the compute/memory efficiency of SAMA, can we improve the few-shot generalization capability by scaling up the network size?"' Since SAMA a variant implicit differentiation, we closely follow the experiment setup in iMAML [51) where proximity to the initialization weights is explicitly enforced by Ig-regularization. The major difference is that iMAML uses a snjjugatt-gradient method, which requires second-order gradient information to compute meta gradients, while adopt SAMA to achieve improved scaling to larger networks with its superior memory/compute efficiency. We conduct preliminary experiments the Omniglot 20-way way 1-/5-shot tasks with the basic 4-layer CNN architecture, while varying the width (hidden size) of the networks to study the effect of the model size the few-shot classification accuracy. The experiment results are provided in Figure below.
Since the inception of MAML
Ornniglot: 20-Way 1-Shot 4-layer
Omniglot 20-Way 5-Shot 4-layer CNN
100.0
99.5
95
99.0
Accuracy
98.5
93
98.0
92
97.5
90
23
2
27
28
2
2
2
96.5
25
2
2
2
2
2
Width (Hidden Size)
Width (Hidden Size)
Figure 4: Few-shot image classification accuracy Omniglot 20- 20-way 1-/5-shot tasks with varying network sizes.
Interestingly, we observe that the increased model size leads to consistent improvements in few-shot classification accuracy. The important question following this observation is "can we apply scaling laws (30 from other tasks (e.g., langugge modeling) to general meta learning beyond few-shot image classification?" Since meta learning involves two optimization problems (meta and base) unlike traditional machine learning problems, it is of now unclear how to define the general concept of "scale" in terms of both model and dataset sizes. We expect that further research in this direction would be critical in systematically studying scalable meta learning.
18

---

E Justification of the Identity Approximation for Base Jacobian
In this section, we aim study the empirical effeet of our identity approximation for base Jacobian on the meta gradient Mometa and the optimal meta solution X", when the true base Jacobian is an identity matrix. As obtaining the closed-form solution of the Hessian is impossible in almost all deep learning problems, we study the soundness the identity approximation of base Jacobian in the simpler "biased regression" setting [19], for which the bilevel optimization formulation is as follows:
8A
= arg min X' wt (A) 1/112
w" (A) arg min X y/1²
3
X|2
Given the above formulation, the closed form solutions for the base Jacobian, the meta-gradient 9x+ and the optimal meta solution A are:
1. Base Jacobian x 31
2. ΩA 3(x X + (X X'w" - x where w² = (x (x)+ -
3. A = (A" A)- where A = BX= b x
We set 3 = 0. and perform 100 meta updates, and measure 1) the cosine similarity between the ground truth d사 and the meta gradient obtained with our approximation (i.e. (SAMA), and 2) the L.2 distance between the current meta parameter X; and the optimal solution X" each time step t. For more thorough analysis, we also compute these two metrics for other meta gradient algorithms that explicitly approximate base Jacobian inverse with conjugate gradient and Neumann series. In Figure we provide the metric obtained from all time steps.
gradient cosibe similarty
1.0
12 optimal solution
Noi Naimann
SAMA
0.8
0.6
0.4
0.2
0.0
updote stops
update stops
Figure 5: Left: cos(ga, 9approz); where dΛ is a ground truth (i.e. closed-form) meta gradient, and Japproz is approximate meta gradient obtained with various popular meta learning algorithms including SAMA. Right: || 12 where X, is the meta parameter after meta updates.
From Figure 5 it can be clearly seen that 1) while slightly less accurate than second-order algorithms like CG, SAMA still achieves a high directional alignment with the ground truth meta-gradient, and 2) SAMA also achieves stable convergence to the optimal solution at comparable speed.
'We used smaller than in the original paper [9 (1 vs 0.1), to amplify the "non-identitines"" of the base Jacobian.
19

---

Extensive Ablation Study
In this section, we perform an extensive ablation study that studies the effectiveness of each component in SAMA (i.e. base Jacobian inverse, algorithmic adaptation for the adaptive optimizer, and efficient distributed training) by comparing test accuracy, GPU memory usage, and throughput against various baseline meta learning algorithms, Our ablation study is performed on AGNews and IMDB datasets from the Wrench benchmark 677, and the result presented below.
Base Jacobian
Algo Adapt
Distributed
Accuracy
Throughput
Memory
Finetuning meta learning baseline)
169.16
7.77
Iterative Diff (e.g. MAML) [44 +2
85.78
28.07
22.94
Conjugate gradient (e.g. iMAML) 51
K
X
86.78
65, 14
22.03
Noumann scries
86.65
67.03
19.70
DARIS for 2 22)
38
41
X
X
86.36
43.69
10,81
137.90
SAMA
89.05
134.56
11.12
SAMA (2 GPUs)
226.27
SAMA (4 GPUs)
89.02
298.28
6.46
Table 8: Ablation results on AGNews
Jacobian
Algo Adapt
Accuracy
Throughput
Memory
Finetuning (po meta learning buseline)
x
16
144.39
6.60
Tterative Diff (e.g. MAML.) [H]] [2]
80.25
24.24
22.03
Conjugate gradient (e.g. iMAML) 51
81.01
$6.27
21.92
Neumann
19.75
DARTS (or Ti T
33 41
80.47
37.53
10.35
SAMA-NA
X
81.92
117.86
9.93
SAMA
84.31
116.94
10.84
SAMA (2 GPL))
D
0
85.18
196.48
7.84
SAMA GPUs)
84.19
263.74
6.39
Table 9: Ablation results on IMDB
From extended ablation study, it can be seen that 1) identity approximation of base Jacobian significantly improves memory/compute efficiency, 2) algorithmic adaptation improves meta learning performance the minimal compute/memorr cost, and 3) communicationopiimizze distributed training further improves compute/memory efficiency.
20


Reviews:
[{"review_id": "rHXVgBeHpT", "reviewer": "Reviewer_2sgt", "summary": "The paper addresses the challenges of scalability in meta learning by introducing SAMA, a novel approach that combines advancements in implicit differentiation algorithms and systems. SAMA demonstrates improvements in computational efficiency and memory consumption compared to other baseline algorithms, and it showcases practical applicability in language and vision domains through experiments with large language models and image classification tasks.", "questions": "I would appreciate it if the authors address the following questions:\n1. How does the inconsistency between the assumed vanilla SGD optimizer and the actual optimizer used in large models affect the computation of the meta gradient? Could you provide more insights into the inconsistencies and their consequences in terms of training instabilities and reduced performance in meta learning?\n2. The paper mentions that solutions proposed in data-centric AI works to improve training data quality often rely on hand-designed heuristics. Could you elaborate on the specific limitations or drawbacks of existing approaches based on hand-designed heuristics? How does the proposed meta learning approach address or overcome these limitations? \n3. In section 4.1, are there any specific techniques or algorithms employed within the meta learning framework to optimize the noisy training data? How does the proposed approach leverage meta learning to adaptively update and improve the quality of the labels generated by weak labeling functions?\n4. What are the key differences between vanilla SGD and adaptive optimizers like Adam in terms of their impact on the fixed point condition (mentioned in section 3.2 line 149)?", "limitations": "The authors have provided a sufficient discussion of the limitations of their work in section 6. It would be beneficial for the authors to further elaborate on these aspects, beyond language models, to ensure a comprehensive analysis of any potential societal impacts.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Originality: The paper introduces a novel approach to address the scalability challenges of meta learning. By combining advancements in implicit differentiation algorithms and systems, it proposes a solution that supports arbitrary optimizers in meta learning while reducing computational burden.\n2. Quality: The paper demonstrates a level of quality in terms of the overall idea and experimental evaluation. By evaluating their method on a few benchmarks, including language models and image classification tasks, the authors provide a good assessment of its performance.\n3. Clarity: The key contributions, such as the introduction of SAMA and its evaluation of various benchmarks, are presented in a concise manner. The background and related work sections also help in properly positioning the paper in the literature.\n4. Significance: The paper addresses the long-standing challenge of scalability in meta learning. It is able to tackle a variety of high-dimensional inductive biases of large-scale learning - for instance in optimizing large language models as a very relevant open problem to the community.", "weaknesses": "In general, the main weaknesses that I noticed are in the design of the experimental section:\n1. Comparison with State-of-the-Art: The paper could benefit from a more comprehensive comparison with state-of-the-art meta learning approaches. Is it possible to compare it to well-established meta-learning algorithms, such as MAML? \n2. Lack of ablations: More importantly, the paper seems to lack an ablation study. Including such comparisons would provide a more thorough understanding of the model's effectiveness and highlight more clearly its advantages over existing methods. It would be appreciated if the authors identify separate components of their method and ablate them w.r.t the full method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your positive review and valuable comments. We strive to address concerns and questions that you raised below and in the **global response**.\n\n### **Ablation Study & SOTA comparison**\nThough we didn’t explicitly frame it as an “ablation study”, the effectiveness of each component in SAMA can be understood from Table 1 & 2 in our paper. Below, we directly discuss each component of SAMA based on Table 1 & 2. We also provide a unified ablation study result table in the global response.\n\n**Base Jacobian inverse**\n\n*tl;dr* Identity approximation of base Jacobian significantly improves memory/compute efficiency (Table 2).\n\nOur baseline algorithms in Table 2 (and Figure 1), Neumann and CG, are both state-of-the-art implicit differentiation meta learning algorithms that attempt to approximate base Jacobian inverse as accurately as possible with multiple Hessian-vector product operations, instead of approximating it with an identity matrix as in SAMA. As shown in Table 2, due to an expensive second-order gradient computation involved in Hessian-vector products, these methods demonstrate much poorer throughput and GPU memory usage than SAMA, both of which are the major bottlenecks in scalable meta learning. In our global response, we additionally demonstrate that this identity approximation also has a minimal impact on accuracy by comparing SAMA against Neumann and CG in terms of accuracy.\n\n**Algorithmic adaptation for adaptive optimizer**\n\n*tl;dr* Algorithmic adaptation significantly improves accuracy (Table 1) at the minimal memory/compute cost (Table 2).\n\nThe main goal of this work is to devise a *(1) memory/compute efficient* meta learning algorithm that *(2) achieves good performance/accuracy*. In Table 1, SAMA consistently achieves better accuracy than SAMA-NA, which lacks algorithmic adaptation. Moreover, Table 2 shows that SAMA achieves a comparable memory/compute efficiency as SAMA-NA. \n\n**Distributed training**\n\n*tl;dr* Both GPU memory usage and throughput improve consistently as computations are distributed across more GPUs (Table 2). \n\nWe agree that it may not be straightforward for readers to clearly understand the effectiveness of each component when the results are spread across two separate tables. Hence, in the global response, we provide a unified table for Wrench experiments, including a comparison with state-of-the-art meta-learning baselines.\n\n### **Questions**\n> Q. How does the inconsistency between the assumed vanilla SGD optimizer and the actual optimizer used in large models affect the computation of the meta gradient? Could you provide more insights into the inconsistencies?\n\nRoughly speaking, once base Jacobian is approximated as identity, the meta gradient formulation follows $g_{meta}=-\\frac{\\partial u}{\\partial \\lambda}\\cdot\\frac{\\partial L_{meta}}{\\partial \\theta^*}=-\\frac{\\partial}{\\partial \\lambda}(u\\cdot\\frac{\\partial L_{meta}}{\\partial \\theta^*})$. Thus, meta gradient descent essentially maximizes the inner product between the base update vector $u$ and meta gradient w.r.t base parameters $\\theta$. This way, performing base updates with $u$ not only decreases the base loss, but also *maximally* decreases the meta loss. We note that the update direction of the base problem $u$ is dependent on its optimizer. Thus, we believe that reflecting this base optimizer choice would lead to improved meta learning performance in the end by better aligning $u$ and $\\frac{\\partial L_{meta}}{\\partial \\theta^*}$.\n\n> Q. The paper mentions that solutions proposed in data-centric AI works to improve training data quality often rely on hand-designed heuristics. \nCould you elaborate on the specific limitations of existing approaches based on hand-designed heuristics? \n\nThe benefits of meta-learning approaches to data-centric AI can be most clearly seen in our dataset pruning experiments (Sec 4.3). Specifically, heuristics-based methods (e.g. EL2N, forgetting) rank each training sample using some heuristics, such as forgetting counts and gradient norm at initialization, *in the hope* that these heuristics would indeed capture the importance weight of each sample. On the other hand, our meta-learning approach directly optimizes the importance weight of each sample in a way that the resulting model trained with these learned importance weights minimizes the original training loss. Indeed, our experiment results in Fig. 3 clearly shows that our meta-learning approach outperforms all heuristics-based methods on both small-/large-scale datasets. \n\n> Q. In section 4.1, are there any specific techniques or algorithms employed within the meta learning framework to optimize noisy training data?\n\nThere are abundant meta-learning works for tackling various data issues (e.g. noisy labels, class imbalance). A few examples are Meta-Weight-Net [50], Learning-to-Reweight [48], and Meta-Label-Correction [60]. We adopt some of the architecture designs from these works, but replace the meta-gradient computation algorithm with SAMA, instead of iterative differentiation or truncated backpropagation, to improve the scalability.\n\n> Q. What are the key differences between vanilla SGD and adaptive optimizers like Adam in terms of their impact on the fixed point condition? \n\nMost gradient-based optimizers, including SGD and Adam, share the same fixed point condition of $\\frac{\\partial L_{base}}{\\partial \\theta^*} = 0$. However, most gradient-based meta learning in practice approximates $\\theta^*$ with only a few gradient updates, and thus this “zero gradient” condition is unlikely met. Therefore, we hypothesize that we in reality need to pay more attention to the alignment between $u$ and $\\frac{\\partial L_{meta}}{\\partial \\theta}$ discussed above, where $u$ is dependent on the base optimizer.\n\nWe hope our response resolved most of your concerns, and helped you evaluate our work more positively. If you have other comments, we are happy to address them in the reviewer-author discussion period.\n"}, {"review_id": "0EA45XqR4Q", "reviewer": "Reviewer_BWd8", "summary": "This paper tries to scale current meta learning algorithm and make scalable meta learning practical. Specially, the authors propose a novel algorithm SAMA from the perspective of algorithm and system, which can support arbitrary optimizers in the base level of meta learning and reduce the computation cost. The experimental results illustrate that the proposed method SAMA can improve the throughput and meanwhile reduce the computation cost in large-scale benchmarks. For example, they try to evaluate the performance of SAMA on text classification with large language models, such as  BERT and RoBERTa. ", "questions": "In section 3.1 and 3.2, you provide two solutions about two important problems. I would like to ask do you try to compare your proposed method with the related works that also try to solve these problems?", "limitations": "NA", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper focus on great direction and how to scale meta learning is very important, especially when large language models are so popular now.\nThe proposed method is also easy to follow, especially the authors try to solve the problem from the perspectives of algorithm and system. ", "weaknesses": "This paper mainly focus on scaling and maybe you should provide more experimemtal analysis about scaling, such as from small mode to large model, from 1 gpu to more gpus. Maybe you can provide more results of your proposed method and baseline about scaling. \nThe most experiments concentrate on nlp tasks. However, we know there are also dome large models on other tasks, such as computer vision. Therefore, maybe you should provide more results on other tasks to illustrate the generality of SAMA.\nI think the contribution from system is a little weak and the workflow can be directly implemented with original PyTorch. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our work, and for the useful feedback. We address the comments and questions raised in your review below and in the **global response**.\n\n### **Additional scalability analysis**\n> **Q.** This paper mainly focus on scaling and maybe you should provide more experimental analysis about scaling, such as from small mode to large model, from 1 gpu to more gpus. Maybe you can provide more results of your proposed method and baseline about scaling.\n\n**A.** We agree with the value of these experiments, and want to highlight the three scalability analyses in our paper — Table 2 & Figure 1 (bottom left): memory/throughput analysis on the noisy finetuning task, Figure 1 (bottom right): memory vs model size analysis on the continued pretraining task, and Figure 4 (Appendix D): scale-accuracy analysis on the image few-shot learning task. Given your suggestion, we additionally performed a more extensive ablation study on two datasets from the Wrench benchmark, and presented the results in the global response. From the results, you can clearly see the effects/benefits of each component in SAMA for scalable meta-learning.\n\n### **Other domains than NLP**\n> **Q.** The most experiments concentrate on nlp tasks. However, we know there are also some large models on other tasks, such as computer vision. Therefore, maybe you should provide more results on other tasks to illustrate the generality of SAMA.\n\n**A.** We want to kindly remind the reviewer that we included two computer vision experiments, namely 1) data pruning on ImageNet/CIFAR datasets and 2) (MAML-like) few-shot image classification, respectively in Sec. 4.3 and Appendix D of our paper. With these two experiments, we tried our best to demonstrate the generality of SAMA across different domains. We admit that there are a lot of other interesting applications and domains to explore, though in this paper we aimed to prioritize the best subset of experiments given the 9-page space limit.\n\n### **Distributed training**\n> **Q.** I think the contribution from system is a little weak and the workflow can be directly implemented with original PyTorch.\n\n**A.** To the best of our knowledge, our work is the first to raise the problem of efficient distributed meta learning as well as provides an initial solution to it. While our distributed training solution may look simple in retrospect, we noticed that most existing works on meta learning are limited to a single-GPU setup. For the remaining few works that combine distributed training and meta learning, we found that existing implementations are either incorrect (e.g. no proper gradient synchronization) or do not provide communication optimization. \nIn contrast, as stated in Sec 3.3, our DDP strategy (1) reduces communication cost by 3 times through sparse gradient synchronization, and (2) hides the remaining communication cost behind the computation through our simple yet smart implementation trick.\n\nIf you have any other comments that could make our paper stronger, we are more than happy to discuss them in the remaining review period. Thanks again for your reviewing effort!\n"}, {"review_id": "aKbRfeyiB2", "reviewer": "Reviewer_fPdM", "summary": "The authors explore the issues impacting the scalability of Gradient-based Meta-Learning (GBML), including high memory/compute costs, algorithmic instability, and poor support for distributed training. The causes identified are: the base Jacobian inversion, the absence of algorithmic adaptation for adaptive optimizers, and the requirement for a custom backward pass of meta gradient computation. Proposed solutions include: approximating the base Jacobian with an identity matrix, expanding the meta Jacobian through the chain rule, and creating a communication strategy leveraging the communication-computation overlap trick.", "questions": "1. Can you further explain why an identity matrix can be used to approximate the base Jacobian?\n2. What is the innovative point of the distributed training proposed in this paper? What are the main benefits brought about by this innovation?\n3. Can you provide some analysis demonstrating the distance of each approximate solution from the optimal solution?\n\nPlease see Weaknesses for details.", "limitations": "Please see Weaknesses.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The application of meta-learning (bi-level optimization) is widespread in various aspects of deep learning, and considering the acceleration of meta-learning is an important direction.\n2. Validating the effectiveness of meta-learning (bi-level optimization) on large-scale datasets and models is important, which can verify the ability of meta-learning to solve practical problems.\n3. The manuscript is written quite clearly and is easy to read.\n", "weaknesses": "1. In the methodology section, the explanation of why an identity matrix can be used to approximate the base Jacobian is unclear and lacks necessary analysis. As for approximating the second-order derivative with the first-order derivative, the method is almost the same as DARTS. Distributed training naturally fits in a scenario where only the first-order derivative is used, and the authors did not provide any special design for it.\n2. In the experiment section, the authors should focus on comparing the difference in training efficiency between the meta-learning acceleration framework proposed in this paper and other approximation methods. However, this is only reflected in Table 2, and there is no comparison with the approximation method proposed by DARTS, which is the method most closely related to this paper.\n3. It would be beneficial for this work to provide some analysis demonstrating the distance of each approximate solution from the optimal solution, as this could provide more valuable guidance.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for the valuable feedback that will improve the quality of our work. We attempt to clarify and address your concerns regarding our work here and in the **global response**.\n\n### **Comparison with DARTS**\nWhile we recognize the similarity, SAMA is different from DARTS in two major aspects.\n1. Due to a lack of algorithmic adaptation in DARTS, SAMA and DARTS differ when an adaptive optimizer is used at the base level. As we pointed out in the paper, recent large models are by default trained with adaptive optimizers, and we empirically demonstrate in Table 1 that lacking this adaptation often leads to noticeable performance degradations. When SGD is used as a base optimizer, SAMA and DARTS become more similar, but they still have another important difference which we discuss in the next paragraph.\n2. As stated in line 138-139, DARTS computes a meta Jacobian at initialization $\\theta$, while SAMA computes it at (approximate) convergence $\\theta^*$. Therefore, DARTS is closer to iterative differentiation while SAMA is implicit differentiation. This difference has two important implications in both memory/compute efficiency for scalable meta learning.\n- Memory: DARTS needs to keep the copies of both the initial parameter $\\theta$ and the most recent parameter $\\theta^*$ while SAMA only requires tracking $\\theta^*$ (reference: the official implementation of DARTS, which indeed separately saves the copy of $\\theta$). Therefore, DARTS incurs additional memory usage, worsening the memory bottleneck issue in scalable meta learning.\n- Compute: Compared to DARTS, SAMA more naturally allows for a larger unroll step, which essentially reduces the frequency of the meta-gradient computation, the most expensive operation in meta-learning. Given that the meta-objective is evaluated at the optimal base solution $\\theta^*$, the quality of meta-gradient naturally depends on the quality of $\\theta^*$ approximation. However, DARTS computes the meta-gradient at $\\theta$, and therefore the approximation error of $\\theta^*$ increases as we use larger unroll steps. Indeed, the original DARTS paper only used an unroll step of 1, while we were able to use an unroll step of 10 for our Transformer experiments (Sec 4.1 & 4.2). This significantly improves the overall training efficiency (i.e. throughput) of meta learning.\nTo more clearly understand these benefits, we provide a quantitative analysis of memory/computation efficiency between SAMA and DARTS in an additional ablation study. The result can be found in Table 1 & 2 in the global response. \n\n### **Justification of identity approximation**\nAs shown in [17], approximating base Jacobian as identity can be understood as preconditioning meta-gradient. In other words, given the optimization path from $\\theta$ to $\\theta^*$, the identity base Jacobian trick in SAMA essentially approximates this optimization path with the *reverse* one step update (i.e. gradient ascent) from $\\theta^*$. Similarly, in DARTS, which also adopts the identity base Jacobian trick, it approximates the optimization path with the one step update (i.e. gradient descent) from $\\theta$. \nFollowing your suggestion, we additionally analyzed the effect of this approximate solution (i.e. the distance of the approximate solution obtained by SAMA from the optimal solution) on the meta-gradient computation and the final optimal meta solution ($\\lambda^*$) in the “biased regression” setting, where the closed-form solution can be obtained analytically. In this experiment, we empirically show that the identity approximation still allows for the accurate estimation of meta-gradients and the optimal meta solution $\\lambda^*$, even when the true base Jacobian is not an identity matrix. More detailed discussion is provided in the global response.\n\n[17] Fung et al., Jfb: Jacobian-free backpropagation for implicit networks. AAAI, 2022.\n\n### **Benefits of SAMA’s distributed training scheme**\nTo the best of our knowledge, our work is first to raise the problem of the communication efficiency in distributed meta learning, and propose an initial solution. Given recent advances in hardware, the communication overhead can quickly become a bottleneck for training efficiency in large-scale learning [31]. As stated in Sec 3.3, one meta gradient computation in SAMA involves 3 backward computations. If implemented naively, gradient synchronization can happen after each backward computation, whereas SAMA performs synchronization only once after the last backward computation. This can roughly reduce the communication cost by three times. In addition, most meta learning implementations including DARTS heavily rely on `torch.autograd.grad` instead of `torch.autograd.backward`. Unfortunately, autograd.grad doesn’t support communication optimization such as a communication-computation overlap, which is essential in large-scale learning as shown in [31]. In summary, in addition to simply applying distributed training to our first-order SAMA algorithm, we perform additional communication optimizations that (1) reduce communication cost by 3 times through sparse gradient synchronization, and (2) hide the remaining communication cost behind the computation through the mixed use of autograd.grad and autograd.backward. In doing so, we were able to maximize training efficiency of distributed meta-learning.\n\n[31] Li et al., Pytorch distributed: Experiences on accelerating data parallel training. VLDB, 2021.\n\nWe hope our response resolved most of your concerns, and helped you evaluate our work more positively. If you have other comments, we are more than happy to address them in the reviewer-author discussion period."}, {"review_id": "aZs5RiMd4j", "reviewer": "Reviewer_YbqA", "summary": "This paper proposes a novel framework that could achieve scalable meta-learning algorithms from the perspectives of algorithms and systems. For algorithms, some approximations are proposed for base Jacobian inverse and adaptive optimizers; for systems, it implements the distributed algorithms to ensure different tasks can be done in parallel. Through extensive experiments in scalable meta-learning experiments, the proposed approach shows advantages over other baselines.", "questions": "1. ablation study experiments (see weakness)\n2. how to understand u in section 3.1?\n3. data reweighting method in 4.1: what is the relationship between w and $\\lambda_r$, c and $\\lambda_c$? Are w and c networks? Also, it is helpful to cite similar papers here:\n  - Learning to Reweight Examples for Robust Deep Learning, ICML 2018\n  - A Nested Bi-level Optimization Framework for Robust Few Shot Learning, AAAI 2022\n4. Some typos:\n- 1) Eq.5: the negative sign \"-\" is missing.\n- 2) line 259: the second $L_{ft}$ should be ${L}_{pt}$.\n- Also, it might be easy to follow if $\\lambda$ is included in the first equation below line 69.", "limitations": "yes.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The motivation is clear and robust. Scalable meta-learning is critical, and this paper might be a good solution.\n2. The presentation is clear, although some technical issues are unclear. Generally, the presentation is good.\n3. Different experiments regarding scalable meta-learning are conducted, which could be helpful to get some insights for the common meta-learning/few-shot learning researchers.", "weaknesses": "1. Although the current experiments are very helpful in exploring scalable experiments under meta-learning, some important ablation study ones are missing.\n- 1) Three issues (section 3) are solved to ensure the scalability of meta-learning. Then which issue is the major one to slow down the process? How could these three issues affect the scalabilities? For example, for the DDP, is that possible to conduct some experiments of DDP on the top of MAML? Most researchers are familiar with MAML, and if DDP experiments on MAML show improved efficiency on common few-shot learning experiments, that would be very helpful and easy to understand the benefits. Similarly, for the other two components or issues.\n- 2) Is that possible to conduct more experiments to compare the current approach with others, such as iMAML, see Figure 2 in the iMAML paper. It is hard to compare the proposed approach with others based on the current experiments. It is better to compare them based on the published experiments.\n\n   - MAML: Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks, ICML 2017\n     \n   - iMAML: Meta-Learning with Implicit Gradients, NeurIPS 2019\n\n2. It's straightforward for section 3.3, but it is a little bit difficult to follow for sections 3.1 and 3.2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for the valuable review. We try our best to address your concerns and questions here and in the **global response**.\n\n### **Ablation study**\n> **Q.** Three issues (section 3) are solved to ensure the scalability of meta-learning. Which issue is the major one to slow down the process? How do these three issues affect scalability? \n\n**A.** Although we didn’t explicitly frame it as an “ablation study”, the effectiveness of each component in SAMA can be understood from Table 1 & 2 in our paper. Below, we directly discuss each component of SAMA based on Table 1 & 2. We also provide a unified table as an ablation study, in the global response above.\n\n**Base Jacobian inverse**\n\n*tl;dr* Identity approximation of Base Jacobian significantly improves memory/compute efficiency (Table 2).\n\nOur baseline algorithms in Table 2 and Figure 1, Neumann and CG, are both state-of-the-art implicit differentiation meta learning algorithms that attempt to approximate base Jacobian inverse as accurately as possible with multiple Hessian-vector products, instead of approximating it with the identity as SAMA. As shown in Table 2, due to expensive second-order gradient computations involved in Hessian-vector products, these methods demonstrate much poorer throughput and GPU memory usage than SAMA, both of which are the major bottlenecks in efficiently scaling meta learning. In our ablation study in the global response, we additionally demonstrate that this identity approximation also has a minimal impact on accuracy as well by comparing SAMA against Neumann and CG in terms of accuracy.\n\n**Algorithmic adaptation for adaptive optimizer**\n\n*tl;dr* Algorithmic adaptation significantly improves accuracy (Table 1) at the minimal memory/compute cost (Table 2).\n\nThe main goal of this work is to devise a *(1) memory/compute efficient* meta learning algorithm that *(2) achieves good performance/accuracy*. In Table 1, SAMA consistently achieves better accuracy than SAMA-NA, which lacks algorithmic adaptation. Moreover, Table 2 shows that SAMA achieves a comparable memory/compute efficiency as SAMA-NA. \n\n**Distributed training**\n\n*tl;dr* Both GPU memory usage and throughput improve consistently as computations are distributed across more GPUs (Table 2).\n\nWe agree that it may not be straightforward for readers to clearly understand the benefit of each component when the results are spread across two separate tables. Hence, in the global comment above, we provide a unified table for Wrench experiments, including a comparison with state-of-the-art meta-learning baselines.\n\n> **Q.** For DDP, is it possible to conduct some experiments of DDP on top of MAML?\n\n**A.** We first want to kindly remind the reviewer that we included (MAML-like) few-shot learning experiments in Appendix D (line 209). While we didn’t study the exact problem you mention, we investigated another interesting question of “can we improve the few-shot learning performance by increasing the model size given SAMA’s improved scalability?”, to which our preliminary answer is “yes”. This question is quite different from a majority of previous works where the model size is fixed and authors attempt to improve few-shot accuracy by designing new algorithms. \n\nAs for the applicability of DDP to MAML, we will separately discuss two different aspects of MAML: 1) application and 2) algorithm. More specifically, MAML solves the *few-shot learning* application of meta-learning with the *iterative differentiation* algorithm.\n- Application: We expect DDP will have a limited impact on few-shot learning applications. DDP improves compute/memory efficiency by distributing samples in a mini-batch across multiple GPUs. However, by the nature of few-shot learning, this application usually has a very small batch size. Thus, distributing such a small batch across multiple GPUs would likely lead to only a modest improvement in compute/memory efficiency. However, meta learning has a lot of other applications besides few-shot learning, such as data optimization and hyperparameter optimization, which don’t necessarily have a small batch size. We expect our DDP scheme will have a meaningful impact on these applications.\n- Algorithm: We believe iterative differentiation in MAML would show a reduced compatibility with DDP compared to SAMA. As stated in Sec 3.1, the first-order nature (no Hessian computation) is crucial in achieving the improved DDP compatibility, but the iterative differentiation algorithm usually involves Hessian computation.\n\n### **Clarity**\n> **Q.** It's straightforward for section 3.3, but it is a little bit difficult to follow for sections 3.1 and 3.2.\n\n**A.** While we are unable to edit the manuscript during the review period this year, we will improve our writing and the overall clarity in the camera ready version of our paper. \n\n> **Q.** How to understand $u$ in section 3.1?\n\n**A.** As mentioned in Eq. (2) (line 96), $u$ is the update function of gradient-based optimization. Here, we provide concrete examples of $u$ for several popular optimizers.\n- SGD: $u_t = g_t = g(\\theta_t;\\lambda_t)\\quad$\n- SGD-M: $u = \\beta m_{t-1} + g_t\\quad$\n- Adam: $u = \\frac{\\beta_1 m_{t-1} + (1 - \\beta_1) g_t}{\\sqrt{\\beta_2 v_{t-1} + (1 - \\beta_2 g_t^2)}}$\n\n> **Q.** Data reweighting method in 4.1: what is the relationship between $w$ and $\\lambda_r$, $c$ and $\\lambda_c$? Are w and c networks? Also, it is helpful to cite similar papers.\n\n**A.** $w/c$ are respectively reweighting and label-correction functions (i.e. neural networks) parameterized by $\\lambda_r/\\lambda_c$. We cited a few relevant data reweighting and label-correction works in line 219, and will also add papers you suggested. \n\n> **Q.** Typos\n\n**A.** Thanks for pointing them out. We will fix them in the revision.\n\nWe hope our response resolved most of your concerns, and helped you evaluate our work more positively. If you have other comments, we are more than happy to address them in the remaining review period."}, {"review_id": "gverqWjlRc", "reviewer": "Reviewer_TSYz", "summary": "The paper \"Making Scalable Meta Learning Practical\" introduces a novel approach called SAMA (Scalable Meta Learning with Arbitrary Optimizers) to address the scalability issues in meta learning. The authors combine advancements in implicit differentiation algorithms and systems to develop SAMA, which supports arbitrary optimizers in meta learning programs while reducing computational burden and utilizing efficient distributed training techniques. In experiments, they focus on data optimization tasks and demonstrate the effect of SMLA on several domains.", "questions": "Please see the weakness.", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1 The paper proposes SAMA, a new method that combines implicit differentiation algorithms and systems to address scalability issues in meta-learning. This innovative approach sets it apart from existing methods.\n\n2. The paper presents compelling experimental results, showing significant throughput increases and memory consumption reductions on single- and multi-GPU setups compared to baseline meta-learning algorithms. These performance improvements validate the effectiveness of SAMA.\n\n3. Well written. This paper is well-written and easy to follow.", "weaknesses": "-- Need more clarification on the application scenarios. \n\n1. This paper claims that the proposed SAMA makes scalable meta-learning practical while the experiments are only carried out in data optimization applications. However, technically, the method should also be applicable to other applications of meta-learning. In practice, whether the method proposed in this paper is effective for other applications?\n\n2. There are a large number of meta-learning algorithms, some of which may not involve the computation of a large-scale Jacobian matrix. Does the proposed method work for such algorithms? The application scope should be better clarified. \n\n-- Unclear description\n\n3. How is \\theta^* obtained? Is it approximated by several steps or does it need to be trained to converge, and if it needs to be trained to converge, does the meta-learning process require more training time? Is such overhead already taken into consideration in the evaluation?\n\n-- Typos and errors\n\n4. There are some spelling errors in this paper that need further correction. Line 259 wrong subscript: the second “ft” should be “pt”. In Equation 5, The second equal sign should also be an approximate equal sign?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for the positive review as well as the helpful feedback. Here, we address each of the questions and comments that you raise.\n\n### **Clarification on the application scenarios**\n\n> **Q1.** This paper claims that the proposed SAMA makes scalable meta-learning practical while the experiments are only carried out in data optimization applications. …. In practice, whether the method proposed in this paper is effective for other applications?\n\n**A.** As the reviewer noted, SAMA can be utilized in other meta learning applications that can be formulated as bilevel optimization, such as few-shot learning, neural architecture search, and hyperparameter optimization. As mentioned in line 209, we also performed the traditional (MAML-like) few-shot learning experiments in Appendix D. While most existing few-shot learning works focus on developing new algorithms given the fixed network size, in our experiment we rather study the question of “can we improve the performance of few-shot learning by increasing the model size given the memory/compute efficiency of SAMA?”, to which our preliminary answer is “yes”.\n\n> **Q2.** There are a large number of meta-learning algorithms, some of which may not involve the computation of a large-scale Jacobian matrix. Does the proposed method work for such algorithms?\n\n**A.** We believe several meta-learning algorithms that avoid the use of a large Jacobian matrix (e.g. Reptile [1]) are specifically designed for few-shot learning, which is one particular application of meta-learning/bilevel optimization. Therefore, it is oftentimes not straightforward to apply these algorithms to a more general class of meta learning applications (e.g. data optimization, neural architecture search) that can be formulated as bilevel optimization. In contrast, SAMA is directly derived from the bilevel optimization formulation of meta learning, and thereby demonstrates a more wide applicability. \nAs for applying some components of SAMA (e.g. Sec 3.3: communication optimization) to other meta learning algorithms, we believe that it is possible. For example, our communication optimization strategy can be easily transferred to DARTS [2]. However, as each meta learning algorithm has a different design, our components may need to be additionally adapted to the target algorithm.\n\n[1] Nichol et al., On first-order meta-learning algorithms. Arxiv, 2018.\n[2] Liu et al., Darts: Differentiable architecture search. ICLR, 2018.\n\n### **Unclear description**\n\n> **Q3.** How is $\\theta^*$ obtained?\n\n**A.** $\\theta^*$ is approximated by several unroll updates instead of by training to convergence. As the reviewer expected, by avoiding full training, we were able to significantly reduce the computation burden (refer to line 97-100). All our experiments are conducted with this several-updates-approximation strategy for a fair comparison. More experiment details including the number of unroll steps are provided in Appendix B.\n\n### **Typos and errors**\n> **Q4.** There are some spelling errors in this paper that need further correction\n\n**A.** Thanks for pointing out typos. We will update these in the camera ready version of our paper.\n\nWe again express our gratitude for your constructive comments, which are very helpful to our paper. If you have any further comments that could make our paper stronger, we are more than happy to discuss them in the remaining review period.\n"}]

Author rebuttal:
We first want to express our gratitude to all reviewers for their reviewing efforts. In our global response, we address two issues raised by reviewers: 1) Ablation study and 2) (empirical) justification of the identity base Jacobian approximation.

### **Ablation Study & SOTA comparison**
While the effectiveness of each component in SAMA (i.e. base Jacobian inverse, algorithmic adaptation for the adaptive optimizer, and efficient distributed training) can be collectively understood from Table 1 & 2, several reviewers asked for a more direct ablation study. Therefore, we provide below a unified table for the ablation study. In detail, our experiment settings are:

- Datasets: 1) AGNews and 2) IMDB from the Wrench benchmark (Sec 4.1)

- Baselines: 1) fine-tuning (no meta-learning), 2) iterative differentiation (e.g. MAML), 3) conjugate gradient (e.g. iMAML), 4) Neumann series, 5) DARTS, 6) SAMA-NA (no algorithmic adaptation)

**Table 1. Ablation results on AGNews**

|                  | Base Jacobian | Algo Adapt | Distributed | Accuracy | Throughput | Memory |
|:----------:|:---------:|:-------:|:-------:|:------:|:-------:|:-----:|
| Finetuning (Baseline)  | x | x | x | 85.79 | 169.16 | 7.77 |
| Iterative Differentiation (MAML) | x | x | x | 85.78 | 28.07 | 22.94 |
| Conjugate gradient (iMAML) | x | x | x | 86.78 | 65.14 | 22.03 |
| Neumann series | x | x | x | 86.65 | 67.03 | 19.70 |
| DARTS | o | x | x | 86.36 | 43.69 | 10.81 |
| SAMA-NA | o | x | x | 86.55 | 137.90 | 10.30 |
| SAMA | o | o | x	| **89.05** | 134.56 | 11.12 |
| SAMA (2 GPUs) | o | o | o | **88.85** | 226.27 | 8.00 |
| SAMA (4 GPUs) | o | o | o | **89.02** | **298.28** | **6.46** |

**Table 2. Ablation results on IMDB**

|                  | Base Jacobian | Algo Adapt | Distributed | Accuracy | Throughput | Memory |
|:--------:|:---------:|:------:|:--------:|:------:|:-------:|:-----:|
| Finetuning (Baseline)  | x | x | x | 78.16 | 144.39 | 6.60 |
| Iterative Differentiation (MAML) | x | x | x | 80.25 | 24.24 | 22.03 |
| Conjugate gradient (iMAML) | x | x | x | 81.01 | 56.27 | 21.92 |
| Neumann series | x | x | x | 79.92 | 57.85 | 19.75 |
| DARTS | o | x | x | 80.47 | 37.53 | 10.35 |
| SAMA-NA | o | x | x | 81.92 | 117.86 | 9.93 |
| SAMA | o | o | x	| **84.31** | 116.94 | 10.84 |
| SAMA (2 GPUs) | o | o | o | **85.18** | 196.48 | 7.84 |
| SAMA (4 GPUs) | o | o | o | **84.19** | **263.74** | **6.39** |

From our extended ablation study, it can be seen that 1) an identity approximation of base Jacobian significantly improves memory/compute efficiency, 2) algorithmic adaptation improves meta learning performance at the minimal compute/memory cost, and 3) our communication-optimized distributed training further improves compute/memory efficiency.

### **Analysis on Identity Approximation of the Base Jacobian**
As obtaining the closed form solution of the Hessian is impossible in almost all deep learning problems, we study the soundness of the identity approximation of base Jacobian in the simpler “biased regression” setting [1], for which the bilevel optimization formulation is as follows:
$$
\lambda^* = \arg\min_{\lambda} \Vert X’w^*(\lambda) - y’\Vert^2
$$

$$
w^*(\lambda) = \arg\min_w \Vert Xw - y\Vert^2 + \beta\Vert w - \lambda \Vert^2
$$
Given the above formulation, the closed-form solutions for the base Jacobian, the meta-gradient $g_{\lambda}$, and the optimal meta solution $\lambda^*$ are:

- base Jacobian $= X^TX + \beta I$
- $g_{\lambda} = \beta (X^TX + \beta I)^{-1}(X’^TX’w^* - X’^Ty’)$, where $w^* = (X^TX + \beta I)^{-1}(X^Ty + \beta\lambda)$
- $\lambda^* = (A^TA)^{-1}A^Tb$, where $A=\beta X’(X^TX + \beta I)^{-1}), b=y’ - X’(X^TX + \beta I)^{-1}X^Ty$

We set $\beta=0.1$ and perform 100 meta updates, and measure 1) the cosine similarity between the ground truth $g_{\lambda}$ and the meta gradient $g_{approx}$ obtained with our approximation, and 2) the L2 distance between the current meta parameter $\lambda_t$ and the optimal solution $\lambda^*$ at each time step $t$. For a more thorough analysis, we also compute these two metrics for other meta gradient algorithms that explicitly approximate base Jacobian inverse with conjugate gradient and Neumann series. In the table below, we provide the metric obtained from several different time steps. *The visual plots for all time steps are provided in the attached pdf file at the bottom of this response.*

**Table 3.** $cos(g_{\lambda}, g_{approx})$ **result**

|     	| t=0	| t=10   | t=20   | t=50   | t=100  |
|---------|--------|--------|--------|--------|--------|
| CG  	| 0.9995 | 0.9994 | 0.9989 | 0.9997 | 0.9319 |
| Neumann | 0.9957 | 0.9952 | 0.9949 | 0.9948 | 0.9363 |
| SAMA	| 0.9843 | 0.9818 | 0.9815 | 0.9861 | 0.9321 |

**Table 4.** $\Vert \lambda^* - \lambda_t\Vert_2$ **result**

|     	| t=0	| t=10   | t=20   | t=50   | t=100  |
|---------|--------|--------|--------|--------|--------|
| CG  	| 3.6752 | 1.8959 | 0.9507 | 0.1953 | 0.0160 |
| Neumann | 3.6972 | 2.6116 | 1.3468 | 0.2823 | 0.0216 |
| SAMA	| 3.6856 | 2.1537 | 0.6952 | 0.1966 | 0.0184 |

From Table 3 & 4 (and plots in the attached PDF), we can clearly see that 1) while slightly less accurate than second-order algorithms like CG, SAMA still achieves a high directional alignment with the ground truth meta-gradient, and 2) SAMA also achieves a stable convergence to the optimal solution at a comparable speed. We hope this result corroborates the soundness of our identity approximation for the base Jacobian.

[1] Grazzi et al., On the Iteration Complexity of Hypergradient Computation. ICML, 2020.

Note: We used smaller $\beta$ than in the original paper (1 vs 0.1), to amplify the “non-identitiness” of the base Jacobian.

We will include all of the above results in the camera ready version of our paper. If reviewers have any further questions regarding these additional experiments, we are happy to discuss them during the author-reviewer discussion period. Thanks again!


Meta-review: The topic of practical meta learning is appreciated, the approach using implicit differentiation is novel and of interest and the demonstrated computational benefits are convincing. The author provided responses to several reviewer comments have been well received, and after author/reviewer discussion all five reviewers are supportive of acceptance. The AC agrees and encourages the authors to include all the promised additions and experiments in the camera-ready paper.
Decision: Accept (poster)