Ultrafast machine learning on FPGAs via Kolmogorov-Arnold Networks

개요

Kolmogorov-Arnold Network (KAN) 아키텍처를 활용하여 FPGA 상에서 초고속 추론 및 온라인 학습을 위한 하드웨어 아키텍처를 설계하는 연구를 다룬다.

주요 내용

* FPGA 기반 머신러닝의 필요성: GPU는 높은 처리량을 제공하지만, 나노초 수준의 초저지연(ultra-low latency) 및 고효율이 요구되는 애플리케이션에는 적합하지 않다. FPGA는 재구성 가능한 디지털 로직 장치로서 이러한 요구사항을 충족하는 맞춤형 하드웨어 가속기에 적합하다.
* 고정 소수점 양자화 (Fixed-point quantization): FPGA 및 디지털 장치는 비트를 기반으로 작동하므로, 신경망 연산에 사용되는 실수(real numbers)를 비트열로 인코딩하는 과정이 필요하며, 고정 소수점 양자화는 이를 위한 한 가지 방법이다.
* Lookup-Table 신경망 (LUT-NNs): FPGA는 디지털 로직을 주로 Lookup Table (LUT)을 통해 구현하며, LUT를 신경망의 기본 요소로 학습시키는 LUT-NNs는 FPGA에 적합하다. 그러나 LUT 자체를 직접 학습시키는 것은 어려우므로, 연속적인 함수를 학습시킨 후 양자화하는 방식을 사용한다.
* Kolmogorov-Arnold Networks (KANs): KAN은 MLP의 가중치와 고정된 활성화 함수를 학습 가능한 활성화 함수로 대체하는 아키텍처이다. 각 엣지가 스칼라 가중치 대신 학습 가능한 단변수 함수($\phi_{q,p}(x)$)를 가지며, 이는 B-spline과 같은 함수 기저의 선형 조합으로 매개변수화되어 학습된다. KAN은 모델 압축(pruning) 용이성, 파라미터 효율성, 해석 가능성 측면에서 MLP보다 잠재적 이점을 제공한다.
* KAN의 LUT-NNs 적용: KAN의 각 활성화 함수를 LUT로 표현함으로써, KAN은 학습 가능한 LUT-NNs를 구축하는 원칙적인 방법을 제공한다. KAN의 단변수 활성화 함수 합산 구조는 다변수 함수 LUT의 지수적 스케일링 문제를 피하며, 작은 유한 도메인에서 정의되어 양자화가 용이하다.
* FPGA에서의 KAN 추론 구현: 학습된 KAN 모델을 FPGA에 배포하여 추론을 수행한다. 고정 소수점 양자화를 적용하고, 각 활성화 함수($\phi_{q,p}$)를 LUT로 계산한 후, 가산 트리(adder tree)를 통해 합산한다.
* FPGA에서의 실시간 온라인 학습: 사전 학습된 모델을 배포하는 것을 넘어, FPGA 자체에서 실시간으로 모델을 업데이트하는 온라인 학습을 구현한다. 입력 데이터 스트리밍, 모델 순방향/역방향 전달, 그래디언트 업데이트가 모두 FPGA 상에서 이루어진다.
* 온라인 학습을 위한 KAN 접근 방식: 온라인 학습을 위해, 학습 가능한 계수($c_{q,p,i}$)를 동적으로 업데이트해야 하므로 LUT에는 B-spline 기저 함수($B_i$)를 저장한다.
* B-spline의 희소성 및 안정성: B-spline 기저 함수의 지역성(locality)은 각 입력 값에 대해 소수의 기저 함수만이 활성화됨을 의미하며, 이는 하드웨어 논리 복잡도를 줄인다. 또한, KAN 활성화 함수의 합은 항상 계수들의 최소값과 최대값 사이에 바운드되므로, 예측 가능한 입력 독립적인 범위 내에서 움직여 고정 소수점 양자화 시 그래디언트 업데이트의 안정성을 높이고 양자화 오차를 줄인다.
* 결과: FPGA 기반 KAN 추론은 기존 KAN-FPGA 구현 대비 2700배의 속도 향상을 달성했으며, 최첨단 신경망 FPGA 가속기와 비교해도 지연 시간 및 리소스 사용량 측면에서 경쟁력을 갖는다. KAN 기반 온라인 학습기는 수 마이크로초 미만의 지연 시간으로 순방향 및 역방향 전달을 수행할 수 있으며, MLPs 대비 우수한 하드웨어 확장성 및 수렴 성능을 보인다.

시사점

KAN 아키텍처는 FPGA와 같은 맞춤형 하드웨어 가속기에서 나노초 수준의 초고속 추론 및 실시간 온라인 학습을 가능하게 하는 유망한 경로를 제공하며, 이는 특히 실시간성이 중요한 애플리케이션에서 큰 가치를 가질 수 있다.