[Paper Review] Mamba: Linear-Time Sequence Modeling with Selective State Spaces

원본 게시글: https://velog.io/@euisuk-chung/Paper-Review-Mamba-Linear-Time-Sequence-Modeling-with-Selective-State-Spaces

최근 딥러닝 아키텍처의 중심에는 트랜스포머가 자리 잡고 있습니다. 대규모 언어 모델(LLM)뿐만 아니라, 그림을 생성하는 데 쓰이는 디퓨전 모델 또한 트랜스포머 구조를 활용하고 있습니다. 이외에도 시계열 분석이나 추천 시스템과 같은 다양한 분야에서 트랜스포머가 핵심적인 역할을 하고 있습니다.

그러나 트랜스포머를 대체할 수 있는 새로운 아키텍처를 모색하려는 연구는 계속되고 있으며, 그 중에서도 특히 주목받고 있는 것이 State Space Model(SSM)입니다. 최근 “Mamba: Linear-Time Sequence Modeling with Selective State Spaces“라는 논문과 그 모델이 공개되면서, SSM이 트랜스포머의 대안으로서 더욱 관심을 끌고 있습니다.

• Mamba 논문 링크 : https://arxiv.org/pdf/2312.00752

추가적으로 벌써 Survey 논문도 벌써 나왔는데요!! 흥미로운 이미지들만 좀 reference용으로 가져오자면 아래와 같습니다. Mamba 이후로 많은 후속 연구 및 Variation들이 빠르게 연구되고 있는 것들을 확인할 수 있습니다.

• 많은 관심을 받고 있는 SSM 모델

  

  이미지 출처. State Space Model for New-Generation Network Alternative to Transformers: A Survey

• 수많은 SSM Variation 모델 : new paradigm shift?!

  

  이미지 출처. State Space Model for New-Generation Network Alternative to Transformers: A Survey

1. Introduction (서론)

서론에서는 Mamba 모델의 필요성과 기존 Transformer 모델의 한계점을 설명하고, Mamba가 이 문제를 어떻게 해결하는지를 소개합니다.

• 기존 Transformer의 문제점: Transformer는 Attention 메커니즘을 기반으로 긴 시퀀스를 처리하는 데 매우 뛰어나지만, 시퀀스 길이에 따라 계산 복잡도가 2차 함수(Quadratic)로 증가하여, 긴 시퀀스 데이터를 다룰 때 매우 비효율적입니다.

  • 이러한 문제를 해결하기 위해 많은 연구자들이 Linear Attention과 같은 다양한 방법을 제안했지만, 대부분은 정보 밀도가 높은 데이터(예: 텍스트)에서 Transformer만큼의 성능을 내지 못했습니다.

• SSM(Structured State Space Models)의 등장: 이러한 한계를 해결하기 위해, Structured State Space Models(SSM)이 등장했습니다.

  • SSM은 재귀적 신경망(RNN)과 합성곱 신경망(CNN)의 이점을 결합한 모델로, 시퀀스의 길이에 비례하는 선형적인 계산 복잡도를 가지고 있어 매우 효율적입니다.
  • 그러나 SSM은 정보 밀도가 높은 텍스트 데이터에서는 Transformer만큼의 성능을 내지 못했습니다.

• Mamba 모델의 등장: Mamba는 선택 메커니즘을 도입한 Selective State Space Model(선택적 상태 공간 모델)을 기반으로 하며, 긴 시퀀스를 다루면서도 Transformer 수준의 성능을 유지하면서도 계산 비용을 줄일 수 있는 모델입니다.

  • 특히, 텍스트, 오디오, 유전체학(genomics) 등의 다양한 데이터 유형에서 매우 우수한 성능을 보여줍니다.

2. State Space Models (상태 공간 모델)

이 장에서는 상태 공간 모델(SSM)의 기본 개념과 작동 원리에 대해 설명합니다.

• SSM의 작동 원리: 상태 공간 모델은 시스템의 입력을 고차원의 잠재 공간(latent space)으로 변환하여 처리하는 방식으로 동작합니다. 이는 수학적으로 연속 시스템의 이산화(Discretization) 과정으로 표현됩니다.

  • 수식으로 표현하자면, 입력 x(t)x(t)x(t)을 잠재 상태 h(t)h(t)h(t)로 변환하고, 이를 통해 출력 y(t)y(t)y(t)을 도출하는 방식입니다.
  • 이때, 각 시점에서 상태 공간의 변화를 나타내는 주요 매개변수 A,B,CA, B, CA,B,C가 주어집니다.h(t)=A⋅h(t−1)+B⋅x(t)h(t) = A \cdot h(t-1) + B \cdot x(t)h(t)=A⋅h(t−1)+B⋅x(t) y(t)=C⋅h(t)y(t) = C \cdot h(t)y(t)=C⋅h(t)

• 연속 시스템에서 이산 시스템으로의 변환: SSM에서는 연속적인 시스템의 매개변수를 이산화(discretization)하여 계산합니다. 이를 통해 모델은 연속 데이터를 효율적으로 처리할 수 있습니다.

  • 이 때, 선형 시간 불변 시스템(LTI, Linear Time-Invariant System)의 개념이 사용됩니다. 이 시스템은 시간에 따라 변하지 않는 선형적인 연산을 수행하므로, 긴 시퀀스 처리에 매우 효율적입니다.

💡 결국 이산화하면 RNN이랑 같은 거 아닌가?

• DSBA 연구실의 천재원 석사생의 PYSR을 보면 이에 대한 답변을 얻을 수 있습니다.

  

✔️ SSM (State Space Model)의 연속성

• SSM은 연속적인 시간 흐름에 따라 시스템의 상태를 모델링하는 방식입니다. 이때 A와 B는 연속적 시스템을 표현하는 중요한 매트릭스들로, 시간에 따른 시스템의 상태 변화를 기술합니다.
  • A: 상태 변화를 결정하는 매트릭스. 이전 상태 Xt−1X_{t-1}Xt−1​에 곱해져서 시스템의 상태가 어떻게 변하는지 정의합니다.
  • B: 입력을 상태로 변환하는 매트릭스. 입력 UtU_tUt​를 받아 상태에 반영하는 역할을 합니다.
  • A̅와 B̅는 SSM에서 이산화 된 버전의 매트릭스들로, 연속적인 시스템을 이산적인 형태로 변환하여 시퀀스 데이터를 처리할 수 있게 만듭니다. A̅와 B̅는 연속적인 SSM 모델의 도함수를 기반으로 이산적 시퀀스 처리에 맞게 변환된 것입니다.
• SSM의 장점은 이러한 연속적인 흐름을 기반으로 시스템의 미세한 변화를 더 잘 모델링할 수 있다는 점입니다.
  • 시간 변화가 연속적인 시스템에서 데이터를 잘 반영할 수 있기 때문에 시스템의 물리적 성질을 더 정확하게 반영할 수 있습니다.

✔️ RNN (Recurrent Neural Network)의 이산화

• RNN은 기본적으로 이산화된 시퀀스 데이터를 처리하는 데 중점을 둡니다. RNN은 각 시간 스텝에서 이전 상태와 현재 입력을 바탕으로 다음 상태를 업데이트합니다.
  • RNN은 연속적인 시간 흐름을 명시적으로 모델링하지 않으며, 이전 상태와 현재 상태 간의 단순한 관계에 의존합니다.
• RNN의 한계는 시간의 연속성을 명확하게 다루지 않기 때문에, 시간에 따른 미세한 변화를 반영하는 데에는 한계가 있을 수 있다는 점입니다.
  • 연속적인 시간 흐름을 반영하지 않는 구조이므로, 물리적 시간 흐름이 중요한 문제에 대해서는 성능이 제한적일 수 있습니다.

