인공지능

Mamba: 시퀀스 모델링 재정의 및 트랜스포머 아키텍처 개선

게재 2023 년 12 월 18 일

아유쉬 미탈 미탈

이 기사에서는 맘바에서 이 혁신적인 상태 공간 모델(SSM)이 시퀀스 모델링을 어떻게 혁신하는지 살펴보겠습니다. Albert Gu와 Tri Dao가 개발한 Mamba는 언어 처리, 유전체학, 오디오 분석과 같은 분야에서 복잡한 시퀀스를 처리하는 효율성이 뛰어납니다. 선택적 상태 공간을 사용한 선형 시간 시퀀스 모델링은 이러한 다양한 양식에서 탁월한 성능을 보장합니다.

우리는 특히 긴 시퀀스의 경우 전통적인 Transformer가 직면한 계산 문제를 극복하는 Mamba의 능력을 탐구할 것입니다. 상태 공간 모델의 선택적 접근 방식을 통해 추론 속도가 빨라지고 시퀀스 길이에 따른 선형 확장이 가능해 처리량이 크게 향상됩니다.

Mamba의 독창성은 빠른 처리 능력, 선택적 SSM 레이어, FlashAttention에서 영감을 받은 하드웨어 친화적인 디자인에 있습니다. 이러한 기능을 통해 Mamba는 변환기 접근 방식을 기반으로 한 모델을 포함하여 많은 기존 모델보다 뛰어난 성능을 발휘하여 기계 학습 분야에서 주목할 만한 발전을 이루었습니다.

트랜스포머 대 맘바

GPT-4와 같은 Transformer는 자연어 처리 분야에서 벤치마크를 설정했습니다. 그러나 시퀀스가 길어지면 효율성이 떨어집니다. Mamba는 긴 시퀀스를 보다 효율적으로 처리하는 능력과 전체 프로세스를 단순화하는 고유한 아키텍처를 통해 앞서 나가고 있습니다.

언어 모델의 텍스트와 같은 데이터 시퀀스를 처리하는 데 능숙한 변환기입니다. 데이터를 순차적으로 처리했던 이전 모델과 달리 Transformers는 전체 시퀀스를 동시에 처리하므로 데이터 내의 복잡한 관계를 포착할 수 있습니다.

그들은 예측을 할 때 모델이 시퀀스의 다른 부분에 집중할 수 있도록 하는 주의 메커니즘을 사용합니다.

이 관심은 입력 데이터에서 파생된 쿼리, 키, 값의 세 가지 가중치 세트를 사용하여 계산됩니다. 시퀀스의 각 요소는 다른 모든 요소와 비교되어 시퀀스의 다음 요소를 예측할 때 각 요소가 받아야 하는 중요성 또는 '주의'를 나타내는 가중치를 제공합니다.

Transformer는 입력 데이터를 처리하는 인코더와 출력을 생성하는 디코더라는 두 가지 주요 블록을 유지합니다. 인코더는 여러 레이어로 구성되며, 각 레이어에는 다중 헤드 자체 주의 메커니즘과 단순하고 위치별로 완전히 연결된 피드포워드 네트워크라는 두 개의 하위 레이어가 포함되어 있습니다. 정규화 및 잔차 연결은 각 하위 계층에서 사용되어 심층 네트워크 훈련에 도움이 됩니다.

디코더에는 인코더와 유사한 두 개의 하위 계층이 있는 레이어도 있지만 인코더의 출력에 대해 다중 헤드 주의를 수행하는 세 번째 하위 계층을 추가합니다. 디코더의 순차적 특성으로 인해 위치에 대한 예측은 자동 회귀 속성을 유지하면서 이전 위치만 고려할 수 있습니다.

Transformers와 달리 Mamba 모델은 다른 접근 방식을 취합니다. Transformers는 더 복잡한 Attention 메커니즘을 사용하여 긴 시퀀스 문제를 처리하는 반면 Mamba는 선택적 상태 공간을 사용하여 더 많은 컴퓨팅을 제공합니다.

다음은 변환기의 작동 방식에 대한 높은 수준의 개요입니다.

