Confusion about the different archetectures between code of VSSBlock and figure in paper of VSSBlock

[Allenem]

VM-UNet/models/vmunet/vmamba.py

Line 492 in acbafd5

x = input + self.drop_path(self.self_attention(self.ln_1(input)))

Hi, Thank you very much for your excellent work!

I want to know why the forward function of class VSSBlock only includes

1 layer norm, SS2D, drop & addition;

rather than what is ilustrated in paper Fig1 (b) :

2 Layer Norm, 3 Linear layer, 1 DW-Conv, 2 Activation, SS2D, Addition & 1 Element-wise production

[YunhengWu-IB]

I also would like to ask the same question!!
Do you know the reason now?

[Allenem]

I also would like to ask the same question!! Do you know the reason now?

I understand the code now after reading the class VSSBlock & class SS2D together.

Both code and Fig1 (b) are same (only drop haven't been drawn in the figure). I add the comments to the fig1(b) here. The image below is easy to show the details. I hope you can understand the code easier after reading my comments 😄 :

[YunhengWu-IB]

Thank you very much for your very nice explanation!!!

[ZhouCong223]

Your explanation and annotations are excellent. Thank you for your response