Kolmogorov-Arnold Transformer

## Abstract

Abstract 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. We demonstrate the advantages of KAT across various tasks, including image recognition, object detection, and semantic segmentation. It consistently enhances performance over the standard transformer architectures of different model sizes. Our code is openly available at https://github.com/Adamdad/kat .

## 1 Introduction

Transformers have become the de facto architecture in deep learning, widely adopted in computer vision [] and natural language processing []. At their core, transformers are built upon two fundamental components: attention modules and multi-layer perceptrons (MLPs). Although significant research has focused on replacing the traditional attention mechanism with alternative operations [, , ], these variants still lean heavily on MLPs. Surprisingly, there have been relatively few efforts [] aimed at enhancing MLPs themselves.

Opening up the box, MLPs are composed of stacked linear layers coupled with non-linear activations. What makes it so popular is that, theoretically, they can approximate any function, assuming that there are enough neurons available [].

However, despite their versatility, MLPs face limitations in modeling complex functions. For example, when using ReLU-like activation, a two-layer MLP may struggle to fit periodic functions. Moreover, employing gradient descent to train these networks often results in prolonged convergence times for high-frequency components [, , ]. These challenges have led researchers to explore alternative, perhaps more expressive architectures than MLPs.

Recently, Kolmogorov-Arnold Networks (KANs) emerged as a powerful alternative. KANs are noted for their theoretical parameter efficiency, potentially requiring fewer parameters to model complex functions []. They are particularly suitable for mathematical or symbolic regression tasks [, , ].
The key to such success is the learnable base function in each input-output pair. Those functions are often parameterized by B-spline curves [, ]. This design allows KANs to approximate more intricate functions through a summation of spline bases.

Given its potential, integrating KAN layers into transformers [] becomes an exciting topic. Such integration may boost the expressiveness and efficiency of transformers, enhancing their competitiveness across a wide range of applications.

Unfortunately, this ambition has been met with limited success. In particular, KANs have been reported to be “10$\times$ slower than MLPs, given the same number of parameters”. Initial attempts to apply KANs to vision recognition tasks have yielded disappointing results. Even on a small scale, these studies have consistently fallen short of matching, let alone surpassing, the performance of traditional architectures. This lack of improvement is often attributed to the limited computational resources and ongoing scalability problems [, , , ].

In a preliminary experiment, we attempted to replace MLP layers in the Vision Transformer (ViT) with KAN layers. It creates a model, which we call ViT+KAN. However, as shown in Figure 1 (Right), this straightforward substitution led to significant challenges when performing ImageNet-scale training, resulting in poor performance.
Scalability, therefore, remains a significant obstacle for KAN-based models.

Motivation and Challenges. Through dedicated analysis, we have identified several key challenges that hinder the effectiveness of KANs in large-scale applications, ultimately limiting their scalability.

- •
(C1) *Base function.* The standard B-spline functions in KANs are not ideal for parallel computing architectures typical of modern GPUs. B-splines require recursive computation, which significantly slows down even the most optimized implementations.
- •
(C2) *Parameter and Computation Inefficiency.* Each unique input-output pair in a KAN requires a distinct set of parameters and base functions. This necessity causes an exponential growth in the number of parameters as the network’s hidden size increases, resulting in substantial computational overhead and scalability issues.
- •
(C3) *Weight initialization.* The weight initialization in KANs is similar to that in MLPs, but it does not meet KANs’ needs for convergence. This mismatch can lead to instability and degraded performance during the training process.

Our Approach. In this paper, we introduce *Kolmogorov–Arnold Transformer* (KAT), which successfully integrates KANs into transformers for large-scale training scenarios such as ImageNet. Beyond simple replacement, We have developed three key innovations (S1-S3) to address these challenges (C1-C3) respectively.

- •
(S1) *Rational activation.* We employ rational function as our base function and provide full CUDA implementation. It aligns better with modern GPU architectures, enhancing computational efficiency and compatibility.
- •
(S2) *Group KAN.* We share function coefficients and base functions among groups of edges. This strategy reduces computational load significantly without sacrificing performance.
- •
(S3) *Variance-preserving initialization*. We carefully initialize weights to maintain consistent variance in activations across the model’s layers. This ensures stability during training and improves the model’s learning dynamics.

By combining all solutions S1-S3, we present a new variant of KAN, called *Group-Rational KAN* (GR-KAN), to replace the MLP in transformer. We show that GR-KAN is computationally efficient, easy to implement, and can be seamlessly integrated into vision transformers, replacing MLP layers to achieve superior performance. Furthermore, our designs allow KAT to load pre-trained weights from ViT models and continue training to achieve even better results.

We empirically validate KAT across a range of vision tasks, including image recognition, object detection, and semantic segmentation.
The results demonstrate that KAT outperforms traditional MLP-based transformers, achieving enhanced performance with comparable computational requirements. As illustrated in Figure 1, KAT-B achieves 82.3% accuracy on ImageNet-1K, surpassing the ViT model of the same size by 3.1%. When initialized with pre-trained weights from ViT, the performance further improves to 82.7%.

The contributions of our paper are threefold. First, we conduct a thorough analysis of the challenges in scaling KAN-based models, particularly focusing on inefficiencies in base functions, parameterization, and weight initialization. Based on this analysis, we propose a set of solutions: rational activation functions tailored for GPU efficiency, Group KAN to reduce computational overhead, and variance-preserving initialization to ensure stable training. Second, leveraging these insights, we introduce the Kolmogorov–Arnold Transformer (KAT) and scale it to ImageNet-level training, successfully integrating KANs into large-scale models. Third, we validate our approach through extensive experiments, showing that KAT not only matches but surpasses the performance of ViT models, all under similar computational requirements.

## 2 Preliminary

### 2.1 Kolmogorov-Arnold representation theorem

The Kolmogorov-Arnold representation theorem [] states that any multivariate continuous function $f$, defined on a bounded domain, can be expressed as a finite composition of continuous univariate functions and addition. Specifically, for a smooth function $f:[0,1]^{n}\to\mathbb{R}$, it can be represented as:

$$ $f(x_{1},\ldots,x_{n})=\sum_{q=1}^{2n+1}\Phi_{q}\left(\sum_{p=1}^{n}\phi_{q,p}( x_{p})\right)$ $$

Here, each function $\phi_{q,p}:[0,1]\to\mathbb{R}$ and $\Phi_{q}:\mathbb{R}\to\mathbb{R}$ are continuous. This means that the (2d+1)(d+1) univariate functions $\Phi_{q}$ and $\phi_{q,p}$ are enough for an exact representation of a d-variate function.

This theorm can be written in matrix form as follows:

$$ $f(\mathbf{x})=\Phi_{\text{out}}\circ\Phi_{\text{in}}\circ\mathbf{x}$ (1) $$

where $\Phi_{\text{in}}$ and $\Phi_{\text{out}}$ are defined as:

$$ $\Phi_{\text{in}}=\begin{bmatrix}\phi_{1,1}(\cdot)&\cdots&\phi_{1,n}(\cdot)\\ \vdots&\ddots&\vdots\\ \phi_{2d+1,1}(\cdot)&\cdots&\phi_{2d+1,d}(\cdot)\end{bmatrix}$ (2) $$

$$ $\Phi_{\text{out}}=\begin{bmatrix}\Phi_{1}(\cdot)&\cdots&\Phi_{2d+1}(\cdot)\end {bmatrix}$ (3) $$

This decomposition illustrates how $f$ can be built from simpler functions, showcasing an essential property of multivariate continuous functions.

### 2.2 Kolmogorov–Arnold Networks

