Low Tensor-Rank Adaptation of Kolmogorov--Arnold Networks

TL;DR

This paper introduces LoTRA, a low tensor-rank adaptation method for fine-tuning Kolmogorov--Arnold networks, improving transfer learning efficiency and model size reduction for scientific and image tasks.

Contribution

The paper develops LoTRA based on Tucker decomposition, providing theoretical analysis for learning rate selection and demonstrating its effectiveness in transfer learning and model compression.

Findings

LoTRA improves transfer learning efficiency for KANs.

Using adaptive learning rates enhances training performance.

Slim KANs reduce model size while maintaining accuracy.

Abstract

Kolmogorov--Arnold networks (KANs) have demonstrated their potential as an alternative to multi-layer perceptions (MLPs) in various domains, especially for science-related tasks. However, transfer learning of KANs remains a relatively unexplored area. In this paper, inspired by Tucker decomposition of tensors and evidence on the low tensor-rank structure in KAN parameter updates, we develop low tensor-rank adaptation (LoTRA) for fine-tuning KANs. We study the expressiveness of LoTRA based on Tucker decomposition approximations. Furthermore, we provide a theoretical analysis to select the learning rates for each LoTRA component to enable efficient training. Our analysis also shows that using identical learning rates across all components leads to inefficient training, highlighting the need for an adaptive learning rate strategy. Beyond theoretical insights, we explore the application of LoTRA for efficiently solving various partial differential equations (PDEs) by fine-tuning KANs. Additionally, we propose Slim KANs that incorporate the inherent low-tensor-rank properties of KAN parameter tensors to reduce model size while maintaining superior performance. Experimental results validate the efficacy of the proposed learning rate selection strategy and demonstrate the effectiveness of LoTRA for transfer learning of KANs in solving PDEs. Further evaluations on Slim KANs for function representation and image classification tasks highlight the expressiveness of LoTRA and the potential for parameter reduction through low tensor-rank decomposition.