FlashKAT：理解并解决 Kolmogorov-Arnold Transformer 中的性能瓶颈
机器学习
2026-02-10 v3 计算机视觉与模式识别
摘要
Kolmogorov-Arnold 网络（KAN）因其更强的表达能力和可解释性，作为多层感知器（MLP）的替代方案日益受到欢迎。尽管如此，KAN 仍存在训练不稳定和计算成本增加导致的速度慢数个数量级的问题，这限制了其在大规模任务中的应用。最近，Kolmogorov-Arnold Transformer（KAT）被提出，通过利用 Group-Rational KAN（GR-KAN）实现了与传统 Transformer 模型（使用 MLP）相当的 FLOPs。不幸的是，尽管 FLOPs 相当，我们的测试表明 KAT 在训练期间仍然慢 123 倍，这表明除了 FLOPs 之外还存在其他性能瓶颈。在本文中，我们进行了一系列实验来理解 KAT 减速的根本原因。我们发现减速可以隔离到内存停顿，更具体地说，是与 GR-KAN 反向传播中低效的梯度累积有关。为了解决这一内存瓶颈，我们提出了 FlashKAT，它通过重构内核来减少对慢速内存的访问以及原子加法的使用。评估结果显示，FlashKAT 相比最先进的 KAT 实现了高达 86.5 倍的训练加速，同时减少了梯度计算中的舍入误差。
引用
@article{arxiv.2505.13813,
title = {FlashKAT: Understanding and Addressing Performance Bottlenecks in the Kolmogorov-Arnold Transformer},
author = {Matthew Raffel and Lizhong Chen},
journal= {arXiv preprint arXiv:2505.13813},
year = {2026}
}
备注
AAAI 2026