Paper: DiffusionNAG: Predictor-guided Neural Architecture Generation with Diffusion Models (ICLR-2024)

Reviewer comment (Theory):
Even though the motivation to the use of diffusion models to neural architecture search is interesting, it seems more appropriate to use discrete diffusion models for discrete random variables instead of using a continuous variables. For instance [1] is an interesting approach which seems to be more correlated to the desired problem instead of continuous diffusion.

Author response:
**W1. Even though the motivation to the use of diffusion models to NAS is interesting, using discrete diffusion models seems more appropriate.** 

- Thank you for mentioning that our motivation for applying the diffusion model to NAS is interesting. Additionally, we appreciate your insightful suggestions regarding the types of diffusion models to address NAS challenges. The most critical aspect of applying the diffusion model to NAS is the ability to generate a **valid** neural architecture (a valid graph) while adhering to the rules of the given search space. Therefore, we believe that **as long as we can generate valid neural architectures, we can utilize both continuous and discrete diffusion models**. As suggested by the reviewer, we agree that discrete diffusion models also hold promise in neural architecture generation, as demonstrated by their excellent performance in image synthesis [5].

- In our case, by referring to recent research works [1, 2, 3, 4] that highlight the superior performance of continuous graph diffusion models in generating undirected graphs, we have adopted continuous diffusion models as our primary framework. Following those prior works, we treat graphs with various node/edge categories as entities embedded in continuous space. The node/edge matrix is represented as a binary matrix through the one-hot encoding of features. We discretize the entries of the architecture matrices using the threshold $0.5$ to obtain discrete 0-1 matrices after generating samples by simulating the reverse diffusion process (please see **Section 2**, the reverse process of the main paper for a more detailed explanation).

- Furthermore, please note that, due to the unique characteristics of neural architectures and the diverse search spaces in NAS tasks, merely applying existing diffusion models to NAS without careful considering these characteristics might result in the failure to generate **valid** neural architectures. To address this challenge, we have developed a novel score network for neural architectures to effectively capture the data distribution of neural architectures within the given search space. As shown in **Table 4** of the main paper, **our proposed diffusion model successfully achieved 100.00% validity in both the NB201 and MBv3 search spaces**, while the existing graph diffusion model, GDSS [2], totally fails as 4.56% and 0.00%, respectively. (For a detailed explanation of the challenges in generating valid neural architectures and our technical contributions in designing the score network, please refer to our response to **W2**.)

- Additionally, in the sections **Q1. The feasibility of applying DiffusionNAG to generate optimal generative models** and **Q2. The feasibility of applying DiffusionNAG to generate sparse foundation models**, we discuss the potential for future research utilizing the discrete diffusion models.

---
**References**

[1] Permutation invariant graph generation via score-based generative modeling, ICML 2020.

[2] Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations, ICML 2022.

[3] Exploring Chemical Space with Score-based Out-of-distribution Generation, ICML 2023.

[4] Graph Generation with Destination-Predicting Diffusion Mixture, Arxiv 2023.

[5] Vector quantized diffusion model for text-to-image synthesis, CVPR 2022.

Reviewer comment (Novelty):
The use of diffusion models for NAS seems to be very incremental, as the proposed approach is a direct application of the standard available diffusion based approaches. What are some of the major hurdles faced by the authors in making this work for NAS?

Author response:
**W2. Distinguishing features from other graph diffusion models.**

**[Technical novelty of DiffusionNAG as a diffusion model for NAS]**
- The graph tasks typically handled by existing diffusion models, such as molecules or community connections, are undirected, and the order of node generation is relatively less critical. However, when dealing with neural architectures, the graph structures become more intricate due to 1) the necessity of reflecting the computational flow from the input layer to the output layer and 2) the requirement to adhere to complex and specific rules that vary across different search spaces (Please refer to the **[Validity criteria]** section below for more detailed explanations regarding validity rules). Therefore, to generate valid graphs representing neural architectures, diffusion models should capture 1) the **dependency between nodes**, reflecting the computational flow, and 2) the **accurate position of each layer** within the overall neural architecture to comply with the rules of a specific search space.

- However, to the best of our knowledge, none of the existing graph diffusion models can effectively capture the directionality of graphs or provide a specialized representation to encode the accurate flow of topological ordering. Consequently, as illustrated in Table 4 of our main paper, directly applying the existing diffusion-based models, such as **GDSS**, leads to difficulties in generating valid architectures. The validity ratios are only **4.56%** and **0.00%** for the NB201 and MBv3 search spaces, respectively.

