FlashKAT: 解读和解决Kolmogorov-Arnold变换器中的性能瓶颈
FlashKAT: Understanding and Addressing Performance Bottlenecks in the Kolmogorov-Arnold Transformer
AI总结：
FlashKAT通过优化内存访问和梯度计算，显著提升了KAT的训练速度，解决了其性能瓶颈问题。
摘要：
Kolmogorov-Arnold网络（KAN）因其更高的表达能力和可解释性，正逐渐成为多层感知机（MLP）的替代方案。然而，KAN由于计算成本的增加，导致训练不稳定且速度慢了多个数量级，限制了其在大规模任务中的应用。最近，Kolmogorov-Arnold变换器（KAT）被提出，通过利用分组理性KAN（GR-KAN）实现了与传统Transformer模型相当的FLOPs。不幸的是，尽管FLOPs相当，我们的测试显示KAT在训练过程中仍然慢了123倍，表明存在其他性能瓶颈，而不仅仅是FLOPs。在本文中，我们进行了一系列实验以了解KAT减速的根本原因。我们发现，减速可以归因于内存 stalls，更具体地说是GR-KAN反向传播过程中梯度累积的低效。为了解决这个内存瓶颈，我们提出了FlashKAT，通过重构内核来最小化对慢速内存的访问和原子加法的使用。评估显示，FlashKAT在最先进的KAT上实现了高达86.5倍的训练加速，同时减少了梯度计算中的舍入误差。
英文摘要：
The Kolmogorov-Arnold Network (KAN) has been gaining popularity as an alternative to the multilayer perceptron (MLP) due to its greater expressiveness and interpretability. Even so, KAN suffers from training instability and being orders of magnitude slower due to its increased computational cost, limiting its applicability to large-scale tasks. Recently, the Kolmogorov-Arnold Transformer (KAT) has been proposed, achieving FLOPs comparable to traditional Transformer models with MLPs by leveraging Group-Rational KAN (GR-KAN). Unfortunately, despite the comparable FLOPs, our testing shows that KAT remains 123x slower during training, indicating that there are other performance bottlenecks beyond FLOPs. In this paper, we conduct a series of experiments to understand the root cause of the slowdown in KAT. We uncover that the slowdown can be isolated to memory stalls, linked more specifically to inefficient gradient accumulations in the backward pass of GR-KAN. To address this memory bottleneck, we propose FlashKAT, which minimizes accesses to slow memory and the usage of atomic adds through a restructured kernel. Evaluations show that FlashKAT achieves up to an 86.5x training speedup over state-of-the-art KAT while reducing rounding errors in gradient computation.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
