PaperlayerView exact revision on arXiv

Source-linked AI summary

VeLO: Training Versatile Learned Optimizers by Scaling Up

Luke Metz, James Harrison, C. Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, Jascha Sohl-Dickstein

arXiv:2211.09760v1cs.LGmath.OCstat.ML

TL;DR

Deep learning still relies on hand-designed optimizers, while learned optimizers face challenges in scaling meta-training and balancing optimizer overhead against performance. VeLO uses a neural-network update rule meta-trained across diverse tasks; it outperforms heavily tuned baselines across more than 80 tasks, with limitations on models beyond approximately 500M parameters.

  • Problem

    Learned optimization lacks the scale of deep learning because diverse representative tasks are difficult to collect and costly to meta-train, while larger learned optimizers must balance per-step overhead against performance.

    Source: §1 · p. 3

  • Method

    VeLO is a neural-network learned optimizer meta-trained at greater scale, using a per-tensor network to generate the weights of a small per-parameter network.

    Source: §1 · p. 3, §3.1 · p. 6

  • Results

    Across more than 80 diverse optimization tasks, VeLO outperforms heavily hyperparameter-tuned baselines; on >85% of tasks it also outperforms the 1,000-trial NAdamW search.

    Source: §4.1.2 · p. 9, §7 · p. 20

  • Takeaways & Limitations

    VeLO requires no hyperparameter tuning and works across a wide variety of neural-network training tasks, supporting the viability of general-purpose learned optimization for deep learning.

    Source: §1 · p. 3

    VeLO’s performance lags behind baselines or decreases beyond approximately 500M parameters, a regime far outside its meta-training distribution.

    Source: §4.4.1 · p. 14, §4.4 · p. 14

Abstract

from arXiv · show

While deep learning models have replaced hand-designed features across many domains, these models are still trained with hand-designed optimizers. In this work, we leverage the same scaling approach behind the success of deep learning to learn versatile optimizers. We train an optimizer for deep learning which is itself a small neural network that ingests gradients and outputs parameter updates. Meta-trained with approximately four thousand TPU-months of compute on a wide variety of optimization tasks, our optimizer not only exhibits compelling performance, but optimizes in interesting and unexpected ways. It requires no hyperparameter tuning, instead automatically adapting to the specifics of the problem being optimized. We open source our learned optimizer, meta-training code, the associated train and test data, and an extensive optimizer benchmark suite with baselines at velo-code.github.io.

1 Introduction

Scaling meta-learning is constrained by the difficulty and cost of collecting diverse tasks, while larger learned optimizers must balance performance against per-step overhead. VeLO addresses this challenge with a more extensively meta-trained learned optimizer and is provided for use with JAX models.

1 Introduction

  • Scaling meta-learning is difficult because representative tasks lack a standardized collection method and realistic tasks can make meta-training computationally costly.The paper also notes increased compute, larger datasets, and engineering effort as challenges of scaling.Source: §1 · p. 3
  • Larger learned optimizers may need fewer iterations but incur greater per-step overhead, so their performance must be balanced against simpler optimizers’ training and tuning opportunities.The paper frames this as a trade-off between optimizer overhead and performance.Source: §1 · p. 3
  • VeLO is meta-trained at far greater scale and performs better with less computational overhead, enabling training of much larger models.The work uses three orders of magnitude more tasks and two orders of magnitude more compute than prior learned-optimizer work.Source: §1 · p. 3

1.1 Try VeLO

  • VeLO is designed for JAX models using Optax, and a Colab notebook trains implemented test problems.The section provides an Optax setup and update example.Source: §1.1 · p. 3

2 Problem Setting: Learned Optimization

The section contrasts fixed-form SGD updates with learned update rules whose parameters are meta-trained for an optimization objective. Learned updates can use richer training information, while meta-training seeks update-rule parameters that work well on tasks.

2 Problem Setting: Learned Optimization

  • SGD updates parameters using the minibatch-loss gradient scaled by a learning rate, its only meta-parameter.The update rule can be written as U_SGD(g; α) = αg.Source: §2 · p. 4
  • A learned optimizer replaces SGD’s fixed-form update with a neural-network update rule parameterized by θ and taking gradients as input.The more expressive rule may train faster, but finding effective meta-parameters is harder.Source: §2 · p. 4
  • Learned update rules may also use parameter values, current loss, and information accumulated across training steps, such as gradient moving averages.The shared computation across parameters can allow the rule to work on networks of different sizes.Source: §2 · p. 4
  • Meta-training adjusts θ to optimize a chosen meta-loss, using methods including backpropagation, reinforcement learning, evolution, or Evolution Strategies.This work uses Evolution Strategies to estimate gradients of a Gaussian-smoothed meta-loss, which smooths sensitivity to θ and addresses instability.Source: §2 · p. 4
  • Meta-training cost can be amortized over a task distribution to learn an optimizer intended to generalize to new tasks.The learned optimizer is applied for N inner-training steps to produce final parameters φN, which depend on θ, before evaluating the meta-loss.Source: §2 · p. 5, Figure 2 · p. 5

