InfinityKAN: Adaptive Variational KAN Architecture
- InfinityKAN is a variational machine learning architecture that extends Kolmogorov–Arnold Networks by adaptively learning an effectively infinite number of basis functions per univariate transformation.
- The method replaces fixed basis count hyperparameters with learnable variables using variational inference, a smooth masking window, and a Poisson prior to dynamically regulate network capacity.
- Empirical evaluations demonstrate competitive performance on regression, image, and graph classification tasks, showcasing robust adaptability and reduced overfitting risk.
InfinityKAN is a variational machine learning architecture that extends Kolmogorov–Arnold Networks (KANs) by adaptively learning an effectively infinite number of basis functions per univariate transformation during training. This adaptive mechanism replaces fixed architectural hyperparameters with variationally optimized variables, rendering the number of active bases (and thus network capacity) a learnable quantity. The model is grounded in the Kolmogorov–Arnold Theorem but overcomes the practical limitations of fixed-basis designs by leveraging variational inference and a novel smooth basis-selection mechanism (Alesiani et al., 3 Jul 2025).
1. Foundations: KANs and the Kolmogorov–Arnold Theorem
The Kolmogorov–Arnold representation states that any continuous bounded function can be represented exactly as
for suitable continuous univariate functions . KANs instantiate this result by parameterizing each inner function as a linear combination of preselected basis functions , i.e.,
with coefficients trained via backpropagation. The value of critically affects expressivity and overfitting trade-offs but must be hand-tuned for each task (Alesiani et al., 3 Jul 2025).
2. Variational Formulation and Inference in InfinityKAN
InfinityKAN tackles the architectural constraint of fixed by casting both weights and basis counts as random variables within a joint probabilistic model. For a dataset 0, the latent variables are 1 (basis coefficients, layer-indexed) and 2 (the canonical rates controlling the number of bases per layer).
The joint distribution is formulated as: 3 with
- 4 (Gaussian prior for weights),
- 5 (Poisson prior for basis count),
- 6 being the standard Gaussian (regression) or softmax (classification) likelihood where the KAN’s active basis count is 7 per layer.
Variational inference is used for posterior approximation: 8 where 9 and 0 for active coefficients. The training objective maximizes the ELBO: 1 Optimization is performed by SGD on variational means 2 (Alesiani et al., 3 Jul 2025).
3. InfinityKAN Architecture and Basis Adaptivity
Each layer in InfinityKAN implements a generalized KAN, but with each univariate function parameterized as a potentially infinite linear combination: 3 For practical training and inference, the series is truncated adaptively using a learnable 4 for each layer, with an additional smooth “window” function: 5 This window enables differentiable masking of inactive bases and supports backpropagation-friendly optimization over 6.
The 7th layer output is computed as: 8 where 9 can be chosen from ReLU-steps, Chebyshev polynomials, or Fourier modes. This generalizes classical KANs without prescribing basis cardinality a priori.
4. Training Algorithm and Optimization Techniques
Training of InfinityKAN proceeds via stochastic gradient descent applied to the ELBO, using the following procedure, as excerpted (Algorithm 1) (Alesiani et al., 3 Jul 2025): 9 Special techniques include:
- A soft window to smooth the effective change in basis count.
- Mean-field variational approximations.
- First-order approximations to the variational expectation. The discrete 0 variable is handled via reparameterization and optimization over its mean parameter 1 (Alesiani et al., 3 Jul 2025).
5. Empirical Performance and Robustness Analyses
InfinityKAN was evaluated on a range of tasks, including synthetic regression, image classification (MNIST, Fashion-MNIST, CIFAR-10/100, EuroSAT), and graph classification (MUTAG, NCI1, PROTEINS, ENZYMES, REDDIT-BINARY). The variant Infinity-GKAN extends these ideas to graph neural architectures. Empirical results establish:
- On regression and classification benchmarks, InfinityKAN matches or improves upon best-tuned fixed-2 KANs and MLPs, adaptively learning basis numbers 3–20 per layer.
- On graph classification, Infinity-GKAN outperforms both fixed-basis KANs and conventional MLP graph encoders.
- Ablation studies across the window function parameters 4, basis choice (ReLU, Chebyshev, Fourier), variational priors, and width/depth show robust, task-adaptive selection of 5 and basis type (Alesiani et al., 3 Jul 2025).
Empirical visualizations (see Figure 1 in the source) corroborate smooth, monotonic evolution and stabilization of the learned 6 during training. The ELBO remains Lipschitz-continuous in 7, indicating stable optimization.
6. Significance and Theoretical Properties
InfinityKAN eliminates the need for manual selection of basis set cardinality in KANs, preserving the full expressive power of the Kolmogorov–Arnold Theorem in a learnable, scalable framework. The architecture is fully compatible with modern backpropagation and variational inference pipelines and is practically competitive with standard MLPs and fixed-KAN architectures across diverse tasks. The use of a Poisson prior and adaptive windows for basis selection allows the architecture to allocate model capacity dynamically, reducing overfitting risk and computation associated with unnecessarily large basis sets.
A plausible implication is that this framework generalizes to any machine learning scenario in which function approximation benefits from infinite or adaptive expansion sets, provided a suitable variational or amortized optimization method can be specified (Alesiani et al., 3 Jul 2025).
7. Limitations and Open Directions
Despite empirical success and theoretical soundness, several open questions remain regarding optimal prior selection for 8, the theoretical limits of adaptively truncated univariate basis expansions, and the transferability of the approach to broader network architectures. Further, the extension of InfinityKAN principles to object types other than scalar function approximation, or to structured data beyond graphs, remains an open area for investigation. Robustness against adversarial basis selection and convergence guarantees under non-Gaussian priors also warrant further study (Alesiani et al., 3 Jul 2025).