입력 처리: 변환기는 먼저 입력 데이터를 모델이 이해할 수 있는 형식으로 인코딩하며, 종종 시퀀스의 각 요소 위치도 통합하는 임베딩을 사용합니다.
주의 메커니즘: 핵심적으로 어텐션 메커니즘은 현재 요소를 이해할 때 입력 시퀀스의 다른 부분에 얼마나 집중해야 하는지를 나타내는 점수를 계산합니다.
인코더-디코더 아키텍처: 변환기 모델은 입력을 처리하는 인코더와 출력을 생성하는 디코더로 구성됩니다. 각각은 입력에 대한 모델의 이해를 구체화하는 여러 레이어로 구성됩니다.
다중 헤드 주의: 인코더와 디코더 모두 내에서 다중 헤드 주의를 사용하면 모델이 서로 다른 표현 공간에서 시퀀스의 서로 다른 부분에 동시에 주의를 기울일 수 있으므로 다양한 컨텍스트에서 학습하는 능력이 향상됩니다.
위치별 피드포워드 네트워크: 주목 후 간단한 신경망은 각 위치의 출력을 개별적으로 동일하게 처리합니다. 이는 잔여 연결을 통해 입력과 결합된 후 레이어 정규화가 수행됩니다.
출력 생성: 그런 다음 디코더는 인코더의 컨텍스트와 지금까지 생성된 내용의 영향을 받아 출력 시퀀스를 예측합니다.

시퀀스를 병렬로 처리하는 변환기의 능력과 강력한 주의 메커니즘은 번역 및 텍스트 생성과 같은 작업에 강력합니다.

대조적으로, Mamba 모델은 선택적 상태 공간을 사용하여 시퀀스를 처리함으로써 다르게 작동합니다. 이 접근 방식은 긴 시퀀스를 처리할 때 Transformer의 계산 비효율성을 해결합니다. Mamba의 디자인은 더 빠른 추론을 가능하게 하고 시퀀스 길이에 따라 선형적으로 확장하여 특히 시퀀스가 점점 길어짐에 따라 더 효율적일 수 있는 시퀀스 모델링을 위한 새로운 패러다임을 설정합니다.

맘바

Mamba를 정말 독특하게 만드는 것은 전통적인 주의 및 MLP 블록에서 벗어났다는 것입니다. 이러한 단순화로 인해 시퀀스 길이에 따라 선형적으로 확장되는 더 가볍고 빠른 모델이 탄생했습니다. 이는 이전 모델과 비교할 수 없는 성과입니다.

Mamba의 주요 기능은 다음과 같습니다.

선택적 SSM: 이를 통해 Mamba는 관련 없는 정보를 필터링하고 관련 데이터에 집중하여 시퀀스 처리를 향상할 수 있습니다. 이러한 선택성은 효율적인 내용 기반 추론에 매우 중요합니다.
하드웨어 인식 알고리즘: Mamba는 최신 하드웨어, 특히 GPU에 최적화된 병렬 알고리즘을 사용합니다. 이 설계를 통해 기존 모델에 비해 더 빠른 계산이 가능하고 메모리 요구 사항이 줄어듭니다.
단순화된 아키텍처: 선택적 SSM을 통합하고 Attention 및 MLP 블록을 제거함으로써 Mamba는 보다 단순하고 동질적인 구조를 제공합니다. 이로 인해 확장성과 성능이 향상됩니다.

Mamba는 언어, 오디오, 유전체학을 포함한 다양한 영역에서 탁월한 성능을 입증했으며 사전 훈련과 영역별 작업 모두에서 탁월합니다. 예를 들어, 언어 모델링에서 Mamba는 더 큰 Transformer 모델의 성능과 일치하거나 그 이상입니다.

Mamba의 코드와 사전 훈련된 모델은 다음 커뮤니티에서 공개적으로 사용할 수 있습니다. GitHub의.

선형 모델의 경우 표준 복사 작업이 간단합니다. 선택적 복사 및 인덕션 헤드에는 LLM을 위한 동적 콘텐츠 인식 메모리가 필요합니다.

구조화된 상태 공간(S4) 모델은 최근 RNN, CNN 및 기존 상태 공간 모델의 특성을 포함하는 유망한 시퀀스 모델 클래스로 등장했습니다. S4 모델은 연속 시스템, 특히 암시적 잠재 상태를 통해 XNUMX차원 기능 또는 시퀀스를 매핑하는 시스템 유형에서 영감을 얻습니다. 딥 러닝의 맥락에서 이는 효율적이고 적응성이 뛰어난 시퀀스 모델을 설계하기 위한 새로운 방법론을 제공하는 중요한 혁신을 나타냅니다.

