分叉附近的状态空间NTK坍缩

# 分岔附近的状态空间NTK坍缩

**来源：** https://arxiv.org/html/2605.12763

**作者** James Hazelden · [email protected] · 华盛顿大学  
**和** Eric Shea-Brown · [email protected] · 华盛顿大学

###### 摘要

在随时间展开的任务中，丰富的特征学习通常要求模型经历分岔，即底层模型动态的定性变化。我们通过经验状态空间神经正切核（sNTK）发展了这些转变附近梯度下降的局部理论。我们的核心发现是：分岔既主导又简化了学习动态——在分岔附近，我们可以将 sNTK 约化为一个秩一算子，对应于经典规范形系统中的学习，从而为局部学习几何提供可解析处理的描述，即使对于高维循环系统也是如此。具体而言，我们给出了一种将 sNTK 分解为分岔相关通道和残差通道的程序，表明在常见的余维1分岔附近，相关通道是一个高度放大的秩一算子。这种放大导致分岔通道主导整个 NTK。因此，分岔局部地扭曲了学习景观，将梯度下降引导到少数关键的动态方向，并使附近的核和损失几何变得可由经典规范形预测。我们在一个学生-教师循环神经网络中说明了这一点：第一个学习到的分岔与 NTK 有效秩的急剧坍缩同时发生，并出现了一个主导参数方向，其受限 NTK 与标量叉形规范形预测的景观紧密匹配。最后，我们表明低秩自然梯度方法可以解决分岔附近产生的学习不稳定性，且相对于 SGD 的开销非常小。

## 1 引言

梯度下降（GD）通过重塑其潜动态直到生成的轨迹解决任务来训练动力学模型。在丰富的特征学习机制中，这通常需要的不仅仅是调整输出：学习必须创建、消除或重新组织不动点和相关的动态主题。经典动力系统理论告诉我们，这种定性变化通过局部分岔发生（guckenheimer1983nonlinear）。先前的工作表明，一些分岔与使用 ReLU 激活函数的循环网络中损失的大幅下降同时发生，表明分岔在优化和动态中都是重要事件（eisenmann2023bifurcations）。然而，该机制与特定架构和分岔类型绑定，并不能更普遍地解释 GD 在动态转变附近的行为。更广泛地说，对于给定的参数化和损失函数，GD 会穿越哪些分岔，回避哪些分岔，以及参数更新在分岔集附近如何表现？我们通过经验状态空间神经正切核（sNTK）来研究这些问题，sNTK 是全局参数到状态雅可比矩阵的 Gram 算子。最近的工具使得 sNTK 对于有限循环模型可解释且可计算（hazelden2026globalempiricalntkselfreferential），使其成为研究分岔附近学习的自然视角。我们表明，在分岔附近，学习变得有效低维。对于余维一分岔，sNTK 约化为一个近似秩一的通道，因此 GD 被引导到一小组动态校正中。这个通道被底层动态放大，使得优化在转变附近变得僵硬且高度各向异性。因此，分岔不仅是动态事件，也是优化瓶颈。此外，这种约化为分岔附近的学习提供了一个简单的解析模型：局部动态由相应规范形的低秩 sNTK 支配，允许在一个简单的低维设置中进行分析，该设置仍然与完整模型在转变附近的行为紧密匹配。

总之，我们的**贡献**是：

- 在局部分岔附近，通用模型的 sNTK 允许加性分解为分岔相关项和残差项，前者是低秩的。
- 对于余维一分岔，相应的规范形预测了转变附近的 sNTK 放大，诱发了有效秩一的局部学习几何，该几何与完整模型中经验可计算的 sNTK 紧密匹配。
- 在学生-教师 RNN 中，叉形分岔与 sNTK 有效秩的急剧下降同时发生，表明低秩自然梯度校正可以稳定此类转变附近的训练。

## 2 学生-教师 RNN 中由分岔引起的 NTK 坍缩

*（见标题下图片描述）*

**图 1：** 学生-教师 RNN 中的 NTK 坍缩。在 SGD 过程中，我们测量 (A) 损失，(B) sNTK 的稳定秩，以及 (C) 学生权重的谱半径。(D) 比较最终的读出动态。A–C 中的虚线对应一个叉形分岔，如 (G) 所示，对应于损失突然下降和 sNTK 有效秩坍缩至 1。(E) 展示了该分岔附近的局部 sNTK 范数放大，与叉形规范形预测相匹配。