• SSM 구조: SSM의 구조는 주로 입력 데이터를 잠재 상태로 변환한 후 이를 기반으로 출력을 도출하는 방식으로, 각 채널이 독립적으로 작동하는 특징이 있습니다. 이로 인해 계산의 병렬화가 가능해져 매우 효율적으로 작동할 수 있습니다.

• SSM 아키텍처 개요 : SSM(상태공간모델) 아키텍처는 독립적인 시퀀스 변환 모델로, 엔드 투 엔드 신경망 아키텍처에 통합될 수 있습니다.

  • SSM 아키텍처는 SSNN(State Space Neural Networks)라고도 하며, 이 경우에는 SSM 레이어가 CNN(합성곱신경망) 레이어와 유사한 역할을 합니다.
  • Introduction에는 아래 잘 알려진 몇 가지 SSM 아키텍처를 간단하게 소개합니다:
    • Linear Attention (Katharopoulos et al. 2020): 자가 주의의 근사로, 재귀성을 포함하므로 일종의 선형 SSM으로 볼 수 있습니다.
    • H3 (Dao, Fu, Saab et al. 2023): 이 모델은 S4를 사용하는 재귀를 일반화하며, SSM이 두 개의 게이트가 있는 연결 사이에 위치하는 형태입니다. H3는 표준 지역 합성을 추가하여 이를 shift-SSM으로 간주합니다.
    • Hyena (Poli et al. 2023): H3와 동일한 아키텍처를 사용하지만, S4 레이어를 MLP 매개변수화된 전역 합성으로 대체합니다.
    • RetNet (Y. Sun et al. 2023): 이 아키텍처는 추가적인 게이트를 더하고, 더 단순한 SSM을 사용하여 다중 헤드 주의(MHA) 변형을 통해 대안적인 병렬 계산 경로를 제공합니다.
    • RWKV (B. Peng et al. 2023): 이 모델은 언어 모델링을 위해 설계된 RNN으로, 다른 선형 주의 근사의 일종인 Attention-free Transformer(S. Zhai et al. 2021)를 기반으로 합니다. 주요 “WKV” 메커니즘은 LTI 재귀를 포함하며, 두 개의 SSM의 비율로 볼 수 있습니다.

(참고) LSSL 및 deepSSM 차원 계산 방식

• 슬라이드에 나온 수식과 제가 쓴 수식이 상이합니다. 저는 앞에 쓴 수식에 맞도록 hidden dim을 h로, time input을 x로 정의했습니다.

1. LTI(Linear Time-Invariant) 시스템의 정의

• LTI 시스템은 시간에 따라 시스템의 특성이 변하지 않으며, 선형적이고 시간 불변적입니다. 즉, 같은 입력이 주어지면 언제든지 동일한 방식으로 처리되고, 시간의 흐름에 따라 시스템의 행동이 달라지지 않는 특성을 가집니다.
  • 이는 두 가지 주요 특성에 의해 정의됩니다:
    • 선형성: 입력과 출력의 관계가 선형입니다. 즉, 입력의 합이 출력의 합으로 선형적으로 변환됩니다.
    • 시간 불변성: 시스템의 상태 변화는 시간에 의존하지 않습니다. 즉, 시간에 따라 시스템의 성능이나 동작이 달라지지 않습니다.

❓ 그렇다면 자연스럽게 드는 의문점은?

• 🤔 : 흠.. 얘를 dmodeld_{\text{model}}dmodel​ 차원으로 어떻게 확장을 시킬까? 하는 생각을 하게 됩니다.
  • 단순하게 생각해보면, 아래 그림의 SSM 모델처럼 입력과 출력이 dmodeld_{\text{model}}dmodel​ 차원으로 확장하는 식으로 생각해볼 수도 있긴 합니다만!⚠️⚠️

  • 확장된 방정식:

    h(t)=A⋅h(t−1)+B⋅x(t)h(t) = A \cdot h(t-1) + B \cdot x(t)h(t)=A⋅h(t−1)+B⋅x(t)

    y(t)=C⋅h(t)y(t) = C \cdot h(t)y(t)=C⋅h(t)

  • 각 요소의 디멘션 변화:

    • h(t)∈Rnh(t) \in \mathbb{R}^nh(t)∈Rn: 히든 스테이트는 여전히 nnn차원을 가집니다.

    • x(t)∈Rdmodelx(t) \in \mathbb{R}^{d_{\text{model}}}x(t)∈Rdmodel​: 입력이 이제 dmodeld_{\text{model}}dmodel​ 차원의 벡터로 확장됩니다.

    • y(t)∈Rdmodely(t) \in \mathbb{R}^{d_{\text{model}}}y(t)∈Rdmodel​: 출력 역시 dmodeld_{\text{model}}dmodel​ 차원의 벡터로 확장됩니다.

    • A∈Rn×nA \in \mathbb{R}^{n \times n}A∈Rn×n: 히든 스테이트의 업데이트를 담당하며 차원은 변하지 않습니다.

    • B∈Rn×dmodelB \in \mathbb{R}^{n \times d_{\text{model}}}B∈Rn×dmodel​: 입력을 히든 스테이트로 변환하는 역할을 하며, dmodeld_{\text{model}}dmodel​ 차원을 처리합니다.

    • C∈Rdmodel×nC \in \mathbb{R}^{d_{\text{model}} \times n}C∈Rdmodel​×n: 히든 스테이트를 출력으로 변환하는 매트릭스.

    

💡 하지만!! 위와 같은 방법은 성립하지 않습니다.

• 결론부터 말씀드리자면, State Space Model(SSM)은 본질적으로 LTI(Linear Time-Invariant) 시스템이며, 이 특성 때문에 입력과 출력이 dmodeld_{\text{model}}dmodel​ 차원으로 확장되었을 때, 차원별로 독립적인 SSM 처리가 필요하게 됩니다.

❓ 왜 그럼 성립하지 않는가

=> SSM은 본질적으로 LTI 시스템

• SSM 자체가 LTI 시스템이기 때문에, 기본적으로 선형성과 시간 불변성이 보장되어야 합니다. 이를 기반으로, SSM은 다음의 두 가지 중요한 특성을 가집니다:

  • 선형성: 시스템의 입력과 출력 사이의 관계는 선형입니다. 즉, 입력이 변하면 출력도 선형적으로 변하며, 이는 시스템의 동작을 결정하는 선형 매트릭스에 의해 제어됩니다.
  • 시간 불변성: 시스템은 시간에 따라 변화하지 않고 항상 동일한 방식으로 동작합니다. 입력이 언제 들어오든, 시스템의 상태와 출력은 동일한 방식으로 계산됩니다.

• 따라서 SSM의 모든 연산은 시간에 따라 변하지 않아야 하며, 입력 차원과 상관없이 동일한 방식으로 동작해야 합니다.

입력과 출력이 dmodeld_{\text{model}}dmodel​ 차원으로 확장될 때의 문제

• 만약 입력과 출력이 dmodeld_{\text{model}}dmodel​ 차원으로 확장된다면, 모든 차원을 동시에 처리하기 위해 하나의 공통된 SSM 매트릭스(AAA, BBB, CCC)를 사용하는 방식으로는 문제가 생길 수 있습니다. 왜냐하면:

  1. SSM은 본질적으로 LTI 시스템이기 때문에, 각 차원은 서로 독립적으로 처리되어야만 선형성과 시간 불변성이 보장됩니다.
  2. 여러 차원을 하나의 SSM 시스템에서 처리하는 방식은, 차원 간 상호작용을 야기할 수 있습니다. 이는 각 차원이 개별적으로 독립적으로 처리되지 않기 때문에, 시간에 따른 입력 처리 방식이 달라질 가능성이 생깁니다.

