Neural Architecture Search by Learning a Hierarchical Search Space

Authors: Mehraveh Javan Roshtkhari, Matthew Toews, Marco Pedersoli
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: Neural Architecture Search, Monte-Carlo Tree Search, Hierarchical Search Space, Hierarchical Clustering

Abstract:
Monte-Carlo Tree Search (MCTS) is a powerful tool for many non-differentiable search related problems such as adversarial games. However, the performance of such approach highly depends on the order of the nodes that are considered at each branching of the tree. If the first branches are not discriminative enough, i.e. they cannot distinguish between promising and deceiving configurations for the final task, the efficiency of the search is exponentially reduced. While in some cases the order of the branching is given as part of the problem (e.g. in chess the sequential order of the moves is defined by the game), in others, such as Neural Architecture Search (NAS), the visiting order of the tree is not important, and only the final architecture matters. In this paper, we study the application of MCTS to NAS for the task of image classification. We analyze several sampling methods and branching alternatives for MCTS and propose to learn the branching by hierarchical clustering of architectures based on their similarity. The similarity is measured by the pairwise distance of output vectors of architectures. Extensive experiments on two challenging benchmarks on CIFAR10 and ImageNet show that MCTS, if provided with a good branching hierarchy, can yield promising solutions more efficiently than other approaches for NAS problems.

Decision: reject
Reviewer scores: 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1730463347441

thread_type:
official_review

<thread_object>
creator: reviewer_c3y4

<content>
summary:
This paper presents a Neural Architecture Search (NAS) approach that leverages Monte Carlo Tree Search (MCTS) with a learned hierarchical search space. Instead of using a non-optimal, pre-defined hierarchical search order, this paper proposes to learn the branching by hierarchical clustering of architectures based on their similarity measured by the pairwise distance of output of architectures. The experiments on CIFAR10 and ImageNet demonstrate that the proposed approach yields better solutions than previous approaches.

strengths:
This paper highlights the shortcomings of the previously used node independence assumption and demonstrates that too restrictive assumptions converge to worse solutions.

weaknesses:
The weaknesses of this paper are the limited novelty, results not significant, and the unclear approach.

First, for the limited novelty and results not significant, this work improves previous works (Wang et al., 2021a; Zhao et al., 2021b) by replacing the model accuracy with the output vector. While the output vector provides more information for clustering architectures, the novelty is limited. Furthermore, the results only slightly outperform previous works (Su et al., 2021a; You et al. 2020; …), which is insignificant. In addition, making early tree nodes more discriminative is highly relevant to the partitioning or splitting problems in decision tree learning, which has been studied by many in the past (Costa and Pedreira, 2023).
Costa, V.G., Pedreira, C.E. Recent advances in decision trees: an updated survey. Artif Intell Rev 56, 4765–4800 (2023).

Second, the presentation has a lot of improvements, especially the approach. The proposed method is ambiguous and does not seem like MCTS. MCTS uses UCT to select child modes; however, the proposed method uses Boltzmann sampling with a UCB-like score as the parameter. The authors should justify whether this design follows the UCT foundations of balancing exploration and exploitation. 
Most importantly, using $Acc(a_i)/n_i$ in Eq. 4 is weird. From the definition, the first term in the formula is the average reward (Eq. 1 in Kocsis & Szepesvári, 2006). However, Eq. 4 further divides the accuracy by the visit count. Since the accuracy $Acc(a_i)$ is already considered the average reward, it makes no sense. If this is not a typo, the authors should justify the correctness of such a design.

Please refer below for more questions.

