TorchKM: GPU-Accelerated Kernel Learning

• TorchKM is a GPU-oriented library for kernel methods, integrating training, cross-validation, and model selection into one efficient pipeline.
• It implements kernel SVM, logistic regression, quantile regression, and distance-weighted discrimination using GPU-friendly linear algebra operations.
• The library leverages a single eigen-decomposition and optional Nyström approximation to reduce computational cost and speed up hyperparameter tuning.

Searching arXiv for TorchKM and closely related cited work. Searching for (Zhang et al., 4 Jun 2026). TorchKM is a GPU-oriented Python library for kernel machines with an integrated model-selection pipeline. Introduced in "TorchKM: A GPU-Oriented Library for Kernel Learning and Model Selection" (Zhang et al., 4 Jun 2026), it implements kernel support vector machines, kernel logistic regression, kernel quantile regression, and kernel Distance-Weighted Discrimination, with GPU acceleration via PyTorch/CUDA, a scikit-learn-style API, and optional Nyström approximation for large $n$. Its central design objective is algorithm–hardware co-design: the full train-and-tune workflow is restructured so that computation is dominated by GPU-friendly linear algebra rather than repeated $O(n^3)$ solves across cross-validation folds and regularization parameters.

1. Design objective and computational scope

TorchKM is designed for kernel methods in settings where statistical performance and model-selection rigor are desired, but the computational cost of exact kernel training and tuning is prohibitive (Zhang et al., 4 Jun 2026). The library targets the full workflow rather than a single model fit: training, cross-validation across many candidate regularization values, probability calibration for SVM via Platt scaling, and low-rank kernel approximation are treated as components of one integrated system.

The underlying computational object is the dense kernel, or Gram, matrix $K \in \mathbb{R}^{n \times n}$, where $K_{ij} = k(x_i, x_j)$. TorchKM’s implementation is organized around operations that map naturally to GPU linear algebra: matrix–vector products $Kv$, eigen-decomposition $K = U \Lambda U^\top$, and parallel kernel evaluation. In full-kernel mode, the library computes exact solutions up to numerical precision; in low_rank=True mode, it applies a customized Nyström approximation once outside the tuning and cross-validation loops, so that all subsequent model fits and validation folds reuse the same approximation.

A defining feature is that model selection is internal to the solver. Rather than coupling a standalone estimator to an outer GridSearchCV loop or an external Python loop over $C$ values, TorchKM reuses expensive spectral quantities across all candidate regularization parameters and all validation folds. This suggests that the library is primarily optimized for regimes where the dominant cost is systematic hyperparameter search rather than isolated point estimation.

2. Learning problems and objective functions

TorchKM implements a family of kernelized predictors with common decision-function structure,

$f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$

For classification, the regularized empirical risk has the generic form

$$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$

with analogous regression formulations using $L(y_i - f(x_i))$ (Zhang et al., 4 Jun 2026).

Kernel SVM is given the most detailed treatment. With labels $O(n^3)$0, TorchKM uses a kernelized primal with hinge loss and Tikhonov regularization on $O(n^3)$1: $O(n^3)$2 The paper also gives the usual $O(n^3)$3-parameterization,

$O(n^3)$4

with

$O(n^3)$5

Kernel logistic regression (TorchKMLogit) uses the same kernelized linear decision function but replaces hinge loss with logistic loss,

$O(n^3)$6

leading to

$O(n^3)$7

Because the objective is smooth, the same spectral and cross-validation machinery applies without finite smoothing.

Kernel quantile regression (TorchKMKQR) targets conditional quantiles $O(n^3)$8 via the pinball loss

$O(n^3)$9

with an objective of the form

$K \in \mathbb{R}^{n \times n}$0

This supports regression and quantile estimation rather than classification.

Distance-Weighted Discrimination (TorchKMDWD) is presented as an alternative kernel classifier, often better in high-dimension/low-sample contexts. The paper characterizes it as a logistic-style large-margin objective in kernel form with the same basic decision-function structure and a convex objective with $K \in \mathbb{R}^{n \times n}$1 penalty, while noting that the full details are not spelled out in the same way as for SVM (Zhang et al., 4 Jun 2026).

3. Spectral solver, smoothing, and GPU linear algebra

The library’s core acceleration strategy is a single eigen-decomposition of the kernel matrix,

$K \in \mathbb{R}^{n \times n}$2

