Mamba通过在线梯度下降模拟实现上下文线性回归，理论证明其收敛速度快。
Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regression
- 基于Mamba结构设计新优化方法，突破非凸问题限制
- 实现指数级收敛，损失上限与Transformer相当
- 揭示Mamba以在线梯度下降机制学习，适合研究模型内部机理者
状态空间模型（SSMs），特别是Mamba，作为具有线性复杂度的长序列建模高效替代方案，近年来在实证中展现出与Transformer相当的上下文学习（ICL）能力，这是大模型的关键能力。然而，对Mamba ICL的理论理解仍不充分，尤其缺乏对基础任务如线性回归ICL的深入分析。本文研究Mamba在该任务上的训练动态，通过针对Mamba结构设计的新技术解决相关非凸优化问题，建立了向ICL解指数收敛的速率，并推导出与Transformer可比的损失界。关键发现是，Mamba能执行一种变体的在线梯度下降来学习上下文中的潜在函数，这与通常认为Transformer通过梯度下降模拟实现ICL的机制不同。理论结果经实验模拟验证。
原文摘要 · Abstract (English)
State-space models (SSMs), particularly Mamba, emerge as an efficient Transformer alternative with linear complexity for long-sequence modeling. Recent empirical works demonstrate Mamba's in-context learning (ICL) capabilities competitive with Transformers, a critical capacity for large foundation models. However, theoretical understanding of Mamba's ICL remains limited, restricting deeper insights into its underlying mechanisms. Even fundamental tasks such as linear regression ICL, widely studied as a standard theoretical benchmark for Transformers, have not been thoroughly analyzed in the context of Mamba. To address this gap, we study the training dynamics of Mamba on the linear regression ICL task. By developing novel techniques tackling non-convex optimization with gradient descent related to Mamba's structure, we establish an exponential convergence rate to ICL solution, and derive a loss bound that is comparable to Transformer's. Importantly, our results reveal that Mamba can perform a variant of \textit{online gradient descent} to learn the latent function in context. This mechanism is different from that of Transformer, which is typically understood to achieve ICL through gradient descent emulation. The theoretical results are verified by experimental simulation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。