Kolmogorov-Arnold Transformer

We thank R-s19p for the nice suggestions.

>>> Q1 GPU running time

>>> A1 That is a great question. In the paper, we report the CUDA running time on an A5000 GPU in Appendix J.1, for each layer with different width and group number. Our implementation is faster than the pure PyTorch version.

Additionally, as requested, here is the full model throughput on an A5000 GPU using FP32. While our model is slightly slower than a pure ViT, it is significantly faster than ViT+KAN.

>>> Q2 Group Number

>>> A2 We thank the reviewer the the nice question. To verify this, we conducted an ablation study using the KAT-Tiny model to explore the impact of different group numbers.

Our results showed that increasing the number of groups improved accuracy slightly up to 8 groups, with no further gains beyond that. Based on these findings, we chose 8 groups as the optimal configuration, providing a good trade-off between simplicity and performance.

The results has been added to revision Appendix J.1.

>>> Q3 GPU information

>>> A3 For the image classification tasks, some experiments were conducted on servers with 8×A5000 GPUs. For the larger KAT-B model, we used servers equipped with 8×H100 GPUs.

>>> Q4 KAN statement

>>> A4 Thanks for the question. The authors of KAN made this statement in its paper, at the section Final takeaway: Should I use KANs or MLPs?. They point out that

Currently, the biggest bottleneck of KANs lies in its slow training. KANs are usually 10x slower than MLPs, given the same number of parameters.

Thank R-s19p so much! We’re glad the concerns have been resolved and appreciate the valuable feedback.

I appreciate that the authors incorporated a study on GPU computation and the group number of KAT. All of my primary concerns have been resolved.

>>> Q3 Hyperparameters for KAT

>>> A3 Thanks for the question.

• Number of groups: As suggested, we conducted an ablation study with the KAT-Tiny model to determine the optimal number of groups.

  The results showed that accuracy improved slightly up to 8 groups, with no further gains beyond that. Based on these findings, we chose to use 8 groups, which offer a good balance between simplicity and performance.

  Group Number	2	4	8	16	32
  KAT-Tiny Top-1	74.2	74.3	74.6	74.7	74.6

• Maximum order of rational: We use the rational order (m=5,n=4), as it is the default setting in the PAU paper. We also conduct a preliminary experiment to confirm this choice

  Order Numer (m,n)	(3,2)	(5,4)	(7,6)
  KAT-Tiny Top-1	74.2	74.6	74.6

  The results indicate that increasing the order beyond (m=5,n=4) provides no additional accuracy gains. It stands as a practical and efficient choice.

The results has been added to Appendix J.1 in the revision.

>>> Q4 Connecting to theorem

>>> A4 Thanks. It is true that the KA theorem states that a multivariate function can be approximated by a composition of univariate functions.

However, we are not using just a single univariate rational function. Instead, we approximate the multivariate function by summing several univariate rational functions. Some of these rational functions may share parameters, but the key idea remains consistent with the KA theorem.

>>> Q5 Typos and missing notation

>>> A5

• MSA and LN: MSA stands for multi-head self attention and LN stands for layer norm.
• Missing text: We change the paragraph title from v.s. Activation Function to GR-KAN v.s. Activation Function.
• Annotation: Sorry for the wrong annotation. We have revise the manuscript.

>>> Q6 Toy example

>>> A6 We truly appreciate the question. Our GR-KAN can indeed be applied to other tasks. In this answer we test on 2 tasks.

• Regression: We first test on fitting the GR-KAN onto some special functions. The functions were selected based on the examples used in the KAN paper. We use the $[2\to5\to1]$ network architecture trained for 1000 epochs with the Adam optimizer and a learning rate of 0.001. The MSE results, shown below, indicate that smaller values are better. GR-KAN achieves the best performance.

• PDE solving: To solve a PDE, we use a one-dimensional damped harmonic oscillator governed by:

$m \frac{{d^2 u}}{{dt^2}} + \mu \frac{{du}}{{dt}} + k u = 0$ where $m$ is the mass. $\mu$ is the damping coefficient. $k$ is the stiffness constant. With the initial condition $u(0) = 1,u'(0) = 0,m = 1$. The exact solution should be $u(t) = e^{-d \cdot t} \cdot (A \cos(\omega t + \phi))$, where $d = \mu / 2 w_0 = \sqrt{k}, w = \sqrt{w_0^2 - d^2}, \phi = \arctan(-d / w)$. We solve this using MLP, KAN, and GR-KAN with a network architecture of $[2 \to 5\to 1]$. MLP performs the worst, KAN achieves the best results but is slow, while GR-KAN trains faster and performs slightly worse than KAN in this experiment.

The results has been incorporated in the revised paper Appendix B.

We sincerely thank Reviewer DyMH for the thoughtful comments and suggestions. We have conducted new experiments to address the questions raised.

>>> Q1 Comparision with ConvNext and Swin

>>> A1 Great question. While our performance is behind ConvNext and Swin, this is due to differences in micro-architecture. The hierarchical designs in ConvNext and Swin are inherently more effective than the plain ViT-style architecture we used, especially for tasks like segmentation.

New Results: Despite this, our method can be extended to hierarchical architectures. We tested replacing MLP layers in Swin and $1 \times 1$ convolutions in ConvNext with GR-KAN layers.

Results on ADE20K show that our GR-KAN consistently outperforms the original MLP layers across these architectures.

Note that, Our primary goal was to demonstrate that KAN outperforms MLP within transformers. To focus on this comparison, we intentionally adopted a simple ViT-like design, as noted on Line 399.

>>> Q2 Detailed ablation study

>>> A2 We truly thank the suggestions from R-DyMH. As suggested, we conducted a detailed ablation study focusing on three components and the KAN with parameter grouping modification. These experiments were performed on the KAT-Tiny model using ImageNet with 300 epochs of training.

We analyzed the effects of removing the following components:

1. Rational base function, replaced with B-splines as used in KAN.
2. Group-wise weight sharing, replaced with distinct parameters for each channel.
3. Proper initialization, replaced with random initialization.

Ablation 1: Base function: We replaced the rational base function with a B-spline implemented in torch using the De Boor-Cox algorithm due to the lack of a pure CUDA implementation for B-splines.

• Observation: The choice of base function had a minor impact on performance and influenced runtime significantly.
• Key Insight: As shown in Exp 3 and Exp 6, using rational functions slightly improved performance over B-splines. While the increase in MAC count was minimal, our pure CUDA implementation for rational functions ran considerably faster than the torch-based B-spline implementation.

Ablation 2: Group-wise computation: We replaced group-wise weight sharing with a setup where each channel had distinct parameters.

• Observation: Group-wise computation proved to be the most significant factor for efficiency.
• Key Insight: According to Exp 4 and Exp 6, group-wise computation reduced training time from 38 hours to 12 hours, with a marginal drop in accuracy from 74.8% to 74.6%.

Ablation 3: Initialization: We replaced proper initialization with the default initialization in torch.

• Observation: Proper initialization was critical for fast and reliable convergence, particularly for rational functions with higher-order terms.
• Key Insight: As shown in Exp 6 and Exp 5, without proper initialization, terms of different orders were initialized at similar scales, causing instability. This issue was more pronounced for rational functions than for B-splines, as evidenced by the performance differences between Exp 2 and Exp 5.

KAN with parameter grouping: In Exp 2, we explored parameter sharing in KAN by applying parameter grouping to the original ViT+KAN architecture. This modification significantly improved the training speed, reducing the time from 43 hours to 20 hours. However, it resulted in a performance drop of 2.7%.

We have incorporated the results in Appendix C.

>>> Q5 Table 1 Caption

>>> A5 We appreciate the suggestion. Table 1 is the forward pass FLOPs for each edge in KAN, using different non-linear functions. We have revise in the paper.

However, comparing this to an analogously-sized MLP is not straightforward because MLPs do not apply non-linear functions on each edge; instead, they add a non-linear function at the end. A more meaningful comparison is the overall FLOPs, which is already provided in Table 2.

>>> Q6 ViT-B

