SSM | Mamba Mixer

• Related Project: Private
• Category: Paper Review
• Date: 2024-04-04

MambaMixer: Efficient Selective State Space Models with Dual Token and Channel Selection

• url: https://arxiv.org/abs/2403.19888
• pdf: https://arxiv.org/pdf/2403.19888
• abstract: Recent advances in deep learning have mainly relied on Transformers due to their data dependency and ability to learn at scale. The attention module in these architectures, however, exhibits quadratic time and space in input size, limiting their scalability for long-sequence modeling. Despite recent attempts to design efficient and effective architecture backbone for multi-dimensional data, such as images and multivariate time series, existing models are either data independent, or fail to allow inter- and intra-dimension communication. Recently, State Space Models (SSMs), and more specifically Selective State Space Models, with efficient hardware-aware implementation, have shown promising potential for long sequence modeling. Motivated by the success of SSMs, we present MambaMixer, a new architecture with data-dependent weights that uses a dual selection mechanism across tokens and channels, called Selective Token and Channel Mixer. MambaMixer connects selective mixers using a weighted averaging mechanism, allowing layers to have direct access to early features. As a proof of concept, we design Vision MambaMixer (ViM2) and Time Series MambaMixer (TSM2) architectures based on the MambaMixer block and explore their performance in various vision and time series forecasting tasks. Our results underline the importance of selective mixing across both tokens and channels. In ImageNet classification, object detection, and semantic segmentation tasks, ViM2 achieves competitive performance with well-established vision models and outperforms SSM-based vision models. In time series forecasting, TSM2 achieves outstanding performance compared to state-of-the-art methods while demonstrating significantly improved computational cost. These results show that while Transformers, cross-channel attention, and MLPs are sufficient for good performance in time series forecasting, neither is necessary.

TL;DR

MambaMixer 및 QSMixer: 선택적 상태 공간 모델을 활용한 효율적인 시퀀스 및 채널 믹싱

1. MambaMixer 및 QSMixer를 통해 선택적 상태 공간 모델(SSM) 기반의 시퀀스 및 채널 믹싱 기법 제안
2. 다중 해상도 컨볼루션과 선택적 채널 믹싱을 사용하여 효율적이고 우수한 모델 구현
3. ImageNet, COCO 등 다양한 벤치마크에서 우수한 성능 달성

[러프하게 이해해보기]

• 트랜스포머의 경우, $L$ 길이의 인퍼런스를 위해서 $L^2$의 연산이 필요해 연산량이 많은 것이 문제였는데 이를 해결하기 위해 시간 혹은 상태의 개념을 도입해 각 시퀀스별로 중요한 토큰을 선별하여 다르게 연산할 수 있는 제어 이론의 SSM을 사용하려고 시도했으나 기존 SSM은 채널 간 정보 흐름을 고려하지 않아 대규모 네트워크에서 안정성 문제가 있었습니다.
• 따라서, 데이터 종속적 파라미터를 사용하여 각 입력에 따라 모델의 동작을 다르게 해서 연산량을 줄이면서 안전성을 확보하고자 하였고, 꽤 좋은 성과를 냈다고 합니다.

(사견) 그 외 LLM에서 Mamba와 Transformer 블록을 혼합하는 하이브리드 블록으로 큰 성과를 낸 Jamba 모델이 등장했지만, 1) 꽤 귀찮은 변형 및 2) 호환성 문제(파인튜닝 및 경량화, 빅테크들의 모델들과 호환성이 떨어짐 등) + 3) CNN, Attention 등장 때와 같이 획기적인 퍼포먼스 증가로 볼 수 있는가? + 4) Perplexity의 감소 문제 등으로 AGI 모델로 사용하기 어려움 + 5) 등장한 지 꽤 됐음에도 아직 상용화까지 간 모델의 부재 + 6) 빅테크 혹은 주요 검증이 추가로 등장하지 않아 리스크가 있을 수 있어 LLM 튜닝에 드는 비용을 투입하기 애매 등으로 추가 검증을 기다리는…

[서론]

