Skip to content

필사 모드: KAN(Kolmogorov-Arnold Networks)论文解析:替代 MLP 的可学习激活函数架构

中文
0%
정확도 0%
💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.
KAN Architecture

引言

深度学习的基本构成单元 Multi-Layer Perceptron(MLP)自 1958 年 Frank Rosenblatt 提出感知机以来,60 多年来一直是神经网络架构的核心。MLP 的结构很清晰:每个神经元(节点)应用固定的激活函数(ReLU、SiLU 等),边(权重)则执行可学习的线性变换。这一简单却强大的结构,凭借 Universal Approximation Theorem 获得了理论上的正当性,又因为矩阵乘法针对 GPU 运算做了优化,在实践中也一直占据主导地位。

然而在 2024 年 4 月,MIT 的 Ziming Liu、Max Tegmark 等人发表的 "KAN: Kolmogorov-Arnold Networks" 论文,向这个持续了 60 年的范式提出了根本性的质疑:激活函数一定要固定在节点上吗?边(权重)一定要是单纯的标量乘法吗?KAN 把这两个假设都推翻了——把激活函数放在边而不是节点上,并将每个激活函数用 B 样条参数化,使其可学习。这一想法的数学依据,来自 1957 年 Andrey Kolmogorov 及其学生 Vladimir Arnold 证明的 Kolmogorov-Arnold 表示定理

KAN 已被 ICLR 2025 接收,在函数拟合、求解偏微分方程(PDE)、符号回归(symbolic regression)等科学探索任务上,用 比 MLP 更少的参数达到了更高的精度。此外,由于学习到的激活函数可以被可视化并进行符号化解释,KAN 在 可解释性(interpretability) 方面也大幅领先于 MLP。

本文将全面剖析 KAN 论文——从数学基础、架构设计、B 样条激活函数的工作原理,到与 MLP 的系统性比较、PyTorch 实现代码、可视化方法、实际应用案例,以及当前存在的局限性。

Kolmogorov-Arnold 表示定理

定理的数学定义

