Multilevel Training for Kolmogorov Arnold Networks

本文通过建立 KAN 与多通道 MLP 之间的等价关系，提出了一种基于样条基函数均匀细化的多尺度训练框架，利用解析几何插值算子构建嵌套架构，从而在物理信息神经网络等任务中实现了比传统方法高出数个数量级的训练精度与效率。

要点

这篇论文讲述了一个关于如何让神经网络“学”得更快、更聪明的故事。

想象一下，你正在教一个学生（神经网络）去画一幅复杂的画，或者去解一道很难的数学题。

1. 传统的困境：在迷雾中摸索

传统的神经网络（叫 MLP）就像是一个在迷雾中摸索的盲人。它有很多层，每一层都在做简单的加法和非线性变换。虽然它能学会很多东西，但训练过程非常慢，就像在迷宫里乱撞，经常走弯路，收敛（学会）得很慢。

2. 新主角登场：KAN 网络

最近出现了一种叫 KAN (Kolmogorov-Arnold Networks) 的新架构。它不像传统网络那样把“激活函数”（比如 ReLU）藏在黑盒子里，而是把激活函数变成了可学习的曲线（就像用橡皮泥捏出各种形状）。

• 优点：KAN 更透明，更容易理解，而且特别适合处理那些有“尖角”或不平滑的复杂函数（比如物理定律中的突变）。
• 缺点：虽然它结构好，但如果训练方法不对，它依然可能像传统网络一样慢，甚至更慢。

3. 核心发现：换个“视角”看问题

作者发现了一个惊人的秘密：KAN 网络其实可以看作是另一种形式的传统网络（多通道 MLP），只是它们用的“语言”（基函数）不同。

• 比喻：这就好比两个人在描述同一个物体。
  • 一个人用**“平滑的波浪线”**（样条函数/Spline）来描述。
  • 另一个人用**“折线”**（ReLU 函数）来描述。
  • 作者发现，只要做一个简单的数学翻译（基变换），这两种描述就是完全等价的。

但是！重点来了：
虽然它们描述的物体一样，但**“学习”的过程（梯度下降）却完全不同**。

• 如果你用“折线”语言（ReLU）去学，优化器会极度偏爱平滑的、简单的形状，而忽略那些复杂的、高频的细节（就像只愿意画圆圈，不愿意画锯齿）。
• 如果你用“波浪线”语言（样条/KAN 原生语言）去学，优化器就能同时关注平滑的大轮廓和尖锐的小细节。

4. 终极武器：多级训练法（Multilevel Training）

既然 KAN 有这种特殊的结构，作者就借鉴了物理学中解决难题的**“多级网格法”**（Multigrid），发明了一种新的训练策略。

这个策略就像“先画草图，再画细节”：

1. 粗网阶段（画草图）：
   先用很少的“节点”（比如只有 4 个控制点）训练一个很简单的 KAN 模型。这时候，模型只能画出大概的轮廓（比如一个圆）。因为它很简单，所以学得飞快。
2. 细网阶段（加细节）：
   把刚才画好的草图，通过一种几何插值的方法，“复制”到一个更精细的模型上（比如把 4 个点变成 8 个，再变成 16 个）。
   • 关键点：这种复制不是乱复制，而是保证**“粗模型的进步不会白费”**。新模型继承了旧模型画好的轮廓。
3. 互补优化（各司其职）：
   • 在粗网上，模型负责解决低频、平滑的大问题（大轮廓）。
   • 在细网上，模型负责解决高频、复杂的小问题（细节、尖角）。
   • 因为 KAN 的数学结构（样条函数的局部性），细网模型能迅速捕捉到粗网模型忽略的细节，而不会去重复做粗网已经做好的事。

5. 结果：快得惊人

作者做了很多实验（比如模拟物理方程、函数回归）：

• 传统方法：要么慢，要么精度差。如果用“折线”语言做多级训练，效果几乎为零，因为细网模型还在重复粗网已经学会的东西，无法利用新的细节能力。
• KAN 多级训练：精度提高了100 倍甚至 1000 倍（几个数量级），而且训练速度极快。

