Learning to Schedule Learning rate with Graph Neural Networks
Authors: Yuanhao Xiong
Li-Cheng Lan
Xiangning Chen
Ruochen Wang
Cho-Jui Hsieh
Venue: ICLR 2022 Poster (2022)
Keywords: learning rate scheduling
graph neural networks

Abstract: 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.

Published as a conference paper at ICLR 2022
LEARNING TO SCHEDULE LEARNING RATE WITH GRAPH NEURAL NETWORKS
Yuanhao Xiong, Li-Cheng Lan, Xiangning Chen, Ruochen Wang, Cho-Jui Hsich Department of Computer Science, UCLA (yhxiong, Iclan, xiangning, c Jas. ucla.eda ruocwangeucla.es
ABSTRACT
Recent decades have witnessed great development of stochastic optimization in training deep neural networks. Learning rate scheduling is of the most im- portant factors that influence the performance of stochastic optimizers like Adam. Traditional methods seek to find relatively proper scheduling among limited number of pre-defined rules and might not accommodate a particular target problem. Instead, we propose a novel Traph-Network-baee -aaedd Scheduler (GNS), aiming at learning a specific scheduling mechanism without restrictions to existing principles. By constructing directed graph for the underlying neural network of the target problem, GNS encodes current dynamics with a graph message passing network and trains agent to control the learning rate accordingly via reinforcement learn- ing. 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-MNiI 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 /attps:///qitiiii com/ 2yt////yh////SS
and CIFARIO for image classification,
INTRODUCTION
Stochastic optimization has achieved great successes in training deep learning models (Robbins Monro 195I), due to the tremendous data size and non-convexity neural networks (Xiong
jet | 2020). However, in contrast to classical and hyperparameter-fres free optimizers, performance of
widely used stochastic algorithms such as SGD and Adam is significantly dependent on the choice of
hyperparameters including batch size, learning rate and momentum coefficients (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 (Cooper al. 202T Xu al., (2019). A fixed and static learning rate throughout the whole training is insufficient to attain an optimal model. Thus, various scheduling algorithms (Gotmare) 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 constant learning rate, in terms of convergence rate well as the final performance (Senior et al. 2013 Gotmare et al. 2018: Loshchilov & Hutter 2016).
As the learning rate scheduling is of vital importance the model performance, how tune it for 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...oolyooonm polynomial or cosine decay, warmup and restart) based on empirical studies and domain knowledge (Xu et 2017, Gotmare et al.) 2018), and apply traditional hyperparameter optimization methods (Li et al. 2017 Snoek et al.J 20122 Bergstra & Bengio (2012) tune the parameters of the rule. However, searching within those pre- drenne principles is restricted since the best learning rate scheduling for particular training problem may not exactly follow any existing rule.
To overcome this restriction, finding an optimal scheduler can be viewed as learaing-to-learn problem (Andrychowicz et al. 2016 Chen et al. 2020), where empirical scheduling principles

---

Published as a conference paper at ICLR 2022
replaced by data-driven neural network alternative. For instance, attempts such as Xu et al. (0199) and Xu et al. (0017) model the learning rate scheduling as a sequential decision process where the learning rate is modified in every decision step, and train policy network with reinforcement learning algorithms. However, these methods sacrifice rich state representations in order to maintain generalization. This implies that woobemmddependett information, such as the architecture of neural networks to be optimized, and the state of intermediate neurons not fully exploited. Therefore, state-of-the-an algorithms such Chen et al. (2020) only represent the current state by the statistics of final layer neurons. A straightforward way to leverage such information is concatenate states in every layer as super vector. However, it would lead to the explosion of the input dimension and fail work in models with different sizes. Therefore is nontrivial 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 -aaed Scheduler (GNS) to overcome the above-mentionel 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 ZZang 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 definition for stabling training and an efficient reward signal collection procedure are designed for our GNS framework.
Our main contributions are summarized as following:
We develop the Graph Network-basee 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 classification and language understanding tasks. Also, an efficient 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 GLUE with RoBERTa-base and -large respectively and outper- forms 84.7 and 86.9 from the best existing learning rate schedule. Note that unlike previous learning- to-leamn methods that mostly are applied to toy examples such as CIFAR-I0 we are the first 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 10 to CIFAR 100) and model structures (from ResNet to VGG, and from RoBERTa-base to RoBERTa-large.
2 RELATED WORK
Learning rate scheduling. This topic has been an open problem in stochastic optimization. In Almerda et al.) (1998) and Baydin et al. (2077). learning rate was regarded as a trainable pa- rameter and modified using its gradient with respect the objective function. Sutton (1992) and Schraudolph et al. (206)) considered a meta-descent algorithm to adjust local learning rate with a meta one. In the era of deep learning, 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 large learning rate for infrequently-updatte parameters. Complimentary to these parameter-wis adaptations, some predefined learning rate schedulers are also observed beneficial to training performance (You et al., 2019, Ge et al. 2019 Schaul et al., 2013 Senior et al. (0013). 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 final performance of deep learning models empirically. However, they still hold additional hyperparameters to be tuned. At the same time, these parametric schedules have limited flexibility 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 find 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
2

---

Published as a conference paper at ICLR 2022
applications, such as Atari games (Gu et al., 2016 Hester 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 ling. 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 (0066) Xu et al. (017)) and Xu et al. (219)) leveraged RI agents to adjust the learning rate to reach final performance whereas their selection of state features, reward and action leads to an inefficient training process, and is difficult 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 al. 2018). A graph network, or a graph-to-ge sopph-tt-gae -o-grrhggghga- 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 jet al. 2018 Sanchez-Gonzale et al.] 2020). The most important component in GNN, Message Passing Neural Networks (MPNN) (Gilmer et al., (017). 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 al. 2016) have previously demonstrated promising results on simulation tasks in the field of physics (Battaglia et al. 2016)]
Sanncer--Gonzaler et al. 2020) and chemistry (Gilmer al., 2017 You al. 2018).
3 REINFORCEMENT LEARNING FOR LEARNING RATE SCHEDULING
In this section, we formulate learning rate scheduling as an RL problem, with the goal to find an optimal control policy to maximize the accumulated reward in the current environment which is defined by the target training problem.
3.1 REINFORCEMENT LEARNING
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 formulated by a tuple M = AS, P. R., 7). Given the state set S and action set A, the reward function R is a mapping of S A - E. The discounted return denoted by R, = Σi= r'r where 7 is a discount factor for future rewards. P is the state transition function, which specifies the probability p(5t+11 18t, de) and is usually determined by the environment itself. To encourage more exploration, we focus on a stochastic policy π (alss), which maps states to action probabilities.
3.2 LEARNING RATE SCHEDULING AS A MDP
For most problems involved in training deep neural networks, standard approach for the objective function minimization is stochastic gradient descent on a certain form (Robbins & Monro 1951 Kingma & Ba 2014 Duchi et al. 201T), which sequentially updates the parameters using gradients Specifically, wy represents network parameters to be trained and ag is the global learning rate, gu is defined by the optimizer (e.g.. 9k can be the SGD or Adam direction). For learning rate scheduling, we primarily consider generating a sequence for the global learning rate o'k.
by Wk+1 = W.
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 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 final performance, metrics like loss or accuracy can serve 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 modifications. By training such an agent modeled as a policy network, we can obtain learning rate scheduler to adjust learning rate dynamically.
4
METHOD
4.1
STATE ENCODING WITH GRAPH NEURAL NETWORKS
The selection of features to represent the environment is of vital significance. Previous work such as (Xu al. 2017) and (Xu al. 2019) aims learn generalized agent that can be applicable to any
3

---

