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

TL;DR

Transformers' dot-product attention constrains expressivity; the paper introduces Neural Attention, replacing the dot product with a nonlinear scoring function on the concatenation of queries and keys, i.e., $AttentionScore_{ij} = w_a^T σ(W_h concat(q_i, k_j) + b_h) + b_a$, to produce attention weights. Key contributions include a theoretical justification for increased representational capacity, a down-projection strategy to control memory, and empirical validation on WikiText-103 (perplexity reduction up to 2.19%) and CIFAR-10/100 (accuracy gains of 3.0 and 4.26 points). The method preserves the $n×n$ attention matrix and is modular enough to plug into existing transformers, with practical overhead managed by limiting application to the first layer and using small $d'$. The results suggest that enhancing attention expressivity yields tangible improvements across NLP and CV tasks, offering a scalable path toward more capable transformer models.

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.

Figures (11)

• Figure 1: A smooth parabolic surface (left) versus a complex surface with sharp edges (right).
• Figure 2: Reshape and broadcast process of matrices $\mathit{Q^{\prime}}$ and $\mathit{K^{\prime}}$ into tensors $\mathbf{Q^{\prime}}$ and $\mathbf{K^{\prime}}$. Prime symbols are left out for better readability.
• Figure 3: Tensor $\mathbf{C}$, created after concatenating tensors $\mathbf{Q^{\prime}}$ and $\mathbf{K^{\prime}}$ from Figure \ref{['fig:reshape_and_broadcast']} along their embedding dimension.
• Figure 4: Embedding vectors $\vec{q}$ (shown in orange) and $\vec{k}$ (shown in blue).
• Figure 5: Dot product calculation between embedding vectors $\vec{q}$ (shown in orange) and $\vec{k}$ (shown in blue).