Mamba 모델이란 무엇인가요?

Mamba 모델은 아마도 트랜스포머가 원래 명성을 얻은 작업인 언어 모델링에서 트랜스포머 모델의 효능에 필적하는 최초의 딥 러닝 아키텍처일 것입니다. 특히, Mamba 아키텍처는 주요 LLM 벤치마크 평가에서 동급 크기의 트랜스포머에 필적하는 성능을 보여주면서도 지연 시간 및 메모리 요구 사항 측면에서 훨씬 더 효율적이라는 것을 보여주었습니다.

Mamba 아키텍처는 Tri Dao와 Albert Gu가 2023년 논문 "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"에서 처음 소개했습니다. 1년 후, 그들은 첫 Mamba 관련 논문에 이어 SSM과 트랜스포머 간의 연결을 탐구하고 Mamba-2라고 명명한 Mamba 아키텍처의 세련되고 훨씬 빠른 버전을 제시하는 내용의 후속 논문을 발표했습니다.

비록 첫 Mamba 논문이 발표된 후 2년이 경과한 시점에도 트랜스포머가 LLM의 주역으로 남아있지만, 이 아키텍처는 점점 더 많은 오픈 소스 모델에 통합되고 있습니다. Mistral AI의 Codestral Mamba와 같은 일부 모델은 순수한 Mamba 모델입니다. AI2I의 Jamba 시리즈와 IBM Granite 4.0을 비롯한 이외 여러 모델들이 어텐션(트랜스포머) 레이어와 SSM(Mamba) 레이어를 모두 포함하는 하이브리드 모델입니다. 성능 기반의 이점 외에도, 비교적 저렴한 하드웨어에서 원활하게 실행되는 Mamba 기반 모델의 확산은 AI 액세스의 대중화를 약속합니다.

SSM은 원래 일부 입력을 기반으로 전기 신호, 날씨 패턴 또는 움직이는 물체의 궤적과 같은 연속 시퀀스의 다음 상태를 예측하도록 설계되었습니다. SSM은 개념적으로나 수학적으로나 2017년 트랜스포머가 도입되기 전 자연어 처리(NLP)를 지배했던 순환 신경망(RNN)뿐만 아니라 컨볼루션 신경망(CNN) 및 은닉 마르코프 모델(HMM)과 같은 다른 머신 러닝 알고리즘과도 연관이 있습니다.

이름에서 알 수 있듯이, SSM은 시스템의 상태와 각 변수의 가능성 범위를 상호 연계하여 설명하는 모든 상태 변수의 수학적 표현인 상태 공간을 모델링하여 동적 시스템의 다음 상태를 예측합니다.

SSM은 아웃풋 시퀀스 y(t)를 예측하기 위해 입력 시퀀스 x(t)를 RNN의 숨겨진 상태와 유사한 잠재된 상태 표현인 h(t)에 매핑합니다. 모든 SSM의 핵심에는 2개의 방정식이 있습니다.

• 상태 방정식,  $h^{\prime }\left(t\right)=A*h\left(t\right)+B*x\left(t\right)$ 
• 출력 방정식:  $y\left(t\right)=C*h\left(t\right)+D*x\left(t\right)$ 

모델의 주요 매개변수는 A, B, C 및 D이며, 일반적으로 가중치 행렬의 형태를 취합니다. 제어 이론과 같이 SSM이 일반적으로 사용되는 분야에서 이러한 행렬은 종종 고정된 것으로 가정됩니다. 즉, 이러한 행렬은 확립된 시스템의 역학을 나타내며, SSM은 원하는 아웃풋 y로 이어지는 입력 x를 찾는 데 사용됩니다.SSM에 대한 보다 현대적인 개념에서 이러한 행렬 자체는 머신 러닝을 통해 최적화되는 매개변수입니다. 딥 러닝 모델에서 이러한 행렬은 신경망의 학습 가능한 가중치로 표현됩니다.

상태 방정식은 상태가 어떻게 변하는지를 설명합니다. 행렬 A의 값은 각 상태 변수를 그대로 두면 시간이 지남에 따라 어떻게 변화하는지를 결정합니다. 행렬 B의 값은 텍스트 시퀀스의 다음 토큰이 각 상태 변수에 영향을 미치는 방식을 결정합니다.

