Reproduce “Mesa-Extrapolation: A Weave Position Encoding Method for Enhanced Extrapolation in LLMs” on OpenResearch

About this paper

Xin Ma, Yang Liu, Jingjing Liu, Xiaoxu Ma

Mesa-Extrapolation, developed by researchers from Tsinghua University, Shanghai AI Lab, and Enn Group, introduces an efficient plug-in method for large language models (LLMs) to effectively process input sequences far exceeding their training lengths. The method, leveraging a novel Stair Position Encoding and chunk-based attention, extends context windows while demonstrating approximately linear memory usage and the fastest inference speeds among competitive approaches, reaching up to 192k tokens for Phi-3-mini.

Key insights

• Successful LLM extrapolation is theorized to depend on maintaining hidden state values within a specific "reasonable range" or above a "threshold H" as sequence length increases.
• Theoretical analysis demonstrates that both traditional NoPE and naive PE methods cause hidden state values to fall below this critical threshold beyond the training length, leading to extrapolation failure.
• A meticulously designed "Weave PE" (like Stair PE) can theoretically maintain crucial hidden state values above the threshold, preventing performance degradation and enabling effective extrapolation.

Summary by alphaXiv.