我们从一个学生-教师 RNN 开始，该 RNN 训练于一个具有二维读出动态的任务（详见附录 B (https://arxiv.org/html/2605.12763#A2)）。教师表现出一个塑造采样轨迹的不动点（FP）结构，学生必须学习重现这些轨迹。图 1 (https://arxiv.org/html/2605.12763#S2.F1) 总结了结果。面板 D 显示了教师动态，由四个稳定 FP 和五个不稳定 FP 组成。最初，学生表现出坍缩到单个 FP 的动态，因此需要多个分岔才能重现教师动态。面板 A–C 中的虚线标记了第一个学习到的分岔，一个叉形（如面板 G 所示），它与损失突然下降同时发生（面板 A），与 eisenmann2023bifurcations 一致。在面板 B 中，我们计算了状态空间 NTK（sNTK）的有效秩，下面将详细描述。在第一个分岔处，该秩急剧坍缩到一，然后在训练后期随着进一步的分岔再次扩大。本文的其余部分分析了这种坍缩现象及其后果。特别是，分岔附近的局部 GD 放大景观（面板 E）很好地由与叉形规范形相关的秩一 sNTK 景观预测（图 2 (https://arxiv.org/html/2605.12763#S3.F2)B），后者可以被精确表征。与表现出分岔附近严格单调增益的更简单的稳定性翻转分岔相比，叉形似乎自我调节，景观在分岔附近达到峰值然后衰减。在附录 C (https://arxiv.org/html/2605.12763#A3) 中，我们展示了一个秩一自然梯度校正器为同一任务产生了更平滑的损失曲线，有效地中和了 NTK（因此 Fisher 信息）的低秩不稳定贡献，相比于 SGD 开销非常小。总的来说，这个实验表明，在有限大小的非线性循环网络中，(1) 学习在分岔附近可能变得高度低维，(2) 这种行为的动力学系统理论中的解析可处理规范形很好地捕捉，以及 (3) 低秩自然梯度可以更稳定地训练此类模型，且额外开销很小。

## 3 规范形使机制显式化

上述分解将局部学习几何约化为分岔相关通道 sNTK\(_g\) 和残差项。关键问题是 sNTK\(_g\) 在临界点附近是否被充分放大，以支配整个 NTK。为了研究这一点，我们使用动力系统理论中的规范形：描述分岔附近动态的低维多项式模型。对于通用系统，一个平滑的坐标变换将分岔附近的局部动态与相应的规范形一致。因此，研究规范形的 NTK 可以揭示也出现在完整网络模型中的局部学习行为。具体而言，我们研究一维规范形 \(h_{t+1} = f(h_t, g)\)，其中 \(g \to g^*\) 诱导一个余维一分岔。在此设置中，相关 NTK 通道是秩一的，因此其范数直接测量沿临界动态方向的学习强度。图 2 (https://arxiv.org/html/2605.12763#S3.F2) 显示，在所有具有代表性的余维一分岔中，该通道在临界点附近被强烈放大。为简化起见，我们这里重点关注两种情况：稳定性翻转和叉形。

*（见标题下图片描述）*

**图 2：** 余维一分岔的秩一 sNTK 放大。我们绘制了 sNTK = \((D_g h)(D_g h)^T\) 的范数，初始条件 \(h_0\) 均匀采样自 \([-0.05, 0.05]\)（蓝色）和 \([-0.1, 0.1]\)（橙色），使用 \(T=30\) 个时间步。在所有情况下，范数在分岔点 \(g^* = 1\)（虚线）附近被强烈放大。稳定性翻转在过临界点后表现出单调增长，而非线性分岔自我调节并显示出峰值放大。

#### 一个线性示例：稳定性翻转

我们从简单线性系统 \(h_{t+1} = g h_t\) 开始。这并非通常意义上的规范形，但它是随 \(g\) 变化而动态发生定性变化的最简单系统。事实上，我们可以显式推导相应的 NTK。记 \(T\) 为底层模型的时间范围，
\[
\| \mathrm{sNTK}_g \|_2 \propto \sum_{t=0}^{T-1} (t+1)^2 g^{2t}.
\]
对于 \(|g| < 1\)，这表现为 \(T^3/3\)；而对于 \(|g| > 1\)，行为变为 \(T^2 g^{2T} / (1-g^2) = O(g^{2T})\)，特征为极端无界的爆炸。因此，在分岔局部，NTK 坍缩为有效秩一并具有巨大范数，由 sNTK\(_g\) 主导，并且随着 GD 进一步推入不稳定区域，这种情况会继续恶化。

#### 叉形分岔

对于叉形规范形 \(h_{t+1} = g h_t - h_t^3\)，出现了相同的不稳定性，但非线性逃逸到稳定分支切断了无界增长（图 2 (https://arxiv.org/html/2605.12763#S3.F2) B）。对于 \(0<g<1\)，范数再次呈现 \(O(T^3/3)\) 行为。对于 \(g>1\)，放大达到峰值然后衰减，因为轨迹被阻尼到额外的稳定 FP \(\pm \sqrt{g-1}\) 上。因此，与稳定性翻转不同，叉形在分岔后表现出受调节的而非严格指数的放大。这正是学生-教师 RNN 中观察到的行为（比较图 1 (https://arxiv.org/html/2605.12763#S2.F1) E 和图 2 (https://arxiv.org/html/2605.12763#S3.F2) B）。

#### 要点

分岔强烈影响学习，因为它们诱导出秩一的 NTK 项，这些项可以主导局部学习几何。一旦发生这种情况，GD 变得高度各向异性：与状态空间中分岔相关方向对齐的误差信号导致底层模型状态的大变化，而正交信号产生小得多的更新。因此，即使是标量规范形模型也能准确预测相应完整高维模型的局部 GD 景观。

## 4 讨论

我们已经证明，分岔不仅对应于模型动态的重大变化，也对应于学习中的强烈且可预测的特征。在余维1转变附近，经验状态空间 NTK 坍缩到单个临界通道上，使得即使在大的参数空间中梯度下降也有效低维。这为通过规范形建立分岔附近学习的解析理论打开了大门。它也暗示了 sNTK 坍缩是超出经典动力系统模型（例如，Transformer 或输入驱动的 MLP）的模型中临界特征学习事件的一个更普遍的特征。当前工作的主要局限性在于它是局部性和解释性的，而非完全预测性的，但这同时指出了自然的下一步：开发在 GD 学习期间检测分岔的实际方法，并利用这些信息更稳定地优化模型。

**致谢** 我们感谢 Alexander Hsu 建议使用低秩自然梯度（yang2020sketchy）。

## 参考文献

（此处省略参考文献列表，因为原文未提供具体内容）

## 附录 A 推导

### A.1 分岔坐标下的局部分解

如同 (hazelden2026globalempiricalntkselfreferential)，sNTK 算子可以写为
\[
\mathrm{sNTK} = \mathcal{P} \mathcal{K} \mathcal{P}^T
\]
这源于对动态 \(h_t = f(h_{t-1}, \theta)\) 的隐式重参数化，形式为 \(\mathcal{F}(h, \theta) = h - f(T_\downarrow h, \theta) = 0\)，其中 \(T_\downarrow\) 是将时间减一的线性算子，\(h \in \mathbb{R}^{B \times T \times N}\) 对应于一个批次大小为 \(B\)、模拟 \(T-1\) 个时间步（在下面 RNN 任务中具体值为 \(B, T, H\)，见附录 B (https://arxiv.org/html/2605.12763#A2)）的所有样本轨迹的隐状态。这里，\(\mathcal{P} = (D_h \mathcal{F})^{-1}, \mathcal{K} = (D_\theta \mathcal{F})(D_\theta \mathcal{F})^*\)。关键在于，参数变化只改变 \(\mathcal{K}\)。如果 \(\phi: \mathbb{R}^m \to \mathbb{R} \times \mathbb{R}^{m-1}\) 是一个坐标微分同胚 \(\phi: \theta \to (g, R)\)，局部于一个余维一分岔 \(\theta^* = (g^*, R^*)\)，那么由于 \(\theta = \phi^{-1}(g, R)\)，根据链式法则，
\[
D_\theta \mathcal{F}(\theta^*) = (D_g \mathcal{F}(g^*), D_R \mathcal{F}(R^*)) \cdot D_\theta \phi(\theta^*)
\]
最后，局部地，我们可以选择 \(\phi\) 在 \(\theta^*\) 处为局部等距，使得 \(D_\theta \phi\) 在 \(\theta^*\) 处为单位矩阵，从而在 \(\theta^*\) 处得到
\[
\mathcal{K} = \mathcal{K}_g + \mathcal{K}_R = D_g \mathcal{F}(g^*) D_g \mathcal{F}(g^*)^T + D_R \mathcal{F}(R^*) D_R \mathcal{F}(R^*)^T
\]
因此，
\[
\mathrm{sNTK} = \mathcal{P}(\mathcal{K}_g + \mathcal{K}_R) \mathcal{P}^T = \mathrm{sNTK}_g + \mathrm{sNTK}_R
\]
在这种局部等距坐标变换 \(\theta \mapsto (g, R)\) 下，得到了 NTK 向一个分岔相关秩一算子和一个秩 \(m-1\) 残差算子的清晰分离。当然，同样的程序可以应用于更高秩的分岔，其中 \(g \in \mathbb{R}^k\)。

## 附录 B 学生-教师任务细节

学生和教师都是使用 tanh 激活的普通 RNN (hochreiter1997long)。我们使用 SGD 训练模型，学习率 \(\eta = 5 \cdot 10^{-3}\)，共 35,000 次 GD 迭代，没有使用动量或梯度裁剪。批次大小为 256，每个批次条目对应一个用于评估学生和教师的不同的初始条件。在这种情况下，每个模型有 64 个隐藏神经元。选择隐神经元 \(h_0\) 和 \(h_1\) 作为模型的读入和读出，即输入和输出权重是固定的，并且学生和教师相同。