3 Methods: Large Scale Optimizer Training

VeLO combines a hierarchical learned-optimizer architecture with a varied task distribution and meta-training choices aimed at performance across tasks. The methods also describe how task diversity, compute costs, and training progression shape the training procedure.

3.1 Learned Optimizer Architecture

  • A per-tensor LSTM and per-parameter MLP form a hierarchy that balances optimizer expressiveness with computational efficiency.The LSTM generates the small MLP’s weight matrices, avoiding extra conditioning inputs to the per-parameter network.Source: §3.1 · p. 6

3.2 Data: A Diverse Distribution of Tasks

  • Meta-training tasks vary model family, dataset, loss, and architectural hyperparameters across families including MLPs, ConvNets, ResNets, Transformers, RNNs, and auto-encoders.Task augmentations modify training dynamics through techniques such as re-parameterization, subspace gradient estimation, asynchrony, and changed precision.Source: §3.2 · p. 6
  • Because task runtimes can differ by more than three orders of magnitude, rejection sampling makes fast tasks more frequent in meta-training.Source: §3.2 · p. 7

3.3 Meta-Training

  • The meta-objective targets final training loss, which the authors report yields lower-loss optimizers than targeting average loss, but can produce non-monotonic intermediate loss.Loss may plateau or even increase during part of training.Source: §3.3 · p. 7
  • Full-length unrolls train each model to completion for each meta-gradient evaluation, making it straightforward to target final loss but less compute-efficient than truncated methods.The method normalizes each task’s gradients to unit length before averaging them across tasks.Source: §3.3 · p. 7
  • An increasing curriculum over training iterations and problem size speeds meta-training, alongside JAX vectorization and compilation and asynchronous data-parallel training on 1K–4K TPU accelerators.Training ran for approximately four weeks with batches increasing from 10K to 40K tasks.Source: §3.3 · p. 7

4 Evaluating Learned Optimizers

Evaluation spans canonical benchmarks and real-world models under varied distribution shifts, and reports both strong transfer and clear failure boundaries. VeLO is compared with tuned hand-designed optimizers across language, vision, and other tasks.

4.1 VeLOdrome and evaluation design

  • VeLOdrome contains 83 canonical, single-accelerator tasks; the evaluation also covers MLCommons workloads, real-world models, and failure cases under distribution shift.Hand-designed baselines range from learning-rate tuning to a 1,000-configuration NAdamW search.Source: §4 · pp. 7–8, §4.1 · p. 8, §4.1.1 · p. 8

4.1.2 Normalized Performance Across Tasks

  • Using performance normalized to tuned Adam, VeLO outperforms all learning-rate-tuned optimizers on every problem and, with one training run, beats the extensive 1,000-trial NAdamW hyperparameter search on >85% of tasks.Source: §4.1.2 · p. 9

4.2 MLCommons Tasks

  • On all but one of six MLCommons workloads, VeLO matches or outperforms tuned Adam’s best trial; evaluations also compare runs using 75% as many iterations.The comparison reports training loss and uses an Adam baseline with searched learning rate and weight decay.Source: §4.2 · pp. 9–10, Figure 4 · p. 10

4.3.1 NERF; 4.3.2 MLP-Mixer

  • On unseen NERF tasks, VeLO beats learning-rate-tuned baselines and matches extensively tuned NAdamW; on unseen MLP-Mixer tasks, it beats tuned baselines and matches or exceeds NAdamW.The MLP-Mixers are trained on 64x64 ImageNet.Source: §4.3.1 · p. 10, §4.3.2 · p. 11, Figure 5 · p. 11

4.3.3 Large-scale Vision Transformers; 4.3.4 Knowledge Distillation

  • On JFT ViT-B models VeLO matches or beats heavily tuned Adam but falls behind on larger ViT-L and ViT-H; on ImageNet distillation, it beats the published baseline across durations.The ViT-H model has 650M parameters; longer distillation runs use a 65K batch size.Source: §4.3.3 · p. 11, §4.3.4 · p. 11, Figure 6 · p. 12

