웹사이트 검색

시각적 인식 작업에서 Mamba 메커니즘의 필요성 평가 -MambaOut

소개

Transformer는 BERT, GPT 시리즈, ViT와 같은 파워업 모델의 중추입니다. 그러나 주의 메커니즘은 2차 복잡성을 가지므로 긴 시퀀스에는 까다롭습니다. 이를 해결하기 위해 선형 복잡성을 지닌 다양한 토큰 믹서가 개발되었습니다.

최근 RNN 기반 모델은 긴 시퀀스에 대한 효율적인 훈련 및 추론으로 주목을 받았으며 대규모 언어 모델의 백본으로서의 가능성을 보여주었습니다.

이러한 기능에 영감을 받아 연구원들은 시각적 인식 작업에서 Mamba를 사용하여 Vision Mamba, VMamba, LocalMamba 및 PlainMamba와 같은 모델을 개발했습니다. 그럼에도 불구하고, 실험에 따르면 상태 공간 모델 또는 비전에 대한 SSM 기반 모델은 최첨단 컨벌루션 및 주의 기반 모델에 비해 성능이 떨어지는 것으로 나타났습니다.

이 최근 논문은 새로운 시각적 Mamba 모델 설계에 중점을 두지 않습니다. 대신에 중요한 연구 질문인 시각적 인식 작업에 Mamba가 필요한가요?를 조사합니다.

맘바란 무엇인가?

Mamba는 Carnegie Mellon University 및 Princeton University의 연구원들이 개발한 딥 러닝 아키텍처로, 특히 긴 시퀀스에 대한 변환기 모델의 한계를 해결하기 위해 설계되었습니다. S4(Structured State Space Sequence) 모델을 사용하여 연속 시간, 반복 및 컨벌루션 모델의 장점을 결합하여 긴 종속성과 불규칙하게 샘플링된 데이터를 효율적으로 처리합니다.

최근 연구자들은 ViT(Vision Transformers)가 사용되는 방식과 유사하게 컴퓨터 비전 작업에 Mamba를 적용했습니다. ViM(Vision Mamba)은 양방향 상태 공간 모델(SSM)을 활용하여 특히 고해상도 이미지에 대한 기존 Transformer의 높은 계산 요구 사항을 해결함으로써 효율성을 향상시킵니다.

맘바 아키텍처

Mamba는 입력을 기반으로 매개변수를 조정하는 고유한 선택 메커니즘을 도입하여 S4 모델을 향상시켜 시퀀스 내의 관련 정보에 집중할 수 있도록 합니다. 이 시변 프레임워크는 계산 효율성을 향상시킵니다.

Mamba는 또한 GPU와 같은 최신 하드웨어에서 효율적인 계산을 위해 하드웨어 인식 알고리즘을 사용하여 성능과 메모리 사용을 최적화합니다. 이 아키텍처는 SSM 설계를 MLP 블록과 통합하여 언어, 오디오, 유전체학을 포함한 다양한 데이터 유형에 적합합니다.

맘바 변종

  • MambaByte: 원시 바이트 시퀀스를 처리하여 토큰화 및 관련 편견을 제거하는 토큰 없는 언어 모델입니다.
  • Mamba Mixture of Experts(MOE): Mixture of Experts를 Mamba와 통합하여 Mamba 및 MOE 계층을 교대로 사용하여 효율성과 확장성을 향상시킵니다.
  • 비전 맘바(ViM): ViM은 시각적 시퀀스 인코딩을 위해 양방향 Mamba 블록을 사용하여 시각적 데이터 처리에 SSM을 적용합니다. 이는 계산 요구를 줄이고 ImageNet 분류, COCO 개체 감지 및 ADE20k 의미 체계 분할과 같은 작업에서 향상된 성능을 보여줍니다.
  • Jamba: AI21 Labs에서 개발한 Jamba는 520억 개의 매개변수와 256,000개 토큰의 컨텍스트 창을 갖춘 하이브리드 변환기이자 Mamba SSM 아키텍처입니다.

데모

모델 작업을 시작하기 전에 저장소를 복제하고 몇 가지 필요한 패키지를 설치합니다.

!pip install timm==0.6.11
!git clone https://github.com/yuweihao/MambaOut.git
!pip install gradio

