Optimizing ML Training with Metagradient Descent

1 Introduction

• Training configuration spans choices such as data, architecture, and initialization, but exhaustive search explores only a small part of the design space.Source: §1 · p. 1

1.1 Contributions

• The proposed recipe makes the objective differentiable, computes scalable metagradients with REPLAY, designs metasmooth training routines, and performs metagradient descent.Source: §1 · p. 1, §1.1 · p. 2
• 2x larger performance improvement than the previous DataComp-small leader is achieved for CLIP pre-training data selection.Source: §1.1 · p. 2
• MGD substantially improves instruction-tuning data selection for Gemma-2B over existing selection methods and full-data training.Source: §1.1 · p. 2
• An order-of-magnitude improvement in DNN data poisoning drops CIFAR-10 accuracy from 92%→78%, versus 91% for the best previous attack [LKY23].Source: §1.1 · p. 2
• MGD efficiently finds a CIFAR-10 learning-rate schedule matching the performance of a schedule found by grid search.Source: §1.1 · p. 2

2.1 What is a metagradient?

• The framework models training as an iterative state sequence, where each optimizer update maps the current state and metaparameters to the next state.For gradient descent, the state contains model parameters; for algorithms such as Adam, it also includes gradient moments.Source: §2.1 · p. 3
• A metagradient is the gradient of a training output with respect to continuous metaparameters controlling aspects of the training setup.The setup may include learning rates, weight decay, datasets through continuous relaxations, or other configurable training choices.Source: §2.1 · p. 3

2.2 Warmup: Metagradients via autodifferentiation

• Naive reverse-mode autodifferentiation is impractical because storing intermediates across all training operations can require terabytes even for toy MNIST metagradient computation.Step-wise autodifferentiation reduces the per-step differentiation burden but still stores every optimizer state, remaining too space-intensive at scale.Source: §2.2 · pp. 4–5
• Step-wise autodifferentiation computes the metagradient by traversing optimizer states backward and summing per-step contributions obtained from one-step differentiations.It first stores the training trajectory, then propagates gradients from the final state back to the initial state.Source: §2.2 · pp. 4–5

2.3 REPLAY

• REPLAY computes exact metagradients for large-scale iterative training using O(k log_k(T)) space and additional replayed training computation.It recursively saves and replays optimizer states, balancing storage against recomputation through a lazy k-ary tree.Source: §2.3 · pp. 5–6
• REPLAY uses deterministic training and a lazy k-ary tree to recover optimizer states in reverse order while deleting states that are no longer needed.The method is presented as a rematerialization technique for general, non-reversible learning algorithms rather than relying on state inversion.Source: §2.3 · pp. 5–6

3 Designing metasmooth training routines

• Standard training routines can produce ±∞-valued and generally unhelpful metagradients because the metaparameter optimization landscape is often non-smooth.This motivates modifying learning algorithms before applying metagradient optimization.Source: §3 · p. 6

3.1 The metaparameter optimization landscape

• Metasmoothness adapts smoothness analysis to metagradients, measuring how rapidly directional derivatives change across metaparameter perturbations.The proposed metric avoids exhaustive searches and metagradient-computation errors, requiring only three training-function evaluations.Source: §3.1 · p. 7
• Empirical metasmoothness compares weighted sign agreement between finite-difference metagradients at nearby metaparameter settings and can be computed with three calls to the learning algorithm.It operates in model-parameter space and yields an interpretable value in [−1, 1].Source: §3.1 · p. 7

3.2 Estimating and improving metasmoothness

