FPGA 上的极速机器学习：通过 Kolmogorov-Arnold Networks (KAN) 实现

本文是对作者硕士论文的高层解读，该论文涉及为使用 Kolmogorov-Arnold Network (KAN) 架构进行极速推理和在线学习而设计硬件架构。文章假设读者熟悉标准机器学习概念，并具备一些硬件和数字电路的基础知识。

背景

为什么在 FPGA 上进行机器学习？

大多数现代机器学习工作负载（无论是训练还是推理）都运行在图形处理单元 (GPU) 上。通过支持高度并行执行模型的硬件架构，GPU 能够以极高的吞吐量对大量数据执行简单操作。这使得它们非常适合涉及大型架构或批量式训练和推理的机器学习问题。

然而，复杂的 GPU 架构无法满足对超低延迟（例如亚微秒级延迟）和高硬件效率有严格要求的应用需求。处理器（如 CPU 和 GPU）在调度、优化指令以及动态访问内存等方面会产生显著的开销。对于具有超低延迟（例如纳秒级）和极高效率要求的极端专业化工作负载，定制硬件加速器是更好的选择。

现场可编程门阵列 (FPGA) 是可重新配置的数字逻辑设备，非常适合这种定制硬件加速风格。FPGA 包含查找表 (LUTs)，通过枚举所有二进制输入组合的输出值来表示数字函数；触发器 (FFs)，用于存储状态；以及其他内存和计算原语。通过重新配置这些组件及其之间的连接，可以设计定制的数字电路，从而实现底层硬件架构与算法的协同设计， enabling 极速机器学习。重要的是，神经网络是直接作为数字逻辑实现的，而不是作为在处理器上顺序执行的指令。

定点量化 (Fixed-point Quantization)

FPGA 和其他数字设备根本基于位（bits）而非连续值进行操作。然而，我们通常将神经网络中的算术运算（如 $\times, +$）视为在实数 $\mathbb R$ 上发生。因此，我们需要将实数编码为位串（bitstrings），这一过程称为量化。加法等运算和乘法随后变为二进制函数。

一种方法是定点量化。定点量化以二进制表示数字，其中某些位（小数位）位于小数点之后。例如，如果我们总共使用 8 位，小数点后有 4 位小数，我们可以表示从 $(-2^7) / 2^4 = -8$ 到 $(2^7 - 1) / 2^4 = 7.9375$ 的 $2^8$ 个值，这些值以 $1/2^4 = 0.0625$ 的增量均匀分布。在此我们假设可表示的范围关于零对称。

在定点量化方案中，我们只能表示固定范围内的一组离散值，这在尝试表示实值时会导致近似误差。资源高效的机器学习的一个重点是最小化这种近似误差（或量化误差），以实现稳定的训练和推理。

基于查找表的神经网络 (LUT-NNs)

FPGA 主要通过查找表 (LUTs) 实现数字逻辑，LUTs 是小型组件，通过存储每个二进制输入组合的输出来表示任意二进制函数。例如，$\text{AND} : {0, 1}^2 \to {0, 1}$ 用一个查找表表示。

因此，将这些表示为查找表的二进制函数作为神经网络的核心原语进行学习的想法是合理的：这样的网络被称为基于查找表的神经网络 (LUT-NN)。然而，通过梯度下降或类似方法学习查找表是很困难的。

为了解决这个问题，回想一下我们可以通过梯度下降学习实值函数 $f: \mathbb R \to \mathbb R$。如果我们对 $b_i$ 输入位和 $b_o$ 输出位执行定点量化，$f$ 就变成了一个二进制函数 $f_Q : {0, 1}^{b_i} \to {0, 1}^{b_o}$。我们可以学习连续的 $f$ 并量化以获得所需的查找表！

要将 $f$ 转换为 LUT，我们将函数域和值域离散化为 $N_i = 2^{b_i}$ 和 $N_o = 2^{b_o}$ 个值。$f_Q$ 的查找表存储每个输入值 $I \in {I_0, I_1, \ldots, I_{N_i-1}}$ 对应的输出。

上述示例函数 $f$（其中 $q_{l-1}$ 和 $q_{l}$ 代表输入和输出的量化）产生了一个带有 LUT 的二进制函数 $f_Q$。

我们也可以将此方法扩展到将多元函数表示为查找表，其中 $f_m: \mathbb{R}^{d_i} \to \mathbb{R}^{d_o}$ 将变成一个具有 $d_i b_i$ 输入位和 $d_o b_o$ 输出位的二进制函数。在这个查找表中，我们为 $2^{d_i b_i}$ 种可能的输入组合中的每一种存储大小为 $d_o b_o$ 的条目。

这对于任何合理的 $d_i$ 来说显然是不切实际的。因此，基于 LUT 的网络将较小的查找表与算术运算相结合，以实现具有表现力、资源高效且易于训练的架构。

核心内容

Kolmogorov-Arnold Networks (KANs)

KAN 用可学习的激活函数替换了多层感知机 (MLP) 架构中的可学习权重和固定激活函数。在本研究中，我们证明 KAN 是高效、表现力强的 LUT-NN 的自然架构。

在 KAN 层中，每条边携带一个可学习的单变量函数，而不是标量权重。对于一个具有 $n_{\mathrm{in}}$ 个输入和 $n_{\mathrm{out}}$ 个输出的 KAN 层，第 $q$ 个输出的激活值为：

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

其中 $\phi_{q,p} \colon \mathbb{R} \to \mathbb{R}$ 是学习的边激活函数。与使用固定 $\sigma$ 的 MLP（其公式为 $y_q = \sigma\left( \sum_{p=1}^{n_{\mathrm{in}}} w_{q,p} x_p + b_q \right)$）相比，KAN 将非线性放在边函数 ${ \phi_{q,p} }$ 中，并保持节点操作为简单的求和。

下一个问题是如何学习 KAN 激活函数 ${ \phi_{q,p} }$。为此，我们将它们参数化为某些函数基的线性组合：

[\phi_{q,p}(x) = \sum_{i=1}^n c_{q,p,i}B_i(x)]

这允许我们将系数 $c_{q,p,i}$ 视为梯度下降的可训练参数。原始 KAN 论文使用 B 样条 (B-splines)，它们形成一个多项式基，既是平滑的又是局部的，即对于任何给定输入，只有基函数的一部分是非零的。此外，B 样条（因此也是激活函数 ${ \phi_{q,p} }$）定义在一个小的有限域上（例如 $[-1, 1]$），这被发现非常重要。

术语说明： B 样条指的是基函数 ${B_i}$，而激活函数指的是学习的 $\phi_{q,p}(x) = \sum_{i=1}^n c_{q,p,i}B_i(x)$。

虽然 KAN 架构的完整行为尚未完全探索，但它相比 MLP 在缩放定律、参数效率和可解释性方面提供了潜在的改进。对于极速机器学习，前两个特性尤其相关。

KAN 作为可训练的 LUT-NN

方法

我们第一篇论文的关键思想是利用 KAN 作为原理，将可学习的单变量激活函数（由 B 样条基组成）直接映射到 FPGA 的查找表 (LUT) 上。

1. 局部性利用 (Locality Exploitation): B 样条的一个关键特性是它们的局部性。对于给定的输入 $x$，只有少数几个基函数 $B_i(x)$ 是非零的。这意味着在量化和实现为 LUT 时，我们不需要为每个可能的输入组合存储完整的查找表，而是可以针对每个激活函数单独处理其非零区间。这