[Paper Review] Neural Optimizer Search with Reinforcement Learning

Irwan Bello, Barret Zoph+2 more|arXiv (Cornell University)|Sep 21, 2017

Advanced Neural Network Applications33 references201 citations

TL;DR

The paper automates discovery of neural network optimizers by training an RNN controller to generate update equations from a DSL, optimizing them with PPO-based reinforcement learning, and shows transferability across tasks.

ABSTRACT

We present an approach to automate the process of discovering optimization methods, with a focus on deep learning architectures. We train a Recurrent Neural Network controller to generate a string in a domain specific language that describes a mathematical update equation based on a list of primitive functions, such as the gradient, running average of the gradient, etc. The controller is trained with Reinforcement Learning to maximize the performance of a model after a few epochs. On CIFAR-10, our method discovers several update rules that are better than many commonly used optimizers, such as Adam, RMSProp, or SGD with and without Momentum on a ConvNet model. We introduce two new optimizers, named PowerSign and AddSign, which we show transfer well and improve training on a variety of different tasks and architectures, including ImageNet classification and Google's neural machine translation system.

Motivation & Objective

• Motivate and automate the design of optimization update rules for deep learning.
• Represent update rules in a domain-specific language to enable flexible composition.
• Use reinforcement learning to optimize update rules based on validation performance.
• Demonstrate transferability of discovered optimizers across architectures and tasks.
• Provide faster, more memory-efficient alternatives to traditional optimizers.

Proposed method

• Train a recurrent neural network controller to output strings in a DSL that describe an optimizer update rule.
• Use PPO (proximal policy optimization) to maximize validation accuracy achieved by a target model trained with the sampled rule.
• Construct a domain-specific language that encodes operands, unary functions, and binary functions to form update equations.
• Apply a distributed training setup to accelerate search, using a small ConvNet for evaluation and five-epoch training as the signal.
• Identify and analyze resulting update rules such as PowerSign and AddSign and a learning-rate decay scheme called linear cosine decay.

Experimental results

Research questions

• RQ1 Can an RL-based controller automatically discover effective optimizer update rules for neural networks?
• RQ2 Do the discovered update rules transfer effectively to larger models and different tasks?
• RQ3 What new update rules and learning-rate schedules emerge from the search, and how do they compare to SGD, Momentum, RMSProp, and Adam?

Key findings

• The controller discovers update rules that can outperform Adam, RMSProp, and SGD with/without Momentum on a small ConvNet at CIFAR-10.
• Two primary families, PowerSign and AddSign, emerge as effective update rules and transfer to larger tasks like Wide ResNet on CIFAR-10, ImageNet, GNMT, and PTB language modeling with varying gains.
• A linear cosine decay (and its variants) often enables faster convergence and allows larger initial learning rates.
• Discovered rules can be more memory-efficient than Adam since PowerSign uses one running average per parameter.
• The approach yields measurable improvements on ImageNet top-1/top-5 and GNMT BLEU when replacing standard optimizers in those pipelines.
• The learned rules remain robust across modest hyperparameter variations and can interpolate toward SGD under certain settings.