[ENH] Kolmogorov Arnold Block for NBeats #1751

Sohaib-Ahmed21 · 2025-01-13T10:35:11Z

Description

This PR adds Kolmogorov Arnold(KAN) Blocks in NBeats and also does refactoring of NBeats. Implementation of KAN blocks' layers is taken from original paper code.

Changes in Structure

Introduced the NBEATSKAN module, which enables usage of KAN blocks within the NBEATS architecture.
Integrated KANLayer logic, implemented in kan_layer.py, which handles KAN-specific operations such as:
- Spline coefficient computation,
- Grid initialization and updates, etc.
Imported KANLayer to submodules.py for block operations, allowing NBEATSKAN to delegate block-level behavior through use_kan=True.
Added the NBEATSAdapter class to encapsulate common methods shared by both NBEATS and NBEATSKAN, including:
- Standard training, forward logic, etc.
- Excludes block initialization (__init__), which is separately defined in each class to maintain architectural flexibility.

GridUpdateCallback

When training KAN-based models, the grid can be iteratively refined during training for better performance.

To support this, logic from the original [pykan](https://github.com/KindXiaoming/pykan) implementation has been adapted to define a custom callback named GridUpdateCallback. This callback automatically updates the grid at specified training steps, improving model accuracy and convergence.

This callback has been tested successfully and demonstrates improved results in practice.

An example usage is provided in:
examples/nbeats_with_kan.py

Checklist

Linked issues (if existing)
Used pre-commit hooks when committing to ensure that code is compliant with hooks. Install hooks with pre-commit install.
To run hooks independent of commit, execute pre-commit run --all-files

Make sure to have fun coding!

…ctored NBeats.

…sting into kan-nbeats

Sohaib-Ahmed21 · 2025-01-13T11:01:06Z

@fkiraly @benHeid this PR is ready for review. Kindly review it. Thanks!

Sohaib-Ahmed21 · 2025-01-21T19:44:49Z

@fkiraly @benHeid can you kindly review/merge it so I integrate NBEATSX modification in NBEATS without conflicts as I have asked @julian-fong and he is not working on NBEATSX.

pytorch_forecasting/models/nbeats/_nbeats.py

benHeid

Not fully reviewed yet. Will continue in the next days. But I share my current comments so that you already receive some feedbacks.

When training KANs, the grid can be iteratively be refined. I wonder, if there is a way to implement this also here. However, this might probably more difficult and require changes to the trainer. So probably out of scope for this PR. Do you have opinions on that?

pytorch_forecasting/models/nbeats/_nbeats.py

benHeid · 2025-01-21T20:38:12Z

pytorch_forecasting/models/nbeats/kan_layer.py

@@ -0,0 +1,528 @@
+import numpy as np


The license of the original implementation is MIT. So in theory it is okay to copy the file. However, please add some credits at the top of the file.

Alternatively, we could think about adding KAN as a dependency.

@fkiraly do you have any additions on that matter?

Okay, I will add the appropriate credits at the top of the file. Additionally, I agree that adding KAN as a dependency—perhaps as a soft dependency—seems like a good idea, especially considering its increasing relevance in time series forecasting.

@fkiraly pinging you again to check if this is okay for you :)

which package are we exactly planning to add as a soft dep?

If it is a single layer, I think copying it over and including the license is perhaps better for now, because we do not have machinery to manage soft dependencies (like scikit-base or similar).

The proposed design in here sktime/enhancement-proposals#39 would allow that, but right now I think this would require a significant amounts of custom code to handle.

Or, is there an easy way that I am not seeing how the soft dependency import would work for part of the NN?

which package are we exactly planning to add as a soft dep?

It is pykan library , reference https://pypi.org/project/pykan/

If it is a single layer, I think copying it over and including the license is perhaps better for now, because we do not have machinery to manage soft dependencies (like scikit-base or similar).

The proposed design in here sktime/enhancement-proposals#39 would allow that, but right now I think this would require a significant amounts of custom code to handle.

Or, is there an easy way that I am not seeing how the soft dependency import would work for part of the NN?

yes it is a single layer, only used in NBEATS. Also the library pykan has much more, but we only need this. I have implemented what you already suggested above. I have copied it over and included the license.