Inspired by the Kolmogorov-Arnold representation theorem, [] define a generalized Kolmogorov-Arnold layer to learn univariate functions on edge, in the form of activation function. Formally, a Kolmogorov-Arnold layer with $d_{\text{in}}$-dimensional inputs and $d_{\text{out}}$-dimensional outputs is illustrated as

$$ $f(\mathbf{x})=\Phi\circ\mathbf{x}=\begin{bmatrix}\sum_{i=1}^{d_{in}}\phi_{i,1} (x_{i})&\dots&\sum_{i=1}^{d_{in}}\phi_{i,d_{out}}(x_{i})\end{bmatrix},\text{ where}\quad\Phi=\begin{bmatrix}\phi_{1,1}(\cdot)&\cdots&\phi_{1,d_{\text{in}}} (\cdot)\\ \vdots&\ddots&\vdots\\ \phi_{d_{\text{out}},1}(\cdot)&\cdots&\phi_{d_{\text{out}},d_{\text{in}}}( \cdot)\end{bmatrix}$ (4) $$

Note that Eq 4 can be seen as a generalized form of Eq 1, such that $\Phi=\Phi_{\text{in}}\circ\Phi_{\text{out}}$. A general KAN network is a stacking of
$L$ layers: given an input vector $\mathbf{x}_{0}\in\mathbb{R}^{d_{0}}$, the output of KAN is $KAN(\mathbf{x}_{0})=\Phi_{L-1}\circ\Phi_{L-2}\dots\circ\Phi_{0}\circ\mathbf{x}
_{0}$.

In practice, [] parameterizes $\Phi$ use a linear combination of SiLU activation [] and a B-spline function

$$ $\displaystyle\phi(x)=w_{b}\texttt{silu}(x)+w_{s}\texttt{spline}(x),$ $\displaystyle\text{where}\quad\texttt{silu}(x)=\frac{x}{1+e^{-x}},$ $\displaystyle\texttt{spline}(x)=\sum_{i}c_{i}B_{i}(x)$ (5) $$

## 3 Why original KAN fails to scale?

This section examines the scalability issues of KAN. We will explore three key factors: the choice of base function, redundant parameters and computation, and initialization problems. These design choices make the vanilla version of KAN resource-intensive and difficult to apply to large-scale models.

##### B-spline is not GPU Friendly.

The use of B-spline functions in KAN layers introduces challenges when implemented on GPUs. First, B-splines are not standard functions within CUDA. Implementing them using pure PyTorch and NumPy results in slower performance on modern GPU devices due to the lack of optimized CUDA support. Second, the localized nature of B-Spline computations complicates their use in parallel GPU processes. Typically, each control point influences only a small adjacent area of the curve. This leads to sparse or recursive computations, a type of operation that GPUs manage less efficiently. Although there are efficient implementation for cubic B-Spline [, , ], scaling these methods to higher orders is not straightforward.

##### Parameter and Computation Inefficiency.

Unlike standard neural networks, KAN employs a learnable base function for each pair of input-output channels. This design inherently leads to an increased parameter count and higher computational demands, especially when scaling up the width and depth of a neural network.

In the standard configuration of KAN, a layer with $d_{i}n$ input and $d_{o}ut$ output channels incorporates an B-spline function for each input-output pair, of order $K$ on $G$ intervals. This results in the network having a total of $(d_{in}\times d_{out})\times(G+K+3)+d_{out}$ learnable parameters. In contrast, a typical MLP only needs $(d_{in}\times d_{out})+d_{out}$ parameters.

In terms of computation, the FLOPs for one sample(^1^11For full computation derivation, please see [].) in B-spline with De Boor-Cox Iterative [] formulation is $\Big{\{}\text{FLOPs of non-linear function}\times d_{in}+(d_{in}\times d_{out}
)\times[9K\times(G+1.5K)+2G-2.5K+3]\Big{\}}$. Meanwhile, the FLOPs for an equivalent MLP layer is merely $\Big{\{}\text{FLOPs of non-linear function}\times d_{out}+2\times(d_{in}\times
d
_{out})\Big{\}}$.

Overall, the parameter size and computational effort of KAN are on the order of $O(G+K)$ and $O(GK)$ times greater than those of a conventional MLP, respectively. This significant increase in complexity is a primary reason why KAN struggles to scale effectively.

Weights are not Properly Initialized. Deep learning heavily relies on good weight initialization to enable trainability and convergence. A fundamental principle is to ensure *variance-preserving*, meaning that the variance of the signal should remain constant as it propagates through multiple layers, whether forward or backward [, , ]. This principle ensures that the activation and gradient maintain stability across layers.

However, in the KAN paper, the initialization strategy deviates from this principle. Specifically, the B-spline coefficients $c_{i}$ are initialized as $\mathcal{N}(0,\sigma^{2})$ with $\sigma=0.1$, and $w_{s}=1$ and $w_{b}\sim U[-\frac{6}{\sqrt{d_{in}+d_{out}}},\frac{6}{\sqrt{d_{in}+d_{out}}}]$ are initialized according to the Xavier initialization []. The combined output variance of the model can be expressed as:

$$ $Var[\phi(x)]=Var[w_{b}\texttt{silu}(x)]+Var[w_{s}\texttt{spline}(x)]=3\mathbb{ E}[\texttt{silu}^{2}(x)]+\mathbb{E}[\texttt{spline}^{2}(x)]$ (6) $$

If we assume the input
$x$ is normally distributed, $x\sim\mathcal{N}(0,\sigma_{x}^{2})$ and consider a zero-th order spline, the variance of $\texttt{spline}(x)$ at any point $x$ is simply:

$$ $\mathbb{E}[\texttt{spline}^{2}(x)]=\sum_{i}c_{i}^{2}Var[B_{i}(x)]=\sigma^{2} \sum_{i}Var[B_{i}(x)]=\sigma^{2}=0.01$ (7) $$

For the SiLU activation function, although exact variance calculations are complex, numerical estimations indicates $\mathbb{E}[\texttt{silu}^{2}(x)]\approx 0.355\sigma_{x}^{2}$. Combining these, we find $Var[\phi(x)]\approx 0.01+1.064\sigma_{x}^{2}\neq Var[x]$.

This indicates that, under zero-th order spline, $Var[\phi(x)]\neq Var[x]$. With higher-order splines, the variance instability might increase. Thus, the default initialization opposes the essential variance-preserving principle.

## 4 Kolmogorov–Arnold Transformer

As discussed earlier, the standard KAN faces three major challenges that limit its use in large, deep neural networks. In this section, we refine its design to better suit modern transformers, allowing us to replace MLP layers with KANs.

### 4.1 Overall Architecture

Just as its name imply, Kolmogorov–Arnold Transformer (KAT) replaces the MLPs in vision transformer [] with KAN layers.

Specifically, for a 2D image $\mathbf{x}\in\mathbb{R}^{H\times W\times C}$, we first flatten it into a 1D sequence, apply patch embedding and positional encoding, and then pass it through a series of KAT layers. At layer $\ell$, the following operations are performed:

$$ $\mathbf{x}_{0}^{(\ell)}=\text{MSA}(\text{LN}(x_{\ell-1}))+\mathbf{x}_{\ell-1}, \quad\ell=1,\dots,L$ (8) $$

$$ $\mathbf{x}_{\ell}=\text{MLP}(\text{LN}(\mathbf{x}_{0}^{(\ell)}))+\mathbf{x}_{0 }^{(\ell)},\quad[\textbf{Transformer}]$ (9) $$

$$ $\mathbf{x}_{\ell}=\text{KAN}(\text{LN}(\mathbf{x}_{0}^{(\ell)}))+\mathbf{x}_{0 }^{(\ell)},\quad[\textbf{KAT}]$ (10) $$

