Kolmogorov-Arnold Transformer
机器学习
2024-09-18 v1 人工智能
计算机视觉与模式识别
神经与进化计算
摘要
Transformer 是现代深度学习的基石。传统上，这些模型依赖多层感知机（MLP）层来混合通道间的信息。在本文中，我们引入了 Kolmogorov-Arnold Transformer（KAT），这是一种用 Kolmogorov-Arnold Network（KAN）层替代 MLP 层以增强模型表达能力和性能的新架构。然而，将 KAN 集成到 Transformer 中并非易事，尤其是在扩大规模时。具体而言，我们确定了三个关键挑战：（C1）基函数。KAN 中使用的标准 B 样条函数未针对现代硬件上的并行计算进行优化，导致推理速度较慢。（C2）参数与计算低效。KAN 为每个输入-输出对需要一个独特的函数，使得计算量极大。（C3）权重初始化。由于 KAN 具有可学习的激活函数，其权重初始化尤为困难，而这对深度神经网络实现收敛至关重要。为了克服上述挑战，我们提出了三个关键解决方案：（S1）有理基。我们用有理函数替代 B 样条函数，以提高与现代 GPU 的兼容性。通过在 CUDA 中实现这一点，我们实现了更快的计算。（S2）分组 KAN。我们通过一组神经元共享激活权重，在不牺牲性能的情况下减少计算负荷。（S3）保持方差的初始化。我们仔细初始化激活权重，以确保激活方差在层间得以维持。凭借这些设计，KAT 能够有效扩展，并轻松超越传统的基于 MLP 的 Transformer。
引用
@article{arxiv.2409.10594,
title = {Kolmogorov-Arnold Transformer},
author = {Xingyi Yang and Xinchao Wang},
journal= {arXiv preprint arXiv:2409.10594},
year = {2024}
}
备注
Code: https://github.com/Adamdad/kat