PaperlayerView exact revision on arXiv

Source-linked AI summary

Lightweight and Versatile Learned Optimization by Recombination of Gradient History

Minyoung Choi, Dalta Imam Maulana, Wanyeong Jung

arXiv:2610.09604v1cs.LG

TL;DR

Coordinate-wise learned optimizers can be expensive, motivating lightweight and versatile designs. This paper learns to recombine averages of gradient history into updates, and reports zero-shot generalization to unseen language, vision, and graph benchmarks with improved validation loss and test accuracy over Adam.

  • Problem

    Coordinate-wise learned optimizers are expensive because they process gradient-derived features for individual parameters, making lightweight and versatile architectures challenging.

    Source: §1 · p. 1

  • Method

    The optimizer dynamically recombines averages of gradient history, with coefficients shared across parameter groups and progressive averaging of older gradients.

    Source: §1 · p. 1, §2.1 · p. 2, §2.2 · p. 3

  • Results

    A single trained policy generalizes zero-shot to unseen language, vision, and graph benchmarks, improving validation loss and test accuracy over Adam.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Results establish recombination of gradient history as a practical design principle for learned optimization.

    Source: §4 · p. 8

Abstract

from arXiv · show

This paper presents a lightweight and versatile learned optimizer that dynamically recombines gradient history, represented as averages over disjoint time spans. The optimizer reduces the prediction space to one scalar coefficient per gradient average, shared by multiple parameters. Progressively averaging older gradients minimizes memory cost of long history, while keeping their contributions independently accessible. A 37k-parameter network trained in 0.87 GPU-hours generalizes zero-shot to unseen tasks, lowering validation loss by 9.1% and 0.4% on BERT-Tiny and GPT-Tiny, and improving test accuracy over Adam by 3.5 %p on a Vision Transformer and by 2.7 %p on average across nine graph models, with FLOPs overhead as low as 0.3%.

1 INTRODUCTION

Learned optimizers can adapt to loss landscapes, but coordinate-wise methods are expensive; this paper learns to recombine gradient-history averages as updates. Its compact representation supports configurable parameter granularity and logarithmic memory growth with history length.

1 INTRODUCTION

  • Coordinate-wise learned optimizers process gradient-derived features for individual parameters; VeLO training requires approximately 4,000 TPU-months.The paper frames lightweight, versatile optimizer design as a challenge.Source: §1 · p. 1
  • Compact group statistics allow configurable parameter granularity and keep older, disjoint time spans independently accessible.A fixed-dimensional representation supports versatility across target networks.Source: §1 · p. 2
  • Progressively averaging older gradients limits memory growth logarithmically with history length.This is one of the techniques proposed to make the optimizer practical.Source: §1 · p. 2
  • The optimizer dynamically recombines averages of past gradients, jointly weighting history directions to construct each update.Gradient history provides a directional basis, as in momentum and Adam.Source: §1 · pp. 1–2
  • A single trained policy generalizes zero-shot to unseen language, vision, and graph benchmarks, improving validation loss and test accuracy over Adam.The paper presents this as its core empirical contribution.Source: §1 · p. 2
  • A 37k-parameter policy is trained in under one GPU-hour on an NVIDIA B200.The introduction highlights the policy's compact size and training cost.Source: §1 · p. 2

2 LEARNED OPTIMIZATION BY DYNAMIC RECOMBINATION OF GRADIENT HISTORY

The method updates parameter groups by recombining disjoint-span gradient averages, with coordinates in each group sharing temporal coefficients. Progressive averaging extends the history, while encoded slot features provide the basis for predicting coefficients.

2.1 DYNAMIC RECOMBINATION OF GRADIENT HISTORY

  • Each parameter group updates by recombining its populated disjoint-span gradient averages, with coordinates in the group sharing a temporal coefficient vector.The group update uses the coefficients associated with its populated history slots.Source: §2.1 · p. 2

2.2 MULTI-RESOLUTION GRADIENT HISTORY

  • Each slot is assigned a tier and can average up to r^(c−1) gradients; as averages age, they carry into subsequent tiers to extend history with few slots.The example uses r = 2, with carry propagating through full slots and the oldest average eventually evicted beyond the final tier.Source: §2.2 · p. 3, Figure 3 · p. 3

2.3 TEMPORAL COEFFICIENT PREDICTION

  • A group encoder fuses history and parameter statistics into fixed-dimensional slot features; self-attention with tier context predicts normalized allocations, then a group scale sets final coefficients.The feature dimension is independent of the number of parameters in the group.Source: §2.3 · p. 3, Figure 4 · p. 4

2.4 EXTENSION ACROSS PARAMETER PARTITIONS

  • For overlapping parameter groups, a separate network processes each partition’s temporally recombined output, then spatial coefficients fuse the outputs into the final parameter step.The evaluation focuses on a single partition; the extension is defined and evaluated in Appendix C.Source: §2.4 · p. 4

2.5 TRAJECTORY-LEVEL OPTIMIZER TRAINING

  • Training minimizes cumulative held-out-loss change over an unrolled trajectory of K updates, with loss scales normalized so meta-tasks contribute equally.Source: §2.5 · p. 4