• 따라서 dmodeld_{\text{model}}dmodel​ 차원의 입력을 하나의 SSM으로 처리하는 방식은 LTI 시스템의 요구 사항을 위배할 수 있습니다. 차원별로 독립적인 처리가 이루어지지 않으면 차원 간의 상호작용이 발생하고, 시간에 따라 결과가 달라질 수 있기 때문입니다.

💡 그렇다면 어떻게 되는가?

=> 각 차원의 SSM을 독립적으로 처리

• 설명:
  • dmodeld_{\text{model}}dmodel​ 차원의 입력 벡터가 주어질 때, 각 차원은 개별적인 SSM을 거칩니다.
  • 즉, dmodeld_{\text{model}}dmodel​개의 SSM이 독립적으로 존재하며, 각각 n×nn \times nn×n 크기의 AAA 매트릭스를 사용하여 히든 스테이트를 업데이트합니다.
  • 이 경우 입력과 출력의 차원을 맞추기 위해 BBB와 CCC의 차원도 dmodeld_{\text{model}}dmodel​에 따라 확장됩니다.
• 최종 디멘션:
  • A∈Rn×n×dmodelA \in \mathbb{R}^{n \times n \times d_{\text{model}}}A∈Rn×n×dmodel​
  • B∈Rn×dmodelB \in \mathbb{R}^{n \times d_{\text{model}}}B∈Rn×dmodel​
  • C∈Rdmodel×nC \in \mathbb{R}^{d_{\text{model}} \times n}C∈Rdmodel​×n

3. Selective State Space Models (선택적 상태 공간 모델)

이 장에서는 선택적 상태 공간 모델(SSSM, Selective State Space Models)에 대해 설명하며, 이를 통해 기존 SSM의 성능을 어떻게 개선할 수 있는지 다룹니다.

📊 Figure 1 설명

1. xtx_txt​ (입력 데이터)

• 입력 데이터 xtx_txt​ (초록)는 시퀀스의 현재 시점에서 들어오는 데이터입니다. 이 데이터는 여러 채널(D)로 나누어져 있고, 각각의 채널이 독립적으로 처리됩니다.
  • 예를 들어, 그림에서는 D=5D = 5D=5로, 5개의 입력 채널을 의미합니다.

1. ht−1h_{t-1}ht−1​ (이전 시점의 상태)

• ht−1h_{t-1}ht−1​는 이전 시점에서 계산된 잠재 상태(latent space)를 의미합니다. 이 잠재 상태는 시간에 따라 이어져 있으며, 이전 시점의 정보가 현재 시점에 영향을 미칩니다.
  • 예를 들어 그림에서는 N=4N = 4N=4의 고차원 공간에서 정의되어 있습니다.

1. Bt,Ct,A,ΔtB_t, C_t, A, \Delta_tBt​,Ct​,A,Δt​ (SSM의 주요 매개변수)

• 이 네 가지 매개변수는 SSM에서 중요한 역할을 합니다. 각 매개변수는 시퀀스 데이터를 처리하고, 잠재 상태 hth_tht​를 업데이트하며, 최종 출력 yty_tyt​를 계산하는 데 사용됩니다.
  • BtB_tBt​: 입력 데이터 xtx_txt​와 상호작용하여 새로운 상태를 만듭니다.
    • (참고 수식) ▶ h(t)=A⋅h(t−1)+B⋅x(t)h(t) = A \cdot h(t-1) + B \cdot x(t)h(t)=A⋅h(t−1)+B⋅x(t)
    • 이때 BtB_tBt​는 입력 의존적이며, 입력에 따라 동적으로 변화합니다.
  • AAA: 잠재 상태 hth_tht​를 업데이트하는 데 필요한 중요한 매개변수입니다. 이전 상태와 현재 입력에 기반하여 새로운 상태를 계산할 때 사용됩니다.
  • CtC_tCt​: 계산된 잠재 상태 hth_tht​를 출력 yty_tyt​로 변환하는 데 사용됩니다.
    • (참고 수식) ▶ y(t)=C⋅h(t)y(t) = C \cdot h(t)y(t)=C⋅h(t)
    • 이를 통해 최종적으로 시퀀스의 각 시점에서 모델의 출력을 얻습니다.

  • Δt\Delta_tΔt​: 이 매개변수는 시간 차원을 조절하는 역할을 합니다.

    • 시간에 따라 상태 공간에서의 변화를 조정하여 모델이 시퀀스를 따라 중요한 정보를 기억하거나 잊도록 도와줍니다.

1. 선택 메커니즘 (Selection Mechanism)

• 선택 메커니즘은 이 모델의 핵심적인 요소로, SSSM에서 입력 데이터를 선택적으로 처리하게 만듭니다.
  • 이 메커니즘은 시퀀스 데이터를 필터링하여 중요한 정보만을 선택하고, 불필요한 정보는 무시할 수 있도록 돕습니다. ✨ 즉, 정보 압축 및 선택적 기억을 수행하는 방식입니다. ✨
  • 그림에서는 이 메커니즘이 입력 xtx_txt​와 상호작용하여 BtB_tBt​와 Δt\Delta_tΔt​를 조절하는 방식으로 표현됩니다. (파랑)

1. GPU 메모리 계층 (GPU Memory Hierarchy)

• SSSM의 중요한 특징 중 하나는 하드웨어 친화적인 알고리즘을 사용하는 것입니다.
  • 이 모델은 GPU의 고속 메모리(SRAM)와 대용량 메모리(HBM) 계층을 효과적으로 활용하여, 잠재 상태의 계산을 최적화합니다.
  • 이 구조를 통해 SSSM은 큰 시퀀스 데이터도 효율적으로 처리할 수 있으며, 계산 자원을 효율적으로 활용하여 더 빠르게 결과를 도출할 수 있습니다.

3.1 Motivation: Selection as a Means of Compression

이 섹션은 선택 메커니즘(selection mechanism)이 왜 중요한지를 설명합니다. 주요 내용은 데이터 압축 과 관련된 문제이며, 선택적 메커니즘이 이를 어떻게 해결하는지를 다룹니다.

• 선택의 필요성: 시퀀스 모델링의 주요 과제 중 하나는 컨텍스트(문맥) 정보를 효과적으로 압축하는 것입니다.
  • Transformer는 정보의 압축을 하지 않고, 모든 정보를 저장하여 처리하는 방식으로 작동하지만, 이로 인해 비효율적인 계산이 발생합니다.
  • 반면, RNN과 같은 재귀 모델은 정보를 압축하여 처리하지만, 압축된 정보가 손실되면 성능이 떨어질 수 있습니다.
• 선택적 압축: SSSM(S4)은 컨텍스트 정보를 선택적으로 압축하여 저장하거나 잊을 수 있습니다.
  • 즉, 시퀀스 내에서 중요한 정보는 기억하고, 불필요한 정보는 무시하는 방식입니다. 이를 통해 정보 압축을 최적화하고, 모델이 시퀀스 전반에 걸쳐 중요한 정보를 놓치지 않도록 합니다.

본문에서는 Figure2를 통해 “정보 선택 및 복사 작업” 또는 “정보 필터링 작업” Task들을 소개하며, 기본적인 LTI 시스템은 선택적 메커니즘과 문맥 처리 능력이 필요한 Selective Copying Task와 Induction Heads Task에 적합하지 않음을 지적합니다.