또한 단계를 실행하고 MambaOut으로 추론을 수행하는 노트북에 액세스하는 데 사용할 수 있는 링크를 추가했습니다.

cd /MambaOut

아래 셀은 Gradio 웹 앱을 실행하는 데 도움이 됩니다.

!python gradio_demo/app.py

RNN과 유사한 모델과 인과적 관심

아래 그림은 메모리 관점에서 인과 주의 메커니즘과 RNN 유사 모델을 설명합니다. 여기서 _xi_는 i번째 단계의 입력 토큰을 나타냅니다.

(a) Causal Attention: 이전 토큰의 모든 키(k)와 값(v)을 메모리로 저장합니다. 메모리는 현재 토큰의 키와 값을 지속적으로 추가하여 업데이트되므로 손실이 발생하지 않습니다. 그러나 시퀀스가 길어질수록 이전 메모리를 현재 토큰과 통합하는 계산 복잡성이 증가합니다. 따라서 주의력은 짧은 시퀀스에서는 잘 작동하지만 긴 시퀀스에서는 어려움을 겪습니다.

(b) RNN 유사 모델: 이전 토큰을 메모리 역할을 하는 고정 크기 숨겨진 상태(h)로 압축합니다. 이 고정된 크기는 RNN 메모리가 본질적으로 손실이 많으며 Attention 모델의 무손실 메모리 용량과 일치할 수 없음을 의미합니다. 그럼에도 불구하고 RNN과 유사한 모델은 시퀀스 길이에 관계없이 이전 메모리를 현재 입력과 병합하는 복잡성이 일정하게 유지되므로 긴 시퀀스를 처리하는 데 탁월합니다.

Mamba는 반복적인 속성으로 인해 인과적 토큰 혼합이 필요한 작업에 특히 적합합니다. 특히 Mamba는 다음과 같은 특징을 지닌 작업에 탁월합니다.

1. 작업에는 긴 시퀀스를 처리하는 작업이 포함됩니다. 2. 해당 작업에는 인과적인 토큰 혼합이 필요합니다.

다음 질문은 시각적 인식 작업의 시퀀스가 매우 길다는 것입니다.

ImageNet의 이미지 분류의 경우 일반적인 입력 이미지 크기는 224x224이므로 패치 크기가 16x16인 토큰 196개가 생성됩니다. 이 숫자는 긴 시퀀스 작업의 임계값보다 훨씬 작으므로 ImageNet 분류는 긴 시퀀스 작업으로 간주되지 않습니다.

이미지 크기가 800x1280인 COCO의 객체 감지 및 인스턴스 분할과 ADE20K의 의미 분할을 위해(ADE20K는 150개의 의미 범주로 구성된 의미론적 분할 작업에 널리 사용되는 데이터 세트입니다. 데이터 세트에는 훈련에 20,000개의 이미지가 포함되어 있습니다. 세트 및 검증 세트에 2,000개의 이미지), 이미지 크기는 512x2048, 숫자 토큰 수는 16x16 패치 크기로 약 4,000개입니다. 4,000개의 토큰이 작은 시퀀스의 임계값을 초과하고 기본 임계값에 가깝기 때문에 COCO 감지 및 ADE20K 분할은 모두 긴 시퀀스 작업으로 간주됩니다.

MambaOut의 프레임워크

MambaOut의 전체 프레임워크

그림 (a)는 시각 인식을 위한 MambaOut의 전체 프레임워크:를 나타냅니다. MambaOut은 시각적 인식을 위해 설계되었으며 ResNet과 유사한 계층적 아키텍처를 따릅니다. 이는 _Di_로 표시되는 서로 다른 채널 크기를 갖는 4개의 스테이지로 구성됩니다. 이러한 계층 구조를 통해 모델은 여러 추상화 수준에서 시각적 정보를 처리할 수 있어 이미지의 복잡한 패턴을 인식하는 능력이 향상됩니다.

(b) 게이트 CNN 블록의 아키텍처: Gated CNN 블록은 MambaOut 프레임워크 내의 구성 요소입니다. SSM(State Space Model)을 포함하지 않는다는 점에서 Mamba 블록과 다릅니다. 두 블록 모두 정보 흐름을 조절하기 위해 게이팅 메커니즘이 있는 CNN(컨볼루션 신경망)을 사용하지만, Gated CNN 블록에 SSM이 없다는 것은 SSM을 통합한 Mamba 블록과 같이 긴 시퀀스와 시간적 종속성을 처리할 수 있는 용량이 동일하지 않음을 의미합니다. 이러한 목적을 위해.

