A Trainable Optimizer

2025/08/03 by Wang, Ruiqi, Klabjan, Diego
#FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · doi:10.48550/arxiv.2508.01764

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.

Citations

• Divergence Results and Convergence of a Variance Reduced Version of ADAM
• SUPER-ADAM: Faster and Universal Framework of Adaptive Gradients
• Learning to Optimize: A Primer and A Benchmark
• The Impact of the Mini-batch Size on the Variance of Gradients in Stochastic Gradient Descent
• Momentum-Based Variance Reduction in Non-Convex SGD
• On the Convergence of Adam and Beyond
• Neural Optimizer Search with Reinforcement Learning
• Learning to learn by gradient descent by gradient descent
• Learning to Optimize
• Adam: A Method for Stochastic Optimization
• SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly\n Convex Composite Objectives
• The MNIST Database of Handwritten Digit Images for Machine Learning Research [Best of the Web]

Related

• Greedy Alignment Principle for Optimizer Selection
• LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics
• Narrowing the Focus: Learned Optimizers for Pretrained Models
• LLM as a Complementary Optimizer to Gradient Descent: A Case Study in Prompt Tuning
• Reparametrizing gradient descent