- To address this issue, we first represented neural architectures as **Directed Acyclic Graphs (DAGs) with an upper triangular adjacency matrix**. Then we designed a novel score network to generate valid neural architectures. As detailed in **Section 2** of the main paper, we employed **transformer blocks** with an attention mask that indicates the dependency between nodes to parameterize the score network. Finally, we incorporated **positional embeddings** for nodes, as shown in **Eq (4)**, to capture the accurate topological ordering of layers in architectures. 

- As shown in **Table 4** of the main paper, our DAG representation and transformer-based score network (without positional embeddings) to capture the computational flow (node dependency) of neural architectures achieved **42.17%** validity in MBv3 search space while the existing graph diffusion model, **GDSS**, totally failed as **0.0%**. Notably, our final score network with positional embedding achieved **100.0% validity**, highlighting the necessity of positional embedding for generating architectures to follow complex rules of MBv3 search space. 

**[Validity criteria]**

- NB201 search space:
  - 1) Single operation per layer: Each layer should correspond to only one type of operation, avoiding the presence of multiple operations for a single layer.
  - 2) Capturing the start and end of the computational flow: In the cell-based NB201 search space, the initial node where the computational flow begins must have the operation corresponding to START, and the final node where the computational flow concludes should have the END operation. Additionally, there should be no operations corresponding to START and END between the initial and final nodes.
- MBv3 search space:
  - 1) Single operation per layer: Each layer should correspond to only one type of operation, avoiding the presence of multiple operations for a single layer.
  - 2) No empty stage: The MBv3 search space defines a fixed macro skeleton with five stages. Each stage must consist of at least two layers to prevent the occurrence of an empty (blocked) stage. This configuration, with a directed edge from the node immediately before it, ensures a continuous main path connecting all nodes.
  - 3) No isolated paths: Within each stage, no layer should lack predecessors, ensuring that there are no isolated paths.

**[Contributions to NAS field]**
- The proposed DiffusionNAG not only introduces a score network for NAS but also makes significant contributions by addressing key limitations in conventional NAS methods, offering efficient search, superior predictor utilization, and easy adaptability across tasks. Specifically, we pioneer a paradigm shift towards a Conditional Neural Architecture Generation (NAG) scheme from NAS, enhancing search efficiency by generating neural architectures aligned with a specific distribution. Furthermore, DiffusionNAG extends the predictor's role for both NAG and evaluation purposes, while the traditional approaches restrict the same predictor to the evaluation role. Additionally, DiffusionNAG can generate architectures for unseen tasks, thereby improving the generalization capability compared with most NAS methods. Through extensive experiments in two critical scenarios, Transferable NAS and BO-NAS, DiffusionNAG consistently demonstrates these contributions successfully by outperforming existing NAS methods.

Reviewer comment (Experiments):
For conditional NAG, the authors can consider convert A_t to A_0 using the appropriate equations depending on the choice of diffusion sampler used and then apply the pre-trained predictor. This can truly enable using a pre-trained predictor which is just a function of A_0 instead of A_t.

Author response:
**W3. A suggestion to use the predictor with input $A_0$ instead of $A_t$.**

- Thank you for suggesting an insightful perspective. As you mentioned, if we can accurately transform $A_t$ to $A_0$ during the architecture sampling process using the equation related to the diffusion sampler we adopted, we can effectively guide the generation of neural architectures using a predictor that takes $A_0$ as input. However, even if we have already defined the forward diffusion process that perturbs $A_0$ to $A_t$ from the diffusion sampler we adopted, due to the stochastic nature of the diffusion process, it is challenging to accurately predict $A_0$ from given $A_t$ during the sampling of neural architectures which simulates the *reverse* diffusion process described in **Eq (2)**. 
- In more detail, as shown in **Appendix Section B.2**, we utilize the SDE equation to model the forward diffusion process for our neural architecture generation framework. Given the SDE describes the process (please refer to the equation below),  we can only know the mean and standard deviation of the neural architecture distribution at each time step, $p(A_t)$. Thus, it is difficult to accurately predict $A_0$ for a given $A_t$ during the architecture sampling process simulating the reverse diffusion process, which starts from noise sampled from the known prior distribution (please refer to **Eq. (11)** and **Eq. (12)** of the **Appendix** in the main paper for a more detailed description).

