Weight-Entanglement Meets Gradient-Based Neural Architecture Search

We thank the reviewer for the detailed comments and suggestions. We respond to each of your questions below:

The motivation behind this work is not clear, and it appears to be a simple combination of existing methods, lacking innovation. What is the necessity and advantage of using single-stage search?

We thank the reviewer for their question. To the best of our knowledge, our work is the first to efficiently and effectively apply single-stage methods to weight-entanglement/foundation-model macro spaces. This makes our approach novel and very useful in practice.

Our work is primarily motivated by two goals:

Firstly, we aim to make single-stage NAS methods amenable to foundation model macro spaces (language models, vision transformers, mobilenets). Note that single stage methods, even recent ones like Lambda-DARTS, have only been evaluated on cell-based spaces which are not very realistic in deep learning deployment scenarios. In our opinion, it is crucial to make single-stage methods more accessible to foundation models to ensure that advancement in these methods do not remain restricted to cell-based spaces. Further, our plug-and-play framework and our publicly available code makes integration evaluation of any single-stage optimizer on these spaces easy.

Secondly, two-stage methods rely on performance proxies from the supernet which might not be highly correlated with the ground truth accuracies in practice. These supernets also face interference [1] during training which exacerbate these issues. We also conduct an analysis on the CKA[2] correlation between feature representations of the best architecture in our single-stage supernet and our two-stage supernet and find that our supernet in general shows better correlations.

[1] Shipard, J., Wiliem, A. and Fookes, C., 2022. Does Interference Exist When Training a Once-For-All Network?. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 3619-3628).

[2] Kornblith, S., Norouzi, M., Lee, H. and Hinton, G., 2019, May. Similarity of neural network representations revisited. In International conference on machine learning (pp. 3519-3529). PMLR.

We thank the reviewer for the detailed comments and suggestions. We respond to each of your questions below:

It would add value to the paper to demonstrate DARTS and/or GDAS…

In the single-stage neural architecture search literature [1] it is well-known that several of the single-stage optimizers like DARTS, GDAS etc.have many failure modes. We started our study on weight entanglement spaces by studying DrNAS, DARTS and GDAS on the two tiny search spaces which we create benchmarks for. We find from initial analysis that DrNAS performs competitively and is robust on both of these spaces in comparison to DARTS and GDAS. Hence we restrict ourselves to DrNAS in our large scale experiments. We have also evaluated DARTS and GDAS on the AutoFormer-T space on CIFAR-100 and presented the results below.

Toy conv macro (CIFAR10)

Toy cell entangled (Fashion MNIST)

AutoFormer-T (CIFAR100)

Since our focus is mainly 2-stage gradient based NAS without the use of zero-cost proxies we restrict ourselves to approaches which operate in this realm (unlike TE-NAS, which uses proxies for pruning).

On the macro search space, it would make sense to consider a macro approach and then introduce the cell-based part of the proposed method. It's not clear what the reference is for measuring the improvement. Nevertheless, the comparison with existing methods is useful.

We would like to clarify that none of the macro-spaces in our experimental setup (i.e. AutoFormer, MobileNet, LLM space) have any cell-based components. We operate on the same search spaces as AutoFormer or OFA only modifying the NAS optimizer itself with our proposed optimizer. The baselines we use are the two-stage counterparts (usually methods which are based on SPOS with some modifications). For completeness we compare against the AutoFormer’s original code, the Once-For-All’s original code as baselines when comparing their two-stage method with ours. For search spaces which do not have two-stage methods available we implement the Single-Path-One-Shot (SPOS) two-stage method in these spaces.

This is a benchmark driven approach …

We compare against DrNAS which improves just with the introduction of weight-entanglement. We don’t compare with other baselines simply because the goal of our work is to show the effective applicability of single-stage methods on weight-entanglement spaces (transformers, mobilenet) instead of achieving SOTA on cell-based spaces (DARTS, NB201). We hope that our flexible plug-and-play framework makes evaluation on new single-stage methods on more realistic and practical spaces like transformers, mobilenet easier.

Clarity

We thank the reviewer for this suggestion. We will provide more details about the method in our updated manuscript.

If the reviewer has any further questions we would be happy to address them.

Evaluation on other optimizers..

