Mamba: The Selective Step Size
Intermediate
The problem, from first principles
Mamba's step size depends on the current input, through a softplus so it stays positive. A large step lets the state absorb the current token, and a small step lets it ignore it. That input-dependent choice is the selection mechanism.
From theory to code
Implement selective_step(delta_raw), mapping a raw value to a positive step size.
Constraints
- Use a stable softplus.
Hints
Hint 1
Compute the log of one plus the exponential of the raw value, using logaddexp for stability.
Keywords
research-paperssequence-modelsstate-spacesequence
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.