Kolmogorov-Arnold Networks를 통한 FPGA 상의 초고속 머신러닝

Field-Programmable Gate Arrays (FPGAs)에 구현된 Kolmogorov-Arnold Networks (KANs)는 마이크로초 미만의 지연 시간으로 머신러닝 추론 및 온라인 학습을 가능하게 합니다. Multi-Layer Perceptrons (MLPs)의 고정된 활성화 함수를 하드웨어 lookup tables (LUTs)에 매핑된 학습 가능한 일변수 함수로 대체함으로써, 이 접근 방식은 이전 KAN-FPGA 구현보다 2700배 빠른 속도를 달성하며 지연 시간 및 자원 효율성 측면에서 최첨단 FPGA 가속기를 능가합니다.

학습 가능한 Lookup-Table 신경망으로서의 KANs

FPGAs에 KANs를 구현하면 신경망을 프로세서의 명령 시퀀스가 아닌 디지털 로직으로 직접 실행할 수 있습니다. 이는 CPU 및 GPU에 내재된 스케줄링 및 메모리 액세스 오버헤드를 제거하며, 이는 나노초 단위의 지연 시간이 필요한 애플리케이션에 매우 중요합니다.

LUT 기반 접근 방식

표준 MLP에서는 학습 가능한 가중치는 스칼라이며 활성화 함수는 고정되어 있습니다. KANs는 네트워크의 엣지에 학습 가능한 일변수 함수를 사용하여 이를 뒤집습니다. 이러한 함수는 모든 이진 입력 조합에 대한 출력 값을 저장하는 lookup tables (LUTs)를 사용하여 FPGAs에 표현됩니다.

다변수 함수를 LUTs로 표현하는 것은 지수적 스케일링으로 인해 계산적으로 비실용적이지만, KANs는 일변수 활성화 함수를 합산합니다. 이 설계는 다음과 같은 이점을 제공합니다:

• Linear Scaling: 입력 차원에 따라 자원 사용량이 지수적이 아닌 선형적으로 증가합니다.
• Efficient Pruning: 0에 가까운 활성화 함수를 제거함으로써 중요하지 않은 네트워크 구성 요소를 제거할 수 있으며, 이를 통해 필요한 LUTs 및 adders의 수를 줄일 수 있습니다.
• Full Domain Coverage: KAN 활성화 함수는 작고 유한한 도메인에 대해 정의되므로, 양자화 과정에서 전체 입력 범위를 커버할 수 있습니다.

추론 구현

고정 모델 추론을 위해, KANs는 소프트웨어(예: PyTorch)에서 학습된 후 FPGAs에 배포됩니다. 이 아키텍처는 활성화 함수가 LUTs를 통해 병렬로 계산되고 adder tree를 사용하여 합산되는 고정 소수점 양자화 방식을 사용합니다. 이 프로세스는 순방향 패스(forward pass)의 매우 높은 속도 실행을 가능하게 합니다.

실시간 FPGA 상의 온라인 학습

온라인 학습은 새로운 데이터가가 도착할 때 모델이 실시간으로 동작을 조정할 수 있게 하며, 이는 핵융합이나 양자 제어와 같은 비정적 시스템에 필수적입니다. 순방향 패스, 역방향 패스(backward pass), 그리고 그래디언트 업데이트를 FPGA 상에 병렬 회로로 직접 구현함으로써, 시스템은 마이크로초 미만의 시간 단위로 모델 파라미터를 업데이트할 수 있습니다.

B-Spline Locality 및 Sparse Updates

온칩(on-chip) 학습을 기능을 위해, 시스템은 최종 학습된 활성화 함수 대신 B-spline 기저 함수를 LUTs에 저장합니다. 이는 계수(coefficients)가 학습 과정 중에 동적으로 업데이트되기 때문입니다.

B-splines는 locality라는 중요한 이점을 제공합니다: 임의의 입력에 대해 기저 함수(basis functions)의 아주 작은 부분 집합(spline order $k$에 의해 결정됨)만이 0이 아닌 값을 가집니다. 결과적으로, 순방향 및 역방향 패스의 하드웨어 로직은 전체 그리드 셀(grid cells) $G$의 수가 아닌 $k$에 따라 스케일링됩니다. 이를 통해 모델은 하드웨어 자원 사용량을 비례적으로 늘리지 않고도 그리드 크기를 늘려 표현력을 높일 수 있습니다.

Bounded Activations를 통한 안정성

고정 소수점 양자화는 MLP의 그래디언트 및 가중치의 넓은 범위의 크기에 대응하는 데 어려움을 겪는 경우가 많습니다. KANs는 B-splines가 유계(bounded)이기 때문에 이 문제를 완해합니다. 임의의 입력 $x$에 대해, 출력은 항상 가장 작은 계수와 가장 큰 계수 사이로 제한됩니다: $\min(c_i) \le \sum B_i(x)c_i \le \max(c_i)$.

이러한 유계성은 활성화 함수와 그래디언트 모두 예측 가능한 범위 내에 머물게 하여, 양자화 오차를를 줄이고 하드웨어에서의 그래디언트 기반 학습의 안정성을 향상시킵니다.

성능 및 하드웨어 스케일링

KAN-based 온라인 학습기는 50,000개 이상의 파라미터를 유지하면서 마이크로초 미만의 지연 시간을 유지할 수 있습니다. MLP와 비교했을 때, KANs는 그리드 크기를 확장할 때 자원 사용량이 거의 일정하게 유지되며, 큐비트 판독, 비정상 제어, 함수 근사화와 같은 벤치마크에서 우수한 수렴성을 보여줍니다.

커뮤니티 인사이트 및 제약 사항

지연 시간 측면에서의 성능 이득은 상당하지만, 이 아키텍처는 LLM과 같은 고처리량, 대규모 모델을 위한 것이 아닙니다. 기술 관찰자들 사이의 논의는 처리량(throughput)보다 지연 시간(latency)에 초점을 맞추고 있음을 강조합니다:

"I've been trying to hit 100,000tokens/s with a 3.28m dumb model, and even this is an order of magnitude too large to benefit. It appears to be focussed more on latency, than throughput."

또한, 이 접근 방식의의 실용적인 유용성은 마이크로초 미만의 응답 시간이 필수적인 특수 목적의 소형 모델 작업에서 가장 두드러집니다.