• 1. Copying Task (Figure2 왼쪽 이미지)

  

  • 문제: Copying Task는 시퀀스 데이터를 기억하고, 특정 위치에서 복사하는 작업입니다.

    • 이때, 입력과 출력 간의 간격이 일정하게 유지됩니다. 모델은 시퀀스의 일정한 간격에 있는 데이터를 기억하고, 그대로 복사하면 됩니다.

  • 해결방법: 이 작업은 매우 간단한 패턴이기 때문에 시간 불변 모델(Time-Invariant Model)로 쉽게 해결할 수 있습니다.

    • 시간 불변 모델은 모든 시점에서 동일한 방식으로 데이터를 처리하는 모델로, 선형 재귀 모델(Linear Recurrence Model)이나 글로벌 합성곱 모델(Global Convolution Model) 같은 방식이 사용될 수 있습니다.
    • 이러한 모델은 입력 간의 일정한 간격을 인식하고, 그에 따라 데이터를 복사하는 데 적합합니다.

  • 결과: Copying Task는 간격이 고정되어 있어, LTI(Linear Time-Invariant) 모델로 쉽게 처리할 수 있는 단순한 작업입니다.

    • 모델은 시간 흐름에 따라 동일한 방식으로 데이터를 복사하여 이 작업을 완벽히 해결할 수 있습니다.

• 2. Selective Copying Task (Figure2 오른쪽 위 이미지)

  

  • 문제: Selective Copying Task는 Copying Task와 달리, 입력과 출력 간의 간격이 일정하지 않고 랜덤하게 변동됩니다.

    • 모델은 시퀀스 내에서 중요한 정보를 선택적으로 기억하고, 나머지 불필요한 정보는 무시해야 합니다.

  • 해결방법: 이 문제를 해결하려면 시간 가변 모델(Time-Varying Model)과 선택적 메커니즘(Selection Mechanism)이 필요합니다.

    • 선택적 메커니즘은 모델이 입력 데이터를 분석하여 중요한 정보를 선택적으로 기억하고, 불필요한 데이터는 무시할 수 있게 만듭니다.
    • Selective State Space Model(SSSM)과 같은 구조를 사용하면, 입력 시퀀스 내에서 어떤 정보가 중요한지 선택하고, 중요한 정보를 선택적으로 기억할 수 있습니다.

  • 결과: Selective Copying Task는 시간 가변적 처리와 선택적 메커니즘을 통해 해결됩니다.

    • 이로 인해 모델은 각 시점에서 중요한 데이터를 선택하고, 불필요한 정보는 무시하여 불규칙한 입력 간격에서도 정확하게 복사 작업을 수행할 수 있습니다.

• 3. Induction Heads Task (Figure2 오른쪽 하단 이미지)

  

  • 문제: Induction Heads Task는 연관 기억(Associative Recall) 문제로, 모델이 이전에 학습된 정보를 바탕으로 문맥(Context)을 이해하고, 문맥에 맞는 출력을 유추해야 합니다.

    • 이 작업에서는 주어진 시퀀스에서 특정 패턴이 주어진 후, 비슷한 패턴이 다시 나올 때 올바른 출력을 예측할 수 있어야 합니다. 이는 문맥 기반의 학습과 회상을 요구합니다.

  • 해결방법: 이 문제를 해결하기 위해서는 모델이 문맥을 학습하고 연관 지어 기억할 수 있어야 합니다.

    • 단순한 Copying 작업과 달리, 모델은 이전 시점의 문맥을 연관 지어 기억하고, 필요한 시점에서 이를 회상하여 적절한 출력을 유추해야 합니다.
    • 선택적 메커니즘과 함께 연관 기억 메커니즘을 사용하면, 모델이 문맥을 기반으로 필요한 정보를 기억하고, 새로운 입력이 들어올 때 그 정보를 다시 사용할 수 있습니다.
    • Selective State Space Model (SSSM)은 이러한 문맥 기반 학습에 적합한 구조를 제공합니다.

  • 결과: Induction Heads Task는 문맥 기반의 연관 기억이 중요한 작업입니다. 선택적 메커니즘과 연관 기억 메커니즘을 통해 모델은 문맥 정보를 학습하고, 새로운 입력과 관련된 패턴을 연관 지어 회상할 수 있게 됩니다. 이는 문맥에 맞는 예측을 가능하게 해줍니다.

정리

• Copying Task: 일정한 간격의 시퀀스 데이터를 단순히 복사하는 작업으로, 시간 불변 모델(LTI)을 통해 쉽게 해결됩니다. 이 작업은 고정된 구조의 LTI 모델이 시간 인식을 필요로 하지만, 정보 선택이 단순하기 때문에 처리할 수 있습니다.
• Selective Copying Task: 입력과 출력 간의 간격이 랜덤하게 변동하는 작업으로, 시간 가변 모델과 선택적 메커니즘을 통해 중요한 정보만 선택적으로 기억함으로써 해결됩니다. 이 작업은 내용 인식이 필수적이므로, 기존 LTI 모델에서는 효과적으로 처리하기 어렵습니다.
• Induction Heads Task: 연관 기억 문제로, 문맥 정보를 학습하고 연관 지어 회상하는 능력이 필요합니다. 선택적 메커니즘과 문맥 기반 기억을 통해 해결할 수 있습니다. 이 작업은 복잡한 상관관계를 이해하고 기억하는 것이 중요하여, 보다 발전된 모델 구조가 요구됩니다.

3.2 Improving SSMs with Selection (선택을 통한 SSM 성능 향상)

이 섹션에서는 선택 메커니즘을 기본 SSM 구조에 통합하여, 모델의 성능을 향상시키는 방법을 설명합니다. SSM (S4)는 고정된 파라미터를 사용하여 간단한 구조로 작동하는 반면, SSM + Selection (S6)는 입력 데이터를 동적으로 처리하여 선택적으로 정보를 강조하거나 무시할 수 있는 보다 복잡한 구조입니다.

• 선택 메커니즘: SSSM(S4+Selection, S6)은 SSM의 주요 매개변수(Δ,B,C\Delta, B, CΔ,B,C)를 입력에 따라 선택적으로 변동시킴으로써 데이터를 처리합니다.

  • 이를 통해 모델은 시퀀스의 중요한 부분을 선택적으로 기억하고, 불필요한 부분은 무시할 수 있습니다.

• 시간 불변성을 포기하고 효율성 극대화: 이 선택 메커니즘은 시간 불변성을 유지하지 않기 때문에, 시간에 따라 동적으로 변화하는 시퀀스를 처리할 수 있지만, 이를 효율적으로 구현하기 위해서는 추가적인 알고리즘이 필요합니다.

기본 SSM (S4)과 SSM + Selection (S6)의 알고리즘을 한번 비교하면서 차이점을 살펴보겠습니다.

기본 SSM 구조 (Algorithm 1: SSM (S4))

• 입력:

  • xxx: 입력 데이터로, 형태는 (B, L, D)입니다.
  • 여기서 B는 배치 크기(batch size), L은 시퀀스 길이, D는 채널 수를 나타냅니다.

• 출력:

  • yyy: 출력 데이터로, 입력과 같은 형태를 갖습니다. 즉, (B, L, D)입니다.

• 매개변수:

  • AAA, BBB, CCC: 이 세 개의 매개변수는 SSM의 핵심 파라미터로, 각각 잠재 상태와 입력 데이터를 변환하는 역할을 합니다.
  • Δ\DeltaΔ: 시간 스케일을 조정하는 매개변수로, 상태 공간 모델에서 시간적 변화와 관련된 역할을 합니다.

