최근의 기계 학습과 인공 지능 프레임워크에서 트랜스포머는 자연 언어 처리의 GPT 시리즈와 BERT, 컴퓨터 비전 태스크의 비전 트랜스포머를 포함한 다양한 도메인에서 가장 널리 사용되는 구성 요소 중 하나입니다. 트랜스포머를 모델 아키텍처에 포함하면 모델 성능이 크게 향상되지만, 트랜스포머의 주목 모듈은 시퀀스 길이와 제곱으로 증가하는 계산 복잡성을 갖습니다. 이러한 계산 복잡성을 해결하기 위해 다양한 모델이 커널화, 히스토리 메모리 압축, 토큰 믹싱 범위 제한 및 저랭크 접근 방식과 같은 다양한 전략을 탐구해 왔습니다. 최근에는 Mamba와 RWKV를 포함한 재귀 신경망과 같은 모델이 대규모 언어 모델에서 약속하는 결과를 보임으로써 주목을 받고 있습니다.

Mamba는 재귀 신경망과 같은 토큰 믹서를 갖는 상태 공간 모델 아키텍처를 갖는 모델家族입니다. 최근에 비전 태스크에 적용된 Mamba는 상태 공간 모델을 통해 Attention 메커니즘의 제곱 복잡성을 해결하기 위해 도입되었습니다. 그러나 Mamba 모델의 성능은 비전 태스크에서 전통적인 Attention 기반 모델과 최신의 컨볼루션 모델에 비해 실망스러운 결과를 보이고 있습니다.

MambaOut은 Mamba 모델의 본질을 조사하고, Mamba가 자동 회귀적이거나 긴 시퀀스를 갖는 태스크에 적합하다고 요약합니다. 그러나 대부분의 비전 태스크는 이러한 특성을 갖지 않기 때문에, MambaOut은 두 가지 가설을 제안합니다. 첫째, 이미지 분류 태스크는 자동 회귀적이거나 긴 시퀀스를 갖지 않기 때문에 상태 공간 모델이 필요하지 않습니다. 둘째, 인스턴스 분할과 의미 분할 및 객체 감지와 같은 태스크는 긴 시퀀스를 갖지만 자동 회귀적이지 않기 때문에 상태 공간 모델이 유용할 수 있습니다.

이 문서는 MambaOut 프레임워크를 깊이 있게 다루고, 그 메커니즘, 방법론, 아키텍처 및 최신 프레임워크와의 비교를 탐구합니다. 따라서 시작해 봅시다.

MambaOut: 비전을 위한 Mamba가 정말 필요할까?

기계 학습 응용과 능력의 발전으로 트랜스포머는 다양한 태스크를 위한 주류 백본으로 등장했습니다. 비전 트랜스포머, GPT 시리즈, BERT 등이 이러한 트랜스포머를 기반으로 합니다. 그러나 트랜스포머의 토큰 믹서 모듈은 시퀀스 길이와 제곱으로 증가하는 계산 복잡성을 갖습니다. 이를 해결하기 위해 다양한 토큰 믹서가 제안되었습니다. 최근에는 재귀 신경망과 같은 모델이 병렬izable 트레이닝과 효율적인 성능을 제공함으로써 주목을 받고 있습니다.

MambaOut은 Mamba 모델의 특성을 조사하고, Mamba가 자동 회귀적이거나 긴 시퀀스를 갖는 태스크에 적합하다고 요약합니다. 그러나 대부분의 비전 태스크는 이러한 특성을 갖지 않기 때문에, MambaOut은 두 가지 가설을 제안합니다. 첫째, 이미지 분류 태스크는 자동 회귀적이거나 긴 시퀀스를 갖지 않기 때문에 상태 공간 모델이 필요하지 않습니다. 둘째, 인스턴스 분할과 의미 분할 및 객체 감지와 같은 태스크는 긴 시퀀스를 갖지만 자동 회귀적이지 않기 때문에 상태 공간 모델이 유용할 수 있습니다.

Mamba에 적합한 태스크는 무엇인가?

Mamba 프레임워크의 토큰 믹서 모듈은 선택적 상태 공간 모델입니다. 이는 4개의 입력 의존 파라미터를 정의합니다. 재귀 특성은 Mamba를 다른 Attention 메커니즘과 구별합니다. 히든 상태는 고정 크기의 메모리로서 과거 정보를 저장합니다. 이는 메모리와 현재 입력을 통합하는 계산 복잡성을 일정하게 유지합니다. 반면에, Attention 레이어는 이전 토큰의 키와 값을 저장하고, 현재 토큰의 키와 값을 추가하여 메모리를 확장합니다. 이는 메모리 크기가 입력 토큰의 수에 따라 증가합니다.

