Learning an Interpretable Learning Rate Schedule via the Option Framework

Abstract

Learning rate is a common and important hyperparameter in many gradient-based optimizers which are used for training machine learning models. Heuristic handcrafted learning rate schedules (LRSs) can work in many practical situations, but their design and tuning is a tedious work, and there is no guarantee that a given handcrafted LRS matches a given problem. Many works have been dedicated to automatically learning an LRS from the training dynamics of the optimization problem, but most of them share a common deficit: they borrow the algorithms designed elsewhere as methods for automatic outer-training, but those methods often lack interpretability in the context of learning an LRS. In this paper, we leverage the option framework, a generalization to the common rein-forcement learning framework, to automatically learn an LRS based on the dynamics of optimization, which takes the idea of temporal abstraction as an underlying interpretation. To meet the requirements of LLRS, the RL state is designed as consisting of the global state and the per-parameter state. We propose a policy architecture which processes these two parts according to their respective structures, and combines them to yield the input for functional computation. Experiments are carried out on classic machine learning tasks and test functions for numerical optimization to demonstrate the effectiveness and the interpretability of our method.