최근 트랜스포머(Transformer)는 언어 모델링, 컴퓨터 비전, 시계열 분석 등 다양한 분야에서 혁신적인 성과를 이루어냈습니다. 트랜스포머의 핵심은 어텐션 메커니즘으로, 이는 데이터의 의존성을 모델링하여 unseen data와 tasks에 대해 일반화할 수 있게 합니다. 그러나 트랜스포머는 긴 시퀀스를 처리할 때 시간 및 공간 복잡도가 이차적으로 증가하는 문제가 있습니다. 이를 해결하기 위해 선택적 상태 공간 모델(SSM)을 적용한 MambaMixer와 QSMixer가 제안되었습니다.

기존 트랜스포머의 문제점

기존 트랜스포머는 모든 토큰에 동일한 연산을 수행하며, 이는 다음과 같은 문제점을 야기합니다.

1. 높은 계산 비용: 모든 토큰에 대해 동일한 양의 계산을 수행하므로, 불필요한 계산 자원이 낭비됩니다.
2. 확장성 문제: 긴 시퀀스를 처리할 때 시간 및 공간 복잡도가 이차적으로 증가하여, 긴 컨텍스트를 효과적으로 학습하기 어렵습니다.

선택적 상태 공간 모델(SSM)과 트랜스포머의 차이점

1. 계산 비용의 최적화
   • 기존 트랜스포머: 모든 토큰에 대해 동일한 계산을 수행합니다. 이는 불필요한 계산 자원을 낭비하게 됩니다.
   • MambaMixer와 QSMixer: 토큰별로 필요한 연산만 수행합니다. 선택적 SSM을 통해 중요하지 않은 토큰에 대해 계산을 줄일 수 있습니다. 이는 계산 비용을 최적화하여 더 효율적인 모델을 만듭니다.
2. 데이터 종속적 파라미터
   • 기존 트랜스포머: 모든 입력에 대해 동일한 파라미터를 사용합니다.
   • MambaMixer와 QSMixer: 데이터 종속적 파라미터를 사용하여 각 입력에 따라 모델의 동작을 다르게 합니다. 이를 통해 각 토큰의 중요도에 따라 연산을 조절할 수 있습니다.

선택적 상태 공간 모델(SSM) 기반 MambaMixer와 QSMixer의 도입

MambaMixer와 QSMixer는 선택적 상태 공간 모델(SSM)을 도입하여 트랜스포머의 이런 문제를 해결하고자 합니다. SSM은 연속 시간 상태 공간 모델을 이산 시간 상태 공간 모델로 변환하여 사용합니다.

[선행 연구]

최근 구조화된 상태 공간 모델(SSM)은 시퀀스 모델링에서 유망한 아키텍처로 떠오르고 있습니다. 그러나 (핵심) 기존 SSM은 채널 간 정보 흐름을 고려하지 않아 대규모 네트워크에서 안정성 문제가 발생하며, 글로벌 정보 모델링 능력이 제한됩니다.

본 논문에서는 MambaMixer 및 QSMixer라는 효율적인 선택적 상태 공간 모델을 제안합니다. 이 모델들은 시퀀스 및 채널 믹싱을 통해 정보를 효율적으로 선택하고 필터링합니다. 이를 통해 성능을 유지하면서도 연산 비용을 절감할 수 있습니다.

[방법]

본 논문에서는 ImageNet, COCO 등 다양한 데이터셋을 사용하여 제안된 모델의 성능을 평가합니다.

• 선택적 상태 공간 모델(SSM)

  SSM은 입력 시퀀스 $x(t) \in \mathbb{R}^L$을 응답 시퀀스 $y(t) \in \mathbb{R}^L$로 매핑하는 선형 시불변 시스템입니다. 이를 위해 잠재 상태 $h(t) \in \mathbb{R}^{N \times L}$, 파라미터 $A \in \mathbb{R}^{N \times N}$, $B \in \mathbb{R}^{N \times 1}$, $C \in \mathbb{R}^{1 \times N}$를 사용합니다.

  다음은 SSM의 기본 수식입니다. \(h'(t) = A h(t) + B x(t)\) \(y(t) = C h(t)\)

  이 시스템을 딥러닝 설정에 맞추기 위해 이산화합니다. \(h_t = \bar{A} h_{t-1} + \bar{B} x_t\) \(y_t = C h_t\) 위 식에서 $\bar{A} = \exp(\Delta A)$ 및 $\bar{B} = (\Delta A)^{-1}\exp(\Delta A) - I) \Delta B$입니다.

