想大规模训练 KAN？现在是 UKAN！
机器学习
2025-10-10 v4
摘要
Kolmogorov-Arnold 网络（KAN）最近作为传统多层感知器的强大替代方案出现。然而，其对预定义有界网格的依赖限制了其在无界域上逼近函数的能力。为解决这一问题，我们提出了无界 Kolmogorov-Arnold 网络（UKAN），一种消除了传统 KAN 对有界网格需求的方法。该方法的关键创新在于一个系数生成器（CG）模型，它能够即时生成无界对称网格上局部所需的 B 样条系数。UKAN 通过将网格组的位置编码输入 CG 模型，将多层感知器与 KAN 耦合，从而在无需数据归一化的情况下实现无界域上的函数逼近。为降低 UKAN 和 KAN 的计算成本，我们引入了一个 GPU 加速库，将 B 样条评估复杂度降低与网格规模成比例的倍数，通过高效的内存管理实现大规模学习，这与 FlashAttention 和 FlashFFTConv 等近期软件进展一致。性能基准测试确认了我们加速的 KAN（warpKAN）和 UKAN 的优越内存与计算效率，相比原始 KAN 实现了 3-30 倍加速和高达 1000 倍的内存缩减。在回归、分类和生成任务上的实验验证了 UKAN 在匹配或超越 KAN 精度方面的有效性。最后，我们将加速 KAN 和 UKAN 均应用于分子性质预测任务，证明了我们优化实现的大规模端到端训练的可行性。
引用
@article{arxiv.2408.11200,
title = {Want to train KANS at scale? Now UKAN!},
author = {Alireza Moradzadeh and Srimukh Prasad Veccham and Lukasz Wawrzyniak and Miles Macklin and Saee G. Paliwal},
journal= {arXiv preprint arXiv:2408.11200},
year = {2025}
}
备注
16 pages, 5 figures, 8 tables