where $\mathbf{x}_{\ell}$ stands for the output feature sequence at the $\ell$ layer.
As illustrated, we replace all two-layer MLPs with two-layer KANs while keeping the attention layers unchanged. Although similar efforts have been made in specific domains [, ], a simple replacement is not enough to achieve scalability in large models.

Most importantly, here, we introduce a special kind *Group-Rational KAN*. We use rational functions as the base function for KAN (Section 4.2) and share parameters between a group of edges (Section 4.3). We also specify the weight initialization scheme to ensure stable training (Section 4.4). Together, these enhancements make KAT more scalable and improve performance.

### 4.2 Rational Base Functions

In our method, we use the rational function [, , , ] as the base function for the KAN layer, instead of the B-spline.

Specifically, we parameterize the function $\phi(x)$ on each edge as rational over polynomials $P(x),Q(x)$ of order $m,n$.

$$ $\phi(x)=wF(x)=w\frac{P(x)}{Q(x)}=w\frac{a_{0}+a_{1}x+\dots+a_{m}x^{m}}{b_{0}+b _{1}x+\dots+b_{n}x^{n}}$ (11) $$

$a_{n}$ and $b_{m}$ are coefficient of the rational function and $w$ is the scaling factor. This function is said to have degree $m/n$. We hope to learn those $a_{n},b_{m}$ and $w$ through end-to-end backpropagation.

To avoid instability caused by poles, where $Q(x)\to 0$ and $\phi(x)\to\pm\infty$, we employ a Safe Padé Activation Unit (PAU) [] as our basis, which is a modified form of the standard rational function

$$ $F(x)=\frac{a_{0}+a_{1}x+\dots+a_{m}x^{m}}{1+|b_{1}x+\dots+b_{n}x^{n}|}$ (12) $$

Why use Rational Function? There are practical and theoretical reasons for selecting rational functions as our base functions.

First, from an efficiency perspective, evaluating polynomials involves simple operations that are highly suitable for parallel computing. This makes rational functions computationally efficient for large-scale models.

Second, from a theoretical perspective, rational functions can approximate a wider range of functions—including those with singularities or sharp variations—more efficiently and accurately than polynomials [, ]. Since B-splines are essentially sums of local polynomials, rational functions offer a theoretical advantage over B-splines for modeling complex behaviors.

Third, from a practical perspective, rational activations have already been successfully used as activation functions in neural networks [, ].

Given these reasons, we adopt rational functions as the base functions in our KAN layers to enhance the model’s expressiveness, stability, and computational efficiency.

Implement Rational Function on GPU. With the rational function, a core contribution in this paper is to implement it efficiently on parralized devices like GPU. In stead of using pytorch with automatic differentiation, we implement it fully with CUDA [].

- •
Similar to [], we compute the explicit gradients of $\frac{\delta F}{\delta a_{m}},\frac{\delta F}{\delta b_{n}}$ and $\frac{\delta F}{\delta x}$
$\frac{\delta F}{\delta a_{m}}=\frac{x^{m}}{Q(x)},\quad\frac{\delta F}{\delta b
_{n}}=ax^{n}\frac{A(x)}{|A(x)|}\frac{P(x)}{Q(x)^{2}},\quad\text{and}\quad\frac
{\delta F}{\delta x}=\frac{\delta P(x)}{\delta x}\frac{1}{Q(x)}-\frac{\delta Q
(x)}{\delta x}\frac{P(x)}{Q^{2}(x)}$
(13)
where $A(x)=b_{1}x+\dots+b_{n}x^{n}$, $\frac{\delta P(x)}{\delta x}=a_{1}+2a_{2}x+ma_{m}x^{m-1}$ and $\frac{\delta Q(x)}{x}=\frac{A(x)}{|A(x)|}(b_{1}+2b_{2}x+nb_{n}x^{n-1})$.
- •
To optimize the evaluation of polynomials, we employ Horner’s method [], which reformulates a polynomial in a nested form to reduce the computation:
$a_{0}+a_{1}x+\dots+a_{m}x^{m}=a_{0}+x(a_{1}+x(a_{2}+x(\dots)))$
(14)
This allows the evaluation of a polynomial of degree n with only $n$ multiplications and $n$ additions. By default, we use $m=5$ and $n=4$.

Through this efficient CUDA implementation, we largely reduce the computation for each evaluation of the base function. As shown in Table 1, with a scalar input, the rational function with the Horner method is much cheaper than the B-spline used in the KAN paper.

**Table 1: Comparison of FLOPs for different functions. Compared to B-spline function. using Horner’s method with the Rational function reduces FLOPs by approximately $9.3\times$ compared to the B-Spline function.**
| Name | FLOPs |
| --- | --- |
| B-Spline (G=3, K=3) | 204 |
| Rational (m=5, n=4) | 46 |
| Rational (m=5, n=4) w Horner | 21 |

### 4.3 Group KAN

Instead of learning a unique base function for each input-output pair, we can share their parameters within a group of edges. It reduces the number of parameters and computation. This kind of parameter sharing [, ] and group-wise computation [, ] have been key techniques in neural network design.

Specifically, we divide the input channels $d_{in}$ into $g$ groups, sharing parameters among $d_{in}/g$ input channels within each group. Figure 2 illustrates the distinctions between the original KAN, our Group KAN, and a standard MLP. Unlike MLPs, which employ non-learnable activations, KAN assigns a unique function to each input-output pair. Group KAN reduces the number of parameters by sharing these functions among a group of edges.

Group-Rational KAN. We combine the rational function of Section 4.2
with group-wise parameters to implement our Group-Rational KAN (GR-KAN). In practice, we share the parameter for the rational function $F$ for each group; however, each edge retains a unique scalar $w$.

Suppose $i$ is the index of the input channel. With $g$ groups, each group contains $d_{g}=d_{in}/g$ channels, where $\lfloor i/d_{g}\rfloor$ is the group index. The operation of GR-KAN on input vector $\mathbf{x}$ can be expressed as

$$ $\displaystyle\texttt{GR-KAN}(\mathbf{x})=\Phi\circ\mathbf{x}=\begin{bmatrix} \sum_{i=1}^{d_{in}}w_{i,1}F_{\lfloor i/d_{g}\rfloor}(x_{i})&\dots&\sum_{i=1}^{ d_{in}}w_{i,d_{out}}F_{\lfloor i/d_{g}\rfloor}(x_{i})\end{bmatrix}$ (15) $$

With a simple rewrite, this can be expressed in matrix form as the product of a weight matrix $\mathbf{W}\in\mathbb{R}^{d_{\text{in}}\times d_{\text{out}}}$ and a input-wise rational function $\mathbf{F}$

$$ $\texttt{GR-KAN}(\mathbf{x})=\mathbf{W}\mathbf{F}(\mathbf{x})=\begin{bmatrix}w_ {1,1}&\cdots&w_{1,d_{\text{in}}}\\ \vdots&\ddots&\vdots\\ w_{d_{\text{out}},1}&\cdots&w_{d_{\text{out}},d_{\text{in}}}\end{bmatrix} \times\begin{bmatrix}F_{\lfloor 1/d_{g}\rfloor}(x_{1})&\dots&F_{\lfloor d_{in} /d_{g}\rfloor}(x_{d_{in}})\end{bmatrix}^{\top}$ (16) $$

As such, we can implement this GR-KAN layer as a group-wise rational function $\mathbf{F}$ followed by a linear layer

$$ $\texttt{GR-KAN}(\mathbf{x})=\texttt{linear}(\texttt{group_rational}(\mathbf{x }))$ (17) $$

