AC-PKAN: Attention-Enhanced and Chebyshev Polynomial-Based Physics-Informed Kolmogorov-Arnold Networks
Hangwei Zhang, Zhimu Huang, Yan Wang
TL;DR
This paper tackles the challenge of solving PDEs with physics-informed neural networks in data-sparse settings by addressing instability and limited expressivity in Chebyshev-based KANs. It introduces AC-PKAN, which combines linear up/down projections, wavelet-activated learnable activations, and an internal attention mechanism to preserve a full-rank Jacobian and enable arbitrary-order PDE approximation. An external Residual-Gradient Attention (RGA) mechanism dynamically reweights loss terms based on residual magnitudes and gradient norms, improving training stability and convergence. Across nine benchmark PDEs in three domains, AC-PKAN consistently matches or outperforms state-of-the-art models like PINNsFormer, with ablation studies validating the importance of internal attention, wavelet activations, and RGA. The work advances PINNs by integrating Chebyshev-based KANs with attention, offering a robust and scalable framework for complex, real-world engineering PDE problems in zero-data or data-sparse regimes.
Abstract
Kolmogorov-Arnold Networks (KANs) have recently shown promise for solving partial differential equations (PDEs). Yet their original formulation is computationally and memory intensive, motivating the introduction of Chebyshev Type-I-based KANs (Chebyshev1KANs). Although Chebyshev1KANs have outperformed the vanilla KANs architecture, our rigorous theoretical analysis reveals that they still suffer from rank collapse, ultimately limiting their expressive capacity. To overcome these limitations, we enhance Chebyshev1KANs by integrating wavelet-activated MLPs with learnable parameters and an internal attention mechanism. We prove that this design preserves a full-rank Jacobian and is capable of approximating solutions to PDEs of arbitrary order. Furthermore, to alleviate the loss instability and imbalance introduced by the Chebyshev polynomial basis, we externally incorporate a Residual Gradient Attention (RGA) mechanism that dynamically re-weights individual loss terms according to their gradient norms and residual magnitudes. By jointly leveraging internal and external attention, we present AC-PKAN, a novel architecture that constitutes an enhancement to weakly supervised Physics-Informed Neural Networks (PINNs) and extends the expressive power of KANs. Experimental results from nine benchmark tasks across three domains show that AC-PKAN consistently outperforms or matches state-of-the-art models such as PINNsFormer, establishing it as a highly effective tool for solving complex real-world engineering problems in zero-data or data-sparse regimes. The code will be made publicly available upon acceptance.