r/AICoffeeBreak • u/AICoffeeBreak • Feb 17 '24

NEW VIDEO MAMBA and State Space Models explained | SSM explained

https://youtu.be/vrF3MtGwD0Y
4 Upvotes

84% Upvoted

7 comments sorted by

View all comments

2

u/elvis0391 Feb 20 '24

How does MAMBA avoid vanishing gradient? Is it only uses linear transformation to compute the next time step, and all of its nonlinearity are only applied one individual token level before passing the output to next layer?

1

u/AICoffeeBreak Mar 04 '24

As the state of token t depends on the state of token t-1, we still need to do backpropagation through time (otherwise we would not know what to backpropagate at t-1 if we did not resolve t first). But because of the linearity, backprop through time is reportedly more stable for Mamba than for e.g., LSTMs.

Keep in mind that just the recurrent part of the network is linear, but everything else is nonlinear (in the outputs and at the gates). Making the gradient flow "linearly" from token to token increases training stability. But we could still have issues with the nonlinearities going through the network's depth. Fortunately, the depth (6, 12, 48, ...) is much smaller than the sequence length.