Softplus Attention with Re-weighting Boosts Length Extrapolation in Large Language Models

Large language models have achieved remarkable success in recent years, primarily due to self-attention. However, traditional Softmax attention suffers from numerical instability and reduced performance as the number of inference tokens increases. This work addresses these issues by proposing a new design principle for attention, viewing it as a two-stage process. The first stage (normalisation) refines standard attention by replacing Softmax with the more numerically stable Softplus followed by $l_{1}$-normalisation. Furthermore, we introduce a dynamic scale factor based on invariance entropy. We show that this novel attention mechanism outperforms conventional Softmax attention, and state-of-the-art Softmax-free alternatives. Our second proposal is to introduce a second processing stage (sharpening) which consists of a re-weighting mechanism that amplifies significant attentional weights while diminishing weaker ones. This enables the model to concentrate more effectively on relevant tokens, mitigating the attention sink phenomenon, and fundamentally improving length extrapolation. This novel, two-stage, replacement for self-attention is shown to ensure numerical stability and dramatically improve length extrapolation, maintaining a nearly constant validation loss at 16$\times$ the training length while achieving superior results on challenging long-context retrieval tasks and downstream benchmarks. Furthermore, symbolic regression experiments demonstrate that our method enables models to recover Newton's gravitational law from orbital trajectory sequences, providing evidence that appropriate attention mechanisms are crucial for foundation models to develop genuine physical world models. Our code is available at https://github.com/iminfine/freeattn .

Large language models often struggle when they are asked to read much longer text than they saw during training. A key reason is the attention mechanism, the part of the model that decides which earlier words or tokens are most relevant. In this work, we redesign attention so that it first produces stable scores and then sharpens them to focus more strongly on important information. Our method, called LSSAR, helps models trained on short text work better on much longer inputs without changing the position encoding or retraining from scratch. We show that this improves long-context language modeling and retrieval tasks. We also test the method on a physics-inspired problem, where a model observes orbital trajectories and predicts gravitational forces. The results suggest that better attention mechanisms can help models learn meaningful structure from data, such as the inverse-square pattern of gravity, rather than only improving benchmark scores. This may help future models reason more reliably over long documents and scientific data.