SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Systems compatibility Given that a great deal of systems support in machine learning, such as communication-computation overlap [35], has been developed for first-order gradient methods, avoiding explicit computations of higher-order gradient information including Hessian-vector products was an 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 computation of higher-order gradient information [34, 38, 37], many of these algorithms either assume the use of a 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.
In this section, we discuss various experiment details such as hyperparameters, baselines, and compute resources used for our experiments in Section 4. Our experiment codes are available in the supplementary material as well as under the official Betty GitHub repository  $^{4}$ .
B.1 Noisy Finetuning of Large Language Models
Hyperparameters We ran training for 1000 iterations on TREC/SemEval/IMDB/ChemProt/Yelp/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 4 below.
model	optimizer	init_lr	lr_scheduler	wdecay	dataset	unroll step	SAMA α
Base	BERT-base	Adam	1e-5	cosine	0	WRENCH train set (with majority voting)	10	1.0
Meta (Reweight)	2-layer MLP	Adam	1e-5	None	0	WRENCH dev set	N/A	N/A
Meta (Correct)	2-layer MLP	Adam	1e-5	None	0	WRENCH dev set	N/A	N/A
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 1 NVIDIA RTX 2080Ti GPU for the main experiment, and 4 NVIDIA Tesla V100 GPUs for the throughput-memory analysis in Table 2 and Figure 1.
B.2 Continued Pretraining of Large Language Models
Hyperparameters We ran training for 100 epochs with a batch size of 16, a maximum sequence length of 256, and the hyperparameters in Table 5 below.
Table 4: Hyperparameters for noisy finetuning of large language models experiments.
model	optimizer	init_lr	lr_scheduler	wdecay	dataset	unroll step	SAMA α
Base (Downstream)	RoBERTa-base	Adam	2e-5	linear decay + warmup linear (warmup proportion 0.6)	0	train split of ChemProt/HyperPartisan/ACL-ARC/SciERC	10	0.3
Base (Auxiliary)	RoBERTa-base	Adam	2e-5	linear decay + warmup linear (warmup proportion 0.6)	0	train split of ChemProt/HyperPartisan/ACL-ARC/SciERC	10	0.3
Meta	2-layer MLP	Adam	1e-5	None	0	train split of ChemProt/HyperPartisan/ACL-ARC/SciERC	N/A	N/A
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 domain-specific data on top of the pretrained RoBERTa-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 [11] performs masked language modeling with task specific data and downstream text classification training simultaneously in a multitask fashion through two different heads.
Compute Resources We used 1 NVIDIA Tesla V100 GPU for the main experiment, and 1 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 6 below. After pruning data based on the meta learning result, we ran ImageNet-1k
4https://github.com/leopard-ai/betty/tree/main/examples
training for 120 epochs with the learning rate decayed by 10 at epochs [40, 80] following the set up in DynaMS [63].
model	optimizer	init_lr	lr_scheduler	wdecay	dataset	unroll step	SAMA α
Base	ResNet-50	SGD	1e-1	None	1e-4	ImageNet-1k train set	2	1.0
Meta	2-layer MLP	Adam	1e-5	None	0	ImageNet-1k train set	N/A	N/A
For the CIFAR-10 data pruning experiment, we ran meta learning for 50 epochs with a batch size of 128, and configuration in Table 7 below. 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].
Table 6: Hyperparameters for ImageNet-1k data pruning experiments
model	optimizer	init_lr	lr_schedule	wdecay	dataset	unroll step	SAMA α
Base	ResNet-18	SGD	1e-1	None	5e-4	CIFAR-10 train set	2	1.0
Meta	2-layer MLP	Adam	1e-5	None	0	CIFAR-10 train set	N/A	N/A
Table 7: Hyperparameters for CIFAR-10 data pruning experiments
Baselines We adopt EL2N [47], GraNd [47], DynaMS [63] as our baselines for the ImageNet-1k experiments and GraNd [47], forgetting [61], margin [8] for the CIFAR-10 experiments. In detail, EL2N/GraND [47] respectively select samples with large L2-loss/gradient-norm values, forgetting [47] chooses samples that are frequently forgotten during training, and margin [8] chooses samples with least confidence. While these baselines are considered static pruning, DynaMS [63] falls under the category of dynamic pruning where data to be pruned change during training. Dynamic pruning may see the whole training data across different epochs, making a fair comparison difficult. Surprisingly, despite being a static pruning algorithm, SAMA-based data pruning still achieves a better performance than DynaMS.
Compute Resources We used 4 NVIDIA Tesla V100 GPUs for Imagenet-1k data pruning meta learning experiments and 1 NVIDIA RTX 2080Ti GPU for CIFAR-10 experiments.
Additional Information We measured the uncertainty  $\mathcal{U}$  via the difference between the predictions of the current model and the exponentially-moving-averaged model.
C Algorithmic Adaptation for Adam Optimizer
Since the Adam optimizer [32] has been the most popular optimizer to 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 in Adam as  $m$  and  $v$  respectively, and the learning rate as  $\gamma$ .
$$
\begin{array}{l} \frac {\partial u _ {a d a m}}{\partial g} = \frac {\partial u}{\partial g} \left(\gamma \frac {\beta_ {1} m + (1 - \beta_ {1}) g}{\sqrt {\beta_ {1} v + (1 - \beta_ {1}) g ^ {2}} + \epsilon}\right) \\ = \gamma \frac {(1 - \beta_ {1}) \beta_ {2} v - (1 - \beta_ {1}) \beta_ {2} m g + (1 - \beta_ {1}) \epsilon \sqrt {\beta_ {1} v + (1 - \beta_ {1}) g ^ {2}}}{\sqrt {\beta_ {1} v + (1 - \beta_ {1}) g ^ {2}} \left(\sqrt {\beta_ {1} v + (1 - \beta_ {1}) g ^ {2}} + e\right) ^ {2}} \\ \approx \gamma \frac {(1 - \beta_ {1}) \beta_ {2} v - (1 - \beta_ {1}) \beta_ {2} m g}{\sqrt {\beta_ {1} v + (1 - \beta_ {1}) g ^ {2}} \left(\sqrt {\beta_ {1} v + (1 - \beta_ {1}) g ^ {2}} + e\right) ^ {2}} \quad (\text {b e c a u s e} \epsilon \ll 1) \\ \end{array}
$$
Adaptation matrices can be similarly derived for other adaptive optimizers.
D The Effect of Scaling in Model-Agnostic Meta Learning
Since the inception of MAML [14], a myriad of algorithms have been proposed to improve few-shot image classification while assuming a fixed network architecture. In contrast, here we shift our focus from the algorithm to 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 is a variant of implicit differentiation, we closely follow the experiment setup in iMAML [51], where proximity to the initialization weights is explicitly enforced by  $L_{2}$ -regularization. The major difference is that iMAML uses a conjugate-gradient-based method, which requires second-order gradient information to compute meta gradients, while we adopt SAMA to achieve improved scaling to larger networks with its superior memory/compute efficiency. We conduct preliminary experiments on the Omniglot 20-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 on the few-shot classification accuracy. The experiment results are provided in Figure 4 below.
Figure 4: Few-shot image classification accuracy on Omniglot 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., language 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 as 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.
E Justification of the Identity Approximation for Base Jacobian
In this section, we aim to study the empirical effect of our identity approximation for base Jacobian on the meta gradient  $\frac{\partial L_{\text{meta}}}{\partial \lambda}$  and the optimal meta solution  $\lambda^{*}$ , when the true base Jacobian is not an identity matrix. 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 [19], for which the bilevel optimization formulation is as follows:
$$
\lambda^ {*} = \arg \min  _ {\lambda} \| X ^ {\prime} w ^ {*} (\lambda) - y ^ {\prime} \| ^ {2}
$$
$$
w ^ {*} (\lambda) = \arg \min _ {w} \| X w - y \| ^ {2} + \beta \| w - \lambda \| ^ {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:
1. Base Jacobian  $= X^T X + \beta I$
2.  $g_{\lambda} = \beta (X^{T}X + \beta I)^{-1}(X^{\prime T}X^{\prime}w^{*} - X^{\prime T}y^{\prime})$  , where  $w^{*} = (X^{T}X + \beta I)^{-1}(X^{T}y + \beta \lambda)$
3.  $\lambda^{*} = (A^{T}A)^{-1}A^{T}b$ , where  $A = \beta X^{\prime}(X^{T}X + \beta I)^{-1}$ ,  $b = y^{\prime} - X^{\prime}(X^{T}X + \beta I)^{-1}X^{T}y$
We set  $\beta = 0.1^5$  and perform 100 meta updates, and measure 1) the cosine similarity between the ground truth  $g_{\lambda}$  and the meta gradient obtained with our approximation (i.e.  $g_{SAMA}$ ), 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 Figure 5, we provide the metric obtained from all time steps.
Figure 5: Left:  $\cos(g_{\lambda}, g_{\text{approx}})$ , where  $g_{\lambda}$  is a ground truth (i.e. closed-form) meta gradient, and  $g_{\text{approx}}$  is an approximate meta gradient obtained with various popular meta learning algorithms including SAMA. Right:  $\|\lambda_t - \lambda^*\|_2$ , where  $\lambda_t$  is the meta parameter  $\lambda$  after  $t$  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 a stable convergence to the optimal solution at a comparable speed.
<sup>5</sup>We used smaller  $\beta$  than in the original paper [19] (1 vs 0.1), to amplify the "non-identitiness" of the base Jacobian.
F 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 [67], and the result is presented below.
Base Jacobian	Algo Adapt	Distributed	Accuracy	Throughput	Memory
Finetuning (no meta learning baseline)	x	x	x	85.79	169.16	7.77
Iterative Diff (e.g. MAML) [14, 42]	x	x	x	85.78	28.07	22.94
Conjugate gradient (e.g. iMAML) [51]	x	x	x	86.78	65.14	22.03
Neumann series [40]	x	x	x	86.65	67.03	19.70
DARTS (or T1-T2) [38, 41]	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 8: Ablation results on AGNews
Base Jacobian	Algo Adapt	Distributed	Accuracy	Throughput	Memory
Finetuning (no meta learning baseline)	x	x	x	78.16	144.39	6.60
Iterative Diff (e.g. MAML) [14, 42]	x	x	x	80.25	24.24	22.03
Conjugate gradient (e.g. IMAML) [51]	x	x	x	81.01	56.27	21.92
Neumann series [40]	x	x	x	79.92	57.85	19.75
DARTS (or T1-T2) [38, 41]	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
Table 9: Ablation results on IMDB
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.