← Go back home

Wavlet Kolmogorov-Arnold Networks

Sharing two repositories for Wavelet KAN.

July 22, 2024 12:00:00 PM PDT

  • research
  • open-source
Morlet WaveletMorlet WaveletMorlet Wavelet

Hello!

This is a small post to share two repositories that we developed for our research on Wavelet Kolmogorov-Arnold Networks (KANs). These repositories were developed as part of our exploration of KANs for solving partial differential equations (PDEs) in the context of scientific machine learning.

The Repositories

These repositories explore the application of wavelet Kolmogorov-Arnold Networks (KANs) to two specific problems:

1D Viscoplastic Material Response

Using a Recurrent Neural Operator (RNO) and a Transformer.

2D Darcy Flow

Using a Convolutional Neural Network (CNN) a Fourier Neural Operator (FNO).

Key Takeaways

These repositories may still be valuable for those interested in:

  • Julia implementations of wavelet KAN.
  • KAN implementations of RNO, Transformer, and CNN.
  • Convolution operations implemented from scratch in Julia for the KAN models.

Key takeaways from our experiments include:

  • MLP models generally outperformed our limited complexity wavelet-KAN implementations. The wavelet-KANs can definitely be refined and tuned further, e.g. with pruning, and it's difficult to say if MLPs are better than wavelet-KANs given the reduced scale.
  • Hyperparameter tuning for wavelet-KANs was challenging and unstable, even with batch/layer norm.
  • Deeper wavelet-KANs seemingly performed worse, highlighting the need for pruning as suggested in the arXiv: original KAN paper (not implemented here).
  • wavelet-KAN models struggled to fit on the GPU, limiting the size of the Transformer and preventing the implementation of the FNO.

Feel free to explore/use the repositories for more details. Please note the references in the READMEs for the original papers, implementations, and dataset sources.

// <3 Exa

Message signature

-----BEGIN PGP SIGNED MESSAGE-----
Hash: SHA512

Hello!

This is a small post to share two repositories that we developed for our research on Wavelet Kolmogorov-Arnold Networks (KANs). These repositories were developed as part of our exploration of KANs for solving partial differential equations (PDEs) in the context of scientific machine learning.

## The Repositories

- - [github.com/exa-laboratories/Julia-Wav-KAN](https://github.com/exa-laboratories/Julia-Wav-KAN)
- - [github.com/exa-laboratories/wavKAN-conv2D](https://github.com/exa-laboratories/wavKAN-conv2D)

These repositories explore the application of wavelet Kolmogorov-Arnold Networks (KANs) to two specific problems:

##### 1D Viscoplastic Material Response
Using a Recurrent Neural Operator (RNO) and a Transformer.

##### 2D Darcy Flow
Using a Convolutional Neural Network (CNN) a Fourier Neural Operator (FNO).

## Key Takeaways
These repositories may still be valuable for those interested in:

- - Julia implementations of wavelet KAN.
- - KAN implementations of RNO, Transformer, and CNN.
- - Convolution operations implemented from scratch in Julia for the KAN models.

Key takeaways from our experiments include:

- - MLP models generally outperformed our limited complexity wavelet-KAN implementations. The wavelet-KANs can definitely be refined and tuned further, e.g. with pruning, and it's difficult to say if MLPs are better than wavelet-KANs given the reduced scale.
- - Hyperparameter tuning for wavelet-KANs was challenging and unstable, even with batch/layer norm.
- - Deeper wavelet-KANs seemingly performed worse, highlighting the need for pruning as suggested in the [arXiv: original KAN paper](https://arxiv.org/abs/2404.19756) (not implemented here).
- - wavelet-KAN models struggled to fit on the GPU, limiting the size of the Transformer and preventing the implementation of the FNO. 

Feel free to explore/use the repositories for more details. Please note the references in the READMEs for the original papers, implementations, and dataset sources.

// &lt;3 Exa 
-----BEGIN PGP SIGNATURE-----

iHUEARYKAB0WIQTqH4Le1ZB+MwLlm9jtpSZBtyZXuAUCZprIWAAKCRDtpSZBtyZX
uOu0AQDdzv6gOQhn4wAk2krAuJRAMyaEPenaeB7uS+De4D0cZwD/T+SxvMp5E1M0
lhKQDyL+KJfUT0bH4kia5n6IN5Jv7ws=
=WC7O
-----END PGP SIGNATURE-----