PaperlayerView exact revision on arXiv

Source-linked AI summary

Making Scalable Meta Learning Practical

Sang Keun Choe, Sanket Vaibhav Mehta, Hwijeen Ahn, Willie Neiswanger, Pengtao Xie, Emma Strubell, Eric Xing

arXiv:2310.05674v2cs.LGcs.AI

TL;DR

Gradient-based meta learning faces high compute and memory costs, training instability, and limited distributed-training support. SAMA combines implicit-differentiation and systems advances to support adaptive optimizers while avoiding explicit second-order gradient computation. Across large-scale benchmarks, it reports throughput and memory gains and strong data-optimization results.

  • Problem

    Gradient-based meta learning faces substantial compute and memory costs, training instability, and a lack of efficient distributed support, while many methods assume SGD despite large models commonly using adaptive optimizers.

    Source: §1 · p. 1

  • Method

    SAMA combines implicit-differentiation and systems advances, including adaptive-optimizer support and distributed-training techniques, while avoiding explicit second-order gradient computation.

    Source: §1 · p. 2, §3.2 · p. 5

  • Results

    Up to 1.7/4.8× increase in throughput and 2.0/3.8× decrease in memory consumption respectively on single-/multi-GPU setups versus baseline meta-learning algorithms; SAMA-based data optimization improves text accuracy and data-pruning results.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The benchmarks demonstrate the initial potential of scalable meta learning across large-scale learning and data-optimization tasks.

    Source: §1 · p. 2, §6 · p. 10

    The authors could not test SAMA on models with 1B+ parameters because of limited computational resources.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Meta learning optimizes the inductive biases of a machine-learning program for specified objectives, but gradient-based methods face major scalability barriers. The paper introduces SAMA and reports efficiency gains and data-optimization results across large-scale benchmarks.

1 Introduction

  • Meta learning formulates the search for inductive biases, such as training data or neural architecture, as optimization for objectives such as fairness or quick generalization.Applications include hyperparameter and data optimization, neural architecture search, learned optimizers, and few-shot learning.Source: §1 · p. 1
  • Gradient-based meta learning suffers from compute and memory costs, instability, and limited distributed support, while large models commonly use adaptive optimizers rather than SGD.These constraints complicate applying SGD-based methods to large models and scaling beyond single-GPU setups.Source: §1 · pp. 1–2
  • SAMA achieves up to 1.7/4.8× increase in throughput and 2.0/3.8× decrease in memory consumption on single-/multi-GPU setups versus baseline meta-learning algorithms.SAMA-based data optimization also improves text-classification accuracy and achieves state-of-the-art results in small- and large-scale data pruning.Source: Figure 1 · p. 2, §1 · p. 2
  • SAMA combines an identity approximation for the base Jacobian, chain-rule expansion of the meta Jacobian, and a communication strategy using communication-computation overlap.The paper presents these as solutions to base-Jacobian inversion, adaptive-optimizer support, and custom backward-pass distribution challenges.Source: §1 · p. 2

2 Background: Gradient-Based Meta Learning

Gradient-based meta learning treats the problem as bilevel optimization and computes meta-gradients through the dependence of the base solution on meta parameters. This section motivates implicit differentiation as a more efficient alternative to iterative differentiation.

2 Background: Gradient-Based Meta Learning

  • Meta learning is commonly formulated as bilevel optimization, with meta parameters and objectives at the upper level and the machine-learning program represented by a base-level optimization.The meta objective is evaluated at the base optimizer’s solution.Source: §2 · p. 3
  • Gradient-based meta learning uses the chain rule to combine a best-response Jacobian with a direct gradient, enabling high-dimensional meta-optimization.The approach is useful when inductive-bias search spaces, such as data-importance weights, grow with models and datasets.Source: §2 · p. 3
  • The paper adopts implicit differentiation over iterative differentiation because it offers better computation and memory efficiency.Both efficiencies matter to the goal of scalable gradient-based meta learning.Source: §2 · p. 3
  • Implicit differentiation computes the best-response Jacobian using the implicit function theorem and views base optimization through the fixed point of an iterative solver.The solver may be SGD or Adam.Source: §2 · p. 3
  • Exact implicit differentiation requires convergence of the base optimization, so practical methods approximate the solution with a small number of unrolled updates and compute meta-gradients using implicit-differentiation equations.This focuses the scalability challenge on meta-gradient computation, while standard backpropagation has established scaling techniques.Source: §2 · p. 3

