Wav-KAN: Wavelet Kolmogorov-Arnold Networks
Zavareh Bozorgasl, Hao Chen
TL;DR
The paper introduces Wav-KAN, a wavelet-augmented Kolmogorov-Arnold Network designed to improve interpretability and performance over traditional MLPs and Spl-KAN. By replacing fixed activations with learnable wavelet-based functions and leveraging multiresolution analysis via CWT/DWT, the approach aims to capture both high- and low-frequency data components efficiently. Through MNIST experiments, Wav-KAN demonstrates faster training and higher robustness, with specific wavelets (e.g., Mexican hat, DOG) yielding strong test accuracy. The work advocates Wav-KAN as a practical, interpretable neural architecture and outlines future integration with PyTorch and TensorFlow for broader adoption.
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.