解决KAT模型训练慢的核心瓶颈，实现86.5倍提速
FlashKAT: Understanding and Addressing Performance Bottlenecks in the Kolmogorov-Arnold Transformer
- 重构反向传播内核，减少内存访问和原子加法操作
- 实测训练速度提升86.5倍，且梯度计算误差降低
- 适合关注KAN类模型高效训练的算法与系统研究者
Kolmogorov-Arnold网络（KAN）因更强表达力和可解释性受到关注，但其训练不稳定且速度慢，计算开销大。近期提出的Kolmogorov-Arnold Transformer（KAT）通过组有理KAN（GR-KAN）实现与传统Transformer相近的浮点运算量（FLOPs），但测试显示其训练仍慢123倍，表明存在超越FLOPs的性能瓶颈。本文通过实验发现，该延迟主要源于内存停顿，尤其是GR-KAN反向传播中梯度累积效率低下。为此，提出FlashKAT，通过重构核函数最小化对慢速内存的访问和原子加法使用。评估显示，FlashKAT相比当前最优KAT实现最高达86.5倍的训练加速，并有效降低梯度计算中的舍入误差。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。