前言
论文地址:https://arxiv.org/abs/2409.10594
源码地址:https://github.com/Adamdad/kat.git
传统的变压器模型使用多层感知器(MLP)来混合通道间的信息,而本文则使用了科尔莫哥罗德网络(KAN),从而提高了模型的表现力和性能。
KAT 在大规模图像识别任务以及物体检测和语义分割等视觉任务中表现尤为出色;KAN 擅长高效逼近数学函数,理论上有可能以比 MLP 更少的参数为复杂函数建模。然而,将 KAN 集成到变换器中面临着一些技术挑战。
三个具体挑战是
- 基函数问题:KAN 中使用的标准 B-样条函数未针对现代 GPU 进行优化,难以进行并行计算,从而导致计算速度缓慢的问题。
- 参数和计算效率低下:面临的挑战是,KANs 的计算成本非常高,因为它们需要为每对输入输出使用单独的函数。
- 权重初始化问题:与 MLP 不同,KAN 权重的初始化有一个可学习的激活函数,这就要求特别小心地进行初始化,以便收敛。
为了克服这些挑战,KAT 推出了三种解决方案
- 有理基函数:使用有理函数代替 B 样条函数,计算效率高,适合现代 GPU。
- 组 KAN:每组神经元共享激活权重,以减少计算负荷,同时保持性能。
- 方差保持初始化:权重的初始化可保持各层激活的方差,从而确保稳定的学习。
因此,KAT 比传统的基于 MLP 的变压器具有更好的性能。
模型架构
本文提出的 Kolmogorov-Arnold 变换器(KAT)用 Kolmogorov-Arnold 网络(KAN)取代了传统变换器中使用的 MLP(多层感知器)。它是一种新的架构,通过引入以下功能提高了模型的表达能力和性能
KAT 的一个主要特点是采用了多项创新设计,将 KAN 有效地集成到变换器中。具体来说,为了提高 KAN 层的计算效率,它使用了有理函数,而不是传统的 B-样条函数,后者是在 CUDA 上实现的。这就提高了 GPU 的计算速度,使训练更复杂函数的速度可与传统 MLP 相媲美。
此外,为了减少 KAN 层的计算负荷,还采用了 "组 KAN "方法,即多个边缘共享激活函数的权重。这提高了模型的可扩展性,即使是大型模型也能高效运行。此外,权重的初始化设计还能使各层之间的激活方差保持一致。这种设计提高了训练的稳定性,并能更有效地训练模型。
与传统的变换器模型相比,KAT 实现了更高的准确率,尤其是在 ImageNet-1K 数据集的图像分类任务中,KAT-B 模型以 82.3% 的准确率比 ViT 模型高出 3.1%。这些改进使 KAT 成为一种优于基于 MLP 的简单变换器的新方法。
试验
KAT 针对三大视觉任务(图像分类、物体检测和语义分割)进行了实验,并对每项任务的性能进行了评估。
首先,在图像分类方面,我们使用 ImageNet-1K 数据集来比较 KAT 与其他模型(如 ViT、DeiT 等)的性能。KAT 采用了一种名为 GR-KAN 的新型通道混合器,其性能优于传统的 MLP。例如,KAT-S 模型的准确率达到 81.2%,比传统的 DeiT-S 模型高出 2.4%。此外,KAT 的扩展版本 KAT-B 比 ViT-B 模型的准确率高出约 3.1%,显示出 KAT 在模型大小相同的情况下的优越性。
接下来,在物体检测任务中,使用 MS-COCO2017 数据集将 KAT 纳入 Mask R-CNN,以测量物体检测和实例分割的准确性。在该实验中,KAT 的表现再次优于传统的 ViTDet,尤其是对于较小的模型,APbox 提高了 3.0 个百分点。这证明 KAT 在物体检测方面也能提供高效、准确的结果。
最后,语义分割实验使用 ADE20K 数据集测试了 KAT 的性能。在这项任务中,KAT 被用作 UperNet 的骨干,并与其他传统模型进行了比较:KAT-S 与 DeiT-S 相比,mIoU 提高了约 2.4%,在实现更高精度的同时,由于模型较小,性能损失最小。
这些实验结果证实,与传统的 Transformer 架构相比,KAT 具有更好的表达能力和性能。KAT 的计算效率也特别高,与传统方法相比,通过 CUDA 优化,KAT 的计算速度更快。这样的设计表明,KAT 是各种视觉任务的有力选择。
总结
本文的结论表明,Kolmogorov-Arnold 变换器(KAT)是传统的基于 MLP 变换器的一种有前途的替代方案。(KAT 有效地利用了 Kolmogorov-Arnold 网络(KAN)的特性,在视觉任务中表现出色。KAT 有效地利用了 Kolmogorov-Arnold 网络(KAN)的特性,在视觉任务中表现出色。与传统的变换器架构相比,KAT 在保持计算效率的同时提高了准确性。
此外,KAT 在理论和实践上都有超越 MLP 的潜力,预计在未来的研究中会有更多的应用可能性。特别是其灵活的表达能力和通过使用有理函数实现的学习稳定性,为未来的发展提供了途径,有望扩展到视觉以外的任务。