$$p_{t|0}(A_{t}|A_0) = \mathcal{N}\left(A_t;  A_{0} \, \sigma_{\min}^2\left(\frac{\sigma_{\max}}{\sigma_{\min}}\right)^{2t}\mathbf{I}\right)$$

**[Empirical observation]**

- Therefore, compared to generating neural architectures using a predictor with $A_t$ as input, it will be less effective to guide neural architecture generation using an appropriate equation depending on the chosen diffusion sampler and a predictor with $A_0$ as input. We have empirically observed this assumption. We compared the results of conditional generation using a predictor with $A_0$ as input alongside the appropriate equation of the chosen diffusion process, with the results obtained using a predictor with the $A_t$ as input (original version). We conducted experiments on the CIFAR-10 dataset within the NB201 search space under the Transferable NAS scenario and generated 100 neural architectures using a predictor with input $A_t$ and a predictor with input $A_0$, respectively. The statistics analyzed from the generated neural architectures are provided in **Table R1**. As shown in **Table R1**, when guiding neural architecture generation using the predictor with input $A_0$, the max accuracy is **93.8633%**, which is lower than the max accuracy of **94.3733%** obtained when guiding with the predictor with input $A_t$. This observation also holds for mean accuracy and min accuracy. Additionally, upon ranking the 100 generated neural architectures based on their predicted accuracies with the predictor, the average accuracy of the Top-5 neural architectures is **94.0440%** when guided by the predictor with input $A_t$, whereas it is significantly lower at **88.4357%** when guided by the predictor with input $A_0$.

**<Table R1>**
|                                                     | Predictor with input $A_t$ | Predictor with input $A_0$ |
|:---------------------------------------------------:|:--------------------------:|:--------------------------:|
|                   Max Accuracy (%)                  |         **94.3733**        |           93.8633          |
|                  Mean Accuracy (%)                  |         **93.1460**        |           85.5664          |
|                   Min Accuracy (%)                  |         **88.2650**        |           54.5467          |
| Average Accuracy of Top-5  sorted architectures (%) |         **94.0440**        |           88.4357          |

Reviewer comment (Experiments):
One can start with discrete NAR transformers for instance (like MaskGIT) which can model the NAS process and token predictions, and then the produced model will be used for generation which is another variant of cross entropy minimization problem. Since the foundation models are increasing in size day by day it will be interesting to see if their training methods can be used to find optimal sparse architectures.

Author response:
**Q2. The feasibility of applying DiffusionNAG to generate sparse foundation models.**

- To explore efficient sparse architectures for foundation models, we can leverage DiffusionNAG, similar to the research idea discussed in our response to **Q1. The feasibility of applying DiffusionNAG to generate optimal generative models**. Among various potential directions, one illustrative example is to consider DiffusionNAG as a post-training structured pruning approach designed to induce sparsity in Large Language Models (LLMs). This post-training structured pruning method starts from pre-trained models and eliminates redundant architecture elements along with their parameters. We find this approach reasonable, particularly when considering the impracticality of iteratively pre-training or iteratively pruning models at the scale of LLMs with hundreds of billions of parameters.
- Specifically, since the majority of LLMs are transformer-based models, we can explore sparse architectures within the search space that was defined in our response to Q1. Our goal would be to identify the optimal architecture that minimizes the accuracy drops within a resource budget (e.g., memory footprint or inference time). In other words, we aim to find the best architecture that enhances inference speed or reduces memory footprint without compromising accuracy.
- We can address this problem by employing the proposed DiffusionNAG method. First, we can train the score model of DiffusionNAG in a manner similar to that explained in Q1. Secondly, we can formulate a multi-objective function to evaluate the performance of sparse architecture candidates. For example, we may consider existing multi-objectives proposed by ProxylessNAS [8] or MNASNet [9] to find an optimal balance between accuracy and latency. Using the multi-objective function, we can train the predictor to predict the performance of given architecture candidates.
- Guided by the trained predictor, DiffusionNAG can generate an optimal sparse architecture (graph). In the final step, we selectively activate only the architectural components and parameters of LLMs that align with the generated graph. In other words, we prune (deactivate) architectural components such as blocks, MLP layers, and attention layers that do not correspond to the graph. This process results in a resource-efficient model, enhancing the model's sparsity and contributing to improved efficiency in terms of inference latency or memory utilization.

---
**References**

[8] ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware, ICLR 2019.

