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

Area of Science:

• Artificial Intelligence
• Natural Language Processing
• Machine Learning

Background:

• Processing long contexts is crucial for Large Language Models (LLMs), but training demands substantial computational resources.
• Existing sparse attention mechanisms show a performance gap compared to full attention in capturing long-distance information.
• This limitation hinders the effective long-context processing capabilities of LLMs.

Purpose of the Study:

• To propose a novel sparse transformer architecture, the 2-D transformer (2D-former).
• To extend the context windows of pretrained LLMs while reducing GPU memory requirements.
• To bridge the performance gap between sparse and full attention mechanisms for enhanced long-context processing.

Main Methods:

• Introduced a 2-D attention mechanism comprising a long-distance information compressor (LDIC) and a blockwise attention (BA) mechanism.
• LDIC uses convolution to extract blockwise features and compress long-distance information based on block significance.
• BA integrates these features, allowing direct communication between tokens for sparse attention computation.

Main Results:

• The 2D-former requires less than 0.14% additional trainable parameters to extend LLaMA2 7B context to 32k on 4 A100 GPUs.
• Achieved efficient long-context extension with minimal GPU memory and computational time.
• Demonstrated superior performance on both long-context and short-context downstream tasks after fine-tuning on the LongTuning dataset.

Conclusions:

• The 2D-former effectively extends LLM context length with significant reductions in computational demands.
• The architecture is compatible with current acceleration techniques and parameter-efficient fine-tuning (PEFT) methods.
• This approach offers a practical solution for enhancing LLM capabilities in handling extended contexts.