SlowGuess/NeurIPS_2023_OCR huggingface dataset:
From a technical perspective, large-scale data optimization has a 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. 4.2), and scale-agnostic data pruning (Sec. 4.3). While not an application of data optimization, we also include a preliminary analysis on 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 [53] 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:
$$
\lambda^ {*} = (\lambda_ {r} ^ {*}, \lambda_ {c} ^ {*}) = \operatorname * {a r g m i n} _ {\lambda_ {r}, \lambda_ {c}} \mathcal {L} (D _ {c l e a n}; \theta^ {*} (\lambda_ {r}, \lambda_ {c}))
$$
$$
\text{s.t.} \theta^{*}(\lambda_{r},\lambda_{c}) = \operatorname *{argmin}_{\theta}\frac{1}{|D_{noisy}|}\sum_{(x,y)\in D_{noisy}}w(\mathcal{L};\lambda_{r})\cdot \mathcal{L}(f(x;\theta),c(x,y;\lambda_{c}))
$$
where  $w(\cdot ;\lambda_r)$  and  $c(\cdot ;\lambda_c)$  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
<sup>3</sup>https://github.com/leopard-ai/betty
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 1.
Algorithm	TREC	SemEval	IMDB	ChemProt	AGNews	Yelp
Finetune (orig) [67]	-	66.56 (2.31)	83.93 (1.74)	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 (0.08)	87.03 (0.00)	89.22 (0.05)
Finetune (ours)	-	67.93 (2.55)	79.28 (1.78)	78.16 (2.28)	57.35 (1.43)	85.79 (0.49)	84.32 (2.55)
+R	SAMA-NA	74.33 (2.34)	87.11 (1.11)	81.92 (1.74)	60.88 (0.60)	86.83 (0.19)	80.96 (3.04)
+R & C	SAMA-NA	79.00 (2.62)	87.67 (1.36)	80.44 (0.97)	64.05 (0.52)	87.05 (0.39)	80.73 (3.11)
+R	SAMA	85.73 (0.81)	89.67 (0.67)	84.31 (1.86)	76.89 (1.39)	89.05 (0.34)	93.64 (0.40)
+R & C	SAMA	87.93 (1.17)	88.83 (2.32)	85.71 (0.82)	77.78 (0.59)	89.79 (0.27)	93.77 (0.08)
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_{clean}$  in the meta level. Second, given the superior performance of SAMA compared to SAMA-NA, we empirically verify the importance of algorithmic adaptation when an adaptive optimizer is used to train this Transformer base learner.
Additionally, we compare compute/memory efficiency of SAMA to that of two other implicit-differentiation-based meta learning algorithms, namely Neumann Series [40] and conjugate gradient [51]. For fair comparison, we evaluate GPU memory usage (MB) and throughput (samples/sec) on the AGNews dataset from Wrench with a fixed global batch size of 48, and summarize the result in Table 2. Three observations that can be made here are that (1) SAMA is generally more compute/memory efficient than the baselines, (2) the cost of algorithmic adaptation for adaptive optimizers is marginal as expected, and (3) the efficiency
Table 1: 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.
GPUs	Memory	Throughput
Neumann	1	26.0	82.9
CG	1	28.4	82.1
SAMA-NA	1	13.7	144.1
SAMA	1	14.3	142.0
SAMA	2	10.4	241.2
SAMA	4	7.4	396.7
Table 2: Memory and throughput analysis on AGNews with 4 V100 GPUs.
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 F, we provide a 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.
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-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 [64], which could lead to 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 [11] 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:
$$
\lambda^ {*} = \operatorname * {a r g m i n} _ {\lambda} \mathcal {L} _ {f t} (D _ {f t}; \theta^ {*} (\lambda))
$$
$$
\text {s . t .} \theta^ {*} (\lambda) = \underset {\theta} {\operatorname {a r g m i n}} \mathcal {L} _ {f t} (D _ {f t}; \theta) + \frac {1}{| D _ {p t} |} \sum_ {x \in D _ {p t}} w (x; \lambda) \cdot \mathcal {L} _ {p t} (x; \theta)
$$
where  $\mathcal{L}_{ft} / \mathcal{L}_{ft}$  are finetuning/pretraining loss functions,  $D_{ft} / D_{pt}$  are finetuning/pretraining datasets, and  $w(\cdot ;\lambda)$  is the data reweighting network. Following the experiment setup in TARTAN [11], 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-Meta respectively because (1) TAPT consistently underperforms TARTAN-MT [9] and (2) TARTAN-Meta uses additional validation data in the meta level of the downstream tasks, making the comparison unfair. We report our experiment results in Table 3.
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 [24]	84.17 (0.50)	87.23 (3.65)	71.84 (4.78)	80.42 (1.57)	80.92
TARTAN-MT [11]	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.03)	71.63 (1.68)	81.84 (0.08)	83.29
Table 3: Experiment results for auxiliary learning with the continued pretraining task. Following [24], we report test micro-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-/down-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-MT 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 1. 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 a 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-Net (MWN) [58] with four major 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 of the prediction in addition to the loss value as an input to MWN to better estimate importance 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:
$$
\lambda^ {*} = \operatorname * {a r g m i n} _ {\lambda} \mathcal {L} (D _ {t r a i n}; \theta^ {*} (\lambda))
$$
$$
\text {s . t .} \theta^ {*} (\lambda) = \operatorname * {a r g m i n} _ {\theta} \frac {1}{| D _ {t r a i n} |} \sum_ {(x, y) \in D _ {t r a i n}} w (\mathcal {L}, \mathcal {U}; \lambda) \cdot \mathcal {L} (x, y; \theta)
$$
where  $w(\cdot ;\lambda)$  is MWN that takes the loss value  $\mathcal{L}$  and the uncertainty  $\mathcal{U}$  of the training sample  $(x,y)$  as an input and outputs the importance weight. Under this setup, we run meta learning with SAMA for 30 / 50 epochs respectively for ImageNet-1k / CIFAR-10 and obtain the pruning metrics by averaging the importance weights of the last 5 epochs. We compare our 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
Figure 3: Top Left: ImageNet-1k data pruning results with ResNet-50. Reported numbers are relative accuracy compared to full training accuracy (i.e., pruned_acc/full_acc). 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 is obtained from Deepcore [22]. Bottom: Relative time spent in finding data to prune compared to full ImageNet-1k training time.
EL2N1 [47]	EL2N10 [47]	GradNd10 [47]	DynaMS [63]	SAMA (ours)
Relative Search Time	0.16	1.65	1.65	0.16	0.46
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 an 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-based methods. We also note that, while the original MWN [58] encounters the OOM error under our setup of batch_size=256, the throughput analysis with the reduced batch size reveals that efficient distributed training with SAMA on 4 GPUs achieves  $15 - 20 \times$  speed up compared to the original MWN that lacks distributed training support.
Algorithms Two major lines of research in gradient-based meta learning algorithms are iterative and implicit differentiation [19]. Iterative differentiation [14, 15, 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 differentiation, 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, a 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 [14, 51, 71], neural architecture search [38, 69], hyperparameter optimization [15, 16, 40, 42], data optimization [21, 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/declarative programming paradigms.
Systems Compared to algorithms and systems research, there are relatively fewer research efforts in meta learning systems. In an attempt to facilitate research in 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 or 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 to support various systems optimization for large-scale meta learning, we chose to implement SAMA in this framework.
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 a technical level and resolved them by developing SAMA. Tested on multiple benchmarks, SAMA empirically demonstrated its scaling efficiency as well as its capability to optimize a variety of high-dimensional inductive biases of large-scale learning. In future work, we plan to explore two directions. First, given that training extremely large models with  $10\mathrm{B}+$  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-family models. Overall, we hope that our work can serve as a stepping stone for a lot of interesting scalable meta learning research to come.
Limitations & Broader Impacts While SAMA has demonstrated significantly improved compute/memory efficiencies, and stably worked with fairly large models like BERT/RoBERTa, we could not test it on larger models with 1B+ parameters due to a lack of computational resources. Meta learning itself is mostly value-neutral, but there is a chance that practitioners amplify the bias or toxicity of the machine learning programs by enforcing them in the meta objective.
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 IIS2311990, NSF IIS2123952, NSF CNS2008248, NSF BCS2040381, NGA HM04762010002, NIGMS R01GM140467, and Amazon 2023-5007107. WN was supported in part by NSF (#1651565), AFOSR (FA95501910024), ARO (W911NF-21-1-0125), CZ Biohub, and Sloan Fellowship. ES and SVM acknowledge the support in part by DSO National Laboratories.
[1] Amro Kamal Mohamed Abbas, Kushal Tirumala, Daniel Simig, Surya Ganguli, and Ari S Morcos. Semdedup: Data-efficient learning at web-scale through semantic dedduplication. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
[2] Antreas Antoniou, Harrison Edwards, and Amos Storkey. How to 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. learn2learn: A library for meta-learning research. arXiv preprint arXiv: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 Representations, 2023.
[7] Ross M Clarke, Elre Talea Oldewage, and José Miguel Hernández-Lobato. Scalable one-pass optimisation of high-dimensional weight-update hyperparameters by implicit differentiation. In International Conference on Learning Representations, 2022.
[62] Ashish Vavwani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is 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: Dynamic 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 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 pretrained language model with weak supervision: A contrastive-regularized 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, 33:15383-15393, 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 Machine 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. A contrastive rule for meta-learning. Advances in Neural Information Processing Systems, 35:25921-25936, 2022.
A Philosophies behind SAMA & Scalable Meta Learning
Here, we additionally discuss several important design principles and philosophies behind scalable meta learning and SAMA in a 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 budget (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 [14, 51], and adversarial learning [45], it was a natural call for us to investigate the scalability of meta learning algorithms following the spirit of "The Bitter Lesson". Interestingly, such a focus on the scalability of 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 a limited computation budget (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 law) has evolved a new ecosystem of large models and datasets in machine learning over time, which involves both systems and algorithms components. For example, to efficiently leverage the increased computation for large-scale learning, diverse systems techniques, such as data/model/pipeline parallelism 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 connections [27], Adam optimizer [32], self-attention [62], etc. Accordingly, in 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.