• 작동 방식:

  1. AAA, BBB, CCC 파라미터가 설정됩니다.
  2. Δ\DeltaΔ 값이 파라미터로 설정됩니다.
  3. 주어진 Δ\DeltaΔ와 AAA, BBB, CCC 값들을 기반으로 이산화(discretization)가 이루어집니다.
  4. 상태 공간 모델이 시간 불변(Time-invariant)인 재귀(recursion) 또는 합성곱(convolution)을 통해 계산됩니다.

💡 핵심 특징: 이 알고리즘은 시간 불변적 구조로, 입력 데이터를 고정된 방식으로 처리합니다. 이는 동일한 매개변수를 모든 시점에 적용한다는 의미입니다.

⏱ 계산 방식: 선형 재귀 또는 합성곱 연산을 사용하여 시간 불변적 처리만 가능합니다.

🔎 (심화) S4 (SSM) 알고리즘의 텐서 연산 및 차원 변환

• S4는 고정된 파라미터와 이산화 방식을 사용하여 모든 시퀀스에 동일한 연산을 적용하고, 그 결과로 일정한 recurrence 또는 convolution을 수행합니다.

• 입력 텐서 (x): (B, L, D)의 형태를 가지며, B는 배치 크기, L은 시퀀스 길이, D는 각 토큰의 차원을 나타냅니다.

  • 즉, B개의 샘플에 대해 L개의 시퀀스를 처리하며, 각 시퀀스는 D 차원의 벡터로 표현됩니다.

• S4에서는 (D, N)의 고정된 파라미터가 모든 시퀀스에 동일하게 적용됩니다.

  • 
  • 
  • 
  • 

• 파라미터 A, B, C: A, B, C는 모두 (D, N) 형태로 존재하며, 여기서 D는 입력 차원, N은 숨겨진 차원(hidden state)의 크기입니다.

  • A: Structured N×NN \times NN×N 매트릭스로 연산을 담당합니다.
  • B: (D, N) 크기의 파라미터로서, 입력 텐서와 곱해져 새로운 상태(state)를 생성하는 역할을 합니다.
  • C: 역시 (D, N) 크기를 가지며 출력 텐서 생성에 영향을 미칩니다.
• 이산화 (discretization): 연속적 시스템을 이산화하여 연산을 수행하는데, 이때 사용하는 Δ는 (D) 크기의 파라미터입니다. 이 파라미터는 시간 간격을 이산화하여 A, B 매트릭스의 값을 변환합니다.
• S4에서는 고정된 Δ가 사용됩니다. 이산화된 Δ는 각 시퀀스에 대해 각각의 매트릭스 A, B와 곱해져 hth_tht​를 업데이트합니다.
• 최종 출력 y: (B, L, D) 크기의 결과를 반환하며, 이는 time-invariant 방식으로 recurrence나 convolution 연산이 적용됩니다.

✅ (결론) S4는 고정된 파라미터와 이산화 방식을 사용하여 모든 시퀀스에 동일한 연산을 적용하고, 그 결과로 일정한 recurrence 또는 convolution을 수행합니다.

선택적 SSM 구조 (Algorithm 2: SSM + Selection (S6))

• 입력 및 출력:

  • S4와 동일하게, 입력 데이터는 xxx로 (B, L, D)의 형태를 가지며, 출력 데이터도 같은 (B, L, D) 구조입니다.

• 주요 차이점:

  • 선택 메커니즘 적용: 입력 데이터에 따라 파라미터가 변화합니다.
    • 즉, S6에서는 입력 의존적인 선택(selectivity)이 추가되어 시점에 따라 달라지는 방식으로 처리됩니다.
    • 선택적 SSM에서는 BtB_tBt​, CtC_tCt​, Δt\Delta_tΔt​ 파라미터가 입력에 따라 동적으로 변화합니다

• 매개변수 변화:

  1. BBB, CCC: S4에서는 고정된 매개변수였으나, S6에서는 sB(x)s_B(x)sB​(x), sC(x)s_C(x)sC​(x)와 같은 함수로 입력 xxx에 따라 변화합니다.
  2. Δ\DeltaΔ: S4에서는 고정된 값이었으나, S6에서는 sΔ(x)s_{\Delta}(x)sΔ​(x)를 통해 입력에 따라 변화합니다. 이는 모델이 시점에 따라 가변적인 시간 스케일을 적용할 수 있게 만듭니다.

• 작동 방식:

  1. 매개변수 AAA, BBB, CCC와 시간 스케일 Δ\DeltaΔ는 입력 데이터 xxx에 따라 변화합니다.
  2. 각 시점에서 재귀적 계산(recurrence)만 수행되며, 이는 시간에 따라 변화하는 time-varying 모델입니다.

💡 핵심 특징: S6는 시간 가변(time-varying) 구조로, 입력에 따라 매번 다른 방식으로 데이터를 처리할 수 있습니다. 이 선택 메커니즘은 중요한 정보는 기억하고, 불필요한 정보는 무시할 수 있도록 만듭니다.

⏱ 계산 방식: 시간 가변적이기 때문에 재귀적 연산만을 수행하며, 시퀀스 전반에서 중요한 정보를 선택적으로 처리할 수 있습니다.

🔎 (심화) S6 (SSM + Selection) 알고리즘의 텐서 연산 및 차원 변환

• S6는 입력 종속적인 파라미터와 시간 변이에 따른 연산을 사용하여 시퀀스마다 다른 연산이 이루어지며, 이는 동적 모델링에 더 적합합니다.

• 입력 텐서 (x): (B, L, D)로 동일하지만, 입력 값에 따라 다양한 매트릭스 연산이 이루어집니다.

• S6에서는 입력 x에 따라 (B, L, N) 차원의 파라미터들이 각 시퀀스별로 다르게 생성됩니다. 이를 통해 각 시퀀스마다 다른 매핑이 일어납니다.

  • 
  • 

• 파라미터 변환 (sB, sC, sΔ): S6에서는 파라미터들이 입력에 종속적으로 변환됩니다.

  • sB(x): 입력 x에 따라 (B, L, N) 차원으로 변환됩니다. 즉, 각 배치 B의 각 시퀀스 L에 대해, 숨겨진 차원 N을 생성합니다. 이때 기존 S4에서 모든 시퀀스가 동일한 B 매트릭스를 사용했던 것과 달리, 이제는 각 시퀀스마다 다른 B가 적용됩니다.
  • sC(x): 역시 입력에 따라 (B, L, N) 차원으로 변환됩니다.
  • sΔ(x): 입력에 따라 시간 간격을 나타내는 Δ 역시 (B, L, D)로 변환됩니다. 즉, 각 배치와 각 시퀀스마다 다른 Δ 값이 주어집니다.
• 이산화 연산 (discretization): S4와 유사하게 이산화를 통해 연산이 이루어지지만, 여기서는 입력에 종속적으로 변환된 Δ가 사용되므로 더 복잡한 형태의 연산이 발생합니다.

• S6에서는 sΔ(x)가 각 시퀀스별로 다르게 계산됩니다. 이산화된 Δ는 각 시퀀스에 대해 각각의 매트릭스 A, B와 곱해져 hth_tht​를 업데이트합니다.

  • 
  • 

✔️ 이때 각 시퀀스는 고유한 Δ를 가지고 있기 때문에 S6에서는 토큰마다 서로 다른 연산이 적용됩니다.