3 Scaling Meta Learning

The paper attributes meta-learning scalability limits to base-Jacobian inversion, inadequate adaptation to adaptive optimizers, and inefficient custom distributed backward passes. SAMA addresses these with an identity approximation, optimizer-aware differentiation, and a communication-overlap strategy.

3.1 Base Jacobian Inverse

  • Base-Jacobian inversion can require O(n_b^3) computation, while Hessian-vector products add memory and compute costs, noisy-gradient bias, and difficulty using first-order distributed-training optimizations.SAMA approximates the base Jacobian with the identity; the paper describes this as preconditioning the meta-gradient, and reports accurate estimation in its biased-regression analysis even when the true Jacobian is not identity.Source: §3.1 · p. 4
  • The identity approximation can support more than one unrolled update, allowing less frequent expensive meta-gradient computations than methods limited to a single unroll step.The paper notes this distinction from methods that calculate the meta Jacobian at initialization rather than convergence.Source: §3.1 · p. 4

3.2 Algorithmic Adaptation for Adaptive Optimizers

  • Assuming vanilla SGD when base optimization actually uses an adaptive optimizer can produce an incorrect meta-gradient and contribute to instability and reduced performance.The mismatch is particularly relevant because large models are commonly trained with optimizers such as Adam.Source: §3.2 · p. 4
  • SAMA expands the meta Jacobian by the chain rule to account for the optimizer update; for element-wise optimizers, the adaptation matrix is diagonal with O(n_b) computation and memory complexity.The expansion reduces to an identity matrix for SGD, and SAMA uses central differences and associative matrix multiplication for the costly second-order matrix-vector product.Source: §3.2 · p. 5

3.3 Efficient Distributed Training & Implementation

  • Meta-gradient algorithms need custom backward passes that complicate distributed data parallelism, and existing support may still incur communication inefficiency.SAMA performs two backward passes locally, then overlaps computation with communication in the final pass; synchronization occurs only once.Source: §3.3 · pp. 5–6, Figure 2 · p. 5

4 Experiments

The experiments evaluate scalable meta learning through noisy text classification, continued pretraining, and data pruning. Across these applications, SAMA-based data optimization improves task performance, while its pruning method works across dataset scales.

4 Experiments

  • The experiments target data optimization because large-scale tasks can be costly, often use adaptive optimizers, and offer a test of SAMA across domains.Applications include noisy finetuning, continued pretraining, and scale-agnostic pruning.Source: §4 · p. 6

4.1 Noisy Finetuning of Large Language Models

  • SAMA-based data reweighting and label correction improve noisy text-classification accuracy over finetuning and self-training across WRENCH benchmarks, with small clean meta-level data.SAMA also outperforms SAMA-NA, supporting the importance of algorithmic adaptation for an adaptive-optimizer base learner.Source: §4.1 · pp. 6–7
  • SAMA is generally more compute- and memory-efficient than Neumann Series and conjugate gradient, with the efficiency gap widening across multiple GPUs.The comparison uses AGNews with a fixed global batch size of 48; experiments reached up to four V100 GPUs, leaving larger-scale evaluation for future work.Source: §4.1 · p. 7, Table 2 · p. 7

4.2 Continued Pretraining of Large Language Models

  • SAMA-based data optimization improves downstream performance on almost all continued-pretraining datasets compared with the evaluated alternatives.The study uses task-specific auxiliary data for comparison with TARTAN-MT and reports superior memory efficiency as RoBERTa model size increases.Source: §4.2 · p. 8

4.3 Scale-Agnostic Efficient Data Pruning

  • The pruning method replaces hand-designed metrics with meta-learned sample importance weights, using both training and meta levels without additional validation data.It modifies MWN by using loss and prediction uncertainty as inputs, applying SAMA, and enabling distributed training.Source: §4.3 · p. 8
  • GBML-based pruning with SAMA outperforms heuristic methods and works across dataset scales, with test-accuracy improvements at pruning ratios 0.1 and 0.2 on ImageNet-1k.SAMA’s search time is comparable to or shorter than heuristic methods; on four GPUs it achieves 15–20× speedup over original MWN, which lacks distributed training support.Source: §4.3 · pp. 8–9, Figure 3 · p. 9

