arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2609.07681cs.LGcs.CL

关于Mamba中回忆缩放定律:基于哈希的理论与机制研究

On the Recall Scaling Laws in Mamba: A Theoretical and Mechanistic Study via Hashing

  • Blavatnik School of Computer Science and AI, Tel Aviv University(特拉维夫大学布拉瓦尼克计算机科学与人工智能学院)

机构由 AI 辅助整理,请以论文原文为准。

Yuval Koren, Assaf Ben-Kish, Raja Giryes, Lior Wolf, Itamar Zimerman

AI总结:

本文通过机制可解释性研究Mamba中的联想回忆,发现其隐式学习线性哈希函数,并基于此提出回忆缩放定律理论框架,可预测模型维度与回忆性能的关系,实证验证了其准确性。

AI中文摘要:

联想回忆(AR)是从记忆中学习并检索项目之间联系的能力。在自然语言处理中,AR被用作评估Mamba等架构的上下文记忆能力的基准,并已被发现与语言建模性能密切相关。本文从机制可解释性的角度探讨AR,旨在逆向工程Mamba用于执行回忆的确切内部算法。我们的关键见解是,Mamba通过隐式学习线性哈希函数来执行回忆,并且我们识别了实现这一行为的底层电路。基于这些发现,并受相似性保持哈希中理论工具(如Johnson-Lindenstrauss引理)的启发,我们开发了一个用于分析AR的理论框架,称之为回忆缩放定律。给定词汇表大小和上下文中的事实数量,该框架使我们能够(1)预测Mamba实现完美回忆所需的嵌入和状态维度,(2)在给定模型维度的情况下预测回忆成功概率,以及(3)分析多层模型和多头SSM模式。实证结果表明,我们的理论发现是准确且具有预测性的,为AR容量如何随词汇量、状态、嵌入大小和架构扩展提供了见解。

英文摘要:

Associative Recall (AR) is the cognitive ability to learn and retrieve links between items in memory. In NLP, AR is used as a benchmark for evaluating the in-context memory capacity of architectures such as Mamba, and has been found to strongly correlate with language modeling performance. This paper explores AR from the perspective of mechanistic interpretability, aiming to reverse-engineer the exact internal algorithm used by Mamba to perform recall. Our key insight is that Mamba performs recall by implicitly learning linear hash functions, and we identify the low-level circuit that enables this behavior. Building on these findings and inspired by theoretical tools in similarity-preserving hashing, such as the Johnson-Lindenstrauss lemma, we develop a theoretical framework for analyzing AR, which we term Recall Scaling Laws. Given the vocabulary size and the number of facts in context, this framework allows us to (1) predict the embedding and state dimensions required for Mamba to achieve perfect recall, (2) predict recall success probability given the model dimensions, and (3) analyze multi-layer models and multi-head SSM patterns. Empirical results show that our theoretical findings are accurate and predictive, offering insights into how AR capacity scales with vocabulary, state, embedding size, and architecture.

补充信息

↑