computed once per dataset and per kernel hyperparameter. This one-time step has $K \in \mathbb{R}^{n \times n}$3 complexity, but it is then reused across the regularization path and across cross-validation folds, after which the dominant operations become $K \in \mathbb{R}^{n \times n}$4 matrix–vector multiplies (Zhang et al., 4 Jun 2026).

For kernel SVM, the hinge loss is nonsmooth. TorchKM therefore uses the finite smoothing algorithm of Wang and Zou (2022). It introduces a $K \in \mathbb{R}^{n \times n}$5-smoothed hinge loss

$K \in \mathbb{R}^{n \times n}$6

and solves a sequence of smooth problems with decreasing $K \in \mathbb{R}^{n \times n}$7, using proximal gradient with Nesterov acceleration, to converge to the exact SVM solution.

The central linear system is built from

$K \in \mathbb{R}^{n \times n}$8

At iteration $K \in \mathbb{R}^{n \times n}$9, with

$K_{ij} = k(x_i, x_j)$0

the update before spectral simplification is

$K_{ij} = k(x_i, x_j)$1

Naively, recomputing and inverting $K_{ij} = k(x_i, x_j)$2 for each $K_{ij} = k(x_i, x_j)$3 would be the $K_{ij} = k(x_i, x_j)$4 bottleneck.

The spectral formulation removes that bottleneck. With $K_{ij} = k(x_i, x_j)$5, TorchKM defines

$K_{ij} = k(x_i, x_j)$6

and writes $K_{ij} = k(x_i, x_j)$7 in block form as

$K_{ij} = k(x_i, x_j)$8

where

$K_{ij} = k(x_i, x_j)$9

The effect is to reduce each update to a small number of matrix–vector multiplies involving $Kv$0, $Kv$1, and $Kv$2, all of which are well suited to GPU execution. The paper emphasizes that simply moving the same algorithm to a GPU yields limited speedups; the large gains come from this restructuring of the algorithm itself (Zhang et al., 4 Jun 2026).

4. Exact cross-validation and regularization-path model selection

TorchKM’s most distinctive methodological feature is an exact cross-validation reformulation that preserves the full kernel matrix across folds. In classical $Kv$3-fold CV, each fold would require a training problem on a distinct kernel submatrix $Kv$4, implying repeated factorization or inversion of many related but different matrices. TorchKM instead follows a reformulation from Wang and Zou (2022): for fold $Kv$5, define modified responses $Kv$6 by

$Kv$7

Then the fold-specific solution is recovered by solving the full-size problem

$Kv$8

after which the held-out entries are removed to obtain $Kv$9 (Zhang et al., 4 Jun 2026).

Because the same $K = U \Lambda U^\top$0 is used for every fold, the same eigen-decomposition $K = U \Lambda U^\top$1 is also valid for every fold. Cross-validation therefore becomes a matter of modifying $K = U \Lambda U^\top$2, recomputing gradients, and reusing the shared spectral inverse structure. The computational implication is explicit: instead of a naive $K = U \Lambda U^\top$3 pipeline for $K = U \Lambda U^\top$4-fold CV across $K = U \Lambda U^\top$5 candidate regularization parameters, TorchKM performs one $K = U \Lambda U^\top$6 eigen-decomposition followed by $K = U \Lambda U^\top$7 work along the path, with folds handled by shared spectral components.

Hyperparameter tuning is organized around sequences of candidate Cs, typically logarithmically spaced, with internal conversion to $K = U \Lambda U^\top$8 through $K = U \Lambda U^\top$9. Kernel parameters such as RBF width $C$0 are also supported, though kernel-parameter changes remain more expensive because they change the kernel matrix itself. For quantile regression, the quantile level $C$1 is part of the specification; for SVM probability calibration, Platt scaling is fit post hoc using held-out data. Warm starts are used along the regularization path, with the solution at $C$2 initializing the optimization at $C$3.

5. Software interface and implementation

TorchKM follows a scikit-learn-style API and is built on PyTorch; the experiments reported in the paper use PyTorch 2.4.1 with CUDA (Zhang et al., 4 Jun 2026). Device selection is exposed through estimator constructors via device='cuda' or 'cpu', and inputs may be NumPy arrays or PyTorch tensors. The package is open-source, MIT-licensed, available on PyPI and GitHub, and tested with pytest.