• A fixed-menu search over training modifications identifies metasmooth algorithms, with BatchNorm before nonlinearities and final-layer output scaling crucial alongside batch size and width.The considered modifications are not exhaustive.Source: §3.2 · pp. 8–9
• The empirical estimate uses dataset perturbations at z = 0, a random direction, three trained parameter states, finite differences, and parameter-range weighting.This procedure estimates metasmoothness for ResNet-9 trained on CIFAR-10 with perturbations to 1000 training images.Source: §3.2 · p. 8
• Metasmooth models produce qualitatively smoother, more straightforward optimization landscapes than non-metasmooth models.Figure 5 compares perturbations to 1000 training examples and the resulting test-example losses for paired algorithms.Source: §3.2 · p. 9, Figure 5 · p. 9
• The most metasmooth algorithms retain near-optimal accuracy despite no clear metasmoothness–accuracy trend, supporting joint search over smoothness and performance.The paper notes that additional searching might identify more accurate metasmooth models.Source: §3.2 · pp. 9–10
• Metasmoothness strongly predicts metagradient utility: smoother configurations can be optimized more effectively by gradient ascent on validation loss.The experiment uses metagradients computed by REPLAY and measures the resulting change in model loss.Source: Figure 4 · p. 8, §3.2 · p. 10

4 Applications

• Across applications, MGD frames the task as optimization, makes the learning algorithm smooth, and applies first-order metagradient optimization.The learning-rate application illustrates REPLAY’s flexibility for metaparameters that affect models through optimization trajectories rather than directly changing the loss landscape.Source: §4 · p. 10, §4.4 · p. 16

4.1.3 Results

• MGD dataset selection greatly outperforms the previous state-of-the-art on DataComp-small, with only a few optimization steps needed to surpass prior methods.The improvement over the previous state-of-the-art is roughly as large as that method’s improvement over randomly chosen data.Source: §4.1.3 · p. 12, Figure 6 · p. 12

4.2 Selecting instruction-tuning data

• MGD selects strictly better instruction-tuning datasets than training on all data or LESS [XMG+24], improving accuracy by +1.5% on BBH and +0.8% on MMLU.Validation loss improves over optimization steps but shows signs of overfitting while optimizing 270,679 datapoint-count weights from few test samples.Source: §4.2 · p. 13

4.3.1 Setup

• MGD directly optimizes continuous poisoned samples by maximizing validation loss through metagradients, using projected gradient ascent under a fixed poisoning budget.The poisoning attack replaces a fixed ε-fraction of training samples and projects image-label pairs back into the valid sample space.Source: §4.3.1 · p. 14, §4.3.2 · p. 14, §4.3.3 · p. 15

4.3.3 Evaluation

• 14% accuracy degradation reduces CIFAR-10 test accuracy to 78% when corrupting 2.5% of training data, versus less than 1% for GradCancel [LKY23].The attack modifies 1000 of 40,000 training samples; with a non-metasmooth learner, performance still drops from 92.8% to 82.6%.Source: §4.3.3 · p. 15

4.4 Finding a learning rate schedule

• MGD optimizes learning-rate keypoints from a flat initialization and matches grid search over standard schedules despite naïve starting hyperparameters.The schedule is parameterized by evenly spaced keypoints with linear interpolation; the comparison grid contains over 1,000 one-cycle schedules.Source: §4.4 · p. 16, §4.4.2 · p. 16, §4.4.1 · p. 16

4.4.3 Evaluation

• MGD learned the schedule from a highly suboptimal flat learning-rate schedule by minimizing loss on the reserved 90% of the test set.Final schedules were evaluated on the held-out 10% test set across five unseen random seeds.Source: §4.4.3 · p. 17
• 50 MGD steps matched the state-of-the-art onecycle schedule found by grid search over more than 1,000 configurations.The comparison used Figure 10, but grid search can run in parallel whereas MGD steps run sequentially.Source: §4.4.3 · p. 17
• MGD is not advised for low-dimensional hyperparameters already thoroughly optimized by grid search, such as CIFAR-10 learning-rate schedules [SN17] [LA19].Future work could examine high-dimensional hyperparameters, including per-parameter or per-layer learning rates and weight decays, attention hyperparameters, and gradient preconditioners.Source: §4.4.3 · p. 17

5 Discussion