Gated CNN과 Mamba 블록의 주요 차이점은 SSM(State Space Model)이 있다는 점입니다.

MambaOut에서는 ConvNeXt와 유사하게 7x7 커널 크기의 깊이별 컨볼루션이 Gated CNN의 토큰 믹서로 사용됩니다. ResNet과 유사하게 MambaOut은 그림과 같이 각 단계에 Gated CNN 블록을 쌓아 4단계 프레임워크를 사용하여 구축됩니다.

더 나아가기 전에 시각적 인식을 위해 Mamba를 도입해야 할 필요성에 관한 가설이 있습니다. 가설 1: ImageNet의 이미지 분류를 위해 SSM을 도입할 필요는 없습니다. 이 작업은 특성 1이나 특성 2를 충족하지 않기 때문입니다. 가설 2: 이러한 작업은 특성 2를 충족하지 못하더라도 특성 1과 일치하므로 시각적 감지 및 분할을 위한 SSM의 잠재력을 더 자세히 탐구하는 것은 여전히 가치가 있습니다.

훈련

ImageNet의 이미지 분류

  • ImageNet은 130만 개의 훈련 이미지와 50,000개의 검증 이미지를 포함하여 이미지 분류를 위한 벤치마크로 사용됩니다.
  • 훈련은 다양한 데이터 증대 기술 및 정규화 방법을 포함하여 증류 없이 DeiT 체계를 따릅니다.
  • AdamW 최적화 프로그램은 lr=배치 크기/1024 * 10^-3의 학습 속도 확장 규칙을 사용하여 학습에 사용되며 결과적으로 배치 크기 4096에서 학습 속도는 0.004가 됩니다.
  • MambaOut 모델은 PyTorch 및 timm 라이브러리를 사용하여 구현되고 TPU v3에서 교육됩니다.

ImageNet의 이미지 분류 결과

  • SSM을 통합하지 않은 MambaOut 모델은 ImageNet의 모든 모델 크기에서 시각적 Mamba 모델보다 지속적으로 성능이 뛰어납니다.
  • 예를 들어 MambaOut-Small 모델은 84.1%의 상위 1 정확도를 달성하여 MAC의 79%만 필요로 하면서 LocalVMamba-S보다 0.4% 뛰어난 성능을 발휘합니다.
  • 이러한 결과는 가설 1을 뒷받침하며, ImageNet에서 이미지 분류를 위해 SSM을 도입하는 것이 불필요함을 시사합니다.
  • Visual Mamba 모델은 현재 ImageNet의 최첨단 컨볼루션 및 어텐션 모델에 비해 상당히 뒤떨어져 있습니다.
  • 예를 들어 CAFormer-M36은 비슷한 크기의 모든 시각적 Mamba 모델보다 1% 이상의 정확도로 성능이 뛰어납니다.
  • 가설 1에 도전하기 위한 향후 연구에서는 ImageNet에서 최첨단 성능을 달성하기 위해 컨볼루션 및 SSM의 토큰 믹서를 사용하여 시각적 Mamba 모델을 개발해야 할 수도 있습니다.

COCO의 객체 감지 및 인스턴스 분할

  • COCO 2017은 객체 감지 및 인스턴스 분할을 위한 벤치마크로 사용됩니다.
  • MambaOut은 Mask R-CNN 내에서 백본으로 활용되며 ImageNet에서 사전 훈련된 가중치로 초기화됩니다.
  • 훈련은 12세대의 표준 1× 일정을 따르며 훈련 이미지의 크기는 짧은 쪽이 800픽셀, 긴 쪽이 1333픽셀을 초과하지 않도록 조정됩니다.
  • AdamW 최적화 프로그램은 학습률 0.0001, 총 배치 크기 16으로 사용되었습니다.
  • 구현은 PyTorch 및 mmDetection 라이브러리를 사용하여 수행되며 FP16 정밀도를 활용하여 교육 비용을 절감합니다.
  • 실험은 4개의 NVIDIA 4090 GPU에서 수행되었습니다.

