InAttention: Linear Context Scaling for Transformers

TL;DR

This paper modify the decoder-only transformer, replacing self-attention with InAttention, which scales linearly with context length during inference by having tokens attend only to initial states, offering a scalable solution for long-range dependencies in transformer models.

Abstract

VRAM requirements for transformer models scale quadratically with context length due to the self-attention mechanism. In this paper we modify the decoder-only transformer, replacing self-attention with InAttention, which scales linearly with context length during inference by having tokens attend only to initial states. Benchmarking shows that InAttention significantly reduces VRAM usage during inference, enabling handling of long sequences on consumer GPUs. We corroborate that fine-tuning extends context length efficiently, improving performance on long sequences without high training costs. InAttention offers a scalable solution for long-range dependencies in transformer models, paving the way for further optimization.

Figures (11)

• Figure 1: Right: A sliding window mask with literal context length of 1. Outlined in red is the path information might take in a 3-layer transformer model giving an effective context length of 3.
• Figure 2: Next Token Game Demonstration
• Figure 3: Left: Visualized attention pairs for dense attention. Right: Visualized attention pairs for inattention.
• Figure 4: Left: GPTNeoX Model (without parallel residuals). Right: GPTNeoX Model modified to use InAttention. Red edges indicate the path of initial states.
• Figure 5: NazX(N) Family of Models