TruKAN: Towards More Efficient Kolmogorov-Arnold Networks Using Truncated Power Functions

기술 요약: "TruKAN: 절단된 멱함수(Truncated Power Functions)를 이용한 더욱 효율적인 Kolmogorov-Arnold 네트워크를 향하여"

문제 정의

Kolmogorov-Arnold Networks (KANs)는 다층 퍼셉트론(MLP)의 유망한 대안으로 부상하며, 뛰어난 해석 가능성과 단변량 에지별 함수(univariate edge-wise functions) 학습 능력을 제공한다. 그러나 표준 KAN은 학습 가능한 활성화 함수로 B-스플라인(B-spline) 기저 함수에 의존한다. 이러한 B-스플라인을 평가하려면 재귀적인 de Boor–Cox 알고리 asorithm이 필요한데, 이는 계산 비용이 많이 들어 KAN의 학습 속도를 동등한 MLP보다 최대 10배까지 느리게 만드는 원인이 된다. 또한, 복잡한 스플라인 알고리즘과 고정된 그리드 구조에 대한 의존성은 고차원 문제 및 심층 아키텍처에서 확장성 문제를 야기할 수 있으며, 이는 노이즈가 섞인 그래디언트와 불안정한 최적화로 나타난다. SineKAN 및 FourierKAN과 같은 최근의 변형 모델들이 이러한 문제를 해결하려고 시도했으나, 이들은 해석 가능성 측면에서 새로운 트레이드오프를 도입하거나 MNIST와 같은 단순한 합성 데이터셋에 국한되는 한계를 보인다.

방법론

저자들은 canonical KAN 토폴로지와 학습 가능한 활성화 함수를 유지하면서, B-스플라인 기저를 $k$-차 스플라인 이론에서 유도된 절단된 멱함수(truncated power functions) 계열로 대체하는 새로운 아키텍처인 TruKAN을 제안한다.

핵심 아키텍처 구성 요소:

1. 절단된 멱함수 기저 (Truncated Power Basis): B-스플라인 대신, TruKAN은 $(x - t_j)^k_+$ 형태의 절단된 멱함수를 활용한다. 이론적으로 이 함수들이 생성하는 공간은 B-스플라인의 공간과 동일하지만, 그 계산 방식은 다르다. B-스플라인은 컴팩트 서포트(compact support)를 가져 재귀를 통해 구간당 $O(k^2)$의 복잡도를 요구하는 반면, 절단된 멱함수는 전역적 서포트(global support)를 가진다. 그러나 저자들은 실제 적용 시 고해상도 그리드가 항상 필수적인 것은 아니므로, 적은 수의 매듭(knot)만으로도 TruKAN이 효율적으로 작동할 수 있다고 주장한다.
2. 하이브리드 레이어 구조: 각 TruKAN 레이어는 다음 두 가지 구성 요소를 결的に 결합한다:
   • 국소적인 조각별 비선형성(piecewise non-linearities)을 포착하기 위한 절단된 멱함수 성분.
   • 표준 KAN에서 사용되는 보조 SiLU 활성화를 대체하기 위한 다항식 성분(전역적 추세). 이 다항식 항은 매끄럽고 미분 가능한 전역 매핑을 제공하여, 명시적인 전역 비선형성(SiLU 등)의 부재를 보완한다.
3. 매듭 구성 (Knot Configurations): 본 논문은 매듭 배치와 관련하여 두 가지 변형을 소개한다:
   • 공유 매듭 (Shared Knots): 모든 출력 채널이 하나의 매듭 위치 세트를 공유한다.
   • 개별 매듭 (Individual Knots): 각 출력 채널이 자신만의 학습 가능한 매듭 위치를 가진다.
   • 고정된(사전 정의된 구간) 방식과 학습 가능한(순서 및 양수 조건을 제약 조건으로 하여 훈련 중 최적화되는) 매듭 전략을 모두 탐구한다.
4. 훈련 프레임워크:
   • 최적화: 절단된 멱함수를 사용하는 심층 네트워크 훈련의 불안정성(지수적 성장/감소로 인한 그래디언트 폭주 문제)을 해결하기 위해, 저자들은 기존 KAN 연구에서 사용되던 L-BFGS 대신 LookAhead로 감싸진 AdamW를 사용하는 하이브리드 최적화 전략을 채택한다. 이는 확률적 미니 배치 훈련에 더 적합하다.
   • 정규화: 분류기 헤드 내에서의 레이어 정규화(Layer Normalization, 그리고 배치 정규화)의 영향을 조사한다. 기존 KAN에서는 드롭아웃(dropout)과 같은 전통적인 정규화 기법이 혼재된 결과를 보였기 때문이다.
   • 백본 (Backbone): 모델은 표준 선형 분류기 헤드를 MLP, KAN, SineKAN 또는 TruKAN 레이어로 교체한 EfficientNet-V2 프레임워크에 통합된다.

주요 기여