4.3.5 Object Detection with Faster R-CNN; 4.3.6 Multi-Game Decision Transformers

  • VeLO beats the typical piecewise-rate SGD baseline on Faster R-CNN training loss and matches tuned Adam on multi-game Decision Transformers.The Decision Transformer is a 200M model trained on sequences from 46 Atari games.Source: §4.3.5 · p. 12, §4.3.6 · p. 12, Figure 7 · p. 13

4.3.7 Large Language Models; 4.3.8 Chemical Properties of Materials with Graph Networks

  • For 100M-parameter language models, VeLO with weight decay beats LR-tuned Adafactor at small batch sizes but underperforms at large batches; on materials GNNs, it beats Adam at short horizons and reaches lower loss at longer ones.The materials GNN results include VeLO using 46K steps versus 184K baseline steps; its advantage weakens on longer runs.Source: §4.3.7 · p. 13, §4.3.8 · p. 13, Figure 8 · p. 14

4.4.1–4.4.4 Limitations and Failure Cases

  • VeLO performance lags or decreases beyond about 500M parameters and worsens as training approaches or exceeds 200K steps; extended runs also struggle with non-random starting states, while Ant results trail Adam or fail to learn gait.For continuation, increasing the conditioned step count beats other continuation methods most often, but a complete single run is best overall most often.Source: §4.4.1 · p. 14, §4.4.2 · p. 15, §4.4.3 · pp. 15–16

5 Understanding Learned Optimizer Behavior

VeLO uses training-horizon information to shape its behavior, including implicit step-size schedules that vary across tasks and tensors. On the tested Transformers, it also handles critical batch sizes around 10x larger than baseline methods, though this advantage is not shown for larger models.

5 Understanding Learned Optimizer Behavior

  • The section experimentally characterizes VeLO’s behavior because learned optimizers can be more diverse and inscrutable than hand-designed optimizers.The authors describe their behavior as requiring techniques for studying black-box systems.Source: §5 · p. 17

5.1 VeLO Adapts to Training Horizon

  • Near the end of training, VeLO drops loss dramatically after using the target training horizon as an input.The experiments varied inner-training length from 1K to 10K steps for an MLP and ConvNet.Source: §5.1 · p. 17, Figure 12 · p. 18

5.1.1 VeLO Learns an Implicit Step Size Schedule

  • VeLO varies step sizes across tasks and tensors, with differences as large as 6 orders of magnitude, and learns warm-up and decay schedules.These schedules were not explicitly encoded by the authors and arose through meta-training.Source: §5.1.1 · p. 18, Figure 13 · p. 19

5.2 VeLO Can Have a Larger Critical Batch Size than Baseline Optimizers

  • On two Transformers, VeLO’s critical batch size is around 10x larger than baseline methods, but this increase was not observed for the larger models discussed.The comparison holds total examples seen constant while sweeping batch size; the larger models were far outside the meta-training distribution.Source: §5.2 · pp. 18–19

6 Related Work

Related work spans early learned update rules, hierarchical learned optimizers, and methods for meta-training them. VeLO builds on neural-network-parameterized optimizers trained across task distributions, while also relating to task-specific optimizers, hyperparameter controllers, and symbolic rules.

Early learned update rules

  • Learned optimization began with simple update rules [1992] and [2000], then Andrychowicz et al. [2016] trained an RNN optimizer by backpropagating through optimization.Later work explored varied optimizer architectures and meta-training algorithms.Source: §6 · p. 19

Closest related work

  • Closest to VeLO are hyperparameter-free neural optimizers trained across task distributions, including hierarchical methods from Wichrowska et al. [2017] and Metz et al. [2020a].Metz et al. [2020a] used a more realistic task distribution [Metz et al., 2020b] and an improved architecture.Source: §6 · p. 19

Meta-training methods

Task-specific and alternative approaches

  • Other lines include task-specific learned optimizers across many domains, architecture changes and fallback mechanisms, hyperparameter controllers, and symbolic optimizer rules.The cited examples span chemistry [Merchant et al., 2021], robustness [Metz et al., 2019b], and other application areas [p44]; architecture and symbolic-rule examples are detailed in [p45].Source: §6 · p. 20

7 Discussion and Outlook

VeLO combines scaled meta-training with a hierarchical HyperNetwork architecture, and the paper presents it as a hyperparameter-free optimizer that outperforms tuned baselines across more than 80 diverse tasks. The discussion highlights open architectural and training questions, while the appendices describe the optimizer’s state, network components, and architectural comparisons.

