[Paper Review] Learning to Optimize Neural Nets

Ke Li, Jitendra Malik|arXiv (Cornell University)|Mar 1, 2017

Machine Learning and Algorithms29 references84 citations

TL;DR

The paper extends a framework to learn high-dimensional, stochastic optimizers via guided policy search and demonstrates a learned optimizer that outperforms hand-engineered algorithms and generalizes across datasets and architectures.

ABSTRACT

Learning to Optimize is a recently proposed framework for learning optimization algorithms using reinforcement learning. In this paper, we explore learning an optimization algorithm for training shallow neural nets. Such high-dimensional stochastic optimization problems present interesting challenges for existing reinforcement learning algorithms. We develop an extension that is suited to learning optimization algorithms in this setting and demonstrate that the learned optimization algorithm consistently outperforms other known optimization algorithms even on unseen tasks and is robust to changes in stochasticity of gradients and the neural net architecture. More specifically, we show that an optimization algorithm trained with the proposed method on the problem of training a neural net on MNIST generalizes to the problems of training neural nets on the Toronto Faces Dataset, CIFAR-10 and CIFAR-100.

Motivation & Objective

• Motivate learning optimization algorithms as a data-driven tool for neural network training rather than hand-crafted rules.
• Develop an extension of Learning to Optimize suitable for high-dimensional stochastic problems.
• Show that the learned optimizer can generalize to unseen architectures and datasets while remaining robust to gradient stochasticity.

Proposed method

• Formulate optimization as a reinforcement learning problem where the update rule is the policy.
• Use guided policy search (GPS) to learn a time-varying linear policy plus a non-linear stationry policy represented by an RNN. Incorporate a local linearization of dynamics and a quadratic approximation of the cost to solve subproblems with an LQG regulator.
• Impose permutation-invariance (coordinate groups) to exploit neural net structure and enable scalable learning in high dimensions.
• Define state features ocusing on histories of iterates, gradients, and objectives; define observation features for training and testing.
• Apply a convolutional GPS (block-diagonal, coordinate-wise decompositions) to scale learning to neural-net optimization.

Experimental results

Research questions

• RQ1 Can a learned optimization algorithm trained on a shallow neural network generalize to deeper architectures and different datasets?
• RQ2 Is the learned optimizer robust to changes in gradient stochasticity and network architecture?
• RQ3 Does the learned optimizer outperform popular hand-engineered optimizers (e.g., SGD, momentum, Adagrad, RMSProp, Adam) on unseen tasks?
• RQ4 To what extent can structural priors (per-coordinate groups) improve learning efficiency for high-dimensional optimization?
• RQ5 How well does the learned optimizer transfer across datasets with very different statistics (MNIST vs. CIFAR/TFD)?

Key findings

• The learned optimizer consistently descends to the optimum faster than hand-engineered methods across MNIST-derived, TFD, CIFAR-10, and CIFAR-100 tasks.
• Predicted Step Descent generalizes to neural nets with larger input and hidden layers (increasing parameters by ~8x) and remains competitive or superior despite initial oscillations.
• The learned optimizer remains robust under reduced batch sizes (increased gradient stochasticity) and still outperforms alternatives on several datasets.
• An optimizer meta-trained on MNIST can outperform on dissimilar datasets, indicating transferability of the learned optimization strategy.
• Compared to an L2L-based learned optimizer, the proposed method shows superior consistency and stability on unseen tasks and architectures.