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

TL;DR

The paper addresses Softmax attention’s numerical instability and diminishing length extrapolation by introducing a two-stage architecture: a Normalisation Stage (LSSA) that uses Softplus with an invariance-entropy-based dynamic length scale and $l_{1}$-normalization, and a Sharpening Stage (re-weighting) that concentrates attention via a power transform. This yields LSSAR, which demonstrates stable training, robust long-context extrapolation, improved passkey retrieval, and competitive downstream performance, while enabling the discovery of physical laws such as Newton’s inverse-square law from orbital trajectories. The results highlight that appropriate attention inductive biases—not just model size or data quality—are crucial for learning genuine physical models, and that LSSAR remains compatible with optimized attention implementations. Overall, the work provides a practical, two-stage replacement for self-attention that improves stability, length extrapolation, and interpretability across tasks and scales.

Abstract

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.

Figures (12)

• Figure 1: Force predictions for Earth's orbit. Left: true forces. Middle: GPT (LSSAR) correctly captures radial gravitational structure. Right: GPT (Softmax) is incoherent.
• Figure 2: Passkey retrieval accuracy for LSSAR ($p=15$) and standard Softmax attention. Accuracy is averaged over 100 trials with the passkey placed at random positions within the sequence.
• Figure A3: Comparison of LSSA and Softmax attention with varying values of $p$ for the re-weighting mechanism across different sequence lengths.
• Figure A1: Comparison of attention maps from the last layer of GPT-2-124M, showing standard Softmax attention (left) versus LSSAR with $p=15$ (right). Each panel displays the 12 attention heads in a 4x3 grid. For visualisation purposes, attention scores are clamped to the range [0, 0.5]. The input text is: Working from home can be great most days. I enjoy the flexibility and not having to commute in traffic. But sometimes I miss the office interactions with my coworkers and team meetings.
• Figure A2: Comparison of Softmax attention and LSSAR($p=15$) with validation loss extrapolation for GPT-2-45M (left), GPT-2-124M (middle) and GPT-2-355M (right).