언어 모델링에서 현재 상태는 각 토큰 후에 업데이트되는 텍스트 시퀀스의 컨텍스트를 나타냅니다. 그 역할은 트랜스포머 모델에서 KV 캐시의 역할과 동일합니다.

출력 방정식은 현재 상태가 아웃풋에 미치는 영향(행렬 C에 의해 매개된 대로)과 입력이 아웃풋에 직접 영향을 미치는 영향(행렬 D에 의해 매개된 대로)을 설명합니다. 행렬 D는 본질적으로 h(t) 자체의 모델링 외부에 있기 때문에 핵심 행렬 A, B 및 C에 초점을 맞추기 위해 SSM의 다이어그램 및 논의에서 생략되는 경우가 많습니다.

Mamba LLM에서는 출력 방정식을 사용하여 다음 토큰을 생성합니다.

기존 SSM은 연속 입력을 모델링하도록 설계되었지만, 텍스트 시퀀스(및 최신 딥 러닝 모델에서 처리되는 대부분의 기타 데이터 양식)는 불연속 입력입니다. SSM을 사용하여 불연속 시퀀스를 모델링하려면 연속 신호의 일부로 별개의 특정 타임스텝을 표현할 수 있는 수단이 필요합니다.

개념적으로 이산화는 특정 순간에 연속 함수 값을 샘플링하는 것에 해당합니다. 여기에는 각 이산 시간 단계 t에서 스냅샷이 샘플링되거나 "유지"되는 시간을 결정하는 새로운 매개변수(단계 크기 ∆)의 도입이 수반됩니다. ∆에 대한 조정은 데이터의 해상도(시계열 데이터의 경우) 또는 프레임 속도(동영상 데이터의 경우)와 같은 품질을 변경하는 것과 유사합니다. 여러 "이산화" 방법이 있지만, 대부분의 최신 SSM 변형(Mamba 포함)에서는 0차 유지(ZOH) 방법이 사용됩니다.

SSM을 이산화하면 시퀀스 간 작업에 RNN처럼 사용할 수 있습니다. 이산화된 SSM의 매개변수와 방정식은 연속 시간 동등성과 구별하기 위해 일반적으로 RNN에 많이 사용되는 아래 첨자 표기법을 사용하여 재작성됩니다. 이 표기법에서 h_t는 모델이 생성할 업데이트된 상태 공간을 나타내며, h_t-1은 그 이전의 상태, 즉 현재의 상태 공간을 나타냅니다.

 $h_t=\overline{A}{h}_{t-1}+\overline{B}{x}_t$ 
 $y_t=\overline{C}{h}_t$

표준 개별 SSM을 사용하여 텍스트 데이터를 모델링하는 것은 RNN과 공통된 여러 단점으로 인해 실용적이지 않습니다. 이러한 단점 중 두 가지는 2021년 Albert Gu 등이 구조화된 상태 공간 시퀀스 모델(또는 "S4 모델")을 도입함으로써 해결되었습니다.

S4 모델과 대각선 SSM(DSS), 대각선 S4(S4D) 및 H3 모델과 같은 많은 파생 모델의 성공은 Mamba를 위한 길을 열었습니다.

RNN의 특정 인스턴스와 동일한 이산화된 SSM의 이점은 RNN의 추론 속도가 매우 빠르다는 것입니다. 반면 단점은 RNN의 학습 속도가 매우 느리다는 것입니다.

다행히 이산화된 SSM에는 다른 RNN과 구별되는 한 가지 중요한 속성이 있습니다. 즉, 이산화된 SSM은 간단하고 직관적인 곱셈과 덧셈 연산만 사용합니다. S4 논문에서 알 수 있듯이 이러한 단순하고 반복적이며 상호 의존적인 선형 반복은 1차원 컨볼루션 커널로 언롤될 수 있으며, 이 커널은 하나의 단계로 아웃풋 x를 아웃풋 y에 직접 매핑합니다. 이는 빠른 푸리에 변환을 사용하여 매우 효율적으로 계산할 수 있습니다.

