Paper: Understanding MLP-Mixer as a wide and sparse MLP (ICLR-2024)

Reviewer comment (Experiments):
The authors claim that wider Mixer is a an effective expression of MLP-Mixer and RP-Mixer is a more memory-efficient alternative to unstructured sparse-weight MLP. It would be beneficial to apply wider Mixer/RP-Mixer to more state-of-the-art frameworks [1,2,3,4,5], especially considering the more efficient MLP-Mixer approach proposed in [6].

Author response:
Thank you for your helpful suggestion. We especially appreciate your providing numerous references. Our response to your comments is as follows.

> The authors claim that wider Mixer is a an effective expression of MLP-Mixer and RP-Mixer is a more memory-efficient alternative to unstructured sparse-weight MLP. It would be beneficial to apply wider Mixer/RP-Mixer to more state-of-the-art frameworks [1,2,3,4,5], especially considering the more efficient MLP-Mixer approach proposed in [6]

and, 

>Q1, Q2.

As Reviewer suggests, research comparing MLP-Mixer with other Mixer variants is making significant contributions while being abundant. However, what we want to emphasize is that our objective is to elucidate **why the MLP-Mixer performs better than a naive MLP**. Therefore, our comparison is primarily with a naive MLP.

The following explanatory diagrams and comparison tables **introduce the RP-Mixer to bridge MLP and MLP-Mixer**, not to propose a superior MLP-block structure design to MLP-Mixer. As shown in the comparison table, studies like [1,2,3,4,5,6] focus on innovating the structure within blocks for inter/intra-token mixing. In contrast, the RP-Mixer, which flattens tensors and mixes them randomly, is closer to a naive  MLP.

-----

MLP  — ( RP-Mixer)  —  MLP-Mixer   —   [1,2,3,4,5,6]
-----
Figure. Relationship between models.  

| Models  |    where  Blocks operate  on | Destroy/Improve structure of MLP-Mixer  |
|--------|--------|--------|
| MLP |              **whole**                |   **Destroy** |
|  RP-Mixer  |     **whole** |     **Destroy** ( Flatten → Randomly Permute → Reshape)  |
|MLP-MIxer |     inter/intra-tokens|  - |
|MorphMLP[1] |  inter/intra-tokens| Improve (Separate H and W)|
|sMLPNet[2] | inter/intra-tokens| Improve (Separate H and W)|
|As-MLP[3] | inter/intra-tokens| Improve (Axial Shift)|
|CycleMLP[4] |inter/intra-tokens| Improve (Axial Shift)|
|ATM[5]  |  inter/intra-tokens |  Improve (Message Passing)|
|AFF[6] |  inter/intra-tokens|  Improve (Adaptive Filter)|

Table.  Comparison of models on their structure.

The references [1-6] by Reviewer  experimentally reveal how modifying the mixing block can improve accuracy over the MLP Mixer. However, the fundamental question of **why a mixing block is superior to the dense layer of a naive MLP** remains a mystery. Understanding the intrinsic properties of these blocks will likely guide future improvements to Mixers, similar to references [1-6]. For instance, the equality $S=C$ is a quantitative suggestion. While Mixers with such a structure present an interesting problem, as mentioned earlier, our study's scope is to clarify the differences with a naive MLP, so we leave this for subsequent work in the future.

Reviewer comment (Experiments):
The paper lacks a thorough comparison with existing methods in terms of both memory efficiency and performance. It would be valuable to include such comparisons to demonstrate the superiority of wider- Mixer/RP-Mixer.

Author response:
Regarding the comparison with MLP:

- In terms of efficiency, the RP-Mixer is about 10^2 to 10^3 times better than WS-MLP (Table 1 in the revised manuscript).
- The state-of-the-art accuracy for MLP, represented by Nesybur's beta-laaso, is 85% (CIFAR10), whereas our proposed performance is 87.98% (CIFAR10), (Table 2).
- While Scaling MLPs[Bachmann2023] achieves 56% on ImageNet with a naive MLP after various modifications, the RP-Mixer achieves 70% (ImageNet) as shown in Figure 6.