[9] MnasNet: Platform-Aware Neural Architecture Search for Mobile, CVPR 2019.

---
We hope that the discussions and additional experiments in this response further strengthen our argument and address the concerns raised by Reviewer EDN5. Thanks to the constructive comments from the reviewer, we have improved the clarity of the technical contributions of our proposed score network. Furthermore, we are delighted to engage in discussions about the exploration of diverse future research inspired by our approach. Particularly, we appreciate the interesting suggestions: 1) Beyond the existing NAS benchmarks focused on classification models, exploring the research opportunity to generate a generative model using our generative model like self-evolving AI. 2) Investigating the applicability of our work to make the foundational model sparse. We included our discussion regarding the future works in **Section E of the Appendix** in the revision. We are grateful for the time and effort spent reviewing our paper. If there is any additional information we should provide, please let us know.

Reviewer comment (Writing):
I would not say that I have the best time understanding the parts about "Score Network for Neural Architectures": what exactly is being pursued here? if our adjacency matrix is defined to be upper triangular, should it not already be a DAG? What is the goal of the proposed positional embedding? A bit more explanation is welcomed here.

Author response:
**Q. Explain in more detail about the score network for neural architectures.**

- To elaborate, our primary objective is to generate **valid** neural architectures. While an upper triangular adjacency matrix representing a DAG helps in defining a neural architecture, it alone falls short in capturing the intricate data distribution of neural architectures within the diverse search spaces inherent to NAS tasks. The complexity arises due to the need to adhere to specific rules that vary across different search spaces, as outlined in the **[Validity Criteria]** section, as well as the computational flow from input to output layers.

- **[Validity Criteria]** For generated neural architectures to be valid within the MBv3 search space, the following rules must be satisfied:
  - 1) Single operation per layer: Each layer should correspond to only one type of operation, avoiding the presence of multiple operations for a single layer.
  - 2) No empty stage: The MBv3 search space defines a fixed macro skeleton with five stages. Each stage must consist of at least two layers to prevent the occurrence of an empty (blocked) stage. This configuration, with a directed edge from the node immediately before it, ensures a continuous main path connecting all nodes.
  - 3) No isolated paths: Within each stage, no layer should lack predecessors, ensuring that there are no isolated paths.

- Therefore, to generate valid graphs representing neural architectures, the score network should capture 1) the dependency between nodes, reflecting the computational flow, and 2) the accurate position of each layer within the overall neural architecture to comply with the rules of a specific search space. **While a DAG representation using an upper triangular matrix might capture the former condition, it might fail to encode the accurate position of each layer to follow complicated rules.**

- To address these challenges, as detailed in Section 2 of the main paper, we have designed a novel score network, which employs transformer blocks and incorporates an upper triangular matrix describing node dependencies within a neural architecture. Moreover, we have introduced **positional embeddings (as depicted in Eq (4)) to more accurately capture the topological ordering of layers in architectures**, which leads to the generation of valid architectures adhering to specific rules within the given search space.

- Empirically, as shown in Table 4 in the main paper, we encountered difficulties in generating valid neural architectures when positional embeddings were absent in the MBv3 search space (**42.17% validity**). However, upon incorporating positional embeddings to capture the topological ordering of neural architectures, we observed a **100% validity** in the MBv3 search space. These experimental results show that positional embeddings are necessary to effectively capture the data distribution of neural architectures within the given search space.
 
---

If you have any further questions or if there is anything else you would like clarification on, please feel free to let us know. We are more than happy to provide additional explanations. Thank you for your time and effort in reviewing our paper.

Reviewer comment (Experiments):
The strength of the predictor guidance is an important hyperparameter. In this paper, a good performance can be achieved by cross-validation. Could the author provide some ablation studies of how this hyperparameter affects the performance? How could we effectively define the search range for cross-validation?

Author response:
**W1-2. Ablation studies on the impact of guidance strength.**

- To empirically verify how this guidance strength affects the performance and further provide insights for future follow-up studies, we compared the neural architecture distributions generated using our sampling strategy (**balanced exploitation-exploration sampling strategy**), both when the guidance scale was **fixed at a large value** and when it was **fixed at a small value** (Please see **Table R2**). We conducted experiments on the CIFAR-10 dataset within the NB201 search space under the Transferable NAS scenario. The fixed large and small values correspond to the maximum and minimum guidance strengths in our strategy of gradually decreasing guidance strength with a scale 5000, respectively (i.e., 15000, 10000, 5000). As we mentioned in **Section 3.1** in the main paper, after generating neural architectures by each sampling strategy, we rank the generated architecture pools using a meta-learned dataset-aware predictor and also provide the average accuracy of the Top-5 neural architectures.

