Paper: Differential Model Scaling using Differential Topk (ICLR-2024)

Reviewer comment (Evaluation):
Shown in Table 6, the performance gain is marginal.

Author response:
Sorry for the confusion. There are two reasons why our performance gain seems 'marginal'.

1. Compared with the Multi-Shot and One-shot methods, our methods have a much lower search cost, up to several tens, even hundreds of times lower. Therefore, our performance gain is not marginal because we may use much lower search costs to achieve a higher performance.
2. Compared with zero-shot NAS methods, they all use much stronger training recipes than ours, such as distillation and mixup. Therefore, the real improvement gain is more significant than the table shows.

Here, we provide a new table to make the comparison clearer (We only list some methods that may be confusing; please refer to Table 5 in our revised paper for full comparison.)

| Method                   | NAS Type  | Top1     | MACs  | Params | Cost |
| ------------------------ | --------- | -------- | ----- | ------ | ---- |
| Zen-score (+)            | ZeroShot  | 78.0     | 0.41G | 5.7M   | Low  |
| **DMS-EN-B0 (ours)**     | Gradient  | **78.5** | 0.39G | 6.2M   | Low  |
| Zen-score （+）          | ZeroShot  | 79.1     | 0.60G | 7.1M   | low  |
| ModelAmplification-EN-B1 | MultiShot | 79.9     | 0.68G | 8.8M   | High |
| **DMS-EN-B1 (ours)**     | Gradient  | **80.0** | 0.68G | 8.9M   | Low  |
| ScaleNet-EN-B2           | OneShot   | 80.8     | 1.6G  | 11.8M  | High |
| Zen-score (+)            | ZeroShot  | 80.8     | 0.9G  | 19.4M  | Low  |
| ModelAmplification-EN-B2 | MultiShot | 80.9     | 1.0G  | 9.3M   | High |
| BigNAS-XL                | OneShot   | 80.9     | 1.0G  | 9.5M   | High |
| **DMS-EN-B2 (ours)**     | Gradient  | **81.1** | 1.1G  | 9.6M   | Low  |

(+) means the model is trained using much stronger training tricks, such as distillation and mixup. As some methods do not report their search costs, we use "High" and "Low" to provide a rough comparison.

