Vision Mamba: Efficient Visual Representation Learning with BidirectionalState Space Model(2024)

1. 들어가며

  먼저 백본이 뭔지 알아보겠습니다. 딥러닝에서 백본은 입력 데이터(예를들면 이미지)를 받아서 핵심적인 특징(피처)을 추출해 주는 기본적인 네트워크를 말합니다. 이미지 분류를 할 때는 백본이 뽑아낸 특징을 바탕으로 확률을 계산하고, 객체 탐지나 분할 같은 작업에서는 백본에 탐지 헤드나 분할 헤드를 붙여서 객체 위치나 픽셀별 레이블을 예측하죠. 백본 하나가 바뀌면 전체 작업 성능이 좌우될 정도로 중요합니다.
​
   컴퓨터 비전 분야에서 대표 백본의 흐름을 보면 판이 몇 번 크게 바뀌었습니다. 처음에는 ResNet 같은 CNN 계열이 이미지 분류부터 객체 탐지, 의미 분할까지 거의 모든 일을 떠맡았고 최근에는 ViT나 DeiT 같은 비전 트랜스포머가 등장하면서 이미지를 패치 단위로 쪼개서 토큰 시퀀스로 보고 셀프 어텐션(self attention, 자기 주의 메커니즘)을 돌리는 방식이 주류가 됐습니다. 하지만 셀프 어텐션에는 근본적인 한계가 있습니다. 입력 이미지 해상도가 올라가면서 패치 토큰 수가 폭발적으로 늘어나면 셀프 어텐션 연산량과 메모리 사용량이 토큰 길이의 제곱에 비례해서 급격히 증가하기 때문입니다. 해상도 512x512에서는 그나마 버티지만, 1024x1024나 1248x1248 같은 고해상도로 가면 DeiT 같은 모델은 속도가 급격히 떨어지고 GPU 메모리가 터질 지경이 됩니다. 실제로 논문 실험에서도 DeiT-Ti는 1248x1248 해상도에서 메모리 부족(Out of Memory, OOM)으로 실패하지만 비전 맘바는 여유롭게 처리한다는 결과가 있습니다.
​
  이런 상황에서 자연어 처리 쪽에서는 맘바(Mamba)라는 상태공간모델(SSM, State Space Model) 기반 구조가 트랜스포머 대체주로 떠올랐습니다. 맘바의 핵심 장점은 긴 시퀀스를 다룰 때도 연산량이 시퀀스 길이에 선형으로만 늘어나고 하드웨어 최적화(입출력 재구성, 재계산 전략 등)를 통해 실제 GPU에서 효율이 매우 좋다는 점입니다. 언어 모델링에서 트랜스포머와 맞먹는 성능을 내는데다 훨씬 빠르게 돌아가니까요.
​
  비전 맘바 논문이 정확히 여기서 출발합니다. 이미지도 패치 단위로 1차원 시퀀스로 볼 수 있는데, 굳이 셀프 어텐션을 쓰지 않고 맘바 스타일의 SSM만으로 백본 전체를 다시 설계하면 어떨까 하는 발상입니다. 양방향 SSM으로 전역 문맥을 잡고 위치 임베딩(position embedding)으로 공간 정보를 보강하면 비전 트랜스포머와 비슷하거나 더 나은 성능을 낼 수 있지 않을까 하는 것입니다.
​

  저자들은 이 가설을 꽤 과감하게 검증합니다. 이미지넷(ImageNet-1K) 분류, COCO 객체 탐지/인스턴스 분할, ADE20K 의미 분할 등 주요 벤치마크에서 비슷한 크기의 DeiT보다 성능은 더 높고, 특히 고해상도 입력에서는 추론 속도가 2.8배 빠르고 GPU 메모리는 86.8%까지 절감된다고 보고합니다. 단순히 새로운 백본을 내놓은 게 아니라, 비전 백본의 기본 전제를 정면으로 뒤흔드는 작업이라 흥미롭습니다. 과연 다음 세대 백본이 트랜스포머에서 SSM으로 넘어갈 수 있을지 보여주는 중요한 첫걸음입니다.

---