**<Table R2>** Distribution of generated neural architectures across various sampling strategies. All methods generated neural architectures within the same sampling time. 
|                                            | Fixed at a large value (15000) | Balanced Exploitation-exploration sampling  (Ours) | Fixed at a small value (5000) |
|:--------------------------------------------------:|:------------------------------:|:--------------------------------------------------:|:-----------------------------:|
| Average Accuracy of Top-5 sorted architectures (%) |             93.262           |                     **94.156**                    |            93.978            |
|               Standard Deviation (%)               |             0.395             |                       0.658                       |             0.953            |


- As evident from **Table R2**, the average accuracy of Top-5 neural architectures generated with the **Fixed at a large value** strategy  (**93.262%**) is lower than the average accuracy of Top-5 neural architectures (**94.156%**) achieved by **Ours**, possibly due to heavily engaging in exploitation. On the other hand, fixing the guidance strength to one small value allows for greater engagement in exploration, leading to the generation of neural architectures with more diversity. However, since it guides towards regions far from the target sub-region, the average accuracy of the generated Top-5 neural architectures (**93.978%**) is lower than the result obtained by **Ours**, i.e., when the guidance strength is 0, it can be considered equivalent to random sampling. 
- Given the trade-off between exploration and exploitation depending on the guidance strength, we progressively adjust it during the generation of neural architectures. This sampling approach offers the advantage of balancing exploitation-exploration by gradually reducing the guidance strength.

Reviewer comment (Experiments):
This work only reports results on two NAS benchmark search spaces. It is unclear whether it can achieve similar performance on much larger search spaces, such as the ones proposed in [1] and [2], which are widely used in previous works.

Author response:
**W1. It is unclear whether DiffusionNAG can achieve good performance on much larger search spaces.**
- Thank you for your valuable feedback. In response to the reviewer's concerns, we would like to clarify that **we have already conducted extensive experiments in the MobileNetV3 (MBv3) search space ($>10^{19}$), which is larger than DARTS ($\approx 10^{18}$) search space**. Through these experiments, we clearly demonstrated that DiffusionNAG achieved good performance in a larger search space (**Table 1** and **Figure 3** of the main paper).

- **Background**: We provide a more detailed explanation of the MobileNetV3 (MBv3) search space. Recently, like the search spaces of NASNet or DARTS, the MBv3 search space has been widely utilized in the NAS community as one of the large search spaces [1-10]. Specifically, alongside NB201, the MBv3 search space serves as an established benchmark for our main experimental setup, Transferable NAS [4]. Therefore, we validated our approach in the same experimental setup with existing Transferable NAS methods to ensure a fair comparison. Additionally, this search space has been actively employed in Bayesian Optimization (BO)-based NAS [9].
- **Experimental Results**: To save the reviewer’s time, we provide a summary of the experiments conducted for two key predictor-based NAS scenarios in the MBv3 search space as follows:
  - **1) Transferable NAS**: DiffusionNAG achieved state-of-the-art (SOTA) performance on three out of four benchmark datasets (CIFAR-10, CIFAR-100, Aircraft, and Pets) within the MBv3 search space, as compared to the baselines (please refer to Table 1 in the main paper).
  - **2) BO-based NAS**: We have shown that DiffusionNAG is an efficient and effective alternative to existing heuristic acquisition optimization strategies, such as random architecture sampling or architecture mutation, within the MBv3 search space (please refer to the left side and middle side of Figure 3 in the main paper). Specifically, the BO-based NAS approach employing DiffusionNAG as the acquisition optimization strategy surpassed existing BO-based NAS methods in experiments conducted on the ImageNet 1K dataset within the expansive MBv3 search space (please refer to the right side of Figure 3 in the main paper).

