[論文レビュー] ReLU-KAN: New Kolmogorov-Arnold Networks that Only Need Matrix Addition, Dot Multiplication, and ReLU

Qi Qiu, Tao Zhu+3 more|arXiv (Cornell University)|Jun 4, 2024

Parallel Computing and Optimization Techniques被引用数 13

ひとこと要約

ReLU-KAN は KAN の B-spline 基底を ReLU ベースの基底に置き換え、全行列演算を可能にして、GPU の大規模速度向上、適合性の安定性の向上、そして崩壊的忘却耐性の維持を実現します。

ABSTRACT

Limited by the complexity of basis function (B-spline) calculations, Kolmogorov-Arnold Networks (KAN) suffer from restricted parallel computing capability on GPUs. This paper proposes a novel ReLU-KAN implementation that inherits the core idea of KAN. By adopting ReLU (Rectified Linear Unit) and point-wise multiplication, we simplify the design of KAN's basis function and optimize the computation process for efficient CUDA computing. The proposed ReLU-KAN architecture can be readily implemented on existing deep learning frameworks (e.g., PyTorch) for both inference and training. Experimental results demonstrate that ReLU-KAN achieves a 20x speedup compared to traditional KAN with 4-layer networks. Furthermore, ReLU-KAN exhibits a more stable training process with superior fitting ability while preserving the "catastrophic forgetting avoidance" property of KAN. You can get the code in https://github.com/quiqi/relu_kan

研究の動機と目的

• Kolmogorov-Arnold Network (KAN) の基底関数を簡素化して、より高速でGPU 対応の実装を動機付ける。
• PyTorch のようなフレームワークとの統合を容易にする、行列ベース計算を可能にする ReLU ベースの基底を開発する。
• ReLU-KAN がトレーニングを速くし、適合精度を向上させつつ、崩壊的忘却を回避する KAN の特性を維持することを示す。

提案手法

• KAN の B-splines の代替として、簡略化された基底関数 R_i(x) = [ReLU(e_i − x) × ReLU(x − s_i)]^2 × 16/(e_i − s_i)^4 を導入する。
• 全体の基底計算を行列演算として表現し、GPU の並列性を高める。
• 計算を加速するために訓練不可のパラメータを事前生成し、位置符号化に類似した手法を採用する。
• 基底関数の加重和を畳み込み演算として表現し、標準的な深層学習フレームワークに適合させる。
• ReLU-KAN レイヤの PyTorch 実装を、30 行未満のコードで提供する。
• R_i 基底評価が出力に使用される行列 F を形成するような、レイヤー単位の計算パイプラインを導出する。

実験結果

リサーチクエスチョン

• RQ1 KAN の B-spline 基底を ReLU ベースの基底に置換すると GPU でのトレーニング速度が向上するか。
• RQ2 ReLU-KAN は 一変量・多変量関数に対して KAN の適合精度と安定性を維持または向上できるか。
• RQ3 ReLU-KAN は 崩壊的忘却耐性を維持しつつ、スケーラブルなネットワークアーキテクチャを実現できるか。

主な発見

• ReLU-KAN は 5 〜 20 倍のトレーニング速度向上を、単層から三層モデルにかけて達成する。
• より大きなネットワークでは、GPU による speedup が顕著になり、最大で約 20 倍の速度向上が観測される。
• ReLU-KAN は KAN より高い適合精度を達成し、報告された比較では約 2 桁の差で高い精度を示す。
• 評価対象の関数全体で、特に高周波ターゲットに対して、ReLU-KAN の学習は KAN より安定した収束を示す。
• ReLU-KAN は 実験において KAN の崩壊的忘却回避性を保持する。