Paper: Improving MLP Module in Vision Transformer (ICLR-2024)

Reviewer comment (Experiments):
The AGeLU's standalone performance appears suboptimal. Table 2 suggests that the primary performance improvements stem from knowledge distillation or spatial enhancement, with the paper lacking a clear demonstration of the enhanced non-linearity's effectiveness, i.e., the proposed AGeLU.

Author response:
There are some misunderstandings and we are sorry for the unclarity. 

Firstly, **we do not use KD in our proposed method**. As shown in the first paragraph of section 4.3, using KD method can enhance the performance but will also increase the GPU memory during the training process, which is not feasible to use when the size of the model grows larger. Thus, we propose our spatial-enhancement part in section 4.3. The KD result in Table 2 is only proposed for the completeness of the experiment. This is explained more clearly in Table 2 of the revised version.

Secondly, the usefulness of the channel-wise enhancement part is shown in Table 2 and Table 3. The performance gains indeed come from depthwise convolution (spatial-wise enhancement part). However, **the proposed AGeLU+concat (channel-wise enhancement part) copes very well with dwconv** and can reduce FLOPs and parameters while maintaining accuracy to the greatest extent possible. For example, as shown in Table 2, DeiT-Ti with dwconv only increases 0.6% Top1 accuracy compared to original DeiT-Ti model (compare Line 1 and Line 4). However, by adding AGeLU+concat, DeiT-Ti with dwconv can increase 2.1% Top1 accuracy (compare Line 2 and Line 5) while saving over 10% FLOPs and parameters to the original model (Line 1). Thus, **the proposed spatial and channel enhancement part should be treated as a whole** to enhance the nonlinearity from two different aspects, and achieve the goal of reducing computational cost while maintaining accuracy. 

For Table 3, we ablate using the original MLP module (Line 1) and its AGeLU counterpart (Line 2). We also give **the proposed IMLP module with AGeLU+concat (Line 4)** and its GeLU counterpart (Line 3). The result shows that our method of using AGeLU+concat is very useful, and there are some explanations at the end of page 7. We add more explanations below. Line 1 and Line 2 have similar performance since they all have the same degree of nonlinearity according to Eq.2 (the form of original MLP), no matter $\phi=GeLU$ or $\phi$=AGeLU. The reason why Line 4 is much better than Line 3 comes from Eq.6. When using two different AGeLUs, we have different $\alpha_{c’}$ and $\alpha_{c’}'$ to form $t_1$ and $t_1’$ in Eq.6. Thus, the output can be expressed as the linear combination of $C’$ different nonlinear functions. However, when using two GeLU functions, we have $\alpha_{c’}=\alpha’_{c’}=1$. Thus, the output can only be expressed as the linear combination of $C’/2$ different nonlinear functions.

Reviewer comment (Reproducibility):
There is ambiguity regarding whether the kernel size is consistently set to 5 for large-scale models such as DeiT-B or Swin-B, lacking a definitive explanation in the paper.

Author response:
Sorry for the unclarity. In fact, we found that the ratio of the computational cost of the MLP module will increase as the scale of the model becomes larger, thus replacing the original MLP module the IMLP module will reduce more FLOPs and parameters and it is harder to make up for the accuracy drop with kernel size=3. Therefore, we consistently set kernel size to 5 for large-scale models.

Reviewer comment (Theory):
The remarkable performance drops attributed to the addition of GeLU with four times the number of channels raise questions. Since the fully connected (fc) layer entails linear calculations, it appears that the addition operation merely doubles the original output, which requires further clarification.

Author response:
I believe you are talking about the results shown in Table 4. In fact, the experiments in Table 4 use the following equation to replace the original MLP module, i.e., $y=AGeLU_1(x)+AGeLU_2(x)+\cdot\cdot\cdot+AGeLU_{4C}(x)$. The reason we conduct this ablation study 4 is shown at the beginning of section 4.1. Table 4 is a straightforward way to use the combination of $C^′=4C$ different nonlinear functions to replace the original MLP module. However, based on Corollary 1 (3) in the paper, the bias that depends on the input elements makes it challenging to attain a comparable degree of non-linearity by merely combining multiple nonlinear functions, and the classification performance does not match that of using the original MLP module. This is more clearly explained in the second last paragraph of section 5.2 in the revised version.

