Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 
main_review: Strengths:
-- The paper explores an interesting idea -- using GCN to encode all model states to better learn a learning rate schedule for a given model.

Weaknesses:
-- The novelty of the paper is limited. 
-- 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? 
-- 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. 
-- Evaluation on RoBERTa-base/large on GLUE datasets appear to be using sub-optimal baselines. 
-- Might not be easy to apply in practice. 

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. 

First, 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. 

Second, 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?

Third, 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. 

Fourth, 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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
main_review: # Strengths:

* 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. 

* It is great that the paper improves realistic large models in settings such as RoBERTa fine-tuning.

* It is natural and useful, in my opinion, that the authors utilize the graph structure of the computational graph.

# Weaknesses:

*  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.

* 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.

* 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.

* 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?


## Minor:

* In Equation 1 the matrix $\mathbf{W}^l$ is not formally defined.

* 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.

* 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?

* In the beginning of Section 4.2: the rest components <- the remaining components.

* Page 6: estimated advantage estimator <- estimated advantage (redundant "estimator").

* 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?).

* Page 16: e can see that GNS <- We can see that GNS (typo).
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.

-- post-rebuttal

Thank you for your response. I will raise my score to 6.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 
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. 
2. 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?
3. 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.
4. Why not use some other forms of action, for example, adding a residual?
5. I recommend reporting the performance on ImageNet to make it more convincing. 
6. 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_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. 

-- post-rebuttal

My concerns are well-addressed in the rebuttal. I will keep my score.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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. 
main_review: Strengths:

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

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

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

- The paper is well-written and easy to follow and understand. 

Weaknesses:

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



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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: In this paper, the authors study the learning rate scheduling for optimization in training deep neural networks.
It proposes graph-network-based scheduler (GNS) to learn a scheduling mechanism for learning rate.
GNS encodes the neural network by GNN and trains an agent to control the learning rate. 
Extensive experiments are conducted to demonstrate the effectiveness of the proposed method.
main_review: Pros: The proposed approach for the learning rate scheduling is novel. It can achieve better performance than baselines.

Issues:
1. Can the proposed method be applicable to simpler models, such as multilayer perceptrons (MLPs)?

2. For a graph constructed based on a network model, what are the features of nodes?  How to keep the dimensions of node features consistent?

3. The training loss curve should be displayed for better comparison.

4. 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].

[1] Zhen Xu, Andrew M Dai, Jonas Kemp, and Luke Metz. Learning an adaptive learning rate schedule.
arXiv preprint arXiv:1909.09712, 2019.

[2] Chang Xu, Tao Qin, Gang Wang, and Tie-Yan Liu. Reinforcement learning for learning rate control.
arXiv preprint arXiv:1705.11159, 2017.
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.
review_title: null

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)