Neural Attention: A Novel Mechanism for Enhanced Expressive Power in Transformer Models

TL;DR

This paper introduces Neural Attention, a new mechanism replacing dot products with feed-forward networks in transformers, significantly enhancing their expressive power and improving performance across NLP and image classification tasks.

Contribution

Neural Attention is a novel attention mechanism that increases transformer expressivity by replacing dot products with neural networks, with demonstrated improvements in multiple domains.

Findings

Over 2% perplexity reduction on WikiText-103

More than 4 percentage points accuracy gain on CIFAR datasets

Maintains compatibility with existing transformer architectures

Abstract

Transformer models typically calculate attention matrices using dot products, which have limitations when capturing nonlinear relationships between embedding vectors. We propose Neural Attention, a technique that replaces dot products with feed-forward networks, enabling a more expressive representation of relationships between tokens. This approach modifies only the attention matrix calculation while preserving the matrix dimensions, making it easily adaptable to existing transformer-based architectures. We provide a detailed mathematical justification for why Neural Attention increases representational capacity and conduct controlled experiments to validate this claim. When comparing Neural Attention and Dot-Product Attention, NLP experiments on WikiText-103 show a reduction in perplexity of over 2 percent. Similarly, experiments on CIFAR-10 and CIFAR-100 show improvements in accuracy of more than 4 percentage points for image classification tasks. While Neural Attention introduces higher computational demands, we develop techniques to mitigate these challenges, ensuring practical usability without sacrificing the increased expressivity it provides. This work establishes Neural Attention as an effective means of enhancing the predictive capabilities of transformer models across a variety of applications. The code for all experiments is available at https://github.com/awayfromzel/neural-attention-research.