- **Overview of the experiments we carried out**: Finally, we politely ask the reviewer to consider that we conducted extensive experiments as follows:
  - 2 search spaces, including a large search space such as MBv3, which contains $\approx 10^{19}$ neural architectures.
  - 5 different datasets, including a large-scale dataset, ImageNet 1K.
  - 2 key predictor-based NAS scenarios on 2 search spaces.
     - [Transferable NAS scenario] Transfer NAS experiments with a meta-learned dataset-aware predictor for 4 unseen datasets on both search spaces:  Main Paper Tables 1, 2.
     - [BO-based NAS scenario] Improving existing BO-based NAS experiments for 2 datasets and 4 acquisition functions on both search spaces: Main Paper Figures 3, 4 & Appendix Figure 6.
  - 4 ablation studies/analyses in the main paper.
    - Analyzing the distribution of generated neural architectures: Main Paper Table 3 & Figure 2 & Appendix Figure 5.
    - Analyzing the effectiveness of our proposed score network for neural architectures: Main Paper Table 4.
    - Analyzing adaptation ability across different objectives: Main Table 5.
    - Analyzing the distribution of operation types in generated neural architectures across different objectives: Appendix Figure 7.

  We would be grateful if the reviewer could take these factors into account.

***
#### **References**

[1] Searching for mobilenetv3, ICCV 2019.

[2] Once-for-all: Train one network and specialize it for efficient deployment, ICLR 2020.

[3] NSGANetV2: Evolutionary multi-objective surrogate-assisted neural architecture search, ECCV 2020.

[4] Rapid neural architecture search by learning to generate graphs from datasets, ICLR 2021.

[5] HELP: Hardware-adaptive efficient latency prediction for nas via meta-learning, NeurIPS 2021.

[6] Task-Adaptive Neural Network Search with Meta-Contrastive Learning, NeurIPS 2021.

[7] Neural Architecture Transfer, TPAMI 2021.

[8] CompOFA – Compound Once-For-All Networks for Faster Multi-Platform Deployment, ICLR 2021.

[9] Transfer NAS with meta-learned bayesian surrogates, ICLR 2023.

[10] MathNAS: If Blocks Have a Role in Mathematical Architecture Design, NeurIPS 2023.

Reviewer comment (Experiments):
Fig. 3 and 4 compare results on existing AO strategies and various acquisition functions. However, there is no comparison to other methods on ImageNet, as shown in Tab. 1 and 2.

Author response:
**W2. In BO-based NAS scenarios (Figs 3 and 4), there is no comparison to other methods on ImageNet.**

- We believe the reviewer may have overlooked **the right side of Figure 3**, which **contains comparison results with other baseline methods on ImageNet**. Specifically, we conducted a performance comparison between our proposed method and well-known NAS baseline methods such as BOHAMIANN [11], BONAS [12], EVOLUTION [13], LOCAL SEARCH [14], and RANDOM, highlighting our superiority.
- Furthermore, as illustrated in **Table R1** below, to address the reviewer's concern, we enhance the comprehensiveness of our results by introducing new baselines in addition to those featured in the right subfigure of Figure 3, resulting in **7 baseline methods** evaluated on ImageNet 1K. DiffusionNAG consistently outperforms other methods, showcasing superior performance across the same **Trained Archs** (the number of architectures requiring full training for obtaining final accuracy). The final performance reaches **83.55%**, surpassing that of the existing methods. 


**<Table R1>**
|    Method \ Trained Archs    |     40    |     60    |     80    |    100    |
|:----------------------------:|:---------:|:---------:|:---------:|:---------:|
| RANDOM                       |   82.51   |   82.58   |   82.60   |   82.72   |
| GCN Pred.                    |   82.51   |   82.58   |   82.60   |   82.76   |
| EVOLUTION                    |   82.68   |   83.00   |   83.15   |   83.21   |
| LOCAL SEARCH                 |   83.21   |   83.30   |   83.36   |   83.43   |
| BONAS                        |   82.71   |   82.98   |   83.09   |   83.25   |
| DNGO                         |   82.87   |   83.09   |   83.25   |   83.49   |
| BOHAMIANN                    |   82.92   |   83.21   |   83.36   |   83.44   |
| **Ours (DiffusionNAG + BO)** | **83.29** | **83.43** | **83.52** | **83.55** |


- We have updated **the right subfigure of Figure 3** by including additional baseline results in **the revision**. 

We appreciate the reviewer's comments since we believe that these additional results further strengthened our paper.


***
#### **References**
[11] Bayesian optimization with robust Bayesian neural networks, NeurIPS 2016.

[12] Bridging the gap between sample-based and one-shot neural architecture search with bonas, NeurIPS 2020.

[13] Regularized evolution for image classifier architecture search, AAAI 2019.

