Mesa-Extrapolation: A Weave Position Encoding Method for Enhanced Extrapolation in LLMs

TL;DR

This paper introduces Mesa-Extrapolation, a novel weave position encoding method that enhances large language models' ability to extrapolate beyond their training lengths, while reducing memory and computation costs.

Contribution

The paper presents a new weave position encoding technique, Mesa-Extrapolation, with theoretical analysis and experimental validation showing improved extrapolation and efficiency in LLMs.

Findings

Mesa-Extrapolation improves LLM extrapolation performance.

The method reduces memory demand and inference time.

Theoretical analysis supports the effectiveness of weave PE.

Abstract

Large language models (LLMs), although having revolutionized many fields, still suffer from the challenging extrapolation problem, where the inference ability of LLMs sharply declines beyond their max training lengths. In this work, we conduct a theoretical analysis to better understand why No Position Encoding (NoPE) fails outside its effective range, as well as examining the power of Position Encoding (PE) in this context. Our findings reveal that with meticulous weave position, PE can indeed be extended beyond effective range. Our theorems establish that LLMs equipped with weave PE can achieve improved extrapolation performance without additional cost. Furthermore, we introduce a novel weave PE method, Mesa-Extrapolation, which utilizes a chunk-based triangular attention matrix and applies Stair PE to manage the final chunk. This method not only retains competitive performance but also offers substantial benefits such as significantly reduced memory demand and faster inference speed. Extensive experiments validate the effectiveness of Mesa-Extrapolation, demonstrating its potential as a scalable solution to enhancing LLMs applicative reach. Our code is available at \url{https://github.com/soacker/Mesa-Extrapolation}.