L EARNING TO S CHEDULE L EARNING R ATE WITH G RAPH N EURAL N ETWORKS
A BSTRACT
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 inﬂuence the performance of stochastic optimizers like Adam.
Traditional methods seek to ﬁnd a relatively proper scheduling among a limited number of pre-deﬁned rules and might not accommodate a particular target problem.
Instead, we propose a novel Graph-Network-based Scheduler (GNS), aiming at learning a speciﬁc 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 efﬁcient reward collection procedure is leveraged to speed up training.
We evaluate our framework on benchmarking datasets, Fashion-MNIST and CIFAR10 for image classiﬁcation, 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 .
I NTRODUCTION
Stochastic optimization has achieved great successes in training deep learning models (Robbins & Monro, 1951), due to the tremendous data size and non-convexity of neural networks (Xiong et al., 2020).
However, in contrast to classical and hyperparameter-free optimizers, performance of widely used stochastic algorithms such as SGD and Adam is signiﬁcantly dependent on the choice of hyperparameters including batch size, learning rate and momentum coefﬁcients (Xiong et al., 2020;
Sivaprasad et al., 2020; Cooper et al., 2021; Choi et al., 2019).
Learning rate scheduling is one of the most important factors which stochastic optimizers are very sensitive to (Cooper et al., 2021;
Xu et al., 2019).
A rate throughout the whole training is insufﬁcient to attain an optimal model.
Thus, various scheduling algorithms (Gotmare et al., 2018) including polynomial decay, cosine decay and warmup have been developed and each of them is designed in a distinct form.
In practice, models trained with an appropriate scheduling mechanism can outperform those with a constant learning rate, in terms of convergence rate as well as the ﬁnal performance (Senior et al., 2013; Gotmare et al., 2018; Loshchilov & Hutter, 2016).
As the learning rate scheduling is of vital importance to the model performance, how to tune it for a given stochastic optimizer effectively and automatically has been investigated intensively.
Current approaches tackle this problem by assuming a particular learning rate scheduling rule (e.g., polynomial or cosine decay, warmup and restart) based on empirical studies and domain knowledge (Xu et al., 2017; Gotmare et al., 2018), and apply traditional hyperparameter optimization methods (Li et al., 2017; Snoek et al., 2012; Bergstra & Bengio, 2012) to tune the parameters of the rule.
However, searching within those pre-deﬁned principles is restricted since the best learning rate scheduling for a particular training problem may not exactly follow any existing rule.
To overcome this restriction, ﬁnding an optimal scheduler can be viewed as a learning-to-learn problem (Andrychowicz et al., 2016;
Chen et al., 2020), where empirical scheduling principles are replaced by a data-driven neural network alternative.
For instance, attempts such as Xu et al.
(2019) and Xu et al.
(2017) model the learning rate scheduling as a sequential decision process where the learning rate is modiﬁed in every decision step, and train a policy network with reinforcement learning algorithms.
However, these methods sacriﬁce rich state representations in order to maintain generalization.
This implies that problem-dependent information, such as the architecture of neural networks to be optimized, and the state of intermediate neurons is not fully exploited.
Therefore, state-of-the-art algorithms such Chen et al.
(2020) only represent the current state by the statistics of ﬁnal layer neurons.
A straightforward way to leverage such information is to concatenate states in every layer as a super vector.
However, it would lead to the explosion of the input dimension and fail to work in models with different sizes.
Therefore it is nontrivial to obtain a representation of training state that is both informative (being able to encode rich information of architecture and intermediate neurons) and generalizable (can be applied to different architectures).
In this paper, we propose a Graph Network-based Scheduler (GNS) to overcome the above-mentioned challenges.
Different from previous attempts, our GNS makes full use of the problem structure: since deep neural networks can be represented as directed graphs intrinsically Zhang et al.
(2018), we leverage graph networks to encode states into latent features with richer information and such a design allows GNS to be easily transferred to different network architectures at the same time.
Moreover, an action deﬁnition for stabling training and an efﬁcient reward signal collection procedure are designed for our GNS framework.
Our main contributions are summarized as following:
• We develop the Graph Network-based Scheduler (GNS) to search for the best learning rate scheduling rule by representing the target network as a directed graph.
Our dynamic learning rate scheduler can thoroughly capture the intermediate layer information of neural network in both image classiﬁcation and language understanding tasks.
Also, an efﬁcient reward collection procedure is designed to speedup our method.
• We apply the proposed GNS to multiple training tasks.
In particular, GNS achieves average scores of 85.2 and 87.8 on GLUE with RoBERTa-base and -large respectively and outperforms 84.7 and 86.9 from the best existing learning rate schedule.
Note that unlike previous learning-to-learn methods that mostly are applied to toy examples such as CIFAR-10, we are the ﬁrst to apply automatic learning rate scheduling to challenging tasks such as RoBERTa training and improve the performance of those realistic large models.
• We demonstrate that the GNS framework is able to generalize to problems of different datasets (from CIFAR10 to CIFAR100) and model structures (from ResNet to VGG, and from RoBERTa-base to RoBERTa-large).
R ELATED W ORK
Learning rate scheduling.
This topic has been an open problem in stochastic optimization.
In Almeida et al.
(1998) and Baydin et al.
(2017), learning rate was regarded as a trainable parameter and modiﬁed using its gradient with respect to the objective function.
Sutton (1992) and
Schraudolph et al.
(2006) considered a meta-descent algorithm to adjust local learning rate with a meta one.
In the era of deep learning, a number of adaptive optimization algorithms have been proposed, such as RMSProp (Tieleman & Hinton, 2012), AdaGrad (Duchi et al., 2011), and Adam (Kingma & Ba, 2014).
They conduct parameter-wise adaptation to learning rate based on heuristics of the geometry of the traversed objective.
For instance, RMSProp normalizes the learning rate with the magnitude of recent gradients while AdaGrad generates a large learning rate for infrequently-updated parameters.
Complimentary to these parameter-wise adaptations, some predeﬁned learning rate schedulers are also observed beneﬁcial to training performance (You et al., 2019;
Ge et al., 2019;
Schaul et al., 2013; Senior et al., 2013).
Mechanisms like polynomial (Mishra & Sarawadekar, 2019) and cosine decay (Loshchilov & Hutter, 2016) together with a warmup process (Gotmare et al., 2018) improve the convergence rate and ﬁnal performance of deep learning models empirically.
However, they still hold additional hyperparameters to be tuned.
At the same time, these parametric schedules have limited ﬂexibility and may not be optimized for the training dynamics of different high dimensional and non-convex optimization problems (Xu et al., 2017; 2019).
Reinforcement learning.
The goal of reinforcement learning (RL) is to ﬁnd an agent that produces an optimal policy to maximize a cumulative reward (Sutton & Barto, 2018).
Recently, deep reinforcement learning with neural networks as approximation functions has shown great potentials in many applications, such as Atari games (Gu et al., 2016; Hester et al., 2018; Schulman et al., 2017),
Go (Silver et al., 2016; 2017), and training deep learning models (Zoph & Le, 2016; Cubuk et al., 2018).
There are also several attempts to apply RL to learning rate controlling.
For example,
Daniel et al.
(2016) proposed to improve the stability of training w.r.t.
the initial learning rate via reinforcement learning.
Li & Malik (2016), Xu et al.
(2017) and Xu et al.
(2019) leveraged RL agents to adjust the learning rate to reach a ﬁnal performance whereas their selection of state features, reward and action leads to an inefﬁcient training process, and is difﬁcult to scale up to large models such as
BERT (Devlin et al., 2018) and RoBERTa (Liu et al., 2019).
Graph Neural Networks.
Graph Neural Networks (GNN) is a type of neural networks that operate on graph data and structure their computations accordingly (Battaglia et al., 2018).
A graph network, or a graph-to-graph module, can be interpreted as a mapping from an input graph to an output graph with the same structure but potentially different node, edge, and graph-level representations (Battaglia et al., 2018; Sanchez-Gonzalez et al., 2020).
The most important component in GNN, Message
Passing Neural Networks (MPNN) (Gilmer et al., 2017), have been shown effective in learning dynamics for graph-structured settings.
Generally, MPNN utilizes the graph structure to propagate information between nodes via edges, and extract latent features for downstream tasks.
Variations such as GCN (Kipf & Welling, 2016) and Interaction Networks (Battaglia et al., 2016) have previously demonstrated promising results on simulation tasks in the ﬁeld of physics (Battaglia et al., 2016;
Sanchez-Gonzalez et al., 2020) and chemistry (Gilmer et al., 2017; You et al., 2018).
R EINFORCEMENT L EARNING FOR L EARNING R ATE S CHEDULING
In this section, we formulate learning rate scheduling as an RL problem, with the goal to ﬁnd an optimal control policy to maximize the accumulated reward in the current environment which is deﬁned by the target training problem.
3. R EINFORCEMENT L EARNING
Reinforcement learning aims at solving a group of problems, known as Markov Decision Process (MDP) with an agent choosing optimal actions (Sutton & Barto, 2018).
Typically, the MDP problem is formulated by a tuple M = (cid:104)S , A , P , R , γ (cid:105) .
Given the state set S and action set A , the reward function R is a mapping of S ×
A → R .
The discounted return is denoted by R t = (cid:80) ∞ l =
γ l r t + l where γ is a discount factor for future rewards.
P is the state transition function, which speciﬁes the probability p ( s t + | s t , a t ) and is usually determined by the environment itself.
To encourage more exploration, we focus on a stochastic policy π ( a | s ) , which maps states to action probabilities.
3. L EARNING RATE SCHEDULING AS A MDP
For most problems involved in training deep neural networks, a standard approach for the objective function minimization is stochastic gradient descent on a certain form (Robbins & Monro, 1951;
Kingma & Ba, 2014; Duchi et al., 2011), which sequentially updates the parameters using gradients by w k +1 =
w k − α k · g k .
Speciﬁcally, w k represents network parameters to be trained and α k is the global learning rate.
g k is deﬁned by the optimizer (e.g., g k can be the SGD or Adam direction).
For learning rate scheduling, we primarily consider generating a sequence for the global learning rate α k .
We can formulate learning rate scheduling into a sequential decision process.
In each environment of a particular optimization task, states can be described by dynamics of the target network such as a combination of weight and gradient values (Xu et al., 2019; Daniel et al., 2016).
Actions can be any operation that might change the learning rate.
As to rewards, since the ultimate goal is to achieve higher ﬁnal performance, metrics like loss or accuracy can serve as a signal to guide the agent training.
Besides, as training a deep learning model usually requires a large number of iterations, rather than a single optimization step, we execute an action every K steps to avoid instability from too frequent learning rate modiﬁcations.
By training such an agent modeled as a policy network, we can obtain a learning rate scheduler to adjust learning rate dynamically.
M ETHOD 4. S TATE E NCODING WITH G RAPH N EURAL N ETWORKS
The selection of features to represent the environment is of vital signiﬁcance.
Previous work such as (Xu et al., 2017) and (Xu et al., 2019) aims to learn a generalized agent that can be applicable to any network architecture, so they were only able to utilize statistics of the ﬁnal layer to describe the state.
It might work for small scale models, but when dealing with large ones like RoBERTa (Liu et al., 2019), solely using information of the last layer is insufﬁcient to describe the environment.
A straightforward way to obtain a more thorough representation is to utilize dynamics of all layers in a neural network and concatenate them into one super vector.
However, such a representation cannot guarantee promising results as shown in Section 5.6 of ablation study.
The problem is that naively concatenating all features increases the input dimension and does not explore the correlation between layers.
For instance, dynamics between nearby layers should have a higher correlation than layers that are far away.
This concatenation also loses transferrablity to different network structures.
Therefore, to maintain a rich state description and generalization in the meanwhile, we construct a computational graph to depict the state.
We represent a network structure as a directed graph
G =
( E , V ) .
Each node v denotes an operator f v ( · ; w v ) (convolution, linear, etc.) parameterized by w v which generates an output activation tensor.
An edge e uv ∈ E represents the computation ﬂow from node u to node v .
The output for node v is computed by applying its associated computational operator taking all source nodes as the input.
Examples of the graph structure of a neural network are presented in Figure 1(a) and (b).
For each node v with trainable parameters w v , the raw feature x v includes the mean and variance of weights, biases, and absolute gradient values of w v , as well as global dynamics of the previous learning rate and the average training loss in the decision step.
In the end, the environment can be represented by a graph with a feature matrix X ∈ R N × F and the adjacency matrix A ∈ R N × N , where N is the number of nodes and F is the feature dimension.
[Figure]
… … <latexit
[Table]
/E/r53q8MLLmEhSTQWZLgpTjnSMxr+jLpOUaD40BBPJzK2I9LHERJuECiYEd/7lRdIo2+6Z7dyZNCowRR4O4RhOwYVzqMIN1KAOBAbwBC/waiXWs/VmvU9bc9ZsZh/+wPr4Aa8Lkgg=</latexit> h G M <latexit sha1_base64="aqq5cK8RDHifH1TjjeONb2z3+ok=">AAAB7nicbVDLSsNAFL2pr1pfVVciyGARXIWkIFpwUXChG6GCfUAbwmQ6aYdOJmFmIpTQXxDcuFDErd/jzi9x6/Sx0NYDFw7n3Mu99wQJZ0o7zpeVW1peWV3Lrxc2Nre2d4q7ew0Vp5LQOol5LFsBVpQzQeuaaU5biaQ4CjhtBoOrsd98oFKxWNzrYUK9CPcECxnB2kjNvp9d+7cjv1hybGcCtEjcGSlVDx4r9tH3Zc0vfna6MUkjKjThWKm26yTay7DUjHA6KnRSRRNMBrhH24YKHFHlZZNzR+jEKF0UxtKU0Gii/p7IcKTUMApMZ4R1X817Y/E/r53q8MLLmEhSTQWZLgpTjnSMxr+jLpOUaD40BBPJzK2I9LHERJuECiYEd/7lRdIo2+6Z7dyZNCowRR4O4RhOwYVzqMIN1KAOBAbwBC/waiXWs/VmvU9bc9ZsZh/+wPr4Aa8Lkgg=</latexit>
h G M
[Figure]
Given an input architecture, we describe the state as a graph G with the feature matrix X and the adjacency matrix A .
Then we can feed G into an encoder E , modeled as a graph convolutional neural network (GCN) (Kipf & Welling, 2016), one of the most effective techniques to achieve message propagation.
More concretely, hidden features can be extracted by
[Equation]
where H 0 is initialized as X ,
˜ A is the normalized adjacency matrix with a self-loop in each node, and W l is the parameters of the l -th layer of GCN.
However, a deep learning model usually adopts hundreds of operations and results in a large computation graph which is hard to process.
We notice that repeated blocks are stacked in commonly-used structures such as ResNet (He et al., 2016) and Transformer (Devlin et al., 2018; Vaswani et al., 2017) to boost the expressiveness of the network (Zhang et al., 2018;
Zoph et al., 2018), shown in
Figure 1(a) and (b).
Instead of passing messages through the whole graph directly, we separate the graph into a number of blocks as { G i =
( V i , E i ) }
M +1 i
=
and each block has a group of nodes.
M is the number of repeated blocks, and G 0 stands for the input block (e.g., embedding layers) while G M + represents the output block (e.g., classiﬁcation layers).
We can then use a shared encoder across different blocks as shown in Figure 1(c) to conduct information propagation.
This is implemented by a hierarchical GCN as follows: (i) compute latent features of nodes in each block by Eq. 1; (ii) summarize the i -th block as h G i =
AGG ( { h v | v ∈ V i } ) ; (iii) then use h G i as the input node feature for the next block: (iv) Finally, to obtain the graph level information to describe the whole network, we just take the aggregation over each node by h G
= AGG ( { h v | v ∈ V} ) .
Here, we use the mean function as aggregation function to allow generalization to different structures.
Compared with a concatenated state with dynamics in all nodes, representing the environment as a graph and then encoding it via a hierarchical GCN have two primary advantages: (1) GCN can capture the correlation between neighboring nodes and generate a meaningful latent feature for action prediction; (2) a shared encoder can learn message passing within one block and between blocks in the meanwhile and is not constrained to a speciﬁc graph structure.
In detail, the dimension of a concatenated representation varies between different network architectures, for example, RoBERTabase and -large models.
A scheduler learned from such concatenated states fails to work in another environment due to the input mismatch.
On the contrary, by stacking the same encoder along the depth dimension, we can transfer our graph-network-based scheduler trained on one network to the new environment of another different network directly or with slight ﬁne-tuning.
4. E NVIRONMENT DESIGN
We discuss the remaining components of the environment of the target problem involved in neural network optimization in this part besides the state.
An effective set of reward functions and actions is deﬁned to facilitate an efﬁcient and generalizable agent.
Action.
We adopt the action as scaling the previous learning rate by a factor (Xu et al., 2019) considering the wide range of α .
Through that operation, the learning rate can be changed stably and consistently.
More concretely, after each action execution, α t +1 =
α t · (1 + a t ) .
State transition.
As mentioned in Section 3.2, once we modify the learning rate to α t +1 , we execute the update rule of gradient descent for K batches to get the next state s t +1 .
Reward design.
With the goal of the agent to achieve higher ﬁnal performance, there are a series of metrics that can be leveraged, such as validation loss and accuracy.
In Xu et al.
(2019), validation loss at each decision step serves as an immediate reward.
However, loss might not be an ideal indicator for task speciﬁc metrics which we are ultimately concerned with, like accuracy, F1 score and correlation (Venter et al., 2021).
Thus, we choose to use such task-related metrics directly instead of the loss.
On the other hand, since deﬁning an intermediate reward can utilize training information thoroughly and provide direct feedback (Xu et al., 2019), we take the performance difference between two steps as the reward signal:
[Equation]
where P t represents the validation result under the given metric at the t -th decision step.
Whereas, computing a reward in Eq. 2 indicates that we have to conduct a validation procedure every decision step, which is a non-trivial burden.
In order to make training more efﬁcient, we turn to reduce the frequency of validation with a probability p , i.e., the validation is only carried out if the number randomly sampled from Uni (0 , 1) is less than p .
Speciﬁcally, suppose under our principle, we calculate the validation metric P t 1 and P t 2 at the decision step t 1 and t 2 respectively.
The reward signal for each step can be derived by linearly distributing the difference as follows:
[Equation]
4. G RAPH N ETWORK - BASED S CHEDULER
With illustration of the environment design, we can then move to the general framework of our
Graph Network-based Scheduler (GNS).
As shown in Figure 2, at each decision step, GNS takes the observation s t as input, encodes s t into h t , sample an action a t based on the policy network, evaluates the value of s t , then execute the action to render the next state s t +1 .
Encoding
[Table]
/SHk4/9PyDknUFIYdN1PJ7ewuLS8kl8trK1vbG4Vt3fqJk404zUWy1g3A2q4FBGvoUDJm0pzGgaSN4Lh5SRv3HJtRBzd4EhxP6T9SPQEo2ithumkeOyNO8WSW3YzkXnwfqBUIW9f6jD3Xu0UP9rdmCUhj5BJakzLcxX6KdUomOTjQjsxXFE2pH3eshjRkBs/
zdodkwPrdEkv1vZESDL394uUhsaMwsBWhhQHZjabmP9lrQR7534qIpUgj9j0o14iCcZkMjvpCs0ZypEFyrSwvRI2oJoytBsq2CV4syPPQ/2k7J2W3WuvVLm4g0x52IN9OAIPzqACV1CFGjAYwj08wpOjnAfn2XmZluac6Q278EfO6zewQZM0</latexit> s t +
In the procedure of sampling and evaluation, we use the hidden feature h t to evaluate the current state and predict the action.
Denote m s and m f as two different MLPs.
Speciﬁcally, the state value can be computed by V ( s t )
= m s ( h t ) .
Then we calculate the mean of the action and the policy π samples a t from the corresponding normal distribution in Eq.
4 with a ﬁxed deviation σ :
[Equation]
By Eq.
4, we can constrain that the scale factor (1 + a t ) is within the range of [0 . 9 , 1 . 1] at most of the times in order to train the target network stably (Xu et al., 2019).
As to the agent training, we ﬁrst parameterize the value network V (critic), and policy network π (actor) by φ, ϕ respectively for clarity.
Each environment of the target problem has a ﬁxed time horizon T , i.e., the number of total training steps, which is called as one episode.
We adopt Proximal
Policy Optimization (PPO) (Schulman et al., 2017) in the actor-critic style to optimize our scheduling policy, which is the state of the art on-policy RL methods in many continuous control tasks (Achiam, 2018; Berner et al., 2019).
On-policy RL methods only use the experience generated by current policy.
Therefore, the training will be more stable due to the training data distribution.
The key objective function of the policy network is deﬁned as
[Equation]
where ρ t ( ϕ ) is the probability ratio and ˆ A t is the estimated advantage value.
Besides, a learned state-value function V φ ( · ) is leveraged to reduce variance in advantage estimators and can be trained by minimizing the following squared-error loss:
[Equation]
where V targ
[Equation]
= r t + γV ( s t +1 ) .
During the training phase, we run the agent in the environment for one episode, and then update the value and policy network at the end.
We reset environment at the beginning of each episode and repeat the process for multiple episodes.
The detailed algorithm is presented in Appendix A. When applying the scheduler, we ﬁrst sample a number of initial learning rates, run each setting for T pre steps, compute corresponding V ( s T pre ) , and choose the learning rate with the largest value to complete the training.
E XPERIMENTAL R ESULTS
To validate the effectiveness of GNS, we evaluate our method on various tasks in image classiﬁcation and language understanding, and compare it with popular learning rate scheduling rules.
We further investigate the generalization of GNS on different transfer tasks.
In addition, we conduct an ablation study to analyze the state representation and reward collection.
5. E XPERIMENTAL S ETTINGS
In this section, we present our experimental settings.
Further details can be found in Appendix B.
Image classiﬁcation.
We consider two benchmark datasets in image classiﬁcation, FashionMNIST (Xiao et al., 2017) and CIFAR10 (Krizhevsky et al., 2014).
These two datasets are ﬁrst split into the standard training and test sets.
Then we randomly sample 10 k images for each dataset from the training set to construct a validation set.
A ResNet-18 (He et al., 2016) model is trained on
Fashion-MNIST while a ResNet-34 (He et al., 2016) is trained on CIFAR10.
We use Adam (Kingma & Ba, 2014) with the batch size of 128 for 200 epochs to train these two image classiﬁcation tasks.
Language understanding.
For language understanding, we conduct experiments on GLUE (Wang et al., 2019), a benchmark consisting of eight sentence- or sentence-pair tasks.
They are divided into training, validation and test sets and we have no access to ground truth labels of test sets.
On each dataset, we ﬁne-tune a pre-trained RoBERTa model (Liu et al., 2019), use it to generate predictions on the test set, and obtain the corresponding performance through the ofﬁcial GLUE platform.
Both
RoBERTa-base and -large models are evaluated.
No ensembling techniques are adopted for purely validating the effectiveness of our proposed scheduler.
The AdamW (Loshchilov & Hutter, 2017) optimizer is adopted to train RoBERTa models.
Details of other hyperparameters like batch size and episode length for each task are provided in Appendix B.
Baselines.
Three hand-designed scheduling mechanisms (constant, polynomial decay, cosine decay), one gradient-based method Hypergradient (Baydin et al., 2017), and one learning-based method,
SRLS (Xu et al., 2019) are compared with our GNS framework as baselines.
For traditional schedulers with the warmup process, we use the following equation to summarize them:
[Equation]
In detail, α max and α min = µα max are the maximum and minimum learning rate respectively during the training.
T total is the number of total training steps, T warm =
η 1 T total is the number of warmup steps and T decay =
η 2 T total is the number of decay steps.
We tune all hyperparameter combinations within the search space in Appendix B and select the conﬁguration with the best validation performance.
Then we run the experiment 5 times and report the average test performance.
For Hypergraident and
SRLS, we follow training details in Baydin et al.
(2017) and Xu et al.
(2019) respectively.
GNS implementation.
To encode raw observations into latent features, we use a 2-layer GCN (Kipf & Welling, 2016) with 64 dimensional node embeddings in all hidden layers, and apply batch normalization for each layer.
After a shared encoder, each of the value network and policy network has an MLP with a hidden layer of size 32.
Our GNS is trained with PPO (Schulman et al., 2017) for 30 episodes for ResNet models and RoBERTa-base while 10 episodes for RoBERTa-large by ﬁne-tuning the agent learned from RoBERTa-base, by an Adam optimizer of the learning rate 0 .and 0 .
001 for the critic and actor respectively.
5. R ESULTS ON I MAGE C LASSIFICATION
In this section, we present experimental results of image classiﬁcation tasks.
Test accuracy obtained by different schedulers are shown in Table 1.
It can be observed that the constant learning rate scheduling is insufﬁcient to lead to a satisfactory result and is beaten by all other methods.
Although handdesigned schedulers like polynomial and cosine decay and the gradient-based method Hypergradient can make an improvement over a constant learning rate, learning-based methods can still outperform them.
In particular, GNS achieves much better accuracy on both Fashion-MNIST and CIFAR10 than
SRLS, due to a richer state representation from the utilization of the model’s graph structure.
[Table] [Table]
R ESULTS ON L ANGUAGE U NDERSTANDING
Apart from image classiﬁcation, we also focus on language understanding tasks on eight GLUE benchmark datasets.
Two models with different sizes, RoBERTa-base and RoBERTa-large, are evaluated and detailed results are reported in Table 2.
As expected, adjusting learning rate rather than keeping a constant one contributes to the test performance.
In addition, a consistent performance boost can be observed in GNS over all eight datasets for both base and large models.
Speciﬁcally, it achieves a score of 85.2 and 87.8 on average in RoBERTa-base and -large model respectively, while the best baseline in each scenario only reaches 84.5 and 86.9.
It should be noticed that SRLS performs poorly and even cannot surpass the constant scheduling.
It is likely that the simple state design of SRLS (only using information in the last layer) constrains its ability to adjust the learning rate dynamically, especially for large models like RoBERTa.
Thus, SRLS fails to generate a good scheduling to improve the performance.
We plot learning rate scheduling trajectories for three tasks in Figure
MRPC, RTE, and STS
It can be observed that different mechanisms are deployed for distinct tasks.
The learning rate is
[Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
generally adjusted in a similar trend for both RoBERTa-base and large models within a task.
Speciﬁcally, we can see that the learning rate for MRPC is not modiﬁed aggressively.
For RTE, a decayed scheduling is learned while there is a warmup-like procedure followed by a decreasing tendency for
STS-B. Besides, it should be emphasized that the learning rate does not change monotonically, and there are ups and downs during the training to react to the environment observations dynamically.
5. G ENERALIZATION OF GNS
Due to the utilization of hierarchical GCN as the learning rate scheduler described in Section 4.3, our GNS can be transferred to different but similar graphs, e.g., ResNet and VGG (Conneau et al., 2016) structures on CIFAR10.
Thus, we conduct an experiment to evaluate the generalization of the learned scheduler.
We select the policy network trained under the environment of ResNet-34, and apply it directly to the task with VGG-16.
For all other baselines, we simply adopt the best hyperparameter conﬁguration of ResNet-34 to train the VGG-16 network.
Results are shown in the second column of Table 3.
We can see that GNS can still achieve the best performance on the VGG model while the hyperparameter settings of other methods are not suitable for the VGG model.
It indicates that GNS is able to learn the fundamental rule to conduct message passing between nodes and is generalizable to various graphs.
We also evaluate the generalization of GNS on language tasks from RoBERTa-base model to large model in Appendix E.
The excellent transferred performance also explains why only a few episodes are needed when ﬁne-tuning the agent under RoBERTa-large models.
In addition, GNS can be utilized to the same network structure but on different datasets.
[Table]
Speciﬁcally
[Table]
we apply the graph scheduler trained on ResNet-34 for CIFAR10 to training a ResNetnetwork for CIFAR100 and ImageNet.
It can be observed in Table 3 that GNS achieves 72.3% on test accuracy and outperforms 71.2% of the best baseline scheduler with a cosine decay on CIFAR100.
On ImageNet, transferred performance of GNS is 73.1% and is still better than the best baseline of 71.6%.
Such a performance further validates that GNS is generalizable to different datasets.
5. R UNNING TIME COMPARISON
Compared with SRLS in which validation is required in every decision step, GNS can reduce a considerable amount of time with the validation probability p = 0 .
We consider training RoBERTabase in this section.
For instance, when running on MRPC for 5 epochs, we need to make 58 decisions with the number of network updates K = 10 .
The average time of one episode with one NVIDIA 1080Ti GPU of SRLS is 405s while GNS only takes 259s, which decreases the original cost by 30% .
For hand-designed rules, there are 40 hyperparameter settings shown in Appendix B to be considered for RoBERTa ﬁne-tuning while it takes 30 episodes to train the learned scheduler.
The time for running one setting is 223s on average.
Thus, a thorough grid search among these 40 settings consumes approximately 2.5 hours.
On the contrary, GNS takes 2.2 hours and even achieves better results.
Furthermore, ﬁne-tuning the scheduler on RoBERTa-large from the base model only requires 10 episodes and costs much less time than grid search in all conﬁgurations.
5. A BLATION S TUDY
In this part, an ablation study is carried out to corroborate that the design of state representations and reward collection play an important role in the GNS framework.
[Table]
State representations with graph networks.
To show advantages in encoding the state by a graph neural network, we compare our GNS with a variant called concatenated-state scheduler (CSS), where features of all nodes in the state graph are concatenated to form a super representation and then fed into subsequent neural networks for controlling the learning rate.
Compared with the basline SRLS that only uses the ﬁnal layer state, CSS uses all layers but without using GCN to exploit correlation between layers and cannot be applied to different target models.
We can observe from results in Table 4 that using a concatenated state slightly outperforms SRLS but is still insufﬁcient to render as a competitive scheduler as GNS, which shows the importance of encoding by graph neural networks.
Reward collection with a probability.
We also compare the difference between collecting a reward signal at every decision step (the variant is denoted as full-reward scheduler, FRS) and collecting it with a probability.
We report the evaluation metric on the test set in Table 4.
We can see that an instant reward cannot necessarily improve the ﬁnal performance of the learned scheduler.
It is likely that the performance does not improve monotonically and typically oscillates during the training.
For example, in two consecutive decision steps, the latter one gives us a worse validation result, leading to a negative reward computed by Eq. 2 and indicating that the action is not good.
However, this information might be misleading and impede agent training.
Thus, to alleviate such a phenomenon, we design a stochastic reward collection procedure, expecting to obtain more useful signals.
C ONCLUSION
This paper proposes a novel graph-network-based scheduler to control the learning rate dynamically by reinforcement learning.
We are the ﬁrst to take into consideration the structure of the target optimization problem with neural networks, and hereby leverage message passing via the graph to encode the state for predicting the action.
Comprehensive experiments have shown that GNS can consistently achieve performance boosts on benchmark datasets in both image classiﬁcation and language understanding.
However, in this paper we only focus on problems with a moderate time horizon.
How to efﬁciently train a learning rate scheduler for a problem with hundreds of thousands optimization steps such as BERT pre-training with techniques like early stopping and successive halving can be a potential future direction.
E THICS S TATEMENT
To the best of our knowledge, there are no ethical issues with this paper.
R EPRODUCIBILITY S TATEMENT
To ensure reproducibility of our experiments we have provided comprehensive details in the main paper as well as in the appendix.
Speciﬁcally we introduce the experimental settings in Section 5. including selected datasets and implementation details of all compared methods.
Hyperparameter conﬁgurations and tuning details are presented in Appendix B. Moreover, we have also uploaded our source codes as supplementary materials for reference A CKNOWLEDGEMENT
This work is supported by NSF under IIS-2008173, IIS-2048280 and by Army Research Laboratory under agreement number W911NF-20-2-0158.
A GNS A LGORITHM
Input: Value network parameterized by φ , action network parameterized by ϕ , # updates T , decision interval K , prior learning rate distribution D α 1: for episode = 1 , 2 , . . .
do 2: t ← 1 3:
Sample an initial α t from the prior distribution D α 4: for i = 1 to T do 5:
Conduct one optimization step by w t +1 =
w t − α t · g t with α t 6: if i mod K == 0 then 7:
Obtain the graph state representation s t 8:
Sample an action a t and modify the learning rate by α t +1 =
α t · (1 + a t ) 9:
Conduct validation as described in Section 4.
10: t ← t + 1 11: end if 12: end for 13:
Calculate rewards by Eq.
(3) and organize the trajectory as T = { ( s t , a t , r t ) } T//Kt = 14:
Update φ and ϕ using T based on PPO algorithm 15: end for
B A DDITIONAL I MPLEMENTATION D ETAILS
In this section, additional implementation details are provided for reproducing experimental results.
B. S EARCH SPACE FOR HAND - DESIGNED SCHEDULERS
As mentioned in Section 4.
[Table]
we need to conduct hyperparameter optimization on hand-designed learning rate schedulers.
The search spaces for image classiﬁcation and language understanding tasks are presented in Table 5 and 6 respectively, which are selected based on previous work Choi et al.
Liu et al.
; Devlin et al.
[Table] [Table]
B. E XPERIMENTAL SETTINGS ON IMAGE CLASSIFICATION
We present a detailed hyperparameter conﬁguration for two image classiﬁcation tasks in Table
[Table] [Table]
B. E XPERIMENTAL SETTINGS ON GLUE
We further elaborate on details of GLUE benchmarking.
All RoBERTa models in this paper are implemented by Hugging Face 1 Wolf et al.
and pre-trained models are obtained from the corresponding model hub 2 .
We use suggested values of batch size and the number of updates in Liu et al.
for all eight tasks, as shown in Table 8.
In addition, the decision interval and the metric for reward computation of each dataset are provided.
All tasks are optimized by AdamW with the weight decay of 0 . 1 .
1 https:/
[Table]
B. H YPERGRADIENT
We present the details of Hypergradient in Algorithm 2.
We still have two tunable hyperparameters for
Hypergradient, α 0 and β .
The search space for image classiﬁcation is: α 0 ∈
[10 − 1 , 10 − 2 , 10 − 3 ] and β ∈ [10 − 7 , 10 − 8 , 10 − 9 ] .
The search space for language understanding is: α 0 ∈ [1 ×
10 − 5 , 2 × 10 − 5 ] and β ∈ [10 − 7 , 10 − 8 , 10 − 9 ] .
Algorithm 2 Adam with Hypergradient descent (Adam-HD)
Require: : α 0 : initial learning rate
Require: : β : hypergradient learning rate
Require: β 1 , β 2 ∈ [0 , 1) : decay rates for Adam t, m 0 , v 0 , ∇ α u 0 ← 0 , 0 , 0 ,(cid:46) Initialization
Update rule: m t ← β 1 m t − 1 + (1 − β 1 ) g t (cid:46) 1st mom.
estimate v t ← β 2 v t − 1 + (1 − β 2 ) g 2 t (cid:46) 2nd mom.
estimate (cid:98) m t ← m t / (1 − β t 1 ) (cid:46)
Bias correction (cid:98) v t ← v t / (1 − β t 2 ) (cid:46)
Bias correction u t ← − α t (cid:98) m t / ( √ (cid:98)
[Table]
B. F URTHER COMPARISON OF SRLS AND GNS
Due to the property of shared parameters in graph neural networks, our GNS does not introduce many extra parameters and it takes similar time to train GNS compared with training the original policy network in SRLS.
On the other hand, although the efﬁcient reward collection can be applied to SRLS as well, there is still an obvious performance gap between GNS and SRLS, as shown in Table 9.
[Table]
C G RAPH R EPRESENTATION OF R O BERT A
We present a complete graph representation of RoBERTa-base model in Figure 4 for clarity.
There are 12 repeated blocks in total in RoBERTa-base.
[Table] [Table]
/QgaPusEoi04QyM17f06wh0m1xABQ2hRT4a8Rwf0ojk21MgYNO2ck3dGoLPEnZBcYfsxe1j6psWm/VlvhTz2IUAumdY114mwkTCFgksYZOqxhojxHutAzdCA+aAbyeifAd0zSou2Q2UqQDpSf08kzNe673um02fY1dPeUPzPq8XYPmskIohihICPF7VjSTGkw3BoSyjgKPuGMK6EuZXyLlOMo4kwY0Jwp1+eJeWjvHucd67cXOGEjJEmO2SXHBCXnJICuSRFUiKc3JMn8kJerQfr2Xqz3setKWsykyV/YH38ADPvn6M
[Figure]
</latexit
[Figure]
D D ETAILS OF PPO
As mentioned in Section 3.5, we adopt Proximal Policy Optimization (PPO)
Schulman et al.
(2017) in the actor-critic style to train our scheduling policy.
The key objective function of the policy network is deﬁned as
[Equation]
where ρ t ( ϕ ) is the probability ratio and ˆ A t is the estimated advantage estimator.
ˆ A t is computed by generalized advantage estimation (GAE) Schulman et al.
(2015) taking the form of
[Equation]
(9) where δ t =
V targ t −
[Equation]
=
[Equation]
+ γV
[Equation]
− V
[Equation]
s t ) 3 .
This objective can further be augmented by adding an entropy bonus to ensure sufﬁcient exploration, as suggested in Schulman et al.
[Equation]
Then ϕ can be optimized by maximizing Eq.
equation 10.
Besides, a learned state-value function
V φ ( · ) is leveraged to reduce variance in advantage estimators and can be trained by minimizing the following squared-error loss:
[Equation]
where V targ t =
r t + γV
( s t +1 ) .
Here, we use default hyperparameter settings suggested in Schulman et al.
(2017): γ = 0 .
99 , λ = 0 .
95 , (cid:15) = 0 .
2 , c = 0 .
E A DDITIONAL E XPERIMENTAL R ESULTS E. E XPERIMENTS WITH A LONGER HORIZON
We conduct an experiment to train the target model for a longer time horizon.
More concretely,
RoBERTa-base is selected and the number of updates is double for each task as the original one shown in Table 8.
We choose the polynomial learning rate scheduler as a representative example.
As can be observed in Table 10, training longer can only slightly improve the performance.
However, it still cannot match GNS, showing the efﬁciency and effectiveness of our proposed method.
E. G ENERALIZATION ON LANGUAGE TASKS
We also evaluate the generalization of GNS on language tasks, by transferring GNS trained on
RoBERTa-base model to RoBERTa-large model.
Results are reported in Table 11.
We can see that
GNS can still achieve relative satisfactory performance on large models and outperform the constant learning rate, even only trained on RoBERTa-base.
This transferred performance also explains why only a few episodes are needed when ﬁne-tuning the agent under RoBERTa-large models
[Table] [Table]
SGD V . S . A DAM
As many papers also selected SGD as the standard optimizer for image classiﬁcation tasks
[Table]
we carry out an additional experiment to train a ResNet-34 on CIFAR10 with SGD.
As shown in Table
[Table]
we can observe that Adam offers comparable and slightly better performance than SGD, especially for the constant learning rate scheduling.
It is consistent with conclusions in (Choi et al., 2019;
Sivaprasad et al., 2020; Xiong et al., 2020) that Adam can serve as the most practical solution in training neural networks.
Besides
[Table]
GNS is still the best among all scheduling mechanisms with SGD, showing the effectiveness of our proposed method.
[Table]
E. S ENSITIVITY OF HYPERPARAMETERS
[Table]
Furthermore
[Table]
we investigate the impact of two hyperparameters, the learning rate update interval
K and the validation probability p .
Note that for CIFAR
[Table]
the real interval should multiply by 50, i.e., K = 1 denotes 50 training steps of CIFAR10.
Two tasks on CIFAR10 and RTE are selected respectively from image classiﬁcation and language understanding.
It can be seen in Table 13a that when the learning rate is modiﬁed aggressively with K = 1 , 5 or infrequently with K = 20 , the scheduling would degrade the optimization performance, e
[Table]
g., from 94.3 ( K = 10 ) to 92.5 ( K = 1 ).
As to the validation probability
[Table]
the performance of GNS is not very sensitive to different values of p in Table 13b.
Together with results of FRS in Table 4, we set p to a relatively small value 0 .
2 in all experiments without much loss of performance while such a p can reduce the time cost notably as discussed in Section 5.5.
E. I MPACT OF MULTIPLE FACTORS
In this part, we investigate multiple factors such as the features, the encoding mechanism, and the hyperparameter.
These factors affect the encoding and the most straightforward way to evaluate the encoding quality is to compare the performance of downstream tasks.
In our paper the metric is just the performance of the target network trained with the neural scheduler.
To investigate the encoding quality, we consider a list of schedulers with different encoding capability:
• SRLS • SRLS + ALS (All Layer’s States) • SRLS + ALS
+ larger network, where the network is chosen from a 3-layer MLP with hidden size 32 (deeper), and a 2-layer MLP with hidden size 64 (wider).
• SRLS + ALS + larger network + tuned update interval K, K is chosen from [1 , 5 , 10 , 20] .
• GNS • GNS + GAT (Veliˇckovi´c et al., 2017)
Experiments are carried out on two datasets, CIFAR10 (ResNet34) and RTE (RoBERTa-base), and results are shown in Table 14.
We can see that by using a super concatenated vector as the input, the performance is improved slightly, from 93.6 to 93.8 on CIFAR10 and 70.8 to 72.2 on RTE.
A larger network also contributes to the performance gain to some degree but the improvement is limited, compared with the result of GNS, where the encoding scheme is signiﬁcantly different from MLP networks.
Note that performance with a tuned K remains the same and it indicates that our default update interval works the best.
Moreover, after switching to a more powerful graph encoder, GAT, the performance can be further improved, in particular from 94.3 to 94.5 on CIFAR10.
To sum up, leveraging more information (SRLS -> SRLS + ALS) does enhance the scheduler.
Using a better network architecture (a larger network) and hyperparameter conﬁguration also helps but encoding with merely MLPs is not sufﬁcient to make full use of all intermediate layers’ information.
By representing the target network as a directed graph and adopting an encoding scheme of GCN, the encoder quality improves and the corresponding scheduler is much better.
With a stronger encoder
GAT, the ﬁnal performance can be further improved, which validates that the encoding quality matters for the learned scheduler.
[Table]
E. C OMPARISON WITH THE ORIGINAL R O BERT A RESULTS
To make our results more convincing, we provide a comparison with the original results in Liu et al.
(2019) as well as the standard practice (a warmup procedure followed by a linear decay) in Table 15.
It can be observed that GNS performs consistently better on the dev set as well.
we can tell that our proposed GNS still outperforms this standard method, indicating that the improvements are not due to the suboptimal learning rate schedules, but from our better learned policy.
[Table]
E. C OMPARISON WITH CMA-ES
In this part
[Table]
we compare our GNS with a more effective hyperparameter tuning approach.
More concretely
[Table]
we leverage CMA-ES (Hansen, 2016) as another baseline, to search for the learning rate schedule.
We choose a cosine decay method which leads to the best-performance results for almost all the tasks as the base scheduler when searching with CMA-ES.
The corresponding search space is then expanded as shown in the following tables:
[Table]
Tuning range
[Equation]
[Figure]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
50 evaluations are conducted in total for CMA-ES and we also train GNS for 50 episodes for a fair comparison.
Results are demonstrated in Figure 5.
We can see that CMA-ES can ﬁnd a satisfactory hyperparameter conﬁguration quickly in the beginning
[Figure]
but our GNS can reach a better ﬁnal performance.
When we consider the whole search trajectory, it is still safe to say that the proposed GNS is a relatively efﬁcient and effective method even compared with the more powerful baseline CMA-ES
[Figure]
E. A BLATION STUDY ON THE CHOICE OF RAW FEATURES
In order to characterize the training dynamics, we picked some common features adopted in (Xu et al., 2017; 2019), including the mean and variance of weights, biases, and the mean and variance of absolute gradient values, as well as global dynamics of the previous learning rate and the average training loss.
We can roughly group them into three categories: parameter state, gradient state, and global state.
To investigate the inﬂuence of different features, we mask one of these three types of features and only use the remaining two as the raw input each time.
Then we use different input features to train GNS on CIFAR10 and RTE.
Detailed results are presented in the following table.
It can be observed that all three categories of features beneﬁt the quality of GNS.
Parameter and gradient information is more important than the global state since the global state can be reﬂected by the reward and the action.
This ablation study validates our choice of these features.
In the future, it is also interesting to integrate more useful features to further improve the performance, such as the accumulated momentums which are used in adaptive optimizers like Adam.
[Table]
C OMPARISON OF OTHER FORMS OF ACTION
We have tried two other forms of action, adding a residual and directly predicting the learning rate.
For adding a residual, since the learning rate for different tasks might have different magnitudes, we have to design some additional hyperparameters elaborately for each task to constrain the output residual to a speciﬁc range.
As to predicting the learning rate, it is very unstable and usually breaks the optimization.
We also implemented the action of adding a residual by constraining the output to the range [ − 10 − 5 , 10 − 5 ] , and the action of predicting the learning rate by replacing the activation function with Sigmoid on CIFAR10, and the corresponding results are shown as follows:
[Table]
We can see that adding a residual is insufﬁcient to beat our original design of multiplying by a factor.
For predicting the learning rate, it is hard to obtain a useful scheduler and the training of
CIFAR10 failed and was stuck with only 10% accuracy.
However, it is possible that by ﬁne-tuning the hyperparameters for these two forms of action, the obtained scheduler will work.
But multiplying by a factor is the most stable action and needs little tuning effort and that’s why we choose this action.
E. E VALUATION ON MLP S
We include an evaluation of training a 3 layer MLP on Fashion MNIST with the hidden size ofand we report the performance of different optimizers in Table 20.
It can be seen that GNS is also applicable to such a simple scenario, with the best test accuracy of 89.8.
[Table]
E. T RAINING LOSS CURVE
We have included the training loss curve for Fashion MNIST and CIFAR10 in Figure 6.
Since for the
GLUE benchmark tasks, ﬁnal performance is usually the focus in the literature and the loss and the evaluation metric do not necessarily correlate, we do not show their loss curves
[Figure]