2-D Transformer: Extending Large Language Models to Long-Context With Few Memory.

Abstract

The ability of processing long contexts is crucial for large language models (LLMs), but training LLMs with a long-context window requires substantial computational resources. Many sought to mitigate this through the sparse attention mechanism. However, sparse attention faces a noticeable gap compared with full attention in capturing long-distance information, leading to limited long-context processing capabilities. To effectively address this issue, this article proposes a novel sparse transformer architecture called 2-D transformer (2D-former), aimed at extending the context windows of pretrained LLMs while reducing GPU memory requirements. The 2D-former incorporates a 2-D attention mechanism that consists of a long-distance information compressor (LDIC) and a blockwise attention (BA) mechanism. LDIC can self-adaptively extract blockwise representational features by convolution and compress long-distance information into a set of tokens based on the significance of each block. The BA mechanism integrates these features, enabling each token to directly communicate with any of its preceding tokens during the computation of sparse attention. In this way, sparse attention can fully utilize long-distance information to bridge the gap with full attention while greatly reducing computational requirements. The 2D-former only needs to add less than 0.14% of additional trainable parameters to extend the context length of LLaMA2 7B to 32k on 4 A100 GPUs with 40-GB memory. In addition, it is compatible with most current acceleration techniques and parameter-efficient fine-tuning (PEFT) methods. Furthermore, we conduct supervised fine-tuning with 2D-former using our self-collected long-instruction fine-tuning dataset, named LongTuning, which comprises over 11k long-context question-answer (QA) pairs. Experimental results demonstrate that 2D-former achieves efficient long-context extension with minimal GPU memory and computational time consumption, while maintaining superior performance across both downstream long-context and short-context tasks.