Published as a conference paper at ICLR 2022
network architecture, they were only able to utilize statistics of the final layer to describe the state. 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 insufficient to describe the environment.
straightforward way to obtain more thorough representation to utilize dynamics of all layers in neural network and concatenate them into one super vector. However, such representation cannot guarantee promising results as shown in Section [.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 higher correlation than layers that are far away. This concatenation also loses transferrablity to different network structures.
Therefore, to maintain rich state description and generalization in the meanwhile, construct a computational graph to depict the state. We represent a network structure as a directed graph G=,E,V) V). Each node denotes an operator f- w.) (convolution, linear, etc.) parameterized by We which generates an output activation tensor. An edge Euv E E represents the computation flow from node 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 neural network are presented in Figure aa) and (b). For each node with trainable parameters Was the raw feature includes the mean and variance of weights, biases, and absolute gradient values of Wo, 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 E R and the adjacency matrix A E IR N where N is the number of nodes and F is the feature dimension.
coma
adi
E Hit
(fuis W)
(bije e
(
(bujt Virin)
CoTY
CoTe
hdl
lG
(a)
Encoder
Encoder
Encoder
Encoder
Encoder
RAR OR
Inter
Input
Output
ipput
self out
out
Go
Gi
G2
G
Gar
(b)
(c)
Figure 1: (a) A repeated block in ResNet models. (b) repeated block in Transformer models. (c) Illustration of the hierarchical GCN message passing.
Given an input architecture, we describe the state graph G with the feature matrix X and the adjacency matrix A. Then we can feed into an encoder F, 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
H = E(H', A) = (AH'w"),
(I)
where HD is initialized as X, Á is the normalized adjacency matrix with a self-loop in each node, and W' is the parameters of the 1-th layer of GCN.
However, a deep learning model usually adopts hundreds of operations and results in large compu- tation 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 al. 2018), shown in Figure aa) and (b). Instead of passing messages through the whole graph directly, we separate the graph into a number of blocks as (G, = (V,, M++ and each block has group of nodes. M is the number of repeated blocks, and Go stands for the input block (e.g., embedding layers) while GM++ represents the output block (e.g., classification layers). We can then use a shared encoder across different blocks as shown in Figure (Tc) to conduct information propagation. This is implemented by a hierarchical GCN as follows: (i) compute latent features of nodes in each block by Eq. [ (ii) summarize the i-th block as lıg. = AGG (h, /V E V.})); (iii) then use lG; as the input node feature for the next block: (iv) Finally, obtain the graph level information describe the whole network, we just take the aggregation over each node by hG = AGG [ha E V)). Here, we use the mean function aggregation function to allow generalization to different structures.
4

---

Published as a conference paper at ICLR 2022
Compared with 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 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 specific graph structure. In detail, the dimension of concatenated representation varies between different network architectures, for example, RoBERTa- base and -large models. A scheduler learned from such concatenated states fails to work in another environment due the input mismatch. On the contrary, by stacking the same encoder along the scheduler trained on one network to the
depth dimension, we can transfer our graph-network-baat new environment of another different network directly with slight fine-tuning.
4.2 ENVIRONMENT 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 defined to facilitate an efficient and generalizable agent.
Action. We adopt the action as scaling the previous learning rate by factor (Xu al. 2019) considering the wide range of a. Through that operation, the learning rate can be changed stably and consistently. More concretely, after each action execution, 0x++1 Ot," (1 + oz).
State transition. As mentioned in Section 3.2 the update rule of gradient descent for K batches to get the next state st+1-
once we modify the learning rate a++1, we execute
Reward design. With the goal of the agent to achieve higher final performance, there are 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 indicator for task specific metrics which we are ultimately concerned with, like accuracy, F1 score and correlation (Venter al. 2021). Thus, we choose to such task-related metrics directly instead of the loss. On the other hand, since defining 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:
immediate reward. However, loss might not be ideal
Pe+1 Pr,
where P. represents the validation result under the given metric the t-th decision step. Whereas, computing a reward in Eq. iiddicates that we have to conduct validation procedure every decision step, which is a non-trivial burden. In order to make training more efficient, turn to reduce the frequency of validation with a probability p. i.e., the validation is only carried out the number randomly sampled from Uni(O. 1) is less than p. Specifically, suppose under principle, calculate the validation metric Pe and Pt at the decision step ( and f2 respectively. The reward signal for each step can be derived by linearly distributing the difference as follows:
Pto Ps 4
Vi E (2)
(3)
4.3 GRAPH NETWORK-BASER SCHEDULER
With illustration of the environment design, we can then move to the general framework of our Graph Network-based based Scheduler (GNS). As shown in Figure at each decision step, GNS takes the observation 8 as input, encodes s, into /t, sample an action a, based on the policy network, evaluates the value of sc, then execute the action to render the next state 8t+1-
Sampling Evaluation,
Encoding
V(s)) = my (Re)
Execution
Rendering
1
a. ~ ((mm((h).aa))
Oce1 = (1 (1+))))00
8t-1
Figure 2: The transition from 8, to 8t+1-
In the procedure of sampling and evaluation, we use the hidden feature h2 to evaluate the current state and predict the action. Denote m, and my two different MLPs. Specifically, the state value can be computed by (st) = m, (he). Then we calculate the mean of the action and the policy π samples a, from the corresponding normal distribution in Eq. 4 with fixed deviation
04 N(M,,0)), Pt 0.1 (anh(my((aa))
5

---

Published as a conference paper at ICLR 2022
By Eq. 4 constrain that the scale factor (1 az) is within the range of [0.9, 1.1] most of the times in order to train the target network stably (Xu al. 2019).
As to the agent training, we first parameterize the value network V (critic), and policy network ı (actor) by Φ; jp respectively for clarity. Each environment of the target problem has fixed 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 defined
L ((p E,
-,1- +)+ pp((p)
(ap (se) Tia (atls)))
(5)
where pe({p) is the probability ratio and A, 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:
L(Ø) = (,((V((s))) )),
(6)
where 1 = Tt + (s++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 first sample a number of initial learning rates, run each setting for pre steps, compute corresponding V (srpa) and choose the learning rate with the largest value to complete the training.
EXPERIMENTAL RESULTS
To validate the effectiveness of GNS, we evaluate our method on various tasks in image classification and language understanding. and compare with popular learning rate scheduling rules. We further investigate the generalization of GNS on different transfer tasks. In addition, we conduct an ablation study analyze the state representation and reward collection.
5.1 EXPERIMENTAL SETTINGS
In this section, we present our experimental settings. Further details can be found in Appendix B
Image classification. We consider two benchmark datasets in image classification, Fashion- MNIST (Xiao et al. 2017) and CIFARIO (Krizhevsky et al., 2014). These two datasets are first split into the standard training and test sets. Then we randomly sample 10k images for each dataset from the training set to construct a validation set. A ResNet-18 (He et al. (2166) model is trained on Fashion-MNIST while ResNet-34 (He al.) 2016) is trained on CIFAR 10. We Adam (Kingmal 2014) with the batch size of 128 for 200 epochs to train these two image classification tasks.
& Ba
Language understanding, For language understanding, we conduct experiments on GLUE (Wang ct al. 2019). a benchmark consisting cight sentence- or sentence- peeence--iirrree tasks. They are divided into training, validation and test sets and we have access to ground truth labels of test sets. On each dataset, we fine-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 official 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 (017)) optimizer adopted 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-basse based method Hypergradient (Baydin et al. 2017), and one learning- method, SRLS (Xu et al. 2019) compared with our GNS framework as baselines. For traditional schedulers

---

Published as a conference paper at ICLR 2022
with the warmup process, we use the following equation to summarize them:
Tram Omax:
warmup;
[Omax Ominn)
(1 foor Turn
+ Olminn
polynomial;
Ott
(a'max aimin)
1 + COS(
for Tomm iT) + amin
cosine;
Thway Twen
Ommin:
constant.
In detail, O'max and Ommin = pammmx are the maximum and minimum learning rate respectively during the training. Tota is the number of total training steps, Tarm 1/1 Total is the number of warmup steps and Taeeay = 772 is the number of decay steps. We tune all hyperparameter combinations within the search space in Appendix B and select the configuration 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. (0117) and Xu et al. (2019) respectively.
GNS implementation. To encode raw observations into latent features, we use 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 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 fine-tuning the agent learned from RoBERTa-base by an Adam optimizer of the learning rate 0.005 and 0.001 for the critic and actor respectively.
5.2 RESULTS ON IMAGE CLASSIFICATION
In this section, we present experimental results of image classification tasks. Test accuracy obtained by different schedulers are shown T can be observed that constant learning rate scheduling is insufficient to lead to satisfactory result and is beaten by all other methods. Although hand- designed schedulers like polynomial and cosine decay and the gradient- method Hypergradient can make an improvement over a constant learning rate, dearning-basee methods can still outperform them. In particular, GNS achieves much better accuracy on both Mishion-MNIST and CIFARI0 than SRLS, due to richer state representation from the utilization of the model's 's graph structure.
Table 1: Test accuracy on Fashion-MNIST and CIFARIO of different schedulers
Scheduler
Fashion-MNIST
CIFAR10
Constant
93.0 + 0.2
93.0 1. 0.4
Polynomial
93.4 + 0.1
93.5 0.3
Cosine
93.4 0.1
93.6 + 0.4
Hypergradient
93.5 - 0.2
93.7 0.6
SRLS
93.6 + 0.3
93.6 + 0.6
GNS
94.6 - 0.2
94.3 1 0.5
5.3 RESULTS ON LANGUAGE UNDERSTANDING
Apart from image classification, we also focus on language understanding tasks on eight GLUE benchmark datasets. Two models with different sizes, RoBERTa-bor base and RoBERTa-large, evaluated and detailed results are reported Table 2 As expected, adjusting learning rate rather than keeping constant one contributes to the test performance. In addition, consistent performance boost can be observed GNS over all eight datasets for both base and large models. Specifically, it achieves a score of 85.2 and 87.8 on average in RoBERTa-bars and -large model respectively. while the best baseline in each scenario only reaches 84.5 and 86.9. 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 adjust the learning rate dynamically. especially for large models like RoBERTa. Thus. SRLS fails to generate good scheduling improve the performance.
We plot learning rate scheduling trajectories for three tasks in Figure 3 MRPC, RTE, and STS-B. It can be observed that different mechanisms deployed for distinct tasks. The learning rate is
7

---

Published as a conference paper at ICLR 2022
Table 2: Test performance of RoBERTa on GLUE benchmarking of different schedulers.
Scheduler
CoLA
SST-2
MRPC
QQP
STS-B
MNLI
QNLI
RTE
Avg.
Base model:
Constant
60.5 = 0.4
95.1 0.2
86.9 = 0.6
89.0
88.6 0.2
87.0
92.9 -
73.0
84.1
Polynomial
61.2 0.3
95.3
87.2
8.1
89.3
93.0 0.2
73.1
01.1::00 0.4
84.4
Cosine
61.2 0.3
95.4
0.
89.2
894
871
930
73.5
84.5
Hypergradient
61.0 0.4
95.3
87.0 0.4
8.2
8.2
87.1
0.1
73.4
4: 0.5
84.4
SRLS
58.4 1.2
95.3
86.8
0.5
88.4 0.1
89.4 0.3
86.7 1: 0.2
92.3
70.8
0.7
83.5
GNS
63.3 0.5
95.8
0.2
87.7
0.5
89.6
1 0.0
89.7
1 0.1
87.7
1 0.2
93.2
1:
0.0
74.4 1: 0.3
85.2
Large model:
Constant
63.1
87.7
89.2
80.6
Polynomial
64.2 4- 0.8
95.9 + 0.4
88.4
+ 0.7
89.2+
91.1 + 0.2
90.0 + 0.0
94.5 + 0.1
82.0
$2.0--10 0.2
86.9
Cosine
89.3
91.2 0.2
Hypergradient
63.3 + 0.7
95.7
+ 0.4
88.2
0.8
89.2 0.1
91.0 0.3
89.9 + 0.3
94.4
0.2
81.4
0.4
86.6
SRLS
62.3
95.9
91.2
79.8
GNS
65.8 + 0.7
96.7
0.1
89.0
0.9
89.7 + 0.1
92.1 + 0.2
90.4
0.1
94.7
0.1
84.0
0.2
87.8
lrge
large
125
910
1001
1712
1.80
192
Stepp
(a) MRPC
(b) RTE
(c) STS-B
Figure 3: Learning rate trajectories generated by GNS for three tasks.
generally adjusted in similar trend for both RoBERTa-bees base and large models within task. Specifi- cally, we can see that the learning rate for MRPC is not modified aggressively. For RTE, a decayed scheduling is learned while there is warmup-liks procedure followed by decreasing tendency for STS-B. B. Besides, should be emphasized that the learning rate does not change monotonically, and there are ups and downs during the training react to the environment observations dynamically.
5.4 GENERALIZATION GNS
Due to the utilization of hierarchical GCN as the learning rate scheduler described in Section 4.3 our GNS can be transferred different but similar graphs, e.g., ResNet and VGG (Conneau al. 2016) structures on CIFARIO 10. 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 directly to the task with VGG-16. For all other baselines, we simply adopt the best hyperparameter configuration of ResNet-34 to train the VGG-16 network. Results are shown in the second column of 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 language tasks from RoBERTa-base model to large model in Appendix [ The excellent transferred performance also explains why only few episodes are needed when fine-tuning the agent under RoBERTa Ta-lagee models. addition, GNS can utilized to the same network structure but on different datasets.
Table 3: Transferred performance GNS.
Scheduler
ResNet - VGG
CIFARIO - CIFAR100 CIFARIO ImageNet
Constant
91.4
69,3 : 0.3
Polynomial
91.8 0.3
70.5 : 0.4
70.3
Cosine
91.7 = 0.4
71.2 0.2
70.5
Hypergrudhemt
70,9 0.3
SRLS
91.7 = 0.3
71.0 - 0.6
71.6
GNSS
92.5 0.4
72.3 +: 0.3
73.1
Specifically, we apply the graph scheduler trained on ResNet-34 for CIFAR10 to training ResNet-34 network for CIFAR 100 and ImageNet. It can be observed in Table 3 that GNS achieves 72.3% test
8

---

Published as a conference paper at ICLR 2022
accuracy and outperforms 71.2% of the best baseline scheduler with cosine decay CIFAR100. On ImageNet, transferred performance of GNS 73.1 1% and is still better than the best baseline of 71.6%. Such performance further validates that GNS is generalizable to different datasets.
5.5 RUNNING 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.2. We consider training RoBERTa- base in this section. For instance, when running on MRPC for 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 Dopendix(BJto B be considered for RoBERTa fine-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, fine-tuning the scheduler on RoBERTa-large from the base model only requires 10 episodes and costs much less time than grid search in all configurations.
5.6 ABLATION STUDY
In this part, an ablation study is carried out to corroborate that the design of state representations and reward collection play important role in the GNS framework.
Table 4: Test performance of three scheduler variants for RoBERTa-base.
Scheduler
CoLA
SST-2
MRPC
QQP
STS-B
MNLI
QNLI
RTE
Avg.
SRLS
58.4 + 1.2
95.3
0.3
86.8
0.5
88.4
+ 0.1
89,4 0.3
86.7 + 0.2
92.3 + 0.2
70.8
0.7
83.5
CSS
60.0 0.5
94.9
0.1
87.1
+ 0.3
89.2
: 0.2
88.7
0.3
86.4 = 0.1
0..
72.2
= 0.5
83.8
FRS
63.5 0.3
95.9
0,3
87.7
1 0.4
89.4
0.0
89.4 0.2
87.6
0.2
93.2 0.1
74.6
0.4
85.2
GNS
63.3 1: 0.5
95.8
0.2
87.7
1: 0.5
89.6
0.0
89.7 1 0.1
87.7 = 0.2
93.2 1 0.0
74.4
0.3
85.2
State representations with graph networks. To show advantages in encoding the state by graph neural network, we compare our GNS with a variant called concatenated-stee scheduler (CSS), where features of all nodes in the state graph 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 final 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 concatenated state slightly outperforms SRLS but is still insufficient to render competitive scheduler as GNS, which shows the importance of encoding by graph neural networks.
Reward collection with probability. We also compare the difference between collecting reward signal at every decision step (the variant denoted as full-reward scheduler, FRS) and collecting it with probability. We report the evaluation metric on the test set in Table 4 We can see that an instant reward cannot necessarily improve the final performance of the learned scheduler. 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 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, alleviate such a phenomenon, we design a stochastic reward collection procedure, expecting to obtain more useful signals.
likely
6 CONCLUSION
This paper proposes a novel graph-network-bee by reinforcement learning. We are the first 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 classification and language understanding. However, in this paper we only focus on problems with a moderate time horizon. How to efficiently train 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.
scheduler to control the learning rate dynamically
9

---

Published as a conference paper at ICLR 2022
ETHICS STATEMENT
To the best our knowledge, there are no ethical issues with this paper.
REPRODUCIBILITY STATEMENT
To ensure reproducibility of our experiments, have provided comprehensive details in the main paper as well as in the appendix. Specifically, we introduce the experimental settings in Section [.T including selected datasets and implementation details of all compared methods. Hyperparameter configurations and tuning details are presented in Appendix B Moreover, we have also uploaded our source codes as supplementary materials for reference.
ACKNOWLEDGEMENT
This work is supported by NSF under IIS-2008173, IIS-2048280 and by Army Research Laboratory under agreement number W911NF-2002-011
REFERENCES
Joshua Achiam. Spinning Up in Deep Reinforcement Learning. 2018.
Luis B Almeida, Thibault Langlois, José D Amaral, and Alexander Plakhov. Parameter adaptation in stochastic optimization. On-Line Learning in Neural Networks, Publications of the Newton Institute, pp. 111-134, 1998.
Marcin Andrychowiez, Misha Denil, Sergio Gomez, Matthew Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. arXiv preprint (et.:6600004474 2016.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, and Koray kavukcuoglu. Interaction networks for learning about objects, relations and physics. In Proceedings of the 30th International Conference Neural Information Processing Systems, pp. 4509-4517, 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 186.00806..1211,,, 2018.
Atilim Gunes Baydin, Robert Cornish, David Martinez Rubio, Mark Schmidt, and Frank Wood. Online learning rate adaptation with hypergradient descent. arXiv preprint urXiv:1703.04733 2017.
James Bergstra and Yoshua Bengio. Random search for hyper-parameter optimization. The Journal of Machine Learning Research, 13())2281-305, 2012.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemystaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:19/2.06680. 2019.
Tianlong Chen, Weiyi Zhang, Zhou Jingyang, Shiyu Chang, Sijia Liu, Lisa Amini, and Zhangyang Wang. Training stronger baselines for learning to optimize. Advances Neural Information Processing Systems, 33, 2020.
Dami Choi, Christopher Shallue, Zachary Nado, Jaehoon Lee, Chris J Maddison, and George E Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint 2019.
Alexis Conneau, Holger Schwenk, Loic Barrault, and Yann Lecun. Very deep convolutional networks for text classification. arXiv preprint arXiv: 01771 2016.
Feder Cooper, Yucheng Lu, and Christopher De Sa. Hyperparameter optimization is deceiving us, and how stop it. arXiv preprint arXiv:2102.03034, 2021.
10

---

Published as a conference paper at ICLR 2022
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoe V Le. Autoaugment: Learning augmentation policies from data. arXiv preprint 1805:::1800005555 2018.
Christian
Daniel, Jonathan Taylor, and Sebastian Nowozin. Learning step size controllers for robust neural network training. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 30, 2016.
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint
2018.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Rong Ge, Sham M Kakade, Rahul Kidambi, and Praneeth Netrapalli. The step decay schedule: A near optimal, geometrically decaying learning rate procedure for least squares. arXiv preprins arXiv: 12838, 2019.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry, In International Conference Machine Learning, pp. 1263-1272. PMLR, 2017.
Akhilesh Gotmare, Nitish Shirish Keskar, Caiming Xiong, and Richard Socher. A closer look at deep learning heuristics: Learning rate restarts, warmup and distillation. arXiv preprint arXiv: 132433 2018.
Shixiang Gu, Timothy Lillicrap, Ilya Sutskever, and Sergey Levine. Continuous deep q-learning with model-based acceleration. In International Conference on Machine Learning, pp. 2829-2838. PMLR, 2016.
Nikolaus Hansen. The ema evolution strategy: tutorial. arXiv preprint arXiv: 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Todd Hester, Matej Vecerik, Olivier Pietquin, Marc Lanctot, Tom Schaul, Bilal Piot, Dan Horgan, John Quan, Andrew Sendonaris, Ian Osband, et al. Deep q-learning from demonstrations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1rXivv1412.660 2014.
Thomas N Kipf and Max Welling, Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv: 1609 0209000907, 2016.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. The cifar-10 dataset. online: http://www. cs. foronto. edu/kriz/cifar html, 55, 2014.
Ke Li and Jitendra Malik. Learning to optimize. arXiv preprint arXiv: 2016.
Lisha Li, Kevin Jamicson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hyperband: A novel bandit- approach to hyperparameter optimization. The Journal of Machine Learning Research, 18(1):6765-68) 2017."
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 190.:.1977.11... 2019.
llya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.0394 2016.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 171.01111 2017.
11

---

Published as a conference paper at ICLR 2022
Purnendu Mishra and Kishor Sarawadekar. Polynomial learning rate policy with warm restart for deep neural network. In TENCON 2019-2019 IEEE Region 10 Conference (TENCON), pp. 2087-2092. IEEE, 2019.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400-407, 1951.
Alvaro Sanchez-Coozalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter Battaglia. Learning to simulate complex physics with graph networks. International Conference on Machine Learning, pp. 8459- 8468. PMLR, 2020.
Tom Schaul, Sixin Zhang, and Yann LeCun. No more pesky learning rates. In International Conference on Machine Learning, pp. 343-351. PMLR, 2013.
Nicol N Schraudolph, Jin Yu, Douglas Aberdeen, et al. Fast online policy gradient learning with smd gain vector adaptation. Advances neural information processing systems, 18:1185, 2006.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High High-dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.0 02438, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06377 2017.
Andrew Senior, Georg Heigold, Marc' aurelio Ranzato, and Ke Yang. An empirical study of learning rates in deep neural networks for speech recognition. In 2013 IEEE international conference acoustics, speech and signal processing, pp. 6724-6728. IEEE, 2013.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwesse, loannis Antonoglou, Veda Panneershelvam,, Marc Lanctot, et al. Mastering the game go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker. Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):344-55, 2017.
Prabhu Teja Sivaprasad, Florian Mai, Thijs Vogels, Martin Jaggi, and François Fleuret. Optimizer benchmarking needs to account for hyperparameter tuning. In International Conference on Machine Learning, pp. 9036-9045, PMLR, 2020.
Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. In Advances in neural information processing systems, pp. 2951-2959, 2012.
Richard S Sutton. Gain adaptation beats least squares. In Proceedings of the 7th Yale workshop adaptive and learning systems, volume 161 161168, 1992.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
T. Tieleman and G. Hinton. Lecture 6.5 :5-RmsProp: Divide the gradient by running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv: : /00. 736.03762, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv:1710.10003 2017.
Arthur EW Venter, Marthinus W Theunissen, and Marelie H Davel. Pre-interpolation loss behavior in neural networks. In Southern African Conference for Artificial Intelligence Research, pp. 296-309,
Springer, 2021.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. 2019. In the Proceedings of ICLR.
12

---

Published as a conference paper at ICLR 2022
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-a natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Pro- cessing: System Demonstrations, pp. 45, Online, October 2020. Association for Computational Linguistics. URL htt.t:::pp.....c.....m www ac web. www.bb/ggggaaooooooooooooooooo http:/hhhoolyyyyoooo/oe emn p-demoss 6
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.077451 2017.
Yuanhao Xiong, Xuanqing Liu, Li-Cheng Lan, Yang You, Si Si, and Cho-Jui Hsieh. How much progress have made in neural network training? new evaluation protocol for benchmarking optimizers. arXiv preprint arXiv:20/0009889 2020.
Chang Xu, Tao Qin, Gang Wang, and Tie- Yan Liu. Reinforcement learning for learning rate control. arXiv preprint arXiv: 11159, 2017.
Zhen Xu, Andrew Dai, Jonas Kemp, and Luke Metz. Learning an adaptive learning rate schedule. arXiv preprint arXiv: 1991199799 2019.
Jiaxuan You, Bowen Liu, Rex Ying. Vijay Pande, and Jure Leskovec. Graph convolutional policy network for goal-directed molecular graph generation. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp.
2018.
Kaichao You, Mingsheng Long, Jianmin Wang, and Michael I Jordan. How does learning rate decay help modern neural networks? arXiv preprint arXiv: 2019.
Chris Zhang, Mengye Ren, and Raquel Urtasun. Graph hypernetworks for neural architecture search. arXiv preprint arXiv: 1810.05749, 2018
Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning, arXiv preprint arXiv:1611.01578 2016.
Barret Zoph, Vijay Vasudevan, Jonathon Shlens, and Quoc V Le. Learning transferable architectures for scalable image recognition. In Proceedings of the IEFE conference on computer vision and pattern recognition, pp. 8697-8710, 2018.
A GNS ALGORITHM
Algorithm 1 Graph Network-based Scheduler.
Input: Value network parameterized by Φ action network parameterized by p, # updates T, decision interval K, prior learning rate distribution Da
1: for episode 1, 1,2,... do
2: f-1
3: Sample an initial a; from the prior distribution D..
4: for = to T do
5: Conduct one optimization step by W:+1 Wi gt with az
6:
i mod K = 0 then
7:
Obtain the graph state representation 8t
8:
Sample an action a, and modify the learning rate by a:-1 = 02 (1 + at)
9:
Conduct validation as described in Section 4.2
+
11:
end if end for
Calculate rewards by Eq. (3) and organize the trajectory as T = {(5E, Cz, rt) Tt) T//K
13:
Update I and ip using 7 based on PPO algorithm
15:
end for
13

---

Published as a conference paper at ICLR 2022
B ADDITIONAL IMPLEMENTATION DETAILS
In this section, additional implementation details are provided for reproducing experimental results.
B. SEARCH SPACE FOR HAND-DESIGNER SCHEDULERS
As mentioned in Section 4.1, need conduct hyperparameter optimization hand-designed learning rate schedulers. The search spaces for image classification and language understanding tasks presented in Table and respectively. which are selected based previous work [Choi et al. (2019): [Liu et al. (0019) Devlin et al. (2018).
Table 5: Search space for image classification
Hyperparameter
Tuning range
[10-1 10-2 10- 110-1.10 10100-2 10031
Omax
112
[0.5. 1.0]
[1. 2]
Table 6: Search space for language understanding
Hyperparameter
Tuning range
[1 10-5, 2 x 10-9) [0, 0.025, 0.05, 0.075, 0.1] [1,
Omma
m
P
B.2 EXPERIMENTAL SETTINGS ON IMAGE CLASSIFICATION
We present a detailed hyperparameter configuration for two image classification tasks in Table 7
Table 7: Hyperparameter configuration for GLUE benchmarking datasets.
Fashion MNIST
CIFARIO
batch size
128
128
# updates T
78200
62600
decision interval K
500
500
Metric
Acc
Acc
B.3 EXPERIMENTAL SETTINGS ON GLUE
We further elaborate details of GLUE benchmarking. All RoBERTa models in this paper implemented by Hugging Face Wolf et al. (2000) and pre-trained models are obtained from the corresponding model hub We suggested values of batch size and the number of updates in Liu jet al. (2019) for all eight tasks, as shown Table In addition, the decision interval and the metric for reward computation of each dataset are provided. All tasks are optimized by Adam W with the weight decay of 0.1.
htwp:/uwwww.ooooooggggggggggoceo http:////ttaa.ooco/ccnceeaaaaaaa htttp:://ugggnnge co/ roberta-base and
http:://uugingiicccco
roberta-large
14

---

Published as a conference paper at ICLR 2022
Table 8: Hyperparameter configuration for GLUE benchmarking datasets.
CoLA
SST-2
MRPC
QQP
STS-B
MNNI
QNLI
RTE
batch size
16
32
32
32
16
32
32
16
updates 7
2680
10525
575
56855
1800
61360
16370
1560
decision interval K
10
100
10
1000
10
1000
100
10
Metric
Corr
Acc
Acc
Acc
Corr²
Acc
Acc
Acc
Mathew's s correlation.
Pearson correlation.
B.4 HYPERGRADIENT
We present the details of Hypergradient in Algorithm 2 We still have two tunable hyperparameters for
Hypergradient, Okg and 3. The search space for image classification is: ag [10 10
110-1.10-2 10-3 and
E [10 7 10 10-8 10 9]. The search space for language understanding is: ag E [1 10-5,2 x 10
and 3 E 10 10 1.-9).
Algorithm 2 Adam with Hypergradient descent (Adam-HD)
Require: ogy: initial learning rate
Require: 3: hypergradient learning rate
Require: 31 E [0, 1): decay rates for Adam
I, mn, No. V. + 0,0, 0,0
D. Initialization
Update rule:
my - 81 m + (1 (1)
D 1st mom. estimate 2nd mom. estimate
vi +- B2 14 +
82) gi
m, + m4/(1-
31)
> Bias correction
0, + v 85)
Bias correction
Ut +- -at m// (V (v +)
Parameter update
Vaut Mi/ /inn/VVit (v
B.5 FURTHER 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 takes similar time train GNS compared with training the original policy network in SRLS. On the other hand, although the efficient reward collection can be applied to SRLS well, there is still obvious performance gap between GNS and SRLS, as shown Table
Table 9: Test accuracy of SRLS and GNS under reward collection with p 0.2.
Scheduler
Fashion-MNIST
CIFAR10
SRLS GNS
93.5 + 0.4
93.4 + 0.3
94.6 0.2
94.3 0.5
C GRAPH REPRESENTATION OF ROBERTA
We present a complete graph representation of RoBERTa-base model in Figure for clarity. There 12 repeated blocks total in RoBERTa-base.
15

---

Published as a conference paper at ICLR 2022
Grepest X
self att
inter
self att
inter
embed
add
add
add
add
clas sifier
self out
self
out
out
Repeated block Gyepeat
Figure 4: graph representation of RoBERTa-base model.
D DETAILS 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 defined as
To(at (t)
L (1))( E, ((tp:())))...ci((((())))))))
pe(())
(8)
T Tp (1)l
where (r())) is the probability ratio and A, is the estimated advantage estimator. Å, is computed by generalized advantage estimation (GAE) Schulman al, (2015) taking the form of
A = 5, + (TA)8 + ....
87-1,
(9)
where 6, Vrtar V(si) ri + +V Vss, This objective further be augmented by
adding an entropy bonus to ensure sufficient exploration, as suggested in Schulman et al. (0017):
L(p) (xp) entropy|xp](st).
(10)
Then to can be optimized by maximizing Eq. equation 10 Besides, learned state- value function (-(-) leveraged to reduce variance in advantage estimators and can trained by minimizing the following squared-error loss:
( = [VV ( |)
where V-arr Ft + (sL+1) Here, we use default hyperparameter settings suggested Schulman al. 2017 Y 0.99, A 0.95, = 0.2, 0.01.
E
ADDITIONAL EXPERIMENTAL RESULTS
E EXPERIMENTS WITH A LONGER HORIZON
We conduct experiment train the target model for longer time horizon. More concretely, RoBERTa-base 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 can be observed in Table 10 training longer can only slightly improve the performance. However, still cannot match GNS, showing the efficiency and effectiveness of our proposed method.
representative example. As
E.2 GENERALIZATION 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 reported in Table We 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 fine-tuning the agent under RoBERTa-larg models
When ((+)) rather than V8(*t), refer the value and it has gradient
16

---

Published as a conference paper at ICLR 2022
Table 10: Comparison with the scheduler of longer trajectory.
Scheduler
CoLA
SST-2
MRPC
QQP
STS-B
MNLI
QNLI
RTE
Avg.
Polynomial
61.2
95.3
87.2
89.2
89.3
87.3
93.0
73.1
84.4
Polynomial (longer)
61.3
95.8
87.2
89.3
89.3
87.4
92.9
73.5
84,6
GNS
63.3
95.8
87.7
89.6
89.7
87.7
93.2
74.4
85.2
Table 11: Comparison between the constant learning rate and transferred GNS for RoBERTa-large.
Scheduler
MRPC
STS-B
RTF
Avg.
Constant 87.7 = 0.7
90.1 = 0.3
86.1 91.1 = 0.2
80.6 + 0.2
GNS (transferred)
87.9 = 0.8
79.5 - 0.2
86.2
E.3 SGD V.S. ADAM
As many papers also selected SGD as the standard optimizer for image classification tasks, we carry out an additional experiment to train a ResNet-34 on CIFARIO with SGD. As shown in Table I2] 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 al., 2020 Xiong al.] 2020) that Adam can
the most practical solution in training neural networks. Besides, GNS is still the best among all scheduling mechanisms with SGD, showing the effectiveness of our proposed method.
Table 12: Test accuracy with SGD and Adam the base optimizer respectively on CIFAR1O.
Scheduler
SGD
Adam
Constant
92.7 + 0.3
93.0 + 0.4
Polynomial
93.5 + 0.2
93.5 + 0.3
Cosine
1..4:..4
93.6 0.4
Hypergradient
93.6 1.61100 0.5
93.7 1 0.6
SRLS
93.6 = 0.4
93.6 0.6
GNS
94.1 1 0.3
94.3 + 0.5
E.4 SENSITIVITY OF HYPERPARAMETERS
Table 13: Sensitivity of hyperparameters.
(a) Performace different K with fixed 0.2.
(b) Performance on different with K fixed at 10.
RTE CIFAR10
p
CIFAR10
RTE
92.5 = 0.6
71.6 1..6 0.5
0.2
94.3
0.5
74.4 + 0.3
93.0 0.4
72.8
0.5
0.6
94.4
+ 0,4
74.4 + 0.5
10
94.3
0.5
74.4 0.3
1.0
94.4
0.4
74.6 + 0.4
20
93.7 = 0.4
74.0 + 0.3
Furthermore, we investigate the impact of two hyperparameters, the learning rate update interval K and the validation probability Note that for CIFAR1O, the real interval should multiply by 50, ie.. K denotes 50 training steps CIFAR10. Two tasks on CIFAR10 and RTE are selected respectively from image classification and language understanding. It can be seen in Table 33a when the learning rate is modified aggressively with K 1, infrequently with K 20. the scheduling would degrade the optimization performance, e.g., from 94.3 (K 10) to 92.5 (K = 1). As to the validation probability, the performance of GNS is not very sensitive to different values of
17

---

Published as a conference paper at ICLR 2022
in Table 13b Together with results of FRS in Table set p to relatively small value 0.2 in all experiments without much loss performance while such p can reduce the time cost notably as discussed in Section 5.5
E.5
IMPACT OF MULTIPLE FACTORS
In this part, investigate multiple factors such 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 paper the metric is just the performance 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 3-layer MLP with hidden size 32 (deeper), and 2-layer MLP with hidden size 64 (wider).
SRLS ALS + larger network + tuned update interval K, K is chosen from [1,5, 10, 20j.
GNS
GNS + GAT VVeličković et al. 2017)
Experiments are carried out on two datasets, CIFARIO (ResNet34) and RTE (RoBERTa-base), and results shown in Table [4 We performance improved slightly, from 93.6 to 93.8 on CIFAR 10 and 70.8 to 72.2 on RTE. A larger network also contributes to the performance gain some degree but the improvement is limited,, compared with the result of GNS, where the encoding scheme significantly different from MLP networks. Note that performance with tuned K remains the same and it indicates that 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 CIFARIO.
that by using super concatenated vector the input, the
To sum up. leveraging more information (SRLS SRLS + ALS) does enhance the scheduler. Using a better network architecture (a larger network) and hyperparameter configuration also helps but encoding with merely MLPs is not sufficient make full 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 final performance can be further improved, which validates that the encoding quality matters for the learned scheduler.
Table 14: Comparison of schedulers with different configurations.
Scheduler
CIFARI0
RTE
SRLS
93.6 0.6
70.8 0.7
SRLS ALS
93,8
72.2 = 0.5
SRLS ALS + larger network
93.9
72.6 0.4
SRLS ALS + larger network tuned K GNS
93.9
72.6 0.4
94.3 + 0.5
74.4 0.3
GNS + GAT
94.5 + 0.4
74.7 0.3
E.6 COMPARISON WITH THE ORIGINAL ROBERTA RESULTS
To make our results more convincing, we provide comparison with the original results in Liu et al. (2019) as well the standard practice warmup procedure followed by a linear decay) in Table 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.
18

---

Published as a conference paper at ICLR 2022
Table 15: Development performance of RoBERTa-large on GLUE of different schedulers.
Scheduler
CoLA
SST-2
MRPC
QQP
STS-B
MNLI
QNLI
RTE
Avg.
Large model:
Constant
68.8
96.1
90.3
92.0
92.1
90.0
94.4
86.2
88.7
Polynomial
69.1
96.3
90.6
92.2
92.4
90.2
94.8
86.5
89.0
Cosine
69.0
96.4
90.5
92.2
92.3
90.1
94.6
86.7
89.0
Hypergradient
68.6
96.3
90.5
92.0
92.1
90.1
94.5
86.5
88.8
SRLS
67.8
95.9
89.7
91.6
92.1
89.9
94.2
86.0
88.4
GNS
70.5
97.1
91.2
92.7
92.9
90.5
94.9
87.4
89.6
Liu et al. (2019)
68.0
96.4
90.9
92.2
92.4
90.2
94.7
86.6
88.9
Standard
68.8
96.3
90.7
92.1
92.3
90.2
94.7
86.5
88.9
E.7 COMPARISON CMA-ES
In this part, we compare our GNS with more effective hyperparameter tuning approach. More concretely, we leverage CMA-E (Hansen 2016) as another baseline, to search for the learning rate schedule. We choose a cosine decay method which leads to the best best- pest-performance results for almost all the tasks the base scheduler when searching with CMA-ES. The corresponding search space is then expanded as shown in the following tables:
Table 16: Expanded search space for image classification
Hyperparameter
Tuning range
LogUniform(10-11 10-1] LogUniform[10) (10-3, 10-11
Omax
122
Uniform 1.0]
Table 17: Expanded search space for language understanding
Hyperparameter
Tuning range
O'ax
LogUniforrm[109 10-4] Uniform[O.0, 0.1]
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 find satisfactory hyperparameter configuration quickly in the beginning, but GNS can reach better final performance. When we consider the whole search trajectory. is still safe to say that the proposed GNS is relatively efficient and effective method even compared with the more powerful baseline CMA-ES.
CIFARIO
RTF
TENS
01.
08.0
Nenderol Enihalee
Mander Euleaiik
CIFAR10
(b) RTE
Figure 5: Search performance of GNS and CMA-ES.
19

---

Published as a conference paper at ICLR 2022
E.8 ABLATION STUDY ON THE CHOICE OF RAW FEATURES
In order to characterize the training dynamics, we picked some common features adopted in (Xu ct 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 roughly group them into three categories: parameter state, gradient state, and global state. To investigate the influence of different features, we mask one of these three types of features and only the remaining two as the raw input each time. Then features to train GNS on CIFARIO and RTE. Detailed results are presented in the following table. It can be observed that all three categories of features benefit the quality of GNS. Parameter and gradient information more important than the global state since the global state can be reflected by the reward and the action. This ablation study validates our choice of these features. In the future, it 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.
different input
Table 18: Performance of GNS with different input features.
Scheduler
CIFAR10
RTE
GNS w/o parameter state
93.8 + 0.6
73.3 + 0.5
GNS w/o gradient state
93.7 + 0.4
73.5
+ 0.4
GNS w/o global state
94.1 + 0.6
74.0
GNS
94.3 + 0.5
74.4 + 0.3
E.9 COMPARISON OF OTHER F OF ACTION
We have tried 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 design some additional hyperparameters elaborately for each task to constrain the output residual to specific 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 the range [-10-5 10-5], and the action of predicting the learning rate by replacing the activation function with Sigmoid on CIFAR1O, 10, and the corresponding results are shown as follows:
Table 19: Performance of GNS with different forms of action.
Scheduler
CIFAR10
GNS (adding residual) GNS (predicting the learning rate)
90.7 + 0.6
10.0 + 0.0
GNS (multiplying by a factor)
94.3 + 0.5
We can see that adding a residual is insufficient to beat our original design of multiplying by factor. For predicting the learning rate, it is hard obtain useful scheduler and the training of CIFARIO failed and was stuck with only 10% accuracy. However, it is possible that by fine-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. EVALUATION ON MLPs
We include evaluation of training 3 layer MI. on Fashion MNIST with the hidden size of 128 and we report the performance of different optimizers in Table 20 It can be seen that GNS is also applicable to such simple scenario, with the best test accuracy of 89.8.
20

---

Published as a conference paper at ICLR 2022
Table 20: Test accuracy of 3-layer MLP on Fashion MNIST.
Scheduler
Fashion MNIST
Constant
88.8 0.2
Polynomial
89.3 + 0.3
Cosine
89.3 + 0.2
Hypergradient
89.0 0.2
SRLS
89.5 + 0.4
GNS
89.8 0.3
E.11 TRAINING LOSS CURVE
We have included the training loss curve for Fashion MNIST and CIFARIO in Figure Since for the GLUE benchmark tasks, final performance is usually the focus in the literature and the loss and the evaluation metric do necessarily correlate, do not show their loss curves.
Fislion MNIST
CIPARUOO
Dreban
Guueart
Peleumial
Userr
Godee
Thpegraleet
WLA
(a) Fashion MNIST
(b) CIFARIO
Figure Training loss curves for Fashion MNIST and CIFARIO.
21


Reviews:
[{"review_id": "k3C4ZuWc7VI", "reviewer": "Reviewer_RAQV", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. Finding a good learning rate scheduler is an interesting topic. This paper proposes a new way to search for a good learning rate by exploiting the neural network structure with a graph. \n2. In practice, the mean and variance of the weights, biases, and absolute gradient values, etc. are used as the node features. Why choose them as the features?\n3. What is the learning rate for the training of GNS? How to set this learning rate? How this learning rate will affect the performance? Those questions should be investigated.\n4. Why not use some other forms of action, for example, adding a residual?\n5. I recommend reporting the performance on ImageNet to make it more convincing. \n6. The full form of some abbreviations is not declared, which makes it hard to understand for some readers, for example, MRPC, RTE, and STS-B.", "summary_of_the_paper": "The paper proposes a Graph-based learning rate scheduler for neural network optimization. The graph is constructed based on the network structures and trained with reinforcement learning. The method is evaluated based on Image Classification and Natural Language Understanding and achieves superior performance than baselines. ", "summary_of_the_review": "In summary, the paper proposes a novel and effective way to learn a learning rate scheduler. However, some of the investigations such as the node feature and the meta-learning rate are missing. \n\n-- post-rebuttal\n\nMy concerns are well-addressed in the rebuttal. I will keep my score. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "jP1nTqKphcm", "reviewer": "Reviewer_RKqo", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strengths:\n\n* The improvements in the paper are intuitive and consistent across the benchmarks considered. The paper makes an effective presentation of the idea and the experiments are largely sound. \n\n* It is great that the paper improves realistic large models in settings such as RoBERTa fine-tuning.\n\n* It is natural and useful, in my opinion, that the authors utilize the graph structure of the computational graph.\n\n# Weaknesses:\n\n*  The training episodes for the PPO algorithm provide a serious bottleneck for searching for a learning rate scheduling policy. The authors address the problem in Section 5.5, however they do not say how much time it took for the CIFAR experiments. Could you clarify? For example, I am missing the details of what are $T$ and $K$ in Algorithm 1 for the CIFAR experiments. In Section 6 the authors also acknowledge that limitation, but they do not discuss how the potential future direction might look like. It is not clear to me how this method could be useful for long time horizons.\n\n* Focusing on the current experiments, it seems that the empirical argument the paper makes is that using GNS is more optimal than doing a grid search. What if we expand the space and use a more efficient search methods in the hyperparameter space? For example, a natural baseline seems to be CMA-ES (https://arxiv.org/abs/1604.00772). It would be interesting if there is something to say/ show about comparing your method to more efficient alternatives of grid search.\n\n* Is there an ablation on the choices of the raw features $x_v$ used in the GCN? It would be useful to take a look at such a study for future work with GNS.\n\n* I am not sure I understand the last sentence of Section 4, \"During inference, we first sample...\". Could you please clarify what the procedure is and how it is used in your experiments, because I don't think I see that described in Algorithm 1?\n\n\n## Minor:\n\n* In Equation 1 the matrix $\\mathbf{W}^l$ is not formally defined.\n\n* At the end of page 4, please use different enumeration for the steps in the hierarchical GNNs, because (a,b,c) is used in Figure 1 and it becomes difficult to read the text.\n\n* In the top of page 5, \"directly or with slight fine-tuning\": this is vague, could you be more precise what the fine-tuning is?\n\n* In the beginning of Section 4.2: the rest components <- the remaining components.\n\n* Page 6: estimated advantage estimator <- estimated advantage (redundant \"estimator\").\n\n* Algorithm 1 on page 13: Why is Eq. (1) an optimization step, i.e. are you sure you are referencing the right equation?; Section 3.3 is not present in the main text (perhaps a typo?).\n\n* Page 16: e can see that GNS <- We can see that GNS (typo).", "summary_of_the_paper": "This paper proposes a Graph Network-based Scheduler (GNS) method that learns a learning rate scheduling policy for problems with moderate time horizons for optimization. The novelty of the paper consist of representing the optimized neural network as a graph and using a policy that takes GCN representations of the architecture and its weights in a PPO algorithm. The paper also introduces strategies for validation that allow for more efficient searching of the policy. There are extensive experiments on training ResNets on CIFAR-10 and fine-tuning RoBERTA on GLUE. Additional benefits are demonstrated with transferring the learned policies to novel architectures or datasets. Ablation studies underscore the components in the GNS method.", "summary_of_the_review": "My initial recommendation is 5: marginally below the acceptance threshold. I believe that this paper has great potential to become useful to the community. However, in my opinion there are some important questions that need to be clarified. I have listed my comments/ thoughts and suggestions in the Main Review. If you could please address my comments, I could consider increasing my score.\n\n-- post-rebuttal\n\nThank you for your response. I will raise my score to 6.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RcoSq_Pwi0S", "reviewer": "Reviewer_TPKw", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n-- The paper explores an interesting idea -- using GCN to encode all model states to better learn a learning rate schedule for a given model.\n\nWeaknesses:\n-- The novelty of the paper is limited. \n-- The paper fails to provide convincing evidence that the suboptimal accuracy of existing methods is truly a result of using the final layer state, e.g., is it possible that the RL-based optimization scheme itself is suboptimal? \n-- The evaluation of image classification tasks is conducted on toy datasets such Fashion-MNIST/CIFAR10. To be more convincing, the paper should include an evaluation of larger datasets such as ImageNet. \n-- Evaluation on RoBERTa-base/large on GLUE datasets appear to be using sub-optimal baselines. \n-- Might not be easy to apply in practice. \n", "summary_of_the_paper": "This paper introduces a data-driven approach to automatically learn the learning rate schedule for a given dataset/task. The paper claims that existing RL-based adaptive learning rate schedule methods are insufficient to work well for large models such as RoBERTa because they only use model states from the last layer. The paper then proposes to use all layers' hidden states as feature vectors and uses graph neural network GCN to encode those features. The paper claims that such a feature encoding scheme allows reinforcement learning (e.g., PPO) to better learn a learning rate schedule, especially for large models. The paper evaluates its proposed method on both image classification (e.g., Fashion-MNIST and CIFAR-10) and language understanding tasks (finetuning RoBERTa on GLUE) and shows that the proposed method is able to achieve better final accuracy than some heuristic-based learning rate schedules. ", "summary_of_the_review": "The paper explores an interesting idea of using graph neural networks to better encode model states, with a hope to improve the existing RL-based meta-learning method. Despite showing some promising results, there are quite a few concerns about this work. \n\nFirst, the RL-based scheme (e.g., PPO) is very similar to Xu et al. 2019. Therefore, the remaining novelty is really on changing FFN to GCN to encode more model hidden states. \n\nSecond, it seems there could be multiple factors that impact the final accuracy, such as the encoding scheme (e.g., GCN vs. other encoders), the features (last layers vs. all intermediate states), the optimization mechanism (RL vs. alternative methods), the choice of hyperparameters. However, the paper jumps to the conclusion that the encoding scheme and features are suboptimal without giving enough evidence to show that they are indeed the culprit. The paper also changes multiple factors at once, including tuning additional hyperparameters, making it difficult to tell where the improvements actually come from. The end-to-end results are useful but they do not reveal how the proposed method works internally. To be more convincing, it would be better if the paper could provide more insights on how encoding affects the final accuracy. For example, does the quality of the GCN matter when performing the encoding? If so, how do you measure the quality quantitatively? Could the existing method from Xu et al. 2019 still work by slightly adjusting its model architecture, e.g., by increasing the model depth/width or changing the hyperparameters?\n\nThird, the evaluation on RoBERTa-base/large appears to use unoptimized baselines. For example, the RoBERTa-large baseline reaches 88.1 GLUE score on the leaderboard whereas this paper only reports 86.9. If we look at individual datasets, the RoBERTa-large results from this paper are also constantly worse than what has been reported (e.g., 90.0 vs. 90.8 for MNLI, 64.2 vs. 67.8 for CoLA). If we take the publicly reported results as the baseline, the proposed method actually performs worse than the RoBERTa-large baseline. Furthermore, it is also unclear whether the paper follows the standard practice when setting the learning rates schedules when producing the baselines for GLUE datasets. For example, RoBERTa models are often fine-tuned with warmups followed by a linear decay scheduling. However, it does not seem to be the case that the paper includes this basic baseline in the comparison. It raises concerns on whether the improvements in Table 2 are merely an effect of comparing with a set of suboptimal learning rate schedules. To be more convincing, it would be better to compare the proposed method with the best practice of fine-tuning RoBERTa-large. \n\nFourth, given that the proposed method requires training an additional policy network and GCN to learn the learning rate schedule, which adds additional difficulties to the already complex training of large models,  it raises concerns about the usefulness of the proposed method in practice. \n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "IwJO6uzdLv-", "reviewer": "Reviewer_woPs", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n- The authors has done a great job by employing the techniques from RL and graph neural networks to formulate learning rate schedule and find the optimal trajectory. Finding the optimal schedule of learning rate is an interesting and important practical problem which can help improve the performance of the models in all machine learning community.\n\n- Using graph neural networks for computing the features of the nodes is an innovative idea, because it offers two benefits that the other methods lack: 1- it can use all the information in the middle layers as opposed to the last layer, 2- the learned graph convolutional network (GCN) can be reused and generalized to new networks, so this makes it architecture-agnostic. \n\n- The experiments are convincing and complete. It is clear to see the benefits of GNS over all other methods. The authors has shown the performance improvement in different tasks for image classification and language understanding, analyzed the generalization of GNS between different networks, compared the run times, and done an ablation study to show the importance of different parts of the model.\n\n- The paper is well-written and easy to follow and understand. \n\nWeaknesses:\n\n- The authors has used GCN as the model of choice for GNNs. But there is no explanation on why they choose this. I wonder if the authors have also considered other models such as graph attention networks (GAT)? The attention mechanisms have proven effective in performance improvement in many tasks. It would be great if the authors check to see if other GNNs can even outperform GCNs in this specific problem. \n\n\n", "summary_of_the_paper": "This paper formulates the learning the schedule of learning rate as a reinforcement learning (RL) problem and finds the optimal schedule which outperforms the baseline scheduling algorithms such as constant learning rate, polynomial or cosine decay, warmup and restart, gradient-based method like Hypergradient, and another learning-based method SRLS. The paper proposes Graph-Network-based Scheduler (GNS). Rather than looking at the last layer or concatenating different layers of the network, GNS uses graph convolutional networks (GCN) to do message passings over the computational graphs of the model, and as a results, are more efficient, informative, and architecture-agnostic which can easily generalize to other network architectures. The authors has shown that GNS outperforms all other learning rate scheduling algorithms and achieves superior performance in various tasks. ", "summary_of_the_review": "I believe the problem addressed in this paper (how to learn optimal learning rate schedule) is very important to the machine learning community, and the techniques used to solve it (RL and GNNs) are powerful and effective. The authors has formulated the problem as an RL problem, which is the most correct way to do it, and has also thought how to make it architecture-agnostic, generalizable, and versatile in using the hidden layers' information by message passing through the computational graphs. Overall, it seems that this formulation and solution is the best one can do to learn the optimal learning rate schedule, which is validated by the experimental results and the ablation study. Therefore, I believe this paper is very interesting and I recommend it for acceptance. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Hrbv3H73nJ", "reviewer": "Reviewer_NP67", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros: The proposed approach for the learning rate scheduling is novel. It can achieve better performance than baselines.\n\nIssues:\n1. Can the proposed method be applicable to simpler models, such as multilayer perceptrons (MLPs)?\n\n2. For a graph constructed based on a network model, what are the features of nodes?  How to keep the dimensions of node features consistent?\n\n3. The training loss curve should be displayed for better comparison.\n\n4. In the experiment, only some simple hand-designed scheduling mechanisms are compared. The author should compare some baseline methods based on reinforcement learning, such as [1,2].\n\n[1] Zhen Xu, Andrew M Dai, Jonas Kemp, and Luke Metz. Learning an adaptive learning rate schedule.\narXiv preprint arXiv:1909.09712, 2019.\n\n[2] Chang Xu, Tao Qin, Gang Wang, and Tie-Yan Liu. Reinforcement learning for learning rate control.\narXiv preprint arXiv:1705.11159, 2017.", "summary_of_the_paper": "In this paper, the authors study the learning rate scheduling for optimization in training deep neural networks.\nIt proposes graph-network-based scheduler (GNS) to learn a scheduling mechanism for learning rate.\nGNS encodes the neural network by GNN and trains an agent to control the learning rate. \nExtensive experiments are conducted to demonstrate the effectiveness of the proposed method.", "summary_of_the_review": "In general, the paper is written well and easy to follow. However, there exist aforementioned concerns that the authors need to address.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: ### Description 
The paper develops a new automatic scheduler to schedule the learning rate during the training. The scheduler has access to the current training state summarized by certain statistics of weights and gradients in all layers and the loss history. It is trained by reinforcement learning with the reward derived from the progress with respect to the performance measure, such as validation accuracy. The key innovations are the design of the state vector using graph convolutional neural networks and empirical improvements to the reward function. The main claim is that GCNs allow to take into account the architecture of the network to be trained and the state of all layers, which, authors hypothesize and demonstrate experimentally, improves performance and transferability of the scheduler across networks and tasks.

### Decision
The reviewers recommendations after the rebuttal settled on 4 x "marginally above" and one "accept". Respectively, I recommend to accept. I recommend a poster based on the reception by reviewers: the novelty was assessed as limited because the idea of an automatic schedulers based on RL with a similar learning strategy belongs to the prior work. On the strong side, the paper satisfied all requests by reviewers for the experiments, regarding alternative methods, large datasets and ablation studies demonstrating that it is indeed the new architecture that allows to achieve a significant improvement, making it a solid improvement step. Amongst alternative methods the paper considers all viable alternatives: a function-based schedulers, a hyper-gradient method, and a the RL based scheduler, optimized in hyperparameters.

### Discussion
There was no significant non-public discussion. As an additional feedback, let me just share my observations.

What is somewhat unclear in that the paper starts by discussing the directed graph of a feed-forward network, then it proposes to run GCN on it, which is undirected. Then the hierarchical method is proposed, which runs GCN on each block sequentially while taking the aggregated input from the preceding block. This makes it a directed processing method on the level of block. I wonder whether the directed processing is desired or not desired here? Can a sequential processing summarize the network state efficiently on its own, similar to feedforward propagation, without the global averaging proposed?

It was not very clear to me from the paper what is the meaning of the batch size for training the GSN, and respectively what the batch normalization is doing there.

From some 100 mile perspective, it seems to me that whatever efficient optimization can be performed on the validation set, it helps. It does not matter so much what is varied: the learning rate schedule, other hyperparameters or even the network architecture. So in a sense it is not surprising that one can improve. What is more interesting is that the learned schedulers are generalizable / transferable, as demonstrated in Section 5.4 (changing the architecture or going from CIFAR to ImageNet while keeping the scheduler).

The work has done quite a lot on the experimental side with the baseline GCN model that they proposed. It seems to have still lots of potential via different possible enhancements. For example, what authors mentioned, including exponentially weighted running averages of gradients and squared gradients into the features. They already tried GAT instead of GCN as proposed by reviewers. There may be hyperparameters other than those controlling the learning rate. Some such hyperparameters, e.g. momentum, are apparently tightly coupled with the learning rate. The paper does not discuss how to tune them together with GNS. It could be a difficulty. On the other hand, they can be potentially scheduled with the same GNS. 

When considering SGD with momentum (which is not used in the paper), please note that the common use of a momentum parameter $\mu$ actually mixes the learning rate together with the smoothing parameter controlling the exponentially weighted averaging. So if one wants to control the learning rate alone, it is better to implement the gradient smoothing is done in .e.g. Adam, with its hyperparameter $\beta_1$.
Decision: Accept (Poster)