Skip to content

KAN Backbone ​

ForeBlocks ships a Kolmogorov–Arnold Network (KAN) backbone for time-series forecasting in foreblocks.kan. Instead of fixed activation functions on neurons, KAN places learnable, basis-expanded functions on the edges of the network. This implementation patches the input series, expands each patch through one or more orthogonal polynomial families, and optionally routes between families with a token-level mixture-of-experts (MoKAN) router.

The KAN backbone is a specialist subsystem: it is not part of the top-level foreblocks public API. Import it directly from foreblocks.kan.

When to use it ​

  • You want a non-Transformer, non-recurrent backbone with strong inductive bias for smooth/periodic structure.
  • You want to mix several basis families (e.g. Chebyshev + Fourier + wavelet) and let a router pick per token.
  • You are comparing alternative function approximators against the standard forecasting stack.

If your baseline ForecastingModel run is not working yet, start there first — see Getting Started.

Minimal example ​

KANModel consumes a [B, T, C] series and returns a [B, H, C] forecast, where T = context_window, H = target_window, and C = c_in.

python
import torch
from foreblocks.kan import KANModel

model = KANModel(
    c_in=4,              # number of input channels (features)
    context_window=48,   # input length T
    target_window=24,    # forecast horizon H
    patch_len=16,        # patch size
    stride=8,            # patch stride
    d_model=64,
    depth=2,             # number of KAN blocks
    revin=True,          # reversible instance norm on inputs
)

x = torch.randn(8, 48, 4)   # [B, T, C]
y = model(x)                # [B, 24, 4] -> [B, H, C]

Pass a subset via the families argument; per-family hyperparameters are exposed as keyword arguments (for example jacobi_alpha, jacobi_beta, wavelet_num, fourier_base_freq). Fine-grained layer behaviour can be configured with PolyLayerConfig.

Mixture of families (MoKAN routing) ​

When multiple families are active, a token-level router (TokenRouter, configured with RouterConfig) selects the top-k families per token. The relevant KANModel knobs are:

  • top_k — number of families selected per token (default 2)
  • router_temperature, router_hidden — router softmax temperature and width
  • load_balance_coef — auxiliary load-balancing loss weight

Key constructor arguments ​

ArgumentPurpose
c_in, context_window, target_windowinput channels, input length, forecast horizon
patch_len, stride, padding_patchpatching of the input series
d_model, depthhidden width and number of KAN blocks
familiessequence of PolyFamily instances to expand through
revin, affine, subtract_lastreversible instance normalization options
top_k, router_*, load_balance_coefMoKAN routing controls
head_*, block_*, final_normforecast head and block regularization

MIT License