• 최종 출력 y: S6의 최종 출력 역시 (B, L, D) 차원을 가지지만, S4와 달리 입력에 따라 동적으로 변화한 recurrence 방식을 사용합니다. 특히, 각 시퀀스마다 다르게 이산화된 파라미터가 적용되기 때문에 각 토큰에 맞는 연산이 수행됩니다.

✅ (결론) S6는 입력 종속적인 파라미터와 시간 변이에 따른 연산을 사용하여 시퀀스마다 다른 연산이 이루어지며, 이는 동적 모델링에 더 적합합니다.

💬 Time In-Variant? 무슨 뜻이지?

• 시간 불변적 처리라는 개념은 “모델의 매개변수(가중치)가 시간에 따라 변하지 않는다는 것”을 의미합니다.
  • 즉, 모델이 입력 시퀀스의 각 시점(t)에 대해 동일한 매개변수를 사용한다는 것을 의미합니다.
  • 이는 다음과 같은 의미를 가집니다:
    • 정적인 가중치: 기존 SSM 모델에서 사용되는 매개변수 A, B, C 등은 시점마다 고정되어 있습니다. 따라서, 같은 입력에 대해서는 항상 같은 출력을 생성합니다.
      • 예를 들어, 과거의 입력이 미래의 출력에 영향을 미칠 때, 출력에 영향을 주는 가중치가 변하지 않기 때문에 특정 입력에 대해 적합하게 조정되지 않습니다.
    • 입력 의존성 부족: 입력 데이터의 특성에 따라 모델이 동적으로 반응하지 못합니다.
      • 예를 들어, 어떤 특정 입력이 매우 중요할 때 그 입력에 대한 반응을 강화하거나, 반대로 덜 중요한 정보는 무시하는 방법이 없습니다. 그래서, 입력 시퀀스의 맥락이나 중요한 정보에 따라 모델이 학습된 행동을 변화시킬 수 있는 능력이 제한됩니다.

🔎 S4와 S6의 차이점 분석

• 입력 의존성
  • S4: 고정된 파라미터를 사용하여, 모든 시점에서 동일한 계산을 수행합니다. 즉, 모든 시점에서 동일한 방식으로 입력 데이터를 처리합니다.
  • S6: 선택 메커니즘을 통해 입력 데이터 xxx에 따라 매개변수들이 동적으로 변화합니다. 이는 데이터의 특성에 따라 각 시점에서 필요한 정보를 선택적으로 처리할 수 있게 만듭니다.
• 시간 불변성(Time-invariant) vs 시간 가변성(Time-varying)
  • S4: 시간 불변적인 구조로, 동일한 파라미터가 모든 시점에 적용됩니다. 이는 주로 합성곱(convolution)이나 재귀(recursion) 형태로 데이터를 처리합니다.
  • S6: 시간 가변적인 구조로, 입력에 따라 매개변수들이 변화하고, 재귀적 방식으로 계산이 이루어집니다. 이를 통해 시퀀스의 각 시점에서 중요한 정보는 기억하고, 불필요한 정보는 무시하는 선택적 처리가 가능합니다.
• 효율성
  • S4: 시간 불변성을 유지하는 SSM은 계산의 병렬화가 가능하여, 비교적 효율적인 계산을 수행할 수 있습니다.
  • S6: 선택 메커니즘을 추가함으로써 더 많은 계산이 필요할 수 있지만, GPU의 메모리 계층을 활용한 하드웨어 최적화가 가능해져 효율성을 유지합니다. 특히, 입력 데이터의 특성에 따라 동적인 계산을 수행하여 더 높은 성능을 낼 수 있습니다.

3.3 Efficient Implementation of Selective SSMs (효율적인 선택적 SSM 구현)

이 섹션에서는 Selective State Space Model(SSSM)을 하드웨어에서 효율적으로 구현하는 방법을 다룹니다. 특히, GPU 등을 활용하여 메모리 사용과 계산을 최적화하는 방법에 대해 설명하고 있습니다.

3.3.1 Motivation of Prior Models (이전 모델들의 동기)

이 항목에서는 선택적 상태 공간 모델(SSSM)이 나오기 전, 기존 상태 공간 모델(SSM)이 어떻게 동작했는지, 그리고 왜 개선이 필요했는지를 설명합니다.

• 1. 기존 상태 공간 모델의 동작 원리

  • SSM(Structured State Space Model)은 시퀀스 데이터를 처리하기 위해 잠재 상태(latent state)를 사용하는 모델입니다. 이 모델은 시퀀스 내의 정보를 재귀적으로 처리하여, 시점 간의 의존성을 유지하면서 긴 시퀀스 데이터를 효율적으로 처리할 수 있습니다.
  • SSM은 시간 불변적(time-invariant)으로 설계되어, 각 시점에서 동일한 방식으로 데이터를 처리합니다. 이는 선형 재귀적 구조(linear recurrence)나 합성곱 연산(convolution)을 사용하여 계산됩니다.

• 2. 기존 모델의 한계

  • 고정된 파라미터: 기존의 SSM은 모든 시점에서 동일한 파라미터로 데이터를 처리합니다. 이는 입력 데이터의 특성이나 중요도에 따라 가변적으로 처리할 수 없다는 한계를 지니고 있습니다. 즉, 모든 데이터를 동일하게 취급하기 때문에 중요한 정보만 선택적으로 처리하는 기능이 부족합니다.
  • 복잡한 계산: 시퀀스가 길어질수록 계산 복잡도가 크게 증가합니다. 특히 고차원의 상태 공간에서 작업할 경우, 계산 비용이 매우 높아지며 메모리 사용량도 크게 증가합니다.
  • 효율성 문제: SSM은 시퀀스 데이터를 처리할 때 모든 정보를 기억해야 하기 때문에, 메모리 사용량이 매우 크고 계산 시간도 길어집니다. 이는 긴 시퀀스를 처리할 때 효율성이 떨어지는 문제를 초래합니다.

• 3. 개선 필요성

  • 입력에 따라 유동적인 처리가 필요합니다. 기존 모델은 모든 시점에서 동일한 방식으로 데이터를 처리했지만, 입력의 중요도에 따라 선택적으로 데이터를 처리하는 기능이 있으면 더 효율적으로 작동할 수 있습니다.
  • 또한, 메모리와 계산 자원을 더 효율적으로 사용하기 위해, 기존 모델보다 더 적은 자원으로 높은 성능을 낼 수 있는 최적화된 방식이 필요했습니다.

3.3.2 Overview of Selective Scan: Hardware-Aware State Expansion (선택적 스캔: 하드웨어 인식 상태 확장의 개요)

이 항목에서는 Selective Scan의 개념과, 이를 통해 SSSM이 하드웨어 상에서 어떻게 효율적으로 구현될 수 있는지를 설명합니다. 여기서 중요한 개념은 하드웨어의 메모리 계층을 최적화하여 선택적 상태 공간 모델의 성능을 극대화하는 방법입니다.

• 1. Selective Scan의 개념

  • Selective Scan은 시퀀스 내에서 중요한 정보를 선택적으로 처리하고, 불필요한 정보는 무시하는 과정을 의미합니다. 이를 통해 모델은 중요한 정보만을 선택적으로 기억하면서, 불필요한 연산을 줄일 수 있습니다.
  • 시간 가변적(time-varying)이라는 특성을 가진 Selective Scan은 각 시점에서 동적으로 변화하는 상태를 기반으로 연산을 수행합니다.
    • 이때 각 시점에서 입력 데이터를 분석하여 중요한 정보만 선택적으로 처리하기 때문에, 메모리 사용과 계산 자원을 절약할 수 있습니다.