7 Discussion and Outlook

  • VeLO has no hyperparameters and outperforms heavily tuned baselines across more than 80 diverse optimization tasks, including large published models.The authors attribute improvements in generality and performance to increased meta-training compute and dataset size plus architectural changes.Source: §7 · p. 20

7.1 Open Questions; 7.2 Meta-Learned Algorithms are the Future

  • Open questions include richer architectures, conditioning on task computation graphs, targeting validation as well as training loss, reverse-engineering optimizer behavior, and more efficient meta-training.The discussion also argues that meta-learned algorithms may replace hand-designed components as compute and data grow, while noting meta-learning’s compute requirement is around 100 million times that of most supervised tasks.Source: §7.1 · p. 21, §7.2 · p. 21

B.1 Extended Architecture Overview

  • The proposed HyperNetwork architecture uses a tensor-level recurrent network to generate per-parameter network weights, increasing expressivity without extra compute at the lower level.This design addresses the representational limits of small per-parameter optimizers and the overhead of adding conditioning inputs.Source: §B.1 · p. 32

B.2 Optimizer State: Non-Learned Accumulators

  • VeLO tracks iteration, gradient moments, squared gradients, Adafactor-style accumulators, and loss histories at multiple timescales.The loss features include exponential moving averages and running minima.Source: §B.2 · p. 33

B.3 The Tensor-Level Recurrent Network

  • A 512-unit tensor-level LSTM processes progress, loss, moment, and tensor-rank features, with information mixed across tensors.Loss features indicate whether loss is decreasing, diverging, or stable; the progress features encode position relative to the target training horizon.Source: §B.3 · pp. 33–34, Figure 15 · p. 33, Figure 16 · p. 34

B.4 The Parameter-Level Network

  • The tensor-level network outputs a learning-rate control and weights that select per-tensor parameters for the per-parameter MLP, which produces updates.The per-parameter features are normalized by the second moment across the tensor before processing.Source: §B.3 · pp. 34–35, §B.4 · p. 35

B.5 Comparing our Architecture to Hyperparameter Controllers

  • Unlike hyperparameter controllers that adjust a hand-designed optimizer, VeLO’s hypernetwork controls the weights of a neural network.The authors describe this as a more complex version of controller-based parameterization.Source: §B.5 · p. 35

B.6 Experimental Validation of Our HyperNet Compared to Past Work

  • On small multitask problems, the HyperNet optimizer achieved low meta-loss and outperformed the compared learned optimizers, while prior hierarchical methods also performed well but cost more to compute.The comparison covered six optimizers and three random seeds.Source: §B.6 · p. 35, Figure 17 · p. 36

B.7 Understanding Computational Costs of VeLO

VeLO’s optimizer computation has higher fixed and per-parameter costs than hand-designed baselines, but the fixed-cost impact diminishes on larger models. The paper also describes a procedural task distribution spanning varied model and optimization families for meta-training.

B.7 Understanding Computational Costs of VeLO

  • VeLO costs 1.58e-9 seconds per parameter versus SGD’s 3.4e-11, while its constant overhead is about 2.6× higher and dominates below roughly one million parameters.The reported comparison is based on timing measurements and a fitted model.Source: §B.7 · p. 36, Figure 18 · p. 37
  • Increasing the per-parameter MLP size raises cost per parameter roughly linearly, whereas enlarging the per-tensor LSTM raises fixed overhead but leaves per-parameter cost constant.The LSTM overhead’s impact diminishes as parameter count grows, saturating at around one million parameters.Source: §B.7 · p. 36
  • For the PaLM 540B configuration, the authors estimate optimizer computation at about 0.3 seconds per device on TPUv3, compared with a 17.6-second training step on TPUv4.This is an extrapolation using 176 million parameters per device.Source: §B.7 · pp. 37–38

C Data: A Large, Diverse Distribution of Meta-Training Tasks

  • The task distribution combines procedural variation in model, dataset, training loss, and other settings to address the lack of a ready-made task distribution for learned-optimizer training.It includes image-classification and generative models, language models, and learned-optimizer training tasks.Source: C Data: A Large, Diverse Distribution… · pp. 38–39

C.1 A Configuration Language for Tasks

  • Task configurations specify nested objects and parameters separately from constructing tasks, and support static and vectorizable dynamic configuration elements.Shared components include sampled initializers and activation functions; base families include fully connected and convolutional networks, auto-encoders, VAEs, residual networks, RNNs, Transformers, and Vision Transformers.Source: §C.1 · p. 39, §C.2 · p. 40, §C.2.1 · p. 40

