Understanding MLP-Mixer as a Wide and Sparse MLP

TL;DR

This paper reveals that the success of MLP-Mixer architectures stems from their effective wide and sparse structure, which can be understood through Kronecker-product weights and implicit sparse regularization, enhancing performance.

Contribution

It provides a novel theoretical understanding of MLP-Mixer as a wide, sparse MLP, connecting architecture to sparsity properties and empirical performance improvements.

Findings

MLP-Mixers can be expressed as wider MLPs with Kronecker-product weights.

The architecture induces implicit sparse regularization.

Empirical evidence shows similarity to unstructured sparse-weight MLPs.

Abstract

Multi-layer perceptron (MLP) is a fundamental component of deep learning, and recent MLP-based architectures, especially the MLP-Mixer, have achieved significant empirical success. Nevertheless, our understanding of why and how the MLP-Mixer outperforms conventional MLPs remains largely unexplored. In this work, we reveal that sparseness is a key mechanism underlying the MLP-Mixers. First, the Mixers have an effective expression as a wider MLP with Kronecker-product weights, clarifying that the Mixers efficiently embody several sparseness properties explored in deep learning. In the case of linear layers, the effective expression elucidates an implicit sparse regularization caused by the model architecture and a hidden relation to Monarch matrices, which is also known as another form of sparse parameterization. Next, for general cases, we empirically demonstrate quantitative similarities between the Mixer and the unstructured sparse-weight MLPs. Following a guiding principle proposed by Golubeva, Neyshabur and Gur-Ari (2021), which fixes the number of connections and increases the width and sparsity, the Mixers can demonstrate improved performance.