总结

这篇论文的核心思想是：
不要试图用一把锤子（传统训练方法）去敲所有的钉子。
通过理解 KAN 网络独特的数学结构（它像样条曲线一样有局部性），我们设计了一套**“先粗后细、层层递进”**的训练方法。这让神经网络能像人类画家一样，先定大框架，再填细节，从而在解决复杂科学问题（如物理模拟）时，展现出惊人的速度和精度。

一句话概括：
给神经网络换了一套更聪明的“画笔”（样条基），并教它学会了“先画大轮廓，再抠小细节”的多级绘画技巧，从而让它在解决难题时快如闪电。

技术摘要

1. 研究背景与问题 (Problem)

• 传统神经网络的训练瓶颈：多层感知机（MLP）及其变体（如 Transformer）在训练时缺乏结构保证，导致基于梯度的优化算法（如 SGD）收敛缓慢。虽然数值计算领域（如偏微分方程 PDE 求解）中的多级方法（Multigrid Methods）和多重网格技术能高效求解线性/非线性系统（$O(n)$ 复杂度），但将其成功应用于机器学习训练仍面临巨大挑战。
• 现有尝试的局限性：以往将多重网格思想引入深度学习的工作，要么仅通过并行化获得加速（非算法性加速），要么侧重于修改网络架构而非训练算法，未能实现类似数值计算领域的算法级加速。
• 核心难点：机器学习模型缺乏像 PDE 网格那样自然的“粗 - 细”层级结构。在粗网格和细网格上操作通常处于相同的维度空间，难以定义具有良好近似性质的粗化算子（Coarsening）和插值算子（Interpolation），且缺乏互补的优化策略（即粗网格解决低频误差，细网格解决高频误差）。
• KAN 的潜力与缺口：Kolmogorov-Arnold 网络（KANs）通过可学习的激活函数（通常基于样条基函数）提供了比 MLP 更丰富的结构，具有更好的可解释性和对低正则性函数的捕捉能力。然而，目前缺乏针对 KAN 的高效训练策略，特别是如何利用其内在结构实现多级训练。

2. 核心方法论 (Methodology)

本文提出了一种基于样条基函数的 KAN 多级训练框架，主要包含以下三个理论支柱：

2.1 KAN 与多通道 MLP 的基变换等价性

• 理论发现：作者证明了使用样条基函数（B-splines）的 KAN 层，可以通过一个线性基变换（Change-of-Basis），等价于具有**幂次 ReLU 激活函数（Power ReLU, $ReLU^{r-1}$）**的多通道 MLP。
• 变换矩阵 $A^{[r]}$：该变换矩阵 $A^{[r]}$ 具有特殊的代数结构。对于均匀节点间距，它是一个上三角 Toeplitz 矩阵。
• 微分算子联系：该矩阵 $A^{[r]}$ 本质上对应于样条节点上 $r$ 阶导数算子的前向有限差分近似。其转置乘积 $(A^{[r]})^T A^{[r]}$ 则对应于 $2r$ 阶导数算子的有限差分近似。

2.2 基变换对优化几何的影响

• 预条件效应：虽然 KAN（样条基）和 MLP（ReLU 基）作为前向算子是等价的，但它们的梯度下降动力学截然不同。基变换矩阵 $A$ 充当了优化过程中的预条件器（Preconditioner）。
• 频谱分析：
  • ReLU/MLP 基：由于 $(A^T A)$ 的特征值谱跨度极大（从平滑模式到振荡模式），梯度下降会强烈偏向于优化平滑函数（低频模式），而难以捕捉高频/振荡模式。
  • 样条/KAN 基：在自然样条基下，权重对应于具有紧支集（Compact Support）的局部基函数。这使得梯度更新具有局部性，能够高效地优化高频和振荡模式。
• 互补性：这种差异意味着，在粗网格（稀疏节点）上优化 MLP 倾向于平滑解，而在细网格（密集节点）上优化 KAN 则能利用新增的表达能力去修正高频误差。

2.3 多级训练框架与“正确嵌套层级”