questions:
Questions related to the proposed MCTS procedure:
- As the tree is already constructed, does the algorithm still run selection from a single root node and then expand the known tree structure? Or does it just start sampling from the entirely constructed tree?
- The typical MCTS involves several phases (e.g., selection, expansion, simulation, backpropagation) per simulation, while it is unclear how the proposed procedures in Algorithm 1 are linked to these phases.
- It is mentioned that $C(a_i) = Acc(a_i)$ for architecture search in line 340. Does it mean that for supernet training, $C(a_i)$ is set to Eq. 4? It is unclear which parts use $Acc(a_i)$ or Eq. 4 in Algorithm 1.
- In Algorithm, $P_{train}$, $P_{search}$, and Eq. 5 are not defined.
- In line 358, it is mentioned that there is a warm-up period for uniform sampling, which is also not included in the typical MCTS routines (Kocsis & Szepesvári, 2006). As MCTS should already be able to balance exploration and exploitation, what is the purpose of adding such a warm-up period?
- In line 360, why is $C(a_i) < 1$ when nodes are visited?

Other comments related to typos and presentation issues:
- For Figure 1, it is difficult to understand why the subfigures "independent" and "joint" are drawn like this.
- For Figure 2, (b) and (c) use different styles to represent the tree structure, which should be normalized to the same.
- The section title "3.4 Sampling with conditional probabilities: Monte Carlo Tree Search" is confusing as this section does not seem to have any links to (the typical) MCTS.
- For Table 3, "the categorical vector representation" is not included in the discussion.
- Some terms are not consistently used, e.g., "Monte Carlo Tree Search" or "Monte-Carlo Tree Search"; "equation" or "eq." or "Eq.".
- Several typos, e.g., "Fig.3.4(a)", "and T the temperature term", "UTC".
- Placing Algorithm 1 in Appendix B lowers the readability. It would be more appropriate to include it in the main text, especially since the authors often refer to it with "see algorithm 1 in Appendix B".
- In Algorithm 1, the equations should be explicitly stated instead of mentioning "as in Eq. 1."
</content>
</thread_object>

</thread>
<thread>
creation date:
1730480900450

thread_type:
official_review

<thread_object>
creator: reviewer_5vyd

<content>
summary:
This paper challenges the commonly assumed node independence in Neural Architecture Search (NAS), which may limit both efficiency and performance. To address this, the authors propose a Monte Carlo Tree Search (MCTS) method incorporating a learned hierarchical tree structure, built with agglomerative clustering based on model output distances, to improve NAS effectiveness. Experiments are conducted on NAS benchmarks for CIFAR-10 and ImageNet image classification tasks.

strengths:
* The paper introduces an approach by addressing node dependencies to improve NAS efficiency.
* Leveraging the UCT (Upper Confidence bounds applied to Trees) approach, the authors further utilize a learned tree structure to reduce the reliance on manually crafted search space designs.
* The paper provides ablation studies to analyze the effects of the proposed method in more depth.

weaknesses:
* The abstract may benefit from significant revision. Currently, it primarily highlights MCTS fundamentals and suggests general applicability, but the paper is focused on a NAS-specific task that utilizes MCTS and related techniques to enhance NAS performance. The abstract and introduction appear inconsistent in conveying the core contribution and scope.
* The rationale behind using model output distances to construct the tree structure and improve NAS is not clearly discussed, and the method itself lacks detail. This part should be the core of the paper, yet there is minimal explanation in the main text.
* While resource constraints may be a factor, it remains unclear whether the method scales well for large networks, which are particularly relevant in NAS applications. The experiments mainly validate that the learned tree provides slight improvements but do not assess scalability in larger search spaces.

Minor Comments
* Line 306: check around "Fig.3.4(a)"

questions:
* How does this method compare with other state-of-the-art NAS techniques, such as those in the Neural Architecture Transfer (NAT) series?
* What insights or theoretical basis underlie the decision to use model output distances for improving NAS performance? (I already assume this will be addressed in a revision in my score.)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730714660138

thread_type:
official_review

<thread_object>
creator: reviewer_xuxx

<content>
summary:
This paper proposes a method for supernet sampling for neural architecture search using Monte-Carlo Tree Search (MCTS). After an initial phase of supernet training, the method uses similarity distances between architecture outputs and hierarchical clustering to build a search tree, then continue the supernet training by sampling from this tree using MCTS.

strengths:
-The paper is overall well written.
-The methodology is well-explained and the contributions are clearly defined, the paper is well-placed in the literature.
-While not theoretically justified, the idea of learning the Monte-Carlo tree is promising.
-The experimental results are convincing on the ImageNet dataset.