>>> A6 We have revise the ViT-L to ViT-B. Thanks.

>>> Q7 Fig 3

>>> A7 Thanks for the suggestion. We have added the Mean Squared Error (MSE) to the Fig 3 to illustrate how well the rational functions approximate the activation functions.

>>> Q8 Key novelty

>>> A8 Great question. We believe our main novelty lie in presenting a framework-wise solution for integrating KAN into transformers. Additionally, we consider the analysis and techniques we developed to be of considerable new.

• Framework Novelty. We are the first to successfully integrate KAN into a transformer and make it work. The replacement is simple, but no one has make this replacement sucess.
• Analysis Novelty. We are the first to provide a quantitative analysis on why KAN lacks scalability. The challenges discussed in Section 3 are introduced for the first time within the context of KAN design.
• Techniqueal Novelty: Our core technical novelty lies in the development of a group-wise computation for activation function. To the best of our knowledge, it has not been introduced before. Besides, the application of rational functions to KAN and the initialization under such case is relatively new, although some aspects draw inspiration from prior research

>>> Q9 FLOPs vs. Runtime

>>> A9 Thanks for the question. Both metrics work for comparison, but FLOPs are more appropriate here because the implementation difference.

FLOPs measure number of operations independent of implementation. In contrast, runtime comparisons can be misleading unless the code is similarly optimized.

For example, The original KAN uses inefficient NumPy code, whereas we optimized GR-KAN on CUDA. GR-KAN can be $100\times$ faster on GPU, but that’s not a fair comparison. Likewise, comparing GR-KAN’s runtime to a torch MLP isn’t fair, as torch benefits from extensive optimization by thousands of engineers.

Given these factors, FLOPs offer a more meaningful and consistent basis for comparison at this stage.

>>> Q10 Operational plan for the project

>>> A10 This is the right question to ask. Currently, we provide the code as a C++ extension for PyTorch, making it usable in torch for various tasks.

We’re also in contact with the maintainers of the Transformers and timm libraries. We are right now working on integrating KAT into their codebase.

We feel incredibly fortunate to have R-cLED's thoughtful and valuable suggestions.

>>> Q1 Details for Initializing $a_m,b_n$

>>> A1 Great question. Given a ground-truth function $g(\cdot)$ and a parameterized rational function $F(\cdot;\{a_m\},\{b_n\})$, we run a linear least square to determine $\{a_m\},\{b_n\}$. Specifically, we optimize the following function: $\min_{\{a_m\},\{b_n\}} \frac{1}{2} \sum_{i=1}^N (g(x_i)-F(x_i;\{a_m\},\{b_n\}))^2$

We uniformly sample 1000 points $x_i$ from the interval $[-3,3]$. $\{a_m\},\{b_n\}$ are randomly initialized. In practice, we solve this using the Levenberg-Marquardt algorithm, available in the MINPACK package.

In the end, this is easily done by calling scipy.optimize.curve_fit. We have added and discription in the revised manuscript Appendix G.

>>> Q2 Ablation Study

>>> A2 We sincerely appreciate the feedback. As suggested, we conduct a new ablation study on the three factors presented in the paper, including 1) rational base function 2) group-wise weight sharing and 3) proper initialization.

We conducted experiments using the KAT-Tiny model on the ImageNet dataset, training for 300 epochs

• Ablation 1: Base function: We replace rational base function by B-spline used in KAN. Because no pure CUDA implementation for B-spline, we implement in torch with De Boor-Cox algorithm.

  The choice of the base function has a minor impact on performance, but significantly affects runtime. According to Exp 2 and Exp 5, the rational function slightly outperforms the B-spline in terms of performance. While the difference in MAC count is negligible, the B-spline implementation in torch runs significantly slower than our pure CUDA implementation.

• Ablation 2: Group-wise computation: In this experiment, group-wise weight sharing is replaced with distinct parameters for each channel.

  According to Exp 4 and Exp 6, group-wise computation plays a crucial role in efficiency. Sharing parameters within groups slightly reduces accuracy from 74.8 to 74.6. However, it significantly reduces training time from 38 hours to 12 hours.

