PaperlayerView exact revision on arXiv

Source-linked AI summary

Kolmogorov-Arnold Transformer

Xingyi Yang, Xinchao Wang

arXiv:2409.10594v1cs.LGcs.AIcs.CVcs.NE

TL;DR

Transformers commonly use MLP layers, while prior efforts to enhance them have been relatively few and scaling KAN layers into transformers presents challenges. The paper introduces KAT, integrating KANs with Group KAN and variance-preserving initialization, and reports stronger vision-task performance than MLP-based transformers at comparable computational requirements.

  • Problem

    Transformers rely on MLPs, which have received relatively few enhancement efforts, while integrating KAN layers into transformers remains a challenge.

    Source: §1 · p. 2

  • Method

    KAT replaces transformer MLP layers with KANs, combining Group KAN sharing and variance-preserving initialization in its design.

    Source: §1 · pp. 2–3

  • Results

    Across vision tasks, KAT outperforms traditional MLP-based transformers with comparable computational demands and achieves superior ImageNet-1K accuracy.

    Source: §1 · p. 3, §6 · p. 13

  • Takeaways & Limitations

    KAT can load pre-trained ViT weights and continue training to achieve further performance improvement.

    Source: §1 · p. 3

    KAT is evaluated on vision tasks, while extending it to other domains and addressing memory footprint and inference speed remain future challenges.

    Source: §6 · p. 13

Abstract

from arXiv · show

Transformers stand as the cornerstone of mordern deep learning. Traditionally, these models rely on multi-layer perceptron (MLP) layers to mix the information between channels. In this paper, we introduce the Kolmogorov-Arnold Transformer (KAT), a novel architecture that replaces MLP layers with Kolmogorov-Arnold Network (KAN) layers to enhance the expressiveness and performance of the model. Integrating KANs into transformers, however, is no easy feat, especially when scaled up. Specifically, we identify three key challenges: (C1) Base function. The standard B-spline function used in KANs is not optimized for parallel computing on modern hardware, resulting in slower inference speeds. (C2) Parameter and Computation Inefficiency. KAN requires a unique function for each input-output pair, making the computation extremely large. (C3) Weight initialization. The initialization of weights in KANs is particularly challenging due to their learnable activation functions, which are critical for achieving convergence in deep neural networks. To overcome the aforementioned challenges, we propose three key solutions: (S1) Rational basis. We replace B-spline functions with rational functions to improve compatibility with modern GPUs. By implementing this in CUDA, we achieve faster computations. (S2) Group KAN. We share the activation weights through a group of neurons, to reduce the computational load without sacrificing performance. (S3) Variance-preserving initialization. We carefully initialize the activation weights to make sure that the activation variance is maintained across layers. With these designs, KAT scales effectively and readily outperforms traditional MLP-based transformers.

1 Introduction

The paper motivates replacing transformer MLPs with KANs, while identifying scalability and training challenges that have limited their use. It proposes KAT with rational activations, Group KAN, and variance-preserving initialization, and reports stronger vision results at comparable computational requirements.

Motivation

  • Prior efforts to replace MLPs have been limited, and earlier KAN vision models fell short of traditional architectures amid compute and scalability problems.A direct ViT-to-KAN substitution also performed poorly at ImageNet scale.Source: §1 · p. 2
  • KANs offer learnable base functions, but integrating them into transformers has seen limited success despite their potential for expressiveness and efficiency.Earlier work reported KANs as 10× slower than MLPs at the same parameter count.Source: §1 · p. 2

Challenges

  • Scaling KANs poses challenges in GPU suitability of B-splines, per-edge parameter and computation costs, and initialization for convergence.The paper links initialization mismatch to training instability and degraded performance.Source: §1 · p. 2

Our Approach

  • KAT integrates KAN layers into transformers for large-scale training and introduces rational activations, Group KAN, and variance-preserving initialization to address these challenges.Rational activations have a CUDA implementation; Group KAN shares coefficients and base functions among edge groups.Source: §1 · p. 2
  • The combined design, Group-Rational KAN, replaces transformer MLPs and allows KAT to load ViT pre-trained weights for continued training.The paper reports that continued training from ViT weights further improves performance.Source: §1 · p. 3

Results

  • KAT outperforms MLP-based transformers across vision tasks under comparable computational requirements; on ImageNet-1K, KAT-B reaches 82.3% accuracy, 3.1% above same-size ViT.The paper also reports 82.7% performance when KAT is initialized with ViT pre-trained weights.Source: §1 · p. 3