유일한 아쉬움은 전체 입력 시퀀스의 모든 단계를 알고 있을 때만 이러한 계산이 가능하다는 것입니다. 이는 추론 중에는 불가능하지만 훈련 중에는 가능합니다. 따라서 구조화된 SSM은 두 가지 장점을 모두 누릴 수 있습니다. 학습 중에는 CNN으로 매우 효율적으로 운영될 수 있으며, 추론 중에는 RNN으로 매우 효율적으로 운영될 수 있습니다.

대부분의 RNN과 마찬가지로, 표준 SSM은 본질적으로 장거리 종속성을 모델링하는 데 취약합니다. 즉, 단락의 시작과 끝에 있는 단어처럼 멀리 떨어져 있는 시퀀스의 단계 간 관계를 잘 이해하지 못하므로 긴 시퀀스를 모델링하는 데 약합니다.

이 문제를 해결하기 위해 Gu와 그의 공동 저자(그 중 한 명은 Tri Dao)들은 HiPPO(High-order Polynomial Projection Operators의 약자)라는 기술을 사용하여 A 행렬과 B 행렬이 작동하는 방식을 정의했습니다. 이를 위해 직교 다항식에서 파생된 공식을 사용하여 초기 값을 구조화했습니다. 이는 모델 학습이 시작될 때 모델 가중치가 무작위로 초기화되는 표준 머신 러닝 관행과 대조적입니다. S4의 경우, Dao와 Gu는 Legendre 다항식에서 파생된 초기화 체계를 제안했습니다. Dao와 Gu는 "How to Train Your HiPPO."이라는 제목의 후속 논문에서 더 많은 공식을 탐구했습니다.¹

S4 논문에서는 “SSM을 랜덤 행렬 A에서 [HiPPO 행렬]로 간단히 수정하면 순차 MNIST 벤치마크에 대한 성능이 60% 에서 98%로 향상"되었으며, 이를 통해 SSM의 장기 메모리 문제를 효과적으로 해결할 수 있었다고 밝혔습니다. DSS, S5 및 Mamba와 같은 구조화된 SSM의 최신 변형은 A와 B에 대해 서로 다른(많은 경우 더 간단한) 초기화 체계를 사용하지만, 핵심 HiPPO 원칙은 동일하게 유지합니다. 즉, 행렬의 각 값 간에 안정적인 업데이트와 어느 정도의 독립성을 요구하는 대각선 구조를 구현합니다.

Mamba 아키텍처의 핵심은 두 가지 혁신입니다. 첫 번째는 선택적 상태 공간 모델입니다 이 모델은 이전에 트랜스포머 모델에게만 있었던 중요한 능력, 즉 현재의 관련성을 기반으로 과거 입력 기록의 특정 부분에 선택적으로 초점을 맞추거나 이를 무시할 수 있는 능력을 Mamba에 제공합니다. 두 번째는 그래픽 처리 장치(GPU)가 메모리 계층 구조에서 모델의 계산을 처리하는 방식을 최적화하여 속도와 계산 효율성을 극대화하는 알고리즘인 하드웨어 인식 병렬 스캔입니다.

트랜스포머에서 이러한 기능은 현재 입력 토큰과의 관련성을 기반으로 각 이전 토큰의 영향을 강조하거나 덜 강조하는 어텐션 가중치를 조정하는 어텐션 메커니즘에 의해 제공됩니다. 일반 SSM은 전체 입력 기록을 사용하여 입력을 아웃풋에 매핑하도록 명시적으로 설계되어 있습니다.이는 일부 시퀀스 모델링 작업에서는 허용되거나 심지어 바람직하지만, 대부분의 고급 언어 모델링 작업에서는 상당한 핸디캡으로 작용합니다.

입력 기록의 특정 부분을 동적으로 생략하거나 강조할 수 없는 이러한 단점을 해결하기 위해 Dao와 Gu는 "선택적 스캔"을 사용하는 새로운 종류의 상태 공간 모델을 제안했습니다. Mamba 논문에서 저자는 "선택적 SSM은 선택 메커니즘이 있고 스캔으로 계산된 S4 모델이기 때문에 가끔 S6 모델로 축약한다"고 언급했습니다. 이 S6 기반 아키텍처에 'Mamba'라는 별명을 붙인 이유는 무엇보다도 'S'가 모두 뱀의 쉿 소리처럼 들리기 때문입니다.