• Ablation 3: Initialization: In this experiment, the variance preserving initialization is replaced by torch default initialization.

As shown in Exp 6 and Exp 5, initialization is critical for good performance. Without proper initialization, terms of different orders were initialized at similar scales, causing instability. This issue was more important for rational functions than for B-splines, as evidenced by the performance differences between Exp 2 and Exp 5.

>>> Q3 Comparing to (Boulle et al.)

>>> A3 Thanks the reviewer for the question. In fact, we have compared with (Boulle et al.) in our paper, since (Boulle et al.) uses exactly the Padé Activation Unit (PAU). This is shown in Table 6 as part of the ablation study, where KAT outperforms PAU.

>>> Q4 Gradient Calculation (Monila et al.)

>>> A4 We completely agree. The derivations are indeed the same. As suggested, we have moved Equation 10 to the Appendix F in our revised version.

We sincerely thank Reviewer yXCb for valuable comments and suggestions. We have carefully incorporated them to improve our paper.

>>> Q1 Results on more tasks

>>> A1 We truly appreciate the suggestions. As suggested, we add experiment on solving PDE and NLP task.

• GR-KAN for PDE: For PDE solving, we resort to a one-dimensional damped harmonic oscillator. It is governed by the differential equation that

  $m \frac{{d^2 u}}{{dt^2}} + \mu \frac{{du}}{{dt}} + k u = 0$ where $m$ is the mass. $\mu$ is the damping coefficient. $k$ is the stiffness constant. With the initial condition $u(0) = 1,u'(0) = 0,m = 1$. The exact solution is $u(t) = e^{-d \cdot t} \cdot (A \cos(\omega t + \phi))$, where $d = \mu / 2 w_0 = \sqrt{k}, w = \sqrt{w_0^2 - d^2}, \phi = \arctan(-d / w)$. We solved this problem using MLP, KAN, and GR-KAN with a network architecture of $[2 \to 5\to 1]$. KAN achieved the best accuracy but trained slowly, MLP performed the worst, and GR-KAN trained faster than KAN but performed slightly worse in this experiment.

  Model	L2 Error	Train Time
  $w_0=10$
  MLP(GELU)	2.0216e-04	~1min
  GR-KAN	6.3909e-06	~4min
  KAN	1.6125e-08	~20min
  $w_0=50$
  MLP(GELU)	1.1805e-01	~1min
  GR-KAN	5.2515e-02	~4min
  KAN	3.7762e-02	20min

  The results have been incorporated in the revised Appendix B.

• KAT for NLP: Due to the tight schedule during rebuttal, we only do some preliminary experiments. We fine-tuned ALBERT-base on two datasets from GLUE: the Stanford Sentiment Treebank (SST-2) and Multi-Genre NLI (MNLI). We replaced the MLP with the GR-KAN designed in the paper. We inherited all weights from pretrained model as described in Section 4.4.

  The accuracy is reported below. We achieved better performance. We will add full GLUE results once they are completed.

  Model	SST-2	MNLI
  ALBERT-Base	90.3	81.6
  ALBERT-Base+KAT	91.6	83.2

>>> Q2 Ablation Study on the base function

>>> A2 Thank the reviewer for the question. In fact, we have conducted an ablation study. ViT + KAN means that we use B-splines same as the KAN paper and compare with our KAT, as shown in Figure 5. However, this experiment ablated all components together.

Based on the suggestion, we ran a more focused study. We re-implemented the method using radiance base functions (RBF) and B-splines, but keeping group-wise computation and initialization consistent.

Note that this PyTorch implementation is still slower than our optimized CUDA version. Our Rational function delivers the best performance.

>>> Q3 Computation Correction

>>> A3We sincerely thank Reviewer yXCb for the thorough proofreading. As suggested, we have revised both Table 1 and the Appendix I as suggested.

>>> Q4 Multivariate and Univariate

>>> A4 Sorry for the confusion might caused. We mean "learn a univariate function on each edge". Altogether, summation of all univariate functions forms a multivariate function. This terminology is consistent with the KAN paper. We have revised the text in Line 146 to make it clearer.

