Search papers, labs, and topics across Lattice.
This paper introduces WS-KAN, a novel weight-space architecture tailored for Kolmogorov-Arnold Networks (KANs) that leverages the permutation symmetries inherent in KANs. The authors define the KAN-graph, a graph representation of KAN computation, and use it to build WS-KAN, which accounts for the identified symmetries. Empirical evaluation on a diverse set of KANs demonstrates that WS-KAN significantly outperforms structure-agnostic baselines in predicting KAN performance across various tasks.
Unlock the potential of Kolmogorov-Arnold Networks with WS-KAN, a weight-space architecture that understands their hidden symmetries and predicts their performance far better than generic methods.
Weight-space models learn directly from the parameters of neural networks, enabling tasks such as predicting their accuracy on new datasets. Naive methods -- like applying MLPs to flattened parameters -- perform poorly, making the design of better weight-space architectures a central challenge. While prior work leveraged permutation symmetries in standard networks to guide such designs, no analogous analysis or tailored architecture yet exists for Kolmogorov-Arnold Networks (KANs). In this work, we show that KANs share the same permutation symmetries as MLPs, and propose the KAN-graph, a graph representation of their computation. Building on this, we develop WS-KAN, the first weight-space architecture that learns on KANs, which naturally accounts for their symmetry. We analyze WS-KAN's expressive power, showing it can replicate an input KAN's forward pass - a standard approach for assessing expressiveness in weight-space architectures. We construct a comprehensive ``zoo'' of trained KANs spanning diverse tasks, which we use as benchmarks to empirically evaluate WS-KAN. Across all tasks, WS-KAN consistently outperforms structure-agnostic baselines, often by a substantial margin. Our code is available at https://github.com/BarSGuy/KAN-Graph-Metanetwork.
NVIDIA