• 2. 하드웨어-인식 상태 확장 (Hardware-Aware State Expansion)

  • Hardware-Aware State Expansion은 선택적 상태 공간 모델의 계산을 하드웨어 효율성을 고려하여 최적화하는 방법입니다.

  • GPU 메모리 계층 활용: 현대 GPU는 고속 메모리(SRAM)와 대용량 메모리(HBM)를 가지고 있습니다. 선택적 SSM에서는 이러한 메모리 계층을 적절히 활용하여, 자주 사용되는 중요한 데이터는 고속 메모리에 저장하고, 덜 중요한 데이터는 대용량 메모리에 저장함으로써 계산 속도와 메모리 사용을 최적화할 수 있습니다.

    • 고속 SRAM은 즉각적인 데이터 접근을 제공하고, 대용량 HBM은 복잡한 연산에 필요한 대량의 데이터를 저장합니다.
    • 이 조합은 GPU가 그래픽 렌더링, 머신 러닝, 과학적 시뮬레이션 등 데이터 집약적인 작업을 효과적으로 수행할 수 있게 합니다.

Kernel Fusion : Hardware-aware Algorithm

• 처음 등장하는 개념은 아니고, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness“에 나오는 idea라고 합니다.

  

• GPU의 주요 병목 현상은 SRAM과 DRAM 사이의 Copy and PASTE에서 발생하는 것을 확인하였고, 저자는 이러한 memory IO 로 발생하는 병목현상을 줄이기 위하여 kernel fusion을 사용하였습니다.
• Mamba는 계산 자체보다는 메모리 전송 과정에서 병목이 발생하는 GPU의 구조를 고려해 성능을 극대화했습니다.
• 입력 벡터와 가중치 매개변수를 고성능 메모리로 전송한 후 모든 계산을 한 번에 처리하고, 다시 메인 메모리로 데이터를 전송합니다.
  • 이로 인해 데이터 전송 시간은 그대로 유지되면서도 16배 확장된 벡터를 사용하는 데 필요한 추가 계산 시간을 거의 무료로 사용할 수 있게 됩니다.