The estimator classes explicitly mentioned are TorchKMSVC, TorchKMDWD, TorchKMLogit, and TorchKMKQR. The general constructor pattern includes a kernel argument such as "rbf", a list or array of candidate Cs, a cv argument for the number of folds, a device selector, and optional low-rank approximation. For SVM, probability=True enables Platt scaling and makes predict_proba available. In the usage examples, .fit(X, y) performs integrated training and model selection, .predict(X_new) returns predictions, and low-rank operation is enabled by passing low_rank=True to fit.

This interface places TorchKM near the Python model-selection ecosystem while differing from conventional workflows in one important respect: cross-validation is not an outer orchestration layer but part of the estimator’s numerical core. Relative to scikit-learn, LIBSVM, and ThunderSVM, the distinguishing systems idea is therefore not merely GPU execution but solver-level integration of spectral reuse, cross-validation, and regularization-path traversal.

6. Empirical performance, calibration, limitations, and ecosystem position

The paper benchmarks the full train-and-tune pipeline against scikit-learn SVM with GridSearchCV and against ThunderSVM. On synthetic problems with $C$4 or $C$5, $C$6, 10-fold CV, and 50 candidate regularization values, TorchKM attains lower objective values than both baselines while also being faster (Zhang et al., 4 Jun 2026). Representative entries are especially informative. At $C$7, the reported objective/time pairs are $C$8 and $C$9 s for scikit-learn, $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$0 and $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$1 s for ThunderSVM, and $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$2 and $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$3 s for TorchKM. At $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$4, scikit-learn does not finish within 8 hours, ThunderSVM reports objective $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$5 in $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$6 s, and TorchKM reports $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$7 in $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$8 s. The paper summarizes some of these gains as exceeding $f(x) = \sum_{i=1}^n \alpha_i k(x_i, x) + \beta_0.$9 over scikit-learn.

On medium-size LIBSVM classification datasets with full kernels, TorchKM matches or improves accuracy and is approximately $$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$0–$$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$1 faster than ThunderSVM. The reported results are: a7a, accuracy $$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$2 in $$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$3 s for TorchKM versus $$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$4 in $$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$5 s for ThunderSVM; a8a, $$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$6 in $$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$7 s versus $$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$8 in $$\min_{\alpha,\beta_0} \frac{1}{n}\sum_{i=1}^n L\big(y_i (K_i^\top \alpha + \beta_0)\big) + \lambda \alpha^\top K \alpha,$$9 s; and w7a, $L(y_i - f(x_i))$0 in $L(y_i - f(x_i))$1 s versus $L(y_i - f(x_i))$2 in $L(y_i - f(x_i))$3 s.

With integrated Nyström approximation, the reported gains extend to substantially larger datasets. On a9a, w8a, ijcnn1, covtype, and MNIST8m (4 vs 6, $L(y_i - f(x_i))$4M), TorchKM is more accurate than scikit-learn’s Nyström pipeline on all listed datasets and markedly faster. The most extreme example in the table is MNIST8m, where TorchKM reports accuracy $L(y_i - f(x_i))$5 in $L(y_i - f(x_i))$6 s, while scikit-learn Nyström reports $L(y_i - f(x_i))$7 in $L(y_i - f(x_i))$8 s; the paper characterizes this as roughly an $L(y_i - f(x_i))$9 speedup.

Probability calibration is evaluated via Platt scaling on an RBF SVM with $O(n^3)$00 and $O(n^3)$01. The reliability curve is reported to be close to the $O(n^3)$02 line, with Expected Calibration Error $O(n^3)$03 and Brier score $O(n^3)$04. This indicates that predict_proba is intended as a calibrated output rather than a raw margin transformation.

The current version has explicit limitations. Binary classification only is supported natively; multiclass must be handled externally through one-vs-rest or one-vs-one schemes. Full-kernel operation requires $O(n^3)$05 memory for the dense kernel and its spectral representation, so very large $O(n^3)$06 requires low-rank strategies such as Nyström. The paper also notes limited kernel types relative to scikit-learn’s breadth. Future directions identified in the paper include native multiclass kernel DWD and SVM, further enhancements for very large $O(n^3)$07, and calibration methods beyond Platt scaling.

Within the kernel-learning ecosystem, TorchKM is positioned against scikit-learn, LIBSVM, ThunderSVM, GPyTorch, and lower-level kernel acceleration frameworks such as KeOps. The stated distinction is that TorchKM is a turn-key high-level library for standard kernel machines whose main contribution is integrated training and tuning under GPU-oriented spectral reuse, rather than only fast single-fit SVM training or low-level kernel primitives (Zhang et al., 4 Jun 2026).