We have now evaluated DARTS and GDAS on the AutoFormer-T space for CIFAR10 too. Similar to the results on CIFAR100, here too TangleNAS-DrNAS performs best out of the three choices of optimizers. This motivates our choice of gradient-based NAS optimizer in our experiments.

AutoFormer-T (CIFAR10)

We thank the reviewer for the detailed comments and suggestions. We respond to each of your questions below:

It is unclear from the paper what are the FLOPs and parameter sizes of the searched results in the experiments. Are the comparisons done between networks?

We have now updated the manuscript to reflect these numbers (Table 5,6,7,8). In our experimental setup, we perform unconstrained search using both two-stage and single-stage methods (i.e. TangleNAS) and then compare their evaluations. The comparisons are done between individual networks in the supernet.

Furthermore now we also conduct an experiment to search for models with smaller parameter sizes. We add a differentiable parameter penalty to the loss function to enforce this. We then perform the evolutionary search of SPOS setting the parameter size of the tanglenas model as constraint. We find out that the architecture SPOS discovers in this constrained version is still worse than the tanglenas architecture discovered. This shows the ability of our search method to discover better architectures even at smaller parameter budgets in its constrained version.

Any reason why the performance of AutoFormer variants from that paper is not listed in Table 6? Again, there should be the FLOPs and parameter size of each model.

We primarily restrict our evaluation to AutoFormer-T due to compute restrictions. However we have now evaluated our method on the AutoFormer-S space and the results (with resolution 224 images) is in Table-6 of the updated manuscript. Our approach again outperforms the architecture derived by the corresponding two-stage method. We present the table-6 below:

On Table 7, there should be comparisons with some other works which use MobileNetV3 search space (e.g., AtomNAS), together with their FLOPs and parameter sizes. Why are the results of OFA with progressive shrink not used in the table?

We do use the OFA-PS supernet and report the accuracy after evolutionary search using their code. The search space variant we use is the smaller MobileNetV3 space, i.e. ofa_mbv3_d234_e346_k357_w1.0 here . To the best of our knowledge AtomNAS has a very different search space compared to OFA, hence we think that AtomNAS is not really comparable to OFA or our proposed approach.

If the reviewer has any further questions we would be happy to address them.

It would be beneficial to provide a more detailed explanation of the rationale behind the selection of specific search types, optimizers, and supernet types for the comparative evaluation.

In the single-stage neural architecture search literature [1], it is well-known that several of the single-stage optimizers such as DARTS and GDAS have several failure modes. We started our study on weight entanglement spaces by studying DrNAS, DARTS and GDAS on the two tiny search spaces which we create benchmarks for. We find in our initial analysis that DrNAS performs competitively and is robust on both of these spaces in comparison to DARTS and GDAS. Hence we restrict ourselves to DrNAS in our large scale experiments. We have also evaluated DARTS and GDAS on the AutoFormer-T space on CIFAR-100 and presented the results below.

Toy conv macro (CIFAR10)

Toy cell entangled (Fashion MNIST)

AutoFormer-T (CIFAR100)

We choose our search spaces with three primary factors in mind (1) Architecture Diversity: ViTs, LLMs, mobilenets i.e. convolutional space (2) Application Diversity: Classification and Language Modelling) (3) Usefulness: A focus on spaces constructed around foundation model architectures

[1] White, C., Safari, M., Sukthanker, R., Ru, B., Elsken, T., Zela, A., Dey, D. and Hutter, F., 2023. Neural architecture search: Insights from 1000 papers. arXiv preprint arXiv:2301.08727.

To discuss potential limitations ...

We thank the reviewer for this question. In our opinion the primary limitation of our work is the inability to handle multiple architecture constraints like parameter-sizes, FLOPS directly. It is indeed possible to add a constrained secondary penalty to our loss. However , this modification still doesn’t allow one to generate a full Pareto Front of objectives which is useful in many deep learning applications (eg: fairness, robustness, hardware constraints). We leave this extension for future work. Further, like 2-stage methods (OFA, AutoFormer), our approach also requires some amount of engineering effort to shard and combine weights into a mixture.

Furthermore now we also conduct an experiment to search for models with smaller parameter sizes. We add a differentiable parameter penalty to the loss function to enforce this. We then perform the evolutionary search of SPOS setting the parameter size of the TangleNAS model as the constraint. We find that the architecture SPOS discovers in this constrained version is still worse than the TangleNAS architecture discovered. This shows the ability of our search method to discover better architectures even at smaller parameter budgets in its constrained version.

