Học máy siêu tốc trên FPGA với mạng Kolmogorov-Arnold (KAN)

Học máy siêu tốc trên FPGA với mạng Kolmogorov-Arnold (KAN)

Trong bối cảnh các khối lượng công việc học máy hiện đại chủ yếu phụ thuộc vào GPU để xử lý song song, một câu hỏi thú vị được đặt ra: Làm thế nào để đạt được độ trễ siêu thấp (sub-microsecond) mà các kiến trúc GPU phức tạp không thể đáp ứng? Câu trả lời nằm ở việc sử dụng các bộ tăng tốc phần cứng tùy chỉnh, cụ thể là FPGA (Field-Programmable Gate Arrays), kết hợp với một kiến trúc mạng nơ-ron mới nổi gọi là Kolmogorov-Arnold Networks (KAN).

Bài viết này sẽ đi sâu vào cách thức KANs có thể được triển khai hiệu quả trên FPGA thông qua các bảng tra cứu (LUTs), cho phép thực hiện suy luận (inference) và thậm chí là học trực tuyến (online learning) với tốc độ chưa từng có.

Mô tả minh họa về mạng KAN

Tại sao chọn FPGA cho học máy?

GPU là vua của thông lượng (throughput), lý tưởng cho việc huấn luyện theo lô (batch training). Tuy nhiên, chúng chịu chi phí đáng kể từ việc lập lịch chỉ lệnh và truy cập bộ nhớ động. Ngược lại, FPGA là các thiết bị logic kỹ thuật số có thể cấu hình lại, cực kỳ phù hợp cho các khối lượng công việc chuyên biệt yêu cầu độ trễ cực thấp và hiệu suất phần cứng cao.

FPGA hoạt động dựa trên các thành phần cơ bản như Lookup Tables (LUT - Bảng tra cứu) và Flip-Flops (FF). Thay vì thực hiện các chỉ lệnh tuần tự như CPU/GPU, FPGA triển khai mạng nơ-ron trực tiếp dưới dạng mạch logic kỹ thuật số. Điều này cho phép thiết kế kiến trúc phần cứng và thuật toán đồng bộ (co-design) để đạt được tốc độ xử lý ở mức nano-giây.

Mạng Kolmogorov-Arnold (KAN) và LUT-NNs

Khác với MLP (Multi-Layer Perceptron) truyền thống sử dụng các trọng số cố định trên các cạnh và hàm kích hoạt cố định tại các nút, KANs đặt các hàm kích hoạt có thể học được (learnable activation functions) trên các cạnh. Cụ thể, đầu ra của một lớp KAN được tính bằng tổng các hàm đơn biến:

$$y_q = \sum_{p=1}^{n_{\mathrm{in}}} \phi_{q,p}(x_p)$$

Trong đó $\phi_{q,p}$ là các hàm kích hoạt được học.

Để triển khai điều này trên FPGA, chúng ta sử dụng khái niệm LUT-NNs (Lookup-table neural networks). Vì FPGA hoạt động trên bit chứ không phải số thực, chúng ta cần quá trình lượng tử hóa (quantization) để chuyển đổi các hàm liên tục thành các hàm nhị phân có thể lưu trữ trong LUT.

Điểm mạnh của KANs nằm ở việc chúng tổng hợp các hàm đơn biến (univariate functions). Điều này giúp tránh việc mở rộng theo cấp số nhân số lượng LUTs như khi lưu trữ các hàm đa biến, giúp tiết kiệm tài nguyên phần cứng đáng kể.

Quá trình lượng tử hóa hàm liên tục thành LUT

Suy luận siêu tốc trên FPGA

Trong nghiên cứu của mình (được trình bày tại FPGA 2026 với tên gọi KANELÉ), chúng tôi đã chứng minh rằng KANs là một kiến trúc tự nhiên để xây dựng các LUT-NNs có thể huấn luyện được.

Cách tiếp cận rất trực diện:

1. Huấn luyện mô hình KAN trong phần mềm (ví dụ: PyTorch).
2. Lượng tử hóa các hàm kích hoạt đã học thành các LUT cố định.
3. Triển khai trên FPGA, nơi mỗi hàm kích hoạt $\phi_{q,p}(x_p)$ được tính toán song song bằng cách tra cứu bảng.
4. Tổng các kết quả đầu ra được thực hiện bằng một cây cộng (adder tree).

Kết quả cho thấy khung công việc này không chỉ tương đương mà còn vượt trội so với các bộ tăng tốc mạng nơ-ron FPGA hiện có, đạt được tốc độ nhanh hơn 2700 lần so với các triển khai KAN-FPGA trước đây.

Học trực tuyến thời gian thực trên FPGA

Mặc dù suy luận siêu tốc là ấn tượng, nhưng nhiều ứng dụng như điều khiển lượng tử hoặc phản ứng hạt nhân yêu cầu mô hình phải thích nghi với sự thay đổi của hệ thống ngay lập tức. Đây là bài toán của học trực tuyến (online learning).

Thách thức là phải thực hiện cả quá trình truyền tiến (forward pass), truyền ngược (backward pass) và cập nhật gradient ngay trên FPGA. Để làm được điều này, chúng tôi không lưu trữ toàn bộ hàm kích hoạt đã học vào LUT, mà lưu trữ các hàm cơ sở B-spline (B-spline basis functions).

Kiến trúc ECLAIR cho học trực tuyến trên FPGA

Tại sao KANs lại phù hợp để học trên phần cứng?

Có hai thuộc tính quan trọng của KANs (cụ thể là sử dụng B-splines) khiến chúng trở nên lý tưởng cho việc học trên FPGA:

1. Tính cục bộ (Locality): Với bất kỳ đầu vào nào, chỉ một tập con nhỏ các hàm cơ sở B-spline là khác không. Điều này cho phép các cập nhật gradient trở nên thưa thớt (sparse) và hiệu quả về tài nguyên. Chúng ta có thể mở rộng kích thước lưới (grid size) để tăng khả năng biểu diễn mà không làm tăng đáng kể phần cứng cần thiết.
2. Tính giới hạn (Boundedness): Tổng các hàm cơ sở B-spline luôn bằng 1. Do đó, đầu ra và gradient luôn nằm trong phạm vi dự đoán được, nằm giữa các hệ số nhỏ nhất và lớn nhất. Điều này giúp việc lựa chọn phạm vi lượng tử hóa trở nên dễ dàng hơn nhiều so với MLP, nơi các giá trị trung gian có thể biến động mạnh, dẫn đến việc huấn luyện không ổn định trên phần cứng cố định.

Kết luận

Việc kết hợp mạng Kolmogorov-Arnold Networks với FPGA mở ra một kỷ nguyên mới cho học máy tốc độ cao. Các hàm kích hoạt trong KANs ánh xạ tự nhiên sang các bảng tra cứu trên FPGA, cho phép suy luận với độ trễ nano-giây. Hơn nữa, các tính chất độc đáo của B-splines giúp việc học trực tuyến thời gian thực trở nên khả thi và ổn định trên phần cứng.

Trong tương lai, việc khám phá các thuộc tính của KANs trên các bộ tăng tốc phần cứng tùy chỉnh sẽ tiếp tục là một lĩnh vực nghiên cứu đầy hứa hẹn, đặc biệt là khi các ứng dụng AI đòi hỏi tốc độ phản ứng ngày càng khắt khe hơn.