MLP Can Be A Good Transformer Learner

TL;DR

The paper addresses the high computational and memory costs of Vision Transformers by identifying uninformative attention layers through an entropy-based criterion. It introduces NOSE, an entropy-guided selection strategy using transfer entropy to pick attention layers with minimal impact on the final output, and a dilution learning technique to degenerate these layers into identical mappings that are absorbed by subsequent MLPs. Across ImageNet-1k, ADE20k, and CIFAR-100, the method removes up to 40–50% of attention layers in DeiT-B (and other backbones) with no accuracy loss and notable gains in throughput and memory efficiency. The work also analyzes feature expressivity via frequency-domain methods, showing enhanced high-frequency information and generalization, and provides code and supplementary materials to facilitate adoption and further study.

Abstract

Self-attention mechanism is the key of the Transformer but often criticized for its computation demands. Previous token pruning works motivate their methods from the view of computation redundancy but still need to load the full network and require same memory costs. This paper introduces a novel strategy that simplifies vision transformers and reduces computational load through the selective removal of non-essential attention layers, guided by entropy considerations. We identify that regarding the attention layer in bottom blocks, their subsequent MLP layers, i.e. two feed-forward layers, can elicit the same entropy quantity. Meanwhile, the accompanied MLPs are under-exploited since they exhibit smaller feature entropy compared to those MLPs in the top blocks. Therefore, we propose to integrate the uninformative attention layers into their subsequent counterparts by degenerating them into identical mapping, yielding only MLP in certain transformer blocks. Experimental results on ImageNet-1k show that the proposed method can remove 40% attention layer of DeiT-B, improving throughput and memory bound without performance compromise. Code is available at https://github.com/sihaoevery/lambda_vit.

Figures (9)

• Figure 1: Pruning the attention layer from the perspective of entropy. (a) We use entropy to illustrate the information amount carried out by the attention layers and MLP layers (i.e. two FFN layers) in each transformer block of DeiT-B touvron2021training. We observe that the entropy quantity of the bottom blocks is lower than that of the top blocks. We identify a pattern that, the attention layer with low entropy is accompanied by the MLP layers with the entropy quantity at the same level. (b) In the bottom blocks, MLP layers can elicit the information as much as that of the attention layers. On the other hand, they are under-exploited given the low entropy quantity compared to those MLP layers in the top blocks. We thus propose to integrate the uninformative attention layer into its subsequent MLP layer through proper optimization. (c) As a result, our method can reduce 13.7% parameters of DeiT-B and improve 20.5% working load in the same memory budget without performance degradation.
• Figure 2: Interaction of multiple layers. Both figures have the same $x$-axis (#Layer). We use the idea of transfer entropy to measure the interaction on multiple layers. Here, we randomly mask 1$\sim$5 attention layers of a pre-trained DeiT-B. We record the means (bars) and variances (red lines) of model performances in (a) and the corresponding transfer entropies in (b). It is clear that model performance drops as transfer entropy increases in both mean and variance. As a motivator, we aim to remove attention layers with fewer interactions (i.e. transfer entropy).
• Figure 3: NOSE vs. Random selection. (a) NOSE consistently outperforms the random selection on ImageNet-1k. (b) This is because NOSE can identify the attention layers with less interaction with the final output layer, which is reflected by transfer entropy.
• Figure 4: Visualization of the proposed NOSE. For each step, a row visualizes the transfer entropy, normed to [0,1], of each attention layer associated with the final output layer. We use greedy search to select the one with minimum transfer entropy, denoted by the red dashed box, e.g., layer 3 is selected at step 0. The selected layer is denoted by a gray dotted box and is suspended to a state set. In the next step, NOSE repeats this procedure on the rest attention layers considering the previous state. Finally, the attention layer indexed by [0,1,3,4,6] will be integrated into their subsequent MLP layers.
• Figure 5: Visualization of feature frequency. We analyze feature expressivity from the frequency perspective. We apply Discrete Fourier Transform to the output feature of each block, where frequency domain is divided into low, medium, and high components. From blocks 3 to 11, our model encodes more significant high-frequency components compared to DeiT-B, implying superior feature power bai2022improvingguo2023aloft.