C.3.9 Learned Optimizer Training

The paper constructs task distributions and sampling procedures to make diverse meta-training computationally tractable. It meta-trains against end-of-training loss using evolution strategies, task curricula, and task-wise gradient normalization.

C.3.9 Learned Optimizer Training

  • Learned-optimizer training tasks vary optimizer architecture, hidden size, training length, outer batch size, and target image-classification task.The distribution also includes dynamic task configurations.Source: §C.3.9 · pp. 43–44

C.4 Controlling Task Size

  • Task runtimes vary by up to four orders of magnitude, so the authors train configuration-to-runtime and configuration-validity predictors to select feasible tasks efficiently.The runtime predictor is trained on approximately 10,000 sampled tasks and estimates runtimes in milliseconds rather than requiring direct timing and compilation.Source: §C.4 · pp. 44–45

C.5 Task Sampling and Augmentations

  • The sampling procedure selects tasks within sampled runtime thresholds and augments them with transformations such as re-parameterization, reduced batch sizes, lower precision, and delayed gradients.Other augmentations include evolution-strategy gradient estimates, gradient normalization, and directional gradient subspaces.Source: §C.5 · pp. 45–46

D.1 Meta-Objective

  • The meta-objective is expected training loss at the end of inner training, averaged over sampled task configurations, initializations, randomness, training lengths, and evaluation data.Inner-problem lengths are sampled from 200–20K steps early and 200–200K steps later in training.Source: §D.1 · p. 46

D.2 Meta-Gradient Estimation

  • Meta-gradients are estimated with antithetic evolution-strategy samples, sharing task initializations and data batches within pairs to reduce variance.The authors choose non-truncated ES for simplicity and lower communication overhead, and to target end-of-training rather than average loss.Source: §D.2 · p. 47

D.3 Curriculum and Meta-Generalization

  • Curricula over inner-problem length speed meta-training by using shorter runs before reaching the target length, trading reduced computation for bias that longer-run training can mitigate.The experiments target 2,000 steps; the final optimizer targets 200,000 iterations.Source: §D.3 · p. 47

D.4 Multi-Task Training

Multi-task meta-training normalizes task gradients to balance contributions, while the distributed system manages task reuse and gradient staleness. Interactive training changes, especially longer inner problems, improved performance on long-horizon tasks.

D.4 Multi-Task Training

  • Normalizing each task’s meta-gradient to constant norm makes training depend on task-sampling ratios rather than loss scale, but removes a well-defined meta-objective and can yield non-descent directions.Hand-designed loss normalizers performed worse, while a learned baseline-performance normalizer added complexity without outperforming gradient normalization.Source: §D.4 · pp. 47–48

E.1 One Learner, Many Workers

  • A centralized learner receives worker-computed meta-gradients, while workers run tasks on TPU chips; evaluation clusters continuously test checkpoints across many configurations.The evaluation setup reserves around 100 accelerators.Source: §E.1 · p. 48, §E.2 · p. 49

E.3 Task Selection and Staleness of Meta-Gradients

  • Workers reuse compiled static task configurations to avoid repeated multi-minute compilations, but this induces temporal gradient correlation; cycling through eight tasks per machine helps mitigate it.The learner discards gradients past a staleness threshold, with approximately 10 reported as a good balance.Source: §E.3 · pp. 49–50

E.4 Interactive Hyperparameter Tuning

  • The full-scale meta-training run was ad hoc and not easily replicated, and computational cost allowed only one model to be trained at full scale.Training proceeded in four phases using the previous phase’s weights as the starting point.Source: §E.4 · p. 50
  • Shifting the maximum inner-problem length from 20K to 200K slowed meta-training but dramatically improved performance on long-horizon problems.The task distribution also shifted toward larger problems in that phase.Source: §E.4 · p. 51
  • The authors used the 10th percentile of evaluation-task performance to guide development because broad performance across tasks mattered more than very fast training on a few tasks.They report that this criterion often generalized better to larger problems.Source: §E.4 · p. 51

E.5 Areas of Improvement

The appendix describes bottlenecks in meta-training infrastructure and details the evaluation normalization and task sets used to assess VeLO.

