Mamba通过测试时特征学习,高效实现低维目标的上下文学习。
Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning
- 利用测试时特征学习,从上下文例子里直接提取关键特征方向。
- 样本复杂度优于线性Transformer,接近非线性Transformer最优水平。
- 揭示了非线性门控机制是Mamba高效学习的核心驱动力。
Mamba是一种近期提出的线性时间序列模型,因其计算效率高和实证性能强而受到广泛关注。然而,对其内在机制的严格理论理解仍有限。本文聚焦于由低维非线性目标函数定义的任务,对Mamba的上下文学习(ICL)能力进行理论分析。具体研究单指数模型 $y ilde{g_*(oldsymbolβ, oldsymbol{x})}$,其仅依赖于单一相关方向 $oldsymbolβ$(即特征)。我们证明,经梯度法预训练的Mamba可通过测试时特征学习实现高效的上下文学习,直接从上下文样例中提取相关方向。由此建立的测试时样本复杂度优于线性Transformer(被分析为核方法),并接近非线性Transformer,后者已被证明可突破相关统计查询(CSQ)下界,并达到近信息论最优速率。分析揭示了非线性门控机制在Mamba中对特征提取的关键作用,凸显其作为模型实现高效与高性能的根本驱动力。
原文摘要 · Abstract (English)
Mamba, a recently proposed linear-time sequence model, has attracted significant attention for its computational efficiency and strong empirical performance. However, a rigorous theoretical understanding of its underlying mechanisms remains limited. In this work, we provide a theoretical analysis of Mamba's in-context learning (ICL) capability by focusing on tasks defined by low-dimensional nonlinear target functions. Specifically, we study in-context learning of a single-index model $y \approx g_*(\langle \boldsymbolβ, \boldsymbol{x} \rangle)$, which depends on only a single relevant direction $\boldsymbolβ$, referred to as feature. We prove that Mamba, pretrained by gradient-based methods, can achieve efficient ICL via test-time feature learning, extracting the relevant direction directly from context examples. Consequently, we establish a test-time sample complexity that improves upon linear Transformers -- analyzed to behave like kernel methods -- and is comparable to nonlinear Transformers, which have been shown to surpass the Correlational Statistical Query (CSQ) lower bound and achieve near information-theoretically optimal rate in previous works. Our analysis reveals the crucial role of the nonlinear gating mechanism in Mamba for feature extraction, highlighting it as the fundamental driver behind Mamba's ability to achieve both computational efficiency and high performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。