Reviewer comment (Evaluation):
The experimental section should include a thorough evaluation of the proposed method's performance compared to other SOTA methods. This evaluation should consider multiple benchmark datasets and provide statistical analysis to support the claims of superiority.

Author response:
>The experimental section should include a thorough evaluation of the proposed method's performance compared to other SOTA method

To achieve our purpose, a thorough evaluation means comparing MLP with MLP/RP-Mixer across a variety of parameters, widths, and numbers of layers. We have conducted performance comparison experiments using 5000 to 10000 GPU hours, which can be considered a thorough evaluation in this context.

Reviewer comment (Novelty):
The MLP-Mixer is a fairly nieche model that hasn't seen much adaption neither in practice nor as a vessel for theoretical analysis. So I'm under the impression that this work will not be immediately useful to a broader audience. However, the tools of analysis used and the results obtained might be useful for future research in related areas. The authors themselves mention such potential areas on the Conclusion.

Author response:
Thank you for your helpful suggestions.  We updated the manuscript based on that.

> My main question to the authors would be if they have any immediate applications of their findings. 

As (Golubeva, Neyshabur, and Gur-Ari (2021)) suggests, **sparsity and width are two critical factors** determining the performance of DNNs, and they have significantly contributed to the development of architecture and algorithms. We have discovered that MLPmixer is a model that efficiently and memory-friendly balances both these aspects.

Therefore, it has the potential for much adaptation in practice and can serve as a vessel for theoretical analysis. For instance, 
(1)**replacing the MLP block with a Mixer-block** and considering a dual-structured Mixer-block as investigated in the research of **Monarch matrices**  [1-a] seems practically useful,
(2) and it can be utilized for **verifying generalization bounds**[1-b, 1-c] based on width and sparsity, by replacing MLPs with MLP-Mixers. Because of computational efficiency, we can evaluate them on large datasets with practical performance. 

