TorchKM：面向GPU的核学习与模型选择库

# TorchKM：面向 GPU 的核学习与模型选择库

来源：https://arxiv.org/html/2606.06742
\\name张一凯\\emailyikai\-zhang@uiowa\.edu \\addr爱荷华大学统计与精算科学系，爱荷华城，IA 52242，美国\\name丁杰\\emaildingj@umn\.edu \\addr明尼苏达大学统计学院，明尼阿波利斯，MN 55455，美国\\name王博祥\\emailboxiang\-wang@uiowa\.edu \\addr爱荷华大学统计与精算科学系，爱荷华城，IA 52242，美国

###### 摘要

TorchKM 是一个面向核机器（包括支持向量机、核逻辑回归和核分位数回归）的开源库，支持 GPU 加速。该库采用 scikit-learn 风格的 API，并充分利用 GPU 友好的线性代数，通过智能复用矩阵运算来加速完整的训练和模型选择流程。基准测试显示，其预测性能具有竞争力，同时相比标准基线方法实现了显著的加速。代码和文档可在 https://github.com/YikaiZhang95/torchkm 获取，并且该包可通过 PyPI 轻松安装。

关键词：算法-硬件协同设计、交叉验证、GPU 加速、核机器、模型选择、支持向量机

## 1 引言