Mamba는 선택적 상태 공간 모델을 핵심으로 하는 신경망 아키텍처로 이해하는 것이 가장 좋습니다. 간단히 비유하자면, Mamba는와 선택적 SSM의 관계는 트랜스포머 모델과 어텐션 메커니즘의 관계와도 같습니다.

기존 SSM에는 고정된 역학이 있습니다. 즉, 숨겨진 상태가 한 단계에서 다음 단계로 어떻게 진화하는지를 정하는 규칙인 모델 매개변수는 시퀀스의 모든 입력과 모든 단계에서 동일합니다. 이 속성을 선형 시간 불변(LTI)이라고 합니다. SSM에 현재의 컨텍스트를 기반으로 특정 과거 정보의 우선 순위를 선택적으로 지정하거나 우선 순위를 낮출 수 있는 기능을 제공하기 위해 Dao와 Gu는 다른 입력마다 주요 모델 매개변수의 값이 달라지도록 SSM을 재구성했습니다.

보다 구체적으로, 선택적 SSM은 단계 크기 ∆_t와 행렬 B_t 및 C_t를 현재 입력 토큰 x_t의 직접 함수로 만듭니다. 이는 먼저 x_t의 벡터 임베딩을 3개의 평행 선형 투영 레이어, 즉 표준 피드포워드 신경망 레이어(또는 MLP 레이어)를 통해 전달함으로써 달성할 수 있습니다. 이는 병렬 쿼리, 키 및 값 헤드가 트랜스포머 모델에서 입력의 각 Q, K 및 V 벡터를 생성하는 방식과 동일합니다.

선형 투영 네트워크에서 x_t 의 벡터 임베딩에 가중치 및 편향 항을 곱하면 ∆_t , B_t 및 C_t의 결과 값이 산출됩니다. 선형 투영 계층 자체의 가중치 및 편향 항은 대규모 텍스트 샘플 데이터 세트를 이용한 모델 사전 훈련 중에 학습된 다음 (선택적으로) 후속 미세 조정을 통해 구체화됩니다.

• ∆_t의 값은 x_t가 모델이 지금까지 살펴본 컨텍스트에 대한 모델의 기억에 미치는 영향의 크기를 결정합니다. 즉, 숨겨진 상태 h_t-1에서 h_t까지 얼마나 많은 업데이트가 있을 것인지에 대한 것입니다.단계 크기 ∆_t기 클수록 변화가 커지며, 상태 내에 포함된 오래된 정보가 쇠퇴하는 속도, 즉 “잊어버리는” 속도가 빨라집니다. 반대로 단계 크기가 작을수록 업데이트 크기도 작아집니다. 단계 크기가 충분히 작으면 현재 입력은 숨겨진 상태에 전혀 영향을 주지 않습니다.
  
  
• 행렬 B_k에 대한 변경은 현재 입력 토큰이 숨겨진 상태를 업데이트 하는 방식을 결정합니다. 예를 들어, x_t 가 "어제"라는 단어의 토큰이라면 B_t는 이후의 컨텍스트가 아마도 과거와 관련이 있다는 것을 반영하여 상태를 업데이트하는 방식으로 조정될 수 있습니다.
  
  
• 행렬 C_t의 변경 사항은 이 컨텍스트 정보가 모델의 아웃풋 y_t에 미치는 영향을 결정합니다. x_k가 “어제”의 토큰이라는 예제를 계속 살펴보면, C_k는 이후 모델에 의해 출력되는 모든 동사가 과거형으로 활용되도록 하는 방식으로 영향을 받을 수 있습니다.

특히, A 행렬에는 이러한 입력 기반 조정이 이루어지지 않습니다. 그 역할은 S4 모델과 동일하게 과거 입력의 전체 기록을 효율적으로 기억하는 것입니다. 특정 순간에 해당 기록의 어느 부분을 활용할지 결정하는 역할은 B 및 C 행렬에 의해 수행됩니다.