Reviewer comment (Experiments):
Introducing spatial enhancement after the GeLU operation (as 'a' in Figure 4) would help to conclusively demonstrate the impact of enhanced non-linearity in parameter reduction.

Author response:
If my understanding is correct, the result of using spatial enhancement after the GeLU operation is shown in Line 4 of Table 2 (w/ spatial).

Reviewer comment (Experiments):
In Equation (5), AGeLU and AGeLU′ are introduced as two nonlinear functions. It prompts an intriguing inquiry: what would the outcome be if the division was into more parts, say four? A more comprehensive ablation study should be conducted to provide a richer understanding of the behavior and performance of these functions.

Author response:
**A1:** This is a really good question. By dividing the MLP module into more parts, the computational complexity of the model can be reduced, while the degree of non-linearity remains the same which can be derived by simply extending Eq.6 in the main paper. However, as shown at the beginning of section 4.3, the degree of freedom is related to the number of channels in each part, and dividing the MLP module into more parts will decrease the degree of freedom and thus make it hard to achieve a good result. We conduct experiments using two and four nonlinear functions and the results are shown below. We can see that using four nonlinear functions gives a worse classification accuracy compared to using two nonlinear functions, which is consistent with our analysis.

|Model|	FLOPs|	Params|	acc|
|-|-|-|-|
|2 nonlinear functions|	1.10|	5.00|	72.6|
|4 nonlinear functions|	1.01|	4.55|	68.6|

Reviewer comment (Experiments):
The proposed IMLP module has only been experimented with a few models like ViT and Swin, which have been proposed for several years. It raises the question of the module's effectiveness on more recent, higher-accuracy models like iFormer. The validation of the IMLP module across a broader spectrum of models could provide a clearer picture of its versatility and efficacy in current vision transformer landscapes.

Author response:
**A3:** We conduct experiments not only on classic models that have been proposed for several years such as ViT and Swin, but also on the latest published popular models such as PoolFormer and portable models LVT which are all published in 2022. We conduct experiments on iFormer-S (NeurIPS 2022) model and the results are shown below. It shows that our method can reduce FLOPs and parameters while keeping the classification performance on iFormer.

|Model|	FLOPs|	Params	|acc|
|-|-|-|-|
|iFormer-S|	4.848|	19.866|83.4|
|iFormer-S + IMLP|	4.330 (-10.7%)|	17.528 (-11.8%)|	83.3|

Reviewer comment (Evaluation):
My major concern is the effectiveness, we can see some parameter and computational cost savings from Table 1, but this method introduced lots of hardware unfriendly operations like DW conv.

Author response:
**A1:** DW conv operation is **GPU unfriendly rather than hardware unfriendly**, since the optimization of matrix multiplication on GPU is very well and DW conv has insufficient parallelism on the GPU, not fully utilizing the computational ability. However, this drawback can be largely eased on other computing devices such as CPU, NPU and FPGA, and the advantage of IMLP with fewer FLOPs and parameters can appear. In fact, DW conv is widely used in mobile devices and there are many papers published to speed up dw conv [1], [2].

[1] High Performance Depthwise and Pointwise Convolutions on Mobile Devices. AAAI, 2020.

[2] Designing efficient accelerator of depthwise separable convolutional neural network on FPGA. JSA, 2019.

Reviewer comment (Evaluation):
From Table 3, I can't see a solid improvement from AGeLU.

Author response:
**A3:** Table 3 shows that AGeLU is not useful when using it alone (compare Line 1 and Line 2), and **this is not used in our method**. They are listed here for the completeness of this ablation study. In fact, our proposed channel-wise enhancement part with AGeLU+concat has a non-negligible 0.3% accuracy improvement compared to using GeLU+concat (compare Line 4 and Line 3). Thus, we need to **treat the channel-wise enhancement part as a whole and a solid improvement can be seen**. More explanation of Table 3 can be found at the end of page 7.

