https://github.com/stergiosba/kanx

Fast Kolmogorov-Arnold Network in JAX, initial experiments
https://github.com/stergiosba/kanx

Last synced: about 1 month ago
JSON representation

Fast Kolmogorov-Arnold Network in JAX, initial experiments

• Host: GitHub
• URL: https://github.com/stergiosba/kanx
• Owner: stergiosba
• Created: 2024-05-19T00:53:01.000Z (8 months ago)
• Default Branch: main
• Last Pushed: 2024-05-20T02:18:15.000Z (8 months ago)
• Last Synced: 2024-05-21T02:54:30.410Z (8 months ago)
• Language: Python
• Size: 66.4 KB
• Stars: 1
• Watchers: 1
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md

Awesome Lists containing this project

• awesome-kan - KANX - Arnold Network in JAX |  (Library / Theorem)

README

        
# KANX: Fast Implementation (Approximation) of Kolmogorov-Arnold Network in JAX

Work in progress

## Introduction

Fast Kolmogorov-Arnold Network in JAX based on [`fast-kan`](https://github.com/ZiyaoLi/fast-kan) using [`equinox`](https://github.com/patrick-kidger/equinox).

The original implementation of KAN is [`pykan`](https://github.com/KindXiaoming/pykan).

## Installation

```bash

pip install .

pip install -r requirements.txt

```

## Example

KANX comes with an example on MNIST:

```bash

python examples/train_mnist.py

```

## Benchmark

We tested the implementation on MNIST and report the following wall-time for 3000 epochs:

| Architecture    | Wall time (sec)|

| -------- | ------- |

| CPU (i5-1135G7)  | 130.51   |

| CPU (i9-12900K) | 67.85     |

| GPU (RTX 3070 Ti)    | 13.55    |

Plots from the GPU experiment:





More experiments to come...