• 아키텍처 혁신: B-스플라인을 절단된 멱함수로 교체하고 다항식 전역 항을 결합함으로써, KAN의 표현력을 유지하면서도 기저 함수 평가를 단순화하는 모델을 구축했다.
• 해석 가능성 향 향상: 절단된 멱함수를 사용함으로써 각 매듭의 기여도가 명시적이다. 매듭 $t_j$에서의 비제로(non-zero) 계수는 해당 지점에서의 의도적인 $k$차 곡률을 직접적으로 나타내며, 이는 B-스플라인의 중첩된 서포트 방식보다 더 높은 투명성을 제공한다.
• 확장성 및 효율성: TruKAN이 이전 KAN 연구의 제한된 합성 데이터 환경을 넘어, 심층 아키텍처와 복잡한 비전 태스크로 효과적으로 확장될 수 있음을 입증했다.
• 포괄적인 벤치마킹: 저자들은 파라미터 수를 일치시켜 공정한 비교를 보장하기 위해, 네 가지 컴퓨터 비전 데이터셋(CIFAR-10, CIFAR-100, Oxford-IIIT Pets, STL-10)을 사용하여 소규모 및 심층 네트워크 구성 모두에서 엄격한 평가를 수행했다.

실험 결과

평가는 다양한 분류기 크기(소형 및 대형 아키텍처)를 가진 EfficientNet-V2 백본을 사용하여 네 가지 데이터셋에 대해 수행되었다.

• 정확도 및 F1-스코어:
  • 소형 아키텍처: TruKAN 변형 모델들, 특히 레이어 정규화와 공유 매듭(TruKAN-SN)을 적용한 모델이 기준이 되는 KAN 및 SineKAN 모델보다 일관되게 우수하거나 대등한 성능을 보였다. CIFAR-10에서 TruKAN-IN은 가장 높은 정확도(0.8927)를 달성하여 최상의 MLP 베이스라인을 능가했다. STL-10에서 TruGAN-SN은 0.9341의 정확도를 기록하며 정규화된 MLP에 근소한 차이로 뒤처졌다.
  • 대형 아키텍처: TruKAN 모델은 우수한 확장성을 보여주었다. CIFAR-10에서 TruKAN-IN은 더 적은 파라미터를 사용했음에도 불구하고 MLP-N 베이스라인을 능가하는 0.8932의 정확도를 달성했다. TruKAN 변형 모델들은 모든 데이터셋에서 KAN 및 SineKAN 베이스라인을 지속적으로 압도했으며, 종종 가장 강력한 MLP와 대등하거나 더 나은 결과를 냈다.
  • 정규화의 영향: 레이어 정규화는 TruKAN 성능에 결정적이었으며, 특히 개별 매듭 구성에서 과적합을 방지하고 일반화 성능을 크게 향상시켰다.
• 계산 효율성:
  • 훈련 시간: 공유 매듭 방식의 TruKAN(TruKAN-FS)은 훈련 단계당 약 0.0024초를 기록하여, 표준 KAN(0.0087s - 0.0099s) 대비 3~4배의 개선을 이루었다.
  • 메모리 사용량: TruKAN-FS는 피크 GPU 메모리를 118.04 MB만 소비하여, KAN-PBT(570.84 MB) 및 KAN-PBF(691.79 MB)보다 현저히 적은 양을 사용했다. 다만, 개별 매듭 변형(TruKAN-FI)은 확장된 매듭 저장 공간으로 인해 높은 메모리 비용(~2.6 GB)이 발생했다.
  • 파라미터: TruKAN 모델은 일반적으로 KAN 및 SineKAN의 대응 모델보다 적은 학습 가능 파라미터로 대등하거나 더 나은 성능을 달성했다.

의의 및 주장

본 논문은 TruKAN이 계산 효율성과 KAN 원칙 준수 사이의 트레이드오프를 성공적으로 해결했다고 주장한다. 절단된 멱함수와 다항식 전역 항을 사용하여, TruKAN은 다음을 달성한다:

1. 우수한 효율성: 표준 KAN에 비해 훈련 시간과 메모리 사용량을 크게 줄여 복잡한 비전 태스크에 적용 가능하게 만들었다.
2. 표현력 유지: KAN의 이론적 근사 능력을 보존하면서도 실제 데이터셋에서의 정확도를 개선했다.
3. 향상된 해석 가능성: 절단된 멱함수의 명시적 특성은 B-스플라인이나 SineKAN과 같은 전역 스펙트럼 방법보다 학습된 에지 함수와 매듭의 기여도를 더 명확하게 검사할 수 있게 한다.

저자들은 결론적으로 TruKAN이 KAN의 실용적인 발전 방향을 제시하며, 이전 연구들의 제한된 설정 이상의 이점을 보여준다고 밝힌다. 이들은 이 방법이 안정성을 개선했음에도 불구하고, 다른 정규화 기법(예: 드롭아웃)의 통합과 효율성을 더욱 높이기 위한 혼합 정밀도(mixed-precision) 훈련의 잠재적 활용에 대한 과제가 남아 있음을 인정한다. 향후 연구로는 이러한 영역을 탐구하고 강화 학습 및 시계열 예측에 TruKAN을 적용하는 것을 제안한다.