If the reviewer has any follow-up questions we would be happy to address them.

We thank the reviewer for the detailed comments and suggestions. We respond to each of your questions below:

The novelty is limited…MergeNAS

We thank the reviewer for pointing out the "MergeNAS" paper, which we were unaware of at the time of writing the paper. We agree that on cell-based spaces MergeNAS works similarly to TangleNAS and we will update the manuscript to cite MergeNAS. However, MergeNAS does not apply and evaluate on MobileNet or transformer spaces ie. the weight entanglement spaces. We on the contrary extend weight entanglement to linear layers, batchnorm and layernorm layers, and attention blocks and show its effectiveness. Our work mainly aims at achieving two objectives. Firstly, making single-stage NAS methods amenable to the foundation model macro spaces (language models, vision transformers, MobileNets). This, in our opinion, is crucial to make single-stage methods more practically useful. Furthermore, our plug-and-play framework and our publicly available code makes integration and evaluation of any single-stage optimizer on these spaces easy.

Performance improvements are limited compared with the baselines

We think that our improvements on ImageNet 1-k on AutoFormer-T spaces and MobileNet spaces are quite significant. Notably, on the same AutoFormer-T space, TangleNAS finds an architecture with ~ 3.8% higher accuracy. We have now also evaluated TangleNAS on AutoFomer-S space where it achieves an improvement of ~ 0.264% compared to AutoFormer. Further, we have now added evaluations on the GPT-2 search space trained on OpenWebText. We transfer the pre-trained architectures (handcrafted and searched) to the Shakespeare dataset and observe that the architectures we discover outperform the larger handcrafted architecture while being efficient in terms of the number of parameters and latency.

How to entangle non-parameter operations, such as skip or pooling, as they have no weight compared to convolution?

We acknowledge that non-parameter operations cannot be entangled. However, we would like to point out that in this work, we focus primarily on entanglement spaces which do not have skip connections as a choice of operation. These spaces, which are designed around foundation model architectures, are in fact far more realistic and practically useful than cell-based spaces (which usually have skip operations). Our goal through weight entanglement is to save as much of GPU memory as possible when applying single-stage NAS methods.

The search cost of original DrNAS is 0.4 GPU-Days in DARTS search space. Whereas, in Table 4, the search time is 29.4 GPU-Hours. Can the authors explain the reason for the different search cost?

In our configurable framework we implement DrNAS with and without weight sharing for a unified comparison of search time and performance. The reported numbers are for the DrNAS optimizer on the DARTS search space using this unified codebase. The differences in search time are likely due to the differences in the GPU hardware setup (gpu types, cpu cores) we use in comparison to the authors.

The search cost of the original DrNAS is 0.4 GPU-Days…

We would like to thank the reviewer for raising the concern about the total time taken by our method. In response, we investigated the matter further, re-ran our experiments, and uncovered two important findings. Firstly, we misreported the total search time on the DARTS search space for DrNAS. Secondly, we had bottlenecks in our code which made our implementations slower than its original counterpart (such as the suboptimal number of workers in the data loader, and unnecessary iterations over all the modules of the model in every step of the optimization loop). Having fixed these, we are happy to report that our DrNAS with weight entanglement is now faster than the DrNAS baseline published by the authors. For a fair comparison against the authors’ baseline, we ran our code and theirs on the same hardware (RTX-2080 with 12 CPU cores), Python environment, and using the same experimental setup. Additionally, we ran experiments with and without the bottlenecks in our code fixed. The results are summarized in the table below, where WS indicates weight-sharing and WE indicates weight-entanglement.

Please note that DrNAS - WS (with bottleneck) and DrNAS - WE (with bottleneck) are what should have been reported in the paper. We thank the reviewer for giving us the opportunity to identify and fix this error.

Evaluation on other optimizers..

We have now evaluated DARTS and GDAS on the AutoFormer-T space for CIFAR10 too. Similar to the results on CIFAR100, here too TangleNAS-DrNAS performs best out of the three choices of optimizers. This motivates our choice of gradient-based NAS optimizer in our experiments.

AutoFormer-T (CIFAR10)