AI 세계에 새로운 사랑이 생겼습니다. 변신 로봇은 아닙니다.

링크 표

초록 및 1. 서론

2 상태 공간 모델

3 선택적 상태 공간 모델 및 3.1 동기: 압축 수단으로서의 선택

3.2 선택을 통한 SSM 개선

3.3 선택적 SSM의 효율적인 구현

3.4 단순화된 SSM 아키텍처

3.5 선택 메커니즘의 속성

3.6 추가 모델 세부 정보

4 경험적 평가 및 4.1 합성 작업

4.2 언어 모델링

4.3 DNA 모델링

4.4 오디오 모델링 및 생성

4.5 속도 및 메모리 벤치마크

4.6 모델 절제

5 토론

6 결론, 감사의 글 및 참고문헌

토론: 선택 메커니즘

B 관련 작업 및 B.1 S4 변형 및 파생

B.2 SSM 아키텍처

B.3 RNN과의 관계

B.4 선형 주의 및 B.5 긴 컨텍스트 모델

C 선택적 SSM의 역학

D 선택적 SSM을 위한 하드웨어 인식 알고리즘

E 실험 세부 사항 및 추가 결과 및 E.1 합성 작업

E.2 언어 모델링

E.3 DNA 모델링

E.4 오디오 세부 정보

E.5 효율성 벤치마크

추상적인

현재 딥 러닝에서 가장 흥미로운 애플리케이션의 대부분을 구동하는 Foundation 모델은 거의 보편적으로 Transformer 아키텍처와 핵심 어텐션 모듈을 기반으로 합니다. 선형 어텐션, 게이트 합성곱 및 순환 모델, 구조화된 상태 공간 모델(SSM)과 같은 많은 부분 이차 시간 아키텍처가 긴 시퀀스에서 Transformer의 계산 비효율성을 해결하기 위해 개발되었지만 언어와 같은 중요한 모달리티에서 어텐션만큼 좋은 성능을 보이지 못했습니다. 이러한 모델의 주요 약점은 내용 기반 추론을 수행할 수 없다는 것을 확인하고 여러 가지 개선을 했습니다. 첫째, SSM 매개변수를 입력의 함수로 두는 것만으로 이산 모달리티로 약점을 해결하여 모델이 현재 토큰에 따라 시퀀스 길이 차원을 따라 정보를 선택적으로 전파하거나 잊을 수 있습니다. 둘째, 이러한 변경으로 인해 효율적인 합성곱을 사용할 수 없지만 순환 모드에서 하드웨어 인식 병렬 알고리즘을 설계합니다. 이러한 선택적 SSM을 어텐션이나 MLP 블록(Mamba) 없이 단순화된 엔드투엔드 신경망 아키텍처에 통합합니다. Mamba는 빠른 추론(Transformers보다 5배 더 높은 처리량)과 시퀀스 길이의 선형 스케일링을 제공하며, 실제 데이터에서 최대 백만 길이의 시퀀스까지 성능이 향상됩니다. Mamba는 일반적인 시퀀스 모델 백본으로서 언어, 오디오, 유전체학과 같은 여러 모달리티에서 최첨단 성능을 달성합니다. 언어 모델링에서 Mamba-3B 모델은 동일한 크기의 Transformer보다 성능이 뛰어나고 사전 학습과 다운스트림 평가 모두에서 두 배 크기의 Transformer와 일치합니다.

1 서론

기초 모델(FM) 또는 방대한 데이터에 사전 학습된 다음 다운스트림 작업에 맞게 조정된 대규모 모델은 현대 머신 러닝에서 효과적인 패러다임으로 등장했습니다. 이러한 FM의 백본은 종종 시퀀스 모델로, 언어, 이미지, 음성, 오디오, 시계열 및 유전체학과 같은 다양한 도메인의 임의의 입력 시퀀스에서 작동합니다(Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz et al. 2019; Oord et al. 2016; Poli et al. 2023; Sutskever, Vinyals, and Quoc V Le 2014). 이 개념은 특정 모델 아키텍처 선택에 무관하지만, 최신 FM은 주로 단일 유형의 시퀀스 모델, 즉 Transformer(Vaswani et al. 2017)와 핵심 어텐션 계층(Bahdanau, Cho, and Bengio 2015)을 기반으로 합니다. 셀프 어텐션의 효능은 컨텍스트 창 내에서 정보를 밀집하게 라우팅하여 복잡한 데이터를 모델링할 수 있는 능력에 기인합니다. 그러나 이 속성은 근본적인 단점을 가져옵니다. 유한 창 외부의 어떤 것도 모델링할 수 없고 창 길이에 대한 2차 스케일링입니다. 이러한 단점을 극복하기 위해 보다 효율적인 어텐션 변형에 대한 엄청난 양의 연구가 있었지만(Tay, Dehghani, Bahri, et al. 2022), 종종 효과적으로 만드는 바로 그 속성을 희생했습니다. 아직까지 이러한 변형 중 어느 것도 도메인 전체에서 규모에 따라 경험적으로 효과적인 것으로 나타나지 않았습니다. 최근, 구조화된 상태 공간 시퀀스 모델(SSM)(Gu, Goel 및 Ré 2022; Gu, Johnson, Goel 등 2021)이 시퀀스 모델링을 위한 유망한 아키텍처 클래스로 등장했습니다. 이러한 모델은 고전적인 상태 공간 모델(Kalman 1960)에서 영감을 받아 순환 신경망(RNN)과 합성곱 신경망(CNN)의 조합으로 해석할 수 있습니다. 이러한 클래스의 모델은 시퀀스 길이에서 선형 또는 거의 선형적인 스케일링을 사용하여 순환 또는 합성곱으로 매우 효율적으로 계산할 수 있습니다. 또한 특정 데이터 모달리티에서 장거리 종속성(Gu, Dao 등 2020)을 모델링하기 위한 원칙적인 메커니즘을 가지고 있으며 Long Range Arena(Tay, Dehghani, Abnar 등 2021)와 같은 벤치마크를 지배했습니다. 많은 종류의 SSM(Gu, Goel, and Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu, and Berant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; Smith, Warrington, and Linderman 2023)은 오디오와 시각과 같은 연속적인 신호 데이터를 포함하는 도메인에서 성공적이었습니다(Goel et al. 2022; Nguyen, Goel, et al. 2022; Saon, Gupta, and Cui 2023). 그러나 텍스트와 같은 불연속적이고 정보가 밀집된 데이터를 모델링하는 데는 덜 효과적이었습니다.

