Training Deep Physics-Informed Kolmogorov-Arnold Networks

本文提出了残差门控自适应 KANs（RGA KANs），这是一种结合了基函数无关初始化方案与残差门控的新型架构，旨在克服深度物理信息 Kolmogorov-Arnold 网络在训练中的不稳定性和发散问题，从而在多种偏微分方程基准测试中实现卓越的准确性与稳定性。

要点

想象一下，你正试图教一台计算机解决复杂的物理谜题，比如预测热量如何在金属板中扩散，或者水流如何绕过船只。多年来，这项工作的标准工具是一种被称为神经网络（具体来说是物理信息神经网络，即 PINN）的 AI。可以将这些网络想象成一群试图走出迷宫的工人。

最近，一种更聪明的、被称为 KAN（Kolmogorov–Arnold Network）的新型工人被引入了。KAN 就像是那些在工作时可以根据需要更换自己工具的工人，这使得它们具有极高的灵活性和准确性。然而，这里有一个问题：当你试图构建一个非常深的 KAN 团队（即具有许多层工人的“深层架构”）时，这个团队往往会崩溃。它们会变得混乱，信号会丢失，并停止学习。这就像试图通过 20 个人传递一个秘密，到最后，声音已经变成了杂音。

这篇论文介绍了两个主要的修复方案，使深层 KAN 团队能够可靠地工作。

1. “类 Glorot”初始化：设定正确的音量

问题： 当你启动一个新的 KAN 团队时，你必须为他们分配初始的“音量”（在数学上，即他们的初始权重）。旧的方法就像是在盲目猜测音量旋钮；有时声音太小（信号消失），有时声音太大（信号爆炸）。这使得训练深层团队变得不可能。

解决方案： 作者发明了一种新的设定初始音量的方法，称为**“类 Glorot 初始化”**。

• 类比： 想象在广播开始前调频。旧的方法只是随机转动旋钮。新方法则像是使用精密的科学仪器，无论电台播放的是什么类型的音乐（基函数），都能找到信号最清晰的确切频率。
• 结果： 通过这种精确的“调频”，KAN 保持了稳定。它们可以学习更深、更复杂的谜题，而不会迷失方向。在许多测试中，这个简单的修复让 AI 的答案比以前精确了数千倍。

2. RGA KAN：带有“残差门控”的安全网

问题： 即使有了完美的音量设置，一些非常深的团队（尤其是处理像 Allen-Cahn 方程这样棘手的谜题时）仍然会陷入困境。它们会开始学习，但随后会撞上一堵墙并停止进步。

解决方案： 作者构建了一种名为 RGA KAN（Residual-Gated Adaptive KAN，残差门控自适应 KAN）的新架构。他们从一种名为“PirateNet”的前期设计中汲取了灵感，并添加了一个特殊的机制。

• 类比： 想象一场接力赛。在标准的深度网络中，接力棒在奔跑者之间直线传递。如果其中一个奔跑者掉落了接力棒，整个比赛就结束了。
  RGA KAN 在每一步都增加了一个“智能门控”。这个门控就像一名裁判，它可以决定：“我是把接力棒交给下一个奔跑者，还是让当前的奔跑者再多跑一会儿？”
  • “门控”（Alpha 和 Beta）： 这些是可调节的旋钮。在开始时，门控可能是关闭的，让团队作为一个浅层的、简单的群体运行。随着训练的进行，门控会打开，允许团队变得更深，从而应对更难的问题。如果团队开始变得混乱，门控可以稍微关闭以使其稳定。
• 结果： 这个“安全网”允许 AI 根据需要深入发展，而不会崩溃。它能成功地引导整个学习过程，而旧方法则会在中途停滞。

他们如何证明其有效性

研究人员在 九个不同的物理谜题（如热传导方程、流体流动和波动方程）上测试了他们的新系统。

• 竞争： 他们将新的 RGA KAN 与标准的 cPIKAN（旧的 KAN 方法）以及 PirateNet（当前的顶级 MLP 方法）进行了对比。
• 结果： RGA KAN 几乎在所有测试中都胜出了。
  • 准确度： 它的准确度通常高出几个数量级（这意味着其误差仅为其他方法产生的误差的极小部分）。
  • 稳定性： 当其他方法在处理难题时崩溃（发散）并放弃时，RGA KAN 能够继续运行并找到解。
  • 一致性： 无论使用哪种随机起始点，这种新方法都是可靠的。

