Summary
The paper introduces QuadMamba, an enhancement of vision State Space Model (SSM) architectures. At its core is a learnable QuadVSS network block that processes the image input patch at two different resolutions. For every 2x2 coarse window with 4 image patches, the method adaptively learns to process one of the 4 patches at the finer 2x resolution using the differentiable Gumbel-softmax formulation. These block are hierarchically stacked similar to SwinTransformer, but using a window shifting scheme in two different directions which is a better fit for SSM.
The method is demonstrated to obtain good results relative to comparably sized transformer and CNN architectures on Imagenet classification, COCO object detection / instance segmentation, and ADE20k semantic segmentation.
Strengths
- Intuitive idea to learn how to pick an area that requires higher resolution processing and pack it into SSM via Gumbel Softmax. Solid although not particularly novel stacking of QuadVSS blocks into a SwinTransformer-like network architecture.
- Competitive results relative to ViT and CNN, and seems to improve a little bit over other Mamba methods, although the gains there seem quite incremental.
- Ablation over multiple network parameter decisions.
Weaknesses
# Significance
The contribution seems a bit incremental. The overall idea of image traversals into several windows to enforce better locality was already explored before in LocalMamba. The gains in the experimental section over some of the other Mamba network variants (EfficientVMamba, VMamba, Swin-S etc) are not too large. The idea that we do not have to process all areas in high resolution but only 1/4 of them does seem useful, although does not seem strictly limited to SSMs -- e.g. would ViT methods also benefit?
# Experimental results
- It seems that the natural baseline to this method is to compare a network made of VSS modules, as opposed to QuadVSS (the main novelty). Such a comparison was done in Fig 5, however it is not too clear how rigorous it is. It would be helpful to compare more directly VSS / QuadVSS equivalents with same FLOPS, for several different FLOPS thresholds. This does not seem to have been done - it would validate more strongly the fact that allocating parameters selectively to higher resolutions actually helps (as opposed to using the baseline VSS module pyramid).
- Unclear what the training overhead of a QuadVSS block is compared to a VSS block (in terms of memory, compute)?
- As opposed to just a table, it would be helpful to have plots with flops vs quality as axes, with a separate curve in that graph for each model family. Such a plot can make the comparison more obvious, given that different methods have somewhat different amounts of flops.
- In Sec 4.3, it's unclear why for the Tiny model, we stack more blocks in the second stage, while for Small and Base we stack more in the third stage. Any particular reason? Third stage is more standard VSS modules, as opposed to your QuadVSS innovation. Also later in Table 6 yet different stackings are best (2,4,6,2). It's hard to discern the logic in all these choices.
# Clarity and questions
Some details of the approach were not particularly clear to me.
- L175 Do you always pick the same fine-grained patch for all 2x2 windows, or do they vary by window? Not particularly clear from the exposition/notation.
- Also, it appears that 7 patches are picked for each region as per L185-187. Don't you want to have 8 patches (powers of 2 are usually more efficient hardware-wise?)
- The softmax in Eq 7, what labels do you train it on? Is it the final labels? Right now it appears that this part is trained first, and kept fixed.
- (minor) In Eq 6, it seems 'v_local' is aggregated across the whole image, shouldn't it be called 'global' instead?
- (minor) L186: where is 'local adjacency score' defined? This is the first mention of the term.
- (minor) Fig 5: S-QuadVSS is not defined, I assume it's the opposite direction shift, but helps to be explicit.
# Minor Nits and Typos
159: infOrmative
Table 1: VMamaba
606: Uperhead [62] (I believe it should be UperNet)
Questions
- Is your quad idea limited to SSM or would it also work for Vit? Can one learn to downsample specific patches in a Vit?
- Do you have comparisons between networks made of QuadVSS/VSS or only VSS for several same FLOPS budgets? Do you have comparison of training overhead (compute/memory) for networks made of QuadVSS/VSS or only VSS baseline?
- L175 Do you always pick the same fine-grained patch for all 2x2 windows, or do they vary by window? Not particularly clear from the exposition/notation.
- Also, it appears that 7 patches are picked for each region as per L185-187. Don't you want to have 8 patches (powers of 2 are usually more efficient hardware-wise?)
- The softmax in Eq 7, what labels do you train it on? Is it the final labels? Right now it appears that this part is trained first, and kept fixed.