VisionMamba implementation

[AliYoussef97]

Vision Mamba implementation. Forward function contains both forward and backward direction SSM, however, I did not implement the step function as both directions would return a cache while the final output is the combination of both the forward and backward direction.

[alxndrTL]

Thank you for the PR! Have you tried it ?

[AliYoussef97]

@alxndrTL I did pass dummy tensors to the model to check if there are any errors thrown, but I haven't trained the model on specific task as that would take a bit more time to make training script similar to this, but I wouldn't mind making one when I have more time on my hands.

I might make another commit to this PR later today however to make the code a little bit cleaner, mainly:

• It is unneccessary to initialise the backward pass parameters if the model is set as unidirectional.
• I am not a big fan of how I set up ssm and ssm_back, would be better to just set one ssm function and add additional arguments to it such as weights and biases.
• Need to add the option to divide combined output as in here.

I will probably commit those changes in the next hour or so. Commits added.

[alxndrTL]

Thanks! I'll merge this now and add comments in the file as well as in the README

[AliYoussef97]

Thanks! I'll merge this now and add comments in the file as well as in the README

Great! I’ll try and add a training script and the full trainable model sometime this month, similar to this (adding cls token, pos embedding, random sequence flips, etc. This PR is the Vision Mamba Block only).

The Vision Mamba block is the “v2” here (it gets called from the Block function in the first link), which what this PR is based off.)

[alxndrTL]

Hello, i'm refactoring your code in vim.py, if I understood you correctly, I get get rid of all the things related to the cache right ?

[AliYoussef97]

That is correct, I did not implement the step function, only the forward function.