5 Related Work

Related work presents iterative and implicit differentiation as two major gradient-based meta-learning approaches, alongside a range of applications. It also notes that existing systems libraries have narrower task focuses or support implicit differentiation.

Algorithms

  • Iterative differentiation stores intermediate optimization states, so its memory and compute costs grow linearly with unrolling steps; implicit differentiation depends on the final state instead.Implicit methods require base-Jacobian inversion, for which variants use Neumann series, conjugate gradient, or the Nystrom method.Source: §5 · p. 9

Applications

  • Meta learning has been applied to few-shot learning, architecture search, hyperparameter optimization, data optimization, and reinforcement learning.These applications share a bilevel-optimization formulation and relate conceptually to optimal design and inductive or declarative programming.Source: §5 · p. 9

Systems

  • Higher, learn2learn, and TorchMeta focus on few-shot image classification, while JaxOpt and Betty support implicit differentiation.The authors implement SAMA in Betty, whose architecture is designed for systems optimization in large-scale meta learning.Source: §5 · pp. 9–10

6 Conclusion

SAMA combines algorithmic and systems advances to address meta-learning scalability, and experiments show scaling efficiency across multiple benchmarks. Its design emphasizes compatibility with large-scale learning systems and supports arbitrary base-level optimizers.

6 Conclusion

  • SAMA showed scaling efficiency across multiple benchmarks, while models with 1B+ parameters remained untested because of limited computational resources.Future work targets compatibility with systems techniques needed for 10B+ models and large-scale neural architecture search for Transformer-family models.Source: §6 · p. 10

Q. Why do we study scalable meta learning?

  • The paper follows “The Bitter Lesson” by focusing on scalability, unlike work that typically optimizes performance under limited, usually single-GPU, compute budgets.The passage also identifies data optimization [54], hyperparameter optimization [16], few-shot learning [14] [51], and adversarial learning [45] as meta-learning applications.Source: Q. Why do we study scalable meta… · p. 15

Q. What are the design principles behind scalable meta learning?

  • SAMA avoids explicit higher-order gradient computations and is compatible with first-order-gradient systems support, half-precision training, and activation checkpointing.Compatibility with existing systems and algorithms in the large-scale learning ecosystem is a stated design principle.Source: Q. What are the design principles… · p. 15
  • SAMA supports arbitrary base-level optimizers through algorithmic adaptation, unlike approaches that assume naive SGD or use algorithm-specific update rules.Those restrictions are described as hampering algorithm compatibility.Source: Q. What are the design principles… · p. 15

B Experiment Details

The appendix section introduces experiment details, including hyperparameters, baselines, and compute resources.

B Experiment Details

  • Experiment details cover hyperparameters, baselines, and compute resources for the Section 4 experiments; code is available in the supplementary material and official Betty GitHub repository.The passage identifies the repository with a footnote marker but does not provide its contents here.Source: B Experiment Details · p. 16

B.1 Noisy Finetuning of Large Language Models

The appendix reports training settings, baselines, and compute resources for noisy language-model finetuning and continued pretraining experiments.

B.1 Noisy Finetuning of Large Language Models

  • Noisy-finetuning runs used 1,000 iterations, batch size 32, majority voting, and WRENCH datasets; baselines were naive finetuning and COSINE self-training.The main experiment used one RTX 2080Ti GPU; throughput-memory analysis used four Tesla V100 GPUs.Source: §B.1 · p. 16

B.2 Continued Pretraining of Large Language Models

  • Continued-pretraining runs used 100 epochs, batch size 16, and a maximum sequence length of 256; baselines were DAPT and TARTAN-MT.The main experiment used one Tesla V100 GPU, while the memory-versus-model-size analysis used one RTX A6000 GPU.Source: §B.2 · p. 16

B.3 Scale-Agnostic Efficient Data Pruning

The appendix gives data-pruning training settings and baselines, then notes how Adam adaptation matrices are specified and generalized to other adaptive optimizers.

