MambaMixer: Efficient Selective State Space Models with Dual Token and Channel Selection

要約

深層学習の最近の進歩は、データの依存性と大規模な学習能力により、主に Transformer に依存しています。
ただし、これらのアーキテクチャのアテンション モジュールは、入力サイズで 2 次の時間と空間を示し、長いシーケンスのモデリングのスケーラビリティが制限されます。
状態空間モデル (SSM)、より具体的には選択的 SSM (S6) は、ハードウェアを意識した効率的な実装を備えており、長い因果系列モデリングに有望な可能性を示しています。
ただし、チャネルごとに個別のブロックを使用するため、無関係なチャネルをフィルタリングしたり、チャネル間の依存関係をキャプチャしたりすることができません。
MLP、アテンション、または SSM を使用してチャネル間で情報を混合しようとする自然な試みは、大規模ネットワークに対する SSM のトレーニングのさらなる不安定性やパラメータ数のほぼ 2 倍をもたらします。
MambaMixer ブロックは、トークンとチャネルにわたる二重選択メカニズム (選択トークンおよびチャネル ミキサーと呼ばれる) を使用する、データ依存の重みを備えた新しい SSM ベースのアーキテクチャです。
パラメータ数の 2 倍を軽減するために、ハードウェアに優しい実装で準分離可能なカーネルを使用した S6 ブロックの新しい非因果的ヒューリスティックを提示します。
さらに、シーケンスと埋め込みの両方の次元に沿って情報を混合する、QSMixer と呼ばれる MambaMixer の効率的なバリアントを紹介します。
概念実証として、Vision MambaMixer (ViM2) および Vision QSMixer (ViQS) アーキテクチャを設計します。
画像内の空間情報をキャプチャする能力を強化するために、一連の有用な画像スキャンを動的に使用して画像パッチを横断するスイッチ オブ スキャン (SoS) を紹介します。
画像分類、セグメンテーション、オブジェクト検出におけるメソッドのパフォーマンスを評価します。
私たちの結果は、トークンとチャネルの両方を選択的に混合することの重要性を強調し、確立されたビジョンモデル（それぞれ SSM ベースのモデル）を使用した場合の当社の手法の競合（それぞれ優れた）パフォーマンスを示しています。

要約(オリジナル)

Recent advances in deep learning have mainly relied on Transformers due to their data dependency and ability to learn at scale. The attention module in these architectures, however, exhibits quadratic time and space in input size, limiting their scalability for long-sequence modeling. State Space Models (SSMs), and more specifically Selective SSMs (S6), with efficient hardware-aware implementation, have shown promising potential for long causal sequence modeling. They, however, use separate blocks for each channel and fail to filter irrelevant channels and capture inter-channel dependencies. Natural attempt to mix information across channels using MLP, attention, or SSMs results in further instability in the training of SSMs for large networks and/or nearly double the number of parameters. We present the MambaMixer block, a new SSM-based architecture with data-dependent weights that uses a dual selection mechanism across tokens and channels-called Selective Token and Channel Mixer. To mitigate doubling the number of parameters, we present a new non-causal heuristic of the S6 block using quasi-separable kernels with a hardware-friendly implementation. We further present an efficient variant of MambaMixer, called QSMixer, that mixes information along both sequence and embedding dimensions. As a proof of concept, we design Vision MambaMixer (ViM2) and Vision QSMixer (ViQS) architectures. To enhance their ability to capture spatial information in images, we present Switch of Scans (SoS) that dynamically uses a set of useful image scans to traverse image patches. We evaluate the performance of our methods in image classification, segmentation, and object detection. Our results underline the importance of selectively mixing across both tokens and channels and show the competitive (resp. superior) performance of our methods with well-established vision models (resp. SSM-based models).

arxiv情報

提供元, 利用サービス

arxiv.jp, Google