1:15 Albert says that doubling the state size in Mamba 1 doubles the wall clock time. Also, that in Mamba 2 much of the computation is not contingent on the state size. Why the latter? Computation time is constant because it’s one matmul happening in parallel as one step on the GPU?
When Albert says, multiple times, that they avoid to materialize the state, it sounds that they don’t materialize the state at all during the forward pass in training. Does he mean that exactly? Or that they avoid to materialize the full state at once, but materialize the whole state incrementally, chunk by chunk?
From the Mamba paper: Commonly, the model uses the convolutional mode (3) for efficient parallelizable training (where the whole input sequence is seen ahead of time), and switched into recurrent mode (2) for efficient autoregressive inference (where the inputs are seen one timestep at a time). ... Note that the recurrent mode is more flexible than the convolution mode, since the latter (3) is derived from expanding the former (2) (Gu, Goel, and Ré 2022; Gu, Johnson, Goel, et al. 2021). However, this would require computing and materializing the latent state ℎ with shape (B, L, D, N), which is much larger (by a factor of 𝑁 , the SSM state dimension) than the input 𝑥 and output 𝑦 of shape (B, L, D). Thus the more efficient convolution mode was introduced which could bypass the state computation and materializes a convolution kernel (3a) of size only (B, L, D).
