[Paper Review] Population Based Training of Neural Networks

Max Jaderberg, Valentin Dalibard+10 more|arXiv (Cornell University)|Nov 27, 2017

Machine Learning and Data ClassificationComputer Science249 citations

TL;DR

Population Based Training (PBT) jointly optimizes neural network weights and hyperparameters across a population of models, online and asynchronously, to yield adaptive hyperparameter schedules and improved final performance.

ABSTRACT

Neural networks dominate the modern machine learning landscape, but their training and success still suffer from sensitivity to empirical choices of hyperparameters such as model architecture, loss function, and optimisation algorithm. In this work we present \emph{Population Based Training (PBT)}, a simple asynchronous optimisation algorithm which effectively utilises a fixed computational budget to jointly optimise a population of models and their hyperparameters to maximise performance. Importantly, PBT discovers a schedule of hyperparameter settings rather than following the generally sub-optimal strategy of trying to find a single fixed set to use for the whole course of training. With just a small modification to a typical distributed hyperparameter training framework, our method allows robust and reliable training of models. We demonstrate the effectiveness of PBT on deep reinforcement learning problems, showing faster wall-clock convergence and higher final performance of agents by optimising over a suite of hyperparameters. In addition, we show the same method can be applied to supervised learning for machine translation, where PBT is used to maximise the BLEU score directly, and also to training of Generative Adversarial Networks to maximise the Inception score of generated images. In all cases PBT results in the automatic discovery of hyperparameter schedules and model selection which results in stable training and better final performance.

Motivation & Objective

• Motivate the need for joint optimization of model parameters and hyperparameters in neural network training.
• Introduce a simple, asynchronous meta-optimization method that uses a population of models to discover adaptive hyperparameter schedules.
• Demonstrate PBT across reinforcement learning, machine translation, and GANs to show faster convergence and better final performance.
• Show that online hyperparameter adaptation and model selection within a population improve training stability and outcomes.

Proposed method

• Define eval to measure the true objective Q on the current model state.
• Train N models in a population with different hyperparameters h, using step to update parameters θ.
• Periodically evaluate performance p and, when ready, perform exploit (copy weights/hyperparams from better performers) and explore (perturb or resample hyperparameters).
• Operate asynchronously without a central controller, updating the population with θ, h, and p.
• Exploit copies weights/hyperparameters from a better performing member; explore perturbs or resamples hyperparameters to broaden search.
• Apply PBT to RL (UNREAL, Feudal Networks, A3C on StarCraft II), machine translation (Transformers on WMT14 EN-DE), and GANs (Inception score).

Experimental results

Research questions

• RQ1 Can online, asynchronous, population-based optimization outperform traditional sequential or parallel hyperparameter tuning?
• RQ2 Do adaptive, non-stationary hyperparameter schedules discovered by PBT lead to better final performance across RL, MT, and GAN domains?
• RQ3 What are the effects of population size, exploitation strategy, and exploration method on PBT performance?
• RQ4 Is joint optimization of weights and hyperparameters via PBT more robust than fixed hyperparameter schedules or naive random/grid search?
• RQ5 Can PBT enable automatic discovery of effective hyperparameter schedules (e.g., learning rate decay, unroll lengths) during training?

Key findings

• PBT increases final performance over random search baselines across DM Lab, Atari, StarCraft II, MT, and GAN tasks.
• Hyperparameters adapt during training, with automatic discovery of schedules such as learning rate decay and unroll length.
• Exploitation and online model selection propagate good solutions across the population, improving stability and sample efficiency.
• PBT achieves higher final performance even with the same training steps, and can reduce the need for extensive hand-tuning.
• A smaller population can increase variance; populations of 20+ workers provide consistent improvements, with diminishing returns at larger sizes.