In this form,
sharing parameters across each input channel allows direct application of the rational function to the input vector, equivalently applying it across each grouped edge. In this way, GR-KAN functions as a specialized MLP, with 1) learnable non-linear functions, 2) activation preceding the linear layer, and 3) unique activation functions tailored for each group of edges.

Figure: Figure 2: Comparing our Group KAN with vanilla KAN and MLPs. While KAN has unique function on each input-output pairs, Group KAN share these functions at with a groups of edges.
Refer to caption: x2.png

**Table 2: Comparison of parameter counts among different models. *Func FLOPs* stands for the computation of used non-linear activation. In KAN, $K$ represents the order number and $G$ the grid number. For our GR-KAN, $m$ and $n$ indicate the order of polynomials, and $g$ represents the number of groups. Our model, GR-KAN, has a parameter size comparable to a constant increase over the MLP, whereas the KAN model’s parameters scale with $(G+K+3)$.**
| Name | No. Params | FLOPs |
| --- | --- | --- |
| MLP | $d_{in}\times d_{out}+d_{out}$ | $\text{Func FLOPs}\times d_{out}+2\times(d_{in}\times d_{out})$ |
| KAN | $d_{in}\times d_{out}\times(G+K+3)+d_{out}$ | $\text{Func FLOPs}\times d_{in}+(d_{in}\times d_{out})\times[9K\times(G+1.5K)+2<br>G-2.5K+3]$ |
| GR-KAN (Ours) | $d_{in}\times d_{out}+d_{out}+(m+n\times g)$ | $(2m+2n+3)\times d_{in}+2\times(d_{in}\times d_{out})$ |

In experiments, we notice that for rational function, we share the denominator coefficient $b_{n}$ among all groups and use different $a_{m}$ for each group. It gets better performance.

Parameter and Computation Savings. The original KAN requires $d_{in}\times d_{out}$ unique activation functions. Through our grouping strategy, only $g$ unique functions are needed, reducing the parameter count to a constant overhead compared to a standard MLP.

Except the saving on parameter number, this grouping also reduces computational demands. Each input channel computes the activation function $\phi$ once, shared across all corresponding output channels. In contrast, the original KAN requires that each output channel $j$ to independently compute $\phi_{i,j}$. This results in significant computational savings. The comparison of the number of parameters and computation is listed in Table 2.

### 4.4 Variance-Preserving Initialization

In this section, we aim to initialize the values for $a_{m},b_{n}$ and $w$ in Group-Rational KAN to ensure variance-preserving behavior across the network. At its core, we prevent the growth or reduction of activation magnitudes throughout the layers, thereby maintaining stability.

Figure: Figure 3: Example of fitted functions with rational form.
Refer to caption: x3.png

We revisit the analysis from [] and adapt it to KANs. For a GR-KAN layer, the computation for each output $y_{j}$ is given by $y_{j}=\sum_{i=1}^{d_{in}}\phi(x_{i})=\sum_{i=1}^{d_{in}}(w_{i,j}F(x_{i}))+b_{j}$. We assume that all $x_{i}$ are mutually independent [] and uniformly distributed. Here, $w_{i,j}$ follows a normal distribution $\mathcal{N}(0,\sigma_{w}^{2})$ and $b_{j}$ is initialized to zero. The variance of $y_{j}$ can then be described as:

$$ $\displaystyle Var[y]=d_{in}Var[wF(x)]$ (18) $\displaystyle Var[y]=d_{in}Var[w]\mathbb{E}[F(x)^{2}]$ (19) $$

where $x$, $y$, and $w$ represent the random variables of each element in $x_{i}$, $y_{j}$, and $w_{i,j}$ respectively.
When layers are stacked, we aim for the variance of the input-output activations to remain consistent, expressed as:

$$ $Var[x]=d_{in}Var[w]\mathbb{E}[F(x)^{2}]$ (20) $$

Since $F(x)$ is the rational function containing coefficients $a_{m}$ and $b_{n}$, the initialization of $w$ and these coefficients are interdependent—the form of $F(x)$ influences the appropriate initialization of $w$. The crucial step is to calculate $\frac{Var[x]}{\mathbb{E}[F(x)^{2}]}$ and adjust $w$ to maintain consistent activation scaling.

For our rational function defined in Equation 12, computing $\mathbb{E}[F(x)^{2}]$ involves evaluating:

$$ $\displaystyle\mathbb{E}[F(x)^{2}]=\int_{-\infty}^{+\infty}F^{2}(x)f(x)dx=\int_ {-\infty}^{+\infty}(\frac{a_{0}+a_{1}x+\dots+a_{m}x^{m}}{1+|b_{1}x+\dots+b_{n} x^{n}|})^{2}f(x)dx$ (21) $$

where $f(x)$ is the density function of $x$. Unlike activation functions such as ReLU, for which $\mathbb{E}[F(x)^{2}]=\frac{1}{2}Var[x]$, computing $\mathbb{E}[F(x)^{2}]$ for the rational function is challenging due to the lack of a closed-form solution.

Initialize $a,b$ first, then initialize $w$. To make the process manageable, Instead of sampling $w$, $a$, and $b$ jointly, we proceed sequentially. Initially, we determine $a$ and $b$ such that $F$ fits established activations like ReLU, GELU, and Swish. Figure 3 illustrates the fitted functions.

Once $a$ and $b$ are set, we estimate the gain $\alpha=\frac{\mathbb{E}[F(x)^{2}]}{Var[x]}$ numerically, assuming $x\sim\mathcal{N}(0,1)$(^2^22This assumption is justified as the inputs to the KAN layer are normalized using layer normalization as in Equation 10. The LN layers are initialized to have zero bias and a scaling factor of one). The calculated gains, $\alpha$, are documented in Table 3. We use the gain value to initialize $w$ from $\mathcal{N}(0,\frac{\alpha}{d_{in}})$.

Initialize KAT from ViT. In addition to random weight initialization, we can also transfer weights from a pre-trained ViT to our KAT model. This transfer is straightforward for most layers, as KAT can replicate the micro-architecture of ViT, except for the KAN layer.

Figure: Figure 4: One-to-one weight mapping between trained MLP in ViT and GR-KAN in KAT.
Refer to caption: x4.png

For the GR-KAN layer, weight transfer is still feasible, as shown in Figure 4. Because the GR-KAN layer consists of a linear layer and a group-wise rational layer, we can directly load the weights of the linear layer from the MLP in the trained ViT.

For rational layers, the first one is initialized to behave like an identity function, while the second layer is set to approximate the non-linear function used in the original MLP. This approach allows all the weights of the GR-KAN layer to be cloned from a ViT model, ensuring compatibility and efficient initialization.

## 5 Experiments

### 5.1 Experimental Setup

We modify the original ViT [] architecture by substituting its MLP layers with GR-KAN layers. By default, these KAN layers employ a rational function with parameters $m=5$ and $n=4$, and are organized into groups of 8 ($g=8$). Each transformer block contains 2 KAN layers. The first GR-KAN layer’s $a_{m}$ and $b_{n}$ are initialized to fit the identity function, while the second is initialized to mimic the Swish function []. The attention layers are initialized with Mimetic Initialization []. The remainder of the architecture remains unchanged. We intentionally do not use hierarchical architectures [] for simplicity.

Model Variant. We select the configurations of KAT to be identical with those used in ViT [], as
summarized in Table 4. All variants use an input patch size of $16\times 16$.

