[论文解读] Wav-KAN: Wavelet Kolmogorov-Arnold Networks

Zavareh Bozorgasl, Hao Chen|arXiv (Cornell University)|May 21, 2024

Neural Networks and Applications被引用 7

一句话总结

Wav-KAN 将小波变换整合到 Kolmogorov-Arnold 网络中，以提升可解释性、效率和鲁棒性，在某些任务中优于 Spl-KAN 和 MLP。

ABSTRACT

In this paper, we introduce Wav-KAN, an innovative neural network architecture that leverages the Wavelet Kolmogorov-Arnold Networks (Wav-KAN) framework to enhance interpretability and performance. Traditional multilayer perceptrons (MLPs) and even recent advancements like Spl-KAN face challenges related to interpretability, training speed, robustness, computational efficiency, and performance. Wav-KAN addresses these limitations by incorporating wavelet functions into the Kolmogorov-Arnold network structure, enabling the network to capture both high-frequency and low-frequency components of the input data efficiently. Wavelet-based approximations employ orthogonal or semi-orthogonal basis and maintain a balance between accurately representing the underlying data structure and avoiding overfitting to the noise. While continuous wavelet transform (CWT) has a lot of potentials, we also employed discrete wavelet transform (DWT) for multiresolution analysis, which obviated the need for recalculation of the previous steps in finding the details. Analogous to how water conforms to the shape of its container, Wav-KAN adapts to the data structure, resulting in enhanced accuracy, faster training speeds, and increased robustness compared to Spl-KAN and MLPs. Our results highlight the potential of Wav-KAN as a powerful tool for developing interpretable and high-performance neural networks, with applications spanning various fields. This work sets the stage for further exploration and implementation of Wav-KAN in frameworks such as PyTorch and TensorFlow, aiming to make wavelets in KAN as widespread as activation functions like ReLU and sigmoid in universal approximation theory (UAT). The codes to replicate the simulations are available at https://github.com/zavareh1/Wav-KAN.

研究动机与目标

• 推动可解释神经网络的发展，解决 MLP 与 Spl-KAN 在可解释性、训练速度和鲁棒性方面的局限。
• 引入基于小波的 KAN 扩展（Wav-KAN），以捕捉高频和低频数据分量。
• 提出一种带小波激活的多层 KAN 架构，以提高效率和准确性。
• 在 MNIST 上演示该方法并讨论相对于传统激活函数和样条基 KAN 的优点。

提出的方法

• 在层与层之间的边上用可学习的一元小波函数取代传统权重。
• 在 KAN 框架内将连续小波变换和离散小波变换用作激活函数。
• 使用多层 KAN 结构，通过运算符 T_o 进行矩阵样聚合以对激活输出求和。
• 结合批归一化以提高训练速度和准确性。
• 在参数、速度和鲁棒性方面将 Wav-KAN 与 Spl-KAN 和 MLPs 进行比较。
• 证明小波选择（Mexican hat、Morlet、DOG、Shannon）会影响 MNIST 上的性能。

Figure 1: Wav-KAN with arbitrary number of layers (here is Wav-KAN[2,3,2])

实验结果

研究问题

• RQ1 Wav-KAN 是否能在图像分类任务上达到比 Spl-KAN 和 MLP 更高的准确率？
• RQ2 将小波整合入 KAN 是否在保持可解释性的同时提升训练速度和鲁棒性？
• RQ3 不同母小波如何影响 Wav-KAN 在 MNIST 上的性能？
• RQ4 在各层中，Wav-KAN 相对于 Spl-KAN 与 MLP 的参数效率是多少？
• RQ5 Wav-KAN 是否能够有效整合到像 PyTorch 或 TensorFlow 这样的主流框架中？

主要发现

• 在测试配置下，Wav-KAN 相较于 Spl-KAN 在 MNIST 上实现了更快的训练和更高的准确性。
• 基于小波的激活函数能够高效表示局部细节和全局结构。
• 批归一化进一步提高了 Wav-KAN 和 Spl-KAN 的准确性和训练速度。
• 不同母小波显著影响性能，在某些设置中 Shannon 小波及其他某些小波表现不佳。
• 由于小波容量（权重、平移、缩放），Wav-KAN 在同等任务中使用的参数少于 Spl-KAN。
• 该方法被定位为在 KAN 家族中更具可解释性和可扩展性。

Figure 2: Training accuracy of Wav-KAN [28*28,32,10] versus Spl-KAN [28*28,32,10]