Softplus Attention with Re-weighting Boosts Length Extrapolation in Large Language Models (2501.13428v3)

Abstract: LLMs have achieved remarkable success in recent years, primarily due to the implementation of self-attention mechanisms. However, traditional Softmax attention suffers from numerical instability and reduced performance as the length of inference tokens increases. This paper addresses these issues by decomposing the Softmax operation into a non-linear transformation and the $l_1$-norm. We identify the latter as essential for maintaining model performance. By replacing the non-linear transformation with the Softplus activation function and introducing a dynamic scale factor for different token lengths based on invariance entropy, we create a novel attention mechanism with performance better than conventional Softmax attention across various inference lengths. To further improve the length extrapolation ability of the proposed attention mechanism, we introduce a novel re-weighting mechanism that amplifies significant attention weights while diminishing weaker ones, enabling the model to concentrate more effectively on relevant tokens. When combined with our proposed attention mechanism, this approach maintains nearly constant validation loss even at 16$\times$ the training token length, ensures numerical stability, and achieves superior results on downstream benchmarks.

• The paper introduces a novel attention mechanism that replaces Softmax with Softplus and includes a re-weighting strategy to enhance length extrapolation.
• It employs a dynamic length scale and l1-norm normalization to boost numerical stability during long sequence processing.
• Experiments with GPT-2 demonstrate nearly constant validation loss even for sequences up to 16 times longer than those used in training.

Softplus Attention with Re-weighting Boosts Length Extrapolation in LLMs

The paper entitled "Softplus Attention with Re-weighting Boosts Length Extrapolation in LLMs" addresses a critical limitation of the conventional Softmax-based attention mechanisms in LLMs: their reduced performance and numerical instability during length extrapolation across varying token lengths. The authors propose a novel attention mechanism, Length Scaled Softplus Attention with Re-weighting (LSSAR), which has shown an improved capacity to manage longer sequences with enhanced numerical stability.

Key Contributions

The paper makes noteworthy contributions by introducing an attention mechanism that replaces the exponential component of Softmax with the Softplus activation function. This novel approach is augmented by a re-weighting mechanism, enabling the model to focus more precisely on significant attention weights, thereby enhancing performance across long inference sequences. The key technical advancements include:

1. Softmax Decomposition: The authors decompose the Softmax operation into a non-linear transformation and normalization through an l1-norm. This step is crucial for maintaining model performance, as it indicates that non-linearities, rather than non-negativity, play a critical role in sustaining the effectiveness of attention mechanisms.
2. Length Scaled Softplus Attention (LSSA): The proposed LSSA mechanism replaces the exponential function with the Softplus activation function and introduces a dynamic length scale factor. This factor accounts for token length variations, thus enhancing the model's performance for extended sequences.
3. Attention Re-weighting: By applying a power transformation to the attention scores, the model's discriminative capacity is improved. This re-weighting approach amplifies significant attention weights, maintaining performance stability even at sequence lengths up to 16 times the training length.

Experimental Validation

The authors conducted a series of experiments using the GPT-2 architecture. Their results demonstrated that the LSSAR model achieved nearly constant validation loss across token sequence lengths expanding up to 16 times the training sequence length, maintaining superior numerical stability. The authors compared LSSAR performance with several state-of-the-art Softmax-free alternatives, including ReLU and Sigmoid-based approaches. LSSAR outperformed these alternatives, indicating its potential to significantly boost the length extrapolation capabilities of transformers.

Implications and Future Directions

Practically, the introduction of LSSAR and its demonstrated capacity to handle longer sequences effectively translates to potentially more resource-efficient LLMs. Theoretically, the work emphasizes the fundamental role of activation function choice and dynamic scaling in attention mechanisms. The findings suggest avenues for further research into attention models that maintain high performance over extended token sequences while ensuring numerical stability.

Future research may focus on implementing this attention mechanism in larger models, given its potential for scalability without loss of performance in extended contexts. Additionally, it could explore optimizing the computational cost of Softplus and re-weighting implementations to facilitate practical applications in resource-constrained environments.

In conclusion, the paper advances our understanding of attention mechanisms in LLMs, offering a promising approach to prolonging sequence length management without prevalent numerical stability issues. The proposed LSSA and LSSAR models mark significant strides in aligning model performance with the increasing demands for long-sequence processing, an essential requirement in numerous real-world applications.