A Trainable Optimizer

TL;DR

This paper introduces a trainable optimizer framework that learns to optimize, matching SGD's convergence rate with reduced variance and demonstrating faster convergence than benchmarks like ADAM in various settings.

Contribution

The paper presents a novel trainable optimizer framework with a pseudo-linear approximation that achieves efficient convergence and introduces simplified variants for better computational efficiency.

Findings

TO methods converge faster than ADAM in experiments

Pseudo-linear TO matches SGD's convergence rate

Minimal additional computational overhead

Abstract

The concept of learning to optimize involves utilizing a trainable optimization strategy rather than relying on manually defined full gradient estimations such as ADAM. We present a framework that jointly trains the full gradient estimator and the trainable weights of the model. Specifically, we prove that pseudo-linear TO (Trainable Optimizer), a linear approximation of the full gradient, matches SGD's convergence rate while effectively reducing variance. Pseudo-linear TO incurs negligible computational overhead, requiring only minimal additional tensor multiplications. To further improve computational efficiency, we introduce two simplified variants of Pseudo-linear TO. Experiments demonstrate that TO methods converge faster than benchmark algorithms (e.g., ADAM) in both strongly convex and non-convex settings, and fine tuning of an LLM.