Multi-Exit Kolmogorov-Arnold Networks: enhancing accuracy and parsimony

2025-06-03

AI Summary

Kolmogorov–Arnold Networks (KANs) suffer from manual depth selection, optimization difficulties with excessive depth, and lack of adaptive model complexity control. Method: We propose Multi-Exit KAN (ME-KAN), embedding learnable, layer-wise independent prediction heads coupled with a differentiable “learned exit” gating mechanism to enable simultaneous multi-depth supervised training and automatic optimal early-exit decision-making. Contribution/Results: ME-KAN is the first framework to integrate the Kolmogorov–Arnold representation theorem with multi-exit architecture and differentiable exit policies, enabling dynamic identification of the minimal yet effective submodel without sacrificing accuracy. Experiments on synthetic function approximation, dynamical system modeling, and real-world benchmarks demonstrate consistent superiority over standard single-exit KANs. Early exits frequently achieve peak performance, yielding substantial gains in parameter efficiency—up to 62% fewer parameters—and improved generalization accuracy.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Learning to Search
• Multiagent Systems: Multiagent Learning

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

Kolmogorov-Arnold Networks (KANs) uniquely combine high accuracy with interpretability, making them valuable for scientific modeling. However, it is unclear a priori how deep a network needs to be for any given task, and deeper KANs can be difficult to optimize. Here we introduce multi-exit KANs, where each layer includes its own prediction branch, enabling the network to make accurate predictions at multiple depths simultaneously. This architecture provides deep supervision that improves training while discovering the right level of model complexity for each task. Multi-exit KANs consistently outperform standard, single-exit versions on synthetic functions, dynamical systems, and real-world datasets. Remarkably, the best predictions often come from earlier, simpler exits, revealing that these networks naturally identify smaller, more parsimonious and interpretable models without sacrificing accuracy. To automate this discovery, we develop a differentiable"learning to exit"algorithm that balances contributions from exits during training. Our approach offers scientists a practical way to achieve both high performance and interpretability, addressing a fundamental challenge in machine learning for scientific discovery.

Problem

• Determining optimal network depth for task accuracy
• Improving training efficiency in deep KANs
• Balancing model complexity and interpretability

Innovation

• Multi-exit KANs enable predictions at multiple depths
• Deep supervision improves training and model complexity
• Differentiable learning to exit algorithm balances contributions