[Paper Review] Faster Improvement Rate Population Based Training

Valentin Dalibard, Max Jaderberg|arXiv (Cornell University)|Sep 28, 2021

Machine Learning and Data Classification21 references4 citations

TL;DR

This paper proposes Faster Improvement Rate Population Based Training (FIRE PBT), an enhanced PBT algorithm that uses a novel fitness metric based on training improvement rate to avoid greedy hyperparameter choices. By dividing the population into sub-populations optimized for different phases of learning rate schedules, FIRE PBT achieves better final performance than standard PBT and matches hand-tuned learning rate schedules on ImageNet and reinforcement learning benchmarks.

ABSTRACT

The successful training of neural networks typically involves careful and time consuming hyperparameter tuning. Population Based Training (PBT) has recently been proposed to automate this process. PBT trains a population of neural networks concurrently, frequently mutating their hyperparameters throughout their training. However, the decision mechanisms of PBT are greedy and favour short-term improvements which can, in some cases, lead to poor long-term performance. This paper presents Faster Improvement Rate PBT (FIRE PBT) which addresses this problem. Our method is guided by an assumption: given two neural networks with similar performance and training with similar hyperparameters, the network showing the faster rate of improvement will lead to a better final performance. Using this, we derive a novel fitness metric and use it to make some of the population members focus on long-term performance. Our experiments show that FIRE PBT is able to outperform PBT on the ImageNet benchmark and match the performance of networks that were trained with a hand-tuned learning rate schedule. We apply FIRE PBT to reinforcement learning tasks and show that it leads to faster learning and higher final performance than both PBT and random hyperparameter search.

Motivation & Objective

• To address the limitation of Population Based Training (PBT) favoring short-term performance gains over long-term generalization.
• To develop a method that identifies hyperparameter schedules leading to superior final performance by emphasizing rate of improvement during training.
• To enable automated hyperparameter tuning that matches or exceeds the performance of hand-tuned learning rate schedules in image classification and reinforcement learning.
• To reduce reliance on manual tuning by creating a scalable, parallelizable optimization framework that learns effective hyperparameter schedules autonomously.

Proposed method

• FIRE PBT divides the population into disjoint sub-populations, each trained with a fitness function based on the rate of improvement in validation performance.
• The fitness metric evaluates how quickly a model improves per unit of training time, favoring models with steeper initial improvement curves.
• Sub-populations are trained with different hyperparameter schedules, with each focusing on a distinct phase of the learning curve (e.g., high learning rate vs. decayed rate).
• The method uses a truncation selector to allow underperforming members to copy weights and hyperparameters from top-performing peers in the same sub-population.
• Hyperparameters such as learning rate and trust region size are mutated independently using multiplicative noise with predefined scaling factors.
• Evaluators monitor performance and trigger exploitation and exploration steps at regular intervals, ensuring continuous adaptation without disrupting training.

Experimental results

Research questions

• RQ1 Can a fitness metric based on improvement rate outperform standard PBT’s immediate performance-based fitness in hyperparameter optimization?
• RQ2 Does FIRE PBT learn hyperparameter schedules that match or exceed the performance of hand-tuned learning rate schedules on ImageNet?
• RQ3 Can FIRE PBT improve sample efficiency and final performance in reinforcement learning compared to PBT and random search?
• RQ4 Does the sub-population decomposition strategy in FIRE PBT enable better exploration of complex, multi-phase hyperparameter schedules?

Key findings

• On the ImageNet benchmark, FIRE PBT achieved a test accuracy of 76.5%, outperforming PBT’s 71.7% and matching the 76.6% accuracy of a hand-tuned learning rate schedule.
• In reinforcement learning on Ant-v1 and Humanoid-v1, FIRE PBT achieved faster learning and higher final performance than both PBT and random search, as shown by the distribution of top episode returns.
• The use of improvement rate as a fitness metric successfully mitigated PBT’s tendency to prematurely decay the learning rate, avoiding suboptimal final performance.
• FIRE PBT demonstrated robustness across different environments and hyperparameter configurations, maintaining consistent performance gains without requiring additional hyperparameter tuning.