AC-PKAN

[TMLR 2026] AC-PKAN: Attention-Enhanced and Chebyshev Polynomial-Based Physics-Informed Kolmogorov–Arnold Networks

This repository provides a minimal demo for AC-PKAN, a Physics-Informed Kolmogorov–Arnold Network (PINN + KAN family) that integrates:

• Chebyshev Type-I polynomial KAN layers (Cheby1KAN) for memory-efficient function representation.
• Internal (feature-wise) attention to mitigate rank collapse in deep Cheby1KAN stacks.
• External Residual–Gradient Attention (RGA) to adaptively reweight PINN losses, stabilizing optimization under stiff gradients and residual imbalance.

If you are looking for the full experimental pipeline in the paper (9 benchmark tasks, 12 baselines, ablations, etc.), please refer to the paper for exact settings and adapt this repo as a starting point.

Physics-Informed Neural Networks (PINNs) often struggle with:

• multi-scale solutions and spectral bias,
• stiff optimization (exploding/vanishing gradients),
• loss imbalance between PDE residual vs. boundary/initial/data terms,
• failure modes where optimization gets trapped in overly smooth local minima.

Meanwhile, KAN variants can improve expressiveness, but Chebyshev-based KAN stacks can suffer from feature space collapse / rank diminution as depth increases.

AC-PKAN addresses these issues by combining internal attention (architecture-level) and external attention (training-level) in a single PINN framework.

AC-PKAN replaces B-splines in vanilla KANs with Chebyshev Type-I polynomials, forming the Cheby1KAN layer. This removes grid storage and update overhead while retaining strong nonlinear approximation capacity.

A lightweight internal attention module computes two feature views from the embedded input, typically denoted U and V, and injects a layerwise gate $ \alpha^{(\ell)} $ into the Cheby1KAN stack to preserve feature diversity and prevent rank collapse.

Intuition: instead of blindly stacking polynomial expansions (which can collapse representation rank), each layer selectively emphasizes complementary feature channels derived from the input.

PINN training usually optimizes:

• PDE residual loss $L_r$
• data / boundary / initial condition losses $L_d$

AC-PKAN introduces RGA, combining:

• Residual-Based Attention (RBA): pointwise weights that focus learning on high-residual regions.
• Gradient-Related Attention (GRA): termwise scaling based on gradient norms to mitigate stiffness and imbalance.

A practical detail from the paper: the GRA weights can vary over a large dynamic range, so a log transform may be applied when integrating them into the loss to improve stability.

.
├── AC_PKAN_demo.ipynb          # runnable demo / walkthrough
├── Cheby1KANLayer.py           # Cheby1KAN implementation (Chebyshev Type-I polynomial KAN layer)
├── chwbykansattention.py       # internal attention module (U/V gating + layerwise α injection)
├── LICENSE                     # MIT
└── README.md

Tested with a standard scientific Python environment.

git clone https://github.com/fogradio/ACPKAN.git
cd ACPKAN

python -m venv .venv
source .venv/bin/activate  # (Windows) .venv\Scripts\activate

pip install --upgrade pip
pip install numpy scipy matplotlib jupyter
pip install torch

conda create -n acpkan python=3.10 -y
conda activate acpkan
pip install numpy scipy matplotlib jupyter
pip install torch

Notes:

• GPU is optional for the demo.
• If you already have a CUDA-enabled PyTorch installed, the demo should automatically benefit.

jupyter notebook AC_PKAN_demo.ipynb

The notebook is intended to:

• instantiate the AC-PKAN building blocks,
• run a small PINN-style training loop (representative PDE),
• visualize training curves and solution prediction.

This repo is intentionally lightweight. The authoritative usage is shown in AC_PKAN_demo.ipynb, but the typical structure is:

1. Embed input coordinates ((t, x, y, ...)) with a linear projection.
2. Apply internal attention to obtain the layerwise gated representation.
3. Stack Cheby1KAN layers (with residual connections / normalization as used in the paper).
4. Train with PINN losses, optionally adding the RGA controller.

Files to look at:

• Cheby1KANLayer.py for the polynomial KAN layer.
• chwbykansattention.py for the internal attention mechanism.
• AC_PKAN_demo.ipynb for end-to-end usage.

If you find this useful, please cite the paper:

@article{zhang2025ac,
  title={AC-PKAN: Attention-Enhanced and Chebyshev Polynomial-Based Physics-Informed Kolmogorov-Arnold Networks},
  author={Zhang, Hangwei and Huang, Zhimu and Wang, Yan},
  journal={arXiv preprint arXiv:2505.08687},
  year={2025}
}

Released under the MIT License. See LICENSE.