1 Introduction

Transformer

• 장점
  • 복잡한 문맥 모델링에 강함
  • Training Parallelism 가능
• 단점
  • Bottleneck
    • 추론 시에 다음 토큰을 생성할 때, 다시 처음부터 전체 sequence에 대한 attention 계산 필요
  • Quadratic scaling with respect to sequence length
    • 시퀀스 길이에 따라 비용 2차로 증가
  • 빠른 학습, 느린 추론

RNN

• 장점
  • 시퀀스 길이에 따라 비용 linear하게 증가
  • 이전 state와 현재 input으로 다음 state을 생성하기 때문에,
    transformer와 같이 모든 이전 state들과의 연산이 필요 없음
• 단점
  • Training Parallelism 불가
  • 시간이 지남에 따라 앞의 정보를 잊어버리는 경향이 있음
• 느린 학습, 빠른 추론

2 State Space Models

SSM

• Mamba 모델의 뼈대를 이루고 있는 아키텍처
• t 시점의 입력값이 들어왔을 때, hidden state를 업데이트하고 출력으로 변환하는 연속 시퀀스 모델
  
  

State Equation: 상태가 어떻게 변하는지

• A (state transition): 이전 상태가 얼마나 영향을 줄 지 = 기억력
• B (input projection): x(t)가 얼마나 영향을 줄 지 = 입력정보 x의 주입 강도 및 방식
  
  

Output Equation: 상태가 어떻게 출력으로 변환되는지

• C (output projection): h(t)의 어떤 부분을 출력으로 쓸 지 = 출력 생성 시 중요한 정보

Discretization

• SSM은 model based on continuous time domain, 하지만 대부분 discrete input 사용
• 흔히 다룰 수 있는 discrete한 sequence에 적용하도록 Discretization 수행
• 과정: 이산 입력 → 연속 시스템 적용 → 다시 이산 출력
  1. 이산 입력을 연속 신호로 해석: ZOH (zero-order hold)
     • 다음 이산신호를 받을 때까지 이전 이산신호 유지
     • ∆: timestep, 값을 유지하는 시간 길이
       
  2. Discretized A,B → $\bar{A}, \bar{B}$
     • $\bar{A} = \exp(\Delta A), \quad \bar{B} = (\Delta A)^{-1}\bigl(\exp(\Delta A) - I\bigr)\,\Delta B$
  3. State Equation, Output Equation using $\bar{A}, \bar{B}$
     • State Equation
       $h_k=\bar{A}h_{k-1}+\bar{B}x_k$
     • Output Equation
       $y_k = Ch_k$
     • 다시 설명하자면,
       • A: 이전 상태가 얼마나 유지/감쇠되어 전달되는 정도
       • B: 현재 입력이 상태에 반영되는 정도

LTI (Linear Time Invariance)

• Linear
  • 상태 변화가 선형 결합으로 표현됨
  • fast training 가능
• Time Invariance
  • A, B, C가 timestep에 따라 변하지 않음
• Structured SSMs have all been LTI
  • 이유: convolution-based computation는 time-invariant일 때만 가능
• LTI constraint was enforced for computational efficiency
• LTI의 한계: time-invariant → 데이터 특성에 따라 가변적 구조를 표현하기 어려움
• Mamba
  • time-variant이면서도 computational efficiency를 유지해 bottleneck 문제를 해결해보겠다

Computation

• After parameters have been transformed to ‘discrete parameters’, model can be computed in 2 ways

1. Linear Recurrence
   
   $h_k=\bar{A}h_{k-1}+\bar{B}x_k, y_k = Ch_k$

   • 순차적으로 state update
   • 병렬화 어려움
   • "추론" 시에 사용

2. Global Convolution
   

   • 어떤 timestep에서 예측을 하든, 커널 내의 값은 변하지 않음
   • 토큰마다 행해지는 연산이 달라지지 않기 때문에, global하게 적용될 수 있는 커널을 미리 만들어놓을 수 있음
   • 전체 입력 시퀀스에 대해 한 번에 convolution 수행
     병렬화 가능
   • "훈련" 시에 사용

3.1 Motivation: Selection as a Means of Compression

• Sequence Modeling의 근본적인 문제: context를 작은 state로 압축하는 문제
• Efficiency vs Effectiveness Tradeoff
  • Effectiveness : small state
  • Efficiency : containing all necessary information/performance
• LTI 모델인 SSM의 문제
  • 고정된 dynamics (A,B) 사용 → 입력에 따라 달라지지 않음
  • Input-independent
    • 입력에 따라 정보 선택 불가
    • context에서 중요한 정보 선택 불가
• 제안: Selectivity
  • Synthetic Tasks
    1. Selective Copying : 입력 중 중요한 정보만 저장, 나머지는 필터링하는 작업
    2. Induction Heads : context 속 특정 패턴을 보고, 그 다음에 올 출력 맞추는 작업
       

3.2 Improving SSMs with Selection

• Selective Mechanism 구현 방법 : “input-dependent parameter”

• 기존 SSM
  

  • 기존 SSM (S4)의 경우, 고정된 파라미터를 사용해 간단한 구조로 작동
  • Time-invariant Parameters : x값이 무엇이 들어오든, 똑같은 transition을 거침
  • Shape of A
    
  • Shape of B
    

• 기존 SSM의 대안으로, input-dependent parameter 제안
• 각 batch(B)와 각 sequence(L)마다 다른 값 → shape: (B, L, .. )
• B, C, Δ는 함수 $s_B, s_C, s_Δ$연산의 결과
  • 모두 x를 input으로 하여 B, C, Δ를 만들어내는 함수