S4 모델의 역동성

SSM(S4) 이것이 기본적인 구조화된 상태 공간 모델입니다. 순서가 필요합니다 x 그리고 출력을 생성합니다 y 학습된 매개변수 사용 A, B, C및 지연 매개변수 Δ. 변환에는 매개변수를 이산화하고(연속 함수를 이산 함수로 전환) 시간 불변인 SSM 연산을 적용하는 작업이 포함됩니다. 즉, 다른 시간 단계에 따라 변경되지 않습니다.

이산화의 중요성

이산화는 고정된 공식을 통해 연속 매개변수를 이산 매개변수로 변환하여 S4 모델이 연속시간 시스템과의 연결을 유지할 수 있도록 하는 핵심 프로세스입니다. 이는 모델에 해상도 불변성과 같은 추가 속성을 부여하고 적절한 정규화를 보장하여 모델 안정성과 성능을 향상시킵니다. 또한 이산화는 네트워크를 통한 정보 흐름을 관리하는 데 중요한 RNN의 게이팅 메커니즘과 유사합니다.

선형 시불변(LTI)

S4 모델의 핵심 기능은 선형 시불변입니다. 이 속성은 모델의 역학이 모든 시간 단계에 대해 고정된 매개변수를 사용하여 시간이 지나도 일관되게 유지된다는 것을 의미합니다. LTI는 반복 및 컨볼루션의 초석으로, 시퀀스 모델 구축을 위한 간단하면서도 강력한 프레임워크를 제공합니다.

근본적인 한계를 극복하다

S4 프레임워크는 전통적으로 LTI 특성으로 인해 제한되어 왔으며 이는 적응형 역학이 필요한 데이터 모델링에 어려움을 안겨주었습니다. 최근 연구 논문에서는 시간에 따라 변하는 매개변수를 도입하여 LTI의 제약을 제거함으로써 이러한 한계를 극복하는 접근 방식을 제시합니다. 이를 통해 S4 모델은 더욱 다양한 시퀀스와 작업 세트를 처리하여 적용 가능성을 크게 확장할 수 있습니다.

'상태 공간 모델'이라는 용어는 잠재 상태와 관련된 모든 반복 프로세스를 광범위하게 다루며 여러 분야에 걸쳐 다양한 개념을 설명하는 데 사용되었습니다. 딥 러닝의 맥락에서 S4 모델 또는 구조화된 SSM은 복잡한 시퀀스를 모델링하는 기능을 유지하면서 효율적인 계산을 위해 최적화된 특정 클래스의 모델을 나타냅니다.

S4 모델은 엔드투엔드 신경망 아키텍처에 통합되어 독립형 시퀀스 변환 기능을 수행할 수 있습니다. 이는 다양한 신경망 아키텍처에서 시퀀스 모델링을 위한 백본을 제공하는 CNN의 컨볼루션 레이어와 유사하다고 볼 수 있습니다.

SSM 대 SSM + 선택

시퀀스 모델링의 선택성에 대한 동기

구조화된 SSM

이 논문에서는 시퀀스 모델링의 기본 측면이 컨텍스트를 관리 가능한 상태로 압축하는 것이라고 주장합니다. 입력에 선택적으로 집중하거나 필터링할 수 있는 모델은 이러한 압축 상태를 유지하는 보다 효과적인 수단을 제공하여 보다 효율적이고 강력한 시퀀스 모델로 이어집니다. 이러한 선택성은 모델이 시퀀스 차원을 따라 정보가 흐르는 방식을 적응적으로 제어하는 데 필수적이며, 이는 언어 모델링 및 그 이상에서 복잡한 작업을 처리하는 데 필수적인 기능입니다.

선택적 SSM은 매개변수가 입력에 종속되도록 함으로써 기존 SSM을 향상시키며, 이는 이전에는 시불변 모델로는 달성할 수 없었던 적응성을 제공합니다. 이로 인해 효율적인 계산을 위해 더 이상 컨볼루션을 사용할 수 없고 대신 선형 반복 메커니즘에 의존하는 시변 SSM이 생성됩니다. 이는 기존 모델과 상당한 차이가 있습니다.

