Effective PE is vital, particularly for enabling LMs trained on shorter contexts to generalize to significantly longer sequences during inference--a desirable capability known as context length extrapolation. Several PE methods have been proposed to facilitate context length extrapolation, including Sinusoidal embeddings (V aswani, 2017), RoPE (Su et al., 2024), ALiBi (Press et al., 2022), and even the omission Bayesian attention mechanism, hereby called BAM. 2.1 B This dependency is trivially modeled by a scalar Z when the scoring function is additive, as detailed below. If the scoring function of the attention mechanism is additive, i.e., of the form With Theorem 1, we can frame positional encoding as priors to BAM. Lemma 2. ALiBi is a special case of BAM prior where the token position distribution comprises Lemma 3. ALiBi becomes local attention as the relative length |j i| increases. See Appendix B.1, B.2, and B.3. 2.3 A PE We call this new PE method GGD-BAM.

The maximum element of the vector output by the Softmax function approaches zero as the input vector size increases. Transformer-based language models rely on Softmax to compute attention scores, causing the attention distribution to flatten as the context size grows. This reduces the model's ability to prioritize key information effectively and potentially limits its length generalization. To address this problem, we propose Scalable-Softmax (SSMax), which replaces Softmax in scenarios where the input vector size varies. SSMax can be seamlessly integrated into existing Transformer-based architectures. Experimental results in language modeling show that models using SSMax not only achieve faster loss reduction during pretraining but also significantly improve performance in long contexts and key information retrieval. Furthermore, an analysis of attention scores reveals that SSMax enables the model to focus attention on key information even in long contexts. Additionally, although models that use SSMax from the beginning of pretraining achieve better length generalization, those that have already started pretraining can still gain some of this ability by replacing Softmax in the attention layers with SSMax, either during or after pretraining.