그러나 모델이 더 이상 시간 불변이 아니면 전이 커널이 더 이상 일정하지 않기 때문에 훈련 중에 더 이상 컨볼루션 지름길을 사용할 수 없습니다: 선택성 메커니즘의 핵심은 h_t-1에서 h_t로의 전환이 이제 컨텍스트에 종속된다는 것입니다.

대신 Mamba는 영리한 해결 방법을 사용하여 유사한 병렬화 이점을 달성합니다. SSM은 곱셈과 덧셈만 사용하기 때문에 계산은 수학의 친숙한 연관 속성의 적용을 받습니다. 이를 통해 많은 순차적 계산을 병렬 접두사 합 스캔을 통해 GPU에서 병렬로 처리할 수 있는 작고 독립적인 청크로 나눌 수 있습니다.

또한, 결과는 특정한 계층적 방식으로 결합되어 GPU에서 다양한 종류의 하드웨어 메모리를 최적의 효율로 사용할 수 있으며, 현재 최신 LLM에서 보편화되어 있는 Tri Dao가 개발한 FlashAttention 기술과 유사한 원리를 사용합니다.  

Mamba 아키텍처 내에서 S6 모델은 어텐션 메커니즘이 더 큰 "어텐션 블록" 내에서 모듈 역할을 하는 방식과 유사하게 더 큰 "Mamba 블록"의 모듈 역할을 합니다. 이는 S6 모듈과 게이트형 신경망 아키텍처를 결합합니다. Mamba 모델은 일반적으로 모델의 최종 아웃풋 예측을 수행하는 아웃풋 레이어 이전에 여러 Mamba 블록, 즉 신경망의 일련의 연속적 Mamba 레이어로 구성됩니다.

입력의 복사본은 Mamba 블록에 들어가기 전에 잔여 연결로서 엔드로 직접 전송됩니다.Mamba 블록의 내부 작동 목적은 더 큰 컨텍스트의 어느 부분이 해당 입력과 관련이 있는지 결정할 뿐만 아니라, 해당 컨텍스트 정보가 입력의 원래 의미를 얼마나 수정해야 하는지 결정하는 것입니다.

Mamba 블록 내에서 원본 입력 벡터는 다음과 같이 처리됩니다:

• 먼저 입력값은 입력 벡터 자체의 두 배 너비의 선형 레이어를 통과하여 더 높은 차원의 공간에 투영됩니다. 예를 들어, 모델이 원래 각 입력 토큰 x를 512차원 벡터 임베딩으로 표현하는 경우,x에 선형 투영 레이어의 가중치를 곱하면 1024차원 벡터로 확장됩니다.
  
  
• 그런 다음 확장된 벡터가 둘로 나뉩니다. 절반(x_proj 라고 부르겠습니다)은 SSM을 통과하는 경로로 공급되고, 나머지 절반( z_proj 라고 부르겠습니다)은 게이팅 메커니즘을 통과하는 별도의 경로로 공급됩니다. 명확성을 위해 이전 확장 단계는 일반적으로 2개의 병렬 선형 레이어에 의해 수행되는 것으로 표시됩니다.
  
  
• x_proj는 SSM에 도달하기 전에 1차원 컨볼루션 레이어에 공급됩니다. 이 컨볼루션 계층은 로컬 패턴(예: 간단한 동사-주어 쌍과 같은 인접한 토큰 간 종속성)을 추출합니다. 이를 통해 SSM은 장거리 글로벌 종속성에 대한 컨텍스트 이해에 "집중"할 수 있습니다.
  
  
• 컨볼루션 레이어의 아웃풋은 비선형 활성화 함수에 대한 입력 역할을 합니다. 비선형 도입은 모든 신경망의 특징이며, 이를 통해 더 복잡한 패턴을 캡처할 수 있습니다. Mamba 논문은 SiLU(Sigmoid Linear Unit)를 사용합니다. 결과 벡터를 x_act라고 하겠습니다.
  
  
• 한편, 별도의 게이팅 메커니즘 경로에서 z_proj도 비선형 활성화 함수에 입력되어 z_act를 도출합니다.
  
  
• SSM 경로에서x_act ∆_x, B_x , C_x 각각에 대한 값을 생성하는 세 개의 병렬 선형 투영 계층에 공급됩니다.
  
  
• SSM은 이러한 입력 종속 매개변수(그리고 A 및 D 행렬)를 사용하여 상태 공간 업데이트 및 SSM의 아웃풋 y를 계산합니다.
  
  
• SSM의 아웃풋 벡터 y는 이제 게이팅 경로의 아웃풋 벡터 z_act에 의해 요소별로 곱해집니다. 기본적으로 z_act의 각 요소는 오디오 믹싱 콘솔의 볼륨 노브처럼 작동합니다. z_act의 주어진 요소가 0에 가까운 경우, y의 해당 부분과 곱하면 0에 가까운 값이 산출되고 영향이 감소합니다. 반대로, zact의 주어진 요소가 큰 경우, y와 곱하면 해당 컨텍스트 정보의 영향력이 증폭됩니다.
  
  
• 결과 벡터는 원래 크기로 다시 투영됩니다. 이는 원래 입력 벡터의 각 요소에 대해 가중치가 부여된 컨텍스트 업데이트(또는 비업데이트)의 벡터로 이해할 수 있습니다.
  
  
• 마지막으로, 해당 업데이트 벡터는 잔여 연결로 블록 엔드로 바로 전송된 원래 입력 벡터의 복사본에 추가됩니다.
  
  
• 이제 선택적 SSM에서 제공하는 컨텍스트 이해를 반영하도록 원래 입력 벡터가 업데이트되었습니다. 이제 이를 다음 Mamba 레이어로 보내거나 모델의 최종 레이어에서 완전히 업데이트된 벡터가 모델 어휘의 각 단어에 해당하는 각 확률을 출력하는 softmax 함수에 대한 입력 역할을 할 수 있습니다.

Mamba 논문의 원본이 발행된 지 1년 후, Dao와 Gu는 “트랜스포머는 SSM이다: 구조화된 상태 공간 이중성을 통한 일반화된 모델과 효율적인 알고리즘”이라는 논문을 발표했습니다. 이 후속 백서의 주된 기여는 다음 세 가지입니다.

• Mamba와 트랜스포머 사이의 이론적 연결과 두 아키텍처 간 공통의 어휘 탐구
• Mamba 모델의 다양한 설계 선택지에 대한 일련의 설명 및 탐색
• 수정된 아키텍처인 Mamba-2는 이러한 설계 탐색을 통해 정보를 얻고 개선되었습니다.

Mamba-2 알고리즘은 원래 Mamba보다 훨씬 빠르고 구현하기 쉽습니다. 저자들은 약 25줄의 코드로 선택적 SSM을 구현하는 "최소 SSD" 코드 베이스를 제공했습니다.² 이러한 효율성을 통해 Mamba-2는 모델 속도를 늦추지 않고 훨씬 더 큰 숨겨진 상태 차원을 사용할 수 있으므로, 아키텍처로 구축된 더 크고 강력하며 표현력이 풍부한 모델을 만들 수 있습니다. 테스트에서 Mamba-2 모델은 일련의 다운스트림 작업에서 해당 크기의 Mamba 및 트랜스포머 모델과 확실히 일치하거나 이들을 능가했습니다.

논문의 서론에서 언급했듯이, Dao와 Gu의 "주요 목표는 구조화된 SSM과 어텐션 변형 사이의 풍부한 이론적 연결을 개발하는 것"이었습니다. 이는 이 두 가지를 통합하는 새로운 개념적 틀을 탄생시켰으며, Dao와 Gu는 이를 "상태 공간 이중성"(SSD)이라고 불렀습니다.³ 이를 통해 그들은 Mamba가 트랜스포머 아키텍처에 대한 수년간의 탐색 및 최적화로부터 이익을 얻을 수 있는 문을 열었습니다.