• 正确嵌套层级（Properly Nested Hierarchy）：作者定义了一个关键概念，即细网格模型在插值粗网格权重后，必须精确保持粗网格模型的功能输出（$g_f(x; P u^{(c)}) = g_c(x; u^{(c)})$）。
• 几何插值算子：利用样条函数的嵌套性质（$S_r(T) \subset S_r(T')$），作者构建了基于几何网格细化的插值算子 $P$。这使得从粗网格到细网格的权重转移是解析的、快速的，且不需要重新求解最小二乘问题。
• 训练流程：
  1. 在粗网格（少量样条节点）上训练 KAN。
  2. 通过几何插值将权重转移到细网格（节点加密）。
  3. 在细网格上继续训练。
  4. 由于样条基的局部性和互补优化特性，细网格训练能迅速利用新增的表达能力来降低损失，而不会破坏粗网格已学到的低频特征。

3. 主要贡献 (Key Contributions)

1. 理论等价性证明：建立了样条 KAN 与多通道 MLP 之间的严格线性等价关系，揭示了 KAN 内部结构隐含的微分算子性质。
2. 优化动力学分析：从预条件和谱分析角度，解释了为何 KAN 在捕捉低正则性（高振荡）函数方面优于 MLP，并指出 MLP 在多级训练中缺乏互补性（Complementary Relaxation）。
3. 多级训练算法：提出了针对 KAN 的“正确嵌套层级”概念，设计了基于均匀节点细化的快速插值算子，实现了无需昂贵计算的多级训练。
4. 实证验证：在函数回归和物理信息神经网络（PINNs）任务中，证明了该方法的有效性。

4. 实验结果 (Results)

实验对比了不同基（样条 vs ReLU）、不同架构（KAN vs MLP）及不同训练策略（单级 vs 多级）：

• 函数回归任务：
  • 精度提升：在样条基上进行多级训练的 KAN，其均方误差（MSE）比单级训练的 KAN 或同等规模的 MLP 高出1-3 个数量级。
  • 基的对比：在 ReLU 基上进行多级训练，精度几乎没有提升（甚至不如粗网格模型），因为细网格优化被平滑偏好主导，无法利用细网格的高频表达能力。
• 物理信息神经网络 (PINNs) 任务：
  • 2D Poisson 方程：样条 KAN 的多级训练收敛更快，且误差更稳定。ReLU KAN 和多级 MLP 在细化网格后陷入停滞。
  • 1D Burger's 方程（低正则性）：多级样条 KAN 比单级 KAN 或 MLP 的精度提高了2-3 个数量级。
  • Allen-Cahn 方程：多级样条 KAN 成功捕捉了不稳定的固定点和复杂的相变结构，而 ReLU 基和 MLP 未能捕捉到有意义的结构。
• 频谱分析：傅里叶变换分析显示，随着网格细化，样条 KAN 的残差频谱能量向高频扩展，表明模型正在有效学习高频模式；而 ReLU 基的频谱依然狭窄。

5. 意义与结论 (Significance)

• 填补了多级机器学习的空白：本文首次展示了在机器学习领域实现类似数值 PDE 求解中的算法级加速（而非仅仅是并行加速）。关键在于利用 KAN 的样条结构构建了“正确嵌套”的层级和互补的优化动力学。
• 理论指导实践：证明了神经网络架构的设计（如选择样条基而非 ReLU）可以直接决定其是否适合多级优化策略。
• PINNs 的突破：对于求解偏微分方程（特别是涉及低正则性解或复杂物理现象的问题），该方法提供了一种无需复杂技巧（如自适应采样、特殊损失加权）即可实现快速、鲁棒训练的“开箱即用”方案。
• 未来方向：为设计具有可解释性、结构化和高效训练能力的下一代神经网络提供了新的范式。

总结：该论文通过深入分析 KAN 的数学结构，发现其天然具备多级优化的潜力。通过引入基变换理论和几何插值算子，作者成功将数值计算中的多级网格思想移植到 KAN 训练中，实现了在精度和效率上的显著突破，特别是在处理物理驱动的科学计算问题时表现卓越。