[14] Exploring the loss landscape in neural architecture search, UAI, 2021.

Reviewer comment (Experiments):
In Tab. 5, the authors include robust accuracy against the APGD attack, which is good. However, the corruption evaluation is insufficient, considering only glass blur. It would be better to include more comprehensive corruptions, such as those in ImageNet-C or CIFAR-10-C.

Author response:
**W3. It would be better to include more corruptions, such as those in ImageNet-C or CIFAR-10-C (Table 5).**

- Thank you for your valuable comments, especially for the positive feedback on our experiments on generating robust neural architectures against APGD attacks. Additionally, following the reviewer’s suggestion, we conducted further experiments to explore a broader range of corruptions introduced not only in glass blur but also in various corruptions in CIFAR-10-C. As evident in **Table R2**, **we conducted experiments to find robust neural architectures for 14 diverse corruptions at severity level 2 in CIFAR-10-C**. Similar to **Table 5** of the main paper, we demonstrated the capabilities of DiffusionNAG to easily adapt to new tasks without retraining the score network. We achieved this by replacing task-specific predictors in a plug-and-play manner. 
- Specifically, we train a clean predictor, $f_{\text{clean}}$, to predict the clean accuracy of neural architectures using randomly sampled pairs of architectures and their corresponding CIFAR-10 clean accuracy from the Robustness-NB201 benchmark [15]. Afterward, we create the training dataset by including pairs of architectures and their respective robust accuracy for each corruption scenario from the Robustness-NB201 benchmark [15]. Then, we train the target corruption predictor, $f_{\text{corruption}}$ using the training dataset to predict the robust accuracy of neural architectures under specific corruption types. 
- With the guidance of the trained $f_{\text{clean}}$, we generate a pool of 20 architectures ($S_{\text{clean}}$).  Similarly, with the guidance of the task-specific predictors ($f_{\text{corruption}}$), we generate another pool of 20 architectures for each target corruption ($S_{\text{corruption}}$). 
Subsequently, we retrieve the robust accuracy on both $f_{\text{clean}}$-guided and  $f_{\text{corruption}}$-guided architecture pools ($S_{\text{clean}}$ and $S_{\text{corruption}}$) from the Robustness-NB201 benchmark [15] for each corruption scenario.

We report the highest robust accuracy, as below:

**<Table R2>**

|                             |     brightness    |       contrast       |   defocus blur  | elastic transform |       fog      |   frost  | gaussian blur |             |
|-----------------------------|:-----------------:|:--------------------:|:---------------:|:-----------------:|:--------------:|:--------:|:-------------:|:-----------:|
| **$f_{\text{clean}}$**      |        89.3       |         58.9         |       64.6      |        61.0       |    **74.7**    | **60.2** |      30.5     |             |
| **$f_{\text{corruption}}$** |      **89.4**     |       **59.6**       |     **64.9**    |      **61.4**     |      74.4      |   58.9   |    **49.6**   |             |
|                             | **impulse noise** | **jpeg compression** | **motion blur** |    **pixelate**   | **shot noise** | **snow** | **zoom blur** | **average** |
| **$f_{\text{clean}}$**      |        55.2       |         61.6         |       43.7      |        71.7       |      39.7      |   70.4   |      44.9     |     59.0    |
| **$f_{\text{corruption}}$** |      **57.7**     |       **64.6**       |     **49.2**    |      **72.1**     |    **51.0**    | **71.2** |    **52.9**   |   **62.6**  |

- As shown in **Table R2**, the maximum accuracies of the architecture pool generated using the predictor $f_{\text{corruption}}$ trained for the target corruption surpass those obtained using the accuracy predictor $f_{\text{clean}}$ trained for clean accuracy. For **12** out of the 14 corruptions, excluding fog and frost corruption, the maximum accuracy of the architecture pool guided by the $f_{\text{corruption}}$ predictor exceeds that of the pool guided by $f_{\text{clean}}$. The average maximum accuracy increased by **3.6%**, reaching **59.0%** and **62.6%**, respectively. 
- We included the experiments of **Table R2** in **Table 10 of Appendix Section D.5** in the revision. 

---
**Reference**

[15] Neural architecture design and robustness: A dataset, ICLR 2023.

---
We believe that our paper became much stronger after including the additional experiments you have suggested. Thank you for the constructive comments. Please inform us if there are any others you would like us to address. If you believe the concerns are addressed, we would appreciate it if you could adjust the rating accordingly.