Technical Summary: Geometric Kolmogorov–Arnold Networks (GeoKAN)

Problem Statement

Standard deep learning models, including Multilayer Perceptrons (MLPs) and recent Kolmogorov–Arnold Networks (KANs), typically operate within a fixed Euclidean coordinate system. While KANs improve approximation flexibility by replacing fixed scalar weights with learnable univariate functions, they still evaluate these functions directly on the original input coordinates. This limitation becomes critical in scientific machine learning and differential equation solving, where target functions often exhibit highly non-uniform behavior: smooth in some regions but varying rapidly near shocks, thin interfaces, sharp peaks, or boundary layers. In such scenarios, fixed-coordinate representations inefficiently allocate model capacity, often over-resolving smooth regions while under-resolving areas of rapid variation.

Methodology

The authors introduce Geometric Kolmogorov–Arnold Networks (GeoKANs), a family of models that adapt the geometry of the input space alongside the approximation function. The core mechanism involves learning a diagonal Riemannian metric $g(u)$ that warps the input coordinates before basis expansion and feature mixing.

Core Mechanism

Instead of evaluating basis functions on raw inputs $x$, GeoKAN first computes warped coordinates:
$$\xi_i = x_i \sqrt{g_i(x)}$$
where $g_i(x) > 0$ is a learned metric component. This warping stretches regions of rapid variation (increasing effective resolution) and compresses smoother regions. The model then constructs features in these geometry-adapted coordinates.

Variants

The paper proposes three main variants of GeoKAN, differing in how the metric is parameterized and how post-warp features are constructed:

1. GeoKAN-NNMetric: Uses a coupled metric learned from the full incoming representation via a neural network (MetricNet). It combines this with a localized wavelet dictionary (Mexican-hat wavelets) and a volume summary term ($\log \det g$).
2. GeoKAN-γ: Employs a structured, separable metric where each coordinate learns its own scaling independently. Instead of a post-warp dictionary, it constructs explicit geometric features derived from the metric, including the warped coordinate, the local logarithmic derivative of the metric, and the metric magnitude.
3. LM-KAN (Learned-Metric KAN): Designed specifically for physics-informed settings. It uses a coupled learned metric (like GeoKAN-NNMetric) but allows the post-warp basis to be chosen based on the problem structure: Wavelet (LM-KAN-Wav), Radial Basis Function (LM-KAN-RBF), or Fourier (LM-KAN-Fourier).

Physics-Informed Learning

In the physics-informed setting, the authors utilize the LM-KAN variant as the surrogate model within the standard Physics-Informed Neural Network (PINN) framework. The training objective remains a residual-based loss function incorporating the governing differential equation, initial conditions, and boundary conditions. The novelty lies solely in the architecture: the surrogate learns a task-dependent metric and evaluates its localized basis representation in metric-adapted coordinates, rather than changing the variational formulation of the loss.

Key Contributions

1. Geometric Inductive Bias: The paper introduces a framework where the model learns not just the function approximation but also the underlying geometry of the representation space, providing a mechanism to reallocate representational resolution task-dependently.
2. Three Distinct Architectures: The development of GeoKAN-NNMetric, GeoKAN-γ, and LM-KAN offers a spectrum of models ranging from flexible coupled-metric approaches to structured, lightweight metric-feature models.
3. Basis Adaptation: The integration of specific basis families (Wavelet, RBF, Fourier) within the geometric warp allows the model to tailor its inductive bias to specific problem types (e.g., Fourier for oscillatory Helmholtz problems).
4. Comprehensive Benchmarking: The work evaluates these models in two distinct regimes:
   • Supervised Function Approximation: Using matched-capacity benchmarks to isolate intrinsic approximation advantages.
   • Physics-Informed Learning: Solving various differential equations (Allen–Cahn, Burgers', Lorenz system, Helmholtz) to test performance in scientific computing contexts.

Results

Supervised Function Approximation

In matched-capacity benchmarks (approx. $2 \times 10^4$ parameters), GeoKAN variants demonstrated superior performance over MLPs and Efficient-KAN on targets with non-uniform structures:

• High-frequency and localized targets: GeoKAN-NNMetric and LM-KAN-Wav achieved the lowest Mean Squared Errors (MSE) on high-frequency sinusoids, step discontinuities, narrow needles, and multiscale envelope-modulated sinusoids.
• Chirp signals: LM-KAN-RBF outperformed other variants on chirp targets, suggesting RBF bases are better suited for smoothly varying local frequencies.
• Interpretation: The results confirm that the learned metric effectively redistributes resolution, stretching difficult regions and compressing smooth ones, leading to tighter fits without increasing parameter counts.

Physics-Informed Learning

The performance of LM-KAN was compared against standard PIKAN baselines (Efficient-KAN and Wav-KAN) across several differential equations:

• Allen–Cahn Equation (Case 1): Both models achieved high accuracy. PIKAN reached a slightly lower final loss, but LM-KAN trained significantly faster (315s vs. 539s) with fewer parameters.
• Allen–Cahn Equation (Case 2 - Stiff): In the more challenging regime with sharper interfaces, LM-KAN provided better solution reconstruction and lower PDE residuals, though at a higher computational cost (4760s vs. 1281s).
• Burgers' Equation: LM-KAN achieved nearly identical final loss to PIKAN but in substantially less time (206s vs. 295s), despite having more parameters, indicating improved computational efficiency in resolving steep gradients.
• Lorenz System: LM-KAN achieved a lower final loss than PIKAN, demonstrating that geometric adaptation benefits coupled nonlinear ODEs, albeit with longer training times.
• Helmholtz Equation: In this highly oscillatory regime, LM-KAN (using Fourier bases) significantly outperformed the Efficient-KAN baseline. It reached a lower final loss (approx. 45% reduction) and trained more than twice as fast (76s vs. 166s), despite having nearly identical parameter counts. The geometric adaptation proved particularly effective in preserving phase accuracy and wave structure as wavelengths decreased.

Significance and Claims

The paper claims that GeoKAN provides a more flexible and often more effective representational mechanism for scientific machine learning, particularly for problems characterized by sharp transitions, localized structures, or oscillatory complexity.

• Architectural vs. Variational: The authors emphasize that the improvement is architectural. The training framework (residual-based loss) remains standard; the gain comes from the surrogate's ability to learn a geometry-adapted coordinate system.
• Task-Dependent Resolution: The primary significance is the ability to "place capacity where it is most needed." By learning a metric that stretches difficult regions, GeoKAN avoids the inefficiency of fixed-grid representations that must globally resolve local features.
• Modest Claims: The authors acknowledge that GeoKAN is not uniformly superior in every metric. In some cases (e.g., Allen–Cahn Case 1), standard PIKAN achieves lower final loss or faster training. The advantage of GeoKAN is most pronounced in regimes with strong non-uniformity, stiffness, or high-frequency oscillations.
• Future Directions: The paper suggests future work could explore anisotropic (non-diagonal) metrics for higher-dimensional PDEs and further adaptivity in basis selection.

In summary, GeoKAN extends the KAN paradigm by learning the geometry of the input domain, offering a robust approach to handling the non-uniformity inherent in many scientific and differential equation problems.