Learning to Schedule Learning rate with Graph Neural Networks
Learning to Schedule Learning rate with Graph Neural Networks
复制标题
DOI：
--
复制
发表时间：
2022
期刊：
影响因子：
--
通讯作者：
Yuanhao Xiong;Li-Cheng Lan;Xiangning Chen;Ruochen Wang;Cho-Jui Hsieh
中科院分区：
文献类型：
--
作者：
Yuanhao Xiong;Li-Cheng Lan;Xiangning Chen;Ruochen Wang;Cho-Jui Hsieh
近几十年来，随机优化在训练深度神经网络方面取得了很大的发展。学习率调度是影响Adam等随机优化器性能的最重要因素之一。传统的方法试图在有限数量的预定义规则中找到相对合适的调度，并且可能不适应特定的目标问题。相反，我们提出了一种新的基于图网络的调度算法（GNS），旨在学习一种特定的调度机制，而不限制现有的原则。通过为目标问题的底层神经网络构建有向图，GNS用图消息传递网络对当前动态进行编码，并通过强化学习训练代理相应地控制学习速率。建议的调度程序可以捕获的中间层信息，同时能够推广到不同规模的问题。此外，一个有效的奖励收集程序被用来加速训练。我们在基准数据集上评估了我们的框架，Fashion-MNIST和CIFAR 10用于图像分类，GLUE用于语言理解。在训练CNN和Transformer模型时，GNS显示出对流行基线的一致改进。此外，GNS对不同的数据集和网络结构表现出很大的泛化能力。我们的代码可在https://github.com/xyh97/GNS上获得。
Recent decades have witnessed great development of stochastic optimization in training deep neural networks. Learning rate scheduling is one of the most important factors that influence the performance of stochastic optimizers like Adam. Traditional methods seek to find a relatively proper scheduling among a limited number of pre-defined rules and might not accommodate a particular target problem. Instead, we propose a novel Graph-Network-based Scheduler (GNS), aiming at learning a specific scheduling mechanism without restrictions to existing principles. By constructing a directed graph for the underlying neural network of the target problem, GNS encodes current dynamics with a graph message passing network and trains an agent to control the learning rate accordingly via reinforcement learning. The proposed scheduler can capture the intermediate layer information while being able to generalize to problems of varying scales. Besides, an efficient reward collection procedure is leveraged to speed up training. We evaluate our framework on benchmarking datasets, Fashion-MNIST and CIFAR10 for image classification, and GLUE for language understanding. GNS shows consistent improvement over popular baselines when training CNN and Transformer models. Moreover, GNS demonstrates great generalization to different datasets and network structures. Our code is available at https://github.com/xyh97/GNS.