An Innovative Networks in Federated Learning

TL;DR

The paper addresses federated learning under data heterogeneity by introducing Wavelet Kolmogorov-Arnold Networks (Wav-KAN) that fuse wavelet transforms with neural activation functions. Wav-KAN uses both $CWT$ and $DWT$ for multiresolution representations and learns wavelet-based activations parameterized by $(w_i,s_i, au_i)$ at each client, with local training and server-side aggregation. Global updates are obtained by averaging across clients, e.g., $\bar{w} = \frac{1}{N} \sum_{i=1}^N w_i$, $\bar{s} = \frac{1}{N} \sum_{i=1}^N s_i$, $\bar{\tau} = \frac{1}{N} \sum_{i=1}^N \tau_i$, and broadcast back to clients. Experiments on MNIST and additional datasets demonstrate improvements in training/test accuracy, interpretability, and computational efficiency, with wavelet choice strongly influencing performance.

Abstract

This paper presents the development and application of Wavelet Kolmogorov-Arnold Networks (Wav-KAN) in federated learning. We implemented Wav-KAN \cite{wav-kan} in the clients. Indeed, we have considered both continuous wavelet transform (CWT) and also discrete wavelet transform (DWT) to enable multiresolution capabaility which helps in heteregeneous data distribution across clients. Extensive experiments were conducted on different datasets, demonstrating Wav-KAN's superior performance in terms of interpretability, computational speed, training and test accuracy. Our federated learning algorithm integrates wavelet-based activation functions, parameterized by weight, scale, and translation, to enhance local and global model performance. Results show significant improvements in computational efficiency, robustness, and accuracy, highlighting the effectiveness of wavelet selection in scalable neural network design.

Figures (2)

• Figure 1: Training accuracy of Wav-KAN [28*28,64,10] with Mexican hat mother wavelet activation.
• Figure 2: Test accuracy of Wav-KAN [28*28,64,10] with Mexican hat mother wavelet activation.