why GPU memory cost of Mamba2 Block > full self-attention block? And how to reduce this memory cost when training?

[thucz]

I found Mamba2 is much faster than full self-attention block. But I met a memory problem.

I used 12 layers of Mamba2 in the vision task. d_model=128, d_state=16, head_dim = 32, expand = 2.

I found the inner dimension of the middle layer is > d_model while the inner_dim in my self-attention block is equal to d_model.

How to reduce the inner dimension to have a similar or smaller memory cost than full self-attention?

[ScottHoang]

Have you tried reducing the expansion_ratio?