(1-a)[Monarch: Expressive Structured Matrices for Efficient and Accurate Training](https://proceedings.mlr.press/v162/dao22a/dao22a.pdf), ICML2022

(1-b)[Norm-based Generalization Bounds for Sparse Neural Networks](https://openreview.net/forum?id=COPzNA10hZ), NeurIPS2023

(1-c)[Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks](https://proceedings.neurips.cc/paper/2019/file/cf9dc5e4e194fc21f397b4cac9cc3ae9-Paper.pdf), NeurIPS2019

Reviewer comment (Evaluation):
The authors do not mention running time for their variations, it would be nice if they would state both FLOPS and wallclock times.

Author response:
We have added not only the running time but also memory requirements and FLOPs to Table 1 in Section 4.2. As expected, **MLP-Mixer and RP-Mixer were comparable in running time, but SW-MLP was slower than them.** Furthermore, **RP-Mixer required $10^2$ ~   $10^3$ times less memory and FLOPs compared to  SW-MLP.**

Reviewer comment (Novelty):
My main question to the authors would be if they have any immediate applications of their findings? I agree that in theory this demonstrates that you could use this to structure weights that can achieve a very large effective width. The MLP Mixer's MLP fulfills a very specific task (interlacing inter/intra-token processing), I'm unsure how to extrapolate that to other architectures.

Author response:
Thank you for your helpful suggestions.  We updated the manuscript based on that.

> My main question to the authors would be if they have any immediate applications of their findings. 

As (Golubeva, Neyshabur, and Gur-Ari (2021)) suggests, **sparsity and width are two critical factors** determining the performance of DNNs, and they have significantly contributed to the development of architecture and algorithms. We have discovered that MLPmixer is a model that efficiently and memory-friendly balances both these aspects.

Therefore, it has the potential for much adaptation in practice and can serve as a vessel for theoretical analysis. For instance, 
(1)**replacing the MLP block with a Mixer-block** and considering a dual-structured Mixer-block as investigated in the research of **Monarch matrices**  [1-a] seems practically useful,
(2) and it can be utilized for **verifying generalization bounds**[1-b, 1-c] based on width and sparsity, by replacing MLPs with MLP-Mixers. Because of computational efficiency, we can evaluate them on large datasets with practical performance. 

(1-a)[Monarch: Expressive Structured Matrices for Efficient and Accurate Training](https://proceedings.mlr.press/v162/dao22a/dao22a.pdf), ICML2022

(1-b)[Norm-based Generalization Bounds for Sparse Neural Networks](https://openreview.net/forum?id=COPzNA10hZ), NeurIPS2023

(1-c)[Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks](https://proceedings.neurips.cc/paper/2019/file/cf9dc5e4e194fc21f397b4cac9cc3ae9-Paper.pdf), NeurIPS2019

Reviewer comment (Evaluation):
What is the running time of their models compared to the original MLP Mixer formulation?

Author response:
We have added not only the running time but also memory requirements and FLOPs to Table 1 in Section 4.2. As expected, **MLP-Mixer and RP-Mixer were comparable in running time, but SW-MLP was slower than them.** Furthermore, **RP-Mixer required $10^2$ ~   $10^3$ times less memory and FLOPs compared to  SW-MLP.**

Reviewer comment (Writing):
The takeaways from this paper is a bit unclear, especially on how to properly make use of the new insights in this paper to either improve the quality or improve the performance of the existing architectures (i.e. MLP-Mixer).

Author response:
As the reviewer mentioned, making improvements that significantly increase accuracy would indeed be a considerable contribution to research. However, since we are still at a stage where the reasons for the high performance of MLP-Mixer are not fully understood, providing an understanding of the mechanism where sparsity and width are crucial would be our sufficient contribution and takeaway. History has shown that contributing to the internal understanding of existing models, even when it doesn't directly lead to accuracy improvements, is considered a contribution within the community: for instance,

**[Understanding Batch Normalization](https://proceedings.neurips.cc/paper_files/paper/2018/file/36072923bfc3cf47745d704feb489480-Paper.pdf)  NeurIPS 2018.**

**[On the relationship between self-attention and convolutional layers](https://arxiv.org/abs/1911.03584) ICLR2020.**

Our focus is not primarily on improving the existing architecture of MLP-Mixer, but rather on highlighting that using MLP-Mixer is more efficient than using a naive MLP  as the following (1), (2) and (3). 


(1)Replacing fully connected layers of existing models with Mixer blocks** could lead to performance improvements due to increased effective width. As discussed in Section 3.3, this relates to the **memory-efficient matrix representation** of MLPs with Monarch matrices, which is an interesting direction. For example,  in the MLP-Mixer,  if we replace each MLP block itself with a Mixer block then we have **dual structured Mixer blocks, which enlarges the effective width and can improve performance**.    However, proposing a higher-performance architecture would be too much content for a single paper and is left as an open problem for subsequent work. 

(2) As shown in the Appendix, Random permutations prevent overfitting and can be used with the dual Mixer block when we want to reduce undesirable implicit bias by usual Mixing blocks.

(3)MLP is the most fundamental model and has been the subject of many theoretical studies[*]. However, training MLPs to a **practical level** on **large-scale datasets** is difficult from the perspectives of memory and computational cost, making it challenging to validate theories on large datasets.

Through this research, we have discovered that the computationally efficient MLP mixer is similar to wide MLPs. This revelation **opens up the opportunity to experimentally validate previous theoretical studies on MLPs using an equivalent MLP-Mixer on large datasets!**

 ref: * Examples of theories and possible applications.
-   (3.1) MLP is known to fall into the NTK regime at infinite width, as discussed in [Jacot 2018]. **By using MLP-Mixer instead,  we can observe the order of finite width where feature learning occurs relative to the number of samples.**    
- (3.2) MLPs follow scaling laws if we prepare a  large dataset [Bachmann2023]. However, their performance was low (about acc. 50 % on ImageNet).  **Then by using the similarity of MLP-Mixer and MLP, we can observe how sparseness and wideness affect the scaling laws.** 

[[Jacot2018] Arthur, Jacot, et.al, Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://proceedings.neurips.cc/paper_files/paper/2018/file/5a4be1fa34e62bb8a6ec6b91d2462f5a-Paper.pdf), NeurIPS2018.
[[Bachmann2023] Gregor Bachmann, et.al., Scaling MLPs: A Tale of Inductive Bias](https://neurips.cc/virtual/2023/poster/71680), NeurIPS2023.




We hope that you can expect that our work will work as a concrete starting point for why the MLP-Mixer achieves higher performance compared to a standard MLP. 
We hope our explanations have clarified the issues you pointed out.   If so, we would be grateful for an adjustment in the score.

Reviewer comment (Experiments):
The paper falls short in terms of the selection of networks for comparison, thereby resulting in a lack of theoretical support.

Author response:
Thank you for your feedback. We acknowledge the concern regarding the selection of networks for comparison and the perceived lack of theoretical support.  We want to emphasize that our goal is to understand a single model, the MLP-Mixer. To achieve this, we have conducted exhaustive experiments with various parameters (about 10 selections), widths (about 50 selections), and numbers of layers (about 10 selections). In this context, we believe that the selection of networks for comparison is adequate for our purpose.

Additionally,  we note that these experiments required over a total of 5000 GPU hours.

Reviewer comment (Experiments):
There is a lack of experimental evidence to support the memory efficiency and lightweight structure of the RP-Mixer.

Author response:
Thank you for your suggestions. We have added Table 1 in Section 4.2 which shows the experimental evidence of the efficiency of RP-Mixer in memory requirements, FLOPs, and runtime. This successfully verifies our claim that RP-Mxier is comparable to MLP-Mixer, and they are much better than  SW-MLP!

Reviewer comment (Theory):
According to (Magnus and Neudecker, 2019), there appear to be slight mistakes in the theoretical proof section. For instance, formula $J_c^{\top}\left(I_S \otimes V\right) J_c=V^{\top} \otimes I_S$ should actually be $J_c^{\top}\left(I_S \otimes V\right) J_c=V \otimes I_S$.  
Ref: Jan R Magnus and Heinz Neudecker. Matrix differential calculus with applications in statistics and econometrics. John Wiley & Sons, 2019.

Author response:
Thank you very much for catching this confusing typo. We have corrected the omission of the transpose symbol in the formula.

Reviewer comment (Novelty):
The motivation for this article's research appears to be similar to that of (Golubeva and Neyshabur, 2021). Could you clarify how the author's approach to analyzing network width and sparsity differs from that in Article A?

Author response:
To avoid potential misunderstanding between the reviewer and us, let us emphasize that our motivation or starting point of the work is totally different from Golubeva et al 2021.
Our motivation is to understand the mechanism of MLP-Mixer and why it works well, and it revealed that the sparsity structure of MLP-Mixer consists of the Kronecker product with the commutation matrix. In contrast, the motivation of Golubeva et al 2021 is to find the importance of sparseness in wide neural networks (in more detail, usual forward neural networks; naive MLP, and CNNs).   
The most interesting finding that we raised is that the seemingly different work on the sparseness by Golubeva et al 2021 has a hidden connection to the MLP-Mixer (through the Kronecker product expression). In other words, what we showed in the current work is that we can apply the same approach as in  Golubeva et al 2021 to the MLP-Mixer which is a “seemingly dense” network in the matrix expression but sparse in the vector expression.

Reviewer comment (Writing):
In the contribution, why is the RP-Mixer referred to as a computationally demanding yet lightly-structured alternative? This statement seems contradictory.  
Ref: Anna Golubeva, Behnam Neyshabur, and Guy Gur-Ari. Are wider nets better given the same number of parameters? In International Conference on Learning Representations, 2021.

Author response:
Thank you for catching this typo. I have corrected the term “demanding” to be ”undemanding” in the section.

Reviewer comment (Writing):
The paper is not very easy to follow. For one, a lot of notations are not defined, which require the reader to find out from the original MLP mixer paper. Examples are eq (1), eq (2). The plots are also hard to interpret, and more explanation could be better. The general structure of the paper could also be improved, to have a more coherent story. For example, section 3.2 could be merged with section 5.

Author response:
Following your suggestion, we have added an explanation of MLP-Mixer in Section 2.2.   We have added some additional captions and refined the legend in Figure 1. Furthermore, we merged  Section 3.2 to Section 5 as long as possible.

Reviewer comment (Novelty):
The contribution is weak. For example, while it's good to formalize the relationship between MLP mixer and wide MLP, it's not that unexpected. From figure 1(d), it seems like the equivalent MLP under-performs the MLP-Mixer which makes the equivalence argument weak. The random permuted mixers also don't consistently out-perform the vanilla mixer which is also a weak argument.

Author response:
As far as the authors know,  there is no literature that reveals any relationship between MLP-Mixer and standard MLP. Certainly, the transformation of the equation is simple, but we believe that being simple does not mean it lacks novelty.

It seems that our insufficient explanation might cause a misunderstanding of the purpose of Fig. 1(d).  The purpose is to understand  MLP-Mixer’s performance tendency based on width.  What we pointed out from Fig1(d) are the following two understandings:  

(1) MLP-Mixer has the same properties as a standard MLP at usual widths,

(2) MLP-Mixer overcomes the disadvantage of MLPs that occur at extremely large widths. The disadvantage, reported by (Golveba et al., 2021), is from the eigenvalue distribution of weight matrices. This discussion is mathematically intricate, so it was delegated to the Appendix. Moreover, we add an explanation of the eigenvalue distribution in Section 4.2. 

In conclusion, we believe that an understanding of MLP-Mixer can be achieved by distinguishing between its similarities and slight differences from standard MLPs in Fig1(d).

Furthermore, to emphasize the similarity in the usual width,  we have added the performance comparison for small width into Fig1(d).  The results make it more clearly show the similarity between MLP-Mixer and MLP in usual widths.

Our insufficient explanation seems to have caused the reviewer’s confusion.  Our goal is to create an equivalent model with consistent performance, **essentially revealing that the performance of the Mixer is controlled by its width and sparseness.**  Therefore, RP-Mixer does *not* need to overcome the vanilla mixer. 

Let us also emphasize that RP-Mixer is not necessarily better than MLP-Mixer when the depth is limited. This is because random permutations break the structure for sharing tokens between mixing blocks. If we have a sufficiently large depth, we can overcome this decrease of performance as is confirmed in Figure 6.

Reviewer comment (Evaluation):
Do you have an explanation on why RP-Mixers are not consistently better than normal mixer.

Author response:
Our insufficient explanation seems to have caused the reviewer’s confusion.  Our goal is to create an equivalent model with consistent performance, **essentially revealing that the performance of the Mixer is controlled by its width and sparseness.**  Therefore, RP-Mixer does *not* need to overcome the vanilla mixer. 

Let us also emphasize that RP-Mixer is not necessarily better than MLP-Mixer when the depth is limited. This is because random permutations break the structure for sharing tokens between mixing blocks. If we have a sufficiently large depth, we can overcome this decrease of performance as is confirmed in Figure 6.

Reviewer comment (Theory):
The $vec(X)$ in eq(4) should be $vec(WX)$ and $vec(XV)$?

Author response:
The equation represents maps and vec(X) represents the source element of the map.  However,  as the Reviewer said it may be not easy to follow, so we have rewritten it into matrix form in the revised version. Thank you for your suggestion.