Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions¶

会议: NeurIPS 2025
arXiv: 2510.18638
代码: 无
领域: 理论机器学习 / Transformer 理论
关键词: 上下文学习, 马尔可夫链, 线性自注意力, NP-hard, 多目标优化

一句话总结¶

从优化理论角度分析 Transformer 学习马尔可夫动态函数的 ICL 能力：推导单层线性自注意力的全局最优解（闭式表达），证明从扩展参数空间恢复 Transformer 参数是 NP-hard 的，并揭示多层 LSA 等价于预条件多目标优化。

研究背景与动机¶

现有痛点¶

现有痛点：领域现状：In-Context Learning（ICL）是 Transformer 的核心能力之一，即模型通过 prompt 中的输入-输出示例学习新任务。现有 ICL 的理论研究主要集中在 i.i.d. 高斯输入的线性回归任务，已证明 Transformer 可隐式执行梯度下降。

然而，现实中许多任务涉及结构化时序数据——输入之间存在时间依赖关系。本文聚焦一个更具挑战性的设定：从多条马尔可夫链轨迹中学习共享的转移动态，给定新的查询序列预测下一个 token。这个设定引入了两个新挑战： 1. 输入不再是 i.i.d. 的，序列内部存在马尔可夫依赖结构 2. 输入-输出关系不再是确定性线性函数，而是随机的转移概率

核心问题：Transformer 能否有效学习这类动态驱动的函数？其优化和表达力的理论极限在哪里？

方法详解¶

整体框架¶

研究框架基于线性自注意力（LSA）模型：移除 softmax 非线性，合并 key/query 矩阵为 \(Q = W_k^T W_q\)。输入为 \(n\) 条长度 \(d+1\) 的二元马尔可夫链，构成嵌入矩阵 \(Z_0\)，目标是预测查询序列的最后一个 token。训练目标为均方误差的总体损失。

关键设计¶

重参数化与凸化：定义映射 \(\phi(b,A) = X \in \mathbb{R}^{dm}\)（\(m = \frac{(d+2)(d+1)}{2}\)），将原始非凸的 LSA 参数空间映射到扩展空间。在扩展空间中，目标函数 \(\tilde{f}(X)\) 是严格凸的（Lemma 3.1），且全局最小值有闭式解 \(X^* = H^{-1}b\)（Lemma 3.2）。其中 \(H\) 矩阵的结构由马尔可夫转移核的幂次决定。
长度-2 马尔可夫链的全局最优解（Theorem 3.3, 3.4）：对最简情形（长度 2 的马尔可夫链），作者推导出 LSA 参数 \((P,Q)\) 的解析全局最优解。关键发现：最优解中的参数全部非零，形成密集结构——这与 i.i.d. 线性回归中的稀疏结构形成鲜明对比。这反映了马尔可夫输入的时间依赖性使得优化解必须同时编码多个统计关系。
NP-hard 性证明（Theorem 3.5）：对于一般长度 \(d+1\) 的马尔可夫链，虽然扩展空间中的最优 \(X^*\) 可解析求解，但从 \(X^*\) 恢复原始 Transformer 参数 \((b,A)\) 是 NP-hard 的。证明通过归约至双线性可分离性问题（已知 NP 完全）。这揭示了单层 LSA 的根本表达力限制：最优解存在于扩展空间但可能无法被 Transformer 实现。
多层 LSA 的多目标优化解释（Theorem 4.1）：\(L\) 层 LSA 的前向传播等价于一个预条件多目标优化过程，同时最小化平方损失和多个线性目标。具体地，\(w_{l+1}^T = w_l^T - b_l^T(\nabla R_1(w_l)\bar{A}_l + \nabla R_2(w_l)[\cdot])\)，其中 \(R_1\) 包含平方损失和状态交互项，\(R_2\) 强调最后一个参数 \(w_d\) 的作用。

损失函数 / 训练策略¶

总体损失 \(f(\{P_l, Q_l\}) = \mathbb{E}[(\mathtt{TF}_L(Z_0) - y_{n+1})^2]\)，期望遍历所有可能的马尔可夫转移核和初始分布。

实验关键数据¶

主实验（多层 LSA 准确率，不同状态空间大小）¶

| 模型深度 | \(|\mathcal{S}|=2\) | \(|\mathcal{S}|=3\) | \(|\mathcal{S}|=4\) | |---------|------|------|------| | 1 层 LSA | ~0.50 (随机猜测) | ~0.55 | ~0.30 | | 2 层 LSA | ~0.85 | ~0.80 | ~0.60 | | 3 层 LSA | ~0.95 | 0.98 | ~0.85 |

单层 LSA 对所有状态空间大小表现接近随机猜测，验证了 NP-hard 理论的预测；增加层数后准确率显著提升。

消融实验（多目标优化行为验证）¶

层数区间	Generational Distance (GD)	R1 目标值	R2 目标值	说明
前几层	下降	改善	改善	初始阶段有效逼近 Pareto 前沿
中间层	最小	最优	最优	多目标优化效果最佳
深层	上升	恶化	恶化	受限参数空间导致退化

10 层 LSA 在马尔可夫数据上训练后，前向传播确实执行多目标优化，验证了 Theorem 4.1。

关键发现¶

单层 LSA 无法有效学习马尔可夫动态函数，即使最优解已知也无法被实现（NP-hard gap）
增加层数可显著改善性能，与多目标优化解释一致
马尔可夫输入导致最优解的密集结构，与 i.i.d. 输入的稀疏结构形成对比

亮点与洞察¶

理论贡献深刻：NP-hard 性证明揭示了 Transformer 在结构化输入上的根本表达力限制，这是 ICL 理论研究的重要进展
多目标优化视角新颖：将多层 Transformer 前向传播解释为同时优化多个目标（不仅仅是平方损失），提供了对深度必要性的理论理解
i.i.d. vs 马尔可夫的对比清晰：精确刻画了输入结构如何影响最优解的数学形式

局限与展望¶

仅分析线性自注意力，未涉及 softmax attention、MLP、位置编码等实际组件
限于一阶二元马尔可夫链，高阶和大状态空间的情况未覆盖
理论结果与实际 Transformer 训练之间仍有较大 gap
缺少实际 NLP/时序任务上的实验验证

评分¶

新颖性: ⭐⭐⭐⭐⭐ NP-hard 性证明和多目标优化解释均为首次提出
实验充分度: ⭐⭐⭐ 实验以验证理论为主，规模和多样性有限
写作质量: ⭐⭐⭐⭐ 数学推导严谨，但符号较重
价值: ⭐⭐⭐⭐ 对理解 ICL 机制和 Transformer 表达力提供了重要理论视角