[논문 리뷰] Mamba-ND: Selective State Space Modeling for Multi-Dimensional Data

Shufan Li, Harkanwar Singh외 1명|arXiv (Cornell University)|2024. 02. 08.

Neural Networks and Applications인용 수 4

한 줄 요약

Mamba-ND는 1차원 시퀀스에 특화된 Mamba 상태공간 모델을 다차원 데이터로 확장하기 위해 층 간에 행우선 순서로 순서를 번갈아 적용함으로써 이미지, 영상, 기상 예측 작업에서 최신 기술 수준의 성능을 달성하며, 파rameter 수를 크게 줄이고 선형 복잡도를 유지함. 이는 ViT 및 Swin-Transformer과 같은 트랜스포머 모델을 능가하면서 모델 크기를 최대 44.5%까지 감소시킴.

ABSTRACT

In recent years, Transformers have become the de-facto architecture for sequence modeling on text and a variety of multi-dimensional data, such as images and video. However, the use of self-attention layers in a Transformer incurs prohibitive compute and memory complexity that scales quadratically w.r.t. the sequence length. A recent architecture, Mamba, based on state space models has been shown to achieve comparable performance for modeling text sequences, while scaling linearly with the sequence length. In this work, we present Mamba-ND, a generalized design extending the Mamba architecture to arbitrary multi-dimensional data. Our design alternatively unravels the input data across different dimensions following row-major orderings. We provide a systematic comparison of Mamba-ND with several other alternatives, based on prior multi-dimensional extensions such as Bi-directional LSTMs and S4ND. Empirically, we show that Mamba-ND demonstrates performance competitive with the state-of-the-art on a variety of multi-dimensional benchmarks, including ImageNet-1K classification, HMDB-51 action recognition, and ERA5 weather forecasting.

연구 동기 및 목표

• 원래 1차원 시퀀스 전용으로 설계된 Mamba 상태공간 모델을 이미지, 영상, 과학 데이터 세트와 같은 다차원 데이터로 확장하는 것.
• 다차원 입력에 적용할 때 Mamba의 순차적 처리 방식이 특정 데이터 순서를 요구하는 문제를 해결하는 것.
• 이중 방향, N차원, 다중 헤드 설계를 포함한 다양한 다차원 확장 방식을 평가하고 비교하는 것.
• 복잡한 분해나 아키텍처를 도입하지 않고도 1D SSM을 고차원으로 확장하는 데 가장 효과적이고 효율적인 아키텍처를 규명하는 것.
• 단순한 층 간 번갈아 순서 전환 전략이 복잡한 다중 방향 설계를 능가하는 성능과 효율성을 달성할 수 있음을 보여주는 것.

제안 방법

• Mamba-ND는 1D Mamba 레이어를 블랙박스 컴ponent로 활용하여 연속된 레이어 간에 번갈아 가며 행우선 순서로 데이터를 평탄화하여 처리함.
• 모델은 레이어 간에 순서를 번갈아 적용함—예를 들어 행우선, 열우선, 대각선 유사 순서—이를 통해 기존 SSM 메커니즘을 수정하지 않고도 다중 방향 정보 흐름을 가능하게 함.
• Mamba의 선택적 상태공간 메커니즘을 활용하여 시퀀스 길이에 대해 선형 복잡도를 유지함으로써 자기주도 어텐션의 제곱 복잡도를 피함.
• 모델은 분해나 전용 2D/3D SSM 레이어를 필요로 하지 않으며, 대신 각 1D SSM 레이어 전에 입력 데이터의 순서를 재정렬함.
• 각 블록이 다른 순서를 적용하는 블록 수준의 아키텍처를 사용하여 다중 방향으로 공간적 및 시간적 의존성을 포착함.
• 입력 차원에 관계없이 일반화 가능하며, 데이터를 패치의 시퀀스로 간주함으로써 2D(이미지), 3D(영상), 심지어 4D(기상 데이터)에 적용 가능함.

Figure 1: Mamba-ND outperforms Transformers while significantly reducing the number of parameters. On ImageNet-1k, we compare against ViT [ 12 ] . On HMDB-51 [ 22 ] , we compare against Video-Swin [ 28 ] . On ERA5, we compare against Cli-ViT [ 34 ] .

실험 결과

연구 질문

• RQ1 원래 1차원 텍스트 시퀀스에 특화된 Mamba 아키텍처가 이미지 및 영상과 같은 다차원 데이터로 효과적으로 확장될 수 있는가?
• RQ2 층 간에 단순한 번갈아 순서 전환 전략이 더 복잡한 다중 방향 또는 다중 헤드 SSM 설계를 능가하는가?
• RQ3 Mamba-ND의 유효 수신 필드는 단일 방향 및 이중 방향 기준 모델과 비교해 공간적 및 시간적 의존성을 어떻게 더 잘 포착하는가?
• RQ4 분해, 레이어 순서, 깊이와 같은 아키텍처 선택 사항이 성능과 파rameter 효율성에 어떤 영향을 미치는가?
• RQ5 단지 레이어 간 입력 시퀀스 순서를 재정렬하는 최소한의 수정만으로 최신 기술 수준의 트랜스포머와 경쟁 가능한 성능을 달성할 수 있는가?

주요 결과

• ImageNet-1K에서 Mamba-ND는 ViT보다 1.5% 높은 상위 1위 정확도를 기록하면서 파rameter 수를 20.7% 감소시킴.
• HMDB-51 행동 인식 벤치마크에서 Mamba-ND는 3D Swin-Transformer보다 0.9% 높은 정확도를 기록하면서 파rameter 수를 39.0% 감소시킴.
• ERA5 72시간 기상 예측 작업에서 Mamba-ND는 Cli-ViT보다 0.7점 높은 이상치 상관계수(ACC)를 기록하며 파rameter 수를 44.5% 감소시킴.
• 복잡한 설계인 2D+3D 분해, 다중 헤드 SSM, 이중 방향 SSM보다는 단순한 행우선 순서 번갈아 적용 전략이 성능 면에서 뛰어남.
• 유효 수신 필드 시각화 결과 Mamba-ND는 단일 또는 이중 방향 모델보다 더 균일하고 전역적인 민감도 패턴을 보이며, 이는 뛰어난 성능을 설명함.
• 제거 실험 결과 단순한 설계, 예를 들어 번갈아 순서 전략이 더 깊거나 넓은 다중 방향 아키텍처보다 효과적임을 확인함으로써 SSM에서 깊이가 너비보다 더 중요함을 뒷받침함.

Figure 2: Architecture of Mamba-ND: In this figure we visualize Mamba-3D. Given 3D input data, we patchify it into $L$ patches. During this process, we maintain the original 3D structure of the input. This sequence is then passed through $K$ Mamba-ND blocks, each of which consists of a chain of 1-D