A Graph Meta-Network for Learning on Kolmogorov-Arnold Networks

TL;DR

This paper introduces WS-KAN, a novel weight-space architecture tailored for Kolmogorov-Arnold Networks, leveraging their permutation symmetries, and demonstrates its superior performance across diverse tasks.

Contribution

The paper develops WS-KAN, the first weight-space architecture specifically designed for KANs, utilizing their permutation symmetries and graph representation for improved learning.

Findings

WS-KAN outperforms baseline models on multiple tasks.

KAN-graph effectively captures KAN computation structure.

WS-KAN demonstrates strong expressive power.

Abstract

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.