3 EVALUATION

The evaluation tests zero-shot deployment across language, vision, and graph tasks, then examines history-span ablations, temporal behavior, and deployment cost.

3.1 EXPERIMENTAL SETUP

  • The learned optimizer is meta-trained on four small networks and evaluated zero-shot on unseen BERT, GPT, ViT, and GNN architectures, with paired runs sharing initialization and data ordering.Source: §3.1 · p. 4

3.2 ZERO-SHOT GENERALIZATION TO UNSEEN TASKS

  • Across unseen tasks, the optimizer reaches the lowest terminal validation loss on language models, leads seven of nine graph comparisons and their mean, and achieves the highest ViT accuracy.Source: Table 1 · p. 5, §3.2 · p. 6

3.3 ABLATIONS AND BEHAVIORAL ANALYSIS

  • A 63-step history maintains convergence across tasks; full-span history gives lower language validation loss, while both spans beat Adam on vision and yield near-identical graph mean accuracy.Source: §3.3 · pp. 6–7, Table 2 · p. 7
  • After sharing a trajectory through step 32, Adam's branch has terminal training loss 0.30 at step 500 versus 0.25 for learned recombination; learned temporal scores vary across updates 24–55.Source: Figure 8 · p. 7, §3.3 · p. 7

3.4 COMPUTATIONAL COST AND DEPLOYMENT SCALING

  • Training costs 0.87 B200 GPU-hours; all-parameter deployment uses 1.035×/1.166× Adam FLOPs per BERT/GPT step, while 131,071-gradient history occupies 301/466 MB.Source: §3.4 · pp. 7–8, Figure 9 · p. 8

4 CONCLUSION

  • The results establish gradient-history recombination as a practical learned-optimization design principle, despite meta-training only on small classification tasks.Source: §4 · p. 8

A IMPLEMENTATION AND EVALUATION DETAILS

The appendix outlines the evaluation materials and describes the grouping and history-slot configuration used by the optimizer.

A IMPLEMENTATION AND EVALUATION DETAILS

A.1 LEARNED-OPTIMIZER CONFIGURATION

  • Target tensors are partitioned into fixed groups of at most 256 active coordinates; the evaluated configuration uses six history slots with capacities 1, 2, 4, 8, 16, 32 and r = 2.Source: §A.1 · p. 10

A.1.1 HISTORY-SLOT FEATURE DEFINITIONS

The appendix describes per-slot gradient features, their temporal predictor, and the multi-resolution procedure for updating stored history.

A.1.1 HISTORY-SLOT FEATURE DEFINITIONS

  • Each populated slot has six statistics, and a histogram path normalizes slot gradients by gradient RMS before forming coordinate-level feature tuples.Tuple channels are standardized within each slot and clipped to [−5, 5].Source: Table 3 · p. 10, §A.1.1 · p. 10
  • Separate MLPs project slot statistics and histograms before gated fusion; a two-layer Transformer predicts coefficients normalized over populated slots, with a learned group scale multiplying the recombination.Source: §A.1.1 · p. 10

A.1.2 MULTI-RESOLUTION HISTORY UPDATE

  • At each step, the current gradient enters the newest slot and older values are carried forward, replaced, or averaged into higher tiers; values carried beyond the final tier are discarded.Source: §A.1.1 · p. 11, §A.1.2 · p. 11

A.2 OPTIMIZER-TRAINING TASKS AND PROTOCOL

Optimizer training uses image and synthetic classification tasks, with a held-out-loss objective evaluated over unrolled updates.

A.2 OPTIMIZER-TRAINING TASKS AND PROTOCOL

  • Training tasks include normalized CIFAR-10 and Fashion-MNIST without augmentation, plus 20,000 standardized 64-dimensional synthetic examples; each dataset is split 80/20 for model training and held-out meta-objective evaluation.Source: §A.2 · p. 11

A.2.1 NORMALIZED CUMULATIVE HELD-OUT-LOSS OBJECTIVE

  • The meta-objective minimizes cumulative held-out-loss change across K unrolled updates, normalizing loss scales so each meta-task contributes equally.Source: §A.2.1 · p. 11

A.3 TARGET-MODEL TRAINING AND EVALUATION

The evaluation fixes baseline settings across tasks and uses paired runs to compare deployment on language and vision targets. The protocol specifies the target architectures, parameter assignments, and training and validation procedures.

A.3.1 SHARED BASELINE HYPERPARAMETERS

  • Across all evaluation domains, baseline optimizers use fixed hyperparameters and zero weight decay, with no domain-specific sweeps because exhaustive tuning is computationally prohibitive for deployment-performance evaluation.Adam uses a learning rate of 10^-3; SGD and SGDM use 0.01.Source: §A.3.1 · p. 12

A.3.2 LANGUAGE MODELS

  • Both language targets use two layers, width 128, and two heads; learned recombination updates all parameters except embeddings and the output head.Those excluded parameters use Adam, and paired runs share initialization, data, stochastic streams, and a 20-microbatch validation panel.Source: §A.3.2 · p. 12, Table 6 · p. 12

