Neural Optimizer Search with Reinforcement Learning

Irwan Bello; Barret Zoph; Vijay Vasudevan; Quoc V. Le

Source-linked AI summary

Neural Optimizer Search with Reinforcement Learning

Irwan Bello, Barret Zoph, Vijay Vasudevan, Quoc V. Le

arXiv:1709.07417v2Sep 21, 2017cs.AIcs.LGstat.ML

TL;DR

Neural Optimizer Search addresses the difficult problem of designing update rules for non-convex deep-learning training. It uses a recurrent controller trained with reinforcement learning to generate mathematical equations, discovering transferable optimizers that improve results across architectures and tasks.

Problem
Designing effective optimization methods for non-convex deep-learning problems is challenging, although optimizer choice strongly affects training success.
Method
A recurrent controller generates domain-specific-language strings describing mathematical update equations and is trained with reinforcement learning using validation performance.
Results
PowerSign and AddSign achieve competitive performance across image-classification and machine-translation tasks, with improvements up to 0.4% on ImageNet and 0.7 BLEU on WMT 2014 English-to-German.
Takeaways & Limitations
The discovered equation-based update rules can be reused on other optimization tasks, and linear cosine decay generally converges faster than cosine annealing.
Takeaways & Limitations
The DSL can represent only a subset of mathematical equations when the number of iterations is limited.

Abstract

from arXiv · show

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.

1. Introduction

Neural Optimizer Search automates update-rule design with a recurrent controller trained by reinforcement learning. On CIFAR-10 and broader tasks, discovered rules outperform or improve training over common optimizers.

Designing optimizers is difficult because neural-network training is non-convex, despite the major role optimization methods play in model training.
The method uses a recurrent controller to generate update equations and reinforcement learning to maximize validation accuracy after fixed training.The generated rule is evaluated on a particular model architecture for a fixed number of epochs.
On CIFAR-10, discovered update rules outperform Adam, RMSProp, and SGD with or without Momentum on a small ConvNet.
ImageNet top-1 and top-5 accuracy improves by up to 0.4% on a mobile-sized model, while WMT 2014 English-to-German improves by up to 0.7 BLEU.

2. Related Work

Prior work developed many hand-designed, learned, and evolved optimization methods, but Neural Optimizer Search instead searches directly for mathematical update equations using flexible primitive combinations.

Modern optimizers combine minibatch stochastic updates with heuristics estimating diagonal second-order information to accelerate practical convergence.Adam, for example, estimates the gradient mean and variance.
Because neural-network objectives are non-convex, empirical learning-rate heuristics indicate that useful optimization ideas remain undiscovered.
Unlike related methods that generate numerical updates, this work searches among primitive-based mathematical equations for update rules.
Genetic-programming approaches can be slow and heuristic-dependent, leading prior studies to test very small networks, including models with about 100 weights.
The approach adapts reinforcement-learning-based automated model discovery by using flexible primitive combinations to search for optimizers.
Validation accuracy serves as the reward, implicitly searching for optimizers that may also support generalization.

3. Method

The method represents optimizer equations in a domain-specific language and uses a recurrent controller with policy-gradient reinforcement learning to search and evaluate them.

Search framework: Controller-generated strings are mapped to update rules, applied to neural networks, and scored to improve later controller-generated rules.
Domain-specific language: The computation graph expresses common optimizers using primitive inputs such as gradients and running estimates, alongside unary and binary functions.
Domain-specific language: The DSL represents each rule through two operands, two unary functions, and one binary function, whose output is stored or used as the weight update.
Domain-specific language: A limited number of DSL iterations restricts the representable equations, although common optimizers can fit within one iteration given suitable primitives.
Controller and training: The recurrent controller samples strings of length 5n, with n fixed during training, while different softmax weights handle the growing operand bank.
Controller and training: The controller’s reward is target-network validation accuracy after training with a sampled update rule.
Controller and training: Proximal Policy Optimization improves controller convergence over vanilla REINFORCE, and distributed workers further accelerate search.
Controller and training: Searching with a small two-layer convolutional network reduces computation while providing a signal for performance on larger models.

4. Experiments

The experiments define a broad optimizer search space from primitive operands and functions, then train and evaluate sampled update rules on CIFAR-10 ConvNets using distributed workers and held-out validation rewards.

4.1. Search space: The controller combines gradient-related operands, unary functions, binary functions, and training-step decay operands to construct update rules.Available primitives include gradients, moving averages, clipping, dropout, arithmetic operations, restart decay, and annealed noise.
4.1. Search space: 106 to 1011 possible update rules are represented by binary trees with depths of 1 to 4 and strings of lengths 5 to 20.Search-space size depends on the selected subsets of operands and functions.
4.2. Experimental details: The controller is a single-layer LSTM trained with Adam at learning rate 10^-5, minibatch size 5, and an entropy penalty of 0.0015.Its hidden state has size 150, and weights are initialized uniformly between -0.08 and 0.08.
4.2. Experimental details: Sampled optimizers are first tuned over learning rates 10^i for i from -5 to 1, then trained for 5 epochs on CIFAR-10 with validation accuracy returned as reward.Each child network uses a fixed held-out validation set of 5,000 examples, and experiments use 100 distributed CPU workers.
4.2. Experimental details: The controller reward increases as more optimizers are sampled.Figure 4 presents the controller learning curve for this sampling process.