2 Preliminary

The preliminaries introduce the Kolmogorov–Arnold representation theorem and describe KAN layers as learnable univariate functions on edges. A KAN network stacks these layers, with functions parameterized in practice by SiLU and B-spline components.

2.1 Kolmogorov-Arnold representation theorem

  • The theorem states that a multivariate continuous function on a bounded domain can be represented through continuous univariate functions and addition.For a d-variate function, (2d+1)(d+1) univariate functions suffice for an exact representation.Source: §2.1 · p. 3

2.2 Kolmogorov–Arnold Networks

  • KAN layers learn univariate functions on edges, and a KAN network stacks L such layers to map an input vector to an output.In practice, the function is parameterized as a linear combination of SiLU and B-spline functions.Source: §2.2 · pp. 3–4

3 Why original KAN fails to scale?

The section attributes vanilla KAN’s difficulty scaling to GPU-unfriendly B-splines, edge-specific parameter and computation costs, and initialization that does not preserve activation variance. Together these design choices make KAN resource-intensive for large models.

Base function

  • KAN’s B-splines are difficult to run efficiently on GPUs because localized, sparse or recursive computations are poorly suited to parallel processing.Pure PyTorch and NumPy implementations lack optimized CUDA support, and scaling efficient cubic implementations to higher spline orders is not straightforward.Source: §3 · p. 4

Parameter and computation inefficiency

  • Each KAN input-output channel pair has its own learnable base function, increasing parameters and computation as network width and depth grow.KAN parameter size and computational effort are respectively O(G+K) and O(GK) times those of a conventional MLP.Source: §3 · p. 4

Initialization

  • KAN’s default initialization does not preserve activation variance, so output variance differs from input variance under the zero-th-order spline analysis.The section estimates E[silu²(x)] ≈ 0.355σ² and notes that higher-order splines might increase variance instability.Source: §3 · pp. 4–5

4 Kolmogorov–Arnold Transformer

The section develops KAT by replacing ViT MLP layers with Group-Rational KAN and combining rational activations, grouped parameter sharing, and variance-preserving initialization. These designs target the challenges of scaling KANs for transformer use.

4 Kolmogorov–Arnold Transformer

  • The section adapts KAN for transformer use by refining its design to address challenges in scaling KAN-based models.The stated goal is to replace transformer MLP layers with KANs.Source: §4 · p. 5

4.1 Overall Architecture

  • KAT replaces ViT MLP layers with two-layer KANs while retaining the attention layers; its GR-KAN combines rational functions, grouped parameters, and stable-training initialization.The paper also states that KAT can use pre-trained ViT weights and continue training.Source: §4.1 · p. 5

4.2 Rational Base Functions

  • Horner-evaluated rational functions reduce FLOPs by approximately 9.3× versus B-spline functions.The paper uses rational functions as KAN base functions and implements them with CUDA.Source: Table 1 · p. 6, §4.2 · p. 6

4.3 Group KAN

  • Group KAN shares function parameters across groups, reducing the original KAN’s d_in×d_out unique activation functions to g and bringing parameter size to a constant overhead over an MLP.Each edge still retains a unique scalar weight in GR-KAN.Source: §4.3 · pp. 6–7

4.4 Variance-Preserving Initialization

  • Variance-preserving initialization aims to keep activation magnitudes stable across layers by fitting rational functions to established activations, estimating their expected squared output, and using the resulting gain to initialize weights.The expected squared output has no closed-form solution for the rational function, so it is calculated numerically.Source: §4.4 · p. 8

5 Experiments

KAT is evaluated on image recognition, object detection, and semantic segmentation, with reported gains over MLP-based transformer baselines. Ablations examine activation accuracy, computational trade-offs, CUDA efficiency, initialization, and learned rational functions.

5.2 Image Recognition

  • KAT-S reaches 81.2% accuracy on IN-1k, outperforming DeiT-S by 2.4%.Vanilla ViT-T/S + KAN reaches around 63%, while ViT-L + KAN fails to converge.Source: §5.2 · p. 9, Table 5 · p. 10

5.3 Object Detection and Instance Segmentation

  • KAT gains 3.0 APbox on the S-sized model and 1.4 APbox on the L-sized model over ViTDet, with only 1 GFLOPs more computational cost in smaller models.The gains are most pronounced in smaller models.Source: §5.3 · p. 10, Table 6 · p. 10

