배고픈 히포들: 상태공간모델을 통한 언어모델링으로 나아가기

상태공간모델(State space models, SSMs)은 일부 모달리티에서 최첨단의 시계열 모델링 성능을 보여주었으나, 언어 모델링에서는 어텐션 메커니즘에 미치지 못하고 있다. 또한, 시퀀스 길이에 대해 거의 선형적으로 확장되는 반면에 어텐션의 경우 제곱적으로 확장되는 점을 고려하면 SSM은 이론적으로 더 효율적인 구조를 지니고 있으나, 하드웨어 활용도가 낮아 여전히 트랜스포머보다 느리다. 본 논문에서는 SSM과 어텐션 간의 언어 모델링 성능 격차를 이해하는 데 있어 진전을 이루었으며, SSM과 어텐션 간의 하드웨어 효율성 장벽을 줄이는 데 기여하였다. 먼저, 합성 언어 모델링 작업을 통해 SSM과 어텐션 간의 성능 차이를 분석하였다. 그 결과, 기존의 SSM은 시퀀스 내에서 이전 토큰을 재현하거나 시퀀스 전체에 걸쳐 토큰을 비교하는 능력에서 어려움을 겪는다는 점을 발견하였다. 이러한 문제의 언어 모델링에 미치는 영향을 이해하기 위해, 이러한 능력을 특별히 고려해 설계된 새로운 SSM 레이어인 H3를 제안한다. H3는 합성 언어에서는 어텐션과 동등한 성능을 보이며, OpenWebText에서 트랜스포머와의 퍼플렉서티(PPL) 차이가 0.4 이내로 근접한다. 또한, 두 개의 어텐션 레이어를 유지하는 하이브리드 125M 파라미터 H3-어텐션 모델이 OpenWebText에서 트랜스포머보다 1.0 PPL 우수한 성능을 달성한다. 다음으로, 현대 하드웨어에서 SSM의 학습 효율성을 향상시키기 위해 FlashConv를 제안한다. FlashConv는 블록 FFT 알고리즘을 융합하여 길이가 최대 8K까지의 시퀀스에서 효율성을 향상시키며, SSM의 반복적 성질을 활용하는 새로운 상태 전달 알고리즘을 도입하여 더 긴 시퀀스로 확장할 수 있도록 하였다. FlashConv는 장거리 벤치마크에서 2배의 속도 향상을 가져오며, 하이브리드 언어 모델이 트랜스포머보다 텍스트 생성 속도를 2.4배 빠르게 수행할 수 있게 한다. FlashConv를 활용하여 Pile 데이터셋에서 하이브리드 H3-어텐션 모델을 최대 27억 파라미터 규모로 확장하였으며, 초기 결과로서 트랜스포머보다 낮은 퍼플렉서티를 달성했으며, SuperGLUE 벤치마크의 대부분의 작업에서 제로-샷 및 희소-샷 학습에서 트랜스포머를 능가하는 성능을 보였다.