[論文レビュー] QuadMamba: Learning Quadtree-based Selective Scan for Visual State Space Model
QuadMamba は、Vision Mamba の適応のために学習可能なクアドツリーベースの選択スキャンを導入し、分類・検出・セグメンテーションのすべてにおいて直線的複雑さで高性能を達成します。
Recent advancements in State Space Models, notably Mamba, have demonstrated superior performance over the dominant Transformer models, particularly in reducing the computational complexity from quadratic to linear. Yet, difficulties in adapting Mamba from language to vision tasks arise due to the distinct characteristics of visual data, such as the spatial locality and adjacency within images and large variations in information granularity across visual tokens. Existing vision Mamba approaches either flatten tokens into sequences in a raster scan fashion, which breaks the local adjacency of images, or manually partition tokens into windows, which limits their long-range modeling and generalization capabilities. To address these limitations, we present a new vision Mamba model, coined QuadMamba, that effectively captures local dependencies of varying granularities via quadtree-based image partition and scan. Concretely, our lightweight quadtree-based scan module learns to preserve the 2D locality of spatial regions within learned window quadrants. The module estimates the locality score of each token from their features, before adaptively partitioning tokens into window quadrants. An omnidirectional window shifting scheme is also introduced to capture more intact and informative features across different local regions. To make the discretized quadtree partition end-to-end trainable, we further devise a sequence masking strategy based on Gumbel-Softmax and its straight-through gradient estimator. Extensive experiments demonstrate that QuadMamba achieves state-of-the-art performance in various vision tasks, including image classification, object detection, instance segmentation, and semantic segmentation. The code is in https://github.com/VISION-SJTU/QuadMamba.
研究の動機と目的
- Mamba を 2D の局所性を保ちながら視覚タスクに適応させ、マルチグラニュラリティの視覚信号を扱う動機づけ。
- Vision Mamba の局所性を意識した 1D シーケンスを形成する学習可能なクアドツリーベースのスキャンモジュールを提案。
- Gumbel-Softmax による微分可能な分割とシーケンスマスキングを用いて、エンドツーエンドの訓練を可能にする。
- 地域間の情報流を強化するため、全方位ウィンドウシフティングを取り入れた情報特徴を捉える。
- CNN、Transformer、従来の Mamba ベースラインと比較して ImageNet-1k、COCO、ADE20K で強力な性能を示すことを示す。
提案手法
- 画像を 4x4 パッチに分割してトークンへ埋め込む。
- パーティションマップを予測して粗い窓と細かい窓への bi-level クアドツリー分割を導く。
- Gumbel-Softmax ベースのマスキングと Hadamard/要素演算を用いて differentiable な 1D トークン列を構築する。
- QuadVSS ブロック(Mamba ベース)をシフトと多段の Transformer 的構造で適用する。
- 全ブロックに跨る全方位ウィンドウシフティングを取り入れ、局所性と領域横断情報の流れを強化する。
- クアド VSS 深度とチャネルを設定可能な複数のアーキテクチャ版本(Lite/Tiny/Small/Base)を調整する。
実験結果
リサーチクエスチョン
- RQ1学習可能なクアドツリーベースのスキャンは Vision Mamba においてラスタや固定ウィンドウスキャンより 2D の局所性をより良く保持できるのか?
- RQ2適応的でマルチグラニュラリティの分割は視覚タスクにおける長距離モデリングと効率性を改善するのか?
- RQ3全方位ウィンドウシフティングは局在化と各タスクの性能にどのような影響を及ぼすのか?
- RQ4QuadMamba のバリエーションは ImageNet-1k、COCO、ADE20K において CNN、Vision Transformer、および従来の Mamba バックボーンと比較してどうなるのか?
主な発見
| モデル | #Params (M) | FLOPs (G) | Top-1 (%) | Top-5 (%) |
|---|---|---|---|---|
| ResNet-18 | 11.7 | 1.8 | 69.7 | 89.1 |
| ResNet-50 | 25.6 | 4.1 | 79.0 | 94.4 |
| ResNet-101 | 44.7 | 7.9 | 80.3 | 95.2 |
| RegNetY-4G | 20.6 | 4.0 | 79.4 | 94.7 |
| RegNetY-8G | 39.2 | 8.0 | 79.9 | 94.9 |
| RegNetY-16G | 83.6 | 15.9 | 80.4 | 95.1 |
| DeiT-S | 22.1 | 4.6 | 79.8 | 94.9 |
| DeiT-B | 86.6 | 17.6 | 81.8 | 95.6 |
| PVT-T | 13.2 | 1.9 | 75.1 | 92.4 |
| PVT-S | 24.5 | 3.7 | 79.8 | 94.9 |
| PVT-M | 44.2 | 6.4 | 81.2 | 95.6 |
| PVT-L | 61.4 | 9.5 | 81.7 | 95.9 |
| Swin-T | 28.3 | 4.5 | 81.3 | 95.5 |
| Swin-S | 49.6 | 8.7 | 83.3 | 96.2 |
| Swin-B | 87.8 | 15.4 | 83.5 | 96.5 |
| Mamba Vim-Ti | 7 | – | 76.1 | 93.0 |
| Vim-S | 26 | – | 80.5 | 95.1 |
| VMamba-T | 22 | 4.5 | 82.2 | – |
| VMamba-S | 44 | 9.1 | 83.5 | – |
| VMamba-B | 75 | 15.2 | 83.7 | – |
| LocalVim-T | 8 | 1.5 | 76.2 | – |
| LocalVim-S | 28 | 4.8 | 81.2 | – |
| PlainMamba-L1 | 7 | 3.0 | 77.9 | – |
| PlainMamba-L2 | 25 | 8.1 | 81.6 | – |
| PlainMamba-L3 | 50 | 14.4 | 82.3 | – |
| QuadMamba-Li | 5.4 | 0.8 | 74.2 | 92.1 |
| QuadMamba-T | 10 | 2.0 | 78.2 | 94.3 |
| QuadMamba-S | 31 | 5.5 | 82.4 | 95.6 |
| QuadMamba-B | 50 | 9.3 | 83.8 | 96.7 |
- QuadMamba のバリエーションは、同時代の手法より低いまたは同等の FLOPs で高精度を達成する(例: QuadMamba-S: ImageNet-1k における Top-1 82.4%、パラメータ 31M、FLOPs 5.5G)。
- QuadMamba-S および -B は、同程度の複雑さで COCO の物体検出とインスタンスセグメンテーションで複数の CNN、Transformer、Mamba バックボーンを上回る(例: QuadMamba-S: APbox 46.7、APmask 42.4)。
- QuadMamba は ADE20K のセマンティックセグメンテーション mIoU が 47.2%(QuadMamba-S)で、多くのベースラインと同等かそれ以上のコストで競合可能。
- アブレーションは、 coarse-to-fine の多段局所性と多方向ウィンドウシフトが、単純なフラット化や単一方向シフトよりも性能を有意に向上させることを示す。
- ウィンドウ分割スケールと層分布(LP3)は、精度と複雑さの間で最良のトレードオフを生み出す。
より良い研究を、今すぐ始めましょう
論文設計から論文執筆まで、研究時間を劇的に削減しましょう。
クレジットカード登録不要
このレビューはAIが作成し、人間の編集者が確認しました。