Non Linear Sum Attention

Standard Transformer architectures rely on variants of Multi-Head Attention (MHA) followed by a dense Multi-Layer Perceptron (MLP). While effective, standard Transformers enforce a strict disconnect between token-level context mixing and position-wise feature transformation. Furthermore, they exhibit structural inefficiencies on algorithmic tasks, requiring tens of thousands of training epochs to transition from rote memorization to generalization. In this paper, we introduce Non-Linear Sum Attention (NLSA), an architecture that addresses this limitation by delegating position-wise transformation to the value matrix, applying non-linear functions to the final attention output, and taking a gated sum over all heads. We additionally introduce Multi-Head Latent Sum Attention (MLSA) to address the problem of KV cache blowup in NLSA. MLSA also achieves the lowest mean and median grokking times among the tested architectures on Multi-Query Associative Recall, significantly outperforming standard Transformers.