TLDR: ShapKAN is a new framework for pruning Kolmogorov-Arnold Networks (KANs) that uses Shapley values to determine neuron importance. Unlike traditional methods, ShapKAN’s importance scores are consistent even when input data shifts, leading to more stable and effective network compression. This improves KANs’ interpretability and predictive performance, making them more practical for real-world applications by preserving true node contributions and enabling efficient model simplification.
Kolmogorov-Arnold Networks (KANs) represent a significant advancement in neural network design, offering a unique blend of high predictive accuracy and interpretability. Unlike traditional “black-box” neural networks, KANs use learnable spline-based activation functions on their edges, which allows them to recover symbolic representations of underlying functional relationships. This makes them particularly valuable in fields where understanding why a prediction is made is as important as the prediction itself, such as clinical research or scientific discovery.
However, the innovative architecture of KANs also introduces new challenges, especially when it comes to network pruning. Pruning is a technique used to simplify neural networks by removing less important components, making them more efficient and easier to interpret. Traditional pruning methods, which often rely on the magnitude of weights, prove unreliable for KANs because their importance scores can fluctuate significantly with simple shifts in input data. This sensitivity to input coordinate shifts leads to inconsistent pruning decisions, as illustrated in Figure 1 of the original research paper, where a “Vanilla KAN” yields different pruning solutions for shifted input ranges.
To address this critical limitation, researchers have developed a novel framework called ShapKAN. This framework leverages Shapley values, a concept from cooperative game theory, to assess the importance of each node (neuron) in a KAN. Shapley values provide a principled way to quantify each player’s (in this case, each neuron’s) actual contribution to the overall outcome, ensuring that importance rankings remain consistent regardless of how input parameters are scaled or shifted. This “shift-invariant” property is a cornerstone of ShapKAN’s effectiveness.
ShapKAN formalizes neuron attribution as a cooperative game. In this game, neurons in a KAN layer are considered “players” that collaborate to contribute to the model’s prediction power. The Shapley value then fairly allocates this contribution among the neurons. This approach inherently satisfies several fairness axioms, such as ensuring that a neuron contributing nothing receives a zero score, and that the sum of all neuron contributions equals the total value of the layer.
Given the computational complexity of calculating exact Shapley values for large networks, ShapKAN employs efficient approximation methods, primarily permutation sampling. This technique estimates Shapley values by averaging marginal contributions across randomly sampled permutations of neurons. To further enhance efficiency and reduce variance, especially with smaller sample sizes, ShapKAN incorporates antithetic permutation sampling. This ensures statistically reliable importance estimation without imposing restrictive model assumptions.
For multi-layer KANs, ShapKAN adopts a bottom-up greedy pruning algorithm. It estimates Shapley values layer by layer, starting from the bottom, and sequentially prunes neurons with low importance. The framework offers flexible pruning criteria, allowing users to remove neurons based on a ratio of their contribution, a specific number, or a defined threshold. This systematic approach ensures that pruning decisions are well-informed and maintain the network’s integrity.
Extensive experiments on both synthetic and real-world datasets demonstrate ShapKAN’s superior performance. In simulation studies, ShapKAN consistently maintained stable neuron importance rankings and achieved lower test errors compared to the vanilla KAN pruning method, particularly under covariate shift. For instance, in symbolic regression tasks, ShapKAN more accurately recovered the ground-truth functional forms, highlighting its enhanced interpretability.
On real-world benchmarks like Census-income, MNIST, and Airbnb datasets, ShapKAN consistently outperformed both the vanilla KAN and DropKAN (another KAN pruning method) in terms of generalization capacity. While traditional Multi-Layer Perceptrons (MLPs) showed strong performance, ShapKAN achieved competitive accuracy while retaining its crucial advantage of interpretability through symbolic function recovery. This makes ShapKAN a powerful tool for deploying KANs in resource-constrained environments where both performance and understanding are vital.
Also Read:
- PrunedLoRA: Enhancing Low-Rank Adaptation for Large Language Models Through Dynamic Pruning
- Optimizing AI Explanations: A New Approach to Understanding Model Decisions
The development of ShapKAN significantly improves the interpretability advantages of KANs by providing a robust and principled method for network compression. Its shift-invariant attribute scoring enhances confidence in KAN model outcomes for both predictive and non-predictive tasks, paving the way for broader adoption in fields like AI+Science, AI+Health, and AI+Finance. Future work aims to extend ShapKAN to account for higher-order Shapley interactions, further refining the quantification of neuron contributions. For more details, you can read the full research paper here.
