TimeKAN

TimeKAN is a Pytorch implementation of integrating Kolmogorov-Arnold Networks (KAN) for temporal data with recurrent neural network architectures (Currently LSTM and GRU). It is still in an experimental stage, the implementation suffer from exploding gradients and vanishing gradients problems but with careful training it can perform well specifically on non-linear/complex temporal data.

Inspired by:

• Kolmogorov-Arnold Networks
• Chebyshev Polynomial-Based KANs,
• ChebyKAN
• FourierKAN
• TKAN

Full documentation can be found here.

In tKANLSTM, KAN layers replace the output gate, computing $o_t = \sigma(\text{KAN}(W_x x_t + W_h h_{t-1}))$. In tKANGRU, they form the candidate hidden state, $\tilde{h}t = \tanh(\text{KAN}(W_x x_t + W_h (r_t \odot h{t-1}))$. The layer basis functions can be Fourier series, Chebyshev polynomials, or splines.

Here's how it can perform on Rossler system signal:

The table below compares TimeKAN (using tKANLSTM and spline as basic functions) and a standard bidirectional LSTM on three chaotic datasets available in timekan.utils.datasets. Metrics include Mean Absolute Error (MAE) and training time (seconds) until convergence.

Installation

Install TimeKAN via pip:

pip install timekan

Alternatively, clone the repository and install locally:

git clone https://github.com/SamerMakni/timekan.git
cd timekan
pip install .

Requirements: Python >= 3.9, PyTorch >= 2.4.0

Usage

Here’s a simple example training a TKANLSTM on Mackey-Glass data:

import torch
import torch.nn as nn
from timekan.models.tkan_lstm import tKANLSTM
from timekan.utils.datasets import mackey_glass

class RecurrentKAN(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.tkan = tKANLSTM(
            input_dim=input_dim,
            hidden_dim=hidden_dim,
            return_sequences=False,
            bidirectional=True,
            kan_type='fourier',
            sub_kan_configs={'gridsize': 50, 'addbias': True}
        )
        self.regressor = nn.Linear(hidden_dim * 2, 1)

    def forward(self, x):
        features = self.tkan(x)
        return self.regressor(features).squeeze(-1)

x_train, y_train, x_test, y_test = mackey_glass()

model = RecurrentKAN(input_dim=1, hidden_dim=16)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(10):
    model.train()
    optimizer.zero_grad()
    outputs = model(x_train)
    loss = criterion(outputs, y_train)
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch + 1}/10, Training MSE: {loss.item():.4f}")

model.eval()
with torch.no_grad():
    test_outputs = model(x_test)
    test_mse = criterion(test_outputs, y_test).item()
    print(f"Test MSE: {test_mse:.4f}")