**Table 4: Details of KAT model variants.**
| Model | Layers | Hidden Size D | MLP Size | Heads | Params |
| --- | --- | --- | --- | --- | --- |
| KAT-Tiny | 12 | 192 | 768 | 3 | 5.7M |
| KAT-Small | 12 | 384 | 1536 | 6 | 22.1M |
| KAT-Base | 12 | 768 | 3072 | 12 | 86.6M |

### 5.2 Image Recognition

Experiment Setup. We do experiments on ImageNet-1K [59] image classification benchmark. ImageNet-1K is one of the most
widely-used datasets in computer vision which contains
about 1.3M images of 1K classes on training set, and 50K images on validation set.

We mainly follow the hyper-parameters of DeiT []. Specifically, models are trained for 300 epochs at $224^{2}$
resolution. The patch size is set to 16. Data augmentation and regularization techniques include RandAugment [], Mixup [], CutMix [], Random Erasing [], weight
decay, Label Smoothing [] and Stochastic Depth []. We adopt
AdamW [] optimizer with batch size of 1024.

We compare with ViT [] and DeiT [], as we share the same architecture, except for the channel mixer. We also report the results of ViT + KAN [], that simply replacing MLP with standard KAN.

**Table 5: Comparative Analysis of Model Performance and Computational Efficiency on ImageNet-1K. We measure the FLOPs under $224^{2}$ using fvcore package. $*$ indicates that the model is initialized using a pre-trained ViT model, otherwise trained from scratch.**
| Model | Channel Mixer | #Param. | FLOPs | IN-1k Top-1 |
| --- | --- | --- | --- | --- |
| ViT-Ti/16 | MLP | 5.7M | 1.08G | 72.7 |
| DeiT-T | MLP | 5.7M | 1.08G | 72.2 |
| ViT-T + KAN | KAN | 12.8M | 1.78G | 64.9 |
| KAT-T | KAN | 5.7M | 1.13G | 74.6 |
| KAT-T^∗ | KAN | 5.7M | 1.13G | 75.7 |
| ViT-S/16 | MLP | 22.1M | 4.25G | 78.8 |
| DeiT-S | MLP | 22.1M | 4.25G | 79.8 |
| ViT-S + KAN | KAN | 50.4M | 7.05G | 62.9 |
| KAT-S | KAN | 22.1M | 4.35G | 81.2 |
| KAT-S^∗ | KAN | 22.1M | 4.35G | 82.0 |
| ViT-B/16 | MLP | 86.6M | 16.87G | 79.1 |
| DeiT-B | MLP | 86.6M | 16.87G | 81.8 |
| ViT-B + KAN | KAN | 199.8M | 28.04G | NAN |
| KAT-B | KAN | 86.6M | 17.06G | 82.3 |
| KAT-B^∗ | KAN | 86.6M | 17.06G | 82.8 |

Results. Our experimental results demonstrate that the KAT models consistently outperform their counterparts on the IN-1k dataset, as shown in Table 5. Firstly, the integration of GR-KAN in the transformer architectures demonstrates superior performance over traditional MLP-based mixers. For instance, the KAT-S model achieved an accuracy of 81.2%, outperforming the DeiT-S model by 2.4%. This improvement underscores the potential of KAN mixers to enhance model efficacy when properly integrated.

Secondly, the vanilla KAN layer faces scalability issues. ViT-T/S + KAN only achieved an accuracy of around 63%, even with a much higher computational cost. ViT-L + KAN fails to converge, resulting in NAN error. We addressed these scaling challenges as detailed in Section 3, enabling our KAT models to scale successfully.

These findings highlight the efficacy of the KAT approach in balancing computational efficiency with improved performance, suggesting valuable directions for further research in optimizing transformer architectures.

### 5.3 Object Detection and Instance Segmentation

Experimental Setup. We evaluate our approach on the MS-COCO2017 [] dataset, a standard benchmark for object detection and instance segmentation. In our setup, the KAT is employed as the backbone within a ViTDet-based [] Mask R-CNN [] model, initialized with weights pre-trained on ImageNet. We followed the standard $3\times$ training schedule, which consists of 36 epochs. The training images were resized to $800\times 1333$ pixels. The AdamW optimizer [] was used with a learning rate of 0.0001 and a total batch size of 16. Our implementation was based on the PyTorch and MMDetection [] libraries, and we use FP16 precision to reduce training costs. The experiments were carried out on 4 NVIDIA H100 GPUs.

Results. Table 6 compares the performance of different backbones. KAT consistently outperformed other models, particularly in object detection, where it achieved a 3.0 AP${}^{\text{box}}$ gain on the S-sized model and a 1.4 AP${}^{\text{box}}$ gain on the L-sized model compared to ViTDet. The improvements were most pronounced in smaller models, where computational cost increased by only 1 GFLOPs. This shows that KAT offers better accuracy with minimal overhead.

**Table 6: Performance of Mask-RCNN with different backbones on 3$\times$ schedule.**
| Backbone | #Param. | FLOPs | AP${}^{\text{box}}$ | AP${}^{\text{box}}_{50}$ | AP${}^{\text{box}}_{75}$ | AP${}^{\text{mask}}$ | AP${}^{\text{mask}}_{50}$ | AP${}^{\text{mask}}_{75}$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| PVT-Small | 44.1M | - | 43.0 | 65.3 | 46.9 | 39.9 | 62.5 | 42.8 |
| Swin-T | 48M | 267G | 46.0 | 68.1 | 50.3 | 41.6 | 65.1 | 44.9 |
| ConvNeXt-T | 48M | 262G | 46.2 | 67.9 | 50.8 | 41.7 | 65.0 | 44.9 |
| ViT-S | 43.8M | 423G | 44.0 | 66.9 | 47.8 | 39.9 | 63.4 | 42.2 |
| ViTDet-S | 44.5M | 423G | 44.5 | 66.9 | 48.4 | 40.1 | 63.6 | 42.5 |
| KATDet-S | 44.5M | 424G | 47.5 | 69.0 | 51.2 | 41.5 | 65.7 | 44.0 |
| ViT-B | 113.6M | 767G | 45.8 | 68.2 | 50.1 | 41.3 | 65.1 | 44.4 |
| ViTDet-B | 113.6M | 767G | 46.3 | 68.6 | 50.5 | 41.6 | 65.3 | 44.5 |
| KATDet-B | 113.7M | 770G | 47.7 | 69.1 | 51.6 | 41.6 | 65.9 | 44.3 |

### 5.4 Semantic Segmentation

Experiment Setup. We evaluated our KAT model on the ADE20K dataset []. This dataset comprises 150 semantic categories with 20,000 images in the training set and 2,000 in the validation set. For our experiments, we utilized KAT as the backbone for the UperNet framework [], initializing it with ImageNet pre-trained weights. The training was conducted using the AdamW optimizer [] with a learning rate of 0.0001 and a batch size of 16, across 160,000 iterations. Our implementation was carried out using the PyTorch and mmsegmentation libraries, and the experiments were performed on two NVIDIA H100 GPUs. For comparison, we evaluated UperNet with other backbones, including DeiT, Swin Transformer, and ConvNeXt.

**Table 7: Performance of Semantic segmentation with UperNet on ADE20K validation set. Images are cropped to $512\times 512$ for training. The MACs are measured with input size of $512\times 2048$.**
| Backbone | #Param. | FLOPs | mIoU (%) |
| --- | --- | --- | --- |
| Swin-T | 60M | 945G | 45.8 |
| ConvNeXt-T | 60M | 939G | 46.7 |
| DeiT-S | 57M | 1217G | 43.5 |
| KAT-S | 57M | 1219G | 46.1 |
| Swin-B | 121M | 1188G | 49.5 |
| ConvNeXt-B | 122M | 1170G | 49.6 |
| DeiT-B | 142M | 2007G | 47.2 |
| KAT-B | 142M | 2011G | 47.4 |

