[논문 리뷰] Simple linear attention language models balance the recall-throughput tradeoff

Simran Arora, Sabri Eyuboglu외 7명|arXiv (Cornell University)|2024. 02. 28.

Topic Modeling인용 수 4

한 줄 요약

이 논문은 선형 어텐션과 슬라이딩 윈도우 어텐션을 조합한 하이브리드 언어 모델 아키텍처인 Based를 제안한다. 이는 복구 능력과 추론 효율성의 균형을 맞추기 위한 것이다. 윈도우 크기와 특징 차원을 조정함으로써 Based는 메모리 사용을 크게 줄이고 FlashAttention-2 대비 24배 높은 처리량을 달성하면서도 표준 어텐션 수준의 복구 능력을 확보한다. 이는 Mamba와 같은 SOTA 모델보다 복구 중심 작업에서 6.22점 높은 정확도를 기록한다.

ABSTRACT

Recent work has shown that attention-based language models excel at recall, the ability to ground generations in tokens previously seen in context. However, the efficiency of attention-based models is bottle-necked during inference by the KV-cache's aggressive memory consumption. In this work, we explore whether we can improve language model efficiency (e.g. by reducing memory consumption) without compromising on recall. By applying experiments and theory to a broad set of architectures, we identify a key tradeoff between a model's state size and recall ability. We show that efficient alternatives to attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but struggle at recall. We propose BASED a simple architecture combining linear and sliding window attention. By varying BASED window size and linear attention feature dimension, we can dial the state size and traverse the pareto frontier of the recall-memory tradeoff curve, recovering the full quality of attention on one end and the small state size of attention-alternatives on the other. We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points. Implementations of linear attention are often less efficient than optimized standard attention implementations. To make BASED competitive, we develop IO-aware algorithms that enable 24x higher throughput on language generation than FlashAttention-2, when generating 1024 tokens using 1.3b parameter models. Code for this work is provided at: https://github.com/HazyResearch/based.

연구 동기 및 목표

• 자기연쇄 언어 모델에서 복구 성능와 추론 효율성 사이의 상충 관계를 해결하기 위해.
• 다양한 아키텍처에서 모델 상태 크기와 복구 능력 사이의 파레토 경계를 식별하고 탐색하기 위해.
• 표준 어텐션처럼 높은 복구 능력을 유지하면서도 Mamba, RWKV 등의 효율적 대안처럼 메모리 소비를 줄이는 모델을 설계하기 위해.
• 실제 처리량 벤치마크에서 경쟁력을 갖추기 위해 I/O에 민감한 최적화 기법을 개발하기 위해.
• 선형 어텐션과 슬라이딩 윈도우 어텐션을 조합함으로써 개별 구성 요소보다 열등한 성능을 내는 파레토 최적의 솔루션을 만들 수 있는지 확인하기 위해.

제안 방법

• 모델 레이어 내부에서 선형 어텐션과 슬라이딩 윈도우 어텐션을 조합한 하이브리드 아키텍처인 Based를 도입한다.
• 가변적인 슬라이딩 윈도우 크기(예: 64, 128)와 선형 어텐션의 특징 차원을 설정하여 상태 크기와 트레이드오프 행동을 제어한다.
• 생성 과정에서 메모리 액세스 오버헤드를 줄이고 처리량을 향상시키기 위해 I/O에 민감한 알고리즘 최적화를 적용한다.
• Adam, 코즈인 감소, BFloat16 정밀도를 사용한 표준 훈련 레시피를 활용해 최대 1.3B 파라미터 모델을 훈련한다.
• 슬라이딩 윈도우 구성 요소에는 로테이션 위치 인코딩을 사용하고, 선형 구성 요소에는 위치 인코딩을 사용하지 않는다.
• 블록 전용 다항식 분석을 구현하여, 연관 복구 작업을 해결하기 위해 충분한 모델 깊이가 필요하다는 이론적 근거를 제시한다.

Figure 1: Based overview. Combining linear attention with tiny sliding window softmax attention (e.g., 64 or 128 tokens in width) enables improved recall accuracy with limited efficiency overhead vs. smaller tile sizes. ( Left ) Time to execute Cutlass GEMMs ( $y$ ) vs. sliding window attention size

실험 결과

연구 질문

• RQ1 Mamba와 같은 효율적 대안의 작은 상태 크기와 표준 어텐션의 복구 성능을 동시에 확보할 수 있는가?
• RQ2 다양한 아키텍처 유형 간에 모델 상태 크기와 복구 능력 사이의 근본적 트레이드오프는 무엇인가?
• RQ3 선형 어텐션과 슬라이딩 윈도우 어텐션을 조합하면 개별 구성 요소보다 뛰어난 파레토 최적 솔루션을 만들 수 있는가?
• RQ4 I/O 효율적인 알고리즘은 하이브리드 어텐션 모델이 실세계 처리량에서 FlashAttention과 경쟁할 수 있도록 만들 수 있는가?
• RQ5 선형 어텐션과 국소 어텐션 메커니즘의 조합은 개별적으로 사용할 때보다 더 나은 장거리 및 정밀한 토큰 상호작용을 가능하게 하는가?

주요 결과

• 기본 상태 크기에도 불구하고, 실세계 복구 중심 작업에서 Mamba보다 6.22점 높은 정확도를 기록하여 뛰어난 복구 능력을 입증한다.
• 1.3B 파라미터 모델에서, 1024개 토큰을 자동으로 생성할 때 Based는 FlashAttention-2 대비 24배 높은 처리량을 달성한다.
• 하이브리드 아키텍처는 복구-메모리 트레이드오프의 파레토 경계를 성공적으로 탐색하여, 소모적인 효율성 오버헤드만으로 표준 어텐션 수준의 복구 능력을 회복한다.
• 선형 어텐션만으로는 국소 토큰 비교의 정밀도 부족으로 인해 연관 복구 작업에서 실패하고, 슬라이딩 윈도우 어텐션은 윈도우 크기 제한으로 인해 제한된다.
• 이론적 분석을 통해 데이터에 독립적인 BaseConv 모델는 연관 복구를 해결하기 위해 최소 ⌊log(2p)⌋층이 필요하다고 밝혀져, 이러한 모델에서 충분한 깊이의 필요성을 검증한다.
• I/O에 민감한 최적화 기법을 통해, 하이브리드 어텐션 구성 요소가 있음에도 불구하고 Based는 최적화된 어텐션 구현과 동일하거나 이를 초월하는 처리량을 확보한다.

Figure 2: Throughput (memory) - recall tradeoff. $x$ -axis shows state size (bytes) during generation; $y$ -axis shows accuracy on the MQAR recall task [ 1 ] . For each architecture, we train several models varying hyperparameters that affect the recurrent state size ( e.g. model dimension). The plo