Two Stones Hit One Bird: Bilevel Positional Encoding for Better Length Extrapolation

TL;DR

This paper introduces Bilevel Positional Encoding (BiPE), which separately handles within-segment (intra-segment) and between-segment (inter-segment) positional information to improve length extrapolation in Transformer models. By applying absolute positional encoding inside segments and relative positional encoding between segments, BiPE achieves greater parameter efficiency and better extrapolation performance, as supported by theory framed around Bi-NFA and extensive experiments across arithmetic reasoning, long-context benchmarks, and normal-length data. BiPE variants, such as BiPE-RoPE and BiPE-ALiBi, demonstrate superior length extrapolation across multimodal tasks and maintain competitive performance on in-distribution tasks. The work suggests a promising direction for leveraging intrinsic data segmentation and opens avenues for hierarchical extensions to further improve extrapolation in longer contexts and other sequence types.

Abstract

In this work, we leverage the intrinsic segmentation of language sequences and design a new positional encoding method called Bilevel Positional Encoding (BiPE). For each position, our BiPE blends an intra-segment encoding and an inter-segment encoding. The intra-segment encoding identifies the locations within a segment and helps the model capture the semantic information therein via absolute positional encoding. The inter-segment encoding specifies the segment index, models the relationships between segments, and aims to improve extrapolation capabilities via relative positional encoding. Theoretical analysis shows this disentanglement of positional information makes learning more effective. The empirical results also show that our BiPE has superior length extrapolation capabilities across a wide range of tasks in diverse text modalities.

Figures (9)

• Figure 1: Left: The distribution of the token number in one segment with different sequence lengths. Right: The distribution of the number of segments with different sequence lengths. We use the tokenizer of Llama 2 touvron2023llama for tokenization on PG-19 Rae2020Compressive. Full stop"." and newline "\\ n" are used for segmentation. It can be seen that even when the sequence length is around 120k, the token number in most sentences is less than 50, while the number of sentences grows up to 10k.
• Figure 2: Left: The schematic comparison of Standard Positional Encoding (top) and our proposed Bilevel Positional Encoding (BiPE, bottom). BiPE differentiates positions using both intra-segment and inter-segment encodings. Right: Absolute positional encoding is used as Intra-Segment Encoding added to the input embedding and relative positional encoding (e.g., RoPE and ALiBi) is used as Inter-Segment Encoding in the Transformer attention module.
• Figure 3: Accuracy of Transformer models with different positional encoding methods on the Arithmetic task. Our BiPE method consistently performs best on different scales of parameters.
• Figure 4: Language modeling perplexity with varying evaluation sequence lengths for models trained on sequence length 1024.
• Figure 5: Language modeling perplexity with varying evaluation sequence lengths for RoPE and BiPE-RoPE finetuned with YaRN.

Theorems & Definitions (13)

• Theorem 3.1: Lower bound for Transformer with absolute positional encoding to represent NFA
• Theorem 3.2: Upper bound for Transformer with BiPE to represent Bi-NFA
• Lemma 1.1: From feng2023towards
• Lemma 1.2: From feng2023towards
• Lemma 1.3: From feng2023towards
• Lemma 1.4
• proof : Proof of \ref{['lemma:MLP_matrix']}
• Lemma 1.5: From feng2023towards
• Theorem 1.6: Lower Bound for Transformer with APE to Represent NFA
• proof : Proof of \ref{['Thm:classic_PE:app']}