• Metagradients remain more expensive than one ordinary training run because they require backward-over-backward computation and currently lack specialized software.The operation requires 2–3 times the operations of a backward pass and float32/tensorfloat32 computation.Source: §5 · p. 17
• Successful metagradient optimization requires smooth model training, motivating further study of metasmoothness and non-smooth optimization alternatives.Future directions include connections to numerical stability, optimizer and architecture design, and non-smooth optimization methods.Source: §5 · p. 17
• Metagradients apply to data poisoning, training-data selection, and the design of training routines and architectures with first-order methods.Small proxy models and engineering improvements are proposed as ways to reduce computational cost.Source: §5 · p. 18

6 Related work

• The related-work literature uses metagradient, hyper-gradient, and outer gradient as synonymous terms.These terms describe gradients through the model-training process.Source: §6.1 · p. 18

6.1 Calculating metagradients

• Implicit differentiation can efficiently estimate metagradients under strong convexity but requires restrictive assumptions, inverse-Hessian approximations, and loss-expressed metaparameters.It can differentiate weight decay but not learning rate.Source: §6.1 · p. 18
• Automatic differentiation avoids some implicit-method restrictions but faces storage, compute, reversibility, approximation, and numerical-instability challenges.Prior methods use small settings, forward-mode AD, truncated steps, or reversible algorithms to manage these constraints.Source: §6.1 · pp. 18–19

6.2 Applying metagradients

• Metagradients have been applied to distillation, data selection, meta-learning, learning-rate and weight-decay selection, augmentation, architecture search, and data attribution.Existing applications generally either use direct AD in small settings or implicit-function estimates.Source: §6.2 · p. 19

7 Conclusion

• Metagradients were successfully applied to models up to 2B parameters for data selection, Huber data poisoning, and per-iteration learning-rate search.These applications cover CLIP pretraining, instruction fine-tuning, accuracy reduction, and high-dimensional hyperparameter search.Source: §7 · p. 19

A.1 Warmup: Step-wise AD

• Step-wise automatic differentiation stores all T optimizer states and is feasible for small setups because each differentiation call has constant memory overhead.The appendix presents it as a warmup before REPLAY.Source: §A.1 · p. 26

A.2 REPLAY

• REPLAY computes metagradients by reverse-order traversal of optimizer states using deterministic training and a memory–compute trade-off.It uses O(k log_k(T)) space and an additional T log_k(T) model steps.Source: §A.2 · p. 26

A.2.1 Lazy k-ary tree

• The lazy k-ary tree recursively replays training to recover optimizer states in reverse order while deleting states no longer needed.Leaves correspond to training states, and reverse in-order traversal yields them in reverse order.Source: §A.2.1 · pp. 26–27
• The traversal stores at most k log_k(n) optimizer states and requires O(n log_k(n)) replay steps, a multiplicative O(log_k(n)) overhead over training.Storage and compute depend on the tree branching factor k.Source: §A.2.1 · p. 27

B.1 Omitted Figures

• Figure 11 examines factors affecting metasmoothness when training a ResNet-9 on CIFAR-10.The caption refers readers to Section 3 for details.Source: Figure 11 · p. 28
• Additional omitted material includes loss-landscape visualizations and examples labeled non-smooth.The supplied passages do not state the visual comparisons or outcomes.Source: Figure 12 · p. 29, §B.1 · p. 29

• The appendix presents pseudocode for the main DataComp dataset-selection algorithm.Source: C Metagradients for DataComp · p. 30
• It explains how metagradients are applied to CLIP models.Source: C Metagradients for DataComp · p. 30
• It also details adaptations specific to the DataComp setting.Source: C Metagradients for DataComp · p. 30

C.1 Dataset Selection Using MGD

• MGD adapts the CLIP contrastive loss to incorporate dataweights and produces a surrogate batch loss for metagradient optimization.The appendix notes that the simplified formulation omits cross terms, which are included in the adapted loss.Source: §C.1 · pp. 30–31
• MGD relaxes data selection into a differentiable loss whose metagradients with respect to dataweights remain meaningful, including at zero weights.The relaxation recovers the normal CLIP loss at zero dataweights and models single-example additions correctly.Source: §C.1 · p. 30
• The adapted CLIP loss empirically yields meaningful dataweight metagradients and strong dataset-selection results.Source: §C.1 · p. 31

