PyTorch 单文件中实现 Mamba

原文 ↗ Hacker News 讨论 ↗

一个紧凑的、单文件的 PyTorch 版本新型 Mamba 神经网络架构实现，被赞为相较于高度优化的官方代码更清晰、更适合学习的替代方案。参与者解释了 Mamba 的结构化状态空间模型如何通过线性时间序列建模、更低的 VRAM 使用和更快的推理来挑战 Transformer，同时也在讨论可读性、性能技巧以及工具选择（从 Fortran 到 Julia 和 Mojo）之间的取舍。这场讨论凸显了 ML 行话和模型家族演化的速度之快，并为希望理解 Mamba 等架构在大语言模型更广阔版图中位置的实践者分享了资源。

教育性的单文件实现

  • 许多人将这个单文件 PyTorch Mamba 视为一种学习工具。
  • 与高度优化的官方 CUDA 代码相比，这个版本被认为更容易理解得多。
  • 这里之所以强调行数和文件数，是因为目标是概念清晰，而不是速度。

替代实现与语言

  • 文中提到了一种基于 Fortran 的 Mamba/SSM 推理实现；Fortran 因其类似 NumPy 的数组语法、编译后的速度、易于并行以及对线性代数的良好支持而受到称赞，但也因非数值任务笨拙而受到批评。
  • 有人建议考虑 Julia 或 Mojo 等更新的语言，但这位 Fortran 支持者接受这场“上坡战”。
  • 另一个库被提到，它将共享的 transformer 代码抽离出来，使 BERT、LLaMA 和 MPT 等模型可以压缩到大约 100 行。

性能与并行化

  • 评论者指出，顺序的 selective_scan 循环可以通过几个 PyTorch 调用并行化，而且所有输出都可以用一个 einsum 计算出来。
  • 其他人则指出，这可能会损害可读性，而可读性正是该仓库的主要目的；建议可以提供可选的并行路径，或者把简单版本保留为注释。

Mamba / SSM 是什么，以及为什么有趣

  • Mamba 被描述为一种基于状态空间模型（SSM）的架构，目标是与 transformer 竞争：用线性时间的序列建模替代二次复杂度的注意力。
  • 关键想法是：让部分 SSM 参数依赖输入（“选择”机制），同时保持潜在状态转移是线性的，从而既具备表达能力，又能实现高效的并行算法。
  • 它被描述为对早期 SSM / 长卷积工作的延续（HIPPO、S4、Hyena 等），并补上了输入依赖门控等缺失要素。

行话、命名与领域速度

  • 几位发帖者对飞快涌现的行话和新模型名称（Mamba、RetNet、RWKV 等）感到吃力，认为这个领域有很强的潮流驱动和品牌化倾向。
  • 也有人认为短名称很有用，而理解则来自于“旁听”社区对话（HN、subreddit、新闻简报）。
  • 建议阅读词汇表、README 中的参考资料、博客入门文章以及讲解视频等资源。

能力、效率与未解问题

  • 相关说法包括：
    • Mamba 的训练速度更快，推理速度也比同等规模的 transformer 快得多，而且每个 token 的 VRAM 使用更低。
    • 它可能在计算/样本效率上更高，也更适合超长上下文和其他序列领域。
  • 有人报告称，在某个特定测试中，Mamba 使用的 VRAM 约为 RetNet 的 60%，且每个 token 的成本低于 transformer。
  • 也有人指出：
    • 与基于注意力的模型相比，下游应用仍不够成熟。
    • 在大规模条件下，Mamba 与 RetNet 和 RWKV 的对比仍不清楚；上下文长度测试和高效实现都还在演进中。
    • 一个人声称当前的 Mamba 模型“比 GPT-2 更不连贯”，另一人则明确反对并说它们更好；这里没有得出结论。

与 RNN / LSTM 的关系

  • 一个反复出现的问题是 Mamba 与 RNN/LSTM 有何不同。
  • 给出的解释是：它本质上是在潜在空间中的线性 RNN（为了可并行化和稳定性），而输入依赖的转移提供了传统上由非线性带来的表达能力。
  • 提到的一个实际限制是：与 transformer 不同，训练时高效地打包不同长度的序列可能更困难。

代码风格、工具与研究过程

  • 对 einops/einsum 的看法不一：有人觉得它们优雅且高效；也有人认为它们会降低可读性。
  • 大家对极简、独立的“代码草图”抱有热情，认为它们能加快研究速度，并降低实验的“柯尔莫哥洛夫复杂度”。
  • 单文件实现被赞为对庞大、相互依赖的 ML 代码库的一种制衡。