Results. Table 8 summarizes the segmentation results. Overall, KAT demonstrates a competitive improvement over plain ViT-based architectures, achieving a 2.4% improvement over DeiT-S and a 0.2% improvement over DeiT-B. This performance boost comes with a slight increase in computational cost, reflected in the higher FLOPs. Similar to the detection results, KAT shows more significant gains in smaller models. However, it still falls short compared to models with hierarchical architectures, such as ConvNeXt, which benefit from more efficient structural design.

### 5.5 Ablation Study and Analysis

Activation Function. As GR-KAN can be considered as a special kind of MLP with group rational function, we do an ablation study to consider different types of activation for MLP and compare with our GR-KAN. Superficially, we replace the activation function in MLP in ViT-Ti/16 to different kinds, including GELU [], ReLU [], SiLU [], PReLU [] and PAU [], and comparing them with KAT.

Table 8 summarizes the top-1 accuracy on the ImageNet-1k dataset for each activation function with ViT-Ti/16. The results indicate that traditional activation functions like ReLU and GELU perform similarly. Learnable activations like PReLU and PAU show an improvement. Notably, Our KAT-T achieves the highest accuracy at 74.6%, outperforming GELU by 1.9%. This suggests that GR-KAN, as used in KAT-T, can significantly enhance the expressiveness of MLPs in vision transformers.

In addition to accuracy, we analyzed the computational cost of different activations by measuring throughput and peak memory usage on an NVIDIA A5000 GPU (Table 9). All activation functions showed similar peak memory usage. However, our method (KAT-T) showed slightly lower throughput compared to the baseline activations (e.g., ReLU, GELU, and SiLU), which are more efficient. This suggests that while KAT-T offers substantial accuracy improvements, there is a trade-off in computational efficiency, which may be attributed to the increased complexity of rational function computations.

**Table 9: Throughput and Peak memory for different activation on A5000 GPU. Input size is fixed to $[64,1000,512]$.**
| Activation | ReLU | GeLU | SiLU | PReLU | Ours |
| --- | --- | --- | --- | --- | --- |
| Throughput (batch/s) | 2654 | 2643 | 2668 | 2644 | 2313 |
| Peak Memory (M) | 1380 | 1380 | 1380 | 1380 | 1380 |

Benefit of CUDA Implementation. To evaluate the efficiency improvements introduced by our CUDA implementation discussed in Section 4.2, we conducted experiments to measure both forward pass speed and peak memory usage. Specifically, we compared our CUDA implementation against two alternative methods. The first is called *Torch Looped*, which loops over each channel group, applies the rational function, and then concatenates the results. The second is called *Torch Vectorized*. In this method, the input tensor is reshaped according to the channel groups, the rational function is applied in a vectorized manner, and the tensor is reshaped back to its original form. We compare these three implementation on A5000 GPU, under 1) different group number $g\in\{1,2,4,8,16\}$. 2) different input dim $D\in\{128,256,512,1024,2048\}$

Figure: (a) Throughput (batch/s) for Different Group Sizes. Larger the better.
Refer to caption: x5.png

Figure: (a) Throughput (batch/s) for Input Dimension Sizes. Larger the better.
Refer to caption: x7.png

The results, presented in Figure 5 and Figure 6, clearly demonstrate that our CUDA implementation significantly outperforms both the Torch Looped and Torch Vectorized implementations, offering superior speed and memory efficiency.

Rational Initialization. We tested our KAT-T model with different initializations of the rational functions when training from scratch. As shown in Table 10, the “Identity - Swish” initialization achieves the best performance, which we have adopted as our default setting.

**Table 10: Ablation on rational function initialization, with KAT-T.**
| Rational 1 Init. | Rational 2 Init. | IN-1k Top-1 |
| --- | --- | --- |
| Identity | Identity | 69.7 |
| Swish | Swish | 74.4 |
| Identity | GeLU | 74.5 |
| Identity | Swish | 74.6 |

Visualization of Trained Functions

An important aspect to examine is the behavior of the trained rational functions. As shown in Figure 7, we plot the functions for KAT-S with $g=8$ across all 12 layers. The results indicate that within each layer, the rational functions exhibit similar trends, while the functions across different layers tend to differ from one another.

## 6 Conclusion and Future Work

In this work, we introduced the Kolmogorov–Arnold Transformer (KAT), a novel architecture that successfully integrates Kolmogorov-Arnold Networks (KANs) into transformers, addressing key challenges associated with large-scale training scenarios. Our proposed Group-Rational KAN (GR-KAN) variant, with its rational activation functions, group-based parameter sharing, and variance-preserving initialization, demonstrated significant improvements in computational efficiency and scalability. Through extensive experiments on vision tasks, including image recognition, object detection, and semantic segmentation, KAT outperformed traditional MLP-based transformers, achieving superior accuracy on ImageNet1K while maintaining comparable computational demands.

Discussion.
Our study highlights KAT’s potential as a good alternative to MLP-based transformers, especially in large-scale vision tasks. This integration introduces exciting opportunities for broad applications. For example, employing KAT architectures might help development of language models.

However, KAT is not without its challenges. A primary concern is running speed. Even with the CUDA optimized code, the rational function is still slower than plain activation like ReLU and GELU. Another issue is the stability when using rational functions in neural networks. The higher order gradients for $a_{m}$ and $b_{n}$ can become unstable because of their dependence on the input power. Integrating these functions into the backpropagation process could introduce complications.

Additionally, it is important to acknowledge that our GR-KAN represents a hybrid model. On the one hand, GR-KAN is a KAN layer with shared edges and a rational base function. On the other hand, it can be interpret as MLP with a redesigned activation placed before the linear layer. It leverages the computational simplicity of MLPs but maintains some characteristics of KANs. However, GR-KAN is not a pure KAN model. Instead, it merges advantages from both systems to enhance overall functionality.

Future Work.
There are multiple directions of KAT for future research. One potential area of exploration is to find alternative base functions to further improve computational efficiency and compatibility with emerging hardware architectures. Currently, rational functions serve as one option, but other possibilities exist. These include Fourier transformations [], Wavelet transforms [], and Gaussian radial bases [].

Additionally, expanding the applicability of KAT to other domains beyond vision tasks, such as natural language processing or reinforcement learning, could unlock new opportunities for performance gains. Further research could also investigate hybrid models [, ], or adaptive mechanisms for dynamically selecting between KAN and MLP layers based on the complexity of the task, thereby optimizing resource utilization. Finally, addressing the remaining scalability challenges, particularly in terms of memory footprint and inference speed, will be crucial for deploying KAT in real-world applications at scale.

## Acknowledgement

We would like to acknowledge that computational work involved in this research work is partially supported by NUS IT’s Research Computing group using grant numbers
NUSREC-HPC-00001. We thank Weihao Yu, Qiuhong Shen and Runpeng yu for valuable discussions.

## 7 Derivation and Calculation of FLOPs

Given the function:

$$ $F(x)=\frac{a_{0}+a_{1}x+\dots+a_{m}x^{m}}{1+|b_{1}x+\dots+b_{n}x^{n}|}$ $$

### 7.1 Plain Computation

#### Numerator

The numerator is a polynomial of degree $m$:

- •
Multiplications: There are $\frac{m(m+1)}{2}$ multiplications for computing powers of $x$ and $m$ multiplications for coefficients $a_{i}$, giving $\frac{m(m+1)}{2}+m$.
- •
Additions: There are $m$ additions to sum up the polynomial terms.

#### Denominator

