Topic #28: MoM (Mixture-of-Mamba)은 무엇인가?

MoM은 Mixture-of-Mamba 블록을 기반으로 구축되는데, 이 블록은 서로 다른 유형의 입력 데이터에 대해서 별도의 처리 규칙을 적용하면서도 의미가 있는 경우에는 공통의 구성 요소를 공유합니다.

‘Modality-Aware Sparsity’ 메커니즘은 다이나믹한 라우팅 시스템처럼 작동합니다:

1. MoM은 ‘모달리티 마스크’라는 걸 사용해서 서로 다른 유형의 입력 토큰 (텍스트, 이미지 또는 음성)을 구별합니다.

2. 파악된 ‘모달리티’를 기준으로 해서, 적절한 가중치 셋을 활성화합니다. 이 과정이 바로 위에서 이야기한 ’모달리티별 매개변수화 (Modality-Specific Parameterization)’라는 과정이구요.

3. MoM은 서로 다른 모달리티 토큰을 동시에 처리해서, 트레이닝과 추론이 모두 효율적으로 진행될 수 있게 합니다.

‘모달리티를 인식’한다는게 어떤 의미인지 좀 더 자세히 알아보면, MoM은 그 처리 과정에서 Mamba 모델의 중요한 계층들을 멀티모달 데이터 처리를 위해서 조금씩 변경하는데요:

• 입력 계층
  일반적인 Mamba는 모든 입력 데이터를 같은 방식으로 처리하는 반면에, MoM에서는 각각의 데이터 유형마다 별도의 입력 처리 파라미터를 세팅합니다. 마치 각 데이터 유형별로 특화된 ‘번역가’가 있다고 생각하면 좋을 거 같네요.

• 중간 계층
  일반 Mamba는 모든 데이터를 같은 과정으로 처리하지만, MoM에서는 이미지 데이터는 공간적인 관계를 강조해서 처리하고, 텍스트는 문맥적 관계를 강조하는 처리를 하는 등, 모달리티별로 최적화된 변환을 수행합니다.

• 출력 계층
  최종적인 출력 형태를 만드는데 있어서도, MoM은 텍스트 데이터의 출력물은 언어 이해에 중요한 특성을 강조하고, 이미지 데이터라면 시각적인 특성을 강조하는 등의 처리를 한다는 겁니다.

즉, 모델이 각각의 데이터 유형에 가장 맥락에도 맞고 효과적인 처리 방법을 적용해서, 트레이닝도 더 빠르게 되고 연산도 효율적으로 하면서도 다양한 모달리티에 적합한 처리를 하게 한다는 겁니다.

그런데, 공통의 구성 요소가 있다고 했죠? 이건 뭘까요?

상태 전이 (State Transitions), 그리고 컨볼루션 (Convolution)은 데이터 유형과는 무관하기 때문에, 공유된 상태로 유지되는데요. 계산 효율성을 유지하기 위해서도 사용되구요.

• 상태 전이(State Transitions)
  이건 시간에 따른 순차적 정보를 추적하고 업데이트하고, 모든 모달리티에 걸쳐서 장기적인 의존성을 포착하는 데 도움을 주는 메모리 단위처럼 작동합니다.

• 1D 컨볼루션 계층(1D Convolutional Layers)
  이건 시퀀스에서 로컬의 (Local) 패턴을 포착해서, 텍스트, 이미지, 음성, 비디오 등 다양한 모달리티에 걸쳐서 시간에 의존적인 정보를 처리하는 데 도움을 줍니다.

이렇게 설계를 해서, MoM이 서로 다른 모달리티를 잘 처리하면서도 핵심적인 계산은 공유를 해서 처리해서 효율을 높여줍니다.

자, 그럼 이제 MoM이 보여주는 성능, 장점, 그리고 한계점을 알아볼까요?

MoM의 성능과 장점, 한계

MoM의 성능

MoM 모델의 강점을 보여주려고, 연구자들이 세 가지 다른 AI 트레이징 셋업에서 MoM을 테스트했습니다:

• Transfusion – 텍스트와 연속적인 (Continuous) 이미지 데이터를 혼합합니다.

• Chameleon – 텍스트와 이산적인 (Discrete) 이미지 데이터를 혼합하는 세팅입니다.

• 3가지 모달리티 프레임웍 – 텍스트와 이미지 외에 음성을 세 번째 모달리티로 포함하는, 더 복잡한 세팅으로, 그 모두를 이산적 토큰 (Discrete Token)으로 처리합니다.