训练的“秘诀”

论文还测试了不同的“训练策略”（例如调整 AI 对谜题不同部分的关注程度）。他们发现，虽然新架构是主要功臣，但将其与特定的自适应技术（如 RBA 和 RAD）结合使用，会使其变得更加强大。然而，即使没有这些额外的技巧，新架构也远优于旧有的架构。

总结

简单来说，这篇论文指出：

1. 旧的 KAN 虽然很棒，但在变得过深时非常脆弱。
2. 修复方案 #1： 我们找到了更好的启动方式（初始化），使它们不会立即陷入混乱。
3. 修复方案 #2： 我们构建了一个新的“智能门控”系统（RGA KAN），它让 AI 可以安全地加深，就像一个防止它掉下悬崖的安全网。
4. 结果： 这个新系统在解决复杂物理问题方面，比目前的尖端方法更出色、更可靠，且领先幅度巨大。

作者得出结论，虽然由于进行了更复杂的数学运算，他们的系统计算速度稍慢，但其在准确性和稳定性方面的巨大提升，使其在面对其他方法无法解决的困难问题时，显得非常有价值。

技术摘要

技术摘要：训练深度物理信息驱动的 Kolmogorov–Arnold 网络

问题陈述
Kolmogorov–Arnold 网络 (KANs) 已成为物理信息机器学习 (PIML) 中极具前景的替代方案，能够提供比多层感知器 (MLPs) 更强的可解释性和对谱偏差的鲁棒性。具体而言，基于切比雪夫函数的物理信息 KAN (cPIKANs) 由于其相比于 B-样条变体的计算效率优势，已成为一种标准配置。然而，当扩展到深层架构时，cPIKANs 面临显著挑战。实证研究表明，随着网络深度的增加，cPIKANs 会出现训练不稳定和发散的问题，限制了其在复杂偏微分方程 (PDE) 问题中的应用。此外，现有的 KAN 权重初始化方案在很大程度上仍是启发式的，缺乏类似于 MLP 中 Glorot 初始化那样的理论基础。此外，目前还缺乏一个整合了针对 cPIKANs 自适应策略的统一训练流水线，且其在深层机制下的失效原因尚未得到充分理解。

方法论
作者提出了一种两管齐下的方法来解决 cPIKANS 的深度扩展限制：一种新型初始化方案和一个新的深层架构。

1. 基函数无关的类 Glorot 初始化：
   作者推导了一种基于前向和后向传播过程中方差保持的 KAN 权重初始化方案。与以往针对 B-样条的特定启发式方法不同，该方案是“基函数无关”的，这意味着它不假设特定的基函数族。通过分析输出信号及其对输入的梯度的方差，他们推导出了基系数 ($w_{jim}$) 的标准差，该标准差平衡了输入维度 ($d_I$)、输出维度 ($d_O$) 和基函数数量 ($D$) 的贡献。这种方法旨在防止梯度消失或爆炸，模仿了 Glorot 初始化在 MLPs 中的成功。

2. 残差门控自适应 KAN (RGA KANs)：
   意识到仅靠初始化不足以应对所有深层 PDE 设置（例如 Allen–Cahn 方程），作者引入了受 MLP 中 PirateNet 架构启发的 RGA KAN 架构。关键组件包括：

   • 嵌入 (Embedding)： 通过正弦/余弦嵌入强制执行周期性边界条件。
   • 基于正弦的输入层： 一个基于正弦的 KAN 层处理嵌入后的输入，其作用类似于随机傅里叶特征 (RFF) 嵌入。
   • 自适应跳跃连接 (Adaptive Skip Connections)： 核心创新在于堆叠“RGA 模块”。每个模块包含基于切比雪夫的 KAN 层和可学习的门控参数 ($\alpha$ 和 $\beta$)。这些门控在训练期间动态调节网络的有效深度。具体而言，$\alpha$ 控制整个模块的跳跃连接，而 $\beta$ 控制模块内第一层之后的跳跃连接。这使得网络可以从浅层开始（如果初始化为 $\alpha=0$），并逐步加深，或者从深层开始并进行自适应剪枝，从而稳定优化过程。
   • 物理信息输出： 最后一层可以通过最小二乘拟合初始化以近似 PDE 的初始条件。