B.3 Scale-Agnostic Efficient Data Pruning

  • ImageNet-1k pruning ran for 30 epochs before 120-epoch training, while CIFAR-10 pruning ran for 50 epochs before 200-epoch training.ImageNet-1k learning rate was decayed by 10 at epochs 40 and 80; CIFAR-10 used cosine learning-rate decay.Source: §B.3 · pp. 16–17
  • SAMA-based static pruning outperformed dynamic-pruning baseline DynaMS, though the passage notes that dynamic pruning may see all training data across epochs, complicating fair comparison.ImageNet-1k baselines were EL2N, GraNd, and DynaMS; CIFAR-10 baselines were GraNd, forgetting, and margin.Source: §B.3 · p. 17
  • The uncertainty measure U is the difference between predictions from the current model and the exponentially moving averaged model.Source: §B.3 · p. 17

C Algorithmic Adaptation for Adam Optimizer

  • For Adam, the adaptation matrix is specified using its first and second gradient moments, m and v, and learning rate γ; similar matrices can be derived for other adaptive optimizers.Source: C Algorithmic Adaptation for Adam… · p. 18

D The Effect of Scaling in Model-Agnostic Meta Learning

This section investigates whether SAMA’s efficiency can support larger networks in few-shot image classification. In preliminary Omniglot experiments, accuracy improves consistently as model size increases.

D The Effect of Scaling in Model-Agnostic Meta Learning

  • The study asks whether SAMA’s compute and memory efficiency can improve few-shot generalization by scaling network size.It follows the iMAML setup, which explicitly enforces proximity to initialization weights with L2 regularization, but uses SAMA rather than conjugate-gradient-based iMAML.Source: D The Effect of Scaling in… · p. 18
  • Larger networks consistently improve few-shot classification accuracy on Omniglot 20-way 1-/5-shot tasks.The experiments vary network width using a basic 4-layer CNN.Source: D The Effect of Scaling in… · p. 18, Figure 4 · p. 18
  • The authors leave open how to define scale across model and dataset sizes in meta learning, and call for further research beyond few-shot image classification.Meta learning involves both meta and base optimization problems.Source: D The Effect of Scaling in… · p. 18

E Justification of the Identity Approximation for Base Jacobian

This section tests SAMA’s identity approximation for the base Jacobian in biased regression, where the true Jacobian is not an identity matrix. SAMA’s meta-gradients align less accurately than CG’s, yet converge stably to the optimum at a comparable speed.

E Justification of the Identity Approximation for Base Jacobian

  • The evaluation measures the identity approximation’s effects on the meta-gradient and optimal meta solution in biased regression, where the true base Jacobian is non-identity.The setting permits closed-form expressions for the base Jacobian, meta-gradient, and optimal meta solution.Source: E Justification of the Identity… · p. 19
  • The comparison includes SAMA, conjugate-gradient and Neumann-series methods that explicitly approximate the inverse base Jacobian.With β = 0.15, the authors perform 100 meta updates and track gradient alignment and distance from the optimum.Source: E Justification of the Identity… · p. 19
  • SAMA achieves high directional alignment with the ground-truth meta-gradient and stable convergence to the optimum at a speed comparable to other algorithms, despite slightly lower accuracy than CG.Figure 5 compares gradient cosine similarity and distance from the optimal meta parameter over 100 meta updates.Source: E Justification of the Identity… · p. 19, Figure 5 · p. 19

F Extensive Ablation Study

An ablation study on AGNews and IMDB evaluates SAMA’s components using test accuracy, GPU memory use, and throughput. The reported findings attribute efficiency gains to the identity approximation and communication-optimized distributed training, while optimizer adaptation improves meta-learning performance at minimal compute and memory cost.

F Extensive Ablation Study

  • The ablation study evaluates SAMA’s base-Jacobian inverse, adaptive-optimizer adaptation, and distributed training components on AGNews and IMDB.It compares test accuracy, GPU memory usage, and throughput against baseline meta-learning algorithms.Source: F Extensive Ablation Study · p. 20, Table 8 · p. 20, Table 9 · p. 20
  • The identity approximation improves memory and compute efficiency, optimizer adaptation improves meta-learning performance at minimal compute and memory cost, and communication-optimized distributed training further improves efficiency.Source: F Extensive Ablation Study · p. 20
Loading 2310.05674v2…