우리는 선형적으로 시퀀스 길이를 확장하는 동시에 Transformers의 모델링 능력을 달성하기 위해 여러 축에 대한 기존 연구를 개선한 새로운 종류의 선택적 상태 공간 모델을 제안합니다.

선택 메커니즘. 먼저, 우리는 이전 모델의 주요 한계를 파악합니다. 입력에 따라 데이터를 효율적으로 선택하는 능력(즉, 특정 입력에 집중하거나 무시하는 능력). 선택적 복사 및 유도 헤드와 같은 중요한 합성 작업에 기반한 직관을 바탕으로, 우리는 입력에 따라 SSM 매개변수를 매개변수화하여 간단한 선택 메커니즘을 설계합니다. 이를 통해 모델은 관련 없는 정보를 걸러내고 관련 정보를 무기한 기억할 수 있습니다.

하드웨어 인식 알고리즘 . 이 간단한 변경은 모델 계산에 기술적 과제를 제기합니다. 사실, 모든 이전 SSM 모델은 계산 효율성을 위해 시간 및 입력 불변이어야 합니다. 우리는 합성곱 대신 스캔으로 모델을 반복적으로 계산하지만 GPU 메모리 계층의 다른 레벨 간의 IO 액세스를 피하기 위해 확장된 상태를 구체화하지 않는 하드웨어 인식 알고리즘으로 이를 극복합니다. 그 결과 구현은 이론적으로(모든 합성곱 기반 SSM의 경우 의사 선형과 비교하여 시퀀스 길이에서 선형적으로 확장) 및 최신 하드웨어(A100 GPU에서 최대 3배 더 빠름)에서 이전 방법보다 빠릅니다.

아키텍처. 이전 SSM 아키텍처(Dao, Fu, Saab, et al. 2023)의 설계와 Transformers의 MLP 블록을 단일 블록으로 결합하여 이전의 딥 시퀀스 모델 아키텍처를 단순화하여 선택적 상태 공간을 통합한 간단하고 동질적인 아키텍처 설계(Mamba)를 도출합니다. 선택적 SSM과 확장된 Mamba 아키텍처는 시퀀스에서 작동하는 일반 기초 모델의 백본으로 적합한 주요 속성을 가진 완전 반복 모델입니다. (i) 고품질: 선택성은 언어 및 유전체학과 같은 고밀도 모달리티에서 강력한 성능을 제공합니다. (ii) 빠른 학습 및 추론: 계산 및 메모리는 학습 중에 시퀀스 길이에 따라 선형적으로 확장되고 추론 중에 모델을 자기 회귀적으로 펼치는 데는 이전 요소의 캐시가 필요하지 않으므로 단계당 일정한 시간만 필요합니다. (iii) 긴 컨텍스트: 품질과 효율성이 합쳐져 최대 시퀀스 길이 1M의 실제 데이터에서 성능이 향상됩니다.

우리는 Mamba가 사전 훈련 품질과 도메인별 작업 성능 측면에서 일반 시퀀스 FM 백본으로서의 잠재력을 여러 유형의 모달리티와 설정에서 경험적으로 검증했습니다.

• 합성. 대규모 언어 모델의 핵심으로 제안된 복사 및 유도 헤드와 같은 중요한 합성 작업에서 Mamba는 이를 쉽게 해결할 뿐만 아니라 무한정 길게(>1M 토큰) 솔루션을 외삽할 수 있습니다.

• 오디오 및 유전체학. Mamba는 오디오 파형 및 DNA 시퀀스 모델링에서 SaShiMi, Hyena, Transformers와 같은 이전의 최첨단 모델을 능가합니다. 사전 학습 품질과 다운스트림 메트릭(예: 어려운 음성 생성 데이터 세트에서 FID를 절반 이상 줄임) 모두에서 그렇습니다. 두 설정 모두에서 최대 백만 길이의 시퀀스까지 더 긴 컨텍스트로 성능이 향상됩니다.