KAN 的理论基础是 1957 年被证明的 Kolmogorov-Arnold 表示定理(Kolmogorov-Arnold Representation Theorem)。这一定理是对希尔伯特第 13 问题(Hilbert's 13th Problem)的解答,证明了任意多元连续函数都可以表示为单变量连续函数的复合。

用数学语言描述该定理如下。对有界域 [0,1]n[0,1]^n 上的任意连续函数 ff,下式成立:

f(x1,x2,,xn)=q=02nΦq(p=1nϕq,p(xp))f(x_1, x_2, \ldots, x_n) = \sum_{q=0}^{2n} \Phi_q \left( \sum_{p=1}^{n} \phi_{q,p}(x_p) \right)

这里 ϕq,p:[0,1]R\phi_{q,p}: [0,1] \to \mathbb{R}内部函数(inner function),Φq:RR\Phi_q: \mathbb{R} \to \mathbb{R}外部函数(outer function)。关键在于,内部函数和外部函数都是 单变量(univariate)连续函数。也就是说,无论多元函数多么复杂,都可以仅用单变量函数的复合与加法来精确表示。

与 MLP 的 Universal Approximation Theorem 的区别

MLP 的理论基础 Universal Approximation Theorem 保证的是 近似(approximation),而 Kolmogorov-Arnold 表示定理保证的是 精确表示(exact representation)。两个定理的比较如下:

项目Universal Approximation(MLP)Kolmogorov-Arnold(KAN)
保证程度任意精度的近似精确表示
结构固定激活函数 + 可学习权重可学习的单变量函数复合
神经元数量可以无限多2n+12n+1 个内部求和即可
激活函数位置节点(神经元)边(权重)
可解释性黑箱可对单变量函数进行可视化

从神经网络视角的再诠释

从神经网络的视角重新诠释 Kolmogorov-Arnold 表示定理,可以把它看作一个拥有 2 个隐藏层(hidden layer)的网络。第一个隐藏层对每个输入维度应用 n×(2n+1)n \times (2n+1) 个单变量函数 ϕq,p\phi_{q,p},第二个隐藏层则对内部求和应用 2n+12n+1 个外部函数 Φq\Phi_q

然而,原始定理中的 Φq\Phi_qϕq,p\phi_{q,p} 可能并不光滑,甚至具有分形结构,因此难以直接用于实际训练。KAN 论文的核心贡献,是把这一定理 推广为任意深度的网络,并通过将每个单变量函数 用 B 样条参数化,使其真正具备了实用性。

KAN 架构详细分析

MLP 与 KAN 的结构性差异

MLP 与 KAN 的根本差异在于,激活函数与权重的角色发生了对调。

MLP:在边上执行线性变换(wxw \cdot x),在节点上应用非线性激活 σ(z)\sigma(z)

MLP Layer: y=σ(Wx+b)\text{MLP Layer: } \mathbf{y} = \sigma(\mathbf{W}\mathbf{x} + \mathbf{b})

KAN:在边上应用可学习的非线性函数 ϕ(x)\phi(x),节点上只做简单的求和

KAN Layer: yj=i=1ninϕj,i(xi)\text{KAN Layer: } y_j = \sum_{i=1}^{n_{\text{in}}} \phi_{j,i}(x_i)

这里 ϕj,i\phi_{j,i} 是放置在「从第 ii 个输入通往第 jj 个输出」这条边上的可学习单变量函数。

KAN 层架构图

下面这张图展示了 KAN 层的结构。关键在于,每条边都不是单纯的标量权重,而是一个可学习的 B 样条函数。

KAN Layer: [2, 3] (2 inputs -> 3 outputs)

Input Layer          Learnable Activations (Edges)          Output Layer
                         (B-spline functions)

              +------ phi_1,1(x) --------+
              |                           |
  x_1 -------+------ phi_2,1(x) --------+-----> y_1 = phi_1,1(x_1) + phi_1,2(x_2)
              |                           |
              +------ phi_3,1(x) --------+
              |                           |
              |                           +-----> y_2 = phi_2,1(x_1) + phi_2,2(x_2)
              |                           |
  x_2 -------+------ phi_1,2(x) --------+
              |                           |
              +------ phi_2,2(x) --------+-----> y_3 = phi_3,1(x_1) + phi_3,2(x_2)
              |                           |
              +------ phi_3,2(x) --------+

  [Nodes: SUM only]   [Edges: Learnable]    [Nodes: SUM only]

  Total learnable functions: n_in x n_out = 2 x 3 = 6
  Each phi_{j,i} is a B-spline with (G + k) parameters
  where G = grid intervals, k = spline order

多层 KAN 网络

Kolmogorov-Arnold 表示定理的原始结构被限制在 2 层,而 KAN 论文将其扩展到 任意深度。在深度为 LL 的 KAN 网络中,每一层 ll 拥有 nln_l 个输入和 nl+1n_{l+1} 个输出,整个网络可以表示为:

KAN(x)=(ΦL1ΦL2Φ1Φ0)(x)\text{KAN}(\mathbf{x}) = (\Phi_{L-1} \circ \Phi_{L-2} \circ \cdots \circ \Phi_1 \circ \Phi_0)(\mathbf{x})

这里每个 Φl\Phi_l 都是一个 KAN 层,层 ll 中共有 nl×nl+1n_l \times n_{l+1} 个可学习的单变量函数。整个网络的可学习函数总数为:

Total functions=l=0L1nl×nl+1\text{Total functions} = \sum_{l=0}^{L-1} n_l \times n_{l+1}

KAN 的参数量计算

如果每个单变量函数用 GG 个网格区间和 kk 阶 B 样条来参数化,那么每个函数需要 G+kG + k 个系数。因此,整个网络的参数量为:

Total params=l=0L1nl×nl+1×(G+k)\text{Total params} = \sum_{l=0}^{L-1} n_l \times n_{l+1} \times (G + k)

举例来说,对于宽度(width)为 [2,5,1][2, 5, 1]G=5G=5k=3k=3 的 KAN:

  • 第 0 层:2×5×(5+3)=802 \times 5 \times (5+3) = 80 个参数
  • 第 1 层:5×1×(5+3)=405 \times 1 \times (5+3) = 40 个参数
  • 总参数量:120120

相同结构的 MLP [2,5,1][2, 5, 1](2×5+5)+(5×1+1)=21(2 \times 5 + 5) + (5 \times 1 + 1) = 21 个参数。虽然 KAN 的参数更多,但论文主张,就 参数与精度之比 而言,KAN 全面压制了 MLP。

B 样条激活函数

B 样条的定义

B 样条(Basis Spline)是由分段多项式(piecewise polynomial)线性组合而定义出的光滑函数。kk 阶 B 样条基函数由 De Boor 递归公式定义:

Bi,0(x)={1if tix<ti+10otherwiseB_{i,0}(x) = \begin{cases} 1 & \text{if } t_i \leq x < t_{i+1} \\ 0 & \text{otherwise} \end{cases} Bi,k(x)=xtiti+ktiBi,k1(x)+ti+k+1xti+k+1ti+1Bi+1,k1(x)B_{i,k}(x) = \frac{x - t_i}{t_{i+k} - t_i} B_{i,k-1}(x) + \frac{t_{i+k+1} - x}{t_{i+k+1} - t_{i+1}} B_{i+1,k-1}(x)

这里 tit_i节点(knot) 向量的元素。使用 GG 个网格区间和 kk 阶样条,会生成 G+kG + k 个 B 样条基函数,样条函数则表示为这些基函数的线性组合:

spline(x)=i=0G+k1ciBi,k(x)\text{spline}(x) = \sum_{i=0}^{G+k-1} c_i B_{i,k}(x)

这里 cic_i 是可学习的系数。

KAN 中激活函数的结构

放置在 KAN 每条边上的激活函数 ϕ(x)\phi(x) 由两部分组成:

ϕ(x)=wbb(x)+wsspline(x)\phi(x) = w_b \cdot b(x) + w_s \cdot \text{spline}(x)

这里 b(x)b(x)残差函数(residual function),默认取 SiLU(x)=xσ(x)\text{SiLU}(x) = x \cdot \sigma(x)wbw_bwsw_s 分别是对应的缩放系数。残差函数在训练初期提供稳定性,帮助样条部分逐步学习更复杂的模式。

B 样条基函数的实现

下面是用 PyTorch 实现 B 样条基函数的代码:

import torch
import torch.nn as nn
import numpy as np


def compute_bspline_basis(x: torch.Tensor, grid: torch.Tensor, k: int) -> torch.Tensor:
    """用 De Boor 递归公式计算 B 样条基函数。

    Args:
        x: 输入张量 [batch_size, in_features]
        grid: 节点(knot)向量 [in_features, grid_size + 2k + 1]
        k: 样条阶数(通常为 3,即三次样条)

    Returns:
        基函数值 [batch_size, in_features, grid_size + k]
    """
    # 扩展 x 的维度,以便与 grid 比较
    # x: [batch, in_features] -> [batch, in_features, 1]
    x = x.unsqueeze(-1)

    # 0 阶基函数:落在区间内为 1,否则为 0
    bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).float()

    # De Boor 递归:重复 k 次以计算 k 阶基函数
    for degree in range(1, k + 1):
        # 左侧项: (x - t_i) / (t_{i+degree} - t_i)
        left_num = x - grid[:, :-(degree + 1)]
        left_den = grid[:, degree:-1] - grid[:, :-(degree + 1)]
        left = left_num / left_den.clamp(min=1e-8)

        # 右侧项: (t_{i+degree+1} - x) / (t_{i+degree+1} - t_{i+1})
        right_num = grid[:, (degree + 1):] - x
        right_den = grid[:, (degree + 1):] - grid[:, 1:(-degree)]
        right = right_num / right_den.clamp(min=1e-8)

        # 递归组合
        bases = left * bases[:, :, :-1] + right * bases[:, :, 1:]

    return bases  # [batch, in_features, grid_size + k]


# 使用示例
batch_size, in_features = 32, 4
grid_size, k = 5, 3  # 5 个网格区间,三次样条

# 生成均匀网格:在 [-1, 1] 范围内两侧各扩展 k 个
grid_points = torch.linspace(-1, 1, grid_size + 1)
extended_grid = torch.cat([
    grid_points[0:1] - torch.arange(k, 0, -1) * (grid_points[1] - grid_points[0]),
    grid_points,
    grid_points[-1:] + torch.arange(1, k + 1) * (grid_points[1] - grid_points[0])
])
# 复制到 in_features 维度
grid = extended_grid.unsqueeze(0).repeat(in_features, 1)

x = torch.randn(batch_size, in_features)
basis_values = compute_bspline_basis(x, grid, k)
print(f"基函数 shape: {basis_values.shape}")
# 输出: 基函数 shape: torch.Size([32, 4, 8])  # G + k = 5 + 3 = 8

网格更新机制

KAN 的一个重要特点是,训练过程中可以 动态更新网格。初期使用均匀网格,但如果按训练数据的分布重新排布网格,就能大幅提升近似精度。grid_eps 超参数负责控制这一点:

  • grid_eps = 1.0:完全均匀的网格
  • grid_eps = 0.0:按数据样本的分位数(quantile)排布的网格
  • 0 < grid_eps < 1:两个极端之间的插值

逐步增加网格区间数 GG(网格扩展,grid extension)可以提高分辨率,从而表示更复杂的函数。这与在 MLP 中增加神经元数量类似,但优点在于,能够在保留已有训练结果的同时,只提升分辨率。

KAN 与 MLP 的比较分析

系统性比较表

比较项MLPKAN
激活函数位置节点(固定:ReLU、SiLU 等)边(可学习:B 样条)
权重的角色可学习的线性变换 wxw \cdot x无(由激活函数取代)
参数效率低(需要加宽才能提升精度)高(靠网格分辨率提升精度)
缩放规律Nα\ell \propto N^{-\alpha}α\alpha 较小Nα\ell \propto N^{-\alpha}α\alpha 更大
可解释性黑箱可对激活函数进行可视化
符号回归不可训练后可提取符号函数
训练速度快(GPU 矩阵乘法已优化)慢(样条计算带来额外开销)
GPU 并行化出色(批量矩阵乘法)受限(每条边的函数各不相同)
高维输入出色(图像、NLP 等)受限(受维度灾难影响)
主要应用领域通用深度学习科学发现、符号回归

精度缩放比较

KAN 论文中最令人印象深刻的结果之一,是在 神经缩放规律(Neural Scaling Law)上表现出的差异。将测试 RMSE 表示为参数量 NN 的函数时:

RMSENα\text{RMSE} \propto N^{-\alpha}

对 MLP 而言,α\alpha 的上限由 Universal Approximation Theorem 决定;而 KAN 基于 Kolmogorov-Arnold 表示定理,展现出更快的缩放速度(更大的 α\alpha)。在论文针对 5 个特殊函数的实验中,KAN [2,5,1] 用 80 个参数,比 MLP [2,100,1] 的 201 个参数少 2-3 倍,却达到了相当甚至更好的精度。

训练速度与现实考量

KAN 的训练速度比 MLP 慢得多,主要原因如下:

  1. 样条基函数计算:每条边都需要递归计算 G+kG+k 个基函数。
  2. GPU 并行化的局限:MLP 可以通过权重矩阵的一次矩阵乘法完成计算,而 KAN 在每条边上应用不同的函数,因此难以有效利用批量矩阵乘法。
  3. 网格更新开销:训练过程中的网格重新排布会带来额外的计算成本。

根据 "KAN or MLP: A Fairer Comparison"(Yu et al., 2024)论文,在大多数标准机器学习、计算机视觉、NLP、音频处理任务上,MLP 都表现出了更高的平均精度。KAN 展现出优势的,仅限于符号公式表示(symbolic formula representation)等特定科学任务。

实现代码

KAN 层的 PyTorch 实现

下面是用 PyTorch 从零(from scratch)实现 KAN 层的代码:

import torch
import torch.nn as nn
import torch.nn.functional as F


class KANLayer(nn.Module):
    """Kolmogorov-Arnold Network 层。

    在每条边上放置可学习的 B 样条激活函数。
    节点上只执行简单的求和。
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        grid_size: int = 5,
        spline_order: int = 3,
        scale_base: float = 1.0,
        scale_spline: float = 1.0,
        grid_range: tuple = (-1.0, 1.0),
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.grid_size = grid_size
        self.spline_order = spline_order
        self.scale_base = scale_base
        self.scale_spline = scale_spline

        # 生成并扩展均匀网格(两侧各扩展 spline_order 个)
        h = (grid_range[1] - grid_range[0]) / grid_size
        grid = torch.arange(-spline_order, grid_size + spline_order + 1) * h + grid_range[0]
        # grid shape: [grid_size + 2 * spline_order + 1]
        # 扩展维度,使每个 (in, out) 组合都可以拥有独立的网格
        self.register_buffer(
            "grid",
            grid.unsqueeze(0).unsqueeze(0).expand(out_features, in_features, -1)
        )

        # 样条系数:每条边 (grid_size + spline_order) 个
        self.spline_weight = nn.Parameter(
            torch.randn(out_features, in_features, grid_size + spline_order)
            * 0.1
        )

        # 残差(base)权重
        self.base_weight = nn.Parameter(
            torch.randn(out_features, in_features) * 0.1
        )

    def b_splines(self, x: torch.Tensor) -> torch.Tensor:
        """计算 B 样条基函数。

        Args:
            x: [batch_size, in_features]

        Returns:
            基函数值 [batch_size, in_features, grid_size + spline_order]
        """
        # 将 x 扩展为 [batch, 1, in_features, 1],以便与 grid 比较
        x = x.unsqueeze(1).unsqueeze(-1)  # [B, 1, in, 1]

        # grid: [out, in, G+2k+1] -> 便于广播
        grid = self.grid  # [out, in, G+2k+1]

        # 0 阶基函数:落在区间内为 1
        # 使用 grid[:, :, :-1] 与 grid[:, :, 1:]
        # 只使用第一个 out 维度的网格(因为都相同)
        g = grid[0]  # [in, G+2k+1]
        bases = ((x[:, 0] >= g[:, :-1]) & (x[:, 0] < g[:, 1:])).float()
        # bases: [B, in, G+2k]

        for degree in range(1, self.spline_order + 1):
            left_num = x[:, 0] - g[:, :-(degree + 1)]
            left_den = g[:, degree:-1] - g[:, :-(degree + 1)]
            left = left_num / left_den.clamp(min=1e-8)

            right_num = g[:, (degree + 1):] - x[:, 0]
            right_den = g[:, (degree + 1):] - g[:, 1:(-degree)]
            right = right_num / right_den.clamp(min=1e-8)

            bases = left * bases[:, :, :-1] + right * bases[:, :, 1:]

        return bases  # [B, in, G+k]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播:残差函数 + 样条函数。

        Args:
            x: [batch_size, in_features]

        Returns:
            输出 [batch_size, out_features]
        """
        batch_size = x.size(0)

        # 1) 残差部分:SiLU 激活后做线性组合
        base_output = F.silu(x)  # [B, in]
        base_output = torch.einsum("bi,oi->bo", base_output, self.base_weight)
        # base_output: [B, out]

        # 2) 样条部分:B 样条基函数 * 系数
        spline_basis = self.b_splines(x)  # [B, in, G+k]
        spline_output = torch.einsum(
            "big,oig->bo", spline_basis, self.spline_weight
        )
        # spline_output: [B, out]

        # 3) 合并
        return self.scale_base * base_output + self.scale_spline * spline_output

完整的 KAN 网络实现

class KAN(nn.Module):
    """多层 Kolmogorov-Arnold Network。

    Args:
        width: 各层神经元数量的列表。例如:[2, 5, 1]
        grid_size: B 样条网格区间数
        spline_order: B 样条阶数(默认值:3,即三次样条)
    """

    def __init__(
        self,
        width: list,
        grid_size: int = 5,
        spline_order: int = 3,
    ):
        super().__init__()
        self.width = width
        self.layers = nn.ModuleList()

        for i in range(len(width) - 1):
            self.layers.append(
                KANLayer(
                    in_features=width[i],
                    out_features=width[i + 1],
                    grid_size=grid_size,
                    spline_order=spline_order,
                )
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x)
        return x

    def parameter_count(self) -> dict:
        """返回各层的参数量。"""
        info = {}
        total = 0
        for i, layer in enumerate(self.layers):
            layer_params = sum(p.numel() for p in layer.parameters())
            info[f"layer_{i}"] = layer_params
            total += layer_params
        info["total"] = total
        return info


# 使用示例
model = KAN(width=[2, 5, 1], grid_size=5, spline_order=3)
print(f"模型结构: {model.width}")
print(f"参数量: {model.parameter_count()}")

# 前向传播测试
x = torch.randn(16, 2)
y = model(x)
print(f"输入: {x.shape}, 输出: {y.shape}")
# 输出示例:
# 模型结构: [2, 5, 1]
# 参数量: {'layer_0': 90, 'layer_1': 45, 'total': 135}
# 输入: torch.Size([16, 2]), 输出: torch.Size([16, 1])

用于对比的 MLP 实现

实现一个相同结构的 MLP,以进行公平的比较:

class MLP(nn.Module):
    """用于对比的 Multi-Layer Perceptron。

    使用与 KAN 相同的 width 结构,
    但采用固定激活函数(SiLU)和线性权重。
    """

    def __init__(self, width: list):
        super().__init__()
        layers = []
        for i in range(len(width) - 1):
            layers.append(nn.Linear(width[i], width[i + 1]))
            if i < len(width) - 2:  # 排除最后一层
                layers.append(nn.SiLU())
        self.network = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


# MLP vs KAN 参数量比较
mlp = MLP(width=[2, 5, 1])
kan = KAN(width=[2, 5, 1], grid_size=5, spline_order=3)

mlp_params = sum(p.numel() for p in mlp.parameters())
kan_params = sum(p.numel() for p in kan.parameters())

print(f"MLP 参数量: {mlp_params}")
print(f"KAN 参数量: {kan_params}")
print(f"KAN/MLP 参数量比率: {kan_params / mlp_params:.1f}x")
# 输出示例:
# MLP 参数量: 21
# KAN 参数量: 135
# KAN/MLP 参数量比率: 6.4x

训练与可视化

训练循环的实现

KAN 论文推荐使用 LBFGS 优化器,但在大规模数据集上也可以使用 Adam。下面是用于函数近似任务的完整训练循环:

import torch
import torch.nn as nn
import matplotlib.pyplot as plt


def create_dataset(f, n_train=1000, n_test=200, input_dim=2):
    """生成训练/测试数据集。

    Args:
        f: 待近似的目标函数
        n_train: 训练样本数
        n_test: 测试样本数
        input_dim: 输入维度
    """
    x_train = torch.rand(n_train, input_dim) * 2 - 1  # [-1, 1]
    y_train = f(x_train)
    x_test = torch.rand(n_test, input_dim) * 2 - 1
    y_test = f(x_test)
    return x_train, y_train, x_test, y_test


def train_model(model, x_train, y_train, x_test, y_test,
                epochs=500, lr=1e-2, optimizer_type="adam"):
    """训练 KAN 或 MLP 模型。

    Args:
        model: KAN 或 MLP 实例
        optimizer_type: 'adam' 或 'lbfgs'

    Returns:
        训练历史字典
    """
    if optimizer_type == "lbfgs":
        optimizer = torch.optim.LBFGS(
            model.parameters(), lr=lr, max_iter=20,
            history_size=10, line_search_fn="strong_wolfe"
        )
    else:
        optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    criterion = nn.MSELoss()
    history = {"train_loss": [], "test_loss": []}

    for epoch in range(epochs):
        model.train()

        if optimizer_type == "lbfgs":
            def closure():
                optimizer.zero_grad()
                pred = model(x_train)
                loss = criterion(pred, y_train)
                loss.backward()
                return loss
            loss = optimizer.step(closure)
        else:
            optimizer.zero_grad()
            pred = model(x_train)
            loss = criterion(pred, y_train)
            loss.backward()
            optimizer.step()

        # 评估
        model.eval()
        with torch.no_grad():
            train_loss = criterion(model(x_train), y_train).item()
            test_loss = criterion(model(x_test), y_test).item()

        history["train_loss"].append(train_loss)
        history["test_loss"].append(test_loss)

        if (epoch + 1) % 100 == 0:
            print(f"Epoch {epoch+1}/{epochs} | "
                  f"Train MSE: {train_loss:.6f} | "
                  f"Test MSE: {test_loss:.6f}")

    return history


# 目标函数: f(x1, x2) = exp(sin(pi * x1) + x2^2)
target_fn = lambda x: torch.exp(
    torch.sin(torch.pi * x[:, 0:1]) + x[:, 1:2] ** 2
)

x_train, y_train, x_test, y_test = create_dataset(target_fn)

# 训练 KAN
kan_model = KAN(width=[2, 5, 1], grid_size=5, spline_order=3)
kan_history = train_model(kan_model, x_train, y_train, x_test, y_test,
                          epochs=500, lr=1e-2, optimizer_type="adam")

# 训练 MLP(为了对齐参数量而加宽)
mlp_model = MLP(width=[2, 50, 50, 1])
mlp_history = train_model(mlp_model, x_train, y_train, x_test, y_test,
                          epochs=500, lr=1e-3, optimizer_type="adam")

# 可视化损失曲线对比
fig, ax = plt.subplots(1, 1, figsize=(10, 6))
ax.semilogy(kan_history["test_loss"], label="KAN [2,5,1]", linewidth=2)
ax.semilogy(mlp_history["test_loss"], label="MLP [2,50,50,1]", linewidth=2)
ax.set_xlabel("Epoch")
ax.set_ylabel("Test MSE (log scale)")
ax.set_title("KAN vs MLP: Function Approximation")
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("kan_vs_mlp_loss.png", dpi=150)
plt.show()

使用 pykan 库进行训练与可视化

在实际研究或实验中,使用官方 pykan 库会更加方便:

# pip install pykan
from kan import KAN as PyKAN
import torch

# 定义目标函数
target_fn = lambda x: torch.exp(
    torch.sin(torch.pi * x[:, [0]]) + x[:, [1]] ** 2
)

# 生成数据集(使用 pykan 内置工具)
dataset = {
    "train_input": torch.rand(1000, 2) * 2 - 1,
    "test_input": torch.rand(200, 2) * 2 - 1,
}
dataset["train_label"] = target_fn(dataset["train_input"])
dataset["test_label"] = target_fn(dataset["test_input"])

# 创建并训练 KAN 模型
model = PyKAN(width=[2, 5, 1], grid=5, k=3, seed=42)

# 执行训练(使用 LBFGS 优化器)
results = model.fit(dataset, opt="LBFGS", steps=50, lamb=0.01)

# 可视化学习到的激活函数
# 以图形方式展示每条边学习到的 B 样条函数
model.plot()

# 符号回归:尝试从学习到的函数中提取数学公式
# 通过 model.auto_symbolic() 映射为 sin、exp 等符号函数
model.auto_symbolic()
formula = model.symbolic_formula()
print(f"提取出的公式: {formula}")
# 预期输出: exp(sin(pi*x1) + x2^2)

实际应用案例

1. 求解偏微分方程(PDE)

在物理学与工程学领域,KAN 正作为求解偏微分方程的神经网络(Physics-Informed Neural Networks, PINNs)的替代方案受到关注。论文中针对泊松方程 2u=f\nabla^2 u = f 的实验表明,KAN 用比基于 MLP 的 PINN 少 100 倍的参数,达到了相同的精度。

应用领域:

  • 流体力学模拟
  • 热传导方程
  • 波动方程

2. 符号回归(Symbolic Regression)

KAN 最强大的应用,是 从数据中发现数学公式 的符号回归。将学习到的激活函数可视化并与已知的符号函数(sin、cos、exp 等)进行匹配,就能提取出隐藏在数据中的数学关系。

KAN 2.0 论文(Liu et al., 2024)在此基础上进一步扩展,重新发现了以下物理定律:

  • 守恒量(Conserved Quantities):力学系统的能量守恒定律
  • 拉格朗日量(Lagrangians):由变分原理导出的运动方程
  • 对称性(Symmetries):基于诺特定理的对称结构
  • 本构定律(Constitutive Laws):材料的应力-应变关系

3. 纽结理论(Knot Theory)

在数学领域,KAN 同样给出了有意义的结果——成功重新发现了纽结不变量(knot invariants)之间的非平凡关系,这也表明 AI 在纯数学研究中的角色有望进一步扩大。

4. 凝聚态物理学

在识别相变(phase transition)边界的任务中,KAN 学习到的边界比 MLP 更精确,并且通过学习到的激活函数,还能够解释相变的物理意义。

局限性与未来方向

当前的局限性

1. 训练速度

KAN 的训练速度明显慢于 MLP。由于需要计算样条基函数,并且每条边都要应用不同的函数,这种结构特性使其无法充分利用 GPU 的批量矩阵乘法加速。在相同任务上,也有报告显示 KAN 的训练时间是 MLP 的 10 倍以上。

2. 高维输入的可扩展性

随着输入维度升高(例如图像的数万像素),KAN 的性能会下降。在 CIFAR-10 这类标准计算机视觉基准上,KAN 相较于 MLP 并未展现出明显优势,样条参数的调优也变得困难。

3. 对噪声的敏感性

KAN 在干净的函数型数据上表现出色,但在噪声较大的真实数据上,可能比 MLP 更敏感。这是因为样条函数存在对噪声过拟合(overfit)的倾向。

4. 内存占用

由于每条边都需要存储 G+kG + k 个样条系数以及网格信息,与相同结构的 MLP 相比,内存占用相当高。在大规模网络中,这可能成为实际的瓶颈。

5. 标准任务上的表现

在机器学习、计算机视觉、NLP、音频处理等常规任务上,MLP 的平均精度普遍高于 KAN。KAN 真正展现优势的领域,仅限于符号公式表示、科学发现等特定方向。

未来研究方向

1. 高效的 KAN 实现

EfficientKAN、FastKAN、FourierKAN 等多种旨在提升 KAN 计算效率的变体已被提出。研究者们正在探索用切比雪夫多项式或傅里叶基替代 B 样条,能否在提升训练速度的同时,保留 KAN 的优势。

2. 混合架构

结合 MLP 与 KAN 各自优点的混合架构,是一个有前景的方向。例如,有研究尝试在 Transformer 的 FFN 层中用 KAN 替代 MLP,或者只在特定层中选择性地使用 KAN。

3. KAN 2.0 与科学发现

KAN 2.0 论文通过 MultKAN(引入乘法节点)、网格自动优化,以及在多个科学领域的应用,大幅扩展了 KAN 的实用性。这项发表于 Physical Review X 的后续研究表明,KAN 有潜力超越单纯的函数近似工具,成为一种 科学发现工具

故障排查与优化技巧

1. 训练不收敛时

KAN 训练中最常见的问题是不收敛。可以检查以下几点:

# 检查网格范围是否覆盖了输入数据的范围
x_min, x_max = x_train.min().item(), x_train.max().item()
print(f"数据范围: [{x_min:.2f}, {x_max:.2f}]")
# 把网格范围设置得比数据范围略宽一些
# 可调整为 grid_range=(-1.5, 1.5) 等

# 使用 LBFGS 优化器时,必须使用 closure 模式
optimizer = torch.optim.LBFGS(
    model.parameters(),
    lr=0.1,           # LBFGS 通常使用更大的 lr
    max_iter=20,       # 内部迭代次数
    tolerance_grad=1e-7,
    tolerance_change=1e-9,
    line_search_fn="strong_wolfe"  # 稳定的线搜索
)

# 通过调整 grid_eps 来适配网格
# 建议初期设为 0.5 左右,训练后期再降到 0.0

2. 防止过拟合

KAN 的表达能力非常强,因此容易过拟合:

  • 正则化:pykan 中可以用 lamb 参数施加 L1 正则化,也可以使用熵正则化(lamb_entropy)
  • 调整网格大小:初期从较小的 GG(例如 3)开始,等训练稳定后再逐步增大
  • 剪枝(Pruning):训练结束后,移除不必要的边(激活函数接近 0 的那些)

3. 内存优化

# 使用梯度检查点(gradient checkpointing)以节省内存
from torch.utils.checkpoint import checkpoint

class MemoryEfficientKAN(nn.Module):
    def __init__(self, width, grid_size=5, spline_order=3):
        super().__init__()
        self.layers = nn.ModuleList([
            KANLayer(width[i], width[i+1], grid_size, spline_order)
            for i in range(len(width) - 1)
        ])

    def forward(self, x):
        for layer in self.layers:
            # 梯度检查点: 节省内存(反向传播时重新计算)
            x = checkpoint(layer, x, use_reentrant=False)
        return x

4. 通过可视化进行调试

KAN 的一大优势,是可以直接可视化学习到的激活函数。要确认训练是否正常进行,定期绘制每条边的激活函数会很有帮助:

import matplotlib.pyplot as plt
import torch
import numpy as np


def visualize_kan_activations(model, layer_idx=0, input_range=(-1, 1)):
    """可视化 KAN 层学习到的激活函数。

    Args:
        model: KAN 模型实例
        layer_idx: 要可视化的层索引
        input_range: 输入范围
    """
    layer = model.layers[layer_idx]
    n_in = layer.in_features
    n_out = layer.out_features

    # 生成用于可视化的输入
    x_vis = torch.linspace(input_range[0], input_range[1], 200)

    fig, axes = plt.subplots(n_out, n_in, figsize=(4 * n_in, 3 * n_out))
    if n_out == 1:
        axes = axes.reshape(1, -1)
    if n_in == 1:
        axes = axes.reshape(-1, 1)

    for j in range(n_out):
        for i in range(n_in):
            ax = axes[j][i]

            # 计算单个输入维度对应的激活函数
            x_input = torch.zeros(200, n_in)
            x_input[:, i] = x_vis
            with torch.no_grad():
                # 仅提取样条部分
                basis = layer.b_splines(x_input)  # [200, in, G+k]
                spline_out = torch.einsum(
                    "big,g->bi", basis[:, i:i+1, :],
                    layer.spline_weight[j, i, :]
                )

            ax.plot(x_vis.numpy(), spline_out[:, 0].numpy(),
                    linewidth=2, color="blue")
            ax.set_title(f"phi_{j+1},{i+1}(x)", fontsize=10)
            ax.grid(True, alpha=0.3)
            ax.axhline(y=0, color="gray", linestyle="--", alpha=0.5)

    plt.suptitle(f"Layer {layer_idx} Activation Functions", fontsize=14)
    plt.tight_layout()
    plt.savefig(f"kan_layer{layer_idx}_activations.png", dpi=150)
    plt.show()


# 训练后进行可视化
# visualize_kan_activations(kan_model, layer_idx=0)

5. 性能比较时的注意事项

要公平地比较 KAN 与 MLP,需要注意以下几点:

  • 控制参数量:应该以相同的参数量而非相同的网络结构来比较
  • 控制训练时间:以总训练时间(wall-clock time)而非 epoch 数来比较更符合实际
  • 任务适配性:KAN 在低维函数近似上更有优势,MLP 在高维模式识别上更强,比较时应涵盖多种任务
  • 优化器选择:KAN 往往在 LBFGS 下表现最佳,MLP 则在 Adam 下表现最佳,因此为每个模型使用各自最优的优化器才算公平

参考资料

  1. KAN: Kolmogorov-Arnold Networks - Ziming Liu, Yixuan Wang, Sachin Vaidya, Fabian Ruehle, James Halverson, Marin Soljacic, Thomas Y. Hou, Max Tegmark (ICLR 2025) https://arxiv.org/abs/2404.19756

  2. KAN 2.0: Kolmogorov-Arnold Networks Meet Science - Ziming Liu et al. (Physical Review X, 2024) https://arxiv.org/abs/2408.10205

  3. pykan: Official Python Library for KAN - KindXiaoming (GitHub) https://github.com/KindXiaoming/pykan

  4. KAN or MLP: A Fairer Comparison - Runpeng Yu, Weihao Yu, Xinchao Wang (2024) https://arxiv.org/abs/2407.16674

  5. Kolmogorov-Arnold Representation Theorem - Wikipedia https://en.wikipedia.org/wiki/Kolmogorov%E2%80%93Arnold_representation_theorem

  6. pykan Documentation - Kolmogorov Arnold Network Documentation https://kindxiaoming.github.io/pykan/

  7. Kolmogorov-Arnold Networks: The Latest Advance in Neural Networks, Simply Explained - Towards Data Science https://towardsdatascience.com/kolmogorov-arnold-networks-the-latest-advance-in-neural-networks-simply-explained-f083cf994a85/

  8. A Comprehensive and FAIR Comparison Between MLP and KAN - ScienceDirect (2024) https://www.sciencedirect.com/science/article/abs/pii/S0045782524005462

현재 단락 (1/495)

深度学习的基本构成单元 **Multi-Layer Perceptron**(MLP)自 1958 年 Frank Rosenblatt 提出感知机以来,60 多年来一直是神经网络架构的核心。MLP ...

작성 글자: 0원문 글자: 19,837작성 단락: 0/495