5.4 Semantic Segmentation

  • KAT improves over DeiT-S by 2.4% and over DeiT-B by 0.2% on ADE20K segmentation, with slightly higher FLOPs.Gains are larger in smaller models, but KAT still falls short of hierarchical architectures such as ConvNeXt.Source: §5.4 · p. 11, Table 7 · p. 11

5.5 Ablation Study and Analysis

  • With ViT-Ti/16, KAT-T achieves 74.6% top-1 accuracy on ImageNet-1k, outperforming GELU by 1.9%.The comparison tests alternative activations in the MLP against KAT.Source: §5.5 · p. 11, Table 8 · p. 11
  • On an A5000 GPU, KAT-T has slightly lower throughput than ReLU, GELU, and SiLU activations, while peak memory usage is similar.The authors suggest rational-function computation complexity may contribute to the throughput trade-off.Source: §5.5 · p. 11, Table 9 · p. 11
  • The CUDA implementation outperforms Torch Looped and Torch Vectorized in speed and memory efficiency across tested group counts and input dimensions.The comparisons vary group count from 1 to 16 and input dimension from 128 to 2048.Source: §5.5 · p. 12, Figure 5 · p. 12, Figure 6 · p. 12
  • Identity–Swish initialization achieves the best performance among tested rational-function initializations for KAT-T trained from scratch.The authors adopt it as the default setting.Source: §5.5 · p. 12, Table 10 · p. 12
  • Across KAT-S's 12 layers, trained rational functions show similar trends within each layer but differ across layers.The visualization uses group size 8.Source: §5.5 · p. 13

6 Conclusion and Future Work

KAT integrates GR-KAN into transformers and outperforms traditional MLP-based transformers on vision tasks while maintaining comparable computational demands. The paper identifies remaining speed and stability concerns and points to broader domains and deployment challenges for future work.

6 Conclusion and Future Work

  • KAT outperforms traditional MLP-based transformers across image recognition, object detection, and semantic segmentation, with superior ImageNet1K accuracy and comparable computational demands.The conclusion reports this cross-task pattern without supplying additional task-specific values.Source: §6 · p. 13
  • Rational functions remain slower than ReLU and GELU despite CUDA optimization, and higher-order coefficient gradients may become unstable with input powers.The authors say incorporating these functions into backpropagation could introduce complications.Source: §6 · p. 13
  • GR-KAN is a hybrid rather than a pure KAN: it shares edges and uses a rational base, and can also be interpreted as an MLP with a redesigned pre-linear activation.The authors describe it as leveraging MLP computational simplicity while retaining some KAN characteristics.Source: §6 · p. 13
  • Future work includes exploring Fourier, wavelet, and Gaussian radial bases, extending KAT beyond vision, and addressing memory footprint and inference speed for real-world deployment at scale.The paper names natural language processing and reinforcement learning as possible additional domains.Source: §6 · p. 13

7 Derivation and Calculation of FLOPs

The section counts operations for rational-function computation using a plain approach and Horner’s method. For the example with m=5 and n=4, Horner’s method uses fewer total operations.

7.1 Plain Computation

  • For m=5 and n=4, plain computation totals 46 operations: 34 multiplications, 10 summations, 1 absolute value, and 1 division.The numerator is a degree-m polynomial and the denominator uses the absolute value of a degree-n polynomial.Source: Denominator · p. 17, Numerator · p. 17

7.2 Horner’s Method

  • For a rational function computed with Horner’s method, the total is m+n+1 summations, m+n multiplications, 1 absolute value, and 1 division.The example with m=5 and n=4 gives 21 FLOPs.Source: §7.2 · p. 18
  • 21 FLOPs are required with Horner’s method for m=5 and n=4, versus 46 operations for plain computation.The Horner count comprises 9 multiplications, 10 summations, 1 absolute value, and 1 division; the plain count comprises 34 multiplications, 10 summations, 1 absolute value, and 1 division.Source: Denominator · p. 17, §7.2 · p. 18

8 Hyper-parameters for KAT model

The section points to the ImageNet-1k training hyperparameters in Table 11 and identifies the fitted rational functions for KAT-S.

8 Hyper-parameters for KAT model

  • The ImageNet-1k KAT training hyperparameters are presented in Table 11.The supplied passage identifies the table but does not provide its parameter values.Source: §8 · p. 18, Table 11 · p. 18
  • The fitted rational functions are shown for KAT-S with 12 layers and 8 groups.Source: Figure 7 · p. 19

Related papers

Loading 2409.10594v1…