E.5 Areas of Improvement

  • TPU matrix-multiplication-unit utilization stayed below 10%, with memory throughput often limiting computations; GPU kernel-execution overhead was worse for these small, iterative operations.The authors cite hardware mismatch and note that custom GPU kernels with more fused operations might be more efficient.Source: §E.5 · p. 51
  • Phase 2 meta-training diverged in the lower percentiles, so the team rolled back and used a different configuration for phase 3.The figure distinguishes three training configurations by color and shows an exponential moving average in black.Source: Figure 22 · p. 52
  • Increasing the inner-training horizon from 20K to 200K iterations dramatically slowed the rate of inner problems trained because each task took longer.The figure also tracks meta-updates, inner gradients, and learned-optimizer applications over meta-training.Source: Figure 23 · p. 53
  • Cluster-state effects made training dynamics variable, appeared only after long meta-training, and were only partly mitigated by larger outer batches and lower staleness tolerance.A more synchronous infrastructure may help, but tasks vary greatly across machines.Source: §E.5 · p. 53

F.1 Evaluation Set Normalization

  • Evaluation normalization interpolates against the time-varying minimum loss across averaged baseline curves; the main-text baselines are learning-rate-tuned Adam run for 300K iterations.This setup measures speedups up to 30x when the target model runs for 10K iterations.Source: §F.1 · p. 54

F.2 Details on VeLOdrome Evaluation Set

  • VeLOdrome-83 contains 83 tasks from seven neural-network families, while VeLOdrome-109 adds learned-optimization and evolution-strategy tasks.The full set spans ten task families; VeLOdrome-83 excludes learned-optimization and ES-based tasks as uncommon training tasks.Source: §F.2 · pp. 54–55

F.3 MLCommons Tasks

The appendix reports VeLO results across MLCommons and other model families, alongside experiments probing learned step sizes and loss features.

F.3 MLCommons Tasks

  • VeLO beat all 20 Adam trials on ViT and WMT17, beat all Adam trials on DeepSpeech, matched Conformer, but underperformed on ResNet50 validation and ranked near the Adam median on molecular GNN.The speech tasks were not represented in meta-training; the authors hypothesize the molecular-GNN gap reflects its distance from the meta-training distribution.Source: §F.3 · pp. 55–56

F.4.1 NERF; F.4.2 MLP-Mixer; F.4.3 Object Detection

  • Additional evaluations used Blender-family NeRFs, two patch-16 MLP-Mixers, and Faster R-CNN on COCO, trained for 22,500 steps with batch size 64.The object-detection setup used SGD with momentum and a staged learning-rate schedule.Source: §F.4.1 · p. 56, §F.4.2 · p. 56, §F.4.3 · p. 56

F.4.4 Large Language Models

  • Pax’s scan-over-layers constraint required VeLO to optimize each layer independently, without signals from other layers beyond the global loss.The language-model experiments used decoder-only Transformers trained on C4.Source: §F.4.4 · p. 56

F.4.5 Graph NN

  • For graph networks, training-loss performance could match Adam while test performance lagged; clipping input gradient norms to at most 1e-2 fixed the generalization gap.The tested GNNs were out of distribution for meta-training and did not always show monotonic training-loss improvements corresponding to better generalization.Source: §F.4.5 · p. 57

F.5 Details on Optimizers Learn to Make Use of Finite Training Time

  • After 8,000 Adam steps, VeLO continuation immediately spiked, ultimately beat continued Adam, but remained worse than VeLO trained from random initialization for 12,000 steps.This demonstrates a lack of robustness to starting from a non-random state.Source: Figure 26 · p. 58

G.1 Learned Optimizers Take Smaller Steps for Larger Models

  • VeLO shrank its step size as Transformer size increased, and the gap between the smallest and largest step sizes widened midway through training.The authors describe the widening gap as unlike hand-designed optimizers they know.Source: §G.1 · p. 58, Figure 27 · p. 59

G.2 Learned Optimizers Leverage Loss Features

  • When given validation loss instead of training loss, VeLO reduced its step size as overfitting began and appeared to overfit less.The experiment used an ImageMLP task; validation loss was not part of meta-training.Source: §G.2 · pp. 58–59, Figure 28 · p. 59

G.3 Learned Optimizers Learn Implicit Step Size Schedule

  • VeLO’s step sizes rose rapidly, decayed, rose slowly through much of training, then decayed near the end, varying more across tensors and time than Adam or SGDM.The slow mid-training increase differs from the warmup-and-decay patterns the authors recognize from common practice.Source: §G.3 · p. 60, Figure 29 · p. 60

G.4 Sensitivity to Number of Training Iterations With Longer Horizons

These experiments examine VeLO’s behavior beyond its meta-training horizon, across model scales, and on MLCommons workloads and batch sizes.

