Kolmogorov-Arnold Attention: Is Learnable Attention Better For Vision
  Transformers?
Authors: name: Subhajit Maity, name: Killian Hitsman, name: Xin Li, name: Aritra Dutta

Kolmogorov-Arnold networks (KANs) are a remarkable innovation consisting of
learnable activation functions with the potential to capture more complex
relationships from data. Although KANs are useful in finding symbolic
representations and continual learning of one-dimensional functions, their
effectiveness in diverse machine learning (ML) tasks, such as vision, remains
questionable. Presently, KANs are deployed by replacing multilayer perceptrons
(MLPs) in deep network architectures, including advanced architectures such as
vision Transformers (ViTs). In this paper, we are the first to design a general
learnable Kolmogorov-Arnold Attention (KArAt) for vanilla ViTs that can operate
on any choice of basis. However, the computing and memory costs of training
them motivated us to propose a more modular version, and we designed particular
learnable attention, called Fourier-KArAt. Fourier-KArAt and its variants
either outperform their ViT counterparts or show comparable performance on
CIFAR-10, CIFAR-100, and ImageNet-1K datasets. We dissect these architectures'
performance and generalization capacity by analyzing their loss landscapes,
weight distributions, optimizer path, attention visualization, and spectral
behavior, and contrast them with vanilla ViTs. The goal of this paper is not to
produce parameter- and compute-efficient attention, but to encourage the
community to explore KANs in conjunction with more advanced architectures that
require a careful understanding of learnable activations. Our open-source code
and implementation details are available on: https://subhajitmaity.me/KArAt