Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning

📄 arXiv: 2510.12026v2 📥 PDF

作者: Junsoo Oh, Wei Huang, Taiji Suzuki

分类: cs.LG, stat.ML

发布日期: 2025-10-14 (更新: 2025-10-15)

备注: 34 pages

💡 一句话要点

提出Mamba模型以解决低维目标的上下文学习问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 上下文学习 低维目标 非线性特征提取 Mamba模型 样本复杂度 机器学习 序列模型

📋 核心要点

  1. 现有方法在低维非线性目标函数的上下文学习能力上存在理论理解不足的问题。
  2. 论文提出Mamba模型，通过测试时特征学习直接从上下文示例中提取相关特征，实现高效的上下文学习。
  3. Mamba在样本复杂度上优于线性Transformer，并与非线性Transformer相当，展现出接近信息理论最优的学习速率。

📝 摘要（中文）

Mamba是一种新提出的线性时间序列模型，以其计算效率和强大的实证性能受到广泛关注。然而，对其基本机制的理论理解仍然有限。本研究通过聚焦于低维非线性目标函数定义的任务，提供了Mamba在上下文学习（ICL）能力的理论分析。我们证明了经过基于梯度的方法预训练的Mamba能够通过测试时特征学习实现高效的ICL，直接从上下文示例中提取相关方向。我们的分析揭示了Mamba中非线性门控机制在特征提取中的关键作用，强调了其在实现计算效率和高性能方面的根本驱动作用。

🔬 方法详解

问题定义：本论文旨在解决Mamba模型在低维非线性目标函数的上下文学习能力的理论理解不足的问题。现有方法在处理此类任务时，往往缺乏有效的特征提取机制，导致性能受限。

核心思路：论文的核心思路是通过测试时特征学习，使Mamba模型能够从上下文示例中直接提取相关特征，从而实现高效的上下文学习。这种设计使得模型在面对复杂的非线性目标时，依然能够保持高效性和准确性。

技术框架：整体架构包括预训练阶段和测试阶段。在预训练阶段，Mamba模型通过梯度方法进行训练；在测试阶段，模型利用上下文示例进行特征提取，进而实现对目标函数的学习。主要模块包括特征提取模块和非线性门控机制。

关键创新：最重要的技术创新点在于Mamba模型中引入的非线性门控机制，该机制在特征提取中起到了关键作用，使得模型能够有效地捕捉到低维非线性目标的特征。这一创新与现有方法的线性特征提取方式形成了本质区别。

关键设计：在模型设计中，关键参数设置包括学习率和梯度更新策略，损失函数采用了适应性损失函数以提高模型的学习效率。此外，网络结构中引入了非线性激活函数，以增强模型的表达能力。

📊 实验亮点

实验结果表明，Mamba模型在样本复杂度上显著优于线性Transformer，并与非线性Transformer的性能相当。具体而言，Mamba在处理低维非线性目标时，达到了接近信息理论最优的学习速率，展现出强大的实用性和有效性。

🎯 应用场景

该研究的潜在应用领域包括金融预测、医疗诊断和智能控制等需要处理低维非线性目标的场景。Mamba模型的高效上下文学习能力能够在这些领域中提供更快速和准确的决策支持，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.