5 ms·
The innovation is that they managed to write a hardware-aware kernel to make it run fast/efficiently on GPUs. The authors of Mamba are the same authors of Flash
by prefer_to_ask 3y ago
The innovation is that they managed to write a hardware-aware kernel to make it run fast/efficiently on GPUs. The authors of Mamba are the same authors of FlashAttention which was a performance optimization kernel written to reduce IO while computing normal O(N^2) attention. As I understand it, previously SSM models were not as easily parallelizable as the Transformer architecture.
- kelseyfrog 3y agoYes and no. There's an dual connection between ssms and convolutional models if certain constraints are met. Training convolutionally and inferring sequentially seeks a compromise between the two sides. I think we're about to find out the degree to which those constraints impact "easily".
- wityl 3y agoThey were previously parallelizable (via fft), but performed poorly on language modeling tasks. Mamba adds a dependence on the inputs that makes language modeling competitive with transformers, but that prevents using the fft approach. So they switch to a method using parallel prefix scan.