Lateralization MLP: A Simple Brain-inspired Architecture for Diffusion

TL;DR

The paper introduces Lateralization MLP (L-MLP), a brain-inspired architecture that rivals transformers in diffusion tasks while being more efficient, by mimicking human brain lateralization in a simple, scalable MLP design.

Contribution

Proposes the Lateralization MLP (L-MLP), a novel brain-inspired architecture that outperforms other MLP variants and matches transformer performance in diffusion tasks.

Findings

L-MLP outperforms other MLP variants.

L-MLP performs comparably to transformers in diffusion tasks.

L-MLP is highly efficient and effective.

Abstract

The Transformer architecture has dominated machine learning in a wide range of tasks. The specific characteristic of this architecture is an expensive scaled dot-product attention mechanism that models the inter-token interactions, which is known to be the reason behind its success. However, such a mechanism does not have a direct parallel to the human brain which brings the question if the scaled-dot product is necessary for intelligence with strong expressive power. Inspired by the lateralization of the human brain, we propose a new simple but effective architecture called the Lateralization MLP (L-MLP). Stacking L-MLP blocks can generate complex architectures. Each L-MLP block is based on a multi-layer perceptron (MLP) that permutes data dimensions, processes each dimension in parallel, merges them, and finally passes through a joint MLP. We discover that this specific design outperforms other MLP variants and performs comparably to a transformer-based architecture in the challenging diffusion task while being highly efficient. We conduct experiments using text-to-image generation tasks to demonstrate the effectiveness and efficiency of L-MLP. Further, we look into the model behavior and discover a connection to the function of the human brain. Our code is publicly available: \url{https://github.com/zizhao-hu/L-MLP}