Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning¶
会议: NeurIPS 2026
arXiv: 2510.12026
领域: 大模型预训练
关键词: 状态空间模型, Mamba, 上下文学习, 特征学习, 单指标模型
一句话总结¶
本文首次从理论上证明了预训练的 Mamba 能够通过非线性门控机制在测试阶段实现特征学习（Test-Time Feature Learning），在单指标模型上下文学习任务中突破相关统计查询（CSQ）下界，达到与非线性 Transformer 相当的近信息论最优样本复杂度。
研究背景与动机¶
在大语言模型与序列建模领域，以 Mamba 为代表的选择性状态空间模型（Selective State Space Models, SSM）凭借其线性时间复杂度和出色的长序列建模能力，迅速成为 Transformer 架构的强劲竞争者。然而，学界对 Mamba 的理论理解仍极度匮乏，尤其是在上下文学习（In-Context Learning, ICL）机制上，绝大多数理论分析均聚焦于具有二次注意力机制的 Transformer，揭示了 Transformer 能够模拟普通最小二乘法、预条件梯度下降甚至测试期特征学习。Mamba 缺乏全局注意力和显式的键值交互，其处理序列完全依赖带有选择性非线性门控的循环状态空间更新，这引发了一个核心疑问：Mamba 能否在测试期仅凭上下文提示词展现出如深度神经网络般的自适应特征学习能力。
这一疑问的核心矛盾在于：纯线性循环模型在表示能力上受限于固定核方法或线性特征空间，若目标任务由高维环境下的低维非线性单指标模型驱动，缺乏特征自适应能力的模型在测试阶段将遭遇维度灾难，其样本复杂度不可避免地由目标函数的最高多项式次数或相关统计查询（CSQ）下界的“信息指数”（Information Exponent）主导。以往研究虽在实验中观察到 Mamba 具备优良的 ICL 表现，但无法说明其究竟是退化为类似于核回归的浅层模式匹配，还是真正通过测试阶段的表征学习自适应提取了任务相关的低维特征子空间。
针对上述挑战，本文将高维统计学习中的高斯单指标模型引入 Mamba 的上下文学习分析中，深入探究循环递归与选择性门控在优化动力学中的本质作用。核心 idea：证明基于梯度预训练的单层 Mamba 结合 MLP，能够利用选择性非线性门控对上下文标签进行非线性变换，将信息指数压降至生成指数（Generative Exponent），从而在测试阶段自适应提取低维特征方向，以超越核方法与 CSQ 下界的样本复杂度实现高效上下文学习。
方法详解¶
整体框架¶
本文构建了一个端到端的两阶段学习理论分析框架。考虑高斯单指标模型 \(y = g_*(\langle \bm{\beta}, \bm{x} \rangle) + \zeta\)，其中未知特征向量 \(\bm{\beta} \in \mathbb{R}^d\) 位于未知 \(r\) 维子空间单位球面 \(S_r\) 上（\(r \ll d\)），\(\zeta\) 为方差极小的均匀噪声，目标是通过提示词 \([(\bm{x}_1, y_1), \dots, (\bm{x}_N, y_N), \bm{x}]\) 预测查询样本 \(\bm{x}\) 对应的标签 \(y\)。模型由三部分构成：基于二阶多项式基的紧凑输入嵌入层、单层 Mamba 递归选择层以及双层多层感知机（MLP）预测头。
在整个预测流中，输入嵌入将每个上下文对转换为特征向量 \(\bm{z}_l = [\phi(\bm{x}_l)^\top, y_l]^\top\)，Mamba 层通过输入依赖的选择性门控与状态转移矩阵将上下文历史压缩并演化为隐状态，提取出上下文交互后的聚合输出；最后，MLP 层在归一化的 Mamba 输出上逼近未知非线性链接函数 \(g_*\)。该过程构成了一个完全自洽的计算流。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入提示词 Prompt<br/>[(x_1, y_1), ..., (x_N, y_N), x]"] --> B["输入嵌入映射<br/>二阶多项式扩展 φ(x) 与标签拼接"]
B --> C["门控选择性状态更新<br/>非线性衰减与输入依赖增益计算"]
C --> D["测试期特征提取<br/>门控降解信息指数至生成指数"]
D --> E["双层 MLP 链接拟合<br/>逼近单指标非线性映射函数 g_*"]
E --> F["输出预测标量<br/>y_hat ≈ g_*(<β, x>)"]关键设计¶
1. 输入嵌入映射：二阶多项式基的紧凑扩展 针对高维输入中低维指标特征与非线性链接函数的交互，若仅采用原始一阶线性输入嵌入，当链接函数为偶函数时，根据 Stein 引理，一阶相关性将完全消失导致无法提取特征。本文设计了二阶多项式特征映射 \(\phi: \mathbb{R}^d \rightarrow \mathbb{R}^{\tilde{d}}\)（其中 \(\tilde{d} = 1 + d + \frac{d(d+1)}{2}\)）： $$ \phi(\bm{\theta}) = \left[ 1, \bm{\theta}, \frac{\bm{\theta} \odot \bm{\theta} - \bm{1}_d}{\sqrt{2}}, \bm{\theta}[1]\bm{\theta}[2], \dots, \bm{\theta}[d-1]\bm{\theta}[d] \right]^\top $$ 输入矩阵 \(\bm{Z} \in \mathbb{R}^{(\tilde{d}+1) \times (N+1)}\) 将上下文对中的 \(\phi(\bm{x}_j)\) 与标签 \(y_j\) 拼接，并在最后一个查询 token 处将标签置零。该设计为后续单层 Mamba 提取低维非线性特征提供了正交基底支撑，避免了以往方法中依赖深层堆叠 GLU 模块的繁琐构造。
2. 门控选择性状态更新：连续离散化下的等价核形式 Mamba 层的核心在于其连续状态空间选择机制 \(\bm{h}_l = \overline{\bm{A}}_l \bm{h}_{l-1} + \overline{\bm{B}}_l \bm{z}_l\) 以及输出映射 \(\bm{o}_l = \bm{C}_l^\top \bm{h}_l\)。论文在结构化状态空间矩阵 \(\bm{A} = -\bm{I}\) 下，采用输入依赖参数 \(\Delta_l = \mathrm{softplus}(\bm{w}^\top \bm{z}_l + b)\) 进行欧拉离散化，推导出其输出在查询位置的精确闭式表达式： $$ \bm{o}{N+1} = \sum}^N G_{j, N+1}(\bm{Z}) \bm{z}_j \bm{zj^\top \bm{W}B^\top \bm{W}C \bm{z} $$ 其中状态衰减与选择因子由非线性门控累积乘积主导：\(G_{j, N+1}(\bm{Z}) = \sigma(\bm{w}^\top \bm{z}_j + b) \prod_{k=j+1}^{N+1}(1 - \sigma(\bm{w}^\top \bm{z}_k + b))\)。在参数简化设定 \(\bm{W}_B^\top \bm{W}_C = \mathrm{diag}(\bm{\gamma}, 0)\) 下，最后维度的输出精准体现为上下文特征项与查询特征的加权内积： $$ \mathsf{Mamba}(\bm{Z}; \bm{\gamma})[\tilde{d}+1, N+1] = \sum^N G)) $$ 这种形式表明 Mamba 巧妙地将非线性门控与类似线性注意力的双线性内积融为一体。}(\bm{Z}) y_j \phi(\bm{x}_j)^\top (\bm{\gamma} \odot \phi(\bm{x
3. 测试期特征提取：非线性门控实现信息指数的本质压降 这是本文最具突破性的理论机理。在以往基于核回归或线性 Transformer 的设定中，由于模型仅对标签进行线性聚合，根据正交 Hermite 多项式性质，当且仅当采样数达到 \(d^{\Theta(\mathrm{deg}(g_*))}\) 或 \(d^{\Theta(\mathrm{ie}(g_*))}\) 阶时才能捕获有效梯度。而在 Mamba 中，通过设置适当的负偏置 \(b = C_b \log d\)，门控项中的 \(\sigma(y_j/\rho + b)\) 对连续标签执行了非线性的阈值变换。这一变换激活了数据的统计查询（SQ）能力，使得原本高阶的多项式期望项被有效“折叠”，将有效信息指数强制降低至由偶/奇对称性决定的生成指数： $$ \mathrm{ie}(g_* \sigma(g_/\rho + b)) = \mathrm{ge}(g_) \in {1, 2} $$ 其直接结果是，在第一阶段对 \(\bm{\gamma}\) 完成单步梯度更新后，预训练参数 \(\bm{\gamma}^*\) 在真实特征子空间 \(\mathcal{I}\) 内部的坐标幅值远大于外部噪声坐标，并在测试阶段对任意新给出的上下文提示词实现如下低维自适应表征： $$ N^{-1} \mathsf{Mamba}(\bm{Z}; \bm{\gamma}^) \approx P_1 + P_2 \left( \frac{\langle \bm{\beta}, \bm{x} \rangle}{r} \right)^{\mathrm{ge}(g_)} $$ 这在理论上严格确立了 Mamba 在推理阶段无须反向传播即可完成“测试期特征学习”（Test-Time Feature Learning）。
4. 双层 MLP 链接拟合：基于稀疏多项式的函数近似 经过 Mamba 层处理后，高维提示词已被浓缩为标量形式的一阶或二阶投影 \((\langle \bm{\beta}, \bm{x} \rangle / r)^{\mathrm{ge}(g_*)}\)。第二阶段通过优化包含 \(m\) 个隐藏单元的双层 ReLU MLP，利用单项式基展开技术逼近原链接函数 \(g_*\)。得益于生成指数的简化，多项式阶数被收敛，使得有限宽度 \(m = \tilde{\Omega}(r^{4 \mathrm{ge}(g_*)})\) 的两层 MLP 即可对复合函数实现任意精度的拟合，最终通过 Rademacher 复杂度分析得出泛化误差上界。
损失函数 / 训练策略¶
论文采用理论可解的二阶段梯度优化预训练策略（Algorithm 1）： - Stage I（Mamba 层单步梯度更新）：固定未训练的初始 MLP 权重 \(\bm{u}^{(0)} = m^{-1}\bm{1}_m\)，在 \(T_1\) 个预训练任务（上下文长度为 \(N_{\mathrm{pt}}\)）上以均方误差加 \(\ell_2\) 正则化对 Mamba 层特征权重 \(\bm{\gamma}\) 进行单步梯度下降： $$ \bm{\gamma}^* \leftarrow \bm{\gamma}^{(0)} - \eta \nabla_{\bm{\gamma}} \left( L_1(\bm{\gamma}^{(0)}, \bm{u}^{(0)}, \bm{v}^{(0)}, \bm{a}^{(0)}) + \frac{\lambda_1}{2} |\bm{\gamma}|^2 \right) $$ 利用 MLP 的 ReLU 导数跳变特性与门控函数的非线性相互作用，单步更新即可精准捕获低维子空间分布。 - Stage II（MLP 凸优化拟合）：固定 \(\bm{\gamma}^*\)，重新采样内层随机权重 \(\bm{v}^* \sim \mathrm{Unif}(\{\pm 1\}^m), \bm{a}^* \sim \mathrm{Unif}([-1, 1]^m)\)，仅对凸外层权重 \(\bm{u}\) 在后续 \(T_2\) 个任务上做正则化极小化求解： $$ \bm{u}^* \leftarrow \arg\min_{\bm{u}} \left( L_2(\bm{\gamma}^, \bm{u}, \bm{v}^, \bm{a}^*) + \frac{\lambda_2}{2} |\bm{u}|^2 \right) $$ 该凸优化保证了解的全局最优性与有界范数 \(\|\bm{u}^*\| \le \tilde{\mathcal{O}}(r^{3\mathrm{ge}(g_*)/2} m^{-1/2})\)。
实验关键数据¶
主实验：不同算法与架构的样本复杂度理论对比¶
下表汇总了高维高斯单指标模型下，不同测试期回归方法与各典型预训练架构在进行上下文学习时的样本复杂度理论保证。
| 架构 / 学习机制 | 理论依据类别 | 预训练样本复杂度 (\(N_{\mathrm{pt}}\)) | 测试期样本复杂度 (\(N_{\mathrm{test}}\)) | 特征学习机制 |
|---|---|---|---|---|
| 核方法（测试期直接回归） | Kernel Methods | — | \(d^{\Theta(\mathrm{deg}(g_*))}\) | 无特征学习（固定基底） |
| CSQ 算法类（如两层浅层网络） | Correlational SQ | — | \(d^{\Theta(\mathrm{ie}(g_*))}\) | 仅依赖一阶相关性 |
| SQ 算法类（带非线性变换） | Statistical Query | — | \(d^{\Theta(\mathrm{ge}(g_*))}\) | 标签非线性折叠 |
| Linear Transformer (Oko et al., 2024) | ICL 线性注意力 | \(d^{\Theta(\mathrm{ie}(g_*))}\) | \(r^{\Theta(\mathrm{deg}(g_*))}\) | 仅子空间自适应，依赖核回归 |
| Softmax Transformer (Nishikawa et al., 2025) | ICL Softmax 注意力 | \(d^{\Theta(\mathrm{ie}(g_*))}\) | \(r^{\Theta(\mathrm{ge}(g_*))}\) | Softmax 测试期特征学习 |
| Mamba（本文模型） | ICL 门控线性状态空间 | \(\bm{d^{\Theta(\mathrm{ge}(g_*))}}\) | \(\bm{r^{\Theta(\mathrm{ge}(g_*))}}\) | 非线性门控驱动测试期特征学习 |
消融与实证仿真：维度与预训练上下文长度敏感性¶
论文在标准架构（12层 Mamba 与 6层 GPT-2 Transformer，嵌入维度 256）及合成单指标任务（\(g_*(z) = \mathrm{He}_3(z)/\sqrt{6}\)，基准 \(d=32, r=8, N_{\mathrm{pt}}=64\)）上进行了数值验证。
| 实验对比配置 | 变化变量 | Mamba 表现（测试预测均方误差） | Transformer 表现（测试预测均方误差） | 核回归（RBF，内禀子空间输入） | 理论机制验证结论 |
|---|---|---|---|---|---|
| 基准 vs 维度缩减 | \(d=32 \rightarrow d=16\) | 误差几乎无变化（收敛至 \(\approx 0.05\)） | 误差几乎无变化（收敛至 \(\approx 0.05\)） | 明显高于神经网络模型（\(\approx 0.35\)） | 证明预训练完全消除了环境维度 \(d\) 的影响 |
| 内禀维度增加 | \(r=8 \rightarrow r=16\) | 误差轻微上升，随样本数增加平缓下降 | 误差上升趋势与 Mamba 完全一致 | 性能显著恶化 | 验证复杂度仅由内禀维度 \(r\) 决定，且大幅领先核方法 |
| 短预训练上下文消融 | \(\bm{N_{\mathrm{pt}}=64 \rightarrow 16}\) | 维持强适应性（误差稳定收敛） | 性能出现大幅恶化（泛化崩溃） | — | 实证验证 Mamba 在小预训练数据下的样本效率优势 |
关键发现¶
- 非线性门控突破理论屏障：Mamba 在测试期的表现不仅打破了依赖于全维度 \(d\) 的维数灾难，而且测试样本复杂度完全由内禀维度 \(r\) 与生成指数 \(\mathrm{ge}(g_*)\) 控制，成功跨越了 CSQ 算法类所面临的下界屏障。
- 预训练样本效率甚至优于 Transformer：在更苛刻的预训练环境（如 \(N_{\mathrm{pt}}=16\) 的短上下文）下，Transformer 因 Softmax 机制依赖高阶多项式拟合而表现出明显的样本饥渴，而 Mamba 凭借二阶嵌入与非线性门控的协同作用，展现出更稳健的收敛特性，在理论上其预训练样本复杂度降至 \(d^{\Theta(\mathrm{ge}(g_*))}\)。
亮点与洞察¶
- 将选择性门控形式化为隐式非线性统计查询：首次揭示了 Mamba 的 \(\Delta\) 软间隔选择离散化所带来的非线性门控，在功能上等价于对标签进行非线性变换，将高阶信息指数转化为低阶生成指数，从根本上解释了为何线性状态空间不会退化为朴素的线性核回归。
- 解耦环境维度与内禀维度：证明了预训练学习全局特征子空间、推理期学习局部特征方向的“两阶段适应”范式在状态空间模型中同样成立，填补了线性时间序列模型在非线性 ICL 领域的理论空白。
- 可复用的理论分析技巧：将基于 Hermite 正交展开分析高斯单指标模型的方法成功拓展到了循环门控递归网络，为未来分析包含门控注意力的线性混合架构（如 RWKV、GLA、xLSTM）提供了严谨的数学分析模版。
局限与展望¶
- 输入嵌入假设的限制：当前证明依赖于二阶显式多项式基的预设嵌入 \(\phi(\bm{x})\)；若使用标准一阶线性嵌入且目标函数为偶函数，单层架构将失效，未来需向多层自学习特征嵌入展开进一步推导。
- 遗忘机制被简化假设：为实现闭式分析，证明中假定长程历史遗忘项在理论渐近尺度下极小。然而在存在噪声异常点或离群任务时，Mamba 的动态遗忘已被证实至关重要，二者的耦合动态仍需深入探索。
相关工作与启发¶
- vs Oko et al. (2024)：Oko 等人证明了线性 Transformer 能够学习单指标模型，但其测试期样本复杂度受限于多项式次数 \(r^{\Theta(\mathrm{deg}(g_*))}\)，本质上仍属于核方法范畴；本文证明了 Mamba 凭借非线性门控实现了测试期特征学习，复杂度降至 \(r^{\Theta(\mathrm{ge}(g_*))}\)。
- vs Nishikawa et al. (2025)：Nishikawa 等人通过 Softmax 的指数非线性证明了非线性 Transformer 的测试期特征学习；本文则证明了无全局自注意力、纯线性时间复杂度的 Mamba 能够通过截然不同的“递归状态更新+门控”机制达成同等甚至更优的理论保证。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从理论上证明了 Mamba 在测试阶段具备特征学习能力，填补了 SSM 在非线性 ICL 机制上的理论空白。
- 理论充分度: ⭐⭐⭐⭐⭐ 数学证明严密完整，包含完整的优化动力学、泛化界与 Rademacher 复杂度分析，并附有严谨的实证对照。
- 写作质量: ⭐⭐⭐⭐⭐ 结构清晰流畅，问题动机明确，定理与证明主线层层递进。
- 价值: ⭐⭐⭐⭐⭐ 为非 Transformer 架构的上下文学习能力提供了坚实的基础理论支撑，对高效基础模型的设计与理解具有深远意义。