• SSM + Selection
  
  • $s_B, s_C$
    • Input: (B, L, D) 차원의 X , Output: (B, L, N) 차원
    • D차원을 N차원으로 linear mapping하는 함수
    • (B, L) 차원 추가 → 각각의 B 내에, 각각의 L 내에 적용될 수 있는 tensor가 만들어졌음
  • $s_Δ$
    • 각 token에 대해 각각 다른 상수로 discretize
    • $s_Δ$는 shape를 바꾸지 않음 (B,L,D)
    • D차원을 1차원으로 projection → 1차원을 D차원으로 broadcast → parameter 더하고, softplus
  • Δ를 통해 A와 B를 discretize
    • Δ: (B, L, D) , A: (D, N) , B: (B, L, N)
    • 기존 SSM과 동일하게 ZOH로 discretize
    • 결국 discretized A, B는 data-dependent

3.3 Efficient Implementation of Selective SSMs

Motivation of Prior Models

• Efficiency vs Effectiveness Tradeoff
  • Mamba는 더이상 고정된 $\bar{A}, \bar{B},C$를 사용하지 않음
    - Input-Dependent
    - 미리 global하게 적용될 수 있는 kernel 만드는 것 불가능
    - convolution을 통한 training parallelism 불가
    → Efficiency Penalty
    
    

Overview of Selective Scan: Hardware-Aware State Expansion

1. Parallel Scan

• 문제: input-dependent → 미리 global kernel 생성 불가 → 병렬처리 어려움
• 조건: 연산의 associativity
  • 연산의 순서가 중요하지 않음. 연산의 순서와 무관하게 결과 동일해야 함.
• ‘먼저 계산할 수 있는 것들은 먼저 계산하자’
• 시퀀스를 부분적으로 계산하고 반복적으로 결합할 수 있음

2. Kernel Fusion

• 문제: large memory usage, memory complexity → bottleneck
  • (B,L,D)의 3차원으로 운용되는 Transformer에 비해 (B,L,D,N) 4차원의 Mamba는 GPU적 부담이 큼
  • GPU에서의 연산
    
  • 실제 연산을 수행하는 시간보다 HBM→SRAM, SRAM→HBM으로 복사하는 시간이 더 오래 걸림
• 대안: Kernel Fusion-여러 개의 연산을 하나의 연산으로 합쳐줌
  • 문제 발생 원인: 4차원이 생기는 SSM 내부
  • 3D인 A,B,C,Δ을 HBM->SRAM으로 복사
  • SRAM 안에서 바로 discretization+recurrence 모두 수행 (4D state를 HBM에 따로 복사X)
  • SSM의 모든 연산이 끝난 이후의 결과물(3D)만 SRAM→HBM으로 복사

3.4 A Simplified SSM Architecture

• Mamba Block designed by combining H3 and Gated MLP
  
• Mamba Block
  
  • Linear Projection
    • 입력 임베딩 확장 (보통 2배)
  • Convolution
    • 지역 정보 추출
    • 같은 채널을 따라 sequence 방향으로 인접 토큰끼리 정보 교환
  • Selective SSM
    • long-range 정보 추출
    • 모든 이전 정보 기억 X. data-dependent로 선택적으로 중요한 정보만 담아서 연산
  • Gating
    • Conv+SSM의 결과를 gating으로 조절
  • Linear Projection
    • 확장했던 차원을 입력 크기에 맞춰주기

3.5 Properties of Selection Mechanisms

Selective Parameter

• Δ
  • SSM 은 continuous system, Δ의 timestep으로 이산화
  • 의미: 현재 입력에 얼마나 집중할지
  • small Δ : state 유지, 현재 입력 무시
  • large Δ : state reset, 현재 입력 강하게 반영
• A
  • 실제로 모델에 영향을 주는 건 Δ로 이산화된
  • Δ만 selective으로 설정해도 충분히 성능 향상됨. A를 따로 selective하게 만들 필요 없음
• B, C
  • B: 입력 x가 hidden state에 얼마나 반영할 지를 조정
  • C: hidden state가 output에 얼마나 반영할 지를 조정

4 Empirical Evaluation

Synthetic Tasks

1. Selective Copying

• S4: not selective / S6: selective
• model/architecture보다는 selective copying layer가 성능에 더 중요한 요인
  

2. Induction Heads

• 추론 요구 길이를 늘려가면서 성능을 확인하는 그래프
• 학습은 고정된 길이의 시퀀스로 수행하고, 추론에서의 시퀀스 길이를 증가시킴
  

Transformer vs Mamba

Transformer의 문제점

• 비효율적인 긴 시퀀스 처리
• sequence의 길이가 증가하면 비용이 2차 함수 형태로 증가
• 모든 토큰을 고려하여 문맥 파악에 강하지만, long sequence를 처리할 때 bottleneck 발생 (메모리, 시간)

Mamba

• 전체 토큰이 아닌 압축된 hidden state 사용
• 기존 SSM의 문제 : time-invariant → 선택적 압축 불가
• 제안: “selectivity” ← input-dependent time-varying SSM
  • 문제: convolutionalize, 병렬 처리가 불가능 → computational efficiency↓
  • 해결방안: Parallel Scan, Kernel Fusion
    
    

연산량(FLOPs)-성능(perplexity)

Efficiency Benchmarks