[논문 리뷰] ReLU-KAN: New Kolmogorov-Arnold Networks that Only Need Matrix Addition, Dot Multiplication, and ReLU

Qi Qiu, Tao Zhu외 3명|arXiv (Cornell University)|2024. 06. 04.

Parallel Computing and Optimization Techniques인용 수 13

한 줄 요약

ReLU-KAN은 KAN의 B-spline 기저를 ReLU-based 기저로 대체해 전체 행렬 연산을 가능하게 하며, 상당한 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 Networks (KANs)의 더 빠르고 GPU 친화적인 구현을 촉진한다 by simplifying basis functions.
• 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 병렬 처리를 향상시키기 위해 전체 기저 계산을 행렬 연산으로 표현한다.
• 계산 속도를 높이기 위해 훈련 가능하지 않은 매개변수를 미리 생성하는데, 이는 위치 인코딩에 비유된다.
• 기저 함수들의 가중 합을 컨볼루션 연산으로 표현하여 표준 DL 프레임워크에 맞춘다.
• 30줄 미만의 코드로 ReLU-KAN 층의 간결한 PyTorch 구현을 제공한다.
• R_i 기저 평가가 출력에서 사용하는 행렬 F를 형성하는 방법을 보여주는 층별 계산 파이프라인을 도출한다.

실험 결과

연구 질문

• RQ1 KAN의 B-spline 기저를 ReLU 기반 기저로 교체하면 GPU에서의 학습 속도가 개선되는가?
• RQ2 ReLU-KAN이 단변수 및 다변수 함수에 걸쳐 KAN의 적합 정확도와 안정성을 유지하거나 개선할 수 있는가?
• RQ3 ReLU-KAN이 확장 가능한 네트워크 아키텍처를 가능하게 하면서 KAN의 파국적 망각에 대한 저항성을 보존하는가?

주요 결과

• ReLU-KAN은 단일에서 3층 모델에 이르는 학습에서 KAN보다 5배에서 20배 더 빠르다.
• 더 큰 네트워크의 경우 ReLU-KAN의 GPU 속도 향상이 더 두드러져 최대 약 20배의 속도 향상을 관찰한다.
• ReLU-KAN은 KAN보다 더 높은 적합 정확도를 달성하며, 보고된 비교에서 대략 두 자릿수 배의 적합 정확도 차이를 달성한다.
• 평가된 함수들 전반에서 ReLU-KAN의 학습은 KAN보다 더 안정적인 수렴을 보이며, 특히 고주파 타깃에서 그렇다.
• 실험에서 ReLU-KAN은 KAN의 파국적 망각 회피 특성을 유지한다.