TorchKM: A GPU-Oriented Library for Kernel Learning and Model Selection

Abstract: TorchKM is an open-source library for kernel machines, including support vector machines, kernel logistic regression, and kernel quantile regression, with GPU acceleration. The library features a scikit-learn-style API and is designed to exploit GPU-friendly linear algebra, accelerating the full training and model-selection pipeline through intelligent reuse of matrix operations. Benchmarks show competitive predictive performance together with substantial speedups over standard baselines. Code and documentation are available at https://github.com/YikaiZhang95/torchkm, and the package can be easily installed via PyPI.

• The paper introduces TorchKM, a GPU-oriented library that fuses exact cross-validation with a spectral algorithm to drastically reduce kernel computation from O(n^3) to efficient GPU-parallel operations.
• It demonstrates significant improvements in both speed and accuracy on large-scale datasets, outperforming traditional libraries like scikit-learn and ThunderSVM.
• The design integrates a scikit-learn-style API with built-in hyperparameter tuning and Nyström approximations, making the library scalable and practical for real-world applications.

TorchKM: A GPU-Oriented Library for Efficient Kernel Learning and Model Selection

Introduction

TorchKM introduces a principled approach to the acceleration of kernel machines—spanning SVM, kernel logistic regression, and kernel quantile regression—via algorithm–hardware co-design that leverages GPU-optimized linear algebra. Traditional kernel machine frameworks suffer from severe computational overhead, particularly during hyperparameter tuning and model selection, which necessitates repeated cubic-cost kernel solves across folds and parameter grids. Prior GPU-accelerated libraries such as ThunderSVM focus primarily on individual model fitting, leaving out-of-core cross-validation and hyperparameter sweeps as bottlenecks. TorchKM addresses these limitations by fusing exact cross-validation and spectral algorithms within a unified GPU-oriented workflow, optimizing the end-to-end pipeline rather than isolated steps.

Core Algorithmic Innovations

TorchKM’s architecture is founded on two key computational strategies: the exact cross-validation formula and the spectral algorithm for regularization path computation. The exact cross-validation strategy modifies the response vector for each fold instead of recomputing the kernel matrix, thus enabling re-use of matrix operations and reducing redundancy. The spectral algorithm performs a single eigendecomposition of the kernel matrix, subsequently allowing model solutions across a regularization path via efficient matrix-vector multiplications that scale as $O(n^2)$, post-decomposition.

This architecture directly targets the main bottleneck in kernel learning—the matrix inversion of $\mathbf{H}_\lambda(\mathbf{K})$ in SVM training which has $O(n^3)$ cost. TorchKM replaces repeated matrix solves with a single decomposition followed by lower-complexity operations, fully exploiting GPU parallelism.

Figure 1: Run time for standard SVM training and tuning under CPU and GPU computation on simulated data sets as compared to TorchKM. TorchKM achieves more than two orders of magnitude speedup on $n=10,000$, demonstrating the necessity of algorithm-hardware co-design.

Library Architecture and API Design

TorchKM adopts a scikit-learn-style interface, enabling seamless integration into existing Python workflows and lowering adoption barriers for practitioners. The library abstracts device management, permitting computations on either CPU or CUDA-enabled GPU transparently. Its API supports SVM, DWD, kernel logistic regression, and kernel quantile regression within the same pipeline, with hyperparameter search integrated into model fitting. Instead of external refitting across parameter grids and folds, TorchKM’s algorithmic design computes all candidate parameter solutions in a single pass using warm starts and Nesterov’s acceleration.

For large-scale datasets, TorchKM implements a Nystr\"om approximation that compresses the kernel matrix using representative subsamples, further alleviating computation and memory constraints. Nystr\"om is executed outside the parameter and cross-validation loops, allowing efficient reuse and consistency across models.

Numerical Performance and Empirical Results

Benchmarks demonstrate that TorchKM reduces end-to-end run times by over an order of magnitude relative to scikit-learn and ThunderSVM, while consistently attaining lower objective values under identical train/test splits and cross-validation protocols. Most notably, TorchKM completes full training and tuning for $n=20,000$, $p=1,000$ in 129.3 seconds, compared to scikit-learn which failed to terminate within an eight-hour constraint.

TorchKM outperforms ThunderSVM and scikit-learn not only in speed but also in accuracy. On large-scale datasets like MNIST8m (4 vs. 6), TorchKM achieves 0.997 accuracy and a wall-clock run time of 64.65 seconds with Nystr\"om, compared to scikit-learn’s 0.996 accuracy at 5189.03 seconds—an efficiency gain exceeding 80-fold. This substantial speed-up is consistent across datasets, underscoring TorchKM’s scalability.

Probability Calibration and Reliability

TorchKM implements Platt scaling for calibrated probability prediction. Reliability curves on independent test sets show an expected calibration error (ECE) of 0.044 and Brier score of 0.130, reflecting high-quality probabilistic outputs.

Figure 2: Reliability curve for Platt-calibrated probabilities, showing the agreement between predicted probability and empirical outcome frequencies. ECE = 0.044, Brier score = 0.130, indicating well-calibrated predictions.

Practical and Theoretical Implications

Practically, TorchKM empowers kernel machines to be competitive alternatives to gradient-boosted trees and deep models on structured tabular data, thanks to efficient model selection and scalable approximation strategies. The spectral and cross-validation innovations enable kernel methods to retain full predictive potential by allowing thorough hyperparameter exploration without prohibitive computational cost.

Theoretically, TorchKM’s approach is significant for the study of algorithm–hardware co-design: efficiency is achieved not merely by porting computations to GPU, but by restructuring algorithms to create reusable matrix operations that exploit hardware architecture. This stands as a paradigm for developing future libraries in statistical learning and optimization.

Future Directions

In terms of research and development, native multiclass support for kernel machines (e.g., multicategory DWD) is prioritized as a future extension. Further, TorchKM's framework can be adapted to emerging parallel hardware and distributed environments, and its spectral algorithm may be generalized for other convex kernel-based objectives.

Conclusion

TorchKM is a rigorously engineered, open-source library that fundamentally advances kernel learning through GPU-oriented algorithms and integrated model-selection pipelines. It maintains exact training and tuning solutions, surpasses standard baselines in efficiency and predictive accuracy, and sets the stage for scalable kernel methods in modern applications. The algorithm-hardware co-design principle demonstrated by TorchKM is likely to influence future AI infrastructure and accelerate broader adoption of kernel-based statistical models.