• MambaMixer 아키텍처

  MambaMixer는 시퀀스 및 채널 믹싱을 동시에 수행하여 정보를 효율적으로 결합합니다. 선택적 토큰 믹서는 시퀀스 축을 따라 정보를 결합하고, 다중 해상도 컨볼루션과 S6 블록을 사용합니다.

  \[\bar{x} = \sigma (\text{Conv} (\text{Linear} (x)))\] \[\Delta x = \text{Softplus} (\text{Linear}_\Delta (x))\] \[y_{\text{Token}} = \text{SSM} \bar{A}, \bar{B}_x, C_x, \Delta x) \otimes \bar{x}_{\text{Gate}}\]

• 선택적 토큰 믹서

  선택적 토큰 믹서는 시퀀스 축을 따라 정보를 결합합니다. 이를 위해 1D 컨볼루션과 S6 블록을 사용합니다.

  \[\bar{x} = \sigma (\text{Conv} (\text{Linear} (x)))\] \[\Delta x = \text{Softplus} (\text{Linear}_\Delta (x))\] \[y_{\text{Token}} = \text{SSM} \bar{A}, \bar{B}_x, C_x, \Delta x) \otimes \bar{x}_{\text{Gate}}\]

• 선택적 채널 믹서

  선택적 채널 믹서는 각 채널에 대해 독립적인 SSM을 사용하여 채널 간 상호 의존성을 고려합니다. 이를 통해 특정 토큰에 대해 관련 없는 채널을 필터링하고, 모델 축을 따라 정보를 결합합니다. 예를 들어:

  \[\tilde{Y} = Q^{(1)}_{\bar{A}, \bar{B}, \bar{C}} Y^T = Q^{(2)}_{\bar{A}', \bar{B}', \bar{C}'} \left( X \otimes W_1 \text{Conv} \left( \big|_{i=1}^T \text{Conv}_{k_i}(X) \right) \right)\]

• QSMixer: 효율적인 아키텍처

  QSMixer는 준분리 행렬(quasi-separable matrix)을 사용하여 시퀀스 및 채널 믹싱을 수행합니다. 준분리 행렬은 효율적인 행렬-벡터 곱셈 알고리즘을 사용하여 선형 시간 훈련을 가능하게 합니다.

  \(\tilde{Y} = Q^{(1)}_{\bar{A}, \bar{B}, \bar{C}} Y^T = Q^{(2)}_{\bar{A}', \bar{B}', \bar{C}'} \left( X \otimes W_1 \text{Conv} \left( \big|_{i=1}^T \text{Conv}_{k_i}(X) \right) \right)\) 위 식에서 $W_1$과 $W_2$는 학습 가능한 파라미터입니다.

[실험 및 결과]

• 이미지 분류 및 객체 검출 ViM2 및 ViQS 모델은 ImageNet 분류, COCO 객체 검출 등 다양한 비전 태스크에서 우수한 성능을 보였습니다. 특히, 다중 해상도 컨볼루션과 선택적 채널 믹싱을 통해 글로벌 수용 영역을 향상시켰습니다.

• 시계열 예측 시계열 예측에서도 MambaMixer 및 QSMixer는 기존 SSM 기반 모델보다 개선된 성능을 보였습니다. 이는 선택적 채널 믹싱을 통해 관련 없는 정보를 효율적으로 필터링한 결과입니다.

• 성능 분석

  • ImageNet 분류 성능: ViM2와 ViQS는 최고의 혼합기 아키텍처 대비 평균 3.55%, 효율적인 트랜스포머 기반 모델 대비 평균 1.2% 향상된 성능을 보였습니다.
  • 채널 믹싱의 중요성: 데이터 종속적 채널 믹싱을 사용한 모델은 더 개선된 성능을 보였습니다. 이는 각 토큰에 대해 관련 없는 채널을 효율적으로 필터링한 결과입니다.

[결론]

본 논문에서는 MambaMixer 및 QSMixer를 통해 선택적 상태 공간 모델의 효율적이고 우수한 시퀀스 및 채널 믹싱 기법을 제안하였습니다. 다양한 비전 및 시계열 태스크에서 우수한 성능을 보였으며, 이는 선택적 채널 믹싱과 다중 해상도 컨볼루션의 조합 덕분입니다. 본 연구는 트랜스포머 모델의 효율성을 높이기 위한 새로운 방법을 제시하였으며, 향후 다양한 응용 분야에서의 활용 가능성을 제시합니다.

[참고자료 1] 이산화 과정 및 선택적 상태 공간 모델(SSM)에 대한 수학적 논증

• 이산화 과정

  연속적인 상태 공간 모델(SSM)을 이산화하여 딥러닝에 적용하기 위해 다음과 같은 수식을 사용합니다. 연속 시스템은 다음과 같이 정의됩니다.

  \(h'(t) = A h(t) + B x(t)\) \(y(t) = C h(t)\)

  위 식에서 $h(t)$는 시간 $t$에서의 잠재 상태, $x(t)$는 입력, $y(t)$는 출력, $A$, $B$, $C$는 시스템의 파라미터입니다. 이 시스템을 이산화하기 위해 샘플링 시간 간격 $\Delta$를 도입하고, $t$에서의 상태 $h(t)$와 $t + \Delta$에서의 상태 $h(t + \Delta)$ 간의 관계를 구합니다.

  이산화는 제로-오더 홀드(Zero-order Hold) 방법을 사용하여 이루어집니다. 이를 통해 연속 시스템을 이산 시스템으로 변환할 수 있습니다. 이산 시스템은 다음과 같이 표현됩니다.

  \(h_t = \bar{A} h_{t-1} + \bar{B} x_t\) \(y_t = C h_t\)

  위 식에서 $\bar{A}$와 $\bar{B}$는 연속 시스템의 파라미터 $A$와 $B$를 기반으로 계산된 이산 시스템의 파라미터입니다. $\bar{A}$와 $\bar{B}$는 다음과 같이 정의됩니다.

  \(\bar{A} = \exp(\Delta A)\) \(\bar{B} = (\Delta A)^{-1} \exp(\Delta A) - I) \Delta B\)

  위 식에서 $\exp(\Delta A)$는 매트릭스 지수 함수(Matrix Exponential)로, $\Delta A$의 지수입니다. 이는 $\Delta A$의 무한 급수로 표현될 수 있습니다.

  \[\exp(\Delta A) = I + \Delta A + \frac{(\Delta A)^2}{2!} + \frac{(\Delta A)^3}{3!} + \cdots\]

  이 식은 $\bar{A}$와 $\bar{B}$를 통해 이산 시스템에서 연속 시스템의 동작을 근사할 수 있게 해줍니다.

• 선택적 상태 공간 모델(SSM)

  선택적 상태 공간 모델(SSM)은 SSM의 파라미터를 시간에 따라 다르게 하여(즉, 데이터 종속적으로) 관련 있는 컨텍스트를 효과적으로 선택하는 방법입니다. 이를 위해 각 시간 단계에서 입력 $x_t$에 따라 파라미터를 조정합니다.

  \(\bar{B}_t = \text{Linear}_B(x_t)\) \(\bar{C}_t = \text{Linear}_C(x_t)\) \(\Delta_t = \text{Softplus}(\text{Linear}_\Delta(x_t))\)

  위 식에서 $\text{Linear}(\cdot)$는 선형 변환 함수, $\text{Softplus}(\cdot) = \log(1 + \exp(\cdot))$는 활성화 함수입니다. 이 데이터 종속적 파라미터는 모델이 입력 시퀀스의 각 부분에 대해 다른 동작을 하도록 하여 컨텍스트 압축을 효과적으로 수행할 수 있게 합니다.

  이 방법을 사용하면 SSM의 선형 회귀를 다음과 같이 재구성할 수 있습니다.

  \(h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t\) \(y_t = C_t h_t\)

  이 시스템에서 $\bar{A}_t$, $\bar{B}_t$, $C_t$는 각각 시간 $t$에서의 파라미터입니다.

• 선택적 토큰 믹서와 선택적 채널 믹서

  • 선택적 토큰 믹서

    선택적 토큰 믹서는 시퀀스 축을 따라 정보를 결합합니다. 이를 위해 1D 컨볼루션과 S6 블록을 사용합니다. 이 과정은 다음과 같이 정의됩니다.

    \(\bar{x} = \sigma (\text{Conv} (\text{Linear} (x)))\) \(\Delta x = \text{Softplus} (\text{Linear}_\Delta (x))\) \(y_{\text{Token}} = \text{SSM} \bar{A}, \bar{B}_x, C_x, \Delta x) \otimes \bar{x}_{\text{Gate}}\)

  • 선택적 채널 믹서

    선택적 채널 믹서는 각 채널에 대해 독립적인 SSM을 사용하여 채널 간 상호 의존성을 고려합니다. 이를 통해 특정 토큰에 대해 관련 없는 채널을 필터링하고, 모델 축을 따라 정보를 결합합니다. 비선형적으로 훈련 가능한 SSM을 사용하여 채널 믹싱의 표현력을 높입니다. 선택적 채널 믹서의 효율성을 높이기 위해 준분리 행렬(quasi-separable matrix)을 사용합니다.

    \[\tilde{Y} = Q^{(1)}_{\bar{A}, \bar{B}, \bar{C}} Y^T = Q^{(2)}_{\bar{A}', \bar{B}', \bar{C}'} \left( X \otimes W_1 \text{Conv} \left( \big|_{i=1}^T \text{Conv}_{k_i}(X) \right) \right)\]

    위 식에서 $W_1$과 $W_2$는 학습 가능한 파라미터입니다. 준분리 행렬은 효율적인 행렬-벡터 곱셈 알고리즘을 사용하여 선형 시간 훈련을 가능하게 합니다.

[참고자료 2] 이산화 과정 및 선택적 상태 공간 모델(SSM) 상세 설명

• 기본 개념

  연속 시간 시스템에서 이산 시간 시스템으로의 변환은 실제로 연속적인 입력 신호를 이산적인 형태로 샘플링하는 과정입니다. 이를 통해 디지털 컴퓨터에서 처리할 수 있는 형태로 변환할 수 있습니다. 제로 오더 홀드(Zero-Order Hold)는 이산화 과정에서 가장 기본적인 방법 중 하나입니다.

  • 연속 시간 상태 공간 모델 연속 시간 상태 공간 모델은 다음과 같이 정의됩니다.

  \(h'(t) = A h(t) + B x(t)\) \(y(t) = C h(t)\)

  • $h(t)$: 시간 $t$에서의 상태 벡터
  • $x(t)$: 시간 $t$에서의 입력 벡터
  • $y(t)$: 시간 $t$에서의 출력 벡터
  • $A$: 상태 천이 행렬
  • $B$: 입력 행렬
  • $C$: 출력 행렬

• 제로 오더 홀드(Zero-Order Hold)란?

  제로 오더 홀드는 연속 신호를 이산 신호로 변환할 때, 각 샘플링 간격 동안 신호 값을 일정하게 유지하는 방법입니다. 즉, 현재 샘플링 지점에서 다음 샘플링 지점까지 신호 값이 변하지 않는다고 가정합니다.

• 이산화 과정

  연속 시간 상태 공간 모델을 이산화하는 과정은 다음과 같이 이루어집니다.

  1. 샘플링 시간 간격 $\Delta t$를 선택합니다.
  2. 연속 시간 상태 방정식을 이산 시간 상태 방정식으로 변환합니다.

  연속 시간 모델에서 $t$와 $t + \Delta t$ 사이의 상태 변화는 다음과 같이 근사할 수 있습니다.

  \[h(t + \Delta t) \approx h(t) + \Delta t \cdot h'(t)\]

  이 식을 상태 방정식에 대입하면,

  \[h(t + \Delta t) \approx h(t) + \Delta t \cdot (A h(t) + B x(t))\]

  이를 다시 정리하면,

  \[h(t + \Delta t) \approx (I + \Delta t A) h(t) + \Delta t B x(t)\]

  위 식에서 $I$는 단위 행렬로, 이산 시간 모델의 파라미터 $\bar{A}$와 $\bar{B}$를 정의할 수 있습니다.

  \(\bar{A} = I + \Delta t A\) \(\bar{B} = \Delta t B\)

  그러나, 보다 정확한 방법으로는 매트릭스 지수 함수를 사용하는 것입니다. 이를 통해 다음과 같이 이산 시간 모델의 파라미터를 정의할 수 있습니다.

  \(\bar{A} = \exp(A \Delta t)\) \(\bar{B} = \left( \int_{0}^{\Delta t} \exp(A \tau) d\tau \right) B\)

  매트릭스 지수 함수 $\exp(A \Delta t)$는 다음과 같은 무한 급수로 정의됩니다.

  \[\exp(A \Delta t) = I + A \Delta t + \frac{(A \Delta t)^2}{2!} + \frac{(A \Delta t)^3}{3!} + \cdots\]

  이를 통해 $\bar{A}$를 계산할 수 있습니다. $\bar{B}$는 다음과 같이 계산됩니다.

  \[\bar{B} = (A \Delta t)^{-1} \exp(A \Delta t) - I) B\]

• 이산 시간 상태 공간 모델

  최종적으로 이산화된 상태 공간 모델은 다음과 같이 표현됩니다.

  \(h_{t+1} = \bar{A} h_t + \bar{B} x_t\) \(y_t = C h_t\)

  • $h_t$: 이산 시간 $t$에서의 상태 벡터
  • $x_t$: 이산 시간 $t$에서의 입력 벡터
  • $y_t$: 이산 시간 $t$에서의 출력 벡터
  • $\bar{A}$: 이산 시간 상태 천이 행렬
  • $\bar{B}$: 이산 시간 입력 행렬
  • $C$: 출력 행렬 (연속 시간 모델과 동일)

  이로써 연속 시간 상태 공간 모델을 이산 시간 상태 공간 모델로 변환할 수 있습니다. 이 과정은 디지털 컴퓨터에서 연속적인 물리 시스템을 시뮬레이션하거나 제어하는 데 유용합니다.

  • 예시: 단순한 시스템

    간단한 예시로서, 1차 시스템을 고려해 봅시다. 이 시스템의 상태 천이 행렬 $A$와 입력 행렬 $B$는 다음과 같습니다. \(A = \begin{pmatrix} -1 \end{pmatrix}, \quad B = \begin{pmatrix} 1 \end{pmatrix}\)

    샘플링 시간 간격 $\Delta t = 1$초로 가정하면, 이산화된 파라미터 $\bar{A}$와 $\bar{B}$는 다음과 같이 계산됩니다.

    \[\bar{A} = \exp(A \Delta t) = \exp(-1) = 0.3679\] \[\bar{B} = \exp(A \Delta t) - I) A^{-1} B = (0.3679 - 1) (-1)^{-1} = 0.6321\]

    이산화된 상태 방정식은 다음과 같습니다.

    \[h_{t+1} = 0.3679 h_t + 0.6321 x_t\] \[y_t = C h_t\]

• 선택적 상태 공간 모델 (SSM)

  선택적 상태 공간 모델은 이산화된 상태 공간 모델을 기반으로 합니다. 그러나 기존 SSM은 시간 불변적(즉, 데이터 독립적) 파라미터를 사용하여 모든 입력에 대해 동일한 파라미터를 사용합니다. 이를 개선하기 위해, 선택적 SSM은 데이터 종속적 파라미터를 도입합니다. 이는 각 시간 단계에서 입력 $x_t$에 따라 파라미터가 변하는 것을 의미합니다.

  \(\bar{B}_t = \text{Linear}_B(x_t)\) \(\bar{C}_t = \text{Linear}_C(x_t)\) \(\Delta_t = \text{Softplus}(\text{Linear}_\Delta(x_t))\)

  이를 통해 모델은 입력 시퀀스의 각 부분에 대해 다른 동작을 수행할 수 있습니다. 선택적 SSM의 최종 형태는 다음과 같습니다.

  \(h_{t+1} = \bar{A}_t h_t + \bar{B}_t x_t\) \(y_t = C_t h_t\)

  위 식에서 $\bar{A}_t$, $\bar{B}_t$, $C_t$는 각각 시간 $t$에서의 데이터 종속적 파라미터입니다. 이로써 선택적 상태 공간 모델은 시간에 따라 변하는 파라미터를 사용하여 컨텍스트를 효과적으로 압축하고, 데이터의 특성에 따라 동작할 수 있는 모델을 구성할 수 있습니다.