weaknesses:
-The method still requires initial supernet training using uniform sampling before being able to build the tree, which is known to be computationally heavy.
-The overall contribution seems incremental, as it is mainly a new way to construct a Monte-Carlo tree for supernet sampling.
-For the experiment on the pooling dataset, the authors explain that this extremely small search space of 36 architectures is challenging because the initial supernet training shares weights between architectures with different pooling configurations. Given that the proposed method discriminates architectures by comparing the outputs after supernet pre-training, I wonder how the method is able to find a more efficient representation of the tree if the weights themselves are not optimal. Furthermore, the classical sampling methods (uniform, Boltzmann…) are unable to find the best architecture out of 36? How many samples are performed? The results, while in line with the results of [1], seem surprising and the paper could benefit from a more thorough explanation.
-There are several typos and the writing is overall unclear in Section 5.1.
-Is the Boltzmann sampling over UCT in Section 4.2 necessary? The UCT formula already offers a trade-off between exploration and exploitation. If it is necessary, then an ablation study could be useful.
-The following claim : “Different from other works such as Wang et al. (2021a) and Zhao et al. (2021b) that use the model accuracy directly for the tree design, the output vector provides more information for clustering architectures” seems unsupported.
-Building the search tree requires building a hierarchical clustering. As the authors use the pairwise distance matrix of all architectures in the search space over a mini-batch, the complexity of building this hierarchical clustering is O(n^2) complexity. For large search spaces, this could be very inefficient. A comparative complexity analysis of the proposed method would be welcome.

[1] : Javan et al., Balanced Mixture of SuperNets for Learning the CNN Pooling Architecture, 2023

questions:
The paper proposes an interesting idea, is mainly well-written and shows some good results on benchmark datasets. As written in the weaknesses section, there are several avenues for clarification and improvements on the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732825309854

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
New version of manuscript

comment:
We thank the reviewers for their detailed reading of the manuscript and valuable comments. We have responded to each reviewer individually and uploaded a new version of the paper that addresses the reviewers’ concerns. 

The additions and changes are shown in blue in the new version. The summary of changes to the manuscript are the following:


- Added a sentence to the abstract to connect it more to the actual problem we are solving, as suggested by Rev. 5Vyd
- Reshaped the last part of the introduction to explain clearly the difference between our approach and two MCTS related papers and the use of clustering on the network class output
- Changed related works on the same two papers, to better explain the differences
- Updated fig. 2 to have the same letter for the architectures, as suggested by Rev. c3Y4
- Updated Tree design in our method, better explaining the advantages of our clustering for the tree partitioning as asked by Rev. 5Vyd
- Improved Search and Training in our method, to acknowledge that our MCTS does not use all the phases of a standard MCTS, and uses Boltzmann sampling to improve the exploration. Added also more context about what are the actual phases of the training. Rev. c3Y4 asked to move the algorithm in the main paper, but unfortunately we did not have space.
- Updated Eq. 3 and 4, to fix a typo and to be more in line with the algorithm in the supplementary material
- Improved the presentation of the pooling search space in section 5.1 as asked by Rev. xUxx
Added an ablation to compare our clustering with a partitioning based on accuracy as asked by Rev. xUxx
-  Added more explanations for tab. 3 in the text
- Added limitations in conclusion
- Improved algorithm in supplementary material as asked by Rev. c3Y4
- Added complexity analysis to the supplementary as asked by Rev. xUxx
</content>
</thread_object>

</thread>
<thread>
creation date:
1733023778202

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Rebuttal

comment:
We would like to thank the reviewers for their valuable comments which greatly helped us in improving our paper. We have responded to each reviewer individually and uploaded a new version of the paper that addresses the reviewers’ concerns. We would like to summarize our answers to the main concerns raised by reviewers:

***