아래 그림은 개략적으로 “입력 벡터와 가중치 매개변수를 고성능 메모리로 전송한 후 모든 계산을 한 번에 처리하고, 다시 메인 메모리로 데이터를 전송”하는 과정을 도시화한 그림입니다. (Source: https://youtu.be/N6Piou4oYx8)

• Scan Operation: Selective Scan 계산을 병렬로 처리할 수 있도록 설계하여, 시퀀스의 여러 시점을 동시에 처리하는 방식을 사용합니다.
  • 
  • 병렬 스캔 알고리즘(Parellel Scan Operation)을 통해 입력 시퀀스를 동시에 처리하면서, 선택적으로 필요한 정보를 처리하고 나머지는 건너뛰는 방식입니다.
  • 이를 통해 전체 시퀀스를 순차적으로 처리하는 방식보다 더 빠르고 효율적으로 연산을 수행할 수 있습니다.

  • 결합규칙(association rule) 기반 접근법을 사용하여 “먼저 계산할 수 있는 것은 계산해주자!” 라는 간단하지만 강력한 방법을 사용했습니다.

    

  • 이 방법을 사용하면 O(log⁡(n))O(\log(n))O(log(n))의 시간 내에 병렬 처리를 수행할 수 있어 계산 속도가 크게 향상됩니다.
• 3. 시간 가변적 선택 처리
  • Selective SSM은 시간 가변적이기 때문에, 각 시점마다 다른 방식으로 데이터를 처리합니다. 이로 인해, 각 시점에서 중요한 정보를 선택적으로 처리하고, 재귀적 연산(recurrent operation)을 통해 이전 상태를 기반으로 다음 상태를 계산합니다.
  • 각 시점에서 처리되는 데이터의 양을 줄이기 위해, Selective Scan을 통해 입력 데이터를 분석하고 필요한 정보만 선택합니다. 이를 통해 메모리 사용량과 계산량을 줄이고, 연산 속도를 크게 향상시킬 수 있습니다.

3.4 A Simplified SSM Architecture (단순화된 SSM 아키텍처)

• Mamba 아키텍처: SSM을 MLP 블록과 결합하여 간단한 형태의 아키텍처를 만들었습니다. 이 아키텍처는 Transformer처럼 복잡한 구조를 갖지 않으며, 단순하지만 강력한 성능을 보여줍니다.
  • 여러 개의 Mamba 블록을 반복적으로 쌓아 모델의 깊이를 확장할 수 있습니다.
  • Mamba 블록은 현대 신경망의 다층 퍼셉트론(MLP) 블록과 대부분의 SSM 아키텍처의 기초가 되는 H3 블록을 조합한 것입니다.

위 그림은 H3 블록, Gated MLP 블록, 그리고 Mamba 블록의 구조를 비교한 것입니다. 각 블록은 현대 신경망에서 사용되는 구조적 차이점을 시각적으로 보여주며, 이를 통해 Mamba 아키텍처가 어떻게 설계되었는지를 설명합니다. 그림에 대한 주요 해석은 아래와 같습니다.

1. H3 Block

• H3는 고차원 데이터를 효과적으로 처리하기 위해 설계된 블록 구조로, RNN과 CNN의 장점을 결합하여 순차적 데이터를 처리하는 데 강력한 성능을 보여주고 있습니다.
  • 이 블록은 과거의 정보와 현재 입력을 기반으로 미래의 출력을 예측하는 방식으로 작동합니다.
  • H3의 주요 특징은 긴 시퀀스를 효율적으로 처리할 수 있는 능력과, 시퀀스의 모든 요소에 대한 정보를 고려할 수 있는 구조적 특성입니다.
• 구성 요소:
  • SSM: 상태 공간 모델(State Space Model)은 시퀀스 변환을 담당합니다. 이는 주로 재귀적인 특성을 활용하여 시퀀스 데이터를 처리합니다.
  • Conv: 합성곱 층이 추가되어 로컬 정보를 처리합니다. 합성곱은 일반적으로 공간적 연관성을 다루는 데 효과적입니다.
  • 곱셈 게이트(Multiplicative Gate): SSM과 Conv 사이에 곱셈 게이트가 있어 데이터의 흐름을 조절합니다.
• 동작 원리:
  • H3 블록은 SSM과 합성곱 층을 교차 배치하여 각각의 시퀀스 데이터와 로컬 데이터를 처리하며, 곱셈 게이트를 통해 중요한 정보를 통과시키거나 억제할 수 있습니다. 이 구조는 데이터를 처리할 때 유연성을 제공하지만, 계산적으로 복잡할 수 있습니다.

2. Gated MLP

• Gated MLP는 다층 퍼셉트론(MLP)과 곱셈 게이트(Multiplicative Gate)를 결합하여 입력 데이터에 대한 비선형 변환을 수행하는 구조입니다.

  • MLP는 일반적으로 고정된 입력 데이터를 처리하는 데 탁월한 능력을 가지고 있으며, 곱셈 게이트는 데이터의 특성에 따라 중요한 정보를 강조하고, 불필요한 정보는 억제할 수 있도록 도와줍니다.

• 구성 요소:

  • MLP: 다층 퍼셉트론(Multi-Layer Perceptron) 블록은 입력 데이터를 비선형적으로 변환합니다.
  • 곱셈 게이트(Multiplicative Gate): H3와 마찬가지로 곱셈 게이트가 추가되어, 데이터를 선택적으로 처리하는 역할을 합니다.

• 동작 원리:

  • Gated MLP는 MLP와 곱셈 게이트를 결합하여, 입력 데이터에 대한 유연한 변환을 수행합니다. 데이터의 특성에 따라 중요한 정보는 곱셈을 통해 강조되고, 불필요한 정보는 억제될 수 있습니다.
  • 그러나 Gated MLP는 비시퀀스 데이터를 처리하는 데는 효과적이지만, 시퀀스 처리에 필요한 메커니즘(예: SSM)이 포함되어 있지 않다는 점에서 시퀀스 기반 작업에 최적화된 구조는 아닙니다.

3. Mamba Block

• Mamba 블록은 현대 신경망의 다층 퍼셉트론(MLP) 블록과 SSM(Structured State Space Model) 아키텍처에서 중요한 역할을 하는 H3 블록을 결합한 설계입니다.

  • 이 구조는 기존의 MLP와 SSM 블록을 교차하거나 혼합하는 대신, 동일한 Mamba 블록을 동질적으로 반복하는 방식으로 설계되었습니다.

• 구성 요소:

  • SSM: H3 블록과 마찬가지로 상태 공간 모델(SSM)이 존재하여 시퀀스 데이터를 처리합니다.
  • Conv: 합성곱 층이 추가되어, 시퀀스 내의 국소적인 정보 처리에 기여합니다.
  • 활성화 함수(SiLU/Swish): H3와는 다르게, 곱셈 게이트 대신 활성화 함수가 사용됩니다. 이 함수는 비선형성을 추가하여 데이터의 표현력을 높입니다.

• 동작 원리:

  • Mamba 블록은 SSM과 Conv를 결합하여 시퀀스 및 로컬 데이터를 모두 처리합니다.
  • 곱셈 게이트 대신 비선형 활성화 함수(SiLU 또는 Swish)를 사용하여 계산 복잡성을 줄이고 효율성을 높였습니다. 활성화 함수는 데이터의 흐름을 조절하는 데 더 간단한 방식을 사용하며, 게이트의 필요성을 제거하여 더 단순한 구조를 만들었습니다.
  • Mamba 블록은 두 가지 주요 연산(SMM, Conv)을 한 블록 내에서 반복하는 간단하면서도 강력한 구조를 채택했습니다.

• Mamba와 다른 블록의 차이점

  • 복잡성 감소: H3 블록에서는 곱셈 게이트를 사용하여 데이터를 선택적으로 처리하지만, Mamba는 이를 제거하고 활성화 함수로 대체함으로써 계산 복잡성을 줄였습니다.
  • 단순화된 아키텍처: Mamba는 SSM과 Conv를 동일한 블록 내에서 반복적으로 사용하여 매우 균일하고 단순한 아키텍처를 유지하면서도 성능을 최적화했습니다. 이는 전체적으로 더 빠르고 효율적인 계산을 가능하게 합니다.
  • 비선형 활성화: Mamba에서는 곱셈 게이트 대신 SiLU/Swish와 같은 비선형 활성화 함수를 사용하여 데이터의 흐름을 조절합니다. 이는 모델의 표현력을 유지하면서도 계산 복잡성을 줄이는 중요한 변화입니다.

✨ Mamba 블록은 H3 및 Gated MLP와 비교하여 더 단순하고 효율적인 아키텍처를 제공합니다. SSM과 Conv를 적절히 결합하여 시퀀스 데이터와 로컬 정보를 동시에 처리하며, 활성화 함수로 비선형성을 부여하여 모델의 성능을 높였습니다.

3.5 Properties of Selection Mechanisms (선택 메커니즘의 특성)

• 선택 메커니즘의 효과: 선택 메커니즘은 불필요한 데이터를 무시하고, 중요한 데이터를 선택적으로 처리하는 능력을 가집니다. 이를 통해 시퀀스의 긴 문맥을 처리할 때 효율성이 증가하며, 긴 시퀀스에서도 성능 저하가 발생하지 않습니다.
• 변수 간 상호작용: 선택 메커니즘은 시퀀스의 각 요소들이 상호작용하는 방식을 조절하며, 이러한 조절 능력은 특히 텍스트나 DNA와 같은 이산적인 데이터에서 효과적입니다.

3.6 Additional Model Details (추가 모델 세부사항)

• 실수 및 복소수 처리: 선택적 SSM은 복소수와 실수를 모두 처리할 수 있지만, 특정 작업에서는 실수 기반 모델이 더 나은 성능을 보일 수 있습니다.
• 초기화 및 파라미터화: 선택적 파라미터의 초기화 방식에 따라 모델의 성능이 달라지며, 각 파라미터에 대한 자세한 설명을 통해 모델의 안정성을 유지합니다.

4. Empirical Evaluation

Mamba 모델을 다양한 데이터 유형과 시퀀스 길이에서 테스트한 결과를 소개합니다.

4.1 Synthetic Tasks (합성 작업)

• Selective Copying: 선택 메커니즘을 사용한 Mamba는 시퀀스의 중요한 부분을 기억하고 나머지를 무시하는 작업에서 뛰어난 성능을 보입니다. 시퀀스 길이가 매우 길어도 Mamba는 성능 저하 없이 정확도를 유지합니다.
• Induction Heads: LLM의 맥락 학습 능력을 평가하는 Induction Heads 작업에서도 Mamba는 중요한 토큰을 기억하며 성능을 유지합니다. 훈련 시 256 길이의 시퀀스로 학습한 모델이 1백만 길이의 시퀀스에서도 정확한 결과를 보입니다.

4.2 Language Modeling (언어 모델링)

• Mamba의 언어 모델링 성능: Mamba는 텍스트 데이터에서 Transformer와 비슷하거나 더 나은 성능을 보입니다. 특히, 1B 이상의 파라미터를 가진 모델에서는 Transformer와 비슷한 수준의 성능을 내는 첫 번째 선형 시퀀스 모델입니다.

4.3 DNA Modeling (DNA 모델링)

• DNA 시퀀스 처리: Mamba는 긴 시퀀스 데이터를 처리하는 데 뛰어나며, 기존의 Transformer 모델보다 더 나은 성능을 보입니다. DNA의 경우 긴 문맥 의존성이 중요한데, Mamba는 이러한 데이터를 잘 처리합니다.

4.4 Audio Modeling and Generation (오디오 모델링 및 생성)

• 오디오 데이터 처리: Mamba는 오디오 데이터를 처리할 때도 효율적이며, 기존 모델보다 더 긴 시퀀스를 처리할 수 있습니다.

4.5 Speed and Memory Benchmarks (속도 및 메모리 벤치마크)

• 처리 속도: Mamba는 Transformer보다 5배 빠른 추론 속도를 보이며, 긴 시퀀스를 처리할 때 메모리 사용량도 매우 적습니다.

4.6 Model Ablations (모델 에블레이션)

• 파라미터 분석: 선택적 파라미터를 추가할수록 성능이 향상되며, 특히 Delta 파라미터가 모델 성능에 가장 중요한 영향을 미칩니다.

Reference

Paper

• Mamba 논문 : https://arxiv.org/pdf/2312.00752

Blogs

• A Visual Guide to Mamba and State Space Models: https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mamba-and-state

Youtube

• MAMBA from Scratch (https://youtu.be/N6Piou4oYx8)
• Mamba Paper Review, DSBA 천재원 (https://youtu.be/JjxBNBzDbNk)
• Mamba Paper Review, AirLab 이정운 (https://youtu.be/l-dQCTv9wIg)

• Previous
  [Paper Review] Structured State Space Models for Deep Sequence Modeling
• Next
  [Paper Review] A COMPREHENSIVE REVIEW OF YOLO ARCHITECTURES IN COMPUTER VISION: FROM YOLOV1 TO YOLOV8 AND YOLO-NAS