Thanks for your comments. We updated the table (Note the index is 5 in our revised paper, while it's 6 in the original paper).
We also added an accuracy vs MACs plot to the paper, notated as Figure 3 in our revised paper.

Reviewer comment (Evaluation):
accuracy vs MACs Plot with Table 6. The performance gain seems marginal from the table.

Author response:
Sorry for the confusion. There are two reasons why our performance gain seems 'marginal'.

1. Compared with the Multi-Shot and One-shot methods, our methods have a much lower search cost, up to several tens, even hundreds of times lower. Therefore, our performance gain is not marginal because we may use much lower search costs to achieve a higher performance.
2. Compared with zero-shot NAS methods, they all use much stronger training recipes than ours, such as distillation and mixup. Therefore, the real improvement gain is more significant than the table shows.

Here, we provide a new table to make the comparison clearer (We only list some methods that may be confusing; please refer to Table 5 in our revised paper for full comparison.)

| Method                   | NAS Type  | Top1     | MACs  | Params | Cost |
| ------------------------ | --------- | -------- | ----- | ------ | ---- |
| Zen-score (+)            | ZeroShot  | 78.0     | 0.41G | 5.7M   | Low  |
| **DMS-EN-B0 (ours)**     | Gradient  | **78.5** | 0.39G | 6.2M   | Low  |
| Zen-score （+）          | ZeroShot  | 79.1     | 0.60G | 7.1M   | low  |
| ModelAmplification-EN-B1 | MultiShot | 79.9     | 0.68G | 8.8M   | High |
| **DMS-EN-B1 (ours)**     | Gradient  | **80.0** | 0.68G | 8.9M   | Low  |
| ScaleNet-EN-B2           | OneShot   | 80.8     | 1.6G  | 11.8M  | High |
| Zen-score (+)            | ZeroShot  | 80.8     | 0.9G  | 19.4M  | Low  |
| ModelAmplification-EN-B2 | MultiShot | 80.9     | 1.0G  | 9.3M   | High |
| BigNAS-XL                | OneShot   | 80.9     | 1.0G  | 9.5M   | High |
| **DMS-EN-B2 (ours)**     | Gradient  | **81.1** | 1.1G  | 9.6M   | Low  |

(+) means the model is trained using much stronger training tricks, such as distillation and mixup. As some methods do not report their search costs, we use "High" and "Low" to provide a rough comparison.

Thanks for your comments. We updated the table (Note the index is 5 in our revised paper, while it's 6 in the original paper).
We also added an accuracy vs MACs plot to the paper, notated as Figure 3 in our revised paper.

Reviewer comment (Reproducibility):
explain how this loss_resource is designed.

Author response:
Resource constraint loss is widely used in the field of NAS and pruning [1,2]. We just follow the common practice in the field.

Take constraining the number of parameters of a fully connected layer as an example. Suppose the layer has $f_{in} \in R$ input features and $f_{out} \in R$ output features. The structure parameters for its input features and output features are $a_{in} \in [0,1]$ and $a_{out} \in 0,1$, respectively. The resource constraint loss for the layer is defined as $f_{in} \times a_{in} \times f_{out} \times a_{out}$. The resource constraint loss of the total model is the sum of all parameter losses of all layers.

We also detailed the resource constraint loss in Appendix A.1.2. Please refer to it for more details.

[1] Li, Yanyu, et al. "Pruning-as-Search: Efficient Neural Architecture Search via Channel Pruning and Structural Reparameterization."  
[2] Gao, Shangqian, et al. "Disentangled differentiable network pruning." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022.

Reviewer comment (Reproducibility):
The paper does not provide the code or the implementation details of the proposed method, which makes it difficult to reproduce and verify the results.

Author response:
Considering the double-blind review policy, we can not open-source our code before the paper is accepted. However, we will open-source our code after the paper is accepted.

Reviewer comment (Experiments):
The paper lacks a proper control experiment to isolate the effect of the differential topk from other factors, such as the network architecture, the learning rate, and the data augmentation. It is possible that some of the improvements are due to these factors rather than the proposed method.

Author response:
We think we have controlled the factors you mentioned.
For architecture, we evaluate our method always with the same architecture as our baselines.
For learning rate and data augmentation, without extra explanation, we use the same training setting for our method and the baseline methods.
We detail the training setting in Appendix a.6. If you think where we have not controlled the factors, please let us know.

Reviewer comment (Experiments):
The paper introduces too many hyperparameters for the differential topk, such as the temperature, the sparsity ratio, and the regularization coefficient. The paper does not provide a systematic analysis of how these hyperparameters affect the performance and the stability of the method. It is also unclear how to choose these hyperparameters for different tasks and architectures.

Author response:
Sorry for the confusion. We think our hyperparameters are not too many, as there are only two hyperparameters that need to be turned for different tasks and architectures, and the other hyperparameters are fixed.

We detail our hyperparameters as follows:

We first introduce the fixed hyperparameters.

1. Decay for Taylor importance. Taylor importance[1,2] is a well-known method to measure the importance of elements, and the decay of moving average is also widely used in the literature. Therefore, we directly use the value from prior work.
2. Temperature of our diffenretial topk. The temperature is used to polarize [3] the mask of elements. This is also a widely used approach for pruning. Directly selecting a value that can polarize the mask of elements is enough. Thanks to our importance normalization, the temperature can be directly computed by closed-form, detailed in section 3.1.2. Therefore, it's also fixed.
3. Supernet size and sparsity ratio. In our opinion, the supernet size and sparsity ratio are not hyperparameters. They are set according to the user's demand and their resource amount.

Only two hyperparameters need to be turned for different tasks and architectures.
1. Resource constraint loss weight, denoted by $\lambda_{resource}$
2. Learning rate for structure parameters denoted by $lr_{structure}$
   
They are used to control the update of the structure parameters. The update value of a structure parameter is computed by $lr_{strucutre}\times (g_{task} + \lambda_{resource}\times g_{resource})$, where $g_{task}$ and $g_{resource}$ is the gradient of structure parameters with respect to the task loss and resource constraint loss,
We provide an ablation study as follows:

| - | 5e-2 | 5e-3 | 5e-4 | 5e-5 |
|--|------|------|------|------|
| 0.1 | / | / | / | / |
| 1 | 73.0 | 73.1 | 72.9 | / |
| 10 | 72.5 | 72.2 | 72.6 | 70.9 |

The head row is the $lr_{structure}$, while the head column is the $\lambda_{resource}$. '/' indicates that the model cannot reach our target resource constraint.
Obviously, 

1. Smaller $\lambda_{resource}$ is better, as long as the model can reach the target resource constraint. Smaller $\lambda_{resource}$ means that the task loss takes more control of the update of the structure parameter.
2. When $\lambda_{resource}$ is small, the model is not sensitive to the change of $lr_{structure}$. When $\lambda_{resource}$ is large, a relatively large $lr_{structure}$ is better. This is because reaching the target resource constraint can reduce the influence of the resource constraint loss, as resource constraint loss is zero when the model reaches the target resource constraint.

Therefore, the setting of $\lambda_{resource}$ and $lr_{structure}$ is not difficult. We first fix $lr_{structure}$ and turn $\lambda_{resource}$ to a small value and ensure the model can reach the target resource constraint. Then, we turn $lr_{structure}$ to a relatively large value, which makes the model reach the target resource constraint in the first hundreds of iterations. Only observing the resource decrease in the first epoch is enough to set these two hyperparameters.

Compared with other NAS methods, our method uses fewer hyperparameters. For example, ModelAmplification must turn at least five hyperparameters for different tasks and models.

We added the ablation study to Appendix A.5.4 in our revised paper.

[1] Molchanov, Pavlo, et al. "Importance estimation for neural network pruning." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019.  
[2] Humble, Ryan, et al. "Soft masking for cost-constrained channel pruning." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022.  
[3] Neuron-level Structured Pruning using Polarization Regularizer

Reviewer comment (Theory):
In Section 3.1.3, you use a moving average to update the layerwise mask. What is the motivation and benefit of this technique? How does it affect the convergence and stability of the optimization?

Author response:
The moving average is used to make element importance measures more stable. As shown in the following table.
With more stable element importance, we can get a better result.
| Measure Method | Top-1 |
|----------------|-------|
| Taylor without moving average | 72.5 |
| Taylor with moving average | **73.1** |

Taylor importance with moving average is a common technique in the pruning field. Please refer to [1,2] for more details.
We also added the ablation study to Appendix A.5.3 in our revised paper.


[1] Molchanov, Pavlo, et al. "Importance estimation for neural network pruning." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019.  
[2] Humble, Ryan, et al. "Soft masking for cost-constrained channel pruning." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022.

Reviewer comment (Novelty):
As far as I know, gumbel top-k method is also differentiable, why you develop a new differentiable top-k methd?

Author response:
We are sorry for the confusion.
Gumbel top-k [1,2] is completely different from our method. It may be confusing because we use the same term, "top-k", in our paper.

Gumbel top-k is learned to sample 'k' different elements from a set of n elements, where 'k' is a constant.
Our differential topk is learning the 'k' itself, which is a learnable parameter.
Gumbel top-k cannot be used to search the width and depth of a network, while our method can.

We also talk about another method, "Gumbel softmax", which is also usually used in pruning and NAS [3].
"Gumbel softmax" is able to work as a learnable gate to select elements for width and depth. However, as we described in our paper, it models depth and width hyperparameter searching as an element selection problem, making it hard to optimize.

[1] Tan, Haoxian, et al. "Mutually-aware Sub-Graphs Differentiable Architecture Search." arXiv preprint arXiv:2107.04324 (2021).  
[2] Li, Jiaoda, Ryan Cotterell, and Mrinmaya Sachan. "Differentiable subset pruning of transformer heads." Transactions of the Association for Computational Linguistics 9 (2021): 1442-1459.  
[3] Herrmann, Charles, Richard Strong Bowen, and Ramin Zabih. "Channel selection using gumbel softmax." European Conference on Computer Vision. Cham: Springer International Publishing, 2020.

Reviewer comment (Reproducibility):
The open source of the code will help to understand the paper.

Author response:
Considering the double-blind review policy, we can not open-source our code before the paper is accepted. We will open-source our code after the paper is accepted.

Reviewer comment (Writing):
In p.5, the authors demonstrate that "Intuitively, $c_{i}^{'}$ indicates the portion of $c$ values larger than $c_{i}$.". Here should be "smaller" instead of "larger".

Author response:
Thanks for your correction. We have updated the paper to correct it.

Reviewer comment (Experiments):
Since the problem is cast as a NAS problem the search spaces used are not the ones very traditional to NAS (refer to questions)

Author response:
There are two main differences between the search space of our method and that of OFA and Autoformer.

1. What is more fine-grained in OFA's search space: We (and Autoformer) only search for the depth and width of networks, while OFA also searches for the resolution and kernel size. This is because we want to build a general method that can be applied to NLP and CV tasks. The resolution and kernel size are not necessary for NLP tasks.
2. What is more fine-grained in our search space: We can search width and depth with step 1 due to our high search efficiency. However, OFA and Autoformer have to use large steps and min-max limitations to compress their search space. This is an advantage of our method, as our search space is easier to define and much larger.

Our current submission uses the default setting directly, causing different search spaces. To make the comparison more fair, we select Autoformer-T as our baseline and use the exact same supernet and search space from Autoformer-T. The performance and estimated search cost are shown below.

| Method | Search Space | MACs | Top-1 | $Cost_{search}$ |
|--------|--------------|------|-------|-----------------|
| Autoformer-T | AutoFormer | 1.3 G | 74.7 | > 25 GPU days |
| DMS | Autorormer | 1.3 G | **75.2** | 2 GPU days |

Using Autoformer's search space, our method outperforms Autoformer-T by 0.5%  with less than 10% of the search cost. It proves the gain of our search method is from itself.

We have also added this ablation study to Appendix A.5.2 in our revised paper.

Reviewer comment (Experiments):
The search needs to be repeated for every resource constraint and obtaining a Pareto-Front of objectives might be very expensive (unlike methods like OFA[1] which directly approximate the whole Pareto-front)

Author response:
Yes, we agree with your opinion. Our method is not suitable for searching for the Pareto-Front.

However, we want to emphasize our target is to build a general and flexible NAS method. There are several situations where our method is more suitable than OFA.

1. Searching for a small number of models, which is the usual case in the real world. 
2. Searching extremely large models. For example, we prepare to train an LLM with 100B parameters using all our GPUs for 30 days. It's impossible to train a supernet (by OneShot NAS, like OFA), which needs much over 30 days, such as 60 or even 300 days. But using much less than 30 days, such as 1 day or 3 days , to search for it with our method is acceptable.

Last, thank you for your helpful reviews! Looking forward to hearing back from you!

Reviewer comment (Reproducibility):
Since the observation from section 5 shows that initialization from pre-trained models is very useful for differential scaling, did the authors include this in the search time computation? If a method relies on pre-trained models, then ideally the pre-training cost is a part of the total cost incurred. Could the authors clarify the initialization scheme used in each of the tables i.e., table 1, 2, 3?

Author response:
We are sorry for the confusion. 
To clarify, we did not load pretrained models by default. We only load pretrained models in Table 3 (LLM experiments), Table 4 (ablation study), and Table 8 (comparison with other pruning methods). Therefore, we did not include the pretraining cost in Table 1 for our method as we did not use pretraining, and we think the comparison is fair for other methods.
We chose this setting in experiments because, based on our ablation study, we found searching from a large, randomly initialized supernet reaches higher accuracy but lower cost than searching from a small pretrained supernet. Details are in the table below.

| Supernet | $Iinit_{search}$ | $cost_{pretrain}$ | $cost_{search}$ | $cost_{total}$ | Top-1 |
|----------|------------------|-------------------|-----------------|----------------|-------|
| ResNet-50 | Random | 0 | 41 | 41 | 73.1 |
| ResNet-50 | Pretrain | 410 | 41 | 451 | 73.8 |
| ResNet-152 | Random | 0 | 116 | 116 | **74.6** |

We have updated this table to section 5 to make our conclusion more clear.

Reviewer comment (Experiments):
In the appendix the authors compare with one-shot methods like OFA [1]. The comparison, in the reviewer's opinion, is unfair since the search is performed on different search spaces. Could the authors evaluate the method on the exact same search space as OFA? This would help differentiate the gains of the search-space versus the method itself. Similarly, could a comparison be made with the AutoFormer [2] by evaluating the method on its exact search space [2]?

Author response:
There are two main differences between the search space of our method and that of OFA and Autoformer.

1. What is more fine-grained in OFA's search space: We (and Autoformer) only search for the depth and width of networks, while OFA also searches for the resolution and kernel size. This is because we want to build a general method that can be applied to NLP and CV tasks. The resolution and kernel size are not necessary for NLP tasks.
2. What is more fine-grained in our search space: We can search width and depth with step 1 due to our high search efficiency. However, OFA and Autoformer have to use large steps and min-max limitations to compress their search space. This is an advantage of our method, as our search space is easier to define and much larger.

Our current submission uses the default setting directly, causing different search spaces. To make the comparison more fair, we select Autoformer-T as our baseline and use the exact same supernet and search space from Autoformer-T. The performance and estimated search cost are shown below.

| Method | Search Space | MACs | Top-1 | $Cost_{search}$ |
|--------|--------------|------|-------|-----------------|
| Autoformer-T | AutoFormer | 1.3 G | 74.7 | > 25 GPU days |
| DMS | Autorormer | 1.3 G | **75.2** | 2 GPU days |

Using Autoformer's search space, our method outperforms Autoformer-T by 0.5%  with less than 10% of the search cost. It proves the gain of our search method is from itself.

We have also added this ablation study to Appendix A.5.2 in our revised paper.