Hierarchical Associative Memory, Parallelized MLP-Mixer, and Symmetry Breaking

TL;DR

This paper introduces a novel framework combining hierarchical associative memory with MetaFormers to unify Transformer components into a Hopfield network, revealing the importance of symmetry-breaking in model performance.

Contribution

It presents a new theoretical perspective by integrating Krotov's associative memory with MetaFormers, leading to a parallelized MLP-Mixer model and insights into symmetry-breaking effects.

Findings

Symmetric matrices hinder image recognition performance.

Symmetry-breaking improves the effectiveness of MLP-Mixer.

Vanilla MLP-Mixer spontaneously develops symmetry-breaking configurations during training.

Abstract

Transformers have established themselves as the leading neural network model in natural language processing and are increasingly foundational in various domains. In vision, the MLP-Mixer model has demonstrated competitive performance, suggesting that attention mechanisms might not be indispensable. Inspired by this, recent research has explored replacing attention modules with other mechanisms, including those described by MetaFormers. However, the theoretical framework for these models remains underdeveloped. This paper proposes a novel perspective by integrating Krotov's hierarchical associative memory with MetaFormers, enabling a comprehensive representation of the entire Transformer block, encompassing token-/channel-mixing modules, layer normalization, and skip connections, as a single Hopfield network. This approach yields a parallelized MLP-Mixer derived from a three-layer Hopfield network, which naturally incorporates symmetric token-/channel-mixing modules and layer normalization. Empirical studies reveal that symmetric interaction matrices in the model hinder performance in image recognition tasks. Introducing symmetry-breaking effects transitions the performance of the symmetric parallelized MLP-Mixer to that of the vanilla MLP-Mixer. This indicates that during standard training, weight matrices of the vanilla MLP-Mixer spontaneously acquire a symmetry-breaking configuration, enhancing their effectiveness. These findings offer insights into the intrinsic properties of Transformers and MLP-Mixers and their theoretical underpinnings, providing a robust framework for future model design and optimization.