기술 요약: KANLib – 모듈식, 확장 가능하며 빠른 Kolmogorov-Arnold Network 구현체

문제 제기

Kolmogorov-Arnold Networks (KANs)는 고정된 선형 가중치를 학습 가능한 단변량 함수로 대체함으로써 기존의 다층 퍼셉트론(MLP)에 대한 유망한 대안으로 부상했습니다. KAN은 해석 가능성과 표현력 측면에서 이론적 이점을 제공하지만, 실제 도입에는 세 가지 주요 과제가 걸림 быть 있습니다:

1. 높은 계산 비용: 기존 구현체들은 종종 비효율적인 기저 함수(basis function) 평가와 중복된 텐서 연산으로 인해 성능 저하를 겪습니다.
2. 일관되지 않은 기능 지원: 서로 다른 프레임워크(예: PyKAN, EfficientKAN, FastKAN)가 이질적인 기능 세트를 제공하여 엄격한 벤치마킹과 비교를 어렵게 만듭니다.
3. 모듈성 부족: 현재의 구현체들은 특정 아키텍처 설계나 기저 함수에 밀접하게 결합되어 있어, 새로운 변형이나 하이브리드 구조를 탐색하는 데 한계가 있습니다.

방법론

본 논문은 세 가지 주요 KAN 구현체인 PyKAN(Liu 등이 개발한 공식 구현체), EfficientKAN, 그리고 FastKAN의 핵심 혁신을 통합하여 설계된 통합 소프트웨어 프레임워크인 KANLib을 소개합니다.

아키텍처 설계

KANLib은 두 가지 주요 기저 함수 유형을 지원하는 모듈형 아키텍처를 기반으로 구축되었습니다:

• B-splines: 학습 가능한 함수가 B-spline 기저 함수와 잔차 브랜치(일반적으로 SiLU)의 가중 조합으로 구성되는 원래의 KAN 공식을 따릅니다.
• Gaussian Radial Basis Functions (GRBFs): 재귀적인 B-spline 계산을 피하여 계산 복잡도를 줄이기 위해 FastKAN에서 채택된 방식입니다.

주요 기술적 특징

기존 도구들의 파편화 문제를 해결하기 위해, KANLib은 다음 기능들을 단일 PyTorch 환경 내에서 통합합니다:

• 통합 선형 레이어: B-spline과 GRBF 레이어 모두에 대해 일관된 인터페이스를 제공하여 직접적인 아키텍처 비교를 가능하게 합니다.
• 세밀한 제어: 잔차 브랜치($w_b b(x)$) 또는 스플라인 가중치($w_s$)와 같이 표준 KAN 방정식에 정의된 구성 요소를 선택적으로 활성화하거나 비활성화할 수 있습니다.
• 적응형 그리드 재스케일링(Adaptive Grid Rescaling): PyKAN/EfficientKAN에서 발견된 적응형 그리드 메커니즘을 B-spline과 GRBF 레이어 모두에 확장 적용합니다(GRBF의 경우 등간격 그리드로 제한됨).
• 그리드 확장: 학습 중에 학습된 정보를 손실하지 않으면서 스플라인 해상도를 높이는 기능을 지원하며, 이는 EfficientKAN과 FastKAN에서 누락되었던 기능입니다.
• 정규화 전략: 입력값을 기저 함수의 활성 범위로 매핑하기 위해 Layer Normalization을 통합하여, 특정 구성에서 발생하는 적응형 재스케일링의 부재를 보완합니다.

주요 기여

1. 프레임워크 통합: KANLib은 PyKAN, EfficientKAN, FastKAN의 기능 세트를 하나의 확장 가능한 코드베이스로 통합하여, 다양한 기저 함수 및 최적화 전략에 대한 엄격한 벤치마킹을 용이하게 합니다.
2. 계산 효율성: EfficientKAN(예: 입력 텐서 확장 방지)과 FastKAN(예: 직접 GRBF 계산)의 최적화 기법을 통합함으로써, KANLib은 풍부한 기능성을 유지하면서도 경쟁력 있는 추론 속도를 달aff합니다.
3. 아키텍처 유연성: 최소한의 코드 변경만으로 KAN 변형(예: 잔차 브랜치 또는 스플라인 가중치 제거)을 체계적으로 탐색할 수 있게 하여, 표준 공식을 넘어서는 연구를 지원합니다.
4. 재현성: KANLib은 기존 참조 구현체의 예측 동작을 재현하도록 설계되어, 향후 연구를 위한 신뢰할 수 있는 베이스라인 역할을 합니다.