5. Results

The search discovers update rules organized around agreement between a gradient and its moving average, including PowerSign and AddSign, and also produces learning-rate decays that improve convergence relative to cosine decay.

5.1. Discovered optimizers: The controller repeatedly samples sign(g) ∗ sign(m), a binary signal indicating whether the gradient and its moving average agree on one dimension.This quantity appears as a subcomponent of many sampled optimizers.
5.1. Discovered optimizers: PowerSign scales each parameter update by αf(t) or 1/αf(t) according to gradient–moving-average agreement.The discovered family has the form αf(t)∗sign(g)∗sign(m) ∗g.
5.1. Discovered optimizers: AddSign scales each parameter update by α+f(t) or α−f(t) according to gradient–moving-average agreement.The discovered family has the form (α + f(t) ∗sign(g) ∗sign(m)) ∗g.
5.1. Discovered optimizers: When f(t) = 0, PowerSign and AddSign reduce to the usual SGD update, so internal decays interpolate toward SGD late in training.AddSign with α = 1 and no internal decay updates only parameters whose gradient and moving average agree.
5.1. Discovered optimizers: Variants replacing g with m, Adam, or RMSProp perform less well, while larger PowerSign base α typically accelerates convergence with minor final-performance loss.The default definitions use α = e for PowerSign and α = 1 for AddSign, without internal decay.
5.2. Discovered learning rate decays: Linear cosine decay and noisy linear cosine decay typically permit larger initial learning rates and faster convergence than cosine decay.The controller discovers linear cosine decay as ld ∗cd and noisy linear decay as (ld+ϵt)∗cd+0.001.

6. Transferability experiments

The discovered optimizers transfer across architectures and tasks, improving or matching established methods in several settings, though performance depends on training configuration and can lag behind Adam on PTB.

Transferability experiments: The search-generated update rules were reused without retraining the controller across new architectures and tasks.The transfer experiments test rules discovered in an earlier experiment directly on different models and datasets.
Rosenbrock function: On the Rosenbrock function, PowerSign outperformed Adam, RMSProp, and SGD, while nearly matching Momentum.Each optimizer ran for 4000 iterations across four logarithmically spaced learning rates, with the best performance plotted.
CIFAR-10 Wide ResNet: Without learning-rate decay on CIFAR-10 Wide ResNet, some discovered optimizers outperformed common optimizers by a sizable margin, especially with internal decays.The comparison used five-run averages; the corresponding results are reported in Table 1 and Figure 6.
CIFAR-10 Wide ResNet: With learning-rate decay on CIFAR-10 Wide ResNet, the discovered optimizers outperformed Adam, RMSProp, and Momentum and were on par with SGD.The authors also observed faster convergence than Momentum and SGD, although optimizer-result variance may reflect learning-rate selection.
ImageNet classification: On ImageNet, PowerSign achieved 0.4% top1 and 0.4% top5 accuracy improvements over RMSProp, along with faster convergence.PowerSign and AddSign were evaluated on a MobileNet baseline using cosine learning-rate decay.
Neural machine translation: On WMT 2014 English →German, PowerSign improved training perplexity by 0.1 and test BLEU by 0.5 over Adam, while using one running average per parameter instead of two.The GNMT hyperparameters had previously been tuned for Adam, and further tuning was expected to improve the new rule.
Language modeling: On PTB language modeling, PowerSign and AddSign substantially outperformed SGD but were outperformed by Adam.The authors note that hyperparameter ranges were first manually tuned with Adam, which may have affected the other optimizers’ performance.

7. Conclusion

The approach automates optimizer discovery while producing update rules that transfer across optimization tasks. It yields PowerSign and AddSign, plus a learning-rate schedule that improves convergence across reported settings.

7. Conclusion: The method searches for optimizers suited to environments such as 4-bit computation or limited communication.Its search can incorporate the constraints of the optimization environment.
7. Conclusion: Unlike prior learning-to-learn approaches, the discovered equation-form update rules transfer readily to other optimization tasks.
7. Conclusion: PowerSign and AddSign achieve competitive performance across tasks and models from ConvNets for image classification to LSTMs for machine translation.
7. Conclusion: Linear cosine decay generally produces faster convergence than cosine annealing.
7. Conclusion: The discovered update rules provide a way to improve training of deep networks while opening new approaches to update-rule design.

8. Appendix

The appendix gives practical learning-rate guidance for PowerSign and AddSign and defines how Wide-ResNet performance is reported across training runs.

8. Appendix: For default PowerSign and AddSign hyperparameters, SGD's learning rate is usually a good first choice.The defaults are α = e for PowerSign and α = 1 for AddSign.
8. Appendix: When using learning-rate decay, the initial learning rate should be larger than without decay.
8. Appendix: With linear cosine decay, the initial learning rate can be larger than with cosine decay.
8. Appendix: Wide-ResNet results report final and best validation/test accuracy after 300 epochs, selecting each optimizer's best result from seven learning rates.Best Test is measured at the epoch with the highest validation accuracy.