3. 信息瓶颈 (IB) 分析：
   为了理解训练动力学，作者应用了信息瓶颈理论。他们监测了梯度的信噪比 (SNR) 和网络的几何复杂度。他们假设成功的训练需要遍历三个阶段：拟合阶段、扩散阶段和扩散平衡阶段。

4. 统一训练流水线：
   实验使用了一个标准的流水线，其中包含了 PINNs 中常用的自适应技术：基于残差的注意力 (RBA)、基于残差的自适应分布 (RAD)、因果训练以及学习率退火 (LRA)。

核心贡献

• 推导了类 Glorot 初始化： 推导出了一个具有理论依据的基函数无关初始化规则，显著提高了 cPIKANs 在函数拟合和 PDE 任务中的稳定性与精度。
• 引入了 RGA KANs： 一种旨在通过自适应跳跃连接和门控机制来缓解深层 cPIKANs 发散问题的创新深层架构。
• 通过 IB 理论提供的理论见解： 一项分析表明，RGA KANs 能够成功遍历所有三个训练阶段（拟合、扩散、扩散平衡），而基准 cPIKANs 则往往停滞在扩散阶段，无法实现泛化。
• 全面的基准测试： 在九个标准前向 PDE 基准（包括 Burgers'、Allen–Cahn、Korteweg–De Vries、Sine Gordon、Advection、Helmholtz、Poisson、Heat 和 Navier-Stokes 方程）上进行了广泛评估，将 RGA KANs 与参数匹配的 cPIKANs 以及 PirateNets 进行了对比。

结果

• 初始化影响： 所提出的类 Glorot 初始化在函数拟合和 PDE 任务中始终优于默认的 cPIKAN 初始化，通常能将相对 $L_2$ 误差降低数个数量级。在深层网络（如 Burgers' 方程）中，默认初始化会导致发散，而所提出的方案能保持稳定性。
• 架构性能： RGA KANs 展示了比基准 cPIKANs 和 PirateNets 更优越的稳定性和准确性。在 cPIKANs 和 PirateNets 发散的基准测试中（如 Allen–Cahn、Advection、Korteweg–De Vries、Sine Gordon），RGA KANs 收敛到了精确解。
• 误差降低： 在九个 PDE 基准测试中，RGA KANs 始终优于参数匹配的基准模型，通常领先数个数量级。例如，在 Helmholtz 方程中，RGA KANs 达到了 $O(10^{-5})$ 级别的误差，优于 cPIKANs ($O(10^{-3})$) 和 PirateNets ($O(10^{-4})$)。
• 消融研究： 自适应组件（RBA、RAD、因果训练、LRA）的贡献因 PDE 而异。虽然 RGA KANs 具有鲁棒性，但移除特定组件（如 Sine Gordon 的 LRA 或 Advection 的 RAD）可能会导致发散或误差显著增加，凸显了这些策略具有问题依赖性的特征。
• 计算成本： 由于存在门控操作和基函数求值，RGA KANs 通常比 cPIKANs 每轮迭代的计算成本更高。然而，在像 Navier-Stokes 这样复杂的问题中，随着门控机制成为 RGA KANs 和 PirateNets 的主要瓶颈，两者之间的成本差距会缩小。

意义与主张
本文声称，所提出的初始化和 RGA KAN 架构共同解决了深层物理信息 KAN 的关键差距。作者断言，他们的工作提供了第一套针对 cPIKANs 的可扩展深度基准测试，并证明了深层 KAN 可以稳定训练而不发生发散，而这在以往观察到的深层 PINNs 和 cPIKANs 中是一个限制。通过成功导航信息瓶颈阶段，RGA KANs 实现了基准架构所缺乏的泛化能力。作者认为，他们的工作并非针对每个特定 PDE 进行超参数微调的最优解，而是一个稳健、统一的框架，在固定的公平训练流水线下优于现有的最先进架构（PirateNets）和基准 KAN。他们建议，该方法为未来在算子学习和其他 KAN 变体中的应用提供了坚实的基础。