A short project on Mamba: grokking & interpretability — LessWrong

Epistemic status: I've worked on this project for ~20h, on my free time and using only a Colab notebook. I trained a minimalistic implementation of Mamba (details below) on the modular addition task. I found that: Almost all of the Mechanistic Interpretability (MI) efforts I've seen people excited about and the great majority of the techniques I've learned are related to Transformer-based architectures. At the same time, a competitive alternative (Mamba) was recently introduced and later scaled. To me, when coupling these two facts together, a giant gap between capabilities and safety emerges. Thus, I think Mamba provides an interesting use case where we can test whether the more conceptual foundations of MI are solid (i.e., somewhat model-agnostic) and, therefore, whether MI can potentially survive another transformer-like paradigm shift on the race towards AGI. For a bit more of context, Mamba is based on a special version of State Space Models (SSMs): add another S (for Structured)

x A short project on Mamba: grokking & interpretability — LessWrong AI Capabilities Grokking (ML) Interpretability (ML & AI) Neuroscience AI Frontpage 21 A short project on Mamba: grokking & interpretability by Alejandro Tlaie 18th Oct 2024 8 min read 0 21 Epistemic status: I've worked on this project for ~20h, on my free time and using only a Colab notebook . Executive summary I trained a minimalistic implementation of Mamba (details below) on the modular addition task . I found that: This non-transformer-based model can also exhibit grokking (i.e., the model learns to generalise after overfi

related reading

• Towards Monosemanticity: Decomposing Language Models With Dictionary Learningtransformer-circuits.pub
• Transformer Circuits Threadtransformer-circuits.pub
• An Extremely Opinionated Annotated List of My Favourite Mechanistic Interpretability Papers v2 — AI Alignment Forumalignmentforum.org
• Mamba: The Easy Wayjackcook.com
• An Extremely Opinionated Annotated List of My Favourite Mechanistic Interpretability Papers v2 — AI Alignment Forumalignmentforum.org
• Mamba Explainedthegradient.pub
• Mamba No. 5 (A Little Bit Of…) | Sparse Notesjameschen.io
• A Mechanistic Interpretability Analysis of Grokking — AI Alignment Forumalignmentforum.org
• A Pragmatic Vision for Interpretability — AI Alignment Forumalignmentforum.org
• A Visual Guide to Mamba and State Space Modelsnewsletter.maartengrootendorst.com
• The Building Blocks of Interpretabilitydistill.pub
• A Mechanistic Interpretability Analysis of Grokking — LessWronglesswrong.com