상태 공간 모델의 메모리는 손실이 발생할 수 있기 때문에, Attention 메커니즘의 무손실 메모리에 비해 부족합니다. 따라서 Mamba 모델은 짧은 시퀀스를 다루는 태스크에서 약점을 보입니다. 그러나 긴 시퀀스를 다루는 태스크에서는 Mamba 모델이 효율적으로 메모리와 현재 입력을 통합할 수 있습니다.

또한, 상태 공간 모델의 재귀 특성으로 인해 Mamba 모델은 자동 회귀적 태스크에 적합합니다. 이는 다음 그림에서와 같이, 현재 토큰이 이전 토큰의 정보에만 접근할 수 있습니다.

이러한 특성으로 인해 Mamba 모델은 자동 회귀적이거나 긴 시퀀스를 갖는 태스크에 적합합니다.

시각적 인식 태스크, 인과적 토큰 믹싱 코드, 및 매우 긴 시퀀스

비전 태스크는 일반적으로 이해 태스크로 분류되며, 모델은 전체 이미지를 한 번에 볼 수 있습니다. 이는 토큰 믹싱에 대한 제약을 없애며, 추가적인 제약을 가할 경우 모델의 성능을 저하시킬 수 있습니다. 따라서, 비전 태스크에서는 완전 가시적 토큰 믹싱 모드가 적합합니다.

실험적 검증 및 결과

다음 단계는 MambaOut 프레임워크의 가설을 실험적으로 검증하는 것입니다. 다음 그림에서와 같이, Mamba 블록은 게이트드 컨볼루션 신경망 블록을 기반으로 하며, Mamba와 게이트드 CNN 블록의 메타 아키텍처는 메타포머 프레임워크의 토큰 믹서와 MLP의 통합으로 간주될 수 있습니다.

Mamba 블록은 게이트드 컨볼루션 신경망을 상태 공간 모델로 확장하며, 이는 게이트드 CNN과 Mamba 블록을 구별합니다. 또한, 실제 속도를 개선하기 위해 MambaOut 프레임워크는 부분 채널에 대한 깊이 방향 컨볼루션만 수행합니다.

이미지 분류 태스크

이미지넷은 이미지 분류 태스크의 벤치마크로 사용되며, 1,000개 이상의 클래스, 1,300만 개의 트레이닝 이미지, 50,000개의 검증 이미지를 포함합니다. 실험에 사용된 데이터 증강에는 무작위 크기 조정, 믹스업, 색상 조절, 랜덤 삭제, 컷믹스, 랜드 오그먼트가 포함됩니다. 다음 표는 Mamba 모델家族, MambaOut 모델 및 기타 Attention 기반 및 컨볼루션 모델의 성능을 요약합니다. MambaOut 프레임워크는 상태 공간 모델 없이 Mamba 모델을 능가합니다.

예를 들어, MambaOut-Small 모델은 84% 이상의 톱-1 정확도를 반환하며, 이는 가장 가까운 Mamba 경쟁자보다 0.4% 높은 결과입니다. 이는 첫 번째 가설을 강력하게 지지하며, 이미지 분류 태스크에 상태 공간 모델을 도입하는 것이 필요하지 않음을 시사합니다.

객체 감지 및 인스턴스 분할 태스크

COCO는 객체 감지 및 인스턴스 분할 태스크의 벤치마크입니다. MambaOut 프레임워크는 일부 비전 Mamba 모델을 능가하지만, 최신의 비전 Mamba 모델에는 여전히 미치지 못합니다. 이는 Mamba 모델의 장점을 강조하며, 특히 긴 시퀀스를 다루는 비전 태스크에서 유용함을 시사합니다.

최종 생각

Mamba 모델은 자동 회귀적이거나 긴 시퀀스를 갖는 태스크에 적합합니다. MambaOut 프레임워크는 두 가지 가설을 제안합니다. 첫째, 이미지 분류 태스크는 자동 회귀적이거나 긴 시퀀스를 갖지 않기 때문에 상태 공간 모델이 필요하지 않습니다. 둘째, 인스턴스 분할과 의미 분할 및 객체 감지와 같은 태스크는 긴 시퀀스를 갖지만 자동 회귀적이지 않기 때문에 상태 공간 모델이 유용할 수 있습니다. 실험 결과는 MambaOut 프레임워크의 가설을 지지하며, 이는 비전 태스크에서 Mamba 모델의 잠재력을 시사합니다.