Reviewer comment (Evaluation):
Can you offer the actual latencies?

Author response:
Thank you for pointing this out. We list the throughput of LVT [1] on CPU and the proposed method in the following table, and hope this can address your concern.

|Model|	Throughput (imgs/s) on CPU|
|-|-|
|LVT-R1|	81.5	|
|LVT-R1+IMLP|	86.8 (+6.5%)|
|LVT-R2|	79.7|	
|LVT-R2+IMLP|	85.1 (+6.8%)|
|LVT-R3|	74.5|	
|LVT-R3+IMLP|78.9	 (+5.9%)|
|LVT-R4|71.6	|	
|LVT-R4+IMLP|	75.3 (+5.2%)|

[1] Lite Vision Transformer with Enhanced Self-Attention.

Reviewer comment (Evaluation):
The performance enhancements provided by the proposed activation function are marginal and non-existent in some instances. The proposed activation function fails to improve accuracy in larger models like Deit_B and LVT-R4; moreover, it actually leads to a decline in performance for Swin-B and Poolformer-M48.

Author response:
**A2:** The proposed IMLP aims to reduce the computational cost while maintaining the classification performance compared to the baseline model, rather than improve the accuracy only. As shown in **Figure 1**, the proposed method has **a better FLOPs-Accuracy trade-off** among all the baseline methods. This is because the FLOPs can reduce significantly with a slight or no performance drop.

Reviewer comment (Novelty):
The main issue identified by the reviewer is the performance gains of this work appear to depend largely on employing depthwise convolution, which has been already recognized in many prior hybrid architectures. The ablation studies presented in the manuscript further underscore this reliance as well. As a result, the paper's contribution is considered to be very limited.

Author response:
**A3:** This is also a very good question. The performance gains indeed come from depthwise convolution (spatial-wise enhancement part). However, **the proposed AGeLU+concat (channel-wise enhancement part) copes very well with dwconv** and can reduce FLOPs and parameters while maintaining accuracy to the greatest extent possible. For example, as shown in Table 2, DeiT-Ti with dwconv only increases 0.6% Top1 accuracy compared to original DeiT-Ti model (compare Line 1 and Line 4). However, by adding AGeLU+concat, DeiT-Ti with dwconv can increase 2.1% Top1 accuracy (compare Line 2 and Line 5) while saving over 10% FLOPs and parameters to the original model (Line 1). Thus, the proposed **spatial and channel enhancement part should be treated as a whole** to enhance the nonlinearity from two different aspects, and achieve the goal of reducing computational cost while maintaining accuracy. 

Besides, although many papers have used dwconv in their transformer, they do not give a thorough analysis of the reason it works (which is also mentioned at the end of page 1). In this paper, we treated **dwconv as the enhancement of nonlinearity** which explains the usefulness of dwconv in a new perspective.

Reviewer comment (Experiments):
It is speculated that the proposed method's effectiveness relies on KD (Table 2 evidently shows this), which requires a teacher model. Consequently, training budgets may not be preserved equally.

Author response:
A5: This is a misunderstanding and we are sorry for the unclarity. **We do not use KD in our proposed method**. As shown in the first paragraph of section 4.3, using KD method can enhance the performance but will also increase the GPU memory during the training process, which is not feasible to use when the size of the model grows larger. Thus, we propose our spatial-enhancement part in section 4.3. The KD result in Table 2 is only proposed for the completeness of the experiment. This is explained more clearly in Table 2 and the corresponding explanation of the revised version.

Reviewer comment (Theory):
The reviewer acknowledges that while the theories included do enhance the paper, it lacks a crucial explanation—specifically, the rationale behind why AGeLU with concatenation is necessary has not been addressed via theory.

Author response:
**A6:** The reason for using AGeLU with concatenation is shown in Eq.2 (the form of original MLP) and Eq.6 (the form of MLP using AGeLU with concatenation) and also in the last paragraph of section 4.1. I will try to explain it more clearly in the following. 
We can rewrite the last line of Eq.6 into the form:

$\mathbf y'=\mathbf t_2 \mathbf W^e=\left(\sum_{j=1}^{C'}w_{jc}^e\cdot[\beta_{j}{\rm{GeLU}}(\alpha_{j}\sum_{i=1}^{C}w_{i,f(j)}^dx_i+\gamma_{j})+\theta_{j}] \right)_{c=1}^C$

$~~~~~~~~~~~~~~~~~~~=\left(\sum_{j=1}^{C'}w_{jc}^{'e}{\rm{GeLU}}(m_{cj}'x_c+n_{cj}')+\theta_{j}\right)_{c=1}^C,~~~~~~~~Eq.6'$

where $w_{jc}^{'e}=w_{jc}^{e}\cdot\beta_j$, $m_{cj}'=w^d_{c,f(j)}$ and $n_{cj}'={\rm{func}}(x_1,\cdot,\cdot,\cdot,x_{c-1},x_{c+1},\cdot,\cdot,\cdot,x_C)=\sum_{i=1,i\neq c}^C{w_{i,f(j)}^dx_i+\gamma_{j}}$.

This is almost the same as Eq.2. According to Corollary 1, the output of Eq.2 and Eq.6’ can both be expressed as the combination of C’ different nonlinear functions with distinct scales and biases, in which the scales are learnable weights independent to the input and the biases are dependent to the input.

The only difference between Eq.2 and Eq.6’ is that the degree of freedom of  $w_{jc}^{'e}$ in Eq.6’ is halved compared to the original $w_{ij}^a$ in Eq.2. Thus, it is harder to optimize Eq.6’ and may degrade the performance, as shown in the first paragraph in section 4.3.

The above analysis is added near Eq.6 in the revised version.

Reviewer comment (Experiments):
The proposed variant of GeLU is not exclusively applicable to vision transformers. It can also be utilized in architectures like ConvNeXt, which shares similar building blocks, excluding self-attention, where the proposed element could serve as a replacement for standard GeLUs.

Author response:
**A7:** This is true. However, the analysis in this paper is based on a 2-layer MLP module and thus is only applicable to vision transformers. Using it to ConvNeXt and other MLP networks is a really good topic and future work that is worth discussing.

Reviewer comment (Evaluation):
The reviewer highlights Table 3, noting it presents a surprising and crucial result of the study. The authors are requested to provide insights or intuitions into why such an outcome occurred.

Author response:
**A9:** Some of the explanations are shown at the end of page 7. In fact, Line 1 of Table 3 stands for the original MLP module, and Line 2 changes the GELU function to AGeLU. These two architectures have similar performance since they all have the same degree of nonlinearity according to Eq.2 (the form of original MLP), no matter $\phi$=GeLU or $\phi$=AGeLU. 

Line 4 of Table 3 is the proposed IMLP module using 2 different AGeLU functions with concat, and Line 3 uses two GeLU functions with concat. The reason why Line 4 is much better than Line 3 comes from Eq.6. When using two different AGeLUs in Line 4, we have different $\alpha_{c’}$ and $\alpha_{c’}'$ to form $t_1$ and $t_1’$ in Eq.6 in the original paper. Thus, the output can be expressed as the linear combination of $C’$ different nonlinear functions. However, when using two GeLU functions in Line 3, we have $\alpha_{c’}=\alpha_{c’}'=1$. Thus, the output can only be expressed as the linear combination of $C’/2$ different nonlinear functions.

Reviewer comment (Presentation):
The results in Table 4 are not clearly explained.

Author response:
**A10:** Sorry for the unclarity. In fact, the explanation of Table 4 is shown at the beginning of section 4.1. Table 4 is a straightforward way to use the combination of C′ different nonlinear functions to replace the original MLP module. However, based on Corollary 1 (3) in the paper, the bias that depends on the input elements makes it challenging to attain a comparable degree of non-linearity by merely combining multiple nonlinear functions, and the classification performance does not match that of using the original MLP module.

Reviewer comment (Theory):
Why is the additional shortcut needed for the dwconv and BN should follow it subsequently?

Author response:
**A11:**  Many researchers found that learning the residual of the output is easier than learning the original input, and BN can normalize the input for better optimization. We found that using BN and shortcut can achieve a better result than not using it through empirical experiments.