[논문 리뷰] Hungry Hungry Hippos: Towards Language Modeling with State Space Models

Daniel Y. Fu, Tri Dao외 4명|arXiv (Cornell University)|2022. 12. 28.

Topic Modeling인용 수 116

한 줄 요약

본 논문은 언어 모델링에서 어텐션과의 표현력 차이를 해소하기 위해 설계된 새로운 상태 공간 모듈 계층 H3와, 긴 시퀀스에서 SSM의 효율적 학습 및 추론을 가능하게 하는 하드웨어 인식 FFT 기반 가속화인 FlashConv를 소개한다.

ABSTRACT

State space models (SSMs) have demonstrated state-of-the-art sequence modeling performance in some modalities, but underperform attention in language modeling. Moreover, despite scaling nearly linearly in sequence length instead of quadratically, SSMs are still slower than Transformers due to poor hardware utilization. In this paper, we make progress on understanding the expressivity gap between SSMs and attention in language modeling, and on reducing the hardware barrier between SSMs and attention. First, we use synthetic language modeling tasks to understand the gap between SSMs and attention. We find that existing SSMs struggle with two capabilities: recalling earlier tokens in the sequence and comparing tokens across the sequence. To understand the impact on language modeling, we propose a new SSM layer, H3, that is explicitly designed for these abilities. H3 matches attention on the synthetic languages and comes within 0.4 PPL of Transformers on OpenWebText. Furthermore, a hybrid 125M-parameter H3-attention model that retains two attention layers surprisingly outperforms Transformers on OpenWebText by 1.0 PPL. Next, to improve the efficiency of training SSMs on modern hardware, we propose FlashConv. FlashConv uses a fused block FFT algorithm to improve efficiency on sequences up to 8K, and introduces a novel state passing algorithm that exploits the recurrent properties of SSMs to scale to longer sequences. FlashConv yields 2$ imes$ speedup on the long-range arena benchmark and allows hybrid language models to generate text 2.4$ imes$ faster than Transformers. Using FlashConv, we scale hybrid H3-attention language models up to 2.7B parameters on the Pile and find promising initial results, achieving lower perplexity than Transformers and outperforming Transformers in zero- and few-shot learning on a majority of tasks in the SuperGLUE benchmark.

연구 동기 및 목표

• 합성 과제를 사용하여 언어 모델링에서 SSM과 어텐션 간의 표현력 차이를 평가한다.
• 과거 토큰을 기억하고 시퀀스 전반에서 토큰 간 비교를 가능하게 하는 SSM 계층(H3)을 설계한다.
• 표준 벤치마크에서 트랜스포머 성능에 맞서거나 이를 능가하도록 SSM의 하드웨어 효율성을 향상시킨다.
• FlashConv를 활용한 대규모 SSM 기반 모델의 확장 가능한 더 빠른 학습 및 생성을 입증한다.

제안 방법

• 토큰 기억 및 토큰 간 비교를 가능하게 하기 위한 이동(shift) 및 대각 행렬을 갖는 두 개의 SSM 계층과 곱셈 입력 프로젝션을 포함하는 H3를 도입한다.
• 합성 언어 과제에서 H3가 어텐션과 일치하고 OpenWebText에서 트랜스포머와의 격차를 거의 좁히는 것을 보여주며(0.4 perplexity 이내).
• OpenWebText에서 트랜스포머를 1.0 perplexity로 능가하는 두 개의 어텐션 레이어를 가진 하이브리드 H3–어텐션 모델을 제안한다.
• 롱 시퀀스 및 더 큰 모델로 SSM을 확장하기 위해 블록 FFT와 상태 전달 알고리즘을 갖춘 융합형 FFT 기반 합성(convolution)인 FlashConv을 개발한다.
• FlashConv이 긴 시퀀스에서 최대 2×의 속도 향상을 제공하고 트랜스포머보다 텍스트 생성을 2.4× 빠르게 가능하게 함을 보여준다.
• Pile에서 최대 2.7B 파라미터까지 하이브리드 H3–어텐션 모델을 확장하고 perplexity 및 SuperGLUE 성능을 보고한다.
• 비슷한 크기에서 Hybrid H3가 트랜스포머보다 최대 2.4× 빠른 추론 처리량을 보인다는 비교를 제공한다.

실험 결과

연구 질문

• RQ1 합성 과제에서 SSM이 어텐션과 동일한 표현력을 보일 수 있는가?
• RQ2 SSM이 과거 토큰을 기억하고 시퀀스 전반에서 비교를 가능하게 하는 어떤 아키텍처적 변화(예: 이동/대각 SSM 및 곱셈 상호작용)가 필요한가?
• RQ3 현대 가속기에서 SSM과 어텐션 간의 런타임 격차를 줄이기 위해 FlashConv 같은 하드웨어 인식 알고리즘은 어떻게 작동하는가?
• RQ4 하이브리드 H3–어텐션 모델이 표준 벤치마크에서 트랜스포머에 비해 경쟁력 있는 perplexity 및 zero-/few-shot 성능을 달성하는가?
• RQ5 FlashConv가 수십억 파라미터 규모의 SSM 기반 모델로 확장될 수 있으며 현실적인 학습 및 추론 속도 향상을 제공하는가?

주요 결과

• H3가 합성 언어 과제에서 어텐션과 일치하며 OpenWebText에서 트랜스포머와의 격차를 거의 좁힌다(0.4 perplexity 이내).
• 두 개의 어텐션 레이어를 갖춘 하이브리드 H3–어텐션 모델이 OpenWebText에서 트랜스포머를 1.0 perplexity로 능가한다.
• 크기가 125M–2.7B인 Pile에서 학습된 하이브리드 H3–어텐션 모델은 perplexity에서 트랜스포머 기준선을 능가하고 강한 SuperGLUE zero-/few-shot 성능을 보여준다.
• FlashConv는 긴 시퀀스에서 최대 2×의 속도 향상을 제공하고 SSM 기반 생성이 트랜스포머보다 2.4× 빠르도록 한다.
• Pile에서 2.7B 파라미터로 확장된 SSM 기반 모델은 여러 설정에서 트랜스포머보다 더 낮은 perplexity를 달성하고 경쟁력 있는 zero-/few-shot 결과를 보인다.
• FlashConv은 Long Range Arena에서 트랜스포머 베이스라인 및 다른 SSM 방법들과 비교했을 때 최상위 속도를 달성한다.