SSM + 선택(S6) 이 변형에는 매개변수에 입력 종속성을 추가하는 선택 메커니즘이 포함되어 있습니다. B 그리고 C및 지연 매개변수 Δ. 이를 통해 모델은 입력 시퀀스의 특정 부분에 선택적으로 집중할 수 있습니다. x. 매개변수는 선택 사항을 고려하여 이산화되고 SSM 작업은 요소를 순차적으로 처리하여 시간에 따라 동적으로 초점을 조정하는 스캔 작업을 사용하여 시간에 따라 변하는 방식으로 적용됩니다.

Mamba의 성능 하이라이트

Mamba는 모든 평가 결과에서 동급 최고입니다.

성능 측면에서 Mamba는 추론 속도와 정확성 모두에서 뛰어납니다. 이 디자인은 DNA와 오디오 모델링 모두에서 입증된 더 긴 컨텍스트를 더 잘 활용할 수 있도록 하며, 장거리 종속성을 요구하는 복잡한 작업에서 이전 모델보다 성능이 뛰어납니다. 다양한 작업에 대한 제로샷 평가에서도 그 다양성이 강조되어 효율성과 확장성 측면에서 이러한 모델에 대한 새로운 표준을 설정합니다.

맘바 시작하기

Mamba 활용에 관심이 있는 경우 기술 요구 사항에는 Linux OS, NVIDIA GPU, PyTorch 1.12+ 및 CUDA 11.6+가 포함됩니다. 설치에는 다음에서 필요한 패키지를 설치하는 간단한 pip 명령이 포함됩니다. 맘바 저장소. PyTorch 버전에서 호환성 문제가 발생하는 경우 pip와 함께 –no-build-isolation 플래그를 사용하면 도움이 될 수 있습니다. Pile 및 SlimPajama 데이터세트와 같은 광범위한 데이터세트로 훈련된 이러한 모델은 다양한 계산 요구 사항과 성능 벤치마크를 충족하도록 설계되었습니다.

Mamba는 선택적 SSM 계층부터 Mamba 블록 및 전체 언어 모델 구조에 이르기까지 다양한 수준의 인터페이스를 제공합니다. 아키텍처의 주요 모듈인 Mamba 블록은 인과 Conv1d 레이어를 활용하며 신경망 설계에 쉽게 통합될 수 있습니다. 제공된 Python 사용 예제는 Mamba 모델을 인스턴스화하고 이를 통해 데이터를 처리하는 방법을 보여 주며 시스템의 단순성과 유연성을 강조합니다.

사전 훈련된 Mamba 모델은 다음에서 사용할 수 있습니다. 포옹하는 얼굴, 130M에서 2.8B 범위의 매개변수 크기를 가지며 광범위한 Pile 데이터세트와 SlimPajama 데이터세트에 대해 훈련되었습니다. 이러한 모델은 GPT-3의 차원 표준을 준수하면서 다양한 계산 및 성능 요구 사항을 충족하도록 설계되었습니다. 사용자는 이러한 모델에서 높은 처리량과 정확성을 기대할 수 있으므로 Mamba는 언어 모델링을 포함하되 이에 국한되지 않는 다양한 응용 프로그램에서 경쟁력 있는 선택이 됩니다.

맘바의 영향력

Mamba는 정보 밀도가 높은 데이터를 처리하기 위한 Transformer 아키텍처에 대한 강력한 대안을 제공하여 시퀀스 모델링의 도약을 나타냅니다. 그 디자인은 최신 하드웨어의 요구 사항에 맞춰 메모리 사용과 병렬 처리 기능을 모두 최적화합니다. Mamba의 코드베이스와 사전 훈련된 모델의 오픈 소스 가용성을 통해 Mamba는 AI 및 딥 러닝 분야의 연구원과 개발자가 접근할 수 있고 강력한 도구가 됩니다.

관련 주제 :주의 메커니즘 GPT 맘바 변압기

다음 위로

HierSpeech++: 제로샷 음성 합성을 위한 계층적 변이 추론

놓치지 마세요.

NeurIPS 2023의 주요 내용 및 기여

아유쉬 미탈

저는 지난 50년 동안 기계 학습과 딥 러닝의 매혹적인 세계에 몰두했습니다. 저의 열정과 전문 지식은 특히 AI/ML에 중점을 둔 XNUMX개 이상의 다양한 소프트웨어 엔지니어링 프로젝트에 기여하도록 이끌었습니다. 나의 계속되는 호기심은 또한 내가 더 탐구하고 싶은 분야인 자연어 처리로 나를 이끌었습니다.