>>> Q5 Reason for Shared Coefficient

>>> A5 This decision is based on empirical observations. Using different coefficients for different groups, especially for denominators, slightly reduces performance. We attribute this to the sensitivity of division operations.

Hypothesis: Denominator sensitivity. In rational functions, small changes in the denominator can cause significant output variations. Additionally, varying embedding scales across groups can introduce inconsistencies when separate coefficients are used, negatively affecting performance.

We have added a formal discussion in Appendix E to clarify this.

It is a great honour to hear the feedback from Reviewer h5qV. We have answered the questions below and revised the paper accordingly.

>>> Q1 Novelty

>>> A1 Thanks for the question. The reviewer is right; simply replacing the MLP with KAN is straight-forward. However, we believe that our key contributions lie in offering a framework-wise solution for integrating KAN into transformers to make this replacement scalable and practical. Additionally, we consider the analysis and techniques we developed to be of considerably new.

• Framework Novelty. We are the first to successfully integrate KAN into a transformer and scale it up. While the replacement is straightforward, no prior work has succeeded to do this.
• Analysis Novelty. Our work introduces the first quantitative analysis of KAN's scalability challenges, as detailed in Section 3.
• Techniqueal Novelty: Our core technical novelty lies in the development of a group-wise computation for activation function. To the best of our knowledge, it has not been introduced before. Additionally, the use of rational functions in KAN and our tailored initialization method are innovative, though partially inspired by earlier studies.

>>> Q2 Typo

>>> A2 Thank the R-h5qV for the proof-reading. We have edit the $\phi_{in}$ matrix in Equation 2.

>>> Q3 Fewer or More parameters

>>> A3 Apologies for the confusion. These two statements are not contradictory; rather, they apply to different contexts, depending on whether we fix the target function or network width.

• Fit the same function (Theoretical): As mentioned in Line 46, when the target function is fixed, KANs theoretically need fewer parameters to represent it.
• Network width (Practical): In Section 3, we explain that when the network width is fixed, KANs require more parameters than standard MLPs.

Thus, these points are not opposing but are valid in different contexts.

>>> Q4 Variance for higher order B-splines

>>> A4 Thanks for the in-depth question. In Line 197,Higher-order splines exacerbate variance instability means that increasing spline order will excessively smooth out the input signal. It reduces function variation, causing $Var[\phi(x)]$ to become smaller.

• High-order B-spline is smoothing. The B-spline is defined recursively as:

shows that each higher-order basis function $B_{i, p}(t)$ a weighted summation of a weighted average of two lower-order basis functions $B_{i, p-1}(t)$ and $B_{i+1, p-1}(t)$. as the order $p$ increases, the B-spline becomes wider and smoother.

In the extreme case, as $p\to \infty$, the smoothing effect causes the basis functions to become nearly uniform across the domain, i.e., $B_{i, p}(t)\approx C$, a constant. Consequently, as, when $p\to \infty$, the variance of the output converges to zero, $Var[\phi(x)]\to 0$. This extreme smoothing leads to instability in the activation variance, as it effectively flattens all variations.

We have included this explanation in Appendix D for further clarification.

>>> Q5 Basic Novelty

>>> A5 We appreciate the comment. We humbly believe our novelty extends beyond simply replacing B-splines with rational functions. It lies in the complete set of new analysis and the solutions, all introduced for the first time in the KAN context:

1. Rational function. We analyze the limitations of B-splines and replace them with rational functions, including a new CUDA implementation.
2. Group Computation. We identify inefficiencies in KAN, where each edge requires distinct parameters and computations. To address this, we share parameters within edge groups, reducing computational costs.
3. Initialization. We address the instability of previous initialization methods for KAN, proposing a new approach that supports the training of larger models.

These contributions are interdependent and essential. Without (1), the implementation would be slow; without (2), parameter size would be excessive; and without (3), training would not converge.

>>> Q6 Contradiction of arguments

>>> A6 We truly thank you for the question. Please see A3.