G.4 Sensitivity to Number of Training Iterations With Longer Horizons

  • Beyond its meta-training horizon of at most 20K steps, VeLO generalized well on some tasks but showed unstable, non-monotonic curves on others; gradient accumulation or longer-horizon fine-tuning improved stability.Accumulating ten iterations of gradients keeps the horizon the optimizer sees within 10K steps; fine-tuning on 200K-iteration problems also fixes the instability but is expensive.Source: §G.4 · pp. 60–61, Figure 30 · p. 61

G.5 Effects of Scaling Model Size

  • On 10M-parameter Transformers, VeLO optimized better than a 100M-parameter Transformer with tuned Adam; on larger auto-encoder and image-classification models, tuned Adam matched or outperformed VeLO.The authors suspect the largest models are increasingly far from the meta-training distribution.Source: §G.5 · pp. 61–62, Figure 31 · p. 62

G.6 Exploration Into Outputs of Per-Tensor LSTM

G.7 MLCommons: Validation Performance and Hyperparameter Sensitivity

  • Across six MLCommons workloads, validation performance beat baselines on ViT and DeepSpeech, matched WMT17 and Conformer, and lagged on ResNet50; training loss beat 20 Adam trials except on GNN and one ResNet50 task.The authors hypothesize that limited data augmentation contributed to the ResNet50 validation result.Source: §G.7 · p. 62, Figure 34 · p. 65

G.8 MLCommons: Larger Batch Sizes

  • At 8x batch size VeLO outperformed the Adam baseline on all tested tasks except Graph NN; at 32x, both optimizers performed poorly, with VeLO winning on three of six workloads.The 32x comparison favored VeLO on DeepSpeech, Conformer, and ResNet50, while Adam did better on ViT, WMT17, and Graph NN.Source: §G.8 · p. 65

G.9 Additional Tasks for Batch Size Scaling

The appendix extends batch-size comparisons, tests VeLO on a 300M-parameter language model, and reports single-task fine-tuning results.

G.9 Additional Tasks for Batch Size Scaling

  • Across additional models, VeLO matched or outperformed baselines at small-to-medium batch sizes, while learning-rate-tuned Adam often performed best at extremely large batch sizes.The authors speculate the large-batch gap reflects distribution shift between meta-training and evaluation.Source: §G.9 · p. 65, Figure 37 · p. 67

G.10 300M Language Model Trained on 64 Accelerators

  • On a 300M-parameter decoder-only Transformer, VeLO performed comparably to the best learning-rate-tuned AdaFactor trial despite optimizing each layer independently.The model was trained on a 4x4x4 TPUv4 slice; the independent-layer restriction removed communication across layers.Source: §G.10 · p. 66, Figure 38 · p. 68

G.11 Finetuning of the Learned Optimizer on a Single Problem

  • Fine-tuning raised VeLO’s speedup over Adam from 3.2x to 4.2x, and improved performance by about 10% on a model with 4x larger hidden size.The authors suggest fine-tuning on a similar, smaller task could be considerably cheaper.Source: §G.11 · p. 69

G.12 Extended Evaluation on Optimizer Benchmark

Across an extended 109-task benchmark, VeLO remains strong against hand-designed and learned optimizers, though some task types are difficult. Experiments using learned optimizers to train other learned optimizers show mixed results, with outcomes varying by task and optimizer.

G.12 Extended Evaluation on Optimizer Benchmark

  • On the 109-task set, VeLO remains strong overall but struggles on evolution-strategy and learned-optimizer-training tasks.Figures 41 and 42 compare speedups over Adam and an extended set of baselines.Source: §G.12 · p. 69, Figure 41 · p. 70, Figure 42 · p. 70
  • VeLO averages more than 7x faster than learning rate-tuned Adam, while the best learned-optimizer ensemble exceeds 11x; extensive multi-trial tuning remains a higher bar on average.The tables report single- and multi-trial optimizers normalized to learning rate-tuned Adam.Source: §G.12 · p. 71, Table 1 · p. 72, Table 2 · p. 73
  • The benchmark averages a task distribution that could favor different optimizers, and selecting the best stochastic configurations may yield better performance than is attainable.The authors average results over five seeds per hyperparameter configuration and use many batches for final losses to mitigate this effect.Source: §G.12 · pp. 75–76

H Learned Optimizers Training Other Learned Optimizers

  • Pretrained learned optimizers were tested on other learned optimizers trained for single tasks, including MLPs on CIFAR10 or FashionMNIST.The experiments include gradients estimated by evolution strategies and gradients computed with backpropagation.Source: §H.1 · p. 76