The denominator involves the absolute value of a polynomial of degree $n$:

- •
Multiplications: There are $\frac{n(n+1)}{2}$ multiplications for powers of $x$ and $n$ multiplications for coefficients $b_{i}$, giving $\frac{n(n+1)}{2}+n$.
- •
Additions: There are $n$ additions for polynomial terms and 1 additional addition after the absolute value operation.
- •
Absolute value operation: 1 absolute value calculation.

Division. There is 1 division operation for the final computation of $F(x)$.

Total FLOPs.
The total FLOPs for any $m$ and $n$ are:

$$ $\text{Multiplications: }\frac{m(m+1)}{2}+\frac{n(n+1)}{2}+m+n+1,\quad\text{ Additions: }m+n+1,\quad\text{Absolute Value: }1,\quad\text{Division: }1$ $$

In case $m=5$ and $n=4$, there are totally 34 multiplications, 10 summations, 1 absolute value and 1 division. In total 46.

### 7.2 Horner’s Method

Using horner’s method, for a polynomial of order $m$, we need $m$ summations and $m$ multiplications.

Thus, for numerator, we need $m$ summations and $m$ multiplications. For denominator, we need $n+1$ summations and $n$ multiplications. In total, we need $m+n+1$ summation, $m+n$ multiplications, $1$ absolute value, and $1$ division.

In case $m=5$ and $n=4$, there are a total of 21 FLOPs, comprising 9 multiplications, 10 summations, 1 absolute value, and 1 division.

## 8 Hyper-parameters for KAT model

The hyper-parameter for training KAT model on ImageNet-1k is shown in Table 11.

**Table 11: Hyper-parameters of KAT on ImageNet image classification.**
|  | KAT |  |  |
| --- | --- | --- | --- |
|  | Tiny | Small | Base |
| Input resolution | $224^{2}$ |  |  |
| Epochs | 300 |  |  |
| Batch size | 1024 |  |  |
| Optimizer | AdamW |  |  |
| Adam $\epsilon$ | $1\times 10^{-8}$ |  |  |
| Adam ($\beta_{1}$, $\beta_{2}$) | (0.9, 0.999) |  |  |
| Learning rate | $1\times 10^{-3}$ |  |  |
| Learning rate decay | Cosine |  |  |
| Gradient clipping | None |  |  |
| Warmup epochs | 5 |  |  |
| Weight decay | 0.05 |  |  |
| Rand Augment | 9/0.5 |  |  |
| Repeated Augmentation | off |  |  |
| Cutmix | 1.0 |  |  |
| Mixup | 0.8 |  |  |
| Cutmix-Mixup switch prob | 0.5 |  |  |
| Random erasing prob | 0.25 |  |  |
| Label smoothing | 0.1 |  |  |
| Peak stochastic depth rate | 0.1 | 0.1 | 0.4 |
| Random erasing prob | 0.25 |  |  |
| EMA decay rate | 0.9999 |  |  |

Figure: Figure 7: Fitted rational functions for KAT-S model, with 12 layers and 8 groups.
Refer to caption: x9.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Q1) As per my understanding, Kolmogorov Arnold Representation Theorem and the setup mentioned here involves learning a multivariate function, why do *lines 146-147* in the *section 2.2* mention *learn a univariate functions on the edge*. The dimension of $f(x)$ would be 
$d_{out}$ making it multivariate?

Q2) Do we have a theoretical understanding of the behaviour outlined in lines *316-317*? How are the parameters adjusted in case of B-splines or any other type of functions?
soundness: 3
strengths: S1) Existing works using KANs in-place of MLP modules in various models (such as feedforward NNs, CNNs, Transformers etc) have had limited success due to various issues. This work goes a step further and proposes an alternative to the original KAN architecture called Group Rational (GR-KAN) solving the issues for inefficiencies in base functions, parametrizations and base initializations reducing the parameters and FLOPs from the original KAN architecture.

S2) The GR-KAN employs rational activations as the base function which are efficient and suitable for GPU computation replacing the B-Splines which are not optimized for GPU or parallel processing. (which is one of the core strengths of this work amongst related KAN literature)

S3) The KAT architecture improves the performance over standard transformer architectures such as ViT within the same computational budget.
summary: There has been a resurgence in Kolmogorov Arnold Networks (KANs) recently as an effective alternative to MLPs. This work carefully studies the major issues with scaling of the standard KAN architecture and proposes suitable modifications to create what they define as Group-Rational KAN (GR-KAN). The GR-KAN is then used as a replacement for the MLP layer in the transformer architecture in this work, which is then called the Kolmogoorv Arnold Transformer (KAT). The utility of the KAT architecture is demonstrated through experiments on various vision tasks such as Image Recognition, Object Detection, Instance Segmentation and Semantic Segmentation.
weaknesses: W1) This work doesn’t include results on using KANs in solving PDEs (as the original work did, albeit the architecture won’t be transformers for this task) and KATs for language processing (mentioned in future work). The experiments are majorly limited to vision tasks. (Table and Graph Classification results have also been included.)

W2) Ablations on the function which parametrizes the learnable activation are not included. Ablations are mainly with respect to alternative MLPs and fixed activation functions. Ablations could have included alternatives to rational functions, B-splines etc.

**Presentation/Typos/Corrections**

I believe that the *Experiments* section can be shortened or moved to appendix and the *Appendix F Discussion and Future Work* be moved to the main paper as it packs a lot of relevant quality content.

*Appendix D Section D.1*
This is a minor correction, the number of multiplications in exponent will be $\frac{m(m-1)}{2}$ and correspondingly $\frac{n(n-1)}{2}$ instead of $\frac{m(m+1)}{2}$ and $\frac{n(n+1)}{2}$ respectively. 
As an example case for $m=1$, the numerator becomes $a_0 + a_1x$ which involves 0 multiplications in the exponent computation and a single multiplication from the multiplication for the coefficients.
The total multiplications will be in fact $\frac{m(m+1)}{2}$ which is $\frac{m(m-1)}{2}+m$
Note that with this change the values in the Table 1 and the appendix D will change accordingly.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 4
questions: How do you explain the contradiction of saying KANs require less paramters and then saying KAN requires more parameters?
soundness: 3
strengths: Indentified issues with integrating KANs with ViT.
Overcame the recursive computation requirement for using B spline curves in KANs and replaced them with faster rational functions.
Introduced variance preserving initialisation to stabilize training.
Created an efficient an fast CUDA implementation of rational function for faster running in GPU
Improved performance gains by grouping and sharing parameters
Extensive experimetation is provided
summary: The paper introduces the Kolmogorov–Arnold Transformer (KAT), where they replaces MLP layers in transformer with Kolmogorov-Arnold Network (KAN) layers. B-spline functions in KANs were replaced with rational functions to improve compatibility with modern GPUs. Group KAN was used to reduce computations and variance preserving initialization was employed. Introduces Group-rational KANs for intergation into ViT. Extensive experiments on image classsification has shown improved performance of GR-KAN.
weaknesses: Replacing MLP in transformers with KAN is an obvious choice and does not warrant any novelty.
phi_in matrix in eq. 2  Top row right element should show phi_1,d instead of phi_1,n
In line 46, the paper says KANs require fewer parameters than MLPs, then goes on to show in Line 176-189 than KAN requires more paramters.
The violation of variance-preserving nature in higher order B-splines is not clear.
The basic novelty is the use of rational function instead of B-splines

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: 1) Can you explain very specifically the novelty of your method? You break your contribution into three pieces, can you say which portions of these are novel and which are not?
2) Much of the presentation about scalability is presented through FLOPs computation. Is there a reason that you chose this representation, rather than actual training/inference time?
3) A major part of this work is the CUDA implementation of the layer. How do you plan to make this available? Will the layer also be eventually made available through frameworks like pytorch or jax? I appreciate the work that goes into a direct CUDA implementation, but I think it is crucial that there are plans for wide availability in popular frameworks in order to have this work translate to impact in the community.
soundness: 3
strengths: 1) The paper is generally well written: it has a clear narrative, each new section is well-motivated, and for the most part, technical discussions are nicely balanced with takeaways and intuition.
2) The technical analysis and contributions of the paper directly lead to practical improvements. The entire chain of derivations/modifications creatively connect different lines of research (KAN networks + rational neural networks, for example) and add some new techniques as well. In particular, I found the re-expression of the grouped parameters in Section 4.3 into a specialized MLP to be quite elegant.
3) The experiments are comprehensive, all of the important training details are made available, and the results show that using these layers does lead to improvement in accuracy with relatively little increase in training time.
summary: This work presents the Komlogorov-Arnold Transformer, a transformer architecture that replaces the typical transformer layer's concluding 2-layer MLP with (modified) Kolmorogov-Arnold layers. The authors explain why a simple substitution of KAN layers will fail to both perform well and to scale with the hidden dimension of the network, and provide fixes (via different base functions, weight sharing across grouped edges, and improved initialization) that improve performance and speed. Improvements are shown in experiments across a variety of machine learning tasks in vision domains and beyond.
weaknesses: 1) I felt some details about the method and novelty were missing. In my opinion this is mostly clearly true in the initialization section (which seems to be where a lot of the improvements/convergence come from). A lot of Section 4.4 is spent showing that choosing these initialization values is difficult and non-obvious - which I agree with! - but then there is no space left to actually show how you pick values for a and b (I know you say that you "determine a and b such that F fits established activations", but a small algorithm or set of equations would do a lot to disambiguate this).