주목할만한 이점 중 하나는 MHA(멀티 헤드 어텐션)에 해당하는 Mamba의 개발이었습니다. Mamba 블록을 트랜스포머의 다중 “어텐션 헤드”와 유사하게 여러 개의 “Mamba 헤드”로 분할할 수 있습니다. Dao와 Gu가 그룹화된 쿼리 어텐션과 유사하다고 간주한 이 접근 방식의 한 가지 변형은 GPU의 텐서 병렬 처리를 통해 훨씬 더 높은 효율성을 가능하게 합니다.

이들이 원래의 "순차적" Mamba 블록과 대조적으로 병렬 Mamba 블록이라고 명명한 Mamba-2 블록에서는 입력 종속 매개변수 ∆, B 및 C가 초기 투영 계층에서 병렬로 생성됩니다. 특히 B와 C는 x_proj를 전용 선형 레이어로 곱하는 것이 아니라 x_proj 의 일부를 복사하여 도출됩니다. 이러한 병렬 처리는 전체 모델 매개변수를 단순화하고 줄이는 것 외에도 훨씬 더 효율적인 대규모 학습을 가능하게 합니다.⁴

Mamba와 트랜스포머 모두 각자의 강점이 있지만, Mamba 기반 모델은 일반적으로 메모리 사용량 및 속도와 관련된 모든 문제에서 우수합니다. Mamba 논문에 따르면, Mamba는 동급 트랜스포머보다 5배 더 큰 처리량을 제공합니다.

트랜스포머는 믿을 수 없을 정도로 정확하고 다재다능하지만, 어마어마한 계산 리소스를 필요로 합니다. 사전 학습(및 미세 조정) 중에 셀프 어텐션의 메모리 요구 사항은 시퀀스 길이에 따라 2차식으로 확장됩니다. 즉, 시퀀스의 컨텍스트 길이를 두 배로 늘리면 어텐션 메커니즘은 4배의 리소스를 사용합니다. 이 "2차 병목 현상"으로 인해 컨텍스트 창이 증가함에 따라 속도와 메모리 가용성이 저하됩니다. 추론 중 메모리 수요는 선형적으로 확대됩니다.

훈련 중 Mamba 모델의 메모리 사용량은 훈련 중에만 선형적으로 확대됩니다. 더 중요한 점은 추론 중 메모리 사용량이 일정하다는 것입니다. 즉, 모델이 본 토큰 수와 관계없이 SSM은 입력 기록에 대한 고정된 크기의 표현을 유지합니다. 이를 통해 이론적으로 무제한의 컨텍스트 길이가 허용되며, 이 길이는 하드웨어 제한에 의해서만 제한됩니다.

즉, 트랜스포머의 메모리 집약적이고 계산적으로 중복되는 방식에는 나름의 장점이 있습니다. 예를 들어, 연구 결과에 따르면 컨텍스트 내 학습(예: 퓨샷 프롬프트), 복사 또는 긴 맥락 추론)이 필요한 작업에서는 여전히 트랜스포머가 Mamba와 Mamba-2를 모두 앞선다고 합니다.

다행히도 트랜스포머와 Mamba의 각각의 강점은 상호 배타적이지 않습니다. Mamba-2 논문은 하이브리드 모델이 순수 트랜스포머나 SSM 모두를 능가할 수 있다고 제안합니다. 이는 2024년 후반에 NVIDIA 연구를 통해 공식적으로 검증된 개념입니다.⁵ 대체로 하이브리드 모델은 효율성이 높다는 Mamba의 이점과 트랜스포머의 보다 리소스 집약적인 어텐션 메커니즘이 제공하는 뉘앙스 및 컨텍스트 내 학습 성능을 결합한 것으로 보입니다.

이에 대해 더 자세히 알아보기 위해 IBM Research는 Dao, Gu, 그리고 일리노이 대학교 어바나-샴페인 캠퍼스(UIUC)의 Minjia Zhang과 함께 Bamba 및 Bamba V2에서 협업을 진행했습니다. Bamba는 IBM Granite 4.0의 많은 아키텍처 요소에 정보를 제공했습니다.

하이브리드 모델에 대한 연구는 특히 오픈 소스 커뮤니티 내에서 활발한 연구 영역으로 남아 있습니다.