A Framework for Non-Linear Attention via Modern Hopfield Networks

TL;DR

This paper introduces a unified energy-based framework for non-linear attention in transformers using Modern Hopfield Networks, enhancing the modeling of complex token relationships and improving performance.

Contribution

It proposes a novel energy functional that unifies attention mechanisms with Hopfield networks, enabling non-linear attention in transformer models.

Findings

Defines an energy landscape where attention corresponds to gradient descent.

Provides a method for integrating non-linear heads into transformer architectures.

Improves understanding of complex token relationships in sequence modeling.

Abstract

In this work we propose an energy functional along the lines of Modern Hopfield Networks (MNH), the stationary points of which correspond to the attention due to Vaswani et al. [12], thus unifying both frameworks. The minima of this landscape form "context wells" - stable configurations that encapsulate the contextual relationships among tokens. A compelling picture emerges: across $n$ token embeddings an energy landscape is defined whose gradient corresponds to the attention computation. Non-linear attention mechanisms offer a means to enhance the capabilities of transformer models for various sequence modeling tasks by improving the model's understanding of complex relationships, learning of representations, and overall efficiency and performance. A rough analogy can be seen via cubic splines which offer a richer representation of non-linear data where a simpler linear model may be inadequate. This approach can be used for the introduction of non-linear heads in transformer based models such as BERT, [6], etc.