SS ViT: Observing pathologies of multi-layer perceptron weights and re-setting vision transformer

Abstract

Vision Transformer (ViT) usually adopts a columnar or hierarchical structure with four stages, where identical block settings are applied within the same stage. To achieve more nuanced configurations for each ViT block, additional search is conducted to explore stronger architectures. However, the search cost is typically expensive and the results may not be transferable to different ViT architectures. In this paper, we present a DFC module, which exploits two lightweight grouped linear (GL) layers to learn the representations of the expansion layer between two fully connected layers and the nonlinear activation of multi-layer perceptron (MLP), respectively. Afterwards, we introduce the DFC module into vanilla ViT and analyze the learned weights of its GL layers. Interestingly, several pathologies arise even though the GL layers share the same initialization strategy. For instance, the GL layer weights display different patterns across various depths, and the GL1 and GL2 weights have different patterns in the same depth. We progressively compare and analyze these pathologies and derive a specific setting (SS) for ViT blocks at different depths. Experimental results demonstrate that SS generically improves the performance of various ViT architectures, not only enhancing accuracy but also reducing inference time and computational complexity. For example, on ImageNet-1k classification task, SS yields a significant 0.8% accuracy improvement, approximately 12.9% faster inference speed, and 25% fewer floating-point operations (FLOPs) on PVTv2 model. The codes and trained models are available at https://github.com/ICSResearch/SS.