selective_scan_ref deltaB_u doesn't follow Eqn 4 in paper about mamba HOT 5 CLOSED

Yes, we took an approximation that didn't change the performance much but is slightly easier to implement. This is the first order approximation on the discretization coefficient in front of $B$ (i.e. the Euler approximation, but only applied to the $B$ term).

from mamba.

Ahh okay. Thanks for clarifying!

from mamba.

A related question on A_bar, shouldn't it be matrix exponential e^(Delta A)? The first order approximation would be Eye + Delta .* A; In the reference code, element wise exp function is used:

def selective_scan_ref(...):
...
deltaA = torch.exp(torch.einsum('bdl,dn->bdln', delta, A))

Thanks

from mamba.

A is a diagonal matrix so exp is applied element-wise (i.e. applied to the diagonal entries). The off diagonal entries are all zero and not stored.

from mamba.

I see, the network will learn all the eigen values of A, and eigen vectors absorbed by surrounding projections (weights also learned).
Thanks for the wonderful work!

from mamba.