테스트 결과는 상당히 인상적이었는데요. (그러니까 논문으로 내기는 했겠지만요 ^.^;_

Transfusion 세팅

• 이미지 성능:

  • 1.4B 규모에서 MoM은 FLOPs (초당 부동 소수점 연산, 계산 성능 측정 단위)의 34.76%만 사용하면서 손실을 2.2% 감소시킵니다.

  • 더 작은 모델들(760M, 163M)도 비슷한 경향을 보여서, 손실은 2.37% 낮고 FLOPs는 60% 적게 사용합니다.

• 텍스트 성능: MoM은 더 적은 FLOPs를 사용하면서 검증 손실을 개선해 주고, 일반화 성능도 더 높은 것으로 나타납니다. 최고의 정확도에 도달하는 시간도 빨라서, 트레이닝 효율성이 높습니다.

• 전반적인 효율성: MoM은 전반적으로 트레이닝 손실을 낮추면서도 FLOPs를 최대 86.11%까지 줄입니다.

Chameleon 세팅 (이미지, 텍스트 모두 이산적 - Discrete - 토큰으로 처리)

• 이미지 성능: MoM은 FLOPs의 25.9% ~ 42.5%만 사용하면서도 이미지 손실을 최대 3.46%까지 줄입니다.

• 텍스트 성능: 1.5B 규모에서 MoM은 FLOPs의 65.4%만 사용하면서 텍스트 손실을 3.01% 낮춰 줍니다.

3가지 모달리티 프레임웍 (이미지, 텍스트, 음성)

• 음성 성능

  • 443M 규모: FLOPs의 19.2%만 사용하면서 훈련 손실을 7.14% 낮춥니다.

  • 1.5B 규모: FLOPs의 24.8%만으로 손실이 5.75% 낮아집니다.

• 전체 지표

  • 최대 3.67%의 훈련 손실 개선과 함께 FLOPs를 최대 56.2%까지 줄입니다.

  • MoM은 세 가지 모달리티 환경에서 일관되게 성능을 향상시켜 줍니다.

MoM (Mixture-of-Mamba)의 성능 수치에서 보듯이, MoM은 더 적은 FLOPs를 사용하면서도 일관되게 더 나은 정확도를 달성해서, 기존의 Mamba, 그리고 트랜스포머에 비해서 더 효율적인 모델이라는 것을 알 수 있습니다.

MoM의 장점

MoM의 장점을 종합해 보면, 다음과 같습니다:

• MoM 모델이 작동할 때, 특정한 모달리티 데이터에 관련된 매개변수만 활성화되어서, 계산과 훈련 속도 및 비용을 최적화할 수 있습니다.

• ‘Modality-Aware Sparsity’의 이점: MoM이 텍스트, 이미지, 음성 각각의 모달리티 데이터를 처리하는데 특화해서 효율화할 수 있도록 해 주고, 각 모달리티의 고유한 구조에 맞는 적합한 방식으로 작동합니다. 모델의 모든 부분에 이런 방법을 적용하면, 선택적으로 적용하는 경우보다 더 나은 결과를 얻을 수 있습니다.

• 확장이 쉽고 유연함: 디퓨젼 기반의 이미지 학습이라든가, 토큰 기반 처리 같은 다양한 트레이닝 전략에서 전반적으로 잘 작동합니다.

• 세 가지 멀티모달 세팅에서, MoM이 일관되게 전통적인 Dense Model 대비 우수한 성능을 보여줍니다: 텍스트 + 연속 이미지; 텍스트 + 이산 이미지; 텍스트 + 이미지 + 음성 등.

• MoM은 계산 효율성을 개선해서, 성능을 유지하거나 향상시키면서 비용을 최대 65%까지 절감해 줍니다.

• MoM은 더 적은 트레이닝 단계로도 동일한 정확도에 도달하기 때문에, Mamba Dense 및 Flex-Attention Transformer에 비해서 트레이닝 속도가 더 빠릅니다.

• 베이스 모델에 비해서 상당한 수준으로 손실을 줄여줘서, Unseen Data가 들어오더라도 잘. 적응하는, 즉 ‘일반화 성능’이 좋습니다.

• 상대적으로 적은 에너지를 소비해서, AI의 접근성을 높이고 환경에 더 친화적인 기법이라고 할 수 있습니다.

• MoM을 MoE 기법과 결합한다면, 멀티모달 AI 시스템에서 효율성을 추가적으로 개선할 여지가 있습니다.

MoM의 한계점

MoM이 Mamba 모델이 훨씬 더 멀티모달 데이터를 잘 처리하도록 개선을 하기는 했지만, 몇 가지 한계점도 있습니다:

• MoM의 Modality-Aware Sparsity는 특수한 매개변수화 (Parameterization)를 해야 해서, 표준의 SSM이나 트랜스포머에 비해 구현도, 최적화도 좀 더 복잡합니다.

• Mamba 블록에서 여러 가지의 Projection 구성 요소를 분리해야 하기 때문에, 디버깅과 파인튜닝이 좀 더 어려워질 수 있습니다.

• MoM의 트레이닝은 텍스트, 이미지, 음성 Representation 사이에서 신중하게 균형을 잡아야 하는 일이라서, 추가적인 트레이닝 오버헤드가 생길 수 있습니다.

• 음성 처리는 여전히 특수한 토큰화 기법이 필요합니다.

• 이 아키텍처의 효율성은 독립적인 구성 요소를 개선하는 부분보다는 공동 부분의 최적화에 크게 의존합니다.

• MoM이 실제 세계의 어플리케이션에 어느 정도 유의미하게 효과가 있을지는 아직 검증되지 않았습니다.

• GPT와 LLaMA와 같은 더 큰 규모의 트랜스포머 모델들이 특정한 NLP 벤치마크에서 여전히 MoM보다 성능이 우수하죠 - 즉, MoM이 성능에 있어서 잠재적으로 한계가 있을 수도 있다는 점을 시사합니다.

이런 한계점들은 MoM의 향후 개선 분야를 드러내 주는 것이지만, 어쨌든 이런 한계에도 불구하고 MoM은 Mamba가 특히 멀티모달 모델로서 제대로 된 트랜스포머의 대안으로 포지셔닝할 수 있도록 해 주는 진일보한 모델임에는 분명합니다.

맺으며: MoM이 주목받는 이유

자, 여기까지 ’Modality-Aware Sparsity’ 개념을 적용해서 다양한 유형의 입력 데이터를 효율적으로 처리하는, 개선된 상태 공간 모델(SSM)인 MoM (Mixture-of-Mamba)에 대해서 간단히 알아봤습니다. MoM은 멀티모달 AI를 위해서 만든, 더 빠르고, 더 효율적이며, 확장 가능한 모델로, 텍스트, 이미지, 음성 작업 전반에 걸쳐서 계산의 효율성과 트레이닝 효과 등의 측면에서 상당한 이점을 제공해 줍니다.

MoM이 다른 Mamba 기반 접근법 및 Sparse Transformer와 어떻게 다른지, 그 고유한 아키텍처를 확실하게 이해하는 게 중요합니다:

• Sparse Transformer는 어텐션 메커니즘에 Sparsity를 통합하기는 하지만, 작업의 관점에서는 ‘텍스트-이미지 생성’ 등과 같은 특정한 유형의 작업에 전문화하고 집중합니다.

• MoE-Mamba라든가 BlackMamba 같은 다른 MoM과 비슷한 모델들은, MLP 계층을 수정하는 방식으로 Mamba에 Sparsity를 도입하지만, Mamba의 핵심 구조를 변경하지는 않습니다.

MoM은 Mamba 블록 내부에 Modality-Aware Sparsity를 적용하는 다른 접근 방식을 통해서, 다양한 데이터 유형에 걸쳐서 특화된 처리를 할 수 있게 해 주고, 트랜스포머를 넘어서 ‘Sparse’ 아키텍처를 확장합니다. 이 모델은 SSM 영역의 중요한 발전 계기의 하나로, 멀티모달 AI가 장식할 세계에서 SSM의 잠재력을 보여준다고 할 수 있습니다.

앞으로, 특히 MoM과 MoE를 혼합하는 기법이 어떤 훌륭한 성능을 보여줄지 궁금한데요 - 구현이 복잡할 수는 있지만 효과가 크지 않을까 기대하고, 미래에 필요한 강력한 하이브리드 모델의 기초가 될 수도 있습니다.

보너스: 참고자료

1. Mixture-of-Mamba: Enhancing Multi-Modal State-Space Models with Modality-Aware Sparsity by Stanford University, Carnegie Mellon University, FAIR @Meta

2. Mixture-of-Mamba code on GitHub

3. BlackMamba: Mixture of Experts for State-Space Models by Zyphra

4. MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts by Maciej Pióro et al.

5. Mamba: Linear-Time Sequence Modeling with Selective State Spaces by Albert Gu, Tri Dao

6. Jamba-1.5: Hybrid Transformer-Mamba Models at Scale by Avshalom Manevich et al.

7. Cobra: Extending Mamba to Multi-Modal Large Language Model for Efficient Inference by Han Zhao, Siteng Huang

8. MambaVision: A Hybrid Mamba-Transformer Vision Backbone by Ali Hatamizadeh, Jan Kautz

9. Scalable Autoregressive Image Generation with Mamba by Haopeng Li, Xuerui Qiu et al.