COCO의 객체 감지 및 인스턴스 분할 결과

  • MambaOut은 COCO의 개체 감지 및 인스턴스 분할에서 일부 시각적 Mamba 모델보다 성능이 뛰어나지만 VMamba 및 LocalVMamba와 같은 최첨단 시각적 Mamba보다 여전히 뒤떨어져 있습니다.
  • 예를 들어 Mask R-CNN의 백본인 MambaOut-Tiny는 VMamba-T를 1.4APb 및 1.1APm만큼 뒤쫓습니다.
  • 이러한 성능 차이는 가설 2를 뒷받침하는 긴 순서의 시각적 작업에 Mamba를 통합하는 것의 이점을 강조합니다.
  • 그러나 Visual Mamba는 TransNeXt와 같은 최첨단 컨볼루션-어텐션-하이브리드 모델에 비해 여전히 상당한 성능 격차를 보여줍니다. Visual Mamba는 시각적 감지 작업에서 다른 최첨단 모델보다 뛰어난 성능을 발휘하여 효율성을 입증해야 합니다.

ADE20K의 의미론적 분할

  • ADE20K는 훈련 세트에 20,000개의 이미지, 검증 세트에 2,000개의 이미지가 포함된 150개의 의미 카테고리로 구성된 의미 체계 분할 작업의 벤치마크로 사용됩니다.
  • Mamba는 ImageNet의 사전 훈련된 가중치로부터 초기화를 통해 UperNet의 백본으로 활용됩니다.
  • 학습은 학습률 0.0001, 배치 크기 16, 반복 160,000회인 AdamW 최적화 프로그램을 사용하여 수행됩니다.
  • 구현은 PyTorch 및 mmsegmentation 라이브러리를 사용하여 수행되며, 4개의 NVIDIA 4090 GPU에서 수행된 실험과 함께 FP16 정밀도를 활용하여 훈련 속도를 향상시킵니다.

ADE20K의 의미론적 분할 결과

  • COCO의 객체 감지와 유사하게, ADE20K의 의미론적 분할에 대한 성능 추세는 MambaOut이 일부 시각적 Mamba 모델보다 성능이 뛰어나지만 최첨단 Mamba 모델의 결과와 일치할 수 없음을 보여줍니다.
  • 예를 들어 LocalVMamba-T는 단일 규모(SS) 및 다중 규모(MS) 평가 모두에서 MambaOut-Tiny를 0.5mIoU 능가하여 가설 2를 경험적으로 뒷받침합니다.
  • 또한 시각적 Mamba 모델은 SG-Former 및 TransNeXt와 같이 컨볼루션 및 어텐션 메커니즘을 통합하는 고급 하이브리드 모델에 비해 계속해서 눈에 띄는 성능 부족을 나타냅니다.
  • Visual Mamba는 시각적 분할 작업에서 더 강력한 성능을 달성함으로써 긴 시퀀스 모델링의 강점을 더욱 입증해야 합니다.

결론

Mamba 메커니즘은 긴 시퀀스와 자동 회귀 특성이 있는 작업에 가장 적합합니다. Mamba는 긴 시퀀스 특성에 맞는 시각적 감지 및 분할 작업의 가능성을 보여줍니다. ImageNet의 모든 시각적 Mamba 모델을 능가하지만 여전히 최첨단 시각적 Mamba 모델에 비해 뒤떨어지는 MambaOut 모델입니다.

그러나 계산 자원의 한계로 인해 본 논문에서는 시각적 작업에 대한 Mamba 개념을 검증하는 데 중점을 둡니다. 향후 연구에서는 Mamba 및 RNN 개념을 더 자세히 탐구할 수 있을 뿐만 아니라 LLM(대형 언어 모델) 및 LMM(대형 다중 모달 모델)을 위한 RNN 및 Transformer 통합을 통해 잠재적으로 이러한 영역에서 새로운 발전을 이룰 수 있습니다.

참고자료

  • 원본 연구 논문:- MambaOut: 시력을 위해 Mamba가 정말 필요한가요?
  • Vision Mamba: 양방향 상태 공간 모델을 사용한 효율적인 시각적 표현 학습
  • Mamba(딥 러닝 아키텍처)

관련기사: