Competitors to Transformer

[Aidenzich]

• xLSTM
  • Video: xLSTM: The Sequel To The Legendary LSTM
• Mamba
• Jamba
• TTT

[Aidenzich]

Mamba

Selective SSM

What are the benefits of SSSM?

• A selective scan algorithm, which allows the model to filter (ir)relevant information
• A hardware-aware algorithm that allows for efficient storage of (intermediate) results through parallel scan, kernel fusion, and recomputation.

Parallel Scan

NOTE: In Maarten's diagram, $H_n$ should be regarded as $H_n'$ because it is not the final output before the sweep-up phase. This means that in the complete computation process, H_0, like other intermediate states, needs to be further processed during the sweep-up phase.

Code Implement

```python def selective_scan(self, u, delta, A, B, C, D): (b, l, d_in) = u.shape n = A.shape[1] # Discretize continuous parameters (A, B) deltaA = torch.exp(einsum(delta, A, 'b l d_in, d_in n -> b l d_in n')) deltaB_u = einsum(delta, B, u, 'b l d_in, b l n, b l d_in -> b l d_in n') # Initialize x x = torch.zeros((b, d_in, n), device=deltaA.device) ys = [] # Sweep-down phase for i in range(l): x = deltaA[:, i] * x + deltaB_u[:, i] y = einsum(x, C[:, i, :], 'b d_in n, b n -> b d_in') ys.append(y) # Stack results y = torch.stack(ys, dim=1) # Perform sweep-up phase y = y + u * D return y ```

Hardware-aware

Reference

URL	Description
mamba	The code for selective scan in the original repo is implemented at a low level, making it harder to understand.
mamba-minimal	Simplified, readable, annotated code for Mamba; not optimized for speed.
Maarten's blog	The comprehensive introduction by Maarten, the author of BERTopic. The best article I have found 💯