ok, makes sense, as long as the license points to the original source and is provided in full form (assuming it hsa the usual reqiurement to reproduce)

Sohaib-Ahmed21 · 2025-01-21T21:30:23Z

Not fully reviewed yet. Will continue in the next days. But I share my current comments so that you already receive some feedbacks.

Thanks! Will address these reviews soon.

Sohaib-Ahmed21 · 2025-01-21T21:42:35Z

When training KANs, the grid can be iteratively be refined. I wonder, if there is a way to implement this also here. However, this might probably more difficult and require changes to the trainer. So probably out of scope for this PR. Do you have opinions on that?

I'll explore this and share my thoughts.

…casting into kan-nbeats

… while using KAN blocks in NBEATS.

Sohaib-Ahmed21 · 2025-01-23T12:51:30Z

@benHeid I have addressed the reviews. Kindly review the updated PR.

When training KANs, the grid can be iteratively be refined. I wonder, if there is a way to implement this also here. However, this > might probably more difficult and require changes to the trainer. So probably out of scope for this PR. Do you have opinions on > that?

To address this, I have taken logic from original implementation of pykan library and made custom Callback i.e. GridUpdateCallback which updates grid after specified steps. This has been tested and working fine with improved results. Your thoughts will be helpful in this regard.

Sohaib-Ahmed21 · 2025-01-26T16:02:39Z

@benHeid @fkiraly I have addressed the reviews. Kindly review it so I proceed integrating NBEATSX without conflicts.

pytorch_forecasting/models/nbeats/grid_callback.py

pytorch_forecasting/models/nbeats/kan_layer.py

benHeid

I only have one last comment. But I also would like to hear @fkiraly opinion on the license of the kan_layer

Sohaib-Ahmed21 · 2025-07-07T07:04:51Z

@benHeid I have addressed the reviews. Kindly review the updated PR, thanks!

fkiraly

Nice!

Could you kindly also add some explanation in the PR fist post about the changes to the structure, and what goes in the common base class? What is going on with the grid callback?

Also requesting review by @phoeenniixx and @PranavBhatP

Optionally, could you check in the layers module whether we can use that to avoid duplicate layer specifications?

phoeenniixx · 2025-07-07T15:40:03Z

pytorch_forecasting/models/nbeats/_kan_layer.py

+    evaluate x on B-spline bases
+
+    Args:
+    -----


Can you please use Parameters like in sktime and follow the numpydoc style docstrings, No need to add colon (:)

phoeenniixx · 2025-07-07T15:41:43Z

pytorch_forecasting/models/nbeats/_kan_layer.py