核机器（包括支持向量机，SVM，Cortes 和 Vapnik，1995 (https://arxiv.org/html/2606.06742#bib.bib3)）是一类基础的机器学习算法。它们以其强大的预测性能、具有唯一全局解的凸优化框架以及严格的理论基础而著称。核机器在结构化数据应用中仍然具有很强的竞争力，并在许多领域持续定义着最先进水平，这一点在 Cervantes 等人（2020 (https://arxiv.org/html/2606.06742#bib.bib1)）的近期全面评估中得到了证实。

然而，在现代实践中，核机器往往被置于机器学习工具箱的“后座”，主要障碍在于计算成本。例如，SVM 的性能对超参数选择非常敏感，但要确定一个好的配置可能代价高昂。经典实现，如 LIBSVM（Chang 和 Lin，2011 (https://arxiv.org/html/2606.06742#bib.bib2)）、scikit-learn（Pedregosa 等人，2011 (https://arxiv.org/html/2606.06742#bib.bib11)）和 kernlab（Karatzoglou 等人，2004 (https://arxiv.org/html/2606.06742#bib.bib6)），遵循一种标准范式，其中训练和模型选择被视为独立阶段。尽管通过专门算法（如序列最小优化，Platt，1999 (https://arxiv.org/html/2606.06742#bib.bib12)）可以使每个调优参数下的模型拟合变得高效，但调优通常仍通过在外循环中重复拟合参数网格来进行，这通常需要大量的计算成本。因此，核方法在实践中常常调优不足，从而可能无法充分发挥其预测潜力。

降低计算成本的一种有前景的方法是利用现代计算硬件（尤其是图形处理单元 GPU 加速）的进步（Nickolls 等人，2008 (https://arxiv.org/html/2606.06742#bib.bib10)）。在这方面的一个重要突破是 ThunderSVM（Wen 等人，2018b (https://arxiv.org/html/2606.06742#bib.bib19)；Jiang 等人，2021 (https://arxiv.org/html/2606.06742#bib.bib5)），它显著加速了 GPU 上的核 SVM 训练。尽管取得了这一重要进展，但端到端的核机器工作流程可以进一步从超越单次模型拟合的加速中受益，因为跨交叉验证折和调优参数的重复矩阵运算通常主导了整体计算成本。因此，减少这种调优开销对于在核机器计算中充分利用 GPU 加速至关重要。

在本文中，我们介绍 TorchKM，这是一个面向 GPU 的开源库，用于高效的端到端核学习，并集成了模型选择。TorchKM 支持多种回归和分类方法，包括分位数回归、逻辑回归、SVM 和距离加权判别（DWD）。它提供了用户友好的 scikit-learn 风格接口。TorchKM 并非直接将现有的 CPU 求解器移植到 GPU，而是通过算法-硬件协同设计构建：其计算算法旨在利用 GPU 友好的线性代数。其核心是两个关键的算法思想：精确交叉验证公式和谱算法，它们共同加速了完整的训练和调优流程，并且仍然产生精确解而非近似解。

表 1：代表性核学习库的比较。ThunderSVM 是知名的 GPU 加速 SVM 库；TorchKM 专注于集成模型选择流程。TorchKM 与两个代表性库 scikit-learn 和 ThunderSVM 的主要区别总结在表 1 (https://arxiv.org/html/2606.06742#S1.T1) 中。值得注意的是，TorchKM 具有一个集成训练和调优的流水线，这与 ThunderSVM 专注于快速 GPU 加速 SVM 拟合的特点相辅相成。第 2 节描述了 TorchKM 的核心特性，包括 GPU 加速、对除 SVM 之外方法的支持、概率校准以及针对大数据集的可扩展近似。第 3 节接着概述了使这种集成成为可能的精确交叉验证和谱算法。算法细节和更多数值结果在附录中提供。

## 2 包概述和工作流程

TorchKM 采用 scikit-learn 风格的接口，因此已经熟悉 scikit-learn 的用户应该会发现它易于使用。以 SVM 为例，其基本工作流程与 sklearn.svm.SVC 类似：用户实例化 TorchKMSVC，调用 fit 训练 SVM 模型，然后使用 predict 进行预测。

与标准 sklearn.svm.SVC 工作流程的关键区别在于，TorchKMSVC 通过 Cs 接受一系列调优参数，但单个值也是允许的。在标准的 scikit-learn 工作流程中，模型选择通常通过函数（如 GridSearchCV 或 RandomizedSearchCV）在外循环中对候选值进行处理，并为每个值单独进行一次拟合。相比之下，在 TorchKM 中，训练和调优被集成到算法本身中，而不是通过外部重复重新拟合来进行。

TorchKMSVC 的基本用法如下代码片段所示。

import numpy as np

from sklearn.datasets import make_circles

from torchkm.estimators import TorchKMSVC

X, y = make_circles(1200, factor=0.4, noise=0.08, random_state=0)

Cs = np.logspace(3, -3, num=12)

clf = TorchKMSVC(kernel="rbf", Cs=Cs, device='cuda', probability=True)

clf.fit(X, y)

clf.predict(X)

clf.predict_proba(X)

clf.fit(X, y, low_rank=True)

设置 device="cuda" 可启用 GPU 计算，而 device="cpu" 则在无法使用 GPU 时在 CPU 上使用相同的工作流程。对于 SVM，通过 Platt 缩放可通过 predict_proba 获得概率估计。默认情况下，low_rank=False，此时 TorchKM 提供精确的 SVM 解；对于更大的数据集，用户可以选择设置 low_rank=True 以启用 Nyström 方法并获得近似解。

尽管我们以 SVM 为例进行说明，但相同的接口也适用于其他方法，包括用于 DWD 的 TorchKMDWD（Marron 等人，2007 (https://arxiv.org/html/2606.06742#bib.bib8)；Wang 和 Zou，2016 (https://arxiv.org/html/2606.06742#bib.bib14)，2018 (https://arxiv.org/html/2606.06742#bib.bib15)）、用于核逻辑回归的 TorchKMLogit 以及用于核分位数回归的 TorchKMKQR（Koenker 和 Hallock，2001 (https://arxiv.org/html/2606.06742#bib.bib7)；Tang 等人，2026 (https://arxiv.org/html/2606.06742#bib.bib13)）。因此，这个共享 API 为不同的核学习模型提供了一致的工作流程。

clf_DWD = TorchKMDWD(kernel="rbf", Cs=Cs, cv=5, device='cuda')

clf_Logit = TorchKMLogit(kernel="rbf", Cs=Cs, cv=5, device='cuda')

TorchKM 在 MIT 许可下发布，并通过 GitHub 和 PyPI 分发。该包使用 pytest 进行测试，存储库中包含详细的安装说明、教程、API 文档、基准测试说明和贡献指南。

## 3 核心计算算法和基准性能

在核学习中，完整的训练和调优流程计算量很大，因为训练拆分和调优参数都会改变需要求解的线性系统。在进行 KK 折搜索和 LL 个候选值的朴素方法中，这会导致大约 KLKL 次核求解，需要 O(KLn3)O(KLn^3) 次运算。由于这些重复的三次成本运算，仅仅在 GPU 上运行相同的算法可能只会带来微小的改进。（参见附录图 1 中的示例。）因此，TorchKM 背后的核心思想是重新设计算法，将三次成本负担转移到矩阵向量运算上。从而更充分地利用 GPU 的并行架构来加速整个训练和调优流程。

为实现这一点，TorchKM 采用了两种技术，这两种技术都保留了精确解。首先，精确的交叉验证公式通过保持核矩阵不变，仅使用修改后的响应向量来表示每个折。其次，谱算法计算核矩阵的单个特征分解，并通过矩阵向量乘法为不同的调优参数执行所有后续更新。这两种策略共同将上述重复的三次成本运算替换为一次 O(n3)O(n^3) 分解，随后仅需 O(n2)O(n^2) 次运算。完整细节在附录中提供。

表 2：scikit-learn、ThunderSVM 和 TorchKM 的比较，平均超过 50 次独立运行，括号内为标准误差，最佳值以粗体显示。如表 2 (https://arxiv.org/html/2606.06742#S3.T2) 所示，TorchKM 达到了最低的目标值和最短的运行时间。所有方法都使用相同的训练/测试拆分和交叉验证折进行评估。所有目标值都在相同调优参数下根据相同的目标（附录中的公式 (1) (https://arxiv.org/html/2606.06742#A1.E1)）进行评估。时间包括完整的训练和调优流程。在最大的设置中，n=20,000n=20,000 且 p=1,000p=1,000，scikit-learn 甚至在 8 小时限制内未能完成，而 TorchKM 在 129.3 秒内就完成了完整的训练和调优任务。ThunderSVM 在运行时间上相比基于 CPU 的工作流程有所改进。TorchKM 通过集成训练和调优成功进一步提升了性能。

## 4 结论

TorchKM 是一个基于算法-硬件协同设计的开源库，实现了天然适合 GPU 计算的精确交叉验证和谱算法。TorchKM 通过 scikit-learn 风格的接口易于使用，支持 SVM、DWD、逻辑回归和分位数回归。通过减少调优的计算开销，TorchKM 赋能了核学习，并支持其在现代结构化数据应用中持续取得成功。

## 附录

## 附录 A 核心算法

在本附录中，我们介绍库 TorchKM 背后的核心算法：

### A.1 核 SVM

在本节中，我们以核 SVM 为例进行说明。给定训练数据 { (yi,xi) } i=1n\{ (y_i, \mathbf{x}_i) \}_{i=1}^n，核 SVM 可以表述为：

\( \hat{\boldsymbol{\alpha}}, \hat{\beta}_0 \) = argmin \( \boldsymbol{\alpha} \in \mathbb{R}^n, \beta_0 \in \mathbb{R} \) \[ \frac{1}{n} \sum_{i=1}^n \left(1 - y_i (\mathbf{K}_i^\top \boldsymbol{\alpha} + \beta_0) \right)_+ + \lambda \boldsymbol{\alpha}^\top \mathbf{K} \boldsymbol{\alpha} \],  (1)

其中 \( (1-u)_+ = \max\{1-u, 0\} \) 是不可微的 SVM 铰链损失，\( \lambda > 0 \) 是调优参数。

为了解决优化问题 (1) (https://arxiv.org/html/2606.06742#A1.E1)，TorchKM 实现了有限平滑算法 (Wang 和 Zou，2022 (https://arxiv.org/html/2606.06742#bib.bib17))，该算法将铰链损失函数转化为一系列平滑优化问题，使用 \( \delta \) 平滑铰链损失：

\[
L_\delta(u) = \begin{cases} 1-u & u \leq 1-\delta, \\ \frac{1}{4\delta}[u-(1+\delta)]^2 & 1-\delta < u < 1+\delta, \\ 0 & u \geq 1+\delta, \end{cases}
\]

并获得精确的 SVM 解。

然后使用近端梯度下降求解平滑后的问题。更新公式为：

\[
\binom{\beta_0^{(k+1)}}{\boldsymbol{\alpha}^{(k+1)}} - \binom{\beta_0^{(k)}}{\boldsymbol{\alpha}^{(k)}} = -\mathbf{H}_\lambda^{-1}(\mathbf{K}) \binom{\mathbf{1}^\top \boldsymbol{z}^{(k)}}{\mathbf{K} \boldsymbol{z}^{(k)} + 2\lambda \mathbf{K} \boldsymbol{\alpha}^{(k)}},

\]

其中 \( \mathbf{H}_\lambda(\mathbf{K}) = 2\lambda \mathbf{K} + \frac{1}{n\kappa} \mathbf{K} \mathbf{K} \) 并且 \( \boldsymbol{z}^{(k)} = (z_1, z_2, \cdots, z_n)^\top \) 且每个 \( z_i = y_i L_\delta'[y_i(\beta_0^{(k)} + \mathbf{K}_i \boldsymbol{\alpha}^{(k)})]/n \)。

从上述更新公式可以看出，训练核 SVM 的主要计算瓶颈在于计算 \( \mathbf{H}_\lambda^{-1}(\mathbf{K}) \)，其复杂度为 \( \mathcal{O}(n^3) \)，即使在 GPU 加速下仍然昂贵。此外，用于模型选择的交叉验证以及沿 \( \lambda \) 路径计算解都需要重复拟合模型。

为了说明算法-硬件协同设计的好处，我们进行了一项模拟研究，比较了三种实现下的完整 SVM 训练和调优流程：CPU 上的标准近端梯度求解器、相同求解器在 GPU 上以及 TorchKM。我们考虑了从 Hastie 等人 (2009 (https://arxiv.org/html/2606.06742#bib.bib4)) 的高斯混合模型中生成的模拟数据集。令 \( \mu_+, \mu_- \in \mathbb{R}^p \) 的条目在坐标的不相交一半上等于 \( \mu \)，在其他地方为零。对于每个重复，我们抽样中心 \( \mu_k^+ \sim \mathcal{N}(\mu_+, I) \) 和 \( \mu_k^- \sim \mathcal{N}(\mu_-, I) \) 对于 \( k=1,\ldots,10 \)。然后正例从等权重混合 \( \frac{1}{10} \sum_{k=1}^{10} \mathcal{N}(\mu_k^+, \sigma I) \) 中抽取，负例从 \( \mu_k^- \) 侧的类似混合中抽取。我们设置 \( \mu=2 \)，\( \sigma=3 \)，并考虑 \( n \in \{1000,10000\} \) 且 \( p=10 \)。对于每个数据集，我们计算了 50 个正则化参数下的解路径，并使用 5 折交叉验证进行模型选择。如图 1 (https://arxiv.org/html/2606.06742#A1.F1) 所示，标准求解器的 GPU 加速减少了运行时间，但跨交叉验证折和调优参数的重复计算仍然是一个主要瓶颈。相比之下，TorchKM 通过在整个训练和调优流程中复用矩阵计算，实现了大幅度的加速。当 \( n=10,000 \) 时，TorchKM 快了两个数量级以上。