MexiKAN: Mexican Hat wavelet-enhanced KANs for fine-grained signal modeling

abstract

Kolmogorov–Arnold Networks (KANs) have recently gained attention for their strong representational capacity. However, their reliance on smooth B-spline activations fundamentally limits their ability to model discontinuous or abrupt signal variations. To overcome this limitation, we propose MexiKAN, a novel KAN variant that integrates the Mexican Hat wavelet into the activation design. The hat-shaped profile of the Mexican Hat wavelet accurately modeling local discontinuities and impulsive patterns, while its spectral energy concentration in the mid-to-low frequency range effectively suppresses high-frequency noise. This combination yields superior robustness and generalization compared to standard KANs. Extensive experiments across diverse tasks demonstrate that MexiKAN achieves higher accuracy and robustness than state-of-the-art baselines.

Project Layout

MexiKAN/
|-- examples/
|   `-- train_mnist.py
|-- src/
|   `-- mexikan/
|       |-- __init__.py
|       |-- mexihat.py
|       |-- models.py
|       |-- spline.py
|       `-- utils.py
|-- LICENSE
|-- pyproject.toml
|-- requirements.txt
`-- README.md

The KAN layer and its spline utilities live inside the mexikan package, so the project does not depend on an external kan package at import time.

Installation

pip install -r requirements.txt
pip install -e .

Train On MNIST

python examples/train_mnist.py --epochs 20 --batch-size 32

Run on a specific GPU:

CUDA_VISIBLE_DEVICES=0 python examples/train_mnist.py --device cuda

Run a single seed:

python examples/train_mnist.py --seeds 0

Change model width:

python examples/train_mnist.py --hidden-dims 29
python examples/train_mnist.py --hidden-dims 64 32

Outputs

By default, results are written to:

runs/mnist_mexikan/
|-- metrics.csv
|-- summary.csv
`-- seed_*/model.pth

The console prints:

• train accuracy
• test accuracy
• macro F1
• runtime per seed
• mean and standard deviation across seeds