2) The suite of experimental results is impressive, but there are comparatively few ablations studies. I am left wondering which parts of your method are responsible for the improvement over standard ViTs. This is particularly disappointing because you clearly break down your improvements step-by-step, so it is easy to imagine what such an ablation would look like. I realize that some of your modifications are meant to address scalability, so a large set of experiments might be infeasible. But something like accuracy & total training time of a ViT-T with different parts of your method (i.e., with/without your initialization scheme, with/without grouping, etc) seems like it should be possible.

3) Similarly to the above: after switching the base functions to rationals, the GR-KAN layer in practice resembles the rational layers from Boulle et al much more than it resembles the original KAN layers. But there are no direct comparisons to those; I would also find this experiment useful.

More minor pieces of feedback:

4) On L248: I don't think it's quite right to say that this gradient computation is "Similar to Monila et al"; it is identical to their calculation. I also don't think the details of this calculation are used anywhere later in the paper, unless I am mistaken, so they could also be relegated to the appendix.
5) Table 1: This caption could use more detail - is this just the forward pass, or forward and backward passes? Also, a comparison to an analogously-sized MLP layer would be additionally useful.
6) L433: "ViT-L + KAN fails to converge" - I think this should be ViT-B?
7) Figure 3: I think I understand what this figure is trying to communicate (i.e., that the choice of values for $a$ and $b$ do a good job of fitting the activation functions), but in practice it just ends up looking like... graphs of activation functions. Something quantitative would probably be more effective is communicating how well the rationals approximate the activation functions, and it would take up less space as well.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - In Eq. 5, 6, and 7 what are MSA and LN? 
- Missing text in paragraph name in line 485 
- In Tab. 4 KATDet-S is reported as best (41.5) but ConvNext and SwinT are higher (41.7 and 41.6)
- I feel like the change you propose can be applied to KANs also outside of the transformer architecture. Perhaps it could be useful to add a comparison between KANs and GR-KANs on toy problems
soundness: 3
strengths: - The topic is highly relevant, as literature has not focused much on improving the MLP aspect of transformers
- KANs have emerged as an interesting alternative to MLPs, however, they present issues (e.g. scaling) which make them unfavorable w.r.t. MLPs
- The proposed approach is scalable, faster, and requires less parameters than KANs, making it usable in practical settings and on large-scale datasets such as ImageNet-1k or MSCOCO
- The experimental evaluation comprises different tasks such as classification, object detection and segmentation.
summary: In this work, authors propose Kolgomorov-Arnold Transformer (KAT) as an improvement to standard transformer architecture, which uses MLP networks, by replacing MLPs with KAN networks. Authors propose a number of improvement over vanilla KAN architectures to make them scalable in practical settings and achieve better results: 1.) they swap b-spline functions with rational function for speed 2.) they make use of parameters sharing (grouping) 3.) they employ a variance preserving initialization.
Results show that KAT achieves better results than ViT on certain benchmarks.
weaknesses: Here are my doubts about this work: 
- While the results for ImageNet are convincing, the proposed KAT still lags behind traditional architectures such as ConvNext or even transformers such as SwinB in segmentation. 
- A proper ablation study is missing (e.g. combinations of KAN + rational function + grouping + initialization). It is difficult to evaluate the contribution of each proposed change quantitatively. Also I wonder how simple KAN perform with parameter grouping in terms of MACS/FLOPS and results? 
- It is not clear how you choose the hyperparameters for KAT such as number of groups and maximum order of rational function (m)
- (More of a philosophical doubt) KAN are based on the Kolmogorov-Arnold representation theorem, which states that a multivariate function can be approximated by a composition of univariate functions, while your proposed method only uses one irrational function. Can you still say it's based on the Kolmogorov-Arnold theorem? 
- (minor) some typos and missing notation are present in the paper (see questions below)

---

confidence: 4
contribution: 4
initial_rating: 6
presentation: 3
questions: * Could you specify where in the report it states that KAN is "10x slower than MLPs, given the same number of parameters"?
soundness: 4
strengths: * Comparative experiments are conducted on image classification, object detection, instance segmentation, and semantic segmentation.
* Challenges and solutions are clearly explained, supported by experimental results.
* The proposed method is sufficiently general to be applicable across different transformer architectures.
summary: The author proposes and implements Kolmogorov-Arnold Transformer (KAN) architecture. Especially, they propose rational activation functions, group KAN, and variance-preserving initialization, demonstrating improved performance over conventional vision transformers across various tasks.
weaknesses: * In Tab 1, given the efficient CUDA implementation, experimental results of GPU computation are needed. Specifically, the inference time on a specified GPU should be provided.
* The choice of the number of groups in group-rational KAN is not discussed. What is the reason for the current choice of group number? An experiment analyzing the effect of different group numbers would be beneficial for better understanding.
* For consistency, it would be helpful to include the GPU information used for image classification task as well.

Meta-review: The paper introduces the Kolmogorov–Arnold Transformer (KAT), which replaces traditional MLP layers in transformers with Kolmogorov-Arnold Network (KAN) layers to enhance model performance. The authors address three main challenges: the inefficiency of B-spline functions, the computational load of unique functions for each input-output pair, and the difficulty of weight initialization. 
The KAT architecture demonstrates improved performance across various vision tasks such as image recognition, object detection, and semantic segmentation. Concretely, the authors demonstrate the benefits over the base transformers ViT and DeiT in the seminal ImageNet-1k, prototyping their architecture there. The KAN achieves an improved performance with the same or less FLOPs. One major drawback of the paper is that it does not compare with other architectures, e.g., MLP-Mixers or their variants, Mamba variants, etc. There is a wealth of architectures the last few years, including papers published in all major ML conferences in 2024, which I hope the authors can include in the camera-ready paper.

Decision: Accept (Poster)