+    ):
+        """'
+        Initialize a KANLayer
+


Can you move this docstring above just after the class definition?

phoeenniixx · 2025-07-07T15:43:05Z

pytorch_forecasting/models/nbeats/_nbeatskan.py

+        shown to consistently outperform N-BEATS.
+
+        Args:
+            stack_types: One of the following values: “generic”, “seasonality" or


Please use numpydoc style doctsring

pytorch_forecasting/models/nbeats/_kan_layer.py

phoeenniixx · 2025-07-07T17:11:15Z

pytorch_forecasting/models/nbeats/_nbeatskan.py

+        forecasting <http://arxiv.org/abs/1905.10437>`_. The network has (if
+        used as ensemble) outperformed all other methods including ensembles of
+        traditional statical methods in the M4 competition. The M4 competition is
+        arguably the most important benchmark for univariate time series forecasting.


I think it would be good to mention the paper for KAN blocks here and how it works (just an overview) and how it is different from the original NBeats

Mentioned the reference with details.

pytorch_forecasting/models/nbeats/_nbeatskan.py

pytorch_forecasting/models/nbeats/_nbeatskan_pkg.py

…ferences, and extend NBEATSKAN test cases.

Sohaib-Ahmed21 · 2025-07-10T11:45:18Z

Could you kindly also add some explanation in the PR fist post about the changes to the structure, and what goes in the common base class? What is going on with the grid callback?

Updated the first post to include explanations on the structure changes, common base class, and grid callback as requested.

Optionally, could you check in the layers module whether we can use that to avoid duplicate layer specifications?

I checked it and currently blocks are reused from sub_modules.py and are not duplicated.

Sohaib-Ahmed21 · 2025-07-12T07:44:11Z

@fkiraly @benHeid kindly review/merge this pull request, thanks!

phoeenniixx · 2025-07-12T18:24:07Z

Optionally, could you check in the layers module whether we can use that to avoid duplicate layer specifications?

I checked it and currently blocks are reused from sub_modules.py and are not duplicated.

I think we can also move some of the layers (like KAN layer) to the layers right? I think these layers can be reused by other models as well?

Also, a question for @fkiraly: Imo, we could also move networks present in sub_modules to layers, or is it better to keep them here along with NBeats only, as other models might not use them?

Sohaib-Ahmed21 · 2025-07-12T19:30:29Z

I think we can also move some of the layers (like KAN layer) to the layers right? I think these layers can be reused by other models as well?

Yes, in the future, KANLayer might be reused...I think then GridUpdateCallback should also be moved to layers. Thoughts?

phoeenniixx · 2025-07-12T19:52:23Z

I think we can also move some of the layers (like KAN layer) to the layers right? I think these layers can be reused by other models as well?

Yes, in the future, KANLayer might be reused...I think then GridUpdateCallback should also be moved to layers. Thoughts?

Yes this can also be added!
I think we need to decide which parts of this PR are reusable and can be moved to layers

But this is a callback and not essentially a neural layer, so I think moving it to layers may not be a good idea, rather a similar folder for callbacks? Not sure, what do you think?

Sohaib-Ahmed21 · 2025-07-12T19:59:27Z

Yes this can also be added! I think we need to decide which parts of this PR are reusable and can be moved to layers

But this is a callback and not essentially a neural layer, so I think moving it to layers may not be a good idea, rather a similar folder for callbacks? Not sure, what do you think?

What I was essentially thinking was that KANLayer in terms of my experience is not used much across the library...so the current implementation also sounds logical to me. Maybe we can ask from @benHeid also about what to include or the current implementation sounds good?

phoeenniixx · 2025-07-12T20:08:05Z

What I was essentially thinking was that KANLayer in terms of my experience is not used much across the library...so the current implementation also sounds logical to me.

I agree, it is not used anywhere right now, but we may need it in future, then it will be handy to have it in layers - A single API point for reusable layers.

As you said, maybe @benHeid, @fkiraly can also provide us a new perspective...

fkiraly · 2025-07-13T18:02:26Z

Also, a question for @fkiraly: Imo, we could also move networks present in sub_modules to layers, or is it better to keep them here along with NBeats only, as other models might not use them?

I think ultimately all layers should move to layers, although moving layers not changed or affected by this PR is out of scope. But I would move new layers, or layers that now get imported by both nbeats variants to there.

Sohaib-Ahmed21 added 3 commits January 13, 2025 11:39

Refactored NBeats and added comments for KAN block and NBeats.

3bab673

Refactored NBeats and added comments for KAN block and NBeats.

acfa626

End to end integrated Kolmogorov Arnold Networks in NBeats. Also refa…

41d7403

…ctored NBeats.

Sohaib-Ahmed21 requested review from benHeid, fkiraly, fnhirwa, geetu040, jdb78, pranavvp16, XinyuWuu and yarnabrina as code owners January 13, 2025 10:35

Sohaib-Ahmed21 and others added 3 commits January 13, 2025 15:41

Merge branch 'main' into kan-nbeats

53fb126

Resolved import error.

594102d

Merge branch 'kan-nbeats' of github.com:Sohaib-Ahmed21/pytorch-foreca…

45c63f6

…sting into kan-nbeats

benHeid reviewed Jan 21, 2025

View reviewed changes

pytorch_forecasting/models/nbeats/_nbeats.py Outdated Show resolved Hide resolved

benHeid requested changes Jan 21, 2025

View reviewed changes

Sohaib-Ahmed21 added 2 commits January 22, 2025 04:18

Merge branch 'main' of https://github.com/Sohaib-Ahmed21/pytorch-fore…

88de705

…casting into kan-nbeats

Refactored NBEATS and added support for grid updation during training…

c8ccfaf

… while using KAN blocks in NBEATS.

Sohaib-Ahmed21 requested a review from benHeid January 23, 2025 13:06

Refactored comments.

348da97

benHeid reviewed Feb 1, 2025

View reviewed changes

pytorch_forecasting/models/nbeats/grid_callback.py Outdated Show resolved Hide resolved

benHeid reviewed Feb 1, 2025

View reviewed changes

pytorch_forecasting/models/nbeats/kan_layer.py Outdated Show resolved Hide resolved

benHeid reviewed Feb 1, 2025

View reviewed changes

Merge branch 'sktime:main' into kan-nbeats

09facba

Sohaib-Ahmed21 and others added 4 commits July 5, 2025 01:02

Merge branch 'sktime:main' into kan-nbeats

2a4d3ec

Address deprecated typing classes

14ca66f

Refactor code with proper docstrings and cleaner structure

89a9a4f

Refactor examples in docstring

0c43448

fkiraly changed the title ~~Kolmogorov Arnold Block for NBeats~~ [ENH] Kolmogorov Arnold Block for NBeats Jul 6, 2025

Include NBEATSKAN package container

2da4d13

Sohaib-Ahmed21 requested a review from benHeid July 7, 2025 05:11

fkiraly requested changes Jul 7, 2025

View reviewed changes

phoeenniixx reviewed Jul 7, 2025

View reviewed changes

pytorch_forecasting/models/nbeats/_kan_layer.py Show resolved Hide resolved

phoeenniixx reviewed Jul 7, 2025

View reviewed changes

pytorch_forecasting/models/nbeats/_nbeatskan.py Show resolved Hide resolved

phoeenniixx reviewed Jul 7, 2025

View reviewed changes

pytorch_forecasting/models/nbeats/_nbeatskan_pkg.py Outdated Show resolved Hide resolved

Refactor and enhance docstrings to follow NumPy style, include KAN re…

eb9c79d

…ferences, and extend NBEATSKAN test cases.

Sohaib-Ahmed21 requested review from fkiraly and phoeenniixx July 10, 2025 11:45

Sohaib-Ahmed21 and others added 2 commits July 10, 2025 16:45

Merge branch 'main' into kan-nbeats

cc819ff

Merge branch 'main' into kan-nbeats

5a31a58

Merge branch 'main' into kan-nbeats

58e14d8

[ENH] Kolmogorov Arnold Block for NBeats #1751

Are you sure you want to change the base?

[ENH] Kolmogorov Arnold Block for NBeats #1751

Uh oh!

Conversation

Sohaib-Ahmed21 commented Jan 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes in Structure

GridUpdateCallback

Checklist

Uh oh!

Sohaib-Ahmed21 commented Jan 13, 2025

Uh oh!

Sohaib-Ahmed21 commented Jan 21, 2025

Uh oh!

Uh oh!

benHeid left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Sohaib-Ahmed21 commented Jan 21, 2025

Uh oh!

Sohaib-Ahmed21 commented Jan 21, 2025

Uh oh!

Sohaib-Ahmed21 commented Jan 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Sohaib-Ahmed21 commented Jan 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

benHeid left a comment

Choose a reason for hiding this comment

Uh oh!

Sohaib-Ahmed21 commented Jul 7, 2025

Uh oh!

fkiraly left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Sohaib-Ahmed21 commented Jul 10, 2025

Uh oh!

Sohaib-Ahmed21 commented Jul 12, 2025

Uh oh!

phoeenniixx commented Jul 12, 2025

Uh oh!

Sohaib-Ahmed21 commented Jul 12, 2025

Uh oh!

phoeenniixx commented Jul 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Sohaib-Ahmed21 commented Jan 13, 2025 •

edited

Loading

Sohaib-Ahmed21 commented Jan 23, 2025 •

edited

Loading

Sohaib-Ahmed21 commented Jan 26, 2025 •

edited

Loading

fkiraly left a comment •

edited

Loading

phoeenniixx commented Jul 12, 2025 •

edited

Loading

phoeenniixx commented Jul 12, 2025 •

edited

Loading