KAN 模型的 FPGA 加速器探索

前言

最近我开始关注一种新的神经网络结构：Kolmogorov-Arnold Network，也就是 KAN。

和传统的多层感知机 MLP 不同，KAN 的一个重要特点是：它不再只是在节点上使用固定的激活函数，而是把可学习的函数放在边上。换句话说，KAN 中的“连接”不再只是一个简单的权重，而可以表示成一个可学习的一元函数。

这个特点让 KAN 在模型表达能力和可解释性方面具有一定吸引力，但同时也带来了新的计算问题：如果每一条边都对应一个可学习函数，那么推理时就不再只是简单的乘加运算，还需要对这些函数进行计算或查表。

这让我产生了一个想法：

KAN 是否适合在 FPGA 上进行硬件加速？

这篇文章主要记录我对 KAN FPGA 加速器的初步理解、设计思路和后续计划。

KAN 的基本计算特点

在传统 MLP 中，一个神经元的计算通常可以写成：

\[y = \sigma \left( \sum_i w_i x_i + b \right)\]

其中，$w_i$ 是权重，$x_i$ 是输入，$\sigma$ 是激活函数。

而在 KAN 中，每条边上不再只是一个标量权重，而是一个可学习的一元函数。可以简单理解为：

\[y_j = \sum_i \phi_{i,j}(x_i)\]

其中，$\phi_{i,j}$ 表示从输入 $x_i$ 到输出 $y_j$ 的边函数。

也就是说，KAN 的核心计算从传统的“乘加”变成了“函数映射 + 求和”。

这对硬件设计来说非常关键，因为它意味着 KAN 的加速器不能只照搬传统神经网络加速器中的 MAC 阵列，而需要重点考虑如何高效实现边函数 $\phi(x)$。

为什么考虑 FPGA 加速 KAN

FPGA 适合做定制化的数据通路设计，尤其适合探索一些还没有完全固定下来的新模型结构。

对于 KAN 来说，我认为 FPGA 加速有几个值得研究的点。

1. 边函数计算可以定制化

KAN 中的边函数通常可以用样条函数、查找表或者分段近似的方式表示。

如果使用 FPGA，可以根据具体的函数形式设计专门的数据通路，比如：

• LUT 查表
• 分段线性近似
• B-spline 基函数计算
• 定点量化计算
• 流水线求和结构

相比通用处理器，FPGA 可以更灵活地针对这些计算进行优化。

2. 适合探索低精度计算

KAN 的函数计算可能不一定需要高精度浮点数。如果能把部分计算转成定点数，就可以降低资源消耗，提高吞吐率。

这和近似计算、低功耗设计也有一定联系。

例如可以探索：

• 8-bit / 16-bit 定点量化
• 函数值查表量化
• 样条系数量化
• 输入区间离散化
• 误差与资源之间的权衡

3. 可以研究新的 AI 加速器结构

传统 AI 加速器很多都围绕矩阵乘法和卷积计算展开，而 KAN 的计算模式有所不同。

这意味着 KAN 加速器可能需要新的结构设计，比如：

• 边函数计算单元
• 函数参数存储结构
• 多输入并行查表结构
• 分段函数流水线
• 面向 KAN layer 的调度方式

这对 AI 加速器设计来说是一个很有意思的方向。

初步硬件架构设想

目前我对 KAN FPGA 加速器的初步设想如下：

```mermaid flowchart TD A[Input Buffer] –> B[Function Evaluation Unit] B –> C[Partial Sum Accumulator] C –> D[Output Buffer]

E[Function Parameter Memory] --> B
F[Control Unit] --> A
F --> B
F --> C
F --> D