C.2 Scaling MGD for CLIP and DataComp

• Metagradients are made scalable through manual checkpointing, minibatched embedding computation, and parallelization across GPUs.These choices address the memory cost of CLIP forward passes and their gradients.Source: §C.2 · p. 31
• Data is optimized by selecting and rewriting shards, while retaining the possibility that datapoints removed at one step can return later.The implementation uses eight base shards and computes metagradients over one shard at a time.Source: §C.2 · p. 31
• Top- and bottom-metagradient updates perform better empirically than randomly sampling points to add or remove.Shard sizes are subsequently fixed so sign-based updates do not cause them to grow indefinitely.Source: §C.2 · p. 31
• Eight-way loading preserves randomness and reduces duplicate datapoints within CLIP batches, avoiding contrastive-loss problems from repeated examples.Source: §C.2 · p. 31
• Deleting stale rewritten shards keeps total disk usage constant in the number of optimization steps.The algorithm retains only the eight most recently written shards.Source: §C.2 · p. 31

• DataComp-small selects from 12.8 million candidate samples while fixing the CLIP architecture, learning algorithm, and 12.8 million-sample training budget.Source: §C.3 · p. 32
• The experiments match DataComp-small hyperparameters except for a positive εroot added for numerical stability, with εroot and the metagradient step chosen to maximize metasmoothness.Source: §C.3 · p. 32
• The candidate pool contains only around 80% of the original DataComp pool because broken or changed URLs were removed.All reported results use this incomplete subset.Source: §C.3 · p. 32
• Disjoint target and validation sets were created where possible so metagradients used target data while evaluation measured generalization on validation data.Target splits could not be constructed for every task, and some datasets were excluded or not used for metagradients.Source: §C.3 · pp. 32–33

D Selecting IFT data

• MGD selects instruction-fine-tuning data for LoRA adaptation of Gemma-2B from four combined datasets.Source: D Selecting IFT data · p. 34
• MGD selects among iterates and hyperparameters using top validation-loss steps followed by full-training-set accuracy because the validation set is small.The procedure is described as akin to Pareto optimization.Source: D Selecting IFT data · p. 34
• MGD is compared against training on the full dataset and LESS-T data selection.Source: D Selecting IFT data · p. 34
• The MGD runs jointly select learning rate and εroot because both substantially affect training smoothness and metagradient quality.Much larger or smaller maximum learning rates, and much larger εroot values, produce non-smooth training.Source: D Selecting IFT data · p. 34

IFT results

• MGD validation loss improves over metagradient steps, but the optimized-sample loss increasingly separates from validation loss, indicating overfitting with too many iterates.Source: Figure 13 · p. 35

E.1 Background on Gradient Cancelling attack

• Gradient Cancelling constructs a malicious parameter target and then poisons data so retraining reaches it.The target is generated independently, while poison values are optimized to induce a near-zero gradient at that target.Source: §E.1 · p. 36
• The parameter-generation phase produces θp with low test accuracy while keeping it close to the clean-training parameter θ0.The target may be generated using Gradient-based Parameter Corruption (GradPC) [LKY23].Source: §E.1 · p. 36
• The poison-crafting phase minimizes the combined clean-and-poisoned loss gradient evaluated at θp.This optimization aims to make the gradient at the malicious parameter approach zero.Source: §E.1 · p. 36

E.2 Metasmooth hyperparameters

• The ResNet-9 [Jor24] CIFAR-10 poisoning experiments use normalization, random horizontal flips, and two-pixel random translations as augmentations.The cited table contains the experiment’s hyperparameters.Source: Table 5 · p. 36

F LR optimization

• The learning-rate optimization grid searches all 528 combinations of the listed hyperparameter values.Figure 14 presents graphs of the learned learning-rate schedules.Source: Table 6 · p. 37, Figure 14 · p. 37