A.3.3 VISION-TRANSFORMER PROTOCOL

  • The CIFAR-100 Vision Transformer uses six blocks, width 256, eight heads, 8×8 patches, and dropout 0.1, with the learned optimizer updating all parameters.Training uses 45,000 examples for 50 epochs; the terminal state is evaluated once on the test set.Source: §A.3.3 · p. 12

A.3.4 GRAPH MODELS

The appendix specifies graph-model datasets and training budgets, deployment-cost measurements, and language-model parameter-assignment comparisons. The supplied passages describe the protocols and comparisons but do not report their outcome values.

A.3.4 GRAPH MODELS

  • The graph evaluation crosses GCN, GraphSAGE, and GAT with Cora, CiteSeer, and PubMed, using 100 updates for the first two datasets and 500 for PubMed.The public Planetoid split is used; training, validation, and test masks are kept distinct.Source: §A.3.4 · pp. 12–13

A.4 RESOURCE MEASUREMENTS

  • Figures 10 and 11 report FLOPs and time per step as ratios to Adam, respectively by history span and language-model parameter assignment.Deployment profiles exclude data loading, validation, logging, and checkpointing; language-model timing replaces only the Adam update interval.Source: Figure 10 · p. 13, Figure 11 · p. 13, §A.4 · pp. 13–14

B.1.1 LANGUAGE-MODEL PARAMETER ASSIGNMENTS

  • Table 7 and Figure 12 compare all-tensor learned recombination with complementary hybrid parameter assignments, using the same trained policy rather than separately training every assignment.The hybrid arms use Adam for tensors not routed to learned recombination.Source: §B.1.1 · p. 14, Table 7 · p. 14, Figure 12 · p. 14

B.1.2 GRAPH TRAINING TRAJECTORIES

The section compares graph-optimizer trajectories and examines grouping, history depth, and temporal behavior. It also introduces a spatio-temporal extension that combines candidate updates formed from multiple tensor partitions.

B.1.2 GRAPH TRAINING TRAJECTORIES

  • Figure 13 compares full-span learned optimization with Adam at equal cumulative FLOPs, while Figure 14 shows validation trajectories used for checkpoint selection without test accuracy.The figures show three-seed means and seedwise ranges; Figure 14 marks each seed's minimum validation loss.Source: §B.1.2 · p. 14, Figure 13 · p. 15, Figure 14 · p. 16

B.2.1 PARAMETER GROUPING

  • All tested grouping trajectories remain finite, but effects vary by task and metric, and the results do not identify a preferred grouping.Matched alternative-grouping policies use one optimizer-training seed.Source: §B.2.1 · p. 17, Table 8 · p. 17

B.2.2 HISTORY-DEPTH CONTROLS

  • Across nine graph-task cells, aggregate accuracy is similar for 63-step and full-span histories, while the preferred span varies; history depth is fixed before validation and test evaluation.The comparison uses three paired target seeds.Source: §B.2.2 · p. 17, Table 9 · p. 17

B.2.3 TEMPORAL COEFFICIENTS AND STORED-HISTORY CONTENT

  • On GAT/PubMed, temporal contributions, slot magnitudes, and signed contribution fractions show distinct patterns for source-attention, dense-weight, and output-bias groups.Figure 15 follows the full trajectory and distinguishes groups of 64, 256, and 3 coordinates.Source: §B.2.3 · p. 18, Figure 15 · p. 18

C SPATIO-TEMPORAL EXTENSION

  • The spatio-temporal extension forms candidate updates from multiple tensor partitions and uses a second predictor to combine them into the final update.This extends temporal recombination across spatial tiers.Source: C SPATIO-TEMPORAL EXTENSION · p. 18

C.1 PYRAMID GROUPING

Pyramid grouping applies temporal recombination to multiple partitions of a tensor, then combines the partition-specific candidate updates. The evaluated construction matches the representative row-grouped policy on nine unseen graph tasks, while the comparison does not isolate a grouping advantage.

C.1 PYRAMID GROUPING

  • A shared temporal predictor recombines gradient history within each retained partition, and a separate spatial predictor uses Softmax weights to combine their candidate updates.The illustrated hierarchy organizes configurations by total subdivision depth.Source: §C.1 · pp. 18–19
  • 76.51% mean test accuracy for pyramid grouping compares with 76.44% for row grouping and 73.78% for Adam across nine unseen graph tasks.The two learned policies were trained independently, so the comparison shows both constructions function under the protocol, not that one grouping causes a gain.Source: §C.1 · p. 19, Table 10 · p. 19
  • The evaluated construction retains whole-tensor groups, 16 contiguous first-axis groups, and 16 contiguous groups along the flattened trailing axis; groups under 64 coordinates are omitted.For tensors of rank at least two, the representation uses the first axis and the flattened product of remaining axes.Source: §C.1 · pp. 18–19
Loading 2610.09604v1…