실험 결과

저자들은 California Housing 데이터셋(20,640개의 샘 샘플과 8개의 입력 특성을 가진 회귀 작업)을 통해 KANLib을 PyKAN, EfficientKAN, FastKAN과 비교 평가하였습니다.

예측 성능

• B-spline 변형: KANLib (B-spline)은 B-spline 구현체 중 가장 우수한 성능을 보였으며, RMSE 0.5376 ± 0.0044 및 $R^2$ 0.7852 ± 0.0035를 기록했습니다. 이는 PyKAN (RMSE 0.5434)과 EfficientKAN (RMSE 0.5463)보다 뛰어난 성능으로, KANLib이 기존의 스플라인 기반 동작을 성공적으로 재현하고 약간 개선했음을 입증합니다.
• GRBF 변형: KANLib (GRBF)은 RMSE 0.5471 ± 0.0078을 달성하여 FastKAN (RMSE 0.5506)과 유사한 결과를 보였으며, 이는 본 프레임워크가 Gaussian RBF 기반 KAN을 재현할 수 있음을 확인시켜 줍니다.

계산 효율성

• 추론 시간:
  • PyKAN이 가장 느렸습니다 (212.19 ms).
  • EfficientKAN과 **KANLib (B-spline)**은 PyKAN 대비 약 32.7% 감소한 추론 시간(~142 ms)을 기록했습니다.
  • **KANLib (GRBF)**는 80.58 ms를 기록하여 PyKAN 대비 **62.0%**의 개선을 이루었습니다.
  • FastKAN은 매우 특화되고 유연성이 낮은 구현 덕분에 가장 빨랐습니다 (41.02 ms).
• 파라미터 수: KANLib 모델은 일반적으로 스플라인 계수, 잔차 가중치, 스플라인 가중치를 포함하여 4,050개의 파라미터를 가졌으나, FastKAN은 스플라인 가중치 메커니즘의 생략으로 인해 3,811개로 더 적었습니다.

아키텍처 탐색

레이어 구조를 변화시킨(예: "No Residual", "No Spline Weight", "Plain") 실험 결과:

• 잔차 브랜치나 스플라인 가중치를 제거하면 파라미터 수가 줄어들고(3,510개까지) 추론 시간도 단축되었습니다(GRBF 모델의 경우 최대 약 16.8% 빠름).
• 예측 성능은 이러한 변형 전반에 걸쳐 경쟁력을 유지하였으며, RMSE 편차가 대체로 작게 나타났습니다. 이는 KANLib이 상당한 정확도 손실 없이 유연한 아키텍처 설계를 지원함을 나타냅니다.

의의 및 주장

본 논문은 KANLib을 모든 기존 방법론보다 예측 정확도가 높은 상태(SOTA)를 달성하는 모델이 아니라, 미래 연구를 위한 견고한 토대로 정의합니다. 그 주요 의의는 다음과 같습니다:

• 격차 해소: 파편화된 구현체들을 통합함으로써 KAN의 이론적 잠재력과 실질적이고 확장 가능한 응용 사이의 간극을 메웁니다.
• 확장성 제공: 모듈형 아키텍처를 제공함으로써 연구자들이 새로운 기저 함수, 그리드 전략 및 하이브리드 구조를 테스트하는 데 드는 장벽을 낮춥니다.
• 표준화: KAN의 비교 연구를 저해하는 일관되지 않은 기능 지원 문제를 해결하고, KAN을 평가할 수 있는 일관된 환경을 제공합니다.

저자들은 KANLib이 (특히 적응형 그리드 로직 포함으로 인한 GRBF 계산 측면에서) 매우 특화된 FastKAN에 비해 약간의 계산 오버헤드가 발생할 수 있지만, KAN 아키텍처의 전체 스펙트럼을 탐색하는 데 필요한 유연성을 제공한다고 결론지었습니다. 향후 과제로 Gaussian RBF 계산을 최적화하여 속도 격차를 줄이고, 시계열 센서 데이터를 위한 1D 컨볼루션 지원을 확장하는 것을 제시하였습니다.