H.1 Learned Optimizers Which Train a Single Problem

  • When training learned optimizers for single tasks, the method matched or beat Adam in two of three evolution-strategy setups but diverged in the third; backpropagation results also matched or beat Adam, with unstable curves.The evolution-strategy runs did not surpass extensive NAdamW tuning.Source: §H.1 · p. 76

H.2 Multi-Task, Small Variation Across Tasks

Small-task experiments expose different outcomes across learned-optimizer architectures and training scales. The paper also documents practical constraints, released resources, and open directions for optimizer design and interpretation.

H.2 Multi-Task, Small Variation Across Tasks

  • On a diverse small MLP task mixture, the per-parameter MLP matched tuned Adam, while the hierarchical hypernetwork remained behind well-tuned Adam despite improved stability with a fourfold larger batch.The small-batch hypernetwork run became unstable and diverged; the tasks used resized, black-and-white CIFAR10, FashionMNIST, MNIST, and SVHN images.Source: §H.2 · p. 77

H.3 Large scale: Self-Hosting Optimizer, Replicating Work with a Learned Optimizer

  • At large scale, VeLO’s hypernetwork diverged, while the Metz et al. [2020a] hierarchical optimizer kept improving, a promising but preliminary result.The authors note that they could not conduct careful studies of the computationally expensive setting.Source: §H.3 · p. 77

I.1 Considerations When Trying VeLO

  • VeLO can underperform on out-of-distribution tasks, be unreliable beyond 200K iterations, require about 2x Adam memory, and compile up to 4x longer.Gradient accumulation is the documented workaround for longer horizons; the authors also describe distributed-training techniques that can improve optimizer performance.Source: §I.1 · p. 78

I.2 Released Optimizers and Learning Curves

  • The release includes pretrained optimizer weights, a guard-railed optimizer with optional weight decay, and over a million training curves across more than 100 tasks.The guard rails use gradient accumulation to keep the learned optimizer under 150K steps.Source: §I.2 · p. 78

J Open Questions for Future Work

  • The authors identify further research opportunities despite the compute demands of large-scale meta-training.They point to multiple areas for follow-up work.Source: J Open Questions for Future Work · p. 78

J.1 Learned Optimizer Architectures

  • Future architectures could use more inner-training signals, approximate second-order features, or configurable objectives beyond a single user-specified training horizon.The current optimizer takes the mean gradient; the authors report difficulty meta-training models with broader test-time configuration.Source: §J.1 · p. 79

J.2 Learned Optimizer Introspection

  • The learned optimizers perform well, but their mechanisms remain difficult to understand, limiting efforts to study them theoretically or distill them into simpler optimizers.Source: §J.2 · p. 79

J.3 Meta-Training

The discussion identifies compute, task-distribution design, and gradient staleness as constraints on meta-training, then examines learned optimizers training other optimizers. These experiments show mixed results, alongside limitations in specialization, long-horizon extrapolation, and worst-case performance.

J.3 Meta-Training

  • Broad meta-training requires thousands of accelerators for over a month, while task distributions are heuristic and may include unrepresentative models or omit important areas.Small-scale training can run on 4 TPU in a couple hours; task-distribution design also controls meta-training cost.Source: §J.3 · p. 79
  • A four-order-of-magnitude spread in task runtime makes slower tasks use older meta-training weights; the impact is uncertain, and large outer-gradient batches are used to mitigate staleness.Source: §J.3 · p. 80

J.4 Usability of Learned Optimizers

  • Although hyperparameter-free optimization makes experimentation cheaper, fine-tuning for specialization can be computationally costly with naive evolution strategies.Source: §J.4 · p. 80
  • VeLO performs well on around 90–95% of tested tasks, but performance can vary or diverge farther from its meta-training distribution, and long-horizon extrapolation is poor.The authors identify better worst-case performance and longer-horizon generalization as areas needing more work.Source: §J.4 · p. 80
  • With Hyper LOpt, an outer batch size of 128 improves stability, but the learned optimizer remains close to Adam’s best learning rate without outperforming Adam.Source: Figure 50 · p. 86
  • When retraining itself, the pre-trained optimizer matches learning-rate-tuned Adam on more than 50% of tasks after modest training but does not surpass it, and meta-loss diverges after 500 inner iterations.Training took 1.2 days.Source: Figure 51 · p. 86
  • Retraining a hierarchical RNN LOpt takes 3 days and matches learning-rate-tuned Adam on more than 50% of tasks, but does not achieve better performance.Source: Figure 52 · p. 87
Loading 2211.09760v1…