2. SSM과 맘바

  비전 맘바를 이해하려면 먼저 상태공간모델(SSM)과 맘바가 뭔지 알아야 합니다. 이게 비전 맘바의 핵심 엔진이니까요. SSM은 원래 제어공학이나 신호처리에서 나온 개념인데, 딥러닝에서는 긴 시퀀스(예: 문장이나 이미지 패치들)를 효율적으로 처리하는 방법으로 쓰입니다. 기본 아이디어는 시퀀스 입력 x(t)를 숨겨진 상태 h(t)를 통해 출력 y(t)로 변환하는 것입니다.. 수식으로 쓰면 h′(t)=Ah(t)+Bx(t), y(t)=Ch(t) 이런 식입니다. 여기서 A, B, C는 매개변수고, 연속시간 시스템을 이산시간으로 바꿔서 실제 네트워크에 적용합니다. 초기 SSM(S4 등)은 이걸 컨볼루션처럼 빠르게 계산했지만, 맘바는 더 발전시켰습니다.

​

  맘바의 특별한 점은 세 가지입니다. 첫째, 선택성(selectivity). 매개변수 A, B, C를 입력에 따라 동적으로 바꿔서, 중요한 부분에만 집중하게 합니다. 셀프 어텐션처럼 데이터 의존적인 전역 문맥을 잡으면서도 계산량은 선형으로 유지하죠. 둘째, 하드웨어 최적화. GPU 메모리 구조(HBM과 SRAM)를 고려해서, 매개변수를 미리 SRAM으로 올려놓고 재사용합니다. 입출력 횟수를 O(BMEN)에서 O(BME+EN)으로 줄여서 실제 속도를 엄청 높입니다. 논문에서도 이게 고해상도 이미지에서 DeiT보다 2.8배 빠른 이유라고 설명합니다.​ 셋째, 재계산(recomputation) 전략. 훈련 중 중간 상태를 저장하지 않고 역전파 때 다시 계산해서 메모리를 아낍니다. 긴 시퀀스에서 메모리 폭발을 막아줍니다.

 

  SSM과 맘바의 계산 복잡도는 셀프 어텐션과 비교하면 확연합니다. 시퀀스 길이 M, 차원 D일 때, 셀프 어텐션은 O(M2D+MD2) SSM(맘바)은 O(MDN) 입니다. 여기서 N은 고정된 작은 값(논문에서 16)이라, M이 크면(고해상도 이미지처럼) SSM이 압도적으로 우위입니다.​ 비전 쪽에서 SSM을 써본 적은 있지만 그 대부분 CNN이나 트랜스포머와 섞은 하이브리드 형태였습니다. 비전 맘바는 처음으로 순수 SSM만으로 제네릭 백본을 만들고자 했습니다. 이미지를 패치로 쪼개서 1D 시퀀스로 보고 단방향만으로는 공간 정보가 부족하니 양방향으로 스캔하는 식으로 비전 특성을 맞췄습니다. 이게 바로 Vim 블록의 핵심입니다.

---

3. 비전 맘바(Vim) 구조 자세히 보기

  비전 맘바는 이미지를 잘게 쪼갠 다음 그 쪼개진 조각들을 하나의 긴 시퀀스로 보고 상태공간모델로 처리하는 구조입니다. 겉으로 보면 비전 트랜스포머와 비슷하지만, 안쪽에서 셀프 어텐션 대신 SSM을 쓰는 점이 가장 큰 차이입니다.​

  입력 이미지는 먼저 작은 패치로 잘게 나눕니다. 논문에서는 보통 한 패치 크기를 16×16으로 잡고, 이미지를 이 패치들로 채운 뒤 각 패치를 하나의 벡터로 평평하게 펴서 이어 붙입니다. 이렇게 하면 원래 2차원 구조였던 이미지가 한 줄로 길게 늘어선 패치 시퀀스로 바뀝니다. 이 벡터들은 선형 변환을 거쳐 공통 차원의 토큰으로 변환되고, 여기에 위치 정보를 담는 위치 임베딩을 더해 줍니다. 그리고 이미지 전체를 대표하는 정보를 모으기 위해 클래스 토큰을 하나 끼워 넣습니다.​

  이제 이 토큰 시퀀스를 여러 층의 Vim 블록에 통과시킵니다. Vim 블록 안에서는 우선 입력 시퀀스를 정규화하고, 선형 변환으로 x와 z라는 두 갈래로 나눕니다. x는 실제 내용을 처리하는 경로이고, z는 나중에 게이트 역할을 하는 경로입니다. x에는 1차원 합성곱을 적용해 주변 토큰 정보까지 섞은 뒤, 상태공간모델을 두 방향으로 돌립니다. 한 번은 앞에서 뒤로, 한 번은 뒤에서 앞으로 읽으면서 전역 문맥을 양방향으로 반영하는 식입니다. 이렇게 해서 앞 방향 출력과 뒤 방향 출력이 나오면, 여기에 z에 비선형 함수를 씌운 값을 곱해 게이트를 적용하고, 둘을 더해 하나의 출력으로 만듭니다. 마지막으로 선형 변환을 거친 뒤, 원래 입력 토큰에 더해 주는 잔차 연결까지 포함하면 Vim 블록 한 층이 완성됩니다. 이런 블록을 여러 번 반복하면 점점 더 깊고 추상적인 시각 표현이 쌓입니다.​

  마지막 층까지 통과한 뒤에는 클래스 토큰을 꺼내 정규화하고, 작은 완전연결 신경망을 붙여 최종 분류 결과를 얻습니다. 논문에서는 이 클래스 토큰을 시퀀스 맨 앞이나 맨 뒤가 아니라 중간에 배치하는 설계를 실험했는데, 가운데에 둘 때 이미지넷 분류 성능이 가장 좋았다고 보고합니다. 이는 상태공간모델이 순차적으로 정보를 흘려보내는 특성과, 이미지에서 중요한 내용이 중앙 근처에 많이 모이는 경향이 함께 작용한 결과로 볼 수 있습니다. 정리하면, 비전 맘바는 이미지를 패치 시퀀스로 본다는 점에서는 비전 트랜스포머와 같지만, 전역 문맥을 계산하는 핵심 연산을 셀프 어텐션 대신 양방향 SSM으로 바꿔 놓은 백본이라고 할 수 있습니다. 이 덕분에 토큰 수가 많아져도 계산량과 메모리가 선형에 가깝게 늘어나 고해상도 이미지에서 특히 큰 이점을 보입니다.

 

---

4. 성능과 효율: DeiT랑 직접 비교해 보면

  이 논문이 제일 자신 있게 보여주는 부분이 바로 “성능은 안 밀리면서 훨씬 가볍다”는 주장입니다. 비교 대상은 같은 비전 트랜스포머 계열인 DeiT입니다.
​
  먼저 이미지넷(ImageNet-1K) 분류 성능부터 보면, 비전 맘바의 작은 모델인 Vim-Ti와 Vim-S가 같은 급의 DeiT-Ti, DeiT-S보다 정확도가 더 높습니다. 예를 들어 Vim-Ti는 파라미터 수가 비슷한 DeiT-Ti보다 탑1 정확도가 몇 퍼센트포인트 이상 높게 나오고, 중간 크기인 Vim-S도 DeiT-S보다 성능이 더 좋습니다. 여기에 시퀀스를 더 길게 사용하는 추가 미세 조정(롱 시퀀스 파인튜닝)을 한 버전에서는 정확도가 한 번 더 상승해서, 더 큰 모델인 DeiT-B와 비슷한 수준까지 올라갑니다.
​

  분류 말고 실제 응용에 더 가까운 다운스트림 작업에서도 결과가 비슷하게 나옵니다. COCO 데이터셋으로 객체 탐지와 인스턴스 분할을 해 보면, 같은 탐지 프레임워크 안에서 백본만 바꿨을 때 Vim을 쓴 모델이 DeiT를 쓴 모델보다 박스 AP, 마스크 AP가 전반적으로 더 높습니다. 특히 중간 크기 물체나 큰 물체에서 차이가 더 뚜렷하게 나와서, 전역 문맥을 잡는 능력이 더 좋다는 쪽으로 해석할 수 있습니다. ADE20K 의미 분할 실험에서도 마찬가지로, UperNet 같은 동일한 분할 구조에서 Vim 백본이 DeiT 백본보다 mIoU가 더 높게 나옵니다.
​
  효율 면에서는 차이가 더 극적입니다. 해상도가 낮을 때, 예를 들어 512×512 정도에서는 두 모델의 속도와 메모리 차이가 아직 크지 않습니다. 그런데 해상도를 640, 738, 1024, 1248 이런 식으로 올릴수록 DeiT의 연산량과 메모리는 토큰 수의 제곱에 가깝게 늘어나고, 비전 맘바는 시퀀스 길이에 거의 선형으로만 늘어납니다. 논문에서 보고한 수치를 보면, 1248×1248 해상도에서 특징을 추출하는 배치 추론 기준으로 Vim-Ti는 DeiT-Ti보다 약 2.8배 정도 빠르고, GPU 메모리는 86.8퍼센트까지 덜 쓰는 것으로 나옵니다. DeiT는 이 해상도에서 아예 메모리 부족으로 실패하는 설정도 있는데, Vim은 끝까지 버텨 주는 식입니다.
​
  이걸 계산 복잡도 관점에서 다시 보면 더 이해가 쉽습니다. 셀프 어텐션은 토큰 수를 M, 차원을 D라고 했을 때, 대략 M²에 비례해서 계산량이 늘어나는데, 상태공간모델은 M에 비례하는 정도로만 늘어납니다. 비전 맘바는 바로 이 차이를 실제 고해상도 이미지에서 실험으로 보여 준 셈입니다. 그래서 논문에서는 “트랜스포머 방식의 표현력을 그대로 가져가면서, 고해상도에서도 실용적으로 돌릴 수 있는 다음 세대 백본 후보”라는 식으로 비전 맘바의 위치를 잡고 있습니다.

---

5. 이 논문을 어떻게 봐야 할까: 장점과 아쉬운 점

  이 논문이 제일 강하게 주장하는 장점은 두 가지로 정리할 수 있습니다. 하나는 “셀프 어텐션 없이도 비전 트랜스포머급 성능을 낼 수 있다”는 점이고, 다른 하나는 “고해상도에서 연산량과 메모리 면에서 훨씬 유리하다”는 점입니다. 실제로 이미지넷, COCO, ADE20K 같은 대표 벤치마크에서 같은 급의 DeiT보다 성능이 더 좋게 나오고, 1248×1248 같은 높은 해상도에서는 속도는 2배 이상 빠르고 메모리는 훨씬 덜 쓰는 결과를 보여 줍니다. 이 정도면 “트랜스포머 말고도 다른 길이 열렸다”고 말할 수 있는 수준입니다.
​
  또 하나 중요한 포인트는, 이 모델이 하이브리드가 아니라 “순수 상태공간 기반 비전 백본”이라는 점입니다. 기존에 SSM을 비전에 쓴 연구들은 보통 합성곱이나 트랜스포머와 섞어서 일부만 바꾸는 정도였는데, 비전 맘바는 아예 셀프 어텐션을 빼고 SSM만으로 전체 백본을 설계했습니다. 그 과정에서 단방향 SSM의 한계를 보완하려고 양방향 구조를 넣고, 위치 임베딩과 클래스 토큰 배치 방식도 여러 실험으로 다듬었습니다. 이런 설계 실험들 덕분에 “언어에서 잘 되는 Mamba를 그대로 가져온 게 아니라, 비전 도메인에 맞게 진짜로 맞춰서 튜닝했다”는 느낌을 줍니다.
​
  물론 아쉬운 점도 있습니다. 논문에서 주로 비교 대상은 같은 계열의 DeiT에 집중돼 있고, 요즘 많이 쓰이는 더 복잡한 백본들(예를 들면 다양한 변형 비전 트랜스포머나 대규모 사전학습 모델들)과의 직접 비교는 상대적으로 적습니다. 또, 상태공간 기반 구조가 실제로 얼마나 쉽게 구현되고, 프레임워크나 라이브러리에서 얼마나 잘 지원되는지도 아직은 트랜스포머만큼 검증돼 있지는 않습니다. 결국 이 논문은 “완전히 새로운 표준”이라기보다, 앞으로 더 큰 사전학습이나 다양한 응용에서 검증해 볼 만한 강력한 후보를 던져 준 단계라고 보는 게 적당해 보입니다.
​
  이런 점들을 종합하면, 비전 맘바는 “이미지를 시퀀스로 본다”는 비전 트랜스포머의 관점을 유지하면서도, 그 안의 핵심 엔진을 셀프 어텐션에서 상태공간모델로 갈아끼운 실험이라고 할 수 있습니다. 고해상도 이미지나 긴 시퀀스를 다뤄야 하는 비전·멀티모달 작업에서, 앞으로 “트랜스포머 대신 SSM을 써 보는 것”이 진지하게 고려될 수 있다는 걸 꽤 설득력 있게 보여 준 논문입니다.
​