**Experimental Results:**  The experiments in our paper were carried out on two CIFAR10 NAS benchmarks and ImageNet mobilenet-like search space. The distribution of the accuracy of architectures in those benchmarks is narrow, with many architectures performing well (for example for distribution of architectures in NAS-Bench-Macro see figure 7.c in Su et al., 2021a). Therefore, numerical gains in terms of accuracy may seem small. However, our method is able to find architectures close to the best (on CIFAR10) or the best (on NAS-Bench-Macro in table 4) and significantly reduce the search cost on ImageNet compared to similar methods.

**Contribution:** We understand reviewers' concerns about the proposed contribution. However, some of those issues were due to a not accurate presentation of the related work. In particular, reviewers considered (Wang et al. 2021a) and (Zhao at al. 2021b) comparable to our work, but with a different strategy to partition the search space with a tree. However, those two papers do not learn the recognition model while finding the optimal architecture as we do. Instead they use MCTS on a given and fixed recognition model (Uniformly trained CNN for (Wang et al. 2021a) and BenchNas-201 for (Zhao at al. 2021b)).

This makes the problem we our tackle quite different than those two and that require different solution and not really comparable. For instance, the choice of using a clustering approach for the tree (in contrast to using directly the model accuracy as in (Su et al., 2021a; Wang et al. 2021a)) is due to the fact that the initial model has poor performance and accuracy is a poor proxy for the real model performance. Thus, partitioning with that would lead to suboptimal results. We added some explanations about that in the presentation of the method and an additional experiment to confirm this intuition. Instead, considering proximity in an unsupervised manner, our clustering algorithm takes into account the model similarities and differences without considering accuracy and leads to better results.
The only paper performing the joint learning of the recognition model and search of the optimal architecture as us is (Su et al., 2021a), and we analyze and compare explicitly and thoroughly with them.  

**Limitations:** Our model in its current form works best for search space of roughly <10k, due to the quadratic complexity of the similarity matrix. We have analyzed this in Appendix C.5, showing that for small search spaces the main computational cost is the computation of the recognition model output, which is linear in the number of architectures. For larger search space, the complexity of the similarity matrix is an issue, not much in terms of computation, but more in terms of memory storage.

In general, there are possible ways to deal with larger search spaces. For instance, on ImageNet tests, we used FLOPs as a simple training-free proxy for performance (similar to Su et al., 2021a) to prune the search space. Other techniques such as zero-cost proxies (Abdelfattah et al., 2022) can be combined with our method to eliminate unpromising architectures.

**references** 

Su et al., 2021a: Prioritized Architecture Sampling with Monte-Carlo Tree Search, CVPR 2021.

Wang et al. 2021a: Sample-efficient neural architecture search by learning actions for monte carlo tree search, 2021.

Zhao at al. 2021b: Multi-objective optimization by learning space partitions, ICLR 2022.

Abdelfattah et al., 2022: Zero-Cost Proxies for Lightweight NAS, ICLR 2022.

***

We hope our answers satisfy reviewers' questions, doubts and comments and look forward to further discussions and clarification. We also hope the reviewers take into account the improved manuscript and responses during the rebuttal and reconsider their scores.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733158550896

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Request for feedback and revaluation of the scores!

comment:
Dear reviewers and AC,
we answered all your comments and clarification, updated the paper and provided a summary of the most critical points. 
We believe the new version of the paper is substantially improved and we would like your feedback and comments on that before the end of the discussion period, to give us a chance to raise your scores!
Thank you!
</content>
</thread_object>

</thread>


Metareview:
metareview:
In this work a MTCS approach to neural architecture search. The method uses similarity distances between the architecture outputs and a hierarchical clustering algorithm to build a search tree. The experimental evaluation is conducted on CIFAR and ImageNet classification tasks. 

Some reviewer’s concerns included the expensive training of the supernet procedure, and the writing of the abstract that focuses on MCTS which shows discrepancy with the general theme of the paper that is centered around NAS. Finally, the main issue raised by the reviewers was the limited novelty and unconvincing experimental evaluation showing sufficient gains for the proposed approach. This is not mitigated by a clear explanation of the approach that justifies the novel algorithmic choices and may mitigate concerns about the performance gap.