2Mamba2Furious: 복잡도는 선형, 정확도는 경쟁력

초록

선형 어텐션 트랜스포머는 효율성으로 인해 소프트맥스 어텐션의 강력한 대안으로 부상했습니다. 그러나 선형 어텐션은 일반적으로 표현력이 낮아 소프트맥스 어텐션 대비 정확도가 감소하는 경향이 있습니다. 본 연구는 소프트맥스 어텐션과 선형 어텐션 간의 정확도 격차를 해소하기 위해 강력한 선형 어텐션 변종인 Mamba-2를 개선합니다. 먼저 Mamba-2를 가장 기본적이고 중요한 구성 요소로 단순화하여 어떤 설계 선택이 높은 정확도를 달성하는지 평가합니다. 이 단순화된 Mamba 변종(Mamba-2S)을 바탕으로 A-마스크를 개선하고 은닉 상태의 차수를 증가시켜, 소프트맥스 어텐션에 근접한 정확도를 유지하면서 긴 컨텍스트 길이에서 훨씬 더 메모리 효율적인 2Mamba 방법을 제안합니다. 또한 Mamba-2가 소프트맥스 어텐션 정확도를 능가하도록 하는 요소를 분석합니다. 모든 실험에 대한 코드를 공개합니다.

English

Linear attention transformers have become a strong alternative to softmax attention due to their efficiency. However, linear attention tends to be less expressive and results in reduced accuracy compared to softmax attention. To bridge the accuracy gap between softmax attention and linear attention, we manipulate Mamba-2, a very strong linear attention variant. We first simplify Mamba-2 down to its most fundamental and important components, evaluating which specific choices make it most accurate. From this simplified Mamba variant (Mamba-2S), we improve the A-mask and increase the order of the hidden state, resulting in a method, which we call 2Mamba, that is nearly as accurate as softmax attention, yet much more memory efficient for long context lengths. We also investigate elements to Mamba-2 that help surpass softmax attention accuracy. Code is provided for all our experiments

2Mamba2Furious: 복잡도는 선형, 정확도는 경쟁력

2Mamba2Furious: Linear in Complexity, Competitive in Accuracy

초록

Support