[论文解读] Kolmogorov-Arnold Transformer

Xingyi Yang, Xinchao Wang|arXiv (Cornell University)|Sep 16, 2024

Fusion and Plasma Physics Studies被引用 14

一句话总结

本文在视觉变换器中用 Group-Rational Kolmogorov–ArnolD Networks (GR-KAN) 替代 MLP 层，以提升表征能力和效率，实现 ImageNet 规模的训练并超越 ViT/DeiT 基线。

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.

研究动机与目标

• 识别将 KANs 集成到变换器中的可扩展性挑战（基础函数、参数化、初始化）。
• 提出解决方案：有理激活、分组 KAN，以及方差保持初始化。
• 通过在 ViT 类架构中用 GR-KAN 替代 MLP，开发并验证 KAT。
• 在图像分类、目标检测和语义分割任务中展示性能提升。

提出的方法

• 在 KAN 中采用有理函数作为基础激活，并实现基于 CUDA 的梯度以提升效率。
• 引入 Group-Rational KAN (GR-KAN)，跨边缘组共享基础函数以减少参数和计算量。
• 对多项式求值应用 Horner 法以加速 CUDA 执行。
• 使用方差保持初始化以在 GR-KAN 层之间稳定训练。
• 实现从预训练 ViT 的权重迁移，使 KAT 能加载 ViT 权重并进行微调。
• 在 ImageNet-1K、COCO（Mask R-CNN with ViTDet）和 ADE20K（UperNet）上评估 KAT，以证明可扩展性和性能提升。

实验结果

研究问题

• RQ1 在 ImageNet 规模下，GR-KAN 能否在不牺牲收敛性或性能的前提下替代视觉变换器中的 MLP？
• RQ2 与 B-spline KAN 相比，使用具组共享参数的有理激活函数是否在计算效率和精度上有所提升？
• RQ3 在相似计算量下，KAT 在标准视觉任务（分类、检测、分割）上的表现相对于 ViT/DeiT 基线如何？
• RQ4 从 ViT 到 KAT 的预训练迁移在最终精度方面的影响是什么？
• RQ5 关于激活选择和初始化对 KAT 性能影响的消融研究揭示了哪些？

主要发现

模型	通道混合器	#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

• 在 ImageNet-1K 上，KAT 变体在相似 FLOPs 和参数预算下持续超越基于 MLP 的变换器。
• KAT-T 实现 74.6% top-1（ViT-Ti/16 尺度），并在预训练迁移下达到 75.7%，超越 ViT/DeiT 基线。
• KAT-S 在 ImageNet-1K 上无预训练时达到 81.2% top-1，带预训练时为 82.0%，超越 DeiT-S 约 2.4%。
• KAT-B 实现 82.3% top-1，若从 ViT 初始化为 82.8%，超过 ViT-B 与 DeiT-B 基线。
• ViT + KAN 在未采用所提出的可扩展性修复时，无法在 ImageNet 规模训练中收敛，凸显 GR-KAN 设计的必要性（S1-S3）。
• 在检测与分割方面，KAT 主干相较于 ViTDet 与 DeiT 风格主干持续带来收益，且小模型的相对提升更大。