• 언어 모델링. Mamba는 사전 학습 퍼플렉시티와 다운스트림 평가 모두에서 Transformer 수준의 성능을 진정으로 달성한 최초의 선형 시간 시퀀스 모델입니다. 최대 1B 매개변수의 스케일링 법칙을 통해 Mamba가 LLaMa(Touvron et al. 2023)를 기반으로 하는 매우 강력한 최신 Transformer 학습 레시피를 포함한 광범위한 기준선의 성능을 능가한다는 것을 보여줍니다. Mamba 언어 모델은 비슷한 크기의 Transformer에 비해 5배의 생성 처리량을 가지고 있으며 Mamba-3B의 품질은 크기가 두 배인 Transformer와 일치합니다(예: Pythia-3B에 비해 상식적 추론에서 평균 4포인트 더 높고 Pythia-7B를 능가).

모델 코드와 사전 훈련된 체크포인트는 https://github.com/state-spaces/mamba에서 오픈 소스로 제공됩니다.

2 상태 공간 모델

일반 상태 공간 모델. 상태 공간 모델이라는 용어는 매우 광범위한 의미를 가지고 있으며, 단순히 잠재 상태를 가진 모든 반복적 프로세스의 개념을 나타냅니다. 이는 마르코프 결정 프로세스(MDP)(강화 학습(Hafner et al. 2020)), 동적 인과 모델링(DCM)(계산 신경 과학(Friston, Harrison, and Penny 2003)), 칼만 필터(제어(Kalman 1960)), 은닉 마르코프 모델(HMM) 및 선형 동적 시스템(LDS)(머신 러닝), 그리고 대규모 반복적(때로는 합성곱) 모델(딥 러닝)을 포함하여 다양한 분야에서 많은 이질적인 개념을 지칭하는 데 사용되었습니다.

이 논문 전체에서 우리는 "SSM"이라는 용어를 구조화된 SSM 또는 S4 모델 클래스(Gu, Goel 및 Ré 2022; Gu, Gupta 등 2022; Gupta, Gu 및 Berant 2022; Hasani 등 2023; Ma 등 2023; Smith, Warrington 및 Linderman 2023)에만 사용하고 이러한 용어를 서로 바꿔 사용합니다. 편의상 선형 재귀 또는 전역 합성곱 관점에 초점을 맞춘 모델(Y. Li 등 2023; Orvieto 등 2023; Poli 등 2023)과 같은 이러한 모델의 파생 모델을 포함할 수도 있으며 필요한 경우 미묘한 차이를 명확히 할 수 있습니다.

SSM 아키텍처. SSM은 엔드투엔드 신경망 아키텍처에 통합할 수 있는 독립형 시퀀스 변환입니다. (SSM 아키텍처를 SSNN이라고도 부르기도 하는데, 이는 SSM 계층에 대한 CNN과 선형 합성곱 계층의 관계와 같습니다.) 가장 잘 알려진 SSM 아키텍처 중 일부를 논의하며, 그 중 다수는 우리의 기본 기준선 역할도 할 것입니다.

• 선형 주의(Katharopoulos et al. 2020)는 퇴화된 선형 SSM으로 볼 수 있는 재발을 수반하는 자기 주의의 근사입니다.

• H3(Dao, Fu, Saab, et al. 2023)는 이 재귀를 S4를 사용하도록 일반화했습니다. 이는 두 개의 게이트 연결로 샌드위치된 SSM이 있는 아키텍처로 볼 수 있습니다(그림 3). H3는 또한 주 SSM 계층 전에 표준 로컬 합성곱을 삽입하는데, 이를 시프트-SSM으로 프레임화합니다.

• 하이에나(Poli et al. 2023)는 H3와 동일한 아키텍처를 사용하지만 S4 계층을 MLP 매개변수화된 글로벌 합성곱(Romero et al. 2021)으로 대체합니다.

• RetNet(Y. Sun et al. 2023)은 아키텍처에 추가 게이트를 추가하고 더 간단한 SSM을 사용하여 합성곱 대신 멀티 헤드 어텐션(MHA) 변형을 사용하는 대체 병렬화 가능 계산 경로를 허용합니다.

• RWKV(B. Peng et al. 2023)는 또 다른 선형 어텐션 근사(어텐션 프리 트랜스포머(S. Zhai et al. 2021))를 기반으로 언어 모델링을 위해 설계된 최근 RNN입니다. 주요 "WKV" 메커니즘에는 LTI 재귀가 포함되며 두 SSM의 비율로 볼 수 있습니다.

다른 밀접하게 관련된 SSM과 아키텍처는 확장된 관련 작업(부록 B)에서 더 자세히 논의됩니다. 우리는 특히 S5(Smith, Warrington, and Linderman 2023), QRNN(Bradbury et al. 2016), SRU(Lei et al. 